diff --git a/go.mod b/go.mod index 52b1671a..9732b8b6 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.164.0 + github.com/planetscale/planetscale-go v0.165.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 12a75e5e..cb08e6b5 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.164.0 h1:k/Jw1robfLhcGS3f4FtqxQXEDuuddj0bnTE/jbAo4/o= -github.com/planetscale/planetscale-go v0.164.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= +github.com/planetscale/planetscale-go v0.165.0 h1:huSIreetdhlXwzSweuI1rB6hlAyYEWvZULjPmkCwTdg= +github.com/planetscale/planetscale-go v0.165.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/throttler.go b/internal/cmd/branch/vtctld/throttler.go new file mode 100644 index 00000000..b3ce4839 --- /dev/null +++ b/internal/cmd/branch/vtctld/throttler.go @@ -0,0 +1,198 @@ +package vtctld + +import ( + "fmt" + + "github.com/planetscale/cli/internal/cmdutil" + "github.com/planetscale/cli/internal/printer" + ps "github.com/planetscale/planetscale-go/planetscale" + "github.com/spf13/cobra" +) + +// ThrottlerCmd groups the tablet throttler subcommands. The throttler controls +// how aggressively Vitess applies background work (such as Online DDL and +// VReplication) based on replication lag and other metrics. +func ThrottlerCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "throttler ", + Short: "Inspect and configure the tablet throttler", + } + + cmd.AddCommand(ThrottlerStatusCmd(ch)) + cmd.AddCommand(ThrottlerCheckCmd(ch)) + cmd.AddCommand(ThrottlerUpdateConfigCmd(ch)) + + return cmd +} + +// ThrottlerStatusCmd reads the live throttler status from a single tablet. +func ThrottlerStatusCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + tabletAlias string + } + + cmd := &cobra.Command{ + Use: "status ", + Short: "Get the throttler status for a single tablet", + Long: "Get the throttler status for a single tablet, identified by its alias. " + + "Discover tablet aliases with `pscale branch vtctld list-tablets`.", + 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("Fetching throttler status for tablet %s on %s…", + printer.BoldBlue(flags.tabletAlias), progressTarget(ch.Config.Organization, database, branch))) + defer end() + + data, err := client.Vtctld.GetThrottlerStatus(ctx, &ps.VtctldGetThrottlerStatusRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + TabletAlias: flags.tabletAlias, + }) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrettyPrintJSON(data) + }, + } + + cmd.Flags().StringVar(&flags.tabletAlias, "tablet-alias", "", "Alias of the tablet to probe (e.g. \"zone1-0000000100\")") + cmd.MarkFlagRequired("tablet-alias") // nolint:errcheck + + return cmd +} + +// ThrottlerCheckCmd issues a throttler check against a single tablet. +func ThrottlerCheckCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + tabletAlias string + appName string + scope string + skipRequestHeartbeats bool + okIfNotExists bool + } + + cmd := &cobra.Command{ + Use: "check ", + Short: "Issue a throttler check against a single tablet", + Long: "Issue a throttler check against a single tablet, identified by its alias. " + + "Discover tablet aliases with `pscale branch vtctld list-tablets`.", + 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("Checking throttler for tablet %s on %s…", + printer.BoldBlue(flags.tabletAlias), progressTarget(ch.Config.Organization, database, branch))) + defer end() + + req := &ps.VtctldCheckThrottlerRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + TabletAlias: flags.tabletAlias, + AppName: flags.appName, + Scope: flags.scope, + } + if cmd.Flags().Changed("skip-request-heartbeats") { + req.SkipRequestHeartbeats = &flags.skipRequestHeartbeats + } + if cmd.Flags().Changed("ok-if-not-exists") { + req.OkIfNotExists = &flags.okIfNotExists + } + + data, err := client.Vtctld.CheckThrottler(ctx, req) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrettyPrintJSON(data) + }, + } + + cmd.Flags().StringVar(&flags.tabletAlias, "tablet-alias", "", "Alias of the tablet to check (e.g. \"zone1-0000000100\")") + cmd.Flags().StringVar(&flags.appName, "app-name", "", "App to issue the check on behalf of (e.g. \"online-ddl\"). Defaults to the throttler's default app.") + cmd.Flags().StringVar(&flags.scope, "scope", "", "Scope of the check, either \"shard\" or \"self\". Defaults to the throttler's default scope.") + cmd.Flags().BoolVar(&flags.skipRequestHeartbeats, "skip-request-heartbeats", false, "Do not renew the throttler's heartbeat lease while serving this check") + cmd.Flags().BoolVar(&flags.okIfNotExists, "ok-if-not-exists", false, "Return OK even if the requested metric does not exist") + cmd.MarkFlagRequired("tablet-alias") // nolint:errcheck + + return cmd +} + +// ThrottlerUpdateConfigCmd updates the tablet throttler configuration for a +// keyspace. +func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + keyspace string + enabled bool + threshold float64 + } + + cmd := &cobra.Command{ + Use: "update-config ", + Short: "Update the throttler configuration for a keyspace", + Long: "Update the tablet throttler configuration for a keyspace. The throttler is " + + "enabled or disabled with --enabled; this flag is required because there is no " + + "separate \"leave unchanged\" state.", + 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("Updating throttler config for keyspace %s on %s…", + printer.BoldBlue(flags.keyspace), progressTarget(ch.Config.Organization, database, branch))) + defer end() + + req := &ps.VtctldUpdateThrottlerConfigRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Keyspace: flags.keyspace, + Enabled: flags.enabled, + } + if cmd.Flags().Changed("threshold") { + req.Threshold = &flags.threshold + } + + data, err := client.Vtctld.UpdateThrottlerConfig(ctx, req) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return ch.Printer.PrettyPrintJSON(data) + }, + } + + cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Keyspace whose throttler config to update") + cmd.Flags().BoolVar(&flags.enabled, "enabled", false, "Enable (true) or disable (false) the throttler for the keyspace") + cmd.Flags().Float64Var(&flags.threshold, "threshold", 0, "Replication lag threshold in seconds for the default check (defaults to 5.0 server-side when omitted)") + cmd.MarkFlagRequired("keyspace") // nolint:errcheck + cmd.MarkFlagRequired("enabled") // nolint:errcheck + + return cmd +} diff --git a/internal/cmd/branch/vtctld/throttler_test.go b/internal/cmd/branch/vtctld/throttler_test.go new file mode 100644 index 00000000..5ef5c592 --- /dev/null +++ b/internal/cmd/branch/vtctld/throttler_test.go @@ -0,0 +1,184 @@ +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 newThrottlerTestHelper(org string, svc *mock.VtctldService) (*cmdutil.Helper, *bytes.Buffer) { + 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 + }, + } + return ch, &buf +} + +func TestThrottlerStatus(t *testing.T) { + c := qt.New(t) + + org, db, branch := "my-org", "my-db", "my-branch" + + svc := &mock.VtctldService{ + GetThrottlerStatusFn: func(ctx context.Context, req *ps.VtctldGetThrottlerStatusRequest) (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.TabletAlias, qt.Equals, "zone1-0000000100") + return json.RawMessage(`{"keyspace":"commerce","enabled":true}`), nil + }, + } + + ch, _ := newThrottlerTestHelper(org, svc) + + cmd := ThrottlerStatusCmd(ch) + cmd.SetArgs([]string{db, branch, "--tablet-alias", "zone1-0000000100"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetThrottlerStatusFnInvoked, qt.IsTrue) +} + +func TestThrottlerStatus_RequiresTabletAlias(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{} + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerStatusCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch"}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(svc.GetThrottlerStatusFnInvoked, qt.IsFalse) +} + +func TestThrottlerCheck(t *testing.T) { + c := qt.New(t) + + org, db, branch := "my-org", "my-db", "my-branch" + + svc := &mock.VtctldService{ + CheckThrottlerFn: func(ctx context.Context, req *ps.VtctldCheckThrottlerRequest) (json.RawMessage, error) { + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.TabletAlias, qt.Equals, "zone1-0000000100") + c.Assert(req.AppName, qt.Equals, "online-ddl") + c.Assert(req.Scope, qt.Equals, "self") + c.Assert(req.SkipRequestHeartbeats, qt.IsNotNil) + c.Assert(*req.SkipRequestHeartbeats, qt.IsTrue) + return json.RawMessage(`{"response_code":"THROTTLER_RESPONSE_CODE_OK"}`), nil + }, + } + + ch, _ := newThrottlerTestHelper(org, svc) + + cmd := ThrottlerCheckCmd(ch) + cmd.SetArgs([]string{db, branch, + "--tablet-alias", "zone1-0000000100", + "--app-name", "online-ddl", + "--scope", "self", + "--skip-request-heartbeats", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.CheckThrottlerFnInvoked, qt.IsTrue) +} + +func TestThrottlerCheck_OmitsUnsetOptionalBools(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{ + CheckThrottlerFn: func(ctx context.Context, req *ps.VtctldCheckThrottlerRequest) (json.RawMessage, error) { + // Unset bool flags stay nil so the server applies its defaults. + c.Assert(req.SkipRequestHeartbeats, qt.IsNil) + c.Assert(req.OkIfNotExists, qt.IsNil) + return json.RawMessage(`{"response_code":"THROTTLER_RESPONSE_CODE_OK"}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerCheckCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", "--tablet-alias", "zone1-0000000100"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.CheckThrottlerFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig(t *testing.T) { + c := qt.New(t) + + org, db, branch := "my-org", "my-db", "my-branch" + + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.Keyspace, qt.Equals, "commerce") + c.Assert(req.Enabled, qt.IsTrue) + c.Assert(req.Threshold, qt.IsNotNil) + c.Assert(*req.Threshold, qt.Equals, 2.5) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper(org, svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{db, branch, + "--keyspace", "commerce", + "--enabled", + "--threshold", "2.5", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Enabled, qt.IsFalse) + // Threshold stays nil when not provided so the server keeps its default. + c.Assert(req.Threshold, qt.IsNil) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", "--keyspace", "commerce", "--enabled=false"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig_RequiresEnabled(t *testing.T) { + c := qt.New(t) + + svc := &mock.VtctldService{} + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", "--keyspace", "commerce"}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsFalse) +} diff --git a/internal/cmd/branch/vtctld/vtctld.go b/internal/cmd/branch/vtctld/vtctld.go index 9ccf01a9..bff983fa 100644 --- a/internal/cmd/branch/vtctld/vtctld.go +++ b/internal/cmd/branch/vtctld/vtctld.go @@ -23,6 +23,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command { cmd.AddCommand(ListTabletsCmd(ch)) cmd.AddCommand(StartWorkflowCmd(ch)) cmd.AddCommand(StopWorkflowCmd(ch)) + cmd.AddCommand(ThrottlerCmd(ch)) return cmd } diff --git a/internal/mock/vtctld_general.go b/internal/mock/vtctld_general.go index fe94b8d4..8b985259 100644 --- a/internal/mock/vtctld_general.go +++ b/internal/mock/vtctld_general.go @@ -23,6 +23,15 @@ type VtctldService struct { StopWorkflowFn func(context.Context, *ps.VtctldStopWorkflowRequest) (json.RawMessage, error) StopWorkflowFnInvoked bool + GetThrottlerStatusFn func(context.Context, *ps.VtctldGetThrottlerStatusRequest) (json.RawMessage, error) + GetThrottlerStatusFnInvoked bool + + CheckThrottlerFn func(context.Context, *ps.VtctldCheckThrottlerRequest) (json.RawMessage, error) + CheckThrottlerFnInvoked bool + + UpdateThrottlerConfigFn func(context.Context, *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) + UpdateThrottlerConfigFnInvoked bool + GetOperationFn func(context.Context, *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) GetOperationFnInvoked bool } @@ -52,6 +61,21 @@ func (s *VtctldService) StopWorkflow(ctx context.Context, req *ps.VtctldStopWork return s.StopWorkflowFn(ctx, req) } +func (s *VtctldService) GetThrottlerStatus(ctx context.Context, req *ps.VtctldGetThrottlerStatusRequest) (json.RawMessage, error) { + s.GetThrottlerStatusFnInvoked = true + return s.GetThrottlerStatusFn(ctx, req) +} + +func (s *VtctldService) CheckThrottler(ctx context.Context, req *ps.VtctldCheckThrottlerRequest) (json.RawMessage, error) { + s.CheckThrottlerFnInvoked = true + return s.CheckThrottlerFn(ctx, req) +} + +func (s *VtctldService) UpdateThrottlerConfig(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + s.UpdateThrottlerConfigFnInvoked = true + return s.UpdateThrottlerConfigFn(ctx, req) +} + func (s *VtctldService) GetOperation(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) { s.GetOperationFnInvoked = true return s.GetOperationFn(ctx, req)