Add throttle-app, unthrottle-app, and app-metrics to vtctld throttler update-config

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,33 @@ 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 <database> <branch>",
 		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. " +
+			"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 {
+			// Mirrors vtctldclient validateUpdateThrottlerConfig.
+			if cmd.Flags().Changed("app-name") && flags.appName == "" {
+				return fmt.Errorf("--app-name must not be empty")
+			}
+			return nil
+		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			database, branch := args[0], args[1]
@@ -172,11 +188,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 +225,17 @@ 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", 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.MarkFlagRequired("enabled")  // nolint:errcheck
+	cmd.MarkFlagsMutuallyExclusive("unthrottle-app", "throttle-app")
+	cmd.MarkFlagsRequiredTogether("app-name", "app-metrics")
 
 	return cmd
 }

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,15 +171,82 @@ func TestThrottlerUpdateConfig_DisableOmitsThreshold(t *testing.T) {
 	c.Assert(svc.UpdateThrottlerConfigFnInvoked, qt.IsTrue)
 }
 
-func TestThrottlerUpdateConfig_RequiresEnabled(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) {
+	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)
 }

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 {

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