planner: support drop binding by sql digest (#39399)

ref #39199

bindinfo/bind_cache.go

@@ -146,6 +146,23 @@ func (c *bindCache) GetBindRecord(hash, normdOrigSQL, db string) *BindRecord {
 	return nil
 }
 
+// GetBindRecordBySQLDigest gets the BindRecord from the cache.
+// The return value is not read-only, but it shouldn't be changed in the caller functions.
+// The function is thread-safe.
+func (c *bindCache) GetBindRecordBySQLDigest(sqlDigest string) (*BindRecord, error) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	bindings := c.get(bindCacheKey(sqlDigest))
+	if len(bindings) > 1 {
+		// currently, we only allow one binding for a sql
+		return nil, errors.New("more than 1 binding matched")
+	}
+	if len(bindings) == 0 || len(bindings[0].Bindings) == 0 {
+		return nil, errors.New("can't find any binding for `" + sqlDigest + "`")
+	}
+	return bindings[0], nil
+}
+
 // GetAllBindRecords return all the bindRecords from the bindCache.
 // The return value is not read-only, but it shouldn't be changed in the caller functions.
 // The function is thread-safe.

bindinfo/bind_test.go

@@ -1304,3 +1304,87 @@ func TestBindSQLDigest(t *testing.T) {
 		require.Equal(t, res[0][9], sqlDigestWithDB.String())
 	}
 }
+
+func TestDropBindBySQLDigest(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(pk int primary key, a int, b int, key(a), key(b))")
+
+	cases := []struct {
+		origin string
+		hint   string
+	}{
+		// agg hints
+		{"select count(1) from t", "select /*+ hash_agg() */ count(1) from t"},
+		{"select count(1) from t", "select /*+ stream_agg() */ count(1) from t"},
+		// join hints
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ merge_join(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ tidb_smj(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ hash_join(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ tidb_hj(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ inl_join(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ tidb_inlj(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		{"select * from t t1, t t2 where t1.a=t2.a", "select /*+ inl_hash_join(t1, t2) */ * from t t1, t t2 where t1.a=t2.a"},
+		// index hints
+		{"select * from t", "select * from t use index(primary)"},
+		{"select * from t", "select /*+ use_index(primary) */ * from t"},
+		{"select * from t", "select * from t use index(a)"},
+		{"select * from t", "select /*+ use_index(a) */ * from t use index(a)"},
+		{"select * from t", "select * from t use index(b)"},
+		{"select * from t", "select /*+ use_index(b) */ * from t use index(b)"},
+		{"select a, b from t where a=1 or b=1", "select /*+ use_index_merge(t, a, b) */ a, b from t where a=1 or b=1"},
+		{"select * from t where a=1", "select /*+ ignore_index(t, a) */ * from t where a=1"},
+		// push-down hints
+		{"select * from t limit 10", "select /*+ limit_to_cop() */ * from t limit 10"},
+		{"select a, count(*) from t group by a", "select /*+ agg_to_cop() */ a, count(*) from t group by a"},
+		// index-merge hints
+		{"select a, b from t where a>1 or b>1", "select /*+ no_index_merge() */ a, b from t where a>1 or b>1"},
+		{"select a, b from t where a>1 or b>1", "select /*+ use_index_merge(t, a, b) */ a, b from t where a>1 or b>1"},
+		// runtime hints
+		{"select * from t", "select /*+ memory_quota(1024 MB) */ * from t"},
+		{"select * from t", "select /*+ max_execution_time(1000) */ * from t"},
+		// storage hints
+		{"select * from t", "select /*+ read_from_storage(tikv[t]) */ * from t"},
+		// others
+		{"select t1.a, t1.b from t t1 where t1.a in (select t2.a from t t2)", "select /*+ use_toja(true) */ t1.a, t1.b from t t1 where t1.a in (select t2.a from t t2)"},
+	}
+
+	h := dom.BindHandle()
+	// global scope
+	for _, c := range cases {
+		utilCleanBindingEnv(tk, dom)
+		sql := "create global binding for " + c.origin + " using " + c.hint
+		tk.MustExec(sql)
+		h.ReloadBindings()
+		res := tk.MustQuery(`show global bindings`).Rows()
+
+		require.Equal(t, len(res), 1)
+		require.Equal(t, len(res[0]), 11)
+		drop := fmt.Sprintf("drop global binding for sql digest '%s'", res[0][9])
+		tk.MustExec(drop)
+		require.NoError(t, h.GCBindRecord())
+		h.ReloadBindings()
+		tk.MustQuery("show global bindings").Check(testkit.Rows())
+	}
+
+	// session scope
+	for _, c := range cases {
+		utilCleanBindingEnv(tk, dom)
+		sql := "create binding for " + c.origin + " using " + c.hint
+		tk.MustExec(sql)
+		res := tk.MustQuery(`show bindings`).Rows()
+
+		require.Equal(t, len(res), 1)
+		require.Equal(t, len(res[0]), 11)
+		drop := fmt.Sprintf("drop binding for sql digest '%s'", res[0][9])
+		tk.MustExec(drop)
+		require.NoError(t, h.GCBindRecord())
+		tk.MustQuery("show bindings").Check(testkit.Rows())
+	}
+
+	// exception cases
+	tk.MustGetErrMsg(fmt.Sprintf("drop binding for sql digest '%s'", "1"), "can't find any binding for `1`")
+	tk.MustGetErrMsg(fmt.Sprintf("drop binding for sql digest '%s'", ""), "sql digest is empty")
+}

bindinfo/handle.go

@@ -437,6 +437,15 @@ func (h *BindHandle) DropBindRecord(originalSQL, db string, binding *Binding) (d
 	return h.sctx.Context.GetSessionVars().StmtCtx.AffectedRows(), nil
 }
 
+// DropBindRecordByDigest drop BindRecord to the storage and BindRecord int the cache.
+func (h *BindHandle) DropBindRecordByDigest(sqlDigest string) (deletedRows uint64, err error) {
+	oldRecord, err := h.GetBindRecordBySQLDigest(sqlDigest)
+	if err != nil {
+		return 0, err
+	}
+	return h.DropBindRecord(oldRecord.OriginalSQL, strings.ToLower(oldRecord.Db), nil)
+}
+
 // SetBindRecordStatus set a BindRecord's status to the storage and bind cache.
 func (h *BindHandle) SetBindRecordStatus(originalSQL string, binding *Binding, newStatus string) (ok bool, err error) {
 	h.bindInfo.Lock()
@@ -658,6 +667,11 @@ func (h *BindHandle) GetBindRecord(hash, normdOrigSQL, db string) *BindRecord {
 	return h.bindInfo.Load().(*bindCache).GetBindRecord(hash, normdOrigSQL, db)
 }
 
+// GetBindRecordBySQLDigest returns the BindRecord of the sql digest.
+func (h *BindHandle) GetBindRecordBySQLDigest(sqlDigest string) (*BindRecord, error) {
+	return h.bindInfo.Load().(*bindCache).GetBindRecordBySQLDigest(sqlDigest)
+}
+
 // GetAllBindRecord returns all bind records in cache.
 func (h *BindHandle) GetAllBindRecord() (bindRecords []*BindRecord) {
 	return h.bindInfo.Load().(*bindCache).GetAllBindRecords()

bindinfo/session_handle.go

@@ -97,11 +97,25 @@ func (h *SessionHandle) DropBindRecord(originalSQL, db string, binding *Binding)
 	return nil
 }
 
+// DropBindRecordByDigest drop BindRecord in the cache.
+func (h *SessionHandle) DropBindRecordByDigest(sqlDigest string) error {
+	oldRecord, err := h.GetBindRecordBySQLDigest(sqlDigest)
+	if err != nil {
+		return err
+	}
+	return h.DropBindRecord(oldRecord.OriginalSQL, strings.ToLower(oldRecord.Db), nil)
+}
+
 // GetBindRecord return the BindMeta of the (normdOrigSQL,db) if BindMeta exist.
 func (h *SessionHandle) GetBindRecord(hash, normdOrigSQL, db string) *BindRecord {
 	return h.ch.GetBindRecord(hash, normdOrigSQL, db)
 }
 
+// GetBindRecordBySQLDigest return all BindMeta corresponding to sqlDigest.
+func (h *SessionHandle) GetBindRecordBySQLDigest(sqlDigest string) (*BindRecord, error) {
+	return h.ch.GetBindRecordBySQLDigest(sqlDigest)
+}
+
 // GetAllBindRecord return all session bind info.
 func (h *SessionHandle) GetAllBindRecord() (bindRecords []*BindRecord) {
 	return h.ch.GetAllBindRecords()

executor/bind.go

@@ -50,6 +50,8 @@ func (e *SQLBindExec) Next(ctx context.Context, req *chunk.Chunk) error {
 		return e.createSQLBind()
 	case plannercore.OpSQLBindDrop:
 		return e.dropSQLBind()
+	case plannercore.OpSQLBindDropByDigest:
+		return e.dropSQLBindByDigest()
 	case plannercore.OpFlushBindings:
 		return e.flushBindings()
 	case plannercore.OpCaptureBindings:
@@ -85,6 +87,20 @@ func (e *SQLBindExec) dropSQLBind() error {
 	return err
 }
 
+func (e *SQLBindExec) dropSQLBindByDigest() error {
+	if e.sqlDigest == "" {
+		return errors.New("sql digest is empty")
+	}
+	if !e.isGlobal {
+		handle := e.ctx.Value(bindinfo.SessionBindInfoKeyType).(*bindinfo.SessionHandle)
+		err := handle.DropBindRecordByDigest(e.sqlDigest)
+		return err
+	}
+	affectedRows, err := domain.GetDomain(e.ctx).BindHandle().DropBindRecordByDigest(e.sqlDigest)
+	e.ctx.GetSessionVars().StmtCtx.AddAffectedRows(affectedRows)
+	return err
+}
+
 func (e *SQLBindExec) setBindingStatus() error {
 	var bindInfo *bindinfo.Binding
 	if e.bindSQL != "" {

parser/ast/misc.go

@@ -1872,6 +1872,7 @@ type DropBindingStmt struct {
 	GlobalScope bool
 	OriginNode  StmtNode
 	HintedNode  StmtNode
+	SQLDigest   string
 }
 
 func (n *DropBindingStmt) Restore(ctx *format.RestoreCtx) error {
@@ -1882,14 +1883,19 @@ func (n *DropBindingStmt) Restore(ctx *format.RestoreCtx) error {
 		ctx.WriteKeyWord("SESSION ")
 	}
 	ctx.WriteKeyWord("BINDING FOR ")
-	if err := n.OriginNode.Restore(ctx); err != nil {
-		return errors.Trace(err)
-	}
-	if n.HintedNode != nil {
-		ctx.WriteKeyWord(" USING ")
-		if err := n.HintedNode.Restore(ctx); err != nil {
+	if n.OriginNode == nil {
+		ctx.WriteKeyWord("SQL DIGEST ")
+		ctx.WriteString(n.SQLDigest)
+	} else {
+		if err := n.OriginNode.Restore(ctx); err != nil {
 			return errors.Trace(err)
 		}
+		if n.HintedNode != nil {
+			ctx.WriteKeyWord(" USING ")
+			if err := n.HintedNode.Restore(ctx); err != nil {
+				return errors.Trace(err)
+			}
+		}
 	}
 	return nil
 }
@@ -1900,17 +1906,20 @@ func (n *DropBindingStmt) Accept(v Visitor) (Node, bool) {
 		return v.Leave(newNode)
 	}
 	n = newNode.(*DropBindingStmt)
-	origNode, ok := n.OriginNode.Accept(v)
-	if !ok {
-		return n, false
-	}
-	n.OriginNode = origNode.(StmtNode)
-	if n.HintedNode != nil {
-		hintedNode, ok := n.HintedNode.Accept(v)
+	if n.OriginNode != nil {
+		//  OriginNode is nil means we build drop binding by sql digest
+		origNode, ok := n.OriginNode.Accept(v)
 		if !ok {
 			return n, false
 		}
-		n.HintedNode = hintedNode.(StmtNode)
+		n.OriginNode = origNode.(StmtNode)
+		if n.HintedNode != nil {
+			hintedNode, ok := n.HintedNode.Accept(v)
+			if !ok {
+				return n, false
+			}
+			n.HintedNode = hintedNode.(StmtNode)
+		}
 	}
 	return v.Leave(n)
 }

parser/misc.go

@@ -289,6 +289,7 @@ var tokenMap = map[string]int{
 	"DEPTH":                    depth,
 	"DESC":                     desc,
 	"DESCRIBE":                 describe,
+	"DIGEST":                   digest,
 	"DIRECTORY":                directory,
 	"DISABLE":                  disable,
 	"DISABLED":                 disabled,