Repository: pingcap/tidb
File: pkg/executor/softdelete_test.go (Go)

Code before:
		tk.MustExec("insert into t_soft values" + insert)

		tk.MustExec("set @@tidb_translate_softdelete_sql = 'ON'")
		tk.MustExec("delete from t_soft where id in " + deleteIn)
		if rewriteMode {
			// only delete normal table in rewrite mode, and if is not rewrite,
			// skip deleting normal table to keep the same row count in tidb_translate_softdelete_sql=0
			tk.MustExec("delete from t_normal where id in " + deleteIn)
		} else {
			tk.MustExec("set @@tidb_translate_softdelete_sql = 'OFF'")
		}
	}

	type stmtResult struct {
		affectedRows uint64
		message      string
	}

	setupLoadDataReader := func(data string) func() {
		sctx := tk.Session().(sessionctx.Context)
		var readerBuilder executor.LoadDataReaderBuilder = func(_ string) (io.ReadCloser, error) {
			return mydump.NewStringReader(data), nil
		}
		sctx.SetValue(executor.LoadDataReaderBuilderKey, readerBuilder)
		return func() {
			sctx.SetValue(executor.LoadDataReaderBuilderKey, nil)
		}
	}

	mustExecDML := func(t *testing.T, sql string, expect stmtResult) {
		require.Contains(t, sql, "{{t}}", "the sql must contain {{t}} placeholder")
		normalSQL := strings.ReplaceAll(sql, "{{t}}", "t_normal")
		tk.MustExec(normalSQL)
		normalResult := stmtResult{
			affectedRows: tk.Session().AffectedRows(),
			message:      tk.Session().GetSessionVars().StmtCtx.GetMessage(),
		}
		require.Equal(t, expect, normalResult)

		softSQL := strings.ReplaceAll(sql, "{{t}}", "t_soft")
		tk.MustExec(softSQL)
		softResult := stmtResult{
			affectedRows: tk.Session().AffectedRows(),
			message:      tk.Session().GetSessionVars().StmtCtx.GetMessage(),
		}
		require.Equal(t, normalResult, softResult)
	}

	t.Run("insert", func(t *testing.T) {
		prepare("(1, 10), (2, 20)", "(2)", true)
		mustExecDML(t, "INSERT INTO {{t}} VALUES (2, 200), (3, 300)", stmtResult{

Review comment (suggestion):
`setupLoadDataReader` overwrites `LoadDataReaderBuilderKey` and the cleanup sets it to `nil` rather than restoring the previous value. To avoid cross-test/session leakage (especially if future tests set a non-nil builder), consider saving the old value and restoring it in the returned cleanup function.
```suggestion
		oldValue := sctx.Value(executor.LoadDataReaderBuilderKey)
		var readerBuilder executor.LoadDataReaderBuilder = func(_ string) (io.ReadCloser, error) {
			return mydump.NewStringReader(data), nil
		}
		sctx.SetValue(executor.LoadDataReaderBuilderKey, readerBuilder)
		return func() {
			sctx.SetValue(executor.LoadDataReaderBuilderKey, oldValue)
```

Code after:
		tk.MustExec("insert into t_soft values" + insert)

		tk.MustExec("set @@tidb_translate_softdelete_sql = 'ON'")
		tk.MustExec("delete from t_soft where id in " + deleteIn)
		if rewriteMode {
			// In rewrite mode, also delete from the normal table. When not in rewrite mode,
			// we keep tidb_translate_softdelete_sql set to 'OFF' to preserve the row count behavior.
			tk.MustExec("delete from t_normal where id in " + deleteIn)
		} else {
			tk.MustExec("set @@tidb_translate_softdelete_sql = 'OFF'")
		}
	}

	type stmtResult struct {
		affectedRows uint64
		message      string
	}

	setupLoadDataReader := func(data string) func() {
		sctx := tk.Session().(sessionctx.Context)
		var readerBuilder executor.LoadDataReaderBuilder = func(_ string) (io.ReadCloser, error) {
			return mydump.NewStringReader(data), nil
		}
		sctx.SetValue(executor.LoadDataReaderBuilderKey, readerBuilder)
		return func() {
			sctx.SetValue(executor.LoadDataReaderBuilderKey, nil)
		}
	}

	mustExecDML := func(t *testing.T, sql string, expect stmtResult) {
		require.Contains(t, sql, "{{t}}", "the sql must contain {{t}} placeholder")
		normalSQL := strings.ReplaceAll(sql, "{{t}}", "t_normal")
		tk.MustExec(normalSQL)
		normalResult := stmtResult{
			affectedRows: tk.Session().AffectedRows(),
			message:      tk.Session().GetSessionVars().StmtCtx.GetMessage(),
		}
		require.Equal(t, expect, normalResult)

		softSQL := strings.ReplaceAll(sql, "{{t}}", "t_soft")
		tk.MustExec(softSQL)
		softResult := stmtResult{
			affectedRows: tk.Session().AffectedRows(),
			message:      tk.Session().GetSessionVars().StmtCtx.GetMessage(),
		}
		require.Equal(t, normalResult, softResult)
	}

	t.Run("insert", func(t *testing.T) {
		prepare("(1, 10), (2, 20)", "(2)", true)
		mustExecDML(t, "INSERT INTO {{t}} VALUES (2, 200), (3, 300)", stmtResult{