diff --git a/internal/cmd/keyspace/keyspace.go b/internal/cmd/keyspace/keyspace.go index 29119b82..7a59d599 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"` + MaxRollout string `header:"max rollout" json:"max_rollout"` Throttler Throttler `header:"inline" json:"throttler"` orig *ps.Keyspace diff --git a/internal/cmd/keyspace/settings.go b/internal/cmd/keyspace/settings.go index 9be93cfa..233fe6d6 100644 --- a/internal/cmd/keyspace/settings.go +++ b/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 diff --git a/internal/cmd/keyspace/settings_test.go b/internal/cmd/keyspace/settings_test.go index f71ccb55..07cb5f44 100644 --- a/internal/cmd/keyspace/settings_test.go +++ b/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) diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index 478b7f9e..69156f65 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/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 diff --git a/internal/cmd/keyspace/update_settings_test.go b/internal/cmd/keyspace/update_settings_test.go index d9f5f539..80208167 100644 --- a/internal/cmd/keyspace/update_settings_test.go +++ b/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) diff --git a/internal/planetscale/keyspaces.go b/internal/planetscale/keyspaces.go index 0e8b1a23..3d60c78b 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"` + 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 { diff --git a/internal/planetscale/keyspaces_test.go b/internal/planetscale/keyspaces_test.go index 792ec397..993e983c 100644 --- a/internal/planetscale/keyspaces_test.go +++ b/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) +}