ddl: fix panic in TablePartitionArgs v2 (pingcap#56222)

ref pingcap#53930
winoros · Sep 23, 2024 · a74a4a1 · a74a4a1
1 parent fe2c538
commit a74a4a1
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/pkg/ddl/tests/tiflash/ddl_tiflash_test.go b/pkg/ddl/tests/tiflash/ddl_tiflash_test.go
@@ -1430,12 +1430,11 @@ func TestTiFlashReorgPartition(t *testing.T) {
 				// Add the tiflash stores as peers for the new regions, to fullfil the check
 				// in checkPartitionReplica
 				pdCli := s.store.(tikv.Storage).GetRegionCache().PDClient()
-				var dummy []pmodel.CIStr
-				partInfo := &model.PartitionInfo{}
-				_ = job.DecodeArgs(&dummy, &partInfo)
+				args, err := model.GetTablePartitionArgs(job)
+				require.NoError(t, err)
 				ctx := context.Background()
 				stores, _ := pdCli.GetAllStores(ctx)
-				for _, pDef := range partInfo.Definitions {
+				for _, pDef := range args.PartInfo.Definitions {
 					startKey, endKey := tablecodec.GetTableHandleKeyRange(pDef.ID)
 					regions, _ := pdCli.BatchScanRegions(ctx, []pd.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1)
 					for i := range regions {

diff --git a/pkg/meta/model/job_args.go b/pkg/meta/model/job_args.go
@@ -434,7 +434,16 @@ func GetTablePartitionArgs(job *Job) (*TablePartitionArgs, error) {
 		}
 		return args, nil
 	}
-	return getOrDecodeArgsV2[*TablePartitionArgs](job)
+	args, err := getOrDecodeArgsV2[*TablePartitionArgs](job)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	// when it's ActionDropTablePartition job, or roll-backing a ActionAddTablePartition
+	// job, our execution part expect a non-nil PartInfo.
+	if args.PartInfo == nil {
+		args.PartInfo = &PartitionInfo{}
+	}
+	return args, nil
 }
 
 // GetFinishedTablePartitionArgs gets the table partition args after the job is finished.

diff --git a/pkg/meta/model/job_args_test.go b/pkg/meta/model/job_args_test.go
@@ -256,11 +256,21 @@ func TestTablePartitionArgs(t *testing.T) {
 				}
 				if j2.Type != ActionDropTablePartition {
 					require.Equal(t, inArgs.PartInfo, args.PartInfo)
+				} else {
+					require.EqualValues(t, &PartitionInfo{}, args.PartInfo)
 				}
 			}
 		}
 	}
 
+	// for ActionDropTablePartition in V2, check PartInfo is not nil
+	j2 := &Job{}
+	require.NoError(t, j2.Decode(getJobBytes(t, &TablePartitionArgs{PartNames: []string{"a", "b"}},
+		JobVersion2, ActionDropTablePartition)))
+	args, err := GetTablePartitionArgs(j2)
+	require.NoError(t, err)
+	require.EqualValues(t, &PartitionInfo{}, args.PartInfo)
+
 	for _, ver := range []JobVersion{JobVersion1, JobVersion2} {
 		j := &Job{
 			Version: ver,
@@ -297,6 +307,7 @@ func TestTablePartitionArgs(t *testing.T) {
 		args, err := GetTablePartitionArgs(j2)
 		require.NoError(t, err)
 		require.EqualValues(t, partNames, args.PartNames)
+		require.EqualValues(t, &PartitionInfo{}, args.PartInfo)
 	}
 }