util: optimize the performance of restore with db

bindinfo/bind_test.go

@@ -148,7 +148,7 @@ func normalizeWithDefaultDB(c *C, sql, db string) (string, string) {
 	testParser := parser.New()
 	stmt, err := testParser.ParseOneStmt(sql, "", "")
 	c.Assert(err, IsNil)
-	return parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(stmt, "test"))
+	return parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(stmt, "test", ""))
 }
 
 func (s *testSuite) TestBindParse(c *C) {

bindinfo/handle.go

@@ -624,7 +624,7 @@ func (h *BindHandle) CaptureBaselines() {
 			continue
 		}
 		dbName := utilparser.GetDefaultDB(stmt, bindableStmt.Schema)
-		normalizedSQL, digest := parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(stmt, dbName))
+		normalizedSQL, digest := parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(stmt, dbName, bindableStmt.Query))
 		if r := h.GetBindRecord(digest, normalizedSQL, dbName); r != nil && r.HasUsingBinding() {
 			continue
 		}

planner/core/planbuilder.go

@@ -756,12 +756,12 @@ func (b *PlanBuilder) buildSet(ctx context.Context, v *ast.SetStmt) (Plan, error
 func (b *PlanBuilder) buildDropBindPlan(v *ast.DropBindingStmt) (Plan, error) {
 	p := &SQLBindPlan{
 		SQLBindOp:    OpSQLBindDrop,
-		NormdOrigSQL: parser.Normalize(utilparser.RestoreWithDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB)),
+		NormdOrigSQL: parser.Normalize(utilparser.RestoreWithDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB, v.OriginNode.Text())),
 		IsGlobal:     v.GlobalScope,
 		Db:           utilparser.GetDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB),
 	}
 	if v.HintedNode != nil {
-		p.BindSQL = utilparser.RestoreWithDefaultDB(v.HintedNode, b.ctx.GetSessionVars().CurrentDB)
+		p.BindSQL = utilparser.RestoreWithDefaultDB(v.HintedNode, b.ctx.GetSessionVars().CurrentDB, v.HintedNode.Text())
 	}
 	b.visitInfo = appendVisitInfo(b.visitInfo, mysql.SuperPriv, "", "", "", nil)
 	return p, nil
@@ -800,8 +800,8 @@ func (b *PlanBuilder) buildCreateBindPlan(v *ast.CreateBindingStmt) (Plan, error
 
 	p := &SQLBindPlan{
 		SQLBindOp:    OpSQLBindCreate,
-		NormdOrigSQL: parser.Normalize(utilparser.RestoreWithDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB)),
-		BindSQL:      utilparser.RestoreWithDefaultDB(v.HintedNode, b.ctx.GetSessionVars().CurrentDB),
+		NormdOrigSQL: parser.Normalize(utilparser.RestoreWithDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB, v.OriginNode.Text())),
+		BindSQL:      utilparser.RestoreWithDefaultDB(v.HintedNode, b.ctx.GetSessionVars().CurrentDB, v.HintedNode.Text()),
 		IsGlobal:     v.GlobalScope,
 		BindStmt:     v.HintedNode,
 		Db:           utilparser.GetDefaultDB(v.OriginNode, b.ctx.GetSessionVars().CurrentDB),

planner/core/preprocess.go

@@ -315,8 +315,8 @@ func (p *preprocessor) checkBindGrammar(originNode, hintedNode ast.StmtNode, def
 			return
 		}
 	}
-	originSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(originNode, defaultDB))
-	hintedSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(hintedNode, defaultDB))
+	originSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(originNode, defaultDB, originNode.Text()))
+	hintedSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(hintedNode, defaultDB, hintedNode.Text()))
 	if originSQL != hintedSQL {
 		p.err = errors.Errorf("hinted sql and origin sql don't match when hinted sql erase the hint info, after erase hint info, originSQL:%s, hintedSQL:%s", originSQL, hintedSQL)
 	}

planner/optimize.go

@@ -292,7 +292,7 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 		case *ast.SelectStmt, *ast.DeleteStmt, *ast.UpdateStmt, *ast.InsertStmt:
 			var normalizeSQL string
 			if specifiledDB != "" {
-				normalizeSQL = parser.Normalize(utilparser.RestoreWithDefaultDB(x.Stmt, specifiledDB))
+				normalizeSQL = parser.Normalize(utilparser.RestoreWithDefaultDB(x.Stmt, specifiledDB, x.Text()))
 			} else {
 				normalizeSQL = parser.Normalize(x.Text())
 			}
@@ -303,7 +303,7 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 			plannercore.EraseLastSemicolon(x)
 			var normalizeExplainSQL string
 			if specifiledDB != "" {
-				normalizeExplainSQL = parser.Normalize(utilparser.RestoreWithDefaultDB(x, specifiledDB))
+				normalizeExplainSQL = parser.Normalize(utilparser.RestoreWithDefaultDB(x, specifiledDB, x.Text()))
 			} else {
 				normalizeExplainSQL = parser.Normalize(x.Text())
 			}
@@ -328,7 +328,7 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 		}
 		var normalizedSQL, hash string
 		if specifiledDB != "" {
-			normalizedSQL, hash = parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(x, specifiledDB))
+			normalizedSQL, hash = parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(x, specifiledDB, x.Text()))
 		} else {
 			normalizedSQL, hash = parser.NormalizeDigest(x.Text())
 		}

session/bootstrap.go

@@ -1389,7 +1389,7 @@ func updateBindInfo(iter *chunk.Iterator4Chunk, p *parser.Parser, bindMap map[st
 		if err != nil {
 			logutil.BgLogger().Fatal("updateBindInfo error", zap.Error(err))
 		}
-		originWithDB := parser.Normalize(utilparser.RestoreWithDefaultDB(stmt, db))
+		originWithDB := parser.Normalize(utilparser.RestoreWithDefaultDB(stmt, db, bind))
 		if _, ok := bindMap[originWithDB]; ok {
 			// The results are sorted in descending order of time.
 			// And in the following cases, duplicate originWithDB may occur
@@ -1400,7 +1400,7 @@ func updateBindInfo(iter *chunk.Iterator4Chunk, p *parser.Parser, bindMap map[st
 			continue
 		}
 		bindMap[originWithDB] = bindInfo{
-			bindSQL:    utilparser.RestoreWithDefaultDB(stmt, db),
+			bindSQL:    utilparser.RestoreWithDefaultDB(stmt, db, bind),
 			status:     row.GetString(2),
 			createTime: row.GetTime(3),
 			charset:    charset,

util/parser/ast.go

@@ -49,8 +49,76 @@ func (i *implicitDatabase) Leave(in ast.Node) (out ast.Node, ok bool) {
 	return in, true
 }
 
+func findTablePos(s, t string) int {
+	l := 0
+	for i := range s {
+		if s[i] == ' ' || s[i] == ',' {
+			if len(t) == i-l && strings.Compare(s[l:i], t) == 0 {
+				return l
+			}
+			l = i + 1
+		}
+	}
+	if len(t) == len(s)-l && strings.Compare(s[l:], t) == 0 {
+		return l
+	}
+	return -1
+}
+
+// SimpleCases captures simple SQL statements and uses string replacement instead of `restore` to improve performance.
+// See https://github.com/pingcap/tidb/issues/22398.
+func SimpleCases(node ast.StmtNode, defaultDB, origin string) (s string, ok bool) {
+	if len(origin) == 0 {
+		return "", false
+	}
+	insert, ok := node.(*ast.InsertStmt)
+	if !ok {
+		return "", false
+	}
+	if insert.Select != nil || insert.Setlist != nil || insert.OnDuplicate != nil || (insert.TableHints != nil && len(insert.TableHints) != 0) {
+		return "", false
+	}
+	join := insert.Table.TableRefs
+	if join.Tp != 0 || join.Right != nil {
+		return "", false
+	}
+	ts, ok := join.Left.(*ast.TableSource)
+	if !ok {
+		return "", false
+	}
+	tn, ok := ts.Source.(*ast.TableName)
+	if !ok {
+		return "", false
+	}
+	parenPos := strings.Index(origin, "(")
+	if parenPos == -1 {
+		return "", false
+	}
+	if strings.Contains(origin[:parenPos], ".") {
+		return origin, true
+	}
+	lower := strings.ToLower(origin[:parenPos])
+	pos := findTablePos(lower, tn.Name.L)
+	if pos == -1 {
+		return "", false
+	}
+	var builder strings.Builder
+	builder.WriteString(origin[:pos])
+	if tn.Schema.String() != "" {
+		builder.WriteString(tn.Schema.String())
+	} else {
+		builder.WriteString(defaultDB)
+	}
+	builder.WriteString(".")
+	builder.WriteString(origin[pos:])
+	return builder.String(), true
+}
+
 // RestoreWithDefaultDB returns restore strings for StmtNode with defaultDB
-func RestoreWithDefaultDB(node ast.StmtNode, defaultDB string) string {
+func RestoreWithDefaultDB(node ast.StmtNode, defaultDB, origin string) string {
+	if s, ok := SimpleCases(node, defaultDB, origin); ok {
+		return s
+	}
 	var sb strings.Builder
 	// Three flags for restore with default DB:
 	// 1. RestoreStringSingleQuotes specifies to use single quotes to surround the string;

util/parser/ast_test.go

@@ -0,0 +1,70 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package parser_test
+
+import (
+	"testing"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/parser"
+	_ "github.com/pingcap/tidb/types/parser_driver"
+	utilparser "github.com/pingcap/tidb/util/parser"
+)
+
+var _ = Suite(&testASTSuite{})
+
+type testASTSuite struct {
+}
+
+func TestT(t *testing.T) {
+	TestingT(t)
+}
+
+func (s *testASTSuite) TestSimpleCases(c *C) {
+	tests := []struct {
+		sql string
+		db  string
+		ans string
+	}{
+		{
+			sql: "insert into t values(1, 2)",
+			db:  "test",
+			ans: "insert into test.t values(1, 2)",
+		},
+		{
+			sql: "insert into mydb.t values(1, 2)",
+			db:  "test",
+			ans: "insert into mydb.t values(1, 2)",
+		},
+		{
+			sql: "insert into t(a, b) values(1, 2)",
+			db:  "test",
+			ans: "insert into test.t(a, b) values(1, 2)",
+		},
+		{
+			sql: "insert into value value(2, 3)",
+			db:  "test",
+			ans: "insert into test.value value(2, 3)",
+		},
+	}
+
+	for _, t := range tests {
+		p := parser.New()
+		stmt, err := p.ParseOneStmt(t.sql, "", "")
+		c.Assert(err, IsNil)
+		ans, ok := utilparser.SimpleCases(stmt, t.db, t.sql)
+		c.Assert(ok, IsTrue)
+		c.Assert(t.ans, Equals, ans)
+	}
+}