Add customer keyspace rollout concurrency 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"`
+	MaxRollout                              string            `header:"max rollout" json:"max_rollout"`
 	Throttler                               Throttler         `header:"inline" json:"throttler"`
 
 	orig *ps.Keyspace

internal/cmd/keyspace/settings.go

@@ -2,6 +2,7 @@ package keyspace
 
 import (
 	"fmt"
+	"strconv"
 
 	"github.com/planetscale/cli/internal/cmdutil"
 	ps "github.com/planetscale/cli/internal/planetscale"
@@ -53,7 +54,12 @@ func SettingsCmd(ch *cmdutil.Helper) *cobra.Command {
 // toKeyspaceSettings converts a Keyspace API response to a KeyspaceSettings object for display
 func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings {
 	settings := &KeyspaceSettings{
-		orig: ks,
+		MaxRollout: "not set",
+		orig:       ks,
+	}
+
+	if ks.MaxRollout != nil {
+		settings.MaxRollout = strconv.Itoa(*ks.MaxRollout)
 	}
 
 	// Set replication durability constraints if available

internal/cmd/keyspace/settings_test.go

@@ -200,8 +200,12 @@ func TestBuildKeyspaceSettings(t *testing.T) {
 		},
 	}
 
+	maxRollout := 8
+	fullKs.MaxRollout = &maxRollout
+
 	settings := toKeyspaceSettings(fullKs)
 	c.Assert(settings.ReplicationDurabilityConstraintStrategy, qt.Equals, "maximum") // Should be translated
+	c.Assert(settings.MaxRollout, qt.Equals, "8")
 	c.Assert(settings.VReplicationFlags.OptimizeInserts, qt.Equals, true)
 	c.Assert(settings.VReplicationFlags.AllowNoBlobBinlogRowImage, qt.Equals, true)
 	c.Assert(settings.VReplicationFlags.VPlayerBatching, qt.Equals, false)
@@ -218,6 +222,7 @@ func TestBuildKeyspaceSettings(t *testing.T) {
 
 	nilSettings := toKeyspaceSettings(nilKs)
 	c.Assert(nilSettings.ReplicationDurabilityConstraintStrategy, qt.Equals, "not set")
+	c.Assert(nilSettings.MaxRollout, qt.Equals, "not set")
 	c.Assert(nilSettings.VReplicationFlags.OptimizeInserts, qt.Equals, false) // Default values
 	c.Assert(nilSettings.VReplicationFlags.AllowNoBlobBinlogRowImage, qt.Equals, false)
 	c.Assert(nilSettings.VReplicationFlags.VPlayerBatching, qt.Equals, false)

internal/cmd/keyspace/update_settings.go

@@ -21,6 +21,7 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command {
 		vreplicationFlags                *ps.VReplicationFlags
 		throttlerEnabled                 bool
 		throttlerThreshold               float64
+		maxRollout                       int
 		interactive                      bool
 	}
 
@@ -108,7 +109,16 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command {
 				}
 			}
 
-			if !rdcChanged && !vrfChanged && !throttlerChanged {
+			maxRolloutChanged := cmd.Flags().Changed("max-rollout")
+
+			if maxRolloutChanged {
+				if flags.maxRollout < 1 || flags.maxRollout > 32 {
+					return errors.New("--max-rollout must be between 1 and 32")
+				}
+				updateReq.MaxRollout = &flags.maxRollout
+			}
+
+			if !rdcChanged && !vrfChanged && !throttlerChanged && !maxRolloutChanged {
 				end()
 				ch.Printer.Println("No changes were requested. No update performed.")
 				return nil
@@ -131,6 +141,7 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command {
 	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().IntVar(&flags.maxRollout, "max-rollout", 1, "Maximum number of shards to roll out changes to concurrently (1-32).")
 	cmd.Flags().BoolVarP(&flags.interactive, "interactive", "i", false, "Run the command in interactive mode")
 
 	return cmd

internal/cmd/keyspace/update_settings_test.go

@@ -846,6 +846,149 @@ func TestKeyspace_UpdateSettingsCmd_RejectsNegativeThrottlerThreshold(t *testing
 	c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse)
 }
 
+func TestKeyspace_UpdateSettingsCmd_MaxRollout(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()
+	maxRollout := 8
+
+	updatedKs := &ps.Keyspace{
+		ID:         "ks1",
+		Name:       keyspace,
+		CreatedAt:  ts,
+		UpdatedAt:  ts,
+		MaxRollout: &maxRollout,
+	}
+
+	svc := &mock.KeyspacesService{
+		GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) {
+			return &ps.Keyspace{ID: "ks1", Name: keyspace}, nil
+		},
+		UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) {
+			c.Assert(req.MaxRollout, qt.Not(qt.IsNil))
+			c.Assert(*req.MaxRollout, qt.Equals, 8)
+			c.Assert(req.Throttler, qt.IsNil)
+
+			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, "--max-rollout=8"})
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, updatedKs)
+}
+
+func TestKeyspace_UpdateSettingsCmd_RejectsOutOfRangeMaxRollout(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"
+
+	for _, arg := range []string{"--max-rollout=0", "--max-rollout=33"} {
+		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, arg})
+		err := cmd.Execute()
+		c.Assert(err, qt.ErrorMatches, ".*max-rollout must be between 1 and 32")
+		c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse)
+	}
+}
+
+func TestKeyspace_UpdateSettingsCmd_OmittedMaxRolloutIsNotSent(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
+		},
+		UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) {
+			c.Assert(req.MaxRollout, qt.IsNil)
+
+			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=10"})
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue)
+}
+
 func TestKeyspace_ConstraintsToStrategy(t *testing.T) {
 	c := qt.New(t)
 

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"`
+	MaxRollout                       *int                              `json:"max_rollout"`
 	Throttler                        *KeyspaceThrottler                `json:"throttler"`
 	ReadOnlyRegions                  []*ReadOnlyRegionKeyspace         `json:"read_only_regions"`
 }
@@ -176,6 +177,7 @@ type UpdateKeyspaceSettingsRequest struct {
 	ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints,omitempty"`
 	VReplicationFlags                *VReplicationFlags                `json:"vreplication_flags,omitempty"`
 	Throttler                        *KeyspaceThrottler                `json:"throttler,omitempty"`
+	MaxRollout                       *int                              `json:"max_rollout,omitempty"`
 }
 
 type ReplicationDurabilityConstraints struct {

internal/planetscale/keyspaces_test.go

@@ -3,8 +3,10 @@ package planetscale
 import (
 	"context"
 	"encoding/json"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	qt "github.com/frankban/quicktest"
@@ -479,3 +481,41 @@ func TestKeyspaces_UpdateSettings(t *testing.T) {
 	c.Assert(keyspace.VReplicationFlags.VPlayerBatching, qt.Equals, true)
 	c.Assert(keyspace.ReplicationDurabilityConstraints.Strategy, qt.Equals, "maximum")
 }
+
+func TestKeyspaces_UpdateSettingsMaxRollout(t *testing.T) {
+	c := qt.New(t)
+
+	var body string
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, qt.Equals, http.MethodPatch)
+
+		raw, err := io.ReadAll(r.Body)
+		c.Assert(err, qt.IsNil)
+		body = string(raw)
+
+		w.WriteHeader(200)
+		out := `{"type":"Keyspace","id":"thisisanid","name":"planetscale","max_rollout":8}`
+		_, err = w.Write([]byte(out))
+		c.Assert(err, qt.IsNil)
+	}))
+
+	client, err := NewClient(WithBaseURL(ts.URL))
+	c.Assert(err, qt.IsNil)
+
+	ctx := context.Background()
+	maxRollout := 8
+
+	keyspace, err := client.Keyspaces.UpdateSettings(ctx, &UpdateKeyspaceSettingsRequest{
+		Organization: "foo",
+		Database:     "bar",
+		Branch:       "baz",
+		Keyspace:     "qux",
+		MaxRollout:   &maxRollout,
+	})
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(strings.TrimSpace(body), qt.Equals, `{"max_rollout":8}`)
+	c.Assert(keyspace.MaxRollout, qt.Not(qt.IsNil))
+	c.Assert(*keyspace.MaxRollout, qt.Equals, 8)
+}