util: optimize the performance of restore with db (#22910)

bindinfo/bind_test.go

@@ -151,7 +151,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
 		}

go.mod

@@ -80,7 +80,7 @@ require (
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/yaml.v2 v2.3.0 // indirect
-	honnef.co/go/tools v0.1.2 // indirect
+	honnef.co/go/tools v0.1.3 // indirect
 	sourcegraph.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0
 	sourcegraph.com/sourcegraph/appdash-data v0.0.0-20151005221446-73f23eafcf67
 )

go.sum

@@ -466,7 +466,6 @@ github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sasha-s/go-deadlock v0.2.0/go.mod h1:StQn567HiB1fF2yJ44N9au7wOhrPS3iZqiDbRupzT10=
 github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
-github.com/sergi/go-diff v1.0.1-0.20180205163309-da645544ed44 h1:tB9NOR21++IjLyVx3/PCPhWMwqGNCMQEH96A6dMZ/gc=
 github.com/sergi/go-diff v1.0.1-0.20180205163309-da645544ed44/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=
 github.com/shirou/gopsutil v2.19.10+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
 github.com/shirou/gopsutil v3.20.12+incompatible h1:6VEGkOXP/eP4o2Ilk8cSsX0PhOEfX6leqAnD+urrp9M=
@@ -846,8 +845,8 @@ honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
-honnef.co/go/tools v0.1.2 h1:SMdYLJl312RXuxXziCCHhRsp/tvct9cGKey0yv95tZM=
-honnef.co/go/tools v0.1.2/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las=
+honnef.co/go/tools v0.1.3 h1:qTakTkI6ni6LFD5sBwwsdSO+AQqbSIxOauHTTQKZ/7o=
+honnef.co/go/tools v0.1.3/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las=
 k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 sigs.k8s.io/yaml v1.1.0 h1:4A07+ZFc2wgJwo8YNlQpr1rVlgUDlxXHhPJciaPY5gs=

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

@@ -290,15 +290,15 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 		}
 		switch x.Stmt.(type) {
 		case *ast.SelectStmt, *ast.DeleteStmt, *ast.UpdateStmt, *ast.InsertStmt:
-			normalizeSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(x.Stmt, specifiledDB))
+			normalizeSQL := parser.Normalize(utilparser.RestoreWithDefaultDB(x.Stmt, specifiledDB, x.Text()))
 			normalizeSQL = plannercore.EraseLastSemicolonInSQL(normalizeSQL)
 			hash := parser.DigestNormalized(normalizeSQL)
 			return x.Stmt, normalizeSQL, hash, nil
 		case *ast.SetOprStmt:
 			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())
 			}
@@ -321,7 +321,7 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 		if len(x.Text()) == 0 {
 			return x, "", "", nil
 		}
-		normalizedSQL, hash := parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(x, specifiledDB))
+		normalizedSQL, hash := parser.NormalizeDigest(utilparser.RestoreWithDefaultDB(x, specifiledDB, x.Text()))
 		return x, normalizedSQL, hash, nil
 	}
 	return nil, "", "", nil

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)
+	}
+}