From eea1a259e11b93aa82801834cabe7e2581065944 Mon Sep 17 00:00:00 2001 From: yangxin Date: Thu, 7 May 2026 09:59:52 +0800 Subject: [PATCH] Add shard mode support to migration CLI --- internal/cli/serverless/migration/create.go | 47 ++++++-- .../cli/serverless/migration/create_test.go | 9 ++ internal/cli/serverless/migration/template.go | 8 ++ .../v1beta1/serverless/dm.swagger.json | 21 ++++ .../migration/.openapi-generator/FILES | 1 + .../v1beta1/serverless/migration/README.md | 1 + .../serverless/migration/api/openapi.yaml | 20 ++++ ...migration_service_create_migration_body.go | 40 ++++++- .../model_migration_service_precheck_body.go | 40 ++++++- .../migration/model_task_shard_mode.go | 105 ++++++++++++++++++ 10 files changed, 278 insertions(+), 14 deletions(-) create mode 100644 pkg/tidbcloud/v1beta1/serverless/migration/model_task_shard_mode.go diff --git a/internal/cli/serverless/migration/create.go b/internal/cli/serverless/migration/create.go index 6e70f9fc..2f75bbf9 100644 --- a/internal/cli/serverless/migration/create.go +++ b/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. diff --git a/internal/cli/serverless/migration/create_test.go b/internal/cli/serverless/migration/create_test.go index 23461daf..ff9350f0 100644 --- a/internal/cli/serverless/migration/create_test.go +++ b/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!" diff --git a/internal/cli/serverless/migration/template.go b/internal/cli/serverless/migration/template.go index 8ce8ea04..b7ffe7ea 100644 --- a/internal/cli/serverless/migration/template.go +++ b/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", diff --git a/pkg/tidbcloud/v1beta1/serverless/dm.swagger.json b/pkg/tidbcloud/v1beta1/serverless/dm.swagger.json index 98b2a0c0..5c898292 100644 --- a/pkg/tidbcloud/v1beta1/serverless/dm.swagger.json +++ b/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." } } } diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/.openapi-generator/FILES b/pkg/tidbcloud/v1beta1/serverless/migration/.openapi-generator/FILES index 66d4257c..52bcbd39 100644 --- a/pkg/tidbcloud/v1beta1/serverless/migration/.openapi-generator/FILES +++ b/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 diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/README.md b/pkg/tidbcloud/v1beta1/serverless/migration/README.md index 6fe901cc..25fddf33 100644 --- a/pkg/tidbcloud/v1beta1/serverless/migration/README.md +++ b/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 diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/api/openapi.yaml b/pkg/tidbcloud/v1beta1/serverless/migration/api/openapi.yaml index 4e84f8cf..69445afe 100644 --- a/pkg/tidbcloud/v1beta1/serverless/migration/api/openapi.yaml +++ b/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" diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_create_migration_body.go b/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_create_migration_body.go index 1c5f1ab0..a79ee32c 100644 --- a/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_create_migration_body.go +++ b/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_create_migration_body.go @@ -27,7 +27,9 @@ type MigrationServiceCreateMigrationBody struct { // The target database credentials. Target Target `json:"target"` // The migration mode (full+incremental or incremental-only). - Mode TaskMode `json:"mode"` + Mode TaskMode `json:"mode"` + // Optional shard mode for shard DDL coordination. + ShardMode *TaskShardMode `json:"shardMode,omitempty"` AdditionalProperties map[string]interface{} } @@ -150,6 +152,38 @@ func (o *MigrationServiceCreateMigrationBody) SetMode(v TaskMode) { o.Mode = v } +// GetShardMode returns the ShardMode field value if set, zero value otherwise. +func (o *MigrationServiceCreateMigrationBody) GetShardMode() TaskShardMode { + if o == nil || IsNil(o.ShardMode) { + var ret TaskShardMode + return ret + } + return *o.ShardMode +} + +// GetShardModeOk returns a tuple with the ShardMode field value if set, nil otherwise +// and a boolean to check if the value has been set. +func (o *MigrationServiceCreateMigrationBody) GetShardModeOk() (*TaskShardMode, bool) { + if o == nil || IsNil(o.ShardMode) { + return nil, false + } + return o.ShardMode, true +} + +// HasShardMode returns a boolean if a field has been set. +func (o *MigrationServiceCreateMigrationBody) HasShardMode() bool { + if o != nil && !IsNil(o.ShardMode) { + return true + } + + return false +} + +// SetShardMode gets a reference to the given TaskShardMode and assigns it to the ShardMode field. +func (o *MigrationServiceCreateMigrationBody) SetShardMode(v TaskShardMode) { + o.ShardMode = &v +} + func (o MigrationServiceCreateMigrationBody) MarshalJSON() ([]byte, error) { toSerialize, err := o.ToMap() if err != nil { @@ -164,6 +198,9 @@ func (o MigrationServiceCreateMigrationBody) ToMap() (map[string]interface{}, er toSerialize["sources"] = o.Sources toSerialize["target"] = o.Target toSerialize["mode"] = o.Mode + if !IsNil(o.ShardMode) { + toSerialize["shardMode"] = o.ShardMode + } for key, value := range o.AdditionalProperties { toSerialize[key] = value @@ -214,6 +251,7 @@ func (o *MigrationServiceCreateMigrationBody) UnmarshalJSON(data []byte) (err er delete(additionalProperties, "sources") delete(additionalProperties, "target") delete(additionalProperties, "mode") + delete(additionalProperties, "shardMode") o.AdditionalProperties = additionalProperties } diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_precheck_body.go b/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_precheck_body.go index 047968a8..8a11e705 100644 --- a/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_precheck_body.go +++ b/pkg/tidbcloud/v1beta1/serverless/migration/model_migration_service_precheck_body.go @@ -27,7 +27,9 @@ type MigrationServicePrecheckBody struct { // The target database credentials. Target Target `json:"target"` // The migration mode (full+incremental or incremental-only). - Mode TaskMode `json:"mode"` + Mode TaskMode `json:"mode"` + // Optional shard mode for shard DDL coordination. + ShardMode *TaskShardMode `json:"shardMode,omitempty"` AdditionalProperties map[string]interface{} } @@ -150,6 +152,38 @@ func (o *MigrationServicePrecheckBody) SetMode(v TaskMode) { o.Mode = v } +// GetShardMode returns the ShardMode field value if set, zero value otherwise. +func (o *MigrationServicePrecheckBody) GetShardMode() TaskShardMode { + if o == nil || IsNil(o.ShardMode) { + var ret TaskShardMode + return ret + } + return *o.ShardMode +} + +// GetShardModeOk returns a tuple with the ShardMode field value if set, nil otherwise +// and a boolean to check if the value has been set. +func (o *MigrationServicePrecheckBody) GetShardModeOk() (*TaskShardMode, bool) { + if o == nil || IsNil(o.ShardMode) { + return nil, false + } + return o.ShardMode, true +} + +// HasShardMode returns a boolean if a field has been set. +func (o *MigrationServicePrecheckBody) HasShardMode() bool { + if o != nil && !IsNil(o.ShardMode) { + return true + } + + return false +} + +// SetShardMode gets a reference to the given TaskShardMode and assigns it to the ShardMode field. +func (o *MigrationServicePrecheckBody) SetShardMode(v TaskShardMode) { + o.ShardMode = &v +} + func (o MigrationServicePrecheckBody) MarshalJSON() ([]byte, error) { toSerialize, err := o.ToMap() if err != nil { @@ -164,6 +198,9 @@ func (o MigrationServicePrecheckBody) ToMap() (map[string]interface{}, error) { toSerialize["sources"] = o.Sources toSerialize["target"] = o.Target toSerialize["mode"] = o.Mode + if !IsNil(o.ShardMode) { + toSerialize["shardMode"] = o.ShardMode + } for key, value := range o.AdditionalProperties { toSerialize[key] = value @@ -214,6 +251,7 @@ func (o *MigrationServicePrecheckBody) UnmarshalJSON(data []byte) (err error) { delete(additionalProperties, "sources") delete(additionalProperties, "target") delete(additionalProperties, "mode") + delete(additionalProperties, "shardMode") o.AdditionalProperties = additionalProperties } diff --git a/pkg/tidbcloud/v1beta1/serverless/migration/model_task_shard_mode.go b/pkg/tidbcloud/v1beta1/serverless/migration/model_task_shard_mode.go new file mode 100644 index 00000000..80f26004 --- /dev/null +++ b/pkg/tidbcloud/v1beta1/serverless/migration/model_task_shard_mode.go @@ -0,0 +1,105 @@ +/* +TiDB Cloud Starter and Essential API + +TiDB Cloud Starter and Essential API + +API version: v1beta1 +*/ + +// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT. + +package migration + +import ( + "encoding/json" +) + +// TaskShardMode Migration shard mode. - PESSIMISTIC: Use pessimistic shard DDL coordination. - OPTIMISTIC: Use optimistic shard DDL coordination. +type TaskShardMode string + +// List of TaskShardMode +const ( + TASKSHARDMODE_PESSIMISTIC TaskShardMode = "PESSIMISTIC" + TASKSHARDMODE_OPTIMISTIC TaskShardMode = "OPTIMISTIC" +) + +// All allowed values of TaskShardMode enum +var AllowedTaskShardModeEnumValues = []TaskShardMode{ + "PESSIMISTIC", + "OPTIMISTIC", +} + +func (v *TaskShardMode) UnmarshalJSON(src []byte) error { + var value string + err := json.Unmarshal(src, &value) + if err != nil { + return err + } + enumTypeValue := TaskShardMode(value) + for _, existing := range AllowedTaskShardModeEnumValues { + if existing == enumTypeValue { + *v = enumTypeValue + return nil + } + } + + *v = TaskShardMode(value) + return nil +} + +// NewTaskShardModeFromValue returns a pointer to a valid TaskShardMode for the value passed as argument +func NewTaskShardModeFromValue(v string) *TaskShardMode { + ev := TaskShardMode(v) + return &ev +} + +// IsValid return true if the value is valid for the enum, false otherwise +func (v TaskShardMode) IsValid() bool { + for _, existing := range AllowedTaskShardModeEnumValues { + if existing == v { + return true + } + } + return false +} + +// Ptr returns reference to TaskShardMode value +func (v TaskShardMode) Ptr() *TaskShardMode { + return &v +} + +type NullableTaskShardMode struct { + value *TaskShardMode + isSet bool +} + +func (v NullableTaskShardMode) Get() *TaskShardMode { + return v.value +} + +func (v *NullableTaskShardMode) Set(val *TaskShardMode) { + v.value = val + v.isSet = true +} + +func (v NullableTaskShardMode) IsSet() bool { + return v.isSet +} + +func (v *NullableTaskShardMode) Unset() { + v.value = nil + v.isSet = false +} + +func NewNullableTaskShardMode(val *TaskShardMode) *NullableTaskShardMode { + return &NullableTaskShardMode{value: val, isSet: true} +} + +func (v NullableTaskShardMode) MarshalJSON() ([]byte, error) { + return json.Marshal(v.value) +} + +func (v *NullableTaskShardMode) UnmarshalJSON(src []byte) error { + v.isSet = true + return json.Unmarshal(src, &v.value) +}