Add planned-reparent command to pscale branch vtctld

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

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=

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 <database> <branch>",
+		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 <db> <branch> <operation-id>`,
+		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 <database> <branch> <operation-id>",
+		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
+}

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))

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)
+}