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..577c4fe5 --- /dev/null +++ b/internal/cmd/branch/vtctld/planned_reparent.go @@ -0,0 +1,203 @@ +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 PlannedReparentShardCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + keyspace string + shard string + newPrimary string + wait bool + } + + cmd := &cobra.Command{ + Use: "planned-reparent-shard ", + 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 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() + 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 + + cmd.AddCommand(plannedReparentShardStatusCmd(ch)) + + return cmd +} + +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] + + 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..26a1331c 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(PlannedReparentShardCmd(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) +}