From 00e96229c13f0ade6d0f000c361d2df6f240b4a6 Mon Sep 17 00:00:00 2001
From: Yuanjia Zhang <zhangyuanjia@pingcap.com>
Date: Tue, 15 Mar 2022 20:01:51 +0800
Subject: [PATCH] planner: add more test cases for auto-capture and support to
 capture more hints (#33051)

---
 bindinfo/capture_test.go      | 10 +++++++++-
 executor/adapter.go           | 11 +++++++++++
 planner/optimize.go           |  1 +
 sessionctx/stmtctx/stmtctx.go |  4 ++++
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/bindinfo/capture_test.go b/bindinfo/capture_test.go
index 316c1b59cd1da..32d38a4cedc8d 100644
--- a/bindinfo/capture_test.go
+++ b/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()
diff --git a/executor/adapter.go b/executor/adapter.go
index 368d31e18560a..e5faf41437960 100644
--- a/executor/adapter.go
+++ b/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)
 	}
diff --git a/planner/optimize.go b/planner/optimize.go
index 0941e89545fb3..a57a71f4d0c16 100644
--- a/planner/optimize.go
+++ b/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
diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go
index fcc166141e6ff..580618ea95d8f 100644
--- a/sessionctx/stmtctx/stmtctx.go
+++ b/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.