From 01a45732ad5e122ca829484fbdf1402bb4165ee0 Mon Sep 17 00:00:00 2001
From: Yuanjia Zhang <zhangyuanjia@pingcap.com>
Date: Tue, 18 Jun 2024 13:19:16 +0800
Subject: [PATCH] planner: fix the issue that bindings with query-level hint
 can not take effect for replace statements (#54048)

close pingcap/tidb#53834
---
 pkg/bindinfo/session_handle_test.go | 20 ++++++++++++++++++++
 pkg/planner/optimize.go             |  2 +-
 pkg/util/hint/hint.go               | 10 ++++++++++
 pkg/util/hint/hint_processor.go     | 16 ++++++++++++----
 4 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/pkg/bindinfo/session_handle_test.go b/pkg/bindinfo/session_handle_test.go
index ebbe4bd6471fc..ae653ef17297a 100644
--- a/pkg/bindinfo/session_handle_test.go
+++ b/pkg/bindinfo/session_handle_test.go
@@ -16,6 +16,7 @@ package bindinfo_test
 
 import (
 	"context"
+	"fmt"
 	"strconv"
 	"testing"
 	"time"
@@ -381,6 +382,25 @@ func TestDropSingleBindings(t *testing.T) {
 	tk.MustExec("drop table t")
 }
 
+func TestIssue53834(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec(`create table t (a varchar(1024))`)
+	tk.MustExec(`insert into t values (space(1024))`)
+	for i := 0; i < 12; i++ {
+		tk.MustExec(`insert into t select * from t`)
+	}
+	oomAction := tk.MustQuery(`select @@tidb_mem_oom_action`).Rows()[0][0].(string)
+	defer func() {
+		tk.MustExec(fmt.Sprintf(`set global tidb_mem_oom_action='%v'`, oomAction))
+	}()
+	tk.MustExec(`set global tidb_mem_oom_action='cancel'`)
+	tk.MustExec(`create binding using replace into t select /*+ memory_quota(1 mb) */ * from t`)
+	err := tk.ExecToErr(`replace into t select * from t`)
+	require.ErrorContains(t, err, "cancelled due to exceeding the allowed memory limit")
+}
+
 func TestPreparedStmt(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)
diff --git a/pkg/planner/optimize.go b/pkg/planner/optimize.go
index 0f18f94b6f7ab..0f1a005b13df9 100644
--- a/pkg/planner/optimize.go
+++ b/pkg/planner/optimize.go
@@ -294,7 +294,7 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 				core.DebugTraceTryBinding(pctx, binding.Hint)
 			}
 			hint.BindHint(stmtNode, binding.Hint)
-			curStmtHints, _, curWarns := hint.ParseStmtHints(binding.Hint.GetFirstTableHints(), setVarHintChecker, byte(kv.ReplicaReadFollower))
+			curStmtHints, _, curWarns := hint.ParseStmtHints(binding.Hint.GetStmtHints(), setVarHintChecker, byte(kv.ReplicaReadFollower))
 			sessVars.StmtCtx.StmtHints = curStmtHints
 			// update session var by hint /set_var/
 			for name, val := range sessVars.StmtCtx.StmtHints.SetVars {
diff --git a/pkg/util/hint/hint.go b/pkg/util/hint/hint.go
index e940b6f12b5c5..e1683f1f90cfe 100644
--- a/pkg/util/hint/hint.go
+++ b/pkg/util/hint/hint.go
@@ -447,6 +447,16 @@ func ParseStmtHints(hints []*ast.TableOptimizerHint,
 	return
 }
 
+// isStmtHint checks whether this hint is a statement-level hint.
+func isStmtHint(h *ast.TableOptimizerHint) bool {
+	switch h.HintName.L {
+	case "max_execution_time", "memory_quota", "resource_group":
+		return true
+	default:
+		return false
+	}
+}
+
 // IndexJoinHints stores hint information about index nested loop join.
 type IndexJoinHints struct {
 	INLJTables  []HintedTable
diff --git a/pkg/util/hint/hint_processor.go b/pkg/util/hint/hint_processor.go
index 728d0b06f0f3a..dff5712c76f62 100644
--- a/pkg/util/hint/hint_processor.go
+++ b/pkg/util/hint/hint_processor.go
@@ -39,12 +39,20 @@ type HintsSet struct {
 	indexHints [][]*ast.IndexHint          // Slice offset is the traversal order of `TableName` in the ast.
 }
 
-// GetFirstTableHints gets the first table hints.
-func (hs *HintsSet) GetFirstTableHints() []*ast.TableOptimizerHint {
+// GetStmtHints gets all statement-level hints.
+func (hs *HintsSet) GetStmtHints() []*ast.TableOptimizerHint {
+	var result []*ast.TableOptimizerHint
 	if len(hs.tableHints) > 0 {
-		return hs.tableHints[0]
+		result = append(result, hs.tableHints[0]...) // keep the same behavior with prior implementation
 	}
-	return nil
+	for _, tHints := range hs.tableHints[1:] {
+		for _, h := range tHints {
+			if isStmtHint(h) {
+				result = append(result, h)
+			}
+		}
+	}
+	return result
 }
 
 // ContainTableHint checks whether the table hint set contains a hint.