From e43aab367ceec4cf693f582f71faaffab7b10d0c Mon Sep 17 00:00:00 2001 From: Max Englander Date: Sat, 6 Jun 2026 18:51:50 -0400 Subject: [PATCH] Add vtctld set-shard-tablet-control command Co-authored-by: Cursor Signed-off-by: Max Englander --- go.sum | 6 -- .../branch/vtctld/set_shard_tablet_control.go | 99 +++++++++++++++++++ .../vtctld/set_shard_tablet_control_test.go | 66 +++++++++++++ internal/cmd/branch/vtctld/vtctld.go | 1 + internal/mock/vtctld_general.go | 8 ++ 5 files changed, 174 insertions(+), 6 deletions(-) create mode 100644 internal/cmd/branch/vtctld/set_shard_tablet_control.go create mode 100644 internal/cmd/branch/vtctld/set_shard_tablet_control_test.go diff --git a/go.sum b/go.sum index 7b198e69..089e76dc 100644 --- a/go.sum +++ b/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= diff --git a/internal/cmd/branch/vtctld/set_shard_tablet_control.go b/internal/cmd/branch/vtctld/set_shard_tablet_control.go new file mode 100644 index 00000000..95ab37c9 --- /dev/null +++ b/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 ", + 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 +} diff --git a/internal/cmd/branch/vtctld/set_shard_tablet_control_test.go b/internal/cmd/branch/vtctld/set_shard_tablet_control_test.go new file mode 100644 index 00000000..69098e1e --- /dev/null +++ b/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) +} diff --git a/internal/cmd/branch/vtctld/vtctld.go b/internal/cmd/branch/vtctld/vtctld.go index 7b54ba8e..3817c9cf 100644 --- a/internal/cmd/branch/vtctld/vtctld.go +++ b/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)) diff --git a/internal/mock/vtctld_general.go b/internal/mock/vtctld_general.go index 256678af..b6e0952d 100644 --- a/internal/mock/vtctld_general.go +++ b/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)