Add shardMode support to serverless migration CLI

internal/cli/serverless/migration/create.go

@@ -86,7 +86,7 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 			}
 			definitionStr := string(definitionBytes)
 
-			sources, target, mode, err := parseMigrationDefinition(definitionStr)
+			sources, target, mode, shardMode, err := parseMigrationDefinition(definitionStr)
 			if err != nil {
 				return err
 			}
@@ -97,6 +97,7 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 					Sources:     sources,
 					Target:      target,
 					Mode:        mode,
+					ShardMode:   shardMode,
 				}
 				return runMigrationPrecheck(ctx, d, clusterID, precheckBody, h)
 			}
@@ -106,6 +107,7 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 				Sources:     sources,
 				Target:      target,
 				Mode:        mode,
+				ShardMode:   shardMode,
 			}
 
 			resp, err := d.CreateMigration(ctx, clusterID, createBody)
@@ -247,34 +249,39 @@ func shouldPrintPrecheckItem(status *pkgmigration.PrecheckItemStatus) bool {
 	}
 }
 
-func parseMigrationDefinition(value string) ([]pkgmigration.Source, pkgmigration.Target, pkgmigration.TaskMode, error) {
+func parseMigrationDefinition(value string) ([]pkgmigration.Source, pkgmigration.Target, pkgmigration.TaskMode, *pkgmigration.TaskShardMode, error) {
 	trimmed := strings.TrimSpace(value)
 	if trimmed == "" {
-		return nil, pkgmigration.Target{}, "", errors.New("migration config is required; use --config-file")
+		return nil, pkgmigration.Target{}, "", nil, errors.New("migration config is required; use --config-file")
 	}
 	var payload struct {
-		Sources []pkgmigration.Source `json:"sources"`
-		Target  *pkgmigration.Target  `json:"target"`
-		Mode    string                `json:"mode"`
+		Sources   []pkgmigration.Source `json:"sources"`
+		Target    *pkgmigration.Target  `json:"target"`
+		Mode      string                `json:"mode"`
+		ShardMode *string               `json:"shardMode,omitempty"`
 	}
 	stdJson, err := standardizeJSON([]byte(trimmed))
 	if err != nil {
-		return nil, pkgmigration.Target{}, "", errors.Annotate(err, "invalid migration definition JSON")
+		return nil, pkgmigration.Target{}, "", nil, errors.Annotate(err, "invalid migration definition JSON")
 	}
 	if err := json.Unmarshal(stdJson, &payload); err != nil {
-		return nil, pkgmigration.Target{}, "", errors.Annotate(err, "invalid migration definition JSON")
+		return nil, pkgmigration.Target{}, "", nil, errors.Annotate(err, "invalid migration definition JSON")
 	}
 	if len(payload.Sources) == 0 {
-		return nil, pkgmigration.Target{}, "", errors.New("migration definition must include at least one source")
+		return nil, pkgmigration.Target{}, "", nil, errors.New("migration definition must include at least one source")
 	}
 	if payload.Target == nil {
-		return nil, pkgmigration.Target{}, "", errors.New("migration definition must include the target block")
+		return nil, pkgmigration.Target{}, "", nil, errors.New("migration definition must include the target block")
 	}
 	mode, err := parseMigrationMode(payload.Mode)
 	if err != nil {
-		return nil, pkgmigration.Target{}, "", err
+		return nil, pkgmigration.Target{}, "", nil, err
 	}
-	return payload.Sources, *payload.Target, mode, nil
+	shardMode, err := parseMigrationShardMode(payload.ShardMode)
+	if err != nil {
+		return nil, pkgmigration.Target{}, "", nil, err
+	}
+	return payload.Sources, *payload.Target, mode, shardMode, nil
 }
 
 func parseMigrationMode(value string) (pkgmigration.TaskMode, error) {
@@ -290,6 +297,22 @@ func parseMigrationMode(value string) (pkgmigration.TaskMode, error) {
 	return "", errors.Errorf("invalid mode %q, allowed values: %s", value, pkgmigration.AllowedTaskModeEnumValues)
 }
 
+func parseMigrationShardMode(value *string) (*pkgmigration.TaskShardMode, error) {
+	if value == nil {
+		return nil, nil
+	}
+	trimmed := strings.TrimSpace(*value)
+	if trimmed == "" {
+		return nil, errors.Errorf("invalid shardMode %q, allowed values: %s", *value, pkgmigration.AllowedTaskShardModeEnumValues)
+	}
+	normalized := strings.ToUpper(trimmed)
+	mode := pkgmigration.TaskShardMode(normalized)
+	if slices.Contains(pkgmigration.AllowedTaskShardModeEnumValues, mode) {
+		return &mode, nil
+	}
+	return nil, errors.Errorf("invalid shardMode %q, allowed values: %s", *value, pkgmigration.AllowedTaskShardModeEnumValues)
+}
+
 // standardizeJSON accepts JSON With Commas and Comments(JWCC) see
 // https://nigeltao.github.io/blog/2021/json-with-commas-comments.html) and
 // returns a standard JSON byte slice ready for json.Unmarshal.

internal/cli/serverless/migration/create_test.go

@@ -72,6 +72,8 @@ func (suite *CreateMigrationSuite) TestCreateMigration() {
 			return body != nil &&
 				body.DisplayName == displayName &&
 				body.Mode == pkgmigration.TASKMODE_ALL &&
+				body.ShardMode != nil &&
+				*body.ShardMode == pkgmigration.TASKSHARDMODE_PESSIMISTIC &&
 				len(body.Sources) == 1 &&
 				body.Target.User == "migration_user"
 		}),
@@ -98,6 +100,7 @@ func (suite *CreateMigrationSuite) TestCreateMigrationInvalidInputs() {
 	blankPath := suite.writeTempConfig(" ")
 	invalidJSONPath := suite.writeTempConfig("{invalid")
 	invalidModePath := suite.writeTempConfig(`{ "mode": "invalid", "target": {"user":"u","password":"p"}, "sources": [{"sourceType":"MYSQL","connProfile":{"connType":"PUBLIC","host":"h","port":3306,"user":"u","password":"p"}}] }`)
+	invalidShardModePath := suite.writeTempConfig(`{ "mode": "ALL", "shardMode": "invalid", "target": {"user":"u","password":"p"}, "sources": [{"sourceType":"MYSQL","connProfile":{"connType":"PUBLIC","host":"h","port":3306,"user":"u","password":"p"}}] }`)
 
 	tests := []struct {
 		name        string
@@ -124,6 +127,11 @@ func (suite *CreateMigrationSuite) TestCreateMigrationInvalidInputs() {
 			args:        []string{"--cluster-id", "c1", "--display-name", "name", "--config-file", invalidModePath},
 			errContains: "invalid mode",
 		},
+		{
+			name:        "invalid shard mode",
+			args:        []string{"--cluster-id", "c1", "--display-name", "name", "--config-file", invalidShardModePath},
+			errContains: "invalid shardMode",
+		},
 	}
 
 	for _, tt := range tests {
@@ -152,6 +160,7 @@ func (suite *CreateMigrationSuite) writeTempConfig(content string) string {
 func validMigrationConfig() string {
 	return `{
   "mode": "ALL",
+  "shardMode": "PESSIMISTIC",
   "target": {
     "user": "migration_user",
     "password": "Passw0rd!"

internal/cli/serverless/migration/template.go

@@ -31,6 +31,10 @@ const (
 	migrationDefinitionAllTemplate = `{
     // Required migration mode. Use "ALL" for full + incremental.
     "mode": "ALL",
+    // Optional shard DDL coordination mode. Only set this for shard merge
+    // (sharded database/table) scenarios. See:
+    // https://docs.pingcap.com/tidb/stable/dm-shard-merge/
+    // "shardMode": "PESSIMISTIC",
     // Target TiDB Cloud user credentials used by the migration
     "target": {
         "user": "migration_user",
@@ -106,6 +110,10 @@ const (
 	migrationDefinitionIncrementalTemplate = `{
     // Incremental-only mode keeps the source and target in sync
     "mode": "INCREMENTAL",
+    // Optional shard DDL coordination mode. Only set this for shard merge
+    // (sharded database/table) scenarios. See:
+    // https://docs.pingcap.com/tidb/stable/dm-shard-merge/
+    // "shardMode": "PESSIMISTIC",
     // Target TiDB Cloud user credentials used by the migration
     "target": {
         "user": "migration_user",

pkg/tidbcloud/v1beta1/serverless/dm.swagger.json

@@ -824,6 +824,14 @@
               "$ref": "#/definitions/TaskMode"
             }
           ]
+        },
+        "shardMode": {
+          "description": "Optional shard mode for shard DDL coordination.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/TaskShardMode"
+            }
+          ]
         }
       },
       "required": ["displayName", "sources", "target", "mode"]
@@ -862,6 +870,14 @@
               "$ref": "#/definitions/TaskMode"
             }
           ]
+        },
+        "shardMode": {
+          "description": "Optional shard mode for shard DDL coordination.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/TaskShardMode"
+            }
+          ]
         }
       },
       "required": ["displayName", "sources", "target", "mode"]
@@ -1150,6 +1166,11 @@
       "type": "string",
       "enum": ["ALL", "INCREMENTAL"],
       "description": "Migration task mode.\n\n - ALL: Full + incremental migration (all phases).\n - INCREMENTAL: Incremental-only migration (replication)."
+    },
+    "TaskShardMode": {
+      "type": "string",
+      "enum": ["PESSIMISTIC", "OPTIMISTIC"],
+      "description": "Migration shard mode.\n\n - PESSIMISTIC: Use pessimistic shard DDL coordination.\n - OPTIMISTIC: Use optimistic shard DDL coordination."
     }
   }
 }

pkg/tidbcloud/v1beta1/serverless/migration/.openapi-generator/FILES

@@ -36,5 +36,6 @@ model_sub_task_step.go
 model_sync_detail.go
 model_target.go
 model_task_mode.go
+model_task_shard_mode.go
 response.go
 utils.go

pkg/tidbcloud/v1beta1/serverless/migration/README.md

@@ -120,6 +120,7 @@ Class | Method | HTTP request | Description
  - [SyncDetail](docs/SyncDetail.md)
  - [Target](docs/Target.md)
  - [TaskMode](docs/TaskMode.md)
+ - [TaskShardMode](docs/TaskShardMode.md)
 
 
 ## Documentation For Authorization

pkg/tidbcloud/v1beta1/serverless/migration/api/openapi.yaml

@@ -778,6 +778,11 @@ components:
           - $ref: '#/components/schemas/TaskMode'
           description: The migration mode (full+incremental or incremental-only).
           type: object
+        shardMode:
+          allOf:
+          - $ref: '#/components/schemas/TaskShardMode'
+          description: Optional shard mode for shard DDL coordination.
+          type: object
       required:
       - displayName
       - mode
@@ -807,6 +812,11 @@ components:
           - $ref: '#/components/schemas/TaskMode'
           description: The migration mode (full+incremental or incremental-only).
           type: object
+        shardMode:
+          allOf:
+          - $ref: '#/components/schemas/TaskShardMode'
+          description: Optional shard mode for shard DDL coordination.
+          type: object
       required:
       - displayName
       - mode
@@ -1099,4 +1109,14 @@ components:
       - ALL
       - INCREMENTAL
       type: string
+    TaskShardMode:
+      description: |-
+        Migration shard mode.
+
+         - PESSIMISTIC: Use pessimistic shard DDL coordination.
+         - OPTIMISTIC: Use optimistic shard DDL coordination.
+      enum:
+      - PESSIMISTIC
+      - OPTIMISTIC
+      type: string
 x-original-swagger-version: "2.0"