*: Truncate partition failpoint tests

pkg/ddl/partition.go

@@ -2507,6 +2507,18 @@ func (w *worker) onTruncateTablePartition(jobCtx *jobContext, job *model.Job) (i
 		return ver, errors.Trace(dbterror.ErrPartitionMgmtOnNonpartitioned)
 	}
 
+	if job.IsRollingback() {
+		return convertTruncateTablePartitionJob2RollbackJob(jobCtx, job, dbterror.ErrCancelledDDLJob, tblInfo)
+	}
+
+	failpoint.Inject("truncatePartCancel1", func(val failpoint.Value) {
+		if val.(bool) {
+			job.State = model.JobStateCancelled
+			err = errors.New("Injected error by truncatePartCancel1")
+			failpoint.Return(ver, err)
+		}
+	})
+
 	var oldDefinitions []model.PartitionDefinition
 	var newDefinitions []model.PartitionDefinition
 
@@ -2522,11 +2534,19 @@ func (w *worker) onTruncateTablePartition(jobCtx *jobContext, job *model.Job) (i
 		pi.DDLState = job.SchemaState
 		return updateVersionAndTableInfo(jobCtx, job, tblInfo, true)
 	case model.StateWriteOnly:
+		// We can still rollback here, since we have not yet started to write to the new partitions!
 		oldDefinitions, newDefinitions, err = replaceTruncatePartitions(job, jobCtx.metaMut, tblInfo, oldIDs, newIDs)
 		if err != nil {
 			return ver, errors.Trace(err)
 		}
 		preSplitAndScatter(w.sess.Context, jobCtx.store, tblInfo, newDefinitions)
+		failpoint.Inject("truncatePartFail1", func(val failpoint.Value) {
+			if val.(bool) {
+				job.ErrorCount += variable.GetDDLErrorCountLimit() / 2
+				err = errors.New("Injected error by truncatePartFail1")
+				failpoint.Return(ver, err)
+			}
+		})
 		// This work as a flag to ignore Global Index entries from the old partitions!
 		// Used in IDsInDDLToIgnore() for filtering old partitions from
 		// the global index
@@ -2554,6 +2574,13 @@ func (w *worker) onTruncateTablePartition(jobCtx *jobContext, job *model.Job) (i
 		if err != nil || !done {
 			return ver, errors.Trace(err)
 		}
+		failpoint.Inject("truncatePartFail2", func(val failpoint.Value) {
+			if val.(bool) {
+				job.ErrorCount += variable.GetDDLErrorCountLimit() / 2
+				err = errors.New("Injected error by truncatePartFail2")
+				failpoint.Return(ver, err)
+			}
+		})
 		// For the truncatePartitionEvent
 		oldDefinitions = pi.DroppingDefinitions
 		newDefinitions = make([]model.PartitionDefinition, 0, len(oldIDs))
@@ -2562,21 +2589,25 @@ func (w *worker) onTruncateTablePartition(jobCtx *jobContext, job *model.Job) (i
 			newDef.ID = newIDs[i]
 			newDefinitions = append(newDefinitions, newDef)
 		}
-		// TODO: Test injecting failure
 
 		pi.DroppingDefinitions = nil
 		pi.NewPartitionIDs = nil
 		pi.DDLState = model.StateNone
 		pi.DDLAction = model.ActionNone
 
+		failpoint.Inject("truncatePartFail3", func(val failpoint.Value) {
+			if val.(bool) {
+				job.ErrorCount += variable.GetDDLErrorCountLimit() / 2
+				err = errors.New("Injected error by truncatePartFail3")
+				failpoint.Return(ver, err)
+			}
+		})
 		// used by ApplyDiff in updateSchemaVersion
 		args.ShouldUpdateAffectedPartitions = true
 		ver, err = updateVersionAndTableInfo(jobCtx, job, tblInfo, true)
 		if err != nil {
 			return ver, errors.Trace(err)
 		}
-		// Finish this job.
-		job.FinishTableJob(model.JobStateDone, model.StateNone, ver, tblInfo)
 		truncatePartitionEvent := notifier.NewTruncatePartitionEvent(
 			tblInfo,
 			&model.PartitionInfo{Definitions: newDefinitions},

pkg/ddl/placement_policy_test.go

@@ -1956,6 +1956,7 @@ func TestTruncateTablePartitionWithPlacement(t *testing.T) {
 		" PARTITION `p1` VALUES LESS THAN (1000) /*T![placement] PLACEMENT POLICY=`p2` */,\n" +
 		" PARTITION `p2` VALUES LESS THAN (10000) /*T![placement] PLACEMENT POLICY=`p3` */,\n" +
 		" PARTITION `p3` VALUES LESS THAN (100000))"))
+	dom.Reload()
 	checkExistTableBundlesInPD(t, dom, "test", "tp")
 	checkWaitingGCPartitionBundlesInPD(t, dom, checkOldPartitions)
 

pkg/ddl/reorg.go

@@ -932,8 +932,8 @@ func getReorgInfo(ctx *ReorgContext, jobCtx *jobContext, rh *reorgHandler, job *
 			zap.String("startKey", hex.EncodeToString(start)),
 			zap.String("endKey", hex.EncodeToString(end)))
 
-		failpoint.Inject("errorUpdateReorgHandle", func() (*reorgInfo, error) {
-			return &info, errors.New("occur an error when update reorg handle")
+		failpoint.Inject("errorUpdateReorgHandle", func() {
+			failpoint.Return(&info, errors.New("occur an error when update reorg handle"))
 		})
 		err = rh.InitDDLReorgHandle(job, start, end, pid, elements[0])
 		if err != nil {

pkg/ddl/rollingback.go

@@ -316,6 +316,41 @@ func rollingbackExchangeTablePartition(jobCtx *jobContext, job *model.Job) (ver
 	return ver, errors.Trace(err)
 }
 
+func rollingbackTruncateTablePartition(jobCtx *jobContext, job *model.Job) (ver int64, err error) {
+	tblInfo, err := GetTableInfoAndCancelFaultJob(jobCtx.metaMut, job, job.SchemaID)
+	if err != nil {
+		return ver, errors.Trace(err)
+	}
+
+	return convertTruncateTablePartitionJob2RollbackJob(jobCtx, job, dbterror.ErrCancelledDDLJob, tblInfo)
+}
+
+func convertTruncateTablePartitionJob2RollbackJob(jobCtx *jobContext, job *model.Job, otherwiseErr error, tblInfo *model.TableInfo) (ver int64, err error) {
+	if !job.IsRollbackable() {
+		// Only Original state and StateWrite can be rolled back, otherwise new partitions
+		// may have been used and new data would get lost.
+		// So we must continue to roll forward!
+		job.State = model.JobStateRunning
+		return ver, nil
+	}
+	pi := tblInfo.Partition
+	if len(pi.NewPartitionIDs) != 0 || pi.DDLAction != model.ActionNone || pi.DDLState != model.StateNone {
+		// Rollback the changes, note that no new partitions has been used yet!
+		// so only metadata rollback and we can cancel the DDL
+		tblInfo.Partition.NewPartitionIDs = nil
+		tblInfo.Partition.DDLAction = model.ActionNone
+		tblInfo.Partition.DDLState = model.StateNone
+		ver, err = updateVersionAndTableInfo(jobCtx, job, tblInfo, true)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		return ver, nil
+	}
+	// No change yet, just cancel the job.
+	job.State = model.JobStateCancelled
+	return ver, errors.Trace(otherwiseErr)
+}
+
 func convertAddTablePartitionJob2RollbackJob(jobCtx *jobContext, job *model.Job, otherwiseErr error, tblInfo *model.TableInfo) (ver int64, err error) {
 	addingDefinitions := tblInfo.Partition.AddingDefinitions
 	partNames := make([]string, 0, len(addingDefinitions))
@@ -590,8 +625,10 @@ func convertJob2RollbackJob(w *worker, jobCtx *jobContext, job *model.Job) (ver
 		ver, err = rollingbackTruncateTable(jobCtx, job)
 	case model.ActionModifyColumn:
 		ver, err = rollingbackModifyColumn(jobCtx, job)
-	case model.ActionDropForeignKey, model.ActionTruncateTablePartition:
+	case model.ActionDropForeignKey:
 		ver, err = cancelOnlyNotHandledJob(job, model.StatePublic)
+	case model.ActionTruncateTablePartition:
+		ver, err = rollingbackTruncateTablePartition(jobCtx, job)
 	case model.ActionRebaseAutoID, model.ActionShardRowID, model.ActionAddForeignKey,
 		model.ActionRenameTable, model.ActionRenameTables,
 		model.ActionModifyTableCharsetAndCollate,

pkg/ddl/tests/partition/BUILD.bazel

@@ -5,6 +5,7 @@ go_test(
     timeout = "short",
     srcs = [
         "db_partition_test.go",
+        "error_injection_test.go",
         "main_test.go",
         "multi_domain_test.go",
         "placement_test.go",

pkg/ddl/tests/partition/db_partition_test.go

@@ -3679,15 +3679,18 @@ func TestTruncateNumberOfPhases(t *testing.T) {
 	tk.MustExec(`create table t (a int primary key , b varchar(255)) partition by hash(a) partitions 3`)
 	ctx := tk.Session()
 	dom := domain.GetDomain(ctx)
+	dom.Reload()
 	schemaVersion := dom.InfoSchema().SchemaMetaVersion()
 	tk.MustExec(`insert into t values (1,1),(2,2),(3,3)`)
 	tk.MustExec(`alter table t truncate partition p1`)
-	// Without global index, truncate partition should be a single state change
+	dom.Reload()
+	// Without global index, truncate partition could be a single state change
 	require.Equal(t, int64(4), dom.InfoSchema().SchemaMetaVersion()-schemaVersion)
 	tk.MustExec(`drop table t`)
 	tk.MustExec(`create table t (a int primary key , b varchar(255), unique key (b) global) partition by hash(a) partitions 3`)
 	schemaVersion = dom.InfoSchema().SchemaMetaVersion()
 	tk.MustExec(`insert into t values (1,1),(2,2),(3,3)`)
 	tk.MustExec(`alter table t truncate partition p1`)
+	dom.Reload()
 	require.Equal(t, int64(4), dom.InfoSchema().SchemaMetaVersion()-schemaVersion)
 }

pkg/ddl/tests/partition/error_injection_test.go

@@ -0,0 +1,211 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package partition
+
+import (
+	"testing"
+
+	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/pkg/ddl"
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/external"
+	"github.com/stretchr/testify/require"
+)
+
+type InjectedTest struct {
+	Name        string
+	Recoverable bool
+	Rollback    bool
+}
+
+type FailureTest struct {
+	FailpointPrefix string
+	Tests           []InjectedTest
+}
+
+var truncateTests = FailureTest{
+	FailpointPrefix: "truncatePart",
+	Tests: []InjectedTest{
+		{
+			Name:        "Cancel1",
+			Recoverable: false,
+			Rollback:    true,
+		},
+		{
+			Name:        "Fail1",
+			Recoverable: true,
+			Rollback:    true,
+		},
+		{
+			Name:        "Fail2",
+			Recoverable: true,
+			Rollback:    false,
+		},
+		{
+			Name:        "Fail3",
+			Recoverable: true,
+			Rollback:    false,
+		},
+	},
+}
+
+func TestTruncatePartitionListFailuresWithGlobalIndex(t *testing.T) {
+	create := `create table t (a int unsigned primary key nonclustered global, b int not null, c varchar(255), unique index (c) global) partition by list(b) (
+                        partition p0 values in (1,2,3),
+                        partition p1 values in (4,5,6),
+                        partition p2 values in (7,8,9))`
+	alter := `alter table t truncate partition p0,p2`
+	beforeDML := []string{
+		`insert into t values (1,1,1),(2,2,2),(4,4,4),(8,8,8),(9,9,9),(6,6,6)`,
+		`update t set a = 7, b = 7, c = 7 where a = 1`,
+		`update t set b = 3, c = 3 where c = 4`,
+		`delete from t where a = 8`,
+		`delete from t where b = 2`,
+	}
+	beforeResult := testkit.Rows("4 3 3", "6 6 6", "7 7 7", "9 9 9")
+	afterDML := []string{
+		`insert into t values (1,1,1),(5,5,5),(8,8,8)`,
+		`update t set a = 2, b = 2, c = 2 where a = 1`,
+		`update t set a = 1, b = 1, c = 1 where c = 6`,
+		`update t set a = 6, b = 6 where a = 9`,
+		`delete from t where a = 5`,
+		`delete from t where b = 3`,
+	}
+	afterResult := testkit.Rows("1 1 1", "2 2 2", "6 6 9", "7 7 7", "8 8 8")
+	afterRecover := testkit.Rows("1 1 1", "2 2 2", "8 8 8")
+	testDDLWithInjectedErrors(t, truncateTests, create, alter, beforeDML, beforeResult, afterDML, afterResult, afterRecover, "Cancel2")
+}
+
+func TestTruncatePartitionListFailures(t *testing.T) {
+	create := `create table t (a int unsigned primary key, b int not null, c varchar(255)) partition by list(a) (
+                        partition p0 values in (1,2,3),
+                        partition p1 values in (4,5,6),
+                        partition p2 values in (7,8,9))`
+	alter := `alter table t truncate partition p0,p2`
+	beforeDML := []string{
+		`insert into t values (1,1,1),(2,2,2),(4,4,4),(8,8,8),(9,9,9),(6,6,6)`,
+		`update t set a = 7, b = 7, c = 7 where a = 1`,
+		`update t set b = 3, c = 3, a = 3 where c = 4`,
+		`delete from t where a = 8`,
+		`delete from t where b = 2`,
+	}
+	beforeResult := testkit.Rows("3 3 3", "6 6 6", "7 7 7", "9 9 9")
+	afterDML := []string{
+		`insert into t values (1,1,1),(5,5,5),(8,8,8)`,
+		`update t set a = 2, b = 2, c = 2 where a = 1`,
+		`update t set a = 1, b = 1, c = 1 where c = 6`,
+		`update t set a = 6, b = 6, c = 6 where a = 9`,
+		`delete from t where a = 5`,
+		`delete from t where b = 3`,
+	}
+	afterResult := testkit.Rows("1 1 1", "2 2 2", "6 6 6", "7 7 7", "8 8 8")
+	afterRecover := testkit.Rows("1 1 1", "2 2 2", "8 8 8")
+	testDDLWithInjectedErrors(t, truncateTests, create, alter, beforeDML, beforeResult, afterDML, afterResult, afterRecover, "Fail1", "Fail2", "Fail3")
+}
+
+func testDDLWithInjectedErrors(t *testing.T, tests FailureTest, createSQL, alterSQL string, beforeDML []string, beforeResult [][]any, afterDML []string, afterRollback, afterRecover [][]any, skipTests ...string) {
+TEST:
+	for _, test := range tests.Tests {
+		for _, skip := range skipTests {
+			if test.Name == skip {
+				continue TEST
+			}
+		}
+		if test.Recoverable {
+			runOneTest(t, test, true, tests.FailpointPrefix, createSQL, alterSQL, beforeDML, beforeResult, afterDML, afterRecover)
+		}
+		if test.Rollback {
+			runOneTest(t, test, false, tests.FailpointPrefix, createSQL, alterSQL, beforeDML, beforeResult, afterDML, afterRollback)
+		}
+	}
+}
+
+func runOneTest(t *testing.T, test InjectedTest, recoverable bool, failpointName, createSQL, alterSQL string, beforeDML []string, beforeResult [][]any, afterDML []string, afterResult [][]any) {
+	name := failpointName + test.Name
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("set tidb_enable_global_index=true")
+	defer func() {
+		tk.MustExec("set tidb_enable_global_index=default")
+	}()
+	oldWaitTimeWhenErrorOccurred := ddl.WaitTimeWhenErrorOccurred
+	defer func() {
+		ddl.WaitTimeWhenErrorOccurred = oldWaitTimeWhenErrorOccurred
+	}()
+	ddl.WaitTimeWhenErrorOccurred = 0
+	tk.MustExec(createSQL)
+	for _, sql := range beforeDML {
+		tk.MustExec(sql + ` /* ` + name + ` */`)
+	}
+	tk.MustQuery(`select * from t /* ` + name + ` */`).Sort().Check(beforeResult)
+	tOrg := external.GetTableByName(t, tk, "test", "t")
+	idxIDs := make([]int64, 0, len(tOrg.Meta().Indices))
+	for _, idx := range tOrg.Meta().Indices {
+		idxIDs = append(idxIDs, idx.ID)
+	}
+	pids := make([]int64, 0, len(tOrg.Meta().Partition.Definitions))
+	for _, def := range tOrg.Meta().Partition.Definitions {
+		pids = append(pids, def.ID)
+	}
+	oldCreate := tk.MustQuery(`show create table t`).Rows()
+	fullName := "github.com/pingcap/tidb/pkg/ddl/" + name
+	term := "return(true)"
+	if recoverable {
+		// test that it should handle recover/retry on error
+		term = "1*return(true)"
+	}
+	require.NoError(t, failpoint.Enable(fullName, term))
+	err := tk.ExecToErr(alterSQL + " /* " + name + " */")
+	require.NoError(t, failpoint.Disable(fullName))
+	tt := external.GetTableByName(t, tk, "test", "t")
+	pi := tt.Meta().Partition
+	if recoverable {
+		require.NoError(t, err)
+		equal := true
+		for i, pid := range pids {
+			equal = equal && pid == pi.Definitions[i].ID
+		}
+		require.False(t, equal, name)
+		return
+	}
+	require.Error(t, err, "failpoint "+name)
+	require.ErrorContains(t, err, "Injected error by "+name)
+	tk.MustQuery(`show create table t /* ` + name + ` */`).Check(oldCreate)
+	require.Equal(t, len(tOrg.Meta().Partition.Definitions), len(pi.Definitions), name)
+	require.Equal(t, 0, len(pi.AddingDefinitions), name)
+	require.Equal(t, 0, len(pi.DroppingDefinitions), name)
+	require.Equal(t, 0, len(pi.NewPartitionIDs), name)
+	require.Equal(t, len(tOrg.Meta().Indices), len(tt.Meta().Indices), name)
+	for i := range tOrg.Meta().Indices {
+		require.Equal(t, idxIDs[i], tt.Meta().Indices[i].ID, name)
+	}
+	for i, pid := range pids {
+		require.Equal(t, pid, tt.Meta().Partition.Definitions[i].ID, name)
+	}
+	tk.MustExec(`admin check table t /* ` + name + ` */`)
+	tk.MustExec(`update t set b = 7 where a = 9 /* ` + name + ` */`)
+	for _, sql := range afterDML {
+		tk.MustExec(sql + " /* " + name + " */")
+	}
+	tk.MustQuery(`select * from t /* ` + name + ` */`).Sort().Check(afterResult)
+	tk.MustExec(`drop table t /* ` + name + ` */`)
+	// TODO: Check no rows on new partitions
+	// TODO: Check TiFlash replicas
+	// TODO: Check Label rules
+	// TODO: Check bundles
+	// TODO: Check autoIDs
+	// TODO: Check delete_range tables, so no delete request for old partitions in failed alters!
+}