diff --git a/pkg/session/nontransactional.go b/pkg/session/nontransactional.go index d1a87af50f..83abcb9f09 100644 --- a/pkg/session/nontransactional.go +++ b/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 { diff --git a/pkg/session/test/nontransactionaltest/BUILD.bazel b/pkg/session/test/nontransactionaltest/BUILD.bazel index bb580cab74..92608921fe 100644 --- a/pkg/session/test/nontransactionaltest/BUILD.bazel +++ b/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", diff --git a/pkg/session/test/nontransactionaltest/nontransactional_test.go b/pkg/session/test/nontransactionaltest/nontransactional_test.go index 9c5ee8a4de..9507e9c5aa 100644 --- a/pkg/session/test/nontransactionaltest/nontransactional_test.go +++ b/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") +}