planner: calculate plan_digest in place when fetching plan exec info from stmt_stats in SHOW PLAN FOR <SQL> (#60273)

ref #60148

Author: qw4990

pkg/bindinfo/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "//pkg/util/chunk",
         "//pkg/util/hack",
         "//pkg/util/hint",
+        "//pkg/util/intest",
         "//pkg/util/logutil",
         "//pkg/util/parser",
         "//pkg/util/sqlexec",

pkg/bindinfo/binding_auto.go

@@ -20,13 +20,18 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/parser"
+	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/intest"
 	utilparser "github.com/pingcap/tidb/pkg/util/parser"
 	"go.uber.org/zap"
 )
 
+// PlanDigestFunc is used to get the plan digest of this SQL.
+var PlanDigestFunc func(sctx sessionctx.Context, stmt ast.StmtNode) (planDigest string, err error)
+
 // BindingPlanInfo contains the binding info and its corresponding plan execution info, which is used by
 // "SHOW PLAN FOR <SQL>" to help users understand the historical plans for a specific SQL.
 type BindingPlanInfo struct {
@@ -90,7 +95,22 @@ func (ba *bindingAuto) ShowPlansForSQL(currentDB, sqlOrDigest, charset, collatio
 	// read plan info from information_schema.tidb_statements_stats
 	bindingPlans := make([]*BindingPlanInfo, 0, len(bindings))
 	for _, binding := range bindings {
-		pInfo, err := ba.getPlanExecInfo(binding.PlanDigest)
+		if binding.Status == StatusDeleted {
+			continue
+		}
+
+		planDigest := binding.PlanDigest
+		if planDigest == "" {
+			if err := callWithSCtx(ba.sPool, false, func(sctx sessionctx.Context) error {
+				planDigest = getBindingPlanDigest(sctx, binding.Db, binding.BindSQL)
+				return nil
+			}); err != nil {
+				bindingLogger().Error("get plan digest failed",
+					zap.String("bind_sql", binding.BindSQL), zap.Error(err))
+			}
+		}
+
+		pInfo, err := ba.getPlanExecInfo(planDigest)
 		if err != nil {
 			bindingLogger().Error("get plan execution info failed", zap.String("plan_digest", binding.PlanDigest), zap.Error(err))
 			continue
@@ -117,9 +137,13 @@ func (ba *bindingAuto) getPlanExecInfo(planDigest string) (plan *planExecInfo, e
 	if planDigest == "" {
 		return nil, nil
 	}
+	stmtStatsTable := "information_schema.cluster_tidb_statements_stats"
+	if intest.InTest { // don't need to access the cluster table in tests.
+		stmtStatsTable = "information_schema.tidb_statements_stats"
+	}
 	stmtQuery := fmt.Sprintf(`select cast(sum(result_rows) as signed), cast(sum(exec_count) as signed),
        cast(sum(processed_keys) as signed), cast(sum(total_time) as signed), any_value(plan)
-       from information_schema.cluster_tidb_statements_stats where plan_digest = '%v'`, planDigest)
+       from %v where plan_digest = '%v'`, stmtStatsTable, planDigest)
 
 	var rows []chunk.Row
 	err = callWithSCtx(ba.sPool, false, func(sctx sessionctx.Context) error {

pkg/bindinfo/utils.go

@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/hint"
 	"github.com/pingcap/tidb/pkg/util/logutil"
+	utilparser "github.com/pingcap/tidb/pkg/util/parser"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"go.uber.org/zap"
 )
@@ -191,3 +192,36 @@ func newBindingFromStorage(row chunk.Row) *Binding {
 		PlanDigest:  row.GetString(10),
 	}
 }
+
+// getBindingPlanDigest does the best efforts to fill binding's plan_digest.
+func getBindingPlanDigest(sctx sessionctx.Context, schema, bindingSQL string) (planDigest string) {
+	defer func() {
+		if r := recover(); r != nil {
+			bindingLogger().Error("panic when filling plan digest for binding",
+				zap.String("binding_sql", bindingSQL), zap.Reflect("panic", r))
+		}
+	}()
+
+	vars := sctx.GetSessionVars()
+	defer func(originalBaseline bool, originalDB string) {
+		vars.UsePlanBaselines = originalBaseline
+		vars.CurrentDB = originalDB
+	}(vars.UsePlanBaselines, vars.CurrentDB)
+	vars.UsePlanBaselines = false
+	vars.CurrentDB = schema
+
+	p := utilparser.GetParser()
+	defer utilparser.DestoryParser(p)
+	p.SetSQLMode(vars.SQLMode)
+	p.SetParserConfig(vars.BuildParserConfig())
+
+	charset, collation := vars.GetCharsetInfo()
+	if stmt, err := p.ParseOneStmt(bindingSQL, charset, collation); err == nil {
+		if !hasParam(stmt) {
+			// if there is '?' from `create binding using select a from t where a=?`,
+			// the final plan digest might be incorrect.
+			planDigest, _ = PlanDigestFunc(sctx, stmt)
+		}
+	}
+	return
+}

pkg/infoschema/test/clustertablestest/cluster_tables_test.go

@@ -1254,6 +1254,41 @@ func TestStmtSummaryShowPlanForSQL(t *testing.T) {
 	require.Equal(t, result[7], "1") // avg_returned_rows
 }
 
+func TestStmtSummaryShowPlanForSQL2(t *testing.T) {
+	s := new(clusterTablesSuite)
+	s.store, s.dom = testkit.CreateMockStoreAndDomain(t)
+	s.rpcserver, s.listenAddr = s.setUpRPCService(t, "127.0.0.1:0", nil)
+	s.httpServer, s.mockAddr = s.setUpMockPDHTTPServer()
+	s.startTime = time.Now()
+	defer s.httpServer.Close()
+	defer s.rpcserver.Stop()
+	tk := s.newTestKitWithRoot(t)
+	require.NoError(t, tk.Session().Auth(&auth.UserIdentity{Username: "root", Hostname: "%"}, nil, nil, nil))
+
+	tk.MustExec("use test")
+	tk.MustExec(`create table t (a int, b int, c varchar(10), key(a), key(b))`)
+
+	tk.MustExec(`create global binding using select /*+ use_index(t, a) */ a from t where b=1`)
+	tk.MustQuery(`select a from test.t where b=1`).Check(testkit.Rows())
+	tk.MustQuery(`select @@last_plan_from_binding`).Check(testkit.Rows("1"))
+
+	rs := tk.MustQuery(`show plan for "select a from test.t where b=1"`).Rows()[0]
+	require.Contains(t, rs[2], "index:a(a)")
+	require.Equal(t, rs[5], "1") // exec_count
+	tk.MustQuery(`select a from test.t where b=1`).Check(testkit.Rows())
+	rs = tk.MustQuery(`show plan for "select a from test.t where b=1"`).Rows()[0]
+	require.Equal(t, rs[5], "2") // exec_count
+
+	tk.MustExec(`create global binding using select /*+ use_index(t, b) */ a from test.t where b=1`)
+	tk.MustQuery(`select a from test.t where b=1`).Check(testkit.Rows())
+	rs = tk.MustQuery(`show plan for "select a from test.t where b=2"`).Rows()[0]
+	require.Contains(t, rs[2], "index:b(b)")
+	require.Equal(t, rs[5], "1") // exec_count
+	tk.MustQuery(`select a from test.t where b=1`).Check(testkit.Rows())
+	rs = tk.MustQuery(`show plan for "select a from test.t where b=2"`).Rows()[0]
+	require.Equal(t, rs[5], "2") // exec_count
+}
+
 func TestCreateBindingFromHistory(t *testing.T) {
 	s := new(clusterTablesSuite)
 	s.store, s.dom = testkit.CreateMockStoreAndDomain(t)

pkg/planner/optimize.go

@@ -594,11 +594,35 @@ func queryPlanCost(sctx sessionctx.Context, stmt ast.StmtNode) (float64, error)
 	return core.GetPlanCost(pp, property.RootTaskType, optimizetrace.NewDefaultPlanCostOption())
 }
 
+func planDigestFunc(sctx sessionctx.Context, stmt ast.StmtNode) (planDigest string, err error) {
+	ret := &core.PreprocessorReturn{}
+	nodeW := resolve.NewNodeW(stmt)
+	err = core.Preprocess(
+		context.Background(),
+		sctx,
+		nodeW,
+		core.WithPreprocessorReturn(ret),
+		core.InitTxnContextProvider,
+	)
+	if err != nil {
+		return "", err
+	}
+
+	p, _, err := Optimize(context.Background(), sctx, nodeW, sctx.GetDomainInfoSchema().(infoschema.InfoSchema))
+	if err != nil {
+		return "", err
+	}
+	flat := core.FlattenPhysicalPlan(p, false)
+	_, digest := core.NormalizeFlatPlan(flat)
+	return digest.String(), nil
+}
+
 func init() {
 	core.OptimizeAstNode = Optimize
 	core.IsReadOnly = IsReadOnly
 	indexadvisor.QueryPlanCostHook = queryPlanCost
 	bindinfo.GetBindingHandle = func(sctx sessionctx.Context) bindinfo.BindingHandle {
 		return domain.GetDomain(sctx).BindingHandle()
 	}
+	bindinfo.PlanDigestFunc = planDigestFunc
 }