planner: fix the issue that statement-level hints in sub-queries of I…

…nsert/Replace can not take effect (#54083) close #53834
pingcap · Jun 18, 2024 · 9906339 · 9906339
1 parent 9008d23
commit 9906339
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/pkg/bindinfo/session_handle_test.go b/pkg/bindinfo/session_handle_test.go
@@ -395,9 +395,13 @@ func TestIssue53834(t *testing.T) {
 	defer func() {
 		tk.MustExec(fmt.Sprintf(`set global tidb_mem_oom_action='%v'`, oomAction))
 	}()
+
 	tk.MustExec(`set global tidb_mem_oom_action='cancel'`)
+	err := tk.ExecToErr(`replace into t select /*+ memory_quota(1 mb) */ * from t`)
+	require.ErrorContains(t, err, "cancelled due to exceeding the allowed memory limit")
+
 	tk.MustExec(`create binding using replace into t select /*+ memory_quota(1 mb) */ * from t`)
-	err := tk.ExecToErr(`replace into t select * from t`)
+	err = tk.ExecToErr(`replace into t select * from t`)
 	require.ErrorContains(t, err, "cancelled due to exceeding the allowed memory limit")
 }
 

diff --git a/pkg/session/test/session_test.go b/pkg/session/test/session_test.go
@@ -410,11 +410,11 @@ func TestStmtHints(t *testing.T) {
 	val = int64(1) * 1024 * 1024
 	require.True(t, tk.Session().GetSessionVars().MemTracker.CheckBytesLimit(val))
 
-	tk.MustExec("insert /*+ MEMORY_QUOTA(1 MB) */  into t1 select /*+ MEMORY_QUOTA(3 MB) */ * from t1;")
+	tk.MustExec("insert /*+ MEMORY_QUOTA(1 MB) */  into t1 select /*+ MEMORY_QUOTA(1 MB) */ * from t1;")
 	val = int64(1) * 1024 * 1024
 	require.True(t, tk.Session().GetSessionVars().MemTracker.CheckBytesLimit(val))
-	require.Len(t, tk.Session().GetSessionVars().StmtCtx.GetWarnings(), 1)
-	require.EqualError(t, tk.Session().GetSessionVars().StmtCtx.GetWarnings()[0].Err, "[planner:3126]Hint MEMORY_QUOTA(`3145728`) is ignored as conflicting/duplicated.")
+	require.Len(t, tk.Session().GetSessionVars().StmtCtx.GetWarnings(), 2)
+	require.EqualError(t, tk.Session().GetSessionVars().StmtCtx.GetWarnings()[0].Err, "[planner:3126]Hint MEMORY_QUOTA(`1048576`) is ignored as conflicting/duplicated.")
 
 	// Test NO_INDEX_MERGE hint
 	tk.Session().GetSessionVars().SetEnableIndexMerge(true)

diff --git a/pkg/util/hint/hint_processor.go b/pkg/util/hint/hint_processor.go
@@ -91,7 +91,18 @@ func ExtractTableHintsFromStmtNode(node ast.Node, warnHandler hintWarnHandler) [
 	case *ast.InsertStmt:
 		// check duplicated hints
 		checkInsertStmtHintDuplicated(node, warnHandler)
-		return x.TableHints
+		result := make([]*ast.TableOptimizerHint, 0, len(x.TableHints))
+		result = append(result, x.TableHints...)
+		if x.Select != nil {
+			// support statement-level hint in sub-select: "insert into t select /* ... */ ..."
+			// TODO: support this for Update and Delete as well
+			for _, h := range ExtractTableHintsFromStmtNode(x.Select, warnHandler) {
+				if isStmtHint(h) {
+					result = append(result, h)
+				}
+			}
+		}
+		return result
 	case *ast.SetOprStmt:
 		var result []*ast.TableOptimizerHint
 		if x.SelectList == nil {