transaction: batch dml should not be affected by max_execution_time (#64349)

close pingcap/tidb#64282

pkg/session/nontransactional.go

@@ -455,6 +455,12 @@ func buildShardJobs(ctx context.Context, stmt *ast.NonTransactionalDMLStmt, se s
 	// A NT-DML is not a SELECT. We ignore the SelectLimit for selectSQL so that it can read all values.
 	originalSelectLimit := se.GetSessionVars().SelectLimit
 	se.GetSessionVars().SelectLimit = math.MaxUint64
+	originalMaxExecutionTime := se.GetSessionVars().MaxExecutionTime
+	// A NT-DML is not read-only, so we disable max execution time for it.
+	se.GetSessionVars().MaxExecutionTime = 0
+	defer func() {
+		se.GetSessionVars().MaxExecutionTime = originalMaxExecutionTime
+	}()
 	// NT-DML is a write operation, and should not be affected by read_staleness that is supposed to affect only SELECT.
 	rss, err := se.Execute(ctx, selectSQL)
 	se.GetSessionVars().SelectLimit = originalSelectLimit
@@ -529,7 +535,15 @@ func buildShardJobs(ctx context.Context, stmt *ast.NonTransactionalDMLStmt, se s
 		currentStart = *currentStart.Clone()
 	}
 
-	return jobs, nil
+	failpoint.Inject("CheckMaxExecutionTime", func(val failpoint.Value) {
+		if val.(bool) {
+			if se.GetSessionVars().MaxExecutionTime > 0 {
+				err = errors.New("injected max execution time exceeded error")
+			}
+		}
+	})
+
+	return jobs, err
 }
 
 func appendNewJob(jobs []job, id int, start types.Datum, end types.Datum, size int, tracker *memory.Tracker) []job {

pkg/session/test/nontransactionaltest/BUILD.bazel

@@ -8,7 +8,7 @@ go_test(
         "nontransactional_test.go",
     ],
     flaky = True,
-    shard_count = 5,
+    shard_count = 6,
     deps = [
         "//pkg/config",
         "//pkg/metrics",

pkg/session/test/nontransactionaltest/nontransactional_test.go

@@ -509,3 +509,18 @@ func TestNonTransactionalMetrics(t *testing.T) {
 		tk.MustExec("BATCH LIMIT 10 REPLACE INTO t2 SELECT * FROM t1")
 	})
 }
+
+func TestNonTransactionalDmlIgnoreMaxExecutionTime(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("set @@tidb_max_chunk_size=10")
+	tk.MustExec("set @@max_execution_time=1000")
+	tk.MustExec("use test")
+	tk.MustExec("create table t(a int, b int, key(a))")
+	for i := range 100 {
+		tk.MustExec(fmt.Sprintf("insert into t values (%d, %d)", i, i*2))
+	}
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/session/CheckMaxExecutionTime", `return(true)`))
+	defer failpoint.Disable("github.com/pingcap/tidb/pkg/session/CheckMaxExecutionTime")
+	tk.MustExec("batch on a limit 10 update t set b = b + 1 where b > 0")
+}