diff --git a/internal/cmd/keyspace/keyspace.go b/internal/cmd/keyspace/keyspace.go index fb11e622..29119b82 100644 --- a/internal/cmd/keyspace/keyspace.go +++ b/internal/cmd/keyspace/keyspace.go @@ -55,6 +55,7 @@ type Keyspace struct { type KeyspaceSettings struct { ReplicationDurabilityConstraintStrategy string `header:"replication durability constraint strategy" json:"replication_durability_constraint"` VReplicationFlags VReplicationFlags `header:"inline" json:"vreplication_flags"` + Throttler Throttler `header:"inline" json:"throttler"` orig *ps.Keyspace } @@ -73,6 +74,11 @@ type VReplicationFlags struct { VPlayerBatching bool `header:"vplayer batching" json:"vplayer_batching"` } +type Throttler struct { + Enabled bool `header:"throttler enabled" json:"enabled"` + Threshold string `header:"throttler threshold" json:"threshold"` +} + func toKeyspaces(keyspaces []*ps.Keyspace) []*Keyspace { kss := make([]*Keyspace, 0, len(keyspaces)) diff --git a/internal/cmd/keyspace/settings.go b/internal/cmd/keyspace/settings.go index 11a9b4d8..9be93cfa 100644 --- a/internal/cmd/keyspace/settings.go +++ b/internal/cmd/keyspace/settings.go @@ -84,5 +84,15 @@ func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings { } } + if ks.Throttler != nil { + settings.Throttler = Throttler{ + Enabled: ks.Throttler.Enabled != nil && *ks.Throttler.Enabled, + Threshold: "not set", + } + if ks.Throttler.Threshold != nil { + settings.Throttler.Threshold = fmt.Sprintf("%gs", *ks.Throttler.Threshold) + } + } + return settings } diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index e92bef54..478b7f9e 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/internal/cmd/keyspace/update_settings.go @@ -2,7 +2,9 @@ package keyspace import ( "context" + "errors" "fmt" + "strconv" "github.com/charmbracelet/huh" "github.com/planetscale/cli/internal/cmdutil" @@ -17,6 +19,8 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { replicationDurabilityConstraints *ps.ReplicationDurabilityConstraints vreplicationFlags *ps.VReplicationFlags + throttlerEnabled bool + throttlerThreshold float64 interactive bool } @@ -84,7 +88,27 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { } } - if !rdcChanged && !vrfChanged { + throttlerChanged := cmd.Flags().Changed("throttler-enabled") || + cmd.Flags().Changed("throttler-threshold") + + if throttlerChanged { + if updateReq.Throttler == nil { + updateReq.Throttler = &ps.KeyspaceThrottler{} + } + + if cmd.Flags().Changed("throttler-enabled") { + updateReq.Throttler.Enabled = &flags.throttlerEnabled + } + + if cmd.Flags().Changed("throttler-threshold") { + if flags.throttlerThreshold < 0 { + return errors.New("--throttler-threshold must be greater than or equal to 0") + } + updateReq.Throttler.Threshold = &flags.throttlerThreshold + } + } + + if !rdcChanged && !vrfChanged && !throttlerChanged { end() ch.Printer.Println("No changes were requested. No update performed.") return nil @@ -105,6 +129,8 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { cmd.Flags().BoolVar(&flags.vreplicationFlags.OptimizeInserts, "vreplication-optimize-inserts", true, "When enabled, skips sending INSERT events for rows that have yet to be replicated.") cmd.Flags().BoolVar(&flags.vreplicationFlags.AllowNoBlobBinlogRowImage, "vreplication-enable-noblob-binlog-mode", true, "When enabled, omits changed BLOB and TEXT columns from replication events, which reduces binlog sizes.") cmd.Flags().BoolVar(&flags.vreplicationFlags.VPlayerBatching, "vreplication-batch-replication-events", false, "When enabled, sends fewer queries to MySQL to improve performance.") + cmd.Flags().BoolVar(&flags.throttlerEnabled, "throttler-enabled", true, "Pause schema migrations and VReplication workflows when replication lag rises above the threshold.") + cmd.Flags().Float64Var(&flags.throttlerThreshold, "throttler-threshold", 5, "Replication lag in seconds above which migrations and workflows are paused.") cmd.Flags().BoolVarP(&flags.interactive, "interactive", "i", false, "Run the command in interactive mode") return cmd @@ -145,6 +171,10 @@ func setInitialSettings(ctx context.Context, ch *cmdutil.Helper, req *ps.UpdateK req.VReplicationFlags = ks.VReplicationFlags } + if ks.Throttler != nil { + req.Throttler = ks.Throttler + } + return nil } @@ -166,6 +196,20 @@ func updateInteractive(ctx context.Context, ch *cmdutil.Helper, updateReq *ps.Up updateReq.VReplicationFlags = &ps.VReplicationFlags{} } + if updateReq.Throttler == nil { + updateReq.Throttler = &ps.KeyspaceThrottler{} + } + + throttlerEnabled := true + if updateReq.Throttler.Enabled != nil { + throttlerEnabled = *updateReq.Throttler.Enabled + } + + throttlerThreshold := "5" + if updateReq.Throttler.Threshold != nil { + throttlerThreshold = strconv.FormatFloat(*updateReq.Throttler.Threshold, 'g', -1, 64) + } + form := huh.NewForm( // Replication Durability Constraints huh.NewGroup( @@ -200,12 +244,41 @@ func updateInteractive(ctx context.Context, ch *cmdutil.Helper, updateReq *ps.Up Description("When enabled, sends fewer queries to MySQL to improve performance."). Value(&updateReq.VReplicationFlags.VPlayerBatching), ).Title("VReplication").Description("Options for improving performance during deploy requests and workflows"), + + huh.NewGroup( + huh.NewConfirm(). + Title("Enable the throttler?"). + Description("Pauses schema migrations and VReplication workflows when replication lag rises above the threshold."). + Value(&throttlerEnabled), + + huh.NewInput(). + Title("Replication lag threshold (seconds)"). + Description("Migrations and workflows are paused while replication lag is above this value."). + Value(&throttlerThreshold). + Validate(func(s string) error { + v, err := strconv.ParseFloat(s, 64) + if err != nil { + return errors.New("threshold must be a number") + } + if v < 0 { + return errors.New("threshold must be greater than or equal to 0") + } + return nil + }), + ).Title("Throttler"), ).WithTheme(huh.ThemeBase16()) if err := form.Run(); err != nil { return err } + threshold, err := strconv.ParseFloat(throttlerThreshold, 64) + if err != nil { + return err + } + updateReq.Throttler.Enabled = &throttlerEnabled + updateReq.Throttler.Threshold = &threshold + ks, err := updateKeyspaceSettings(ctx, client, updateReq) if err != nil { return err diff --git a/internal/cmd/keyspace/update_settings_test.go b/internal/cmd/keyspace/update_settings_test.go index fe2cccbf..d9f5f539 100644 --- a/internal/cmd/keyspace/update_settings_test.go +++ b/internal/cmd/keyspace/update_settings_test.go @@ -619,6 +619,233 @@ func TestKeyspace_UpdateSettingsCmd_PreserveNilValues(t *testing.T) { c.Assert(buf.String(), qt.JSONEquals, updatedKs) } +func boolPtr(b bool) *bool { return &b } + +// Regression: the API may return no throttler at all. A threshold-only update +// must not send enabled=false and silently turn the throttler off. +func TestKeyspace_UpdateSettingsCmd_ThresholdOnlyWithNoExistingThrottler(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + branch := "main" + keyspace := "sharded" + + ks := &ps.Keyspace{ID: "ks1", Name: keyspace, Throttler: nil} + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return ks, nil + }, + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Throttler, qt.Not(qt.IsNil)) + c.Assert(req.Throttler.Enabled, qt.IsNil) + c.Assert(*req.Throttler.Threshold, qt.Equals, 10.0) + + return ks, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{ + Organization: org, + }, + Client: func() (*ps.Client, error) { + return &ps.Client{ + Keyspaces: svc, + }, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{db, branch, keyspace, "--throttler-threshold=10"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) +} + +func TestKeyspace_UpdateSettingsCmd_DisableThrottler(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + branch := "main" + keyspace := "sharded" + + ts := time.Now() + threshold := 5.0 + + ks := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + Throttler: &ps.KeyspaceThrottler{Enabled: boolPtr(true), Threshold: &threshold}, + } + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + Throttler: &ps.KeyspaceThrottler{Enabled: boolPtr(false), Threshold: &threshold}, + } + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return ks, nil + }, + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Throttler, qt.Not(qt.IsNil)) + c.Assert(req.Throttler.Enabled, qt.Not(qt.IsNil)) + c.Assert(*req.Throttler.Enabled, qt.Equals, false) + c.Assert(req.Throttler.Threshold, qt.Not(qt.IsNil)) + c.Assert(*req.Throttler.Threshold, qt.Equals, 5.0) + + return updatedKs, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{ + Organization: org, + }, + Client: func() (*ps.Client, error) { + return &ps.Client{ + Keyspaces: svc, + }, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{db, branch, keyspace, "--throttler-enabled=false"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, updatedKs) +} + +func TestKeyspace_UpdateSettingsCmd_ThrottlerThresholdOnly(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + branch := "main" + keyspace := "sharded" + + ts := time.Now() + initial := 5.0 + updated := 10.0 + + ks := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + Throttler: &ps.KeyspaceThrottler{Enabled: boolPtr(true), Threshold: &initial}, + } + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + Throttler: &ps.KeyspaceThrottler{Enabled: boolPtr(true), Threshold: &updated}, + } + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return ks, nil + }, + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + // Changing only the threshold must not disable the throttler. + c.Assert(req.Throttler.Enabled, qt.Not(qt.IsNil)) + c.Assert(*req.Throttler.Enabled, qt.Equals, true) + c.Assert(*req.Throttler.Threshold, qt.Equals, 10.0) + + return updatedKs, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{ + Organization: org, + }, + Client: func() (*ps.Client, error) { + return &ps.Client{ + Keyspaces: svc, + }, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{db, branch, keyspace, "--throttler-threshold=10"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, updatedKs) +} + +func TestKeyspace_UpdateSettingsCmd_RejectsNegativeThrottlerThreshold(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + branch := "main" + keyspace := "sharded" + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: keyspace}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{ + Organization: org, + }, + Client: func() (*ps.Client, error) { + return &ps.Client{ + Keyspaces: svc, + }, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{db, branch, keyspace, "--throttler-threshold=-1"}) + err := cmd.Execute() + c.Assert(err, qt.ErrorMatches, ".*throttler-threshold must be greater than or equal to 0") + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse) +} + func TestKeyspace_ConstraintsToStrategy(t *testing.T) { c := qt.New(t) diff --git a/internal/planetscale/keyspaces.go b/internal/planetscale/keyspaces.go index 1b78de6b..0e8b1a23 100644 --- a/internal/planetscale/keyspaces.go +++ b/internal/planetscale/keyspaces.go @@ -24,6 +24,7 @@ type Keyspace struct { UpdatedAt time.Time `json:"updated_at"` VReplicationFlags *VReplicationFlags `json:"vreplication_flags"` ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints"` + Throttler *KeyspaceThrottler `json:"throttler"` ReadOnlyRegions []*ReadOnlyRegionKeyspace `json:"read_only_regions"` } @@ -174,6 +175,7 @@ type UpdateKeyspaceSettingsRequest struct { Keyspace string `json:"-"` ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints,omitempty"` VReplicationFlags *VReplicationFlags `json:"vreplication_flags,omitempty"` + Throttler *KeyspaceThrottler `json:"throttler,omitempty"` } type ReplicationDurabilityConstraints struct { @@ -186,6 +188,11 @@ type VReplicationFlags struct { VPlayerBatching bool `json:"vplayer_batching"` } +type KeyspaceThrottler struct { + Enabled *bool `json:"enabled,omitempty"` + Threshold *float64 `json:"threshold,omitempty"` +} + // KeyspacesService is an interface for interacting with the keyspace endpoints of the PlanetScale API type KeyspacesService interface { Create(context.Context, *CreateKeyspaceRequest) (*Keyspace, error)