From a54a7155028c2b154ce8b113c56fec9342c065e4 Mon Sep 17 00:00:00 2001 From: Max Englander Date: Tue, 28 Jul 2026 16:11:49 -0400 Subject: [PATCH 1/3] Add throttle-app, unthrottle-app, and app-metrics to vtctld throttler update-config. Matches Inventory2 Staging Runbook vtctldclient UpdateThrottlerConfig flags so those steps can run via pscale without pskube. Co-authored-by: Cursor Signed-off-by: Max Englander --- internal/cmd/branch/vtctld/throttler.go | 77 +++++++++++++++++--- internal/cmd/branch/vtctld/throttler_test.go | 66 ++++++++++++++++- internal/planetscale/vtctld_general.go | 15 ++-- internal/planetscale/vtctld_general_test.go | 6 +- 4 files changed, 142 insertions(+), 22 deletions(-) diff --git a/internal/cmd/branch/vtctld/throttler.go b/internal/cmd/branch/vtctld/throttler.go index ed7228d9..b9a4299b 100644 --- a/internal/cmd/branch/vtctld/throttler.go +++ b/internal/cmd/branch/vtctld/throttler.go @@ -2,6 +2,7 @@ package vtctld import ( "fmt" + "time" "github.com/planetscale/cli/internal/cmdutil" ps "github.com/planetscale/cli/internal/planetscale" @@ -141,18 +142,49 @@ func ThrottlerCheckCmd(ch *cmdutil.Helper) *cobra.Command { // keyspace. func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { - keyspace string - enabled bool - threshold float64 + keyspace string + enabled bool + threshold float64 + throttleApp string + throttleAppRatio float64 + throttleAppDuration time.Duration + unthrottleApp string + appName string + appMetrics []string } 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.", + Long: "Update the tablet throttler configuration for a keyspace. " + + "Omit --enabled to leave the keyspace enable state unchanged (for " + + "example when only throttling or unthrottling an app). " + + "--throttle-app and --unthrottle-app are mutually exclusive. " + + "--app-name and --app-metrics are required together.", Args: cmdutil.RequiredArgs("database", "branch"), + PreRunE: func(cmd *cobra.Command, args []string) error { + if flags.throttleApp != "" && flags.unthrottleApp != "" { + return fmt.Errorf("--throttle-app and --unthrottle-app are mutually exclusive") + } + appNameSet := cmd.Flags().Changed("app-name") + appMetricsSet := cmd.Flags().Changed("app-metrics") + if appNameSet != appMetricsSet { + return fmt.Errorf("--app-name and --app-metrics are required together") + } + if appNameSet && flags.appName == "" { + return fmt.Errorf("--app-name must not be empty") + } + + hasMutation := cmd.Flags().Changed("enabled") || + cmd.Flags().Changed("threshold") || + flags.throttleApp != "" || + flags.unthrottleApp != "" || + appNameSet + if !hasMutation { + return fmt.Errorf("at least one of --enabled, --threshold, --throttle-app, --unthrottle-app, or --app-name/--app-metrics is required") + } + return nil + }, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() database, branch := args[0], args[1] @@ -172,11 +204,31 @@ func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { Database: database, Branch: branch, Keyspace: flags.keyspace, - Enabled: flags.enabled, + } + if cmd.Flags().Changed("enabled") { + enabled := flags.enabled + req.Enabled = &enabled } if cmd.Flags().Changed("threshold") { req.Threshold = &flags.threshold } + if flags.throttleApp != "" { + ratio := flags.throttleAppRatio + expireAt := time.Now().Add(flags.throttleAppDuration).UTC().Format(time.RFC3339) + req.Apps = []ps.VtctldThrottledAppConfig{{ + Name: flags.throttleApp, + Ratio: &ratio, + ExpireAt: expireAt, + }} + } + if flags.unthrottleApp != "" { + req.UnthrottleApps = []string{flags.unthrottleApp} + } + if flags.appName != "" { + req.AppCheckedMetrics = map[string][]string{ + flags.appName: flags.appMetrics, + } + } data, err := client.Vtctld.UpdateThrottlerConfig(ctx, req) if err != nil { @@ -189,10 +241,15 @@ func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { } 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.Flags().BoolVar(&flags.enabled, "enabled", false, "Enable (true) or disable (false) the throttler for the keyspace. Omit to leave unchanged.") + cmd.Flags().Float64Var(&flags.threshold, "threshold", 0, "Replication lag threshold in seconds for the default check") + cmd.Flags().StringVar(&flags.throttleApp, "throttle-app", "", "App name to throttle (e.g. \"rowstreamer\")") + cmd.Flags().Float64Var(&flags.throttleAppRatio, "throttle-app-ratio", 0.5, "Ratio to throttle the app specified by --throttle-app (0.00-1.00)") + cmd.Flags().DurationVar(&flags.throttleAppDuration, "throttle-app-duration", 96*time.Hour, "Duration after which the --throttle-app rule expires") + cmd.Flags().StringVar(&flags.unthrottleApp, "unthrottle-app", "", "App name whose throttled-app rule should be removed") + cmd.Flags().StringVar(&flags.appName, "app-name", "", "App name for which to assign checked metrics (requires --app-metrics)") + cmd.Flags().StringSliceVar(&flags.appMetrics, "app-metrics", nil, "Metrics to check for --app-name (e.g. lag,loadavg)") 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 index e2d734e6..9ac3529c 100644 --- a/internal/cmd/branch/vtctld/throttler_test.go +++ b/internal/cmd/branch/vtctld/throttler_test.go @@ -129,7 +129,8 @@ func TestThrottlerUpdateConfig(t *testing.T) { 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.Enabled, qt.IsNotNil) + c.Assert(*req.Enabled, qt.IsTrue) c.Assert(req.Threshold, qt.IsNotNil) c.Assert(*req.Threshold, qt.Equals, 2.5) return json.RawMessage(`{}`), nil @@ -154,8 +155,8 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.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.Enabled, qt.IsNotNil) + c.Assert(*req.Enabled, qt.IsFalse) c.Assert(req.Threshold, qt.IsNil) return json.RawMessage(`{}`), nil }, @@ -170,7 +171,7 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) { c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) } -func TestThrottlerUpdateConfig_RequiresEnabled(t *testing.T) { +func TestThrottlerUpdateConfig_RequiresMutation(t *testing.T) { c := qt.New(t) svc := &mock.VtctldService{} @@ -182,3 +183,60 @@ func TestThrottlerUpdateConfig_RequiresEnabled(t *testing.T) { c.Assert(err, qt.IsNotNil) c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsFalse) } + +func TestThrottlerUpdateConfig_ThrottleAppWithMetrics(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.IsNil) + c.Assert(len(req.Apps), qt.Equals, 1) + c.Assert(req.Apps[0].Name, qt.Equals, "rowstreamer") + c.Assert(req.Apps[0].Ratio, qt.IsNotNil) + c.Assert(*req.Apps[0].Ratio, qt.Equals, 0.5) + c.Assert(req.Apps[0].ExpireAt, qt.Not(qt.Equals), "") + c.Assert(req.AppCheckedMetrics, qt.DeepEquals, map[string][]string{ + "rowstreamer": {"lag", "loadavg"}, + }) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", + "--keyspace", "commerce", + "--throttle-app", "rowstreamer", + "--throttle-app-ratio", "0.5", + "--throttle-app-duration", "96h", + "--app-name", "rowstreamer", + "--app-metrics", "lag,loadavg", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} + +func TestThrottlerUpdateConfig_UnthrottleApp(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.IsNil) + c.Assert(req.UnthrottleApps, qt.DeepEquals, []string{"rowstreamer"}) + return json.RawMessage(`{}`), nil + }, + } + + ch, _ := newThrottlerTestHelper("my-org", svc) + + cmd := ThrottlerUpdateConfigCmd(ch) + cmd.SetArgs([]string{"my-db", "my-branch", + "--keyspace", "commerce", + "--unthrottle-app", "rowstreamer", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) +} diff --git a/internal/planetscale/vtctld_general.go b/internal/planetscale/vtctld_general.go index 5ea3ecfe..28026dc1 100644 --- a/internal/planetscale/vtctld_general.go +++ b/internal/planetscale/vtctld_general.go @@ -164,14 +164,19 @@ type VtctldUpdateThrottlerConfigRequest struct { Keyspace string `json:"keyspace"` // Enabled controls whether the tablet throttler is enabled for the keyspace. - // It is required: the API has no tri-state, so omitting it would disable the - // throttler. - Enabled bool `json:"enabled"` - // Threshold is the threshold for the default throttler check (replication lag - // in seconds). It must be >= 0; the server defaults to 5.0 when omitted. + // Omit to leave the enable state unchanged. + Enabled *bool `json:"enabled,omitempty"` + // Threshold is the threshold for the default throttler check (replication + // lag in seconds). It must be >= 0. Omit to leave the existing threshold + // unchanged. Threshold *float64 `json:"threshold,omitempty"` // Apps configures zero or more per-app throttling rules. Apps []VtctldThrottledAppConfig `json:"apps,omitempty"` + // UnthrottleApps names apps whose throttled-app rules should be removed. + UnthrottleApps []string `json:"unthrottle_apps,omitempty"` + // AppCheckedMetrics assigns metrics to check per app name + // (e.g. {"rowstreamer": ["lag", "loadavg"]}). + AppCheckedMetrics map[string][]string `json:"app_checked_metrics,omitempty"` } type vtctldService struct { diff --git a/internal/planetscale/vtctld_general_test.go b/internal/planetscale/vtctld_general_test.go index bfc2285e..ad28f310 100644 --- a/internal/planetscale/vtctld_general_test.go +++ b/internal/planetscale/vtctld_general_test.go @@ -520,7 +520,7 @@ func TestVtctld_UpdateThrottlerConfig(t *testing.T) { Database: "my-db", Branch: "my-branch", Keyspace: "commerce", - Enabled: true, + Enabled: boolPtr(true), Threshold: float64Ptr(2.5), Apps: []VtctldThrottledAppConfig{ {Name: "online-ddl", Ratio: float64Ptr(0.5)}, @@ -562,7 +562,7 @@ func TestVtctld_UpdateThrottlerConfig_DisableSendsEnabledFalse(t *testing.T) { Database: "my-db", Branch: "my-branch", Keyspace: "commerce", - Enabled: false, + Enabled: boolPtr(false), }) c.Assert(err, qt.IsNil) -} +} \ No newline at end of file From 7e240fc86d02677a1ac63bc2975a8ef718b44c3a Mon Sep 17 00:00:00 2001 From: Max Englander Date: Tue, 28 Jul 2026 16:40:46 -0400 Subject: [PATCH 2/3] Drop UpdateThrottlerConfig mutation-required guard; match vtctldclient defaults. vtctldclient allows keyspace-only calls and uses ratio 1.0 / duration 1h defaults. Use cobra MarkFlagsMutuallyExclusive / RequiredTogether for the same flag constraints. Co-authored-by: Cursor Signed-off-by: Max Englander --- internal/cmd/branch/vtctld/throttler.go | 50 +++++++------------- internal/cmd/branch/vtctld/throttler_test.go | 18 +++++-- 2 files changed, 32 insertions(+), 36 deletions(-) diff --git a/internal/cmd/branch/vtctld/throttler.go b/internal/cmd/branch/vtctld/throttler.go index b9a4299b..1770f3ca 100644 --- a/internal/cmd/branch/vtctld/throttler.go +++ b/internal/cmd/branch/vtctld/throttler.go @@ -142,47 +142,31 @@ func ThrottlerCheckCmd(ch *cmdutil.Helper) *cobra.Command { // keyspace. func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { - keyspace string - enabled bool - threshold float64 - throttleApp string - throttleAppRatio float64 - throttleAppDuration time.Duration - unthrottleApp string - appName string - appMetrics []string + keyspace string + enabled bool + threshold float64 + throttleApp string + throttleAppRatio float64 + throttleAppDuration time.Duration + unthrottleApp string + appName string + appMetrics []string } cmd := &cobra.Command{ Use: "update-config ", Short: "Update the throttler configuration for a keyspace", Long: "Update the tablet throttler configuration for a keyspace. " + - "Omit --enabled to leave the keyspace enable state unchanged (for " + - "example when only throttling or unthrottling an app). " + - "--throttle-app and --unthrottle-app are mutually exclusive. " + + "Omit --enabled to leave the keyspace enable state unchanged. " + + "Flag behavior mirrors vtctldclient UpdateThrottlerConfig: " + + "--throttle-app and --unthrottle-app are mutually exclusive; " + "--app-name and --app-metrics are required together.", Args: cmdutil.RequiredArgs("database", "branch"), PreRunE: func(cmd *cobra.Command, args []string) error { - if flags.throttleApp != "" && flags.unthrottleApp != "" { - return fmt.Errorf("--throttle-app and --unthrottle-app are mutually exclusive") - } - appNameSet := cmd.Flags().Changed("app-name") - appMetricsSet := cmd.Flags().Changed("app-metrics") - if appNameSet != appMetricsSet { - return fmt.Errorf("--app-name and --app-metrics are required together") - } - if appNameSet && flags.appName == "" { + // Mirrors vtctldclient validateUpdateThrottlerConfig. + if cmd.Flags().Changed("app-name") && flags.appName == "" { return fmt.Errorf("--app-name must not be empty") } - - hasMutation := cmd.Flags().Changed("enabled") || - cmd.Flags().Changed("threshold") || - flags.throttleApp != "" || - flags.unthrottleApp != "" || - appNameSet - if !hasMutation { - return fmt.Errorf("at least one of --enabled, --threshold, --throttle-app, --unthrottle-app, or --app-name/--app-metrics is required") - } return nil }, RunE: func(cmd *cobra.Command, args []string) error { @@ -244,12 +228,14 @@ func ThrottlerUpdateConfigCmd(ch *cmdutil.Helper) *cobra.Command { cmd.Flags().BoolVar(&flags.enabled, "enabled", false, "Enable (true) or disable (false) the throttler for the keyspace. Omit to leave unchanged.") cmd.Flags().Float64Var(&flags.threshold, "threshold", 0, "Replication lag threshold in seconds for the default check") cmd.Flags().StringVar(&flags.throttleApp, "throttle-app", "", "App name to throttle (e.g. \"rowstreamer\")") - cmd.Flags().Float64Var(&flags.throttleAppRatio, "throttle-app-ratio", 0.5, "Ratio to throttle the app specified by --throttle-app (0.00-1.00)") - cmd.Flags().DurationVar(&flags.throttleAppDuration, "throttle-app-duration", 96*time.Hour, "Duration after which the --throttle-app rule expires") + cmd.Flags().Float64Var(&flags.throttleAppRatio, "throttle-app-ratio", 1.0, "Ratio to throttle the app specified by --throttle-app (0.00-1.00)") + cmd.Flags().DurationVar(&flags.throttleAppDuration, "throttle-app-duration", time.Hour, "Duration after which the --throttle-app rule expires") cmd.Flags().StringVar(&flags.unthrottleApp, "unthrottle-app", "", "App name whose throttled-app rule should be removed") cmd.Flags().StringVar(&flags.appName, "app-name", "", "App name for which to assign checked metrics (requires --app-metrics)") cmd.Flags().StringSliceVar(&flags.appMetrics, "app-metrics", nil, "Metrics to check for --app-name (e.g. lag,loadavg)") cmd.MarkFlagRequired("keyspace") // nolint:errcheck + cmd.MarkFlagsMutuallyExclusive("unthrottle-app", "throttle-app") + cmd.MarkFlagsRequiredTogether("app-name", "app-metrics") return cmd } diff --git a/internal/cmd/branch/vtctld/throttler_test.go b/internal/cmd/branch/vtctld/throttler_test.go index 9ac3529c..61422547 100644 --- a/internal/cmd/branch/vtctld/throttler_test.go +++ b/internal/cmd/branch/vtctld/throttler_test.go @@ -171,17 +171,27 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) { c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) } -func TestThrottlerUpdateConfig_RequiresMutation(t *testing.T) { +func TestThrottlerUpdateConfig_KeyspaceOnlyAllowed(t *testing.T) { c := qt.New(t) - svc := &mock.VtctldService{} + svc := &mock.VtctldService{ + UpdateThrottlerConfigFn: func(ctx context.Context, req *ps.VtctldUpdateThrottlerConfigRequest) (json.RawMessage, error) { + c.Assert(req.Keyspace, qt.Equals, "commerce") + c.Assert(req.Enabled, qt.IsNil) + c.Assert(req.Threshold, qt.IsNil) + c.Assert(req.Apps, qt.HasLen, 0) + c.Assert(req.UnthrottleApps, qt.HasLen, 0) + c.Assert(req.AppCheckedMetrics, qt.IsNil) + return json.RawMessage(`{}`), nil + }, + } 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) + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue) } func TestThrottlerUpdateConfig_ThrottleAppWithMetrics(t *testing.T) { From 98531e64d07d774aa3acbc4b2c80a802758d5973 Mon Sep 17 00:00:00 2001 From: Max Englander Date: Tue, 28 Jul 2026 17:14:46 -0400 Subject: [PATCH 3/3] Add trailing newline to vtctld_general_test.go. Co-authored-by: Cursor Signed-off-by: Max Englander --- internal/planetscale/vtctld_general_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/planetscale/vtctld_general_test.go b/internal/planetscale/vtctld_general_test.go index ad28f310..738b23fa 100644 --- a/internal/planetscale/vtctld_general_test.go +++ b/internal/planetscale/vtctld_general_test.go @@ -565,4 +565,4 @@ func TestVtctld_UpdateThrottlerConfig_DisableSendsEnabledFalse(t *testing.T) { Enabled: boolPtr(false), }) c.Assert(err, qt.IsNil) -} \ No newline at end of file +}