Add vtctld set-shard-tablet-control command

go.sum

@@ -176,12 +176,6 @@ 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.168.1-0.20260609224250-8001bcf5eba5 h1:z+ao9YVEVP9GHpyB1Gmw48dt2+RQsS+BAOTE1QUigrw=
-github.com/planetscale/planetscale-go v0.168.1-0.20260609224250-8001bcf5eba5/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
-github.com/planetscale/planetscale-go v0.168.1-0.20260610223405-068c67a15c9b h1:VJwfNPEI1sYyo52bhN/f8xxeHoh/BhE/0XWaNy5jYOk=
-github.com/planetscale/planetscale-go v0.168.1-0.20260610223405-068c67a15c9b/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
-github.com/planetscale/planetscale-go v0.168.1 h1:ikGvRC5YlQQiNf7vF9vxi6WBXNs2NokNie0sazMTTcA=
-github.com/planetscale/planetscale-go v0.168.1/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
 github.com/planetscale/planetscale-go v0.168.2-0.20260611211624-109b07535906 h1:n+R4KmP6Ud62T8tvCK9uJ5pzpx4g0NHaRahecOH/C4M=
 github.com/planetscale/planetscale-go v0.168.2-0.20260611211624-109b07535906/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=

internal/cmd/branch/vtctld/set_shard_tablet_control.go

@@ -0,0 +1,99 @@
+package vtctld
+
+import (
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+	"github.com/spf13/cobra"
+)
+
+// SetShardTabletControlCmd updates shard tablet controls via vtctld.
+func SetShardTabletControlCmd(ch *cmdutil.Helper) *cobra.Command {
+	var flags struct {
+		keyspace            string
+		shard               string
+		tabletType          string
+		cells               []string
+		deniedTables        []string
+		remove              bool
+		disableQueryService bool
+	}
+
+	cmd := &cobra.Command{
+		Use:   "set-shard-tablet-control <database> <branch>",
+		Short: "Update shard tablet controls for a branch",
+		Long: "Update live shard tablet controls from the cluster via vtctld, " +
+			"including denied tables used during MoveTables cleanup.",
+		Args: cmdutil.RequiredArgs("database", "branch"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			database, branch := args[0], args[1]
+
+			if flags.keyspace == "" {
+				return fmt.Errorf("keyspace is required")
+			}
+			if flags.shard == "" {
+				return fmt.Errorf("shard is required")
+			}
+			if flags.tabletType == "" {
+				return fmt.Errorf("tablet-type is required")
+			}
+
+			removeSet := cmd.Flags().Changed("remove")
+			disableQueryServiceSet := cmd.Flags().Changed("disable-query-service")
+			if !removeSet && len(flags.deniedTables) == 0 && !disableQueryServiceSet {
+				return fmt.Errorf("must specify at least one of --remove, --denied-tables, or --disable-query-service")
+			}
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			end := ch.Printer.PrintProgress(
+				fmt.Sprintf("Updating tablet controls for %s/%s on %s…",
+					flags.keyspace, flags.shard,
+					progressTarget(ch.Config.Organization, database, branch)))
+			defer end()
+
+			req := &ps.VtctldSetShardTabletControlRequest{
+				Organization: ch.Config.Organization,
+				Database:     database,
+				Branch:       branch,
+				Keyspace:     flags.keyspace,
+				Shard:        flags.shard,
+				TabletType:   flags.tabletType,
+				Cells:        flags.cells,
+				DeniedTables: flags.deniedTables,
+			}
+			if removeSet {
+				req.Remove = &flags.remove
+			}
+			if disableQueryServiceSet {
+				req.DisableQueryService = &flags.disableQueryService
+			}
+
+			data, err := client.Vtctld.SetShardTabletControl(ctx, req)
+			if err != nil {
+				return cmdutil.HandleError(err)
+			}
+
+			end()
+			return ch.Printer.PrettyPrintJSON(data)
+		},
+	}
+
+	cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Keyspace name")
+	cmd.Flags().StringVar(&flags.shard, "shard", "", "Shard name (e.g. \"-\" for unsharded)")
+	cmd.Flags().StringVar(&flags.tabletType, "tablet-type", "", "Tablet type (e.g. rdonly, replica, primary)")
+	cmd.Flags().StringSliceVar(&flags.cells, "cells", nil, "Cells to update (comma-separated)")
+	cmd.Flags().StringSliceVar(&flags.deniedTables, "denied-tables", nil, "Tables to add to or remove from the denylist (comma-separated)")
+	cmd.Flags().BoolVar(&flags.remove, "remove", false, "Remove tablet controls (or specific denied tables when combined with --denied-tables)")
+	cmd.Flags().BoolVar(&flags.disableQueryService, "disable-query-service", false, "Disable query service on the provided tablets")
+	cmd.MarkFlagRequired("keyspace")
+	cmd.MarkFlagRequired("shard")
+	cmd.MarkFlagRequired("tablet-type")
+
+	return cmd
+}

internal/cmd/branch/vtctld/set_shard_tablet_control_test.go

@@ -0,0 +1,66 @@
+package vtctld
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"testing"
+
+	qt "github.com/frankban/quicktest"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/mock"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+)
+
+func TestSetShardTabletControl(t *testing.T) {
+	c := qt.New(t)
+
+	org := "my-org"
+	db := "my-db"
+	branch := "my-branch"
+
+	svc := &mock.VtctldService{
+		SetShardTabletControlFn: func(ctx context.Context, req *ps.VtctldSetShardTabletControlRequest) (json.RawMessage, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Branch, qt.Equals, branch)
+			c.Assert(req.Keyspace, qt.Equals, "commerce")
+			c.Assert(req.Shard, qt.Equals, "-")
+			c.Assert(req.TabletType, qt.Equals, "rdonly")
+			c.Assert(req.DeniedTables, qt.DeepEquals, []string{"customers"})
+			c.Assert(req.Remove, qt.Not(qt.IsNil))
+			c.Assert(*req.Remove, qt.Equals, true)
+			return json.RawMessage(`{}`), nil
+		},
+	}
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Vtctld: svc,
+			}, nil
+		},
+	}
+
+	cmd := SetShardTabletControlCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	cmd.Flags().Set("keyspace", "commerce")
+	cmd.Flags().Set("shard", "-")
+	cmd.Flags().Set("tablet-type", "rdonly")
+	cmd.Flags().Set("denied-tables", "customers")
+	cmd.Flags().Set("remove", "true")
+
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.SetShardTabletControlFnInvoked, qt.IsTrue)
+}

internal/cmd/branch/vtctld/vtctld.go

@@ -22,6 +22,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(ListKeyspacesCmd(ch))
 	cmd.AddCommand(GetRoutingRulesCmd(ch))
 	cmd.AddCommand(GetShardCmd(ch))
+	cmd.AddCommand(SetShardTabletControlCmd(ch))
 	cmd.AddCommand(ListTabletsCmd(ch))
 	cmd.AddCommand(StartWorkflowCmd(ch))
 	cmd.AddCommand(StopWorkflowCmd(ch))

internal/mock/vtctld_general.go

@@ -20,6 +20,9 @@ type VtctldService struct {
 	GetShardFn        func(context.Context, *ps.VtctldGetShardRequest) (json.RawMessage, error)
 	GetShardFnInvoked bool
 
+	SetShardTabletControlFn        func(context.Context, *ps.VtctldSetShardTabletControlRequest) (json.RawMessage, error)
+	SetShardTabletControlFnInvoked bool
+
 	ListTabletsFn        func(context.Context, *ps.ListBranchTabletsRequest) ([]*ps.TabletGroup, error)
 	ListTabletsFnInvoked bool
 
@@ -62,6 +65,11 @@ func (s *VtctldService) GetShard(ctx context.Context, req *ps.VtctldGetShardRequ
 	return s.GetShardFn(ctx, req)
 }
 
+func (s *VtctldService) SetShardTabletControl(ctx context.Context, req *ps.VtctldSetShardTabletControlRequest) (json.RawMessage, error) {
+	s.SetShardTabletControlFnInvoked = true
+	return s.SetShardTabletControlFn(ctx, req)
+}
+
 func (s *VtctldService) ListTablets(ctx context.Context, req *ps.ListBranchTabletsRequest) ([]*ps.TabletGroup, error) {
 	s.ListTabletsFnInvoked = true
 	return s.ListTabletsFn(ctx, req)