planner: add more test cases for auto-capture and support to capture …

…more hints (#33051)

bindinfo/capture_test.go

@@ -943,12 +943,20 @@ func TestCaptureHints(t *testing.T) {
 		{"select * from t use index(b)", "use_index(@`sel_1` `test`.`t` `b`)"},
 		{"select /*+ use_index(b) */ * from t use index(b)", "use_index(@`sel_1` `test`.`t` `b`)"},
 		{"select /*+ use_index_merge(t, a, b) */ a, b from t where a=1 or b=1", "use_index_merge(@`sel_1` `t` `a`, `b`)"},
+		{"select /*+ ignore_index(t, a) */ * from t where a=1", "ignore_index(`t` `a`)"},
 		// push-down hints
+		{"select /*+ limit_to_cop() */ * from t limit 10", "limit_to_cop()"},
+		{"select /*+ agg_to_cop() */ a, count(*) from t group by a", "agg_to_cop()"},
 		// index-merge hints
+		{"select /*+ no_index_merge() */ a, b from t where a>1 or b>1", "no_index_merge()"},
+		{"select /*+ use_index_merge(t, a, b) */ a, b from t where a>1 or b>1", "use_index_merge(@`sel_1` `t` `a`, `b`)"},
 		// runtime hints
+		{"select /*+ memory_quota(1024 MB) */ * from t", "memory_quota(1024 mb)"},
+		{"select /*+ max_execution_time(1000) */ * from t", "max_execution_time(1000)"},
 		// storage hints
-		// query-blocks
+		{"select /*+ read_from_storage(tikv[t]) */ * from t", "read_from_storage(tikv[`t`])"},
 		// others
+		{"select /*+ use_toja(true) */ t1.a, t1.b from t t1 where t1.a in (select t2.a from t t2)", "use_toja(true)"},
 	}
 	for _, capCase := range captureCases {
 		stmtsummary.StmtSummaryByDigestMap.Clear()

executor/adapter.go

@@ -1171,6 +1171,17 @@ func getEncodedPlan(sctx sessionctx.Context, p plannercore.Plan, genHint bool) (
 	}
 	if genHint {
 		hints := plannercore.GenHintsFromPhysicalPlan(p)
+		for _, tableHint := range sctx.GetSessionVars().StmtCtx.OriginalTableHints {
+			// some hints like 'memory_quota' cannot be extracted from the PhysicalPlan directly,
+			// so we have to iterate all hints from the customer and keep some other necessary hints.
+			switch tableHint.HintName.L {
+			case "memory_quota", "use_toja", "no_index_merge", "max_execution_time",
+				plannercore.HintAggToCop, plannercore.HintIgnoreIndex,
+				plannercore.HintReadFromStorage, plannercore.HintLimitToCop:
+				hints = append(hints, tableHint)
+			}
+		}
+
 		hintStr = hint.RestoreOptimizerHints(hints)
 		sctx.GetSessionVars().StmtCtx.SetPlanHint(hintStr)
 	}

planner/optimize.go

@@ -602,6 +602,7 @@ func handleStmtHints(hints []*ast.TableOptimizerHint) (stmtHints stmtctx.StmtHin
 			setVarsOffs = append(setVarsOffs, i)
 		}
 	}
+	stmtHints.OriginalTableHints = hints
 	stmtHints.SetVars = setVars
 
 	// Handle MEMORY_QUOTA

sessionctx/stmtctx/stmtctx.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/parser"
+	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/util/disk"
@@ -263,6 +264,9 @@ type StmtHints struct {
 	HasMaxExecutionTime            bool
 	HasEnableCascadesPlannerHint   bool
 	SetVars                        map[string]string
+
+	// the original table hints
+	OriginalTableHints []*ast.TableOptimizerHint
 }
 
 // TaskMapNeedBackUp indicates that whether we need to back up taskMap during physical optimizing.