Add throttler flags to pscale keyspace update-settings

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

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
 }

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