From 219ea2905305658b9fbfde795e1e73b1e2b37781 Mon Sep 17 00:00:00 2001 From: Nick Van Wiggeren Date: Tue, 31 Mar 2026 12:21:27 +0000 Subject: [PATCH 1/3] Add planned-reparent command to pscale branch vtctld MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `pscale branch vtctld planned-reparent` with two subcommands: - `create` — initiates a planned reparent shard operation, waits for completion by default (--wait=false to return immediately) - `get` — polls the status of an existing operation Usage: pscale branch vtctld planned-reparent create \ --keyspace --shard --new-primary Updates planetscale-go to v0.159.0 for PlannedReparentShardService. --- go.mod | 2 +- go.sum | 4 +- .../cmd/branch/vtctld/planned_reparent.go | 208 ++++++++++++++++++ internal/cmd/branch/vtctld/vtctld.go | 1 + internal/mock/planned_reparent_shard.go | 25 +++ 5 files changed, 237 insertions(+), 3 deletions(-) create mode 100644 internal/cmd/branch/vtctld/planned_reparent.go create mode 100644 internal/mock/planned_reparent_shard.go diff --git a/go.mod b/go.mod index 0c27106c..8e21e860 100644 --- a/go.mod +++ b/go.mod @@ -25,7 +25,7 @@ require ( github.com/mattn/go-shellwords v1.0.12 github.com/mitchellh/go-homedir v1.1.0 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c - github.com/planetscale/planetscale-go v0.157.0 + github.com/planetscale/planetscale-go v0.159.0 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 github.com/spf13/cobra v1.10.2 diff --git a/go.sum b/go.sum index 4cc9b45d..6220a740 100644 --- a/go.sum +++ b/go.sum @@ -176,8 +176,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjL github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2vvqGZLvoQfpaGg/j1fNDr4j03s3PRz4rVY= github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q= -github.com/planetscale/planetscale-go v0.157.0 h1:b0kWxC39F4/FQw2/Y+5/H4tRWUAzvl2ZukimrsTYP7M= -github.com/planetscale/planetscale-go v0.157.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= +github.com/planetscale/planetscale-go v0.159.0 h1:qqyZjG/z5k/w5gihfSwxssVu+mIsRTKqXFIeVJa/7hI= +github.com/planetscale/planetscale-go v0.159.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs= github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs= github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8= diff --git a/internal/cmd/branch/vtctld/planned_reparent.go b/internal/cmd/branch/vtctld/planned_reparent.go new file mode 100644 index 00000000..23256001 --- /dev/null +++ b/internal/cmd/branch/vtctld/planned_reparent.go @@ -0,0 +1,208 @@ +package vtctld + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/planetscale-go/planetscale" + "github.com/spf13/cobra" +) + +var ( + plannedReparentOperationPollInterval = time.Second + plannedReparentOperationTimeoutBuffer = 30 * time.Second + plannedReparentOperationDefaultTimeout = 10 * time.Minute +) + +func PlannedReparentCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "planned-reparent ", + Short: "Manage planned reparent shard operations", + } + + cmd.AddCommand(PlannedReparentCreateCmd(ch)) + cmd.AddCommand(PlannedReparentGetCmd(ch)) + + return cmd +} + +func PlannedReparentCreateCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + keyspace string + shard string + newPrimary string + wait bool + } + + cmd := &cobra.Command{ + Use: "create ", + Short: "Execute a planned reparent shard operation", + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch := args[0], args[1] + + client, err := ch.Client() + if err != nil { + return err + } + + end := ch.Printer.PrintProgress( + fmt.Sprintf("Executing PlannedReparentShard on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() + + operation, err := client.PlannedReparentShard.Create(ctx, &ps.PlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Keyspace: flags.keyspace, + Shard: flags.shard, + NewPrimary: flags.newPrimary, + }) + if err != nil { + return cmdutil.HandleError(err) + } + + if !flags.wait { + end() + return ch.Printer.PrintJSON(map[string]string{"id": operation.ID}) + } + + result, err := waitForPlannedReparentResult(ctx, client, ch.Config.Organization, database, branch, operation) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrettyPrintJSON(result) + }, + } + + cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Keyspace name") + cmd.Flags().StringVar(&flags.shard, "shard", "", "Shard range (e.g., '-80', '80-', or '-' for unsharded)") + cmd.Flags().StringVar(&flags.newPrimary, "new-primary", "", "Tablet alias to promote as the new primary") + cmd.Flags().BoolVar(&flags.wait, "wait", true, "Wait for the operation to complete") + cmd.MarkFlagRequired("keyspace") // nolint:errcheck + cmd.MarkFlagRequired("shard") // nolint:errcheck + cmd.MarkFlagRequired("new-primary") // nolint:errcheck + + return cmd +} + +func PlannedReparentGetCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "get ", + Short: "Get the status of a planned reparent shard operation", + Args: cmdutil.RequiredArgs("database", "branch", "id"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch, id := args[0], args[1], args[2] + + client, err := ch.Client() + if err != nil { + return err + } + + end := ch.Printer.PrintProgress( + fmt.Sprintf("Getting PlannedReparentShard operation on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() + + operation, err := client.PlannedReparentShard.Get(ctx, &ps.GetPlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + ID: id, + }) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrintJSON(operation) + }, + } + + return cmd +} + +func waitForPlannedReparentResult(ctx context.Context, client *ps.Client, organization, database, branch string, operation *ps.VtctldOperation) (json.RawMessage, error) { + result, done, err := plannedReparentOperationResult(operation) + if done || err != nil { + return result, err + } + + request := &ps.GetPlannedReparentShardRequest{ + Organization: organization, + Database: database, + Branch: branch, + ID: operation.ID, + } + + pollCtx, cancel := context.WithTimeout(ctx, plannedReparentOperationTimeout(operation)) + defer cancel() + ticker := time.NewTicker(plannedReparentOperationPollInterval) + defer ticker.Stop() + + for { + select { + case <-pollCtx.Done(): + if errors.Is(pollCtx.Err(), context.DeadlineExceeded) { + return nil, fmt.Errorf("timed out waiting for planned reparent operation %s to finish", operation.ID) + } + + return nil, pollCtx.Err() + case <-ticker.C: + } + + op, err := client.PlannedReparentShard.Get(pollCtx, request) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) { + return nil, fmt.Errorf("timed out waiting for planned reparent operation %s to finish", operation.ID) + } + + return nil, err + } + + result, done, err = plannedReparentOperationResult(op) + if done || err != nil { + return result, err + } + } +} + +func plannedReparentOperationResult(operation *ps.VtctldOperation) (json.RawMessage, bool, error) { + if !operation.Completed { + return nil, false, nil + } + + switch operation.State { + case "completed": + if len(operation.Result) == 0 { + return json.RawMessage(`{}`), true, nil + } + + return operation.Result, true, nil + case "failed", "cancelled": + if operation.Error != "" { + return nil, true, errors.New(operation.Error) + } + + return nil, true, fmt.Errorf("planned reparent operation %s ended in state %q", operation.ID, operation.State) + default: + return nil, true, fmt.Errorf("planned reparent operation %s reached unexpected terminal state %q", operation.ID, operation.State) + } +} + +func plannedReparentOperationTimeout(operation *ps.VtctldOperation) time.Duration { + if operation.Timeout > 0 { + return time.Duration(operation.Timeout)*time.Second + plannedReparentOperationTimeoutBuffer + } + + return plannedReparentOperationDefaultTimeout +} diff --git a/internal/cmd/branch/vtctld/vtctld.go b/internal/cmd/branch/vtctld/vtctld.go index 386150f5..4ed6d73e 100644 --- a/internal/cmd/branch/vtctld/vtctld.go +++ b/internal/cmd/branch/vtctld/vtctld.go @@ -17,6 +17,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command { cmd.AddCommand(VDiffCmd(ch)) cmd.AddCommand(LookupVindexCmd(ch)) cmd.AddCommand(MoveTablesCmd(ch)) + cmd.AddCommand(PlannedReparentCmd(ch)) cmd.AddCommand(ListWorkflowsCmd(ch)) cmd.AddCommand(ListKeyspacesCmd(ch)) cmd.AddCommand(StartWorkflowCmd(ch)) diff --git a/internal/mock/planned_reparent_shard.go b/internal/mock/planned_reparent_shard.go new file mode 100644 index 00000000..4b116b6b --- /dev/null +++ b/internal/mock/planned_reparent_shard.go @@ -0,0 +1,25 @@ +package mock + +import ( + "context" + + ps "github.com/planetscale/planetscale-go/planetscale" +) + +type PlannedReparentShardService struct { + CreateFn func(context.Context, *ps.PlannedReparentShardRequest) (*ps.VtctldOperation, error) + CreateFnInvoked bool + + GetFn func(context.Context, *ps.GetPlannedReparentShardRequest) (*ps.VtctldOperation, error) + GetFnInvoked bool +} + +func (s *PlannedReparentShardService) Create(ctx context.Context, req *ps.PlannedReparentShardRequest) (*ps.VtctldOperation, error) { + s.CreateFnInvoked = true + return s.CreateFn(ctx, req) +} + +func (s *PlannedReparentShardService) Get(ctx context.Context, req *ps.GetPlannedReparentShardRequest) (*ps.VtctldOperation, error) { + s.GetFnInvoked = true + return s.GetFn(ctx, req) +} From f140b6a68dd6587d5fc98ec9796f51cf84563e31 Mon Sep 17 00:00:00 2001 From: Nick Van Wiggeren Date: Tue, 31 Mar 2026 12:28:27 +0000 Subject: [PATCH 2/3] Simplify planned-reparent-shard to a single command Instead of create/get subcommands, use a single command with two modes: # Execute a planned reparent pscale branch vtctld planned-reparent-shard \ --keyspace --shard --new-primary # Check on an existing operation pscale branch vtctld planned-reparent-shard --id This matches how Vitess exposes PlannedReparentShard as a single action rather than a CRUD resource. --- .../cmd/branch/vtctld/planned_reparent.go | 139 ++++++++---------- internal/cmd/branch/vtctld/vtctld.go | 2 +- 2 files changed, 66 insertions(+), 75 deletions(-) diff --git a/internal/cmd/branch/vtctld/planned_reparent.go b/internal/cmd/branch/vtctld/planned_reparent.go index 23256001..7f3cb71d 100644 --- a/internal/cmd/branch/vtctld/planned_reparent.go +++ b/internal/cmd/branch/vtctld/planned_reparent.go @@ -18,30 +18,28 @@ var ( plannedReparentOperationDefaultTimeout = 10 * time.Minute ) -func PlannedReparentCmd(ch *cmdutil.Helper) *cobra.Command { - cmd := &cobra.Command{ - Use: "planned-reparent ", - Short: "Manage planned reparent shard operations", - } - - cmd.AddCommand(PlannedReparentCreateCmd(ch)) - cmd.AddCommand(PlannedReparentGetCmd(ch)) - - return cmd -} - -func PlannedReparentCreateCmd(ch *cmdutil.Helper) *cobra.Command { +func PlannedReparentShardCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { keyspace string shard string newPrimary string wait bool + id string } cmd := &cobra.Command{ - Use: "create ", - Short: "Execute a planned reparent shard operation", - Args: cmdutil.RequiredArgs("database", "branch"), + Use: "planned-reparent-shard ", + Short: "Reparent a shard to a new primary, or check on an existing reparent operation", + Long: `Reparent a shard to a new primary using Vitess PlannedReparentShard. +Both the old and new primaries must be up and running. + +To execute a planned reparent: + pscale branch vtctld planned-reparent-shard \ + --keyspace --shard --new-primary + +To check on an existing operation: + pscale branch vtctld planned-reparent-shard --id `, + Args: cmdutil.RequiredArgs("database", "branch"), RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() database, branch := args[0], args[1] @@ -51,35 +49,15 @@ func PlannedReparentCreateCmd(ch *cmdutil.Helper) *cobra.Command { return err } - end := ch.Printer.PrintProgress( - fmt.Sprintf("Executing PlannedReparentShard on %s\u2026", - progressTarget(ch.Config.Organization, database, branch))) - defer end() - - operation, err := client.PlannedReparentShard.Create(ctx, &ps.PlannedReparentShardRequest{ - Organization: ch.Config.Organization, - Database: database, - Branch: branch, - Keyspace: flags.keyspace, - Shard: flags.shard, - NewPrimary: flags.newPrimary, - }) - if err != nil { - return cmdutil.HandleError(err) + if flags.id != "" { + return getPlannedReparentOperation(ctx, ch, client, database, branch, flags.id) } - if !flags.wait { - end() - return ch.Printer.PrintJSON(map[string]string{"id": operation.ID}) + if flags.keyspace == "" || flags.shard == "" || flags.newPrimary == "" { + return fmt.Errorf("--keyspace, --shard, and --new-primary are required when not using --id") } - result, err := waitForPlannedReparentResult(ctx, client, ch.Config.Organization, database, branch, operation) - if err != nil { - return cmdutil.HandleError(err) - } - - end() - return ch.Printer.PrettyPrintJSON(result) + return runPlannedReparentShard(ctx, ch, client, database, branch, flags.keyspace, flags.shard, flags.newPrimary, flags.wait) }, } @@ -87,48 +65,61 @@ func PlannedReparentCreateCmd(ch *cmdutil.Helper) *cobra.Command { cmd.Flags().StringVar(&flags.shard, "shard", "", "Shard range (e.g., '-80', '80-', or '-' for unsharded)") cmd.Flags().StringVar(&flags.newPrimary, "new-primary", "", "Tablet alias to promote as the new primary") cmd.Flags().BoolVar(&flags.wait, "wait", true, "Wait for the operation to complete") - cmd.MarkFlagRequired("keyspace") // nolint:errcheck - cmd.MarkFlagRequired("shard") // nolint:errcheck - cmd.MarkFlagRequired("new-primary") // nolint:errcheck + cmd.Flags().StringVar(&flags.id, "id", "", "Check status of an existing planned reparent operation") return cmd } -func PlannedReparentGetCmd(ch *cmdutil.Helper) *cobra.Command { - cmd := &cobra.Command{ - Use: "get ", - Short: "Get the status of a planned reparent shard operation", - Args: cmdutil.RequiredArgs("database", "branch", "id"), - RunE: func(cmd *cobra.Command, args []string) error { - ctx := cmd.Context() - database, branch, id := args[0], args[1], args[2] +func runPlannedReparentShard(ctx context.Context, ch *cmdutil.Helper, client *ps.Client, database, branch, keyspace, shard, newPrimary string, wait bool) error { + end := ch.Printer.PrintProgress( + fmt.Sprintf("Executing PlannedReparentShard on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() - client, err := ch.Client() - if err != nil { - return err - } + operation, err := client.PlannedReparentShard.Create(ctx, &ps.PlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Keyspace: keyspace, + Shard: shard, + NewPrimary: newPrimary, + }) + if err != nil { + return cmdutil.HandleError(err) + } - end := ch.Printer.PrintProgress( - fmt.Sprintf("Getting PlannedReparentShard operation on %s\u2026", - progressTarget(ch.Config.Organization, database, branch))) - defer end() - - operation, err := client.PlannedReparentShard.Get(ctx, &ps.GetPlannedReparentShardRequest{ - Organization: ch.Config.Organization, - Database: database, - Branch: branch, - ID: id, - }) - if err != nil { - return cmdutil.HandleError(err) - } + if !wait { + end() + return ch.Printer.PrintJSON(map[string]string{"id": operation.ID}) + } - end() - return ch.Printer.PrintJSON(operation) - }, + result, err := waitForPlannedReparentResult(ctx, client, ch.Config.Organization, database, branch, operation) + if err != nil { + return cmdutil.HandleError(err) } - return cmd + end() + return ch.Printer.PrettyPrintJSON(result) +} + +func getPlannedReparentOperation(ctx context.Context, ch *cmdutil.Helper, client *ps.Client, database, branch, id string) error { + end := ch.Printer.PrintProgress( + fmt.Sprintf("Getting PlannedReparentShard operation on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() + + operation, err := client.PlannedReparentShard.Get(ctx, &ps.GetPlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + ID: id, + }) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrintJSON(operation) } func waitForPlannedReparentResult(ctx context.Context, client *ps.Client, organization, database, branch string, operation *ps.VtctldOperation) (json.RawMessage, error) { diff --git a/internal/cmd/branch/vtctld/vtctld.go b/internal/cmd/branch/vtctld/vtctld.go index 4ed6d73e..26a1331c 100644 --- a/internal/cmd/branch/vtctld/vtctld.go +++ b/internal/cmd/branch/vtctld/vtctld.go @@ -17,7 +17,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command { cmd.AddCommand(VDiffCmd(ch)) cmd.AddCommand(LookupVindexCmd(ch)) cmd.AddCommand(MoveTablesCmd(ch)) - cmd.AddCommand(PlannedReparentCmd(ch)) + cmd.AddCommand(PlannedReparentShardCmd(ch)) cmd.AddCommand(ListWorkflowsCmd(ch)) cmd.AddCommand(ListKeyspacesCmd(ch)) cmd.AddCommand(StartWorkflowCmd(ch)) From e25a2aace55eacd582f2ed66987455eb0d6ad1c5 Mon Sep 17 00:00:00 2001 From: Nick Van Wiggeren Date: Tue, 31 Mar 2026 12:32:48 +0000 Subject: [PATCH 3/3] Use status subcommand instead of --id flag for operation polling The root command runs the reparent (all flags required, no ambiguity): pscale branch vtctld planned-reparent-shard \ --keyspace --shard --new-primary A status subcommand checks on an existing operation: pscale branch vtctld planned-reparent-shard status --- .../cmd/branch/vtctld/planned_reparent.go | 122 +++++++++--------- 1 file changed, 63 insertions(+), 59 deletions(-) diff --git a/internal/cmd/branch/vtctld/planned_reparent.go b/internal/cmd/branch/vtctld/planned_reparent.go index 7f3cb71d..577c4fe5 100644 --- a/internal/cmd/branch/vtctld/planned_reparent.go +++ b/internal/cmd/branch/vtctld/planned_reparent.go @@ -24,21 +24,16 @@ func PlannedReparentShardCmd(ch *cmdutil.Helper) *cobra.Command { shard string newPrimary string wait bool - id string } cmd := &cobra.Command{ Use: "planned-reparent-shard ", - Short: "Reparent a shard to a new primary, or check on an existing reparent operation", + Short: "Reparent a shard to a new primary", Long: `Reparent a shard to a new primary using Vitess PlannedReparentShard. Both the old and new primaries must be up and running. -To execute a planned reparent: - pscale branch vtctld planned-reparent-shard \ - --keyspace --shard --new-primary - -To check on an existing operation: - pscale branch vtctld planned-reparent-shard --id `, +To check on an existing operation, use the "status" subcommand: + pscale branch vtctld planned-reparent-shard status `, Args: cmdutil.RequiredArgs("database", "branch"), RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() @@ -49,15 +44,35 @@ To check on an existing operation: return err } - if flags.id != "" { - return getPlannedReparentOperation(ctx, ch, client, database, branch, flags.id) + end := ch.Printer.PrintProgress( + fmt.Sprintf("Executing PlannedReparentShard on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() + + operation, err := client.PlannedReparentShard.Create(ctx, &ps.PlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Keyspace: flags.keyspace, + Shard: flags.shard, + NewPrimary: flags.newPrimary, + }) + if err != nil { + return cmdutil.HandleError(err) } - if flags.keyspace == "" || flags.shard == "" || flags.newPrimary == "" { - return fmt.Errorf("--keyspace, --shard, and --new-primary are required when not using --id") + if !flags.wait { + end() + return ch.Printer.PrintJSON(map[string]string{"id": operation.ID}) } - return runPlannedReparentShard(ctx, ch, client, database, branch, flags.keyspace, flags.shard, flags.newPrimary, flags.wait) + result, err := waitForPlannedReparentResult(ctx, client, ch.Config.Organization, database, branch, operation) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrettyPrintJSON(result) }, } @@ -65,61 +80,50 @@ To check on an existing operation: cmd.Flags().StringVar(&flags.shard, "shard", "", "Shard range (e.g., '-80', '80-', or '-' for unsharded)") cmd.Flags().StringVar(&flags.newPrimary, "new-primary", "", "Tablet alias to promote as the new primary") cmd.Flags().BoolVar(&flags.wait, "wait", true, "Wait for the operation to complete") - cmd.Flags().StringVar(&flags.id, "id", "", "Check status of an existing planned reparent operation") + cmd.MarkFlagRequired("keyspace") // nolint:errcheck + cmd.MarkFlagRequired("shard") // nolint:errcheck + cmd.MarkFlagRequired("new-primary") // nolint:errcheck + + cmd.AddCommand(plannedReparentShardStatusCmd(ch)) return cmd } -func runPlannedReparentShard(ctx context.Context, ch *cmdutil.Helper, client *ps.Client, database, branch, keyspace, shard, newPrimary string, wait bool) error { - end := ch.Printer.PrintProgress( - fmt.Sprintf("Executing PlannedReparentShard on %s\u2026", - progressTarget(ch.Config.Organization, database, branch))) - defer end() - - operation, err := client.PlannedReparentShard.Create(ctx, &ps.PlannedReparentShardRequest{ - Organization: ch.Config.Organization, - Database: database, - Branch: branch, - Keyspace: keyspace, - Shard: shard, - NewPrimary: newPrimary, - }) - if err != nil { - return cmdutil.HandleError(err) - } - - if !wait { - end() - return ch.Printer.PrintJSON(map[string]string{"id": operation.ID}) - } - - result, err := waitForPlannedReparentResult(ctx, client, ch.Config.Organization, database, branch, operation) - if err != nil { - return cmdutil.HandleError(err) - } +func plannedReparentShardStatusCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "status ", + Short: "Check the status of a planned reparent shard operation", + Args: cmdutil.RequiredArgs("database", "branch", "operation-id"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch, id := args[0], args[1], args[2] - end() - return ch.Printer.PrettyPrintJSON(result) -} + client, err := ch.Client() + if err != nil { + return err + } -func getPlannedReparentOperation(ctx context.Context, ch *cmdutil.Helper, client *ps.Client, database, branch, id string) error { - end := ch.Printer.PrintProgress( - fmt.Sprintf("Getting PlannedReparentShard operation on %s\u2026", - progressTarget(ch.Config.Organization, database, branch))) - defer end() + end := ch.Printer.PrintProgress( + fmt.Sprintf("Getting PlannedReparentShard operation on %s\u2026", + progressTarget(ch.Config.Organization, database, branch))) + defer end() + + operation, err := client.PlannedReparentShard.Get(ctx, &ps.GetPlannedReparentShardRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + ID: id, + }) + if err != nil { + return cmdutil.HandleError(err) + } - operation, err := client.PlannedReparentShard.Get(ctx, &ps.GetPlannedReparentShardRequest{ - Organization: ch.Config.Organization, - Database: database, - Branch: branch, - ID: id, - }) - if err != nil { - return cmdutil.HandleError(err) + end() + return ch.Printer.PrintJSON(operation) + }, } - end() - return ch.Printer.PrintJSON(operation) + return cmd } func waitForPlannedReparentResult(ctx context.Context, client *ps.Client, organization, database, branch string, operation *ps.VtctldOperation) (json.RawMessage, error) {