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

Code before:
	"github.com/pingcap/tidb/pkg/sessionctx"
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/stretchr/testify/require"
)

func TestSoftDeleteAffectedRowsAndStmtMessage(t *testing.T) {
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	tk.MustExec("drop table if exists t_normal")
	tk.MustExec("drop table if exists t_soft")
	tk.MustExec("create table t_normal (id int primary key, v int)")
	tk.MustExec("create table t_soft (id int primary key, v int) softdelete retention 7 day")

	prepare := func(insert, deleteIn string, rewriteMode bool) {
		tk.MustExec("set @@tidb_translate_softdelete_sql = 'OFF'")
		tk.MustExec("delete from t_normal")
		tk.MustExec("delete from t_soft")
		tk.MustExec("insert into t_normal values" + insert)
		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")

Review comment (suggestion):
The comment says "keep the same row count in tidb_translate_softdelete_sql=0", but the code uses the string values 'ON'/'OFF'. Consider updating the comment to match the actual setting (and fix grammar like "if is not rewrite").
```suggestion
			// 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.
```

Code after:
	"github.com/pingcap/tidb/pkg/sessionctx"
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/stretchr/testify/require"
)

func TestSoftDeleteAffectedRowsAndStmtMessage(t *testing.T) {
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	tk.MustExec("drop table if exists t_normal")
	tk.MustExec("drop table if exists t_soft")
	tk.MustExec("create table t_normal (id int primary key, v int)")
	tk.MustExec("create table t_soft (id int primary key, v int) softdelete retention 7 day")

	prepare := func(insert, deleteIn string, rewriteMode bool) {
		tk.MustExec("set @@tidb_translate_softdelete_sql = 'OFF'")
		tk.MustExec("delete from t_normal")
		tk.MustExec("delete from t_soft")
		tk.MustExec("insert into t_normal values" + insert)
		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")