diff --git a/internal/cmd/keyspace/keyspace.go b/internal/cmd/keyspace/keyspace.go index fb11e622..b61d563d 100644 --- a/internal/cmd/keyspace/keyspace.go +++ b/internal/cmd/keyspace/keyspace.go @@ -55,6 +55,8 @@ type Keyspace struct { type KeyspaceSettings struct { ReplicationDurabilityConstraintStrategy string `header:"replication durability constraint strategy" json:"replication_durability_constraint"` VReplicationFlags VReplicationFlags `header:"inline" json:"vreplication_flags"` + DiskScalingStrategy string `header:"disk scaling strategy" json:"disk_scaling_strategy"` + MaxStorageBytes int64 `header:"max storage bytes" json:"max_storage_bytes"` orig *ps.Keyspace } diff --git a/internal/cmd/keyspace/settings.go b/internal/cmd/keyspace/settings.go index 11a9b4d8..1319baf0 100644 --- a/internal/cmd/keyspace/settings.go +++ b/internal/cmd/keyspace/settings.go @@ -84,5 +84,13 @@ func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings { } } + // Set disk autoscaling settings if available + if ks.DiskAutoscaling != nil { + settings.DiskScalingStrategy = ks.DiskAutoscaling.Strategy + settings.MaxStorageBytes = ks.DiskAutoscaling.StorageLimitBytes + } else { + settings.DiskScalingStrategy = "not set" + } + return settings } diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index e92bef54..fad77552 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/internal/cmd/keyspace/update_settings.go @@ -3,20 +3,29 @@ package keyspace import ( "context" "fmt" + "slices" + "strings" "github.com/charmbracelet/huh" + "github.com/spf13/cobra" + "github.com/planetscale/cli/internal/cmdutil" ps "github.com/planetscale/cli/internal/planetscale" "github.com/planetscale/cli/internal/printer" - "github.com/spf13/cobra" ) +// diskScalingStrategies are the disk autoscaling strategies accepted by the +// --disk-scaling-strategy flag. +var diskScalingStrategies = []string{"grow", "disable", "shrink"} + func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { updateReq := &ps.UpdateKeyspaceSettingsRequest{} var flags struct { replicationDurabilityConstraints *ps.ReplicationDurabilityConstraints vreplicationFlags *ps.VReplicationFlags + diskScalingStrategy string + maxStorage int64 interactive bool } @@ -36,6 +45,10 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { updateReq.Branch = branch updateReq.Keyspace = keyspace + if cmd.Flags().Changed("disk-scaling-strategy") && !slices.Contains(diskScalingStrategies, flags.diskScalingStrategy) { + return fmt.Errorf("invalid --disk-scaling-strategy %q, must be one of: %s", flags.diskScalingStrategy, strings.Join(diskScalingStrategies, ", ")) + } + if flags.interactive { return updateInteractive(ctx, ch, updateReq) } @@ -84,7 +97,26 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { } } - if !rdcChanged && !vrfChanged { + // Check if any relevant flags are changing disk autoscaling settings + strategyChanged := cmd.Flags().Changed("disk-scaling-strategy") + maxStorageChanged := cmd.Flags().Changed("max-storage") + daChanged := strategyChanged || maxStorageChanged + + if daChanged { + updateReq.DiskAutoscaling = &ps.DiskAutoscalingUpdate{} + + if strategyChanged { + strategy := flags.diskScalingStrategy + updateReq.DiskAutoscaling.Strategy = &strategy + } + + if maxStorageChanged { + maxStorage := flags.maxStorage + updateReq.DiskAutoscaling.StorageLimitBytes = &maxStorage + } + } + + if !rdcChanged && !vrfChanged && !daChanged { end() ch.Printer.Println("No changes were requested. No update performed.") return nil @@ -105,8 +137,12 @@ 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().StringVar(&flags.diskScalingStrategy, "disk-scaling-strategy", "grow", fmt.Sprintf("The disk autoscaling strategy (%s). 'grow' lets dedicated disks grow automatically up to the storage limit; 'disable' turns autoscaling off; 'shrink' recreates disks at their initial size and then disables autoscaling.", strings.Join(diskScalingStrategies, ", "))) + cmd.Flags().Int64Var(&flags.maxStorage, "max-storage", 0, "The maximum size in bytes that dedicated disks may autoscale to. Required when the strategy is 'grow'.") cmd.Flags().BoolVarP(&flags.interactive, "interactive", "i", false, "Run the command in interactive mode") + _ = cmd.RegisterFlagCompletionFunc("disk-scaling-strategy", cobra.FixedCompletions(diskScalingStrategies, cobra.ShellCompDirectiveNoFileComp)) + return cmd } diff --git a/internal/cmd/keyspace/update_settings_test.go b/internal/cmd/keyspace/update_settings_test.go index fe2cccbf..5241bf05 100644 --- a/internal/cmd/keyspace/update_settings_test.go +++ b/internal/cmd/keyspace/update_settings_test.go @@ -8,6 +8,7 @@ import ( "time" qt "github.com/frankban/quicktest" + "github.com/planetscale/cli/internal/cmdutil" "github.com/planetscale/cli/internal/config" "github.com/planetscale/cli/internal/mock" @@ -619,6 +620,246 @@ func TestKeyspace_UpdateSettingsCmd_PreserveNilValues(t *testing.T) { c.Assert(buf.String(), qt.JSONEquals, updatedKs) } +func TestKeyspace_UpdateSettingsCmd_DiskAutoscaling(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() + + ks := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + } + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + DiskAutoscaling: &ps.DiskAutoscaling{ + Strategy: "grow", + StorageLimitBytes: 8796093022208, + }, + } + + 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.Database, qt.Equals, db) + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.Branch, qt.Equals, branch) + c.Assert(req.Keyspace, qt.Equals, keyspace) + + c.Assert(req.DiskAutoscaling, qt.Not(qt.IsNil)) + c.Assert(req.DiskAutoscaling.Strategy, qt.Not(qt.IsNil)) + c.Assert(*req.DiskAutoscaling.Strategy, qt.Equals, "grow") + c.Assert(req.DiskAutoscaling.StorageLimitBytes, qt.Not(qt.IsNil)) + c.Assert(*req.DiskAutoscaling.StorageLimitBytes, qt.Equals, int64(8796093022208)) + + 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, + "--disk-scaling-strategy=grow", + "--max-storage=8796093022208", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsTrue) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, updatedKs) +} + +func TestKeyspace_UpdateSettingsCmd_DiskAutoscalingInvalidStrategy(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{} + + 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, + "--disk-scaling-strategy=nonsense", + }) + err := cmd.Execute() + c.Assert(err, qt.ErrorMatches, `invalid --disk-scaling-strategy "nonsense", must be one of: grow, disable, shrink`) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse) +} + +// When disk autoscaling flags are omitted, an update triggered by other flags +// must not carry any disk autoscaling settings, so flag defaults don't +// overwrite the keyspace's current settings on the server. +func TestKeyspace_UpdateSettingsCmd_DiskAutoscalingNotSet(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() + + ks := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + ReplicationDurabilityConstraints: &ps.ReplicationDurabilityConstraints{ + Strategy: "available", + }, + } + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + ReplicationDurabilityConstraints: &ps.ReplicationDurabilityConstraints{ + Strategy: "lag", + }, + } + + 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) { + // Only an unrelated flag was passed, so no disk autoscaling + // settings should be sent. + c.Assert(req.DiskAutoscaling, 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, + "--replication-durability-constraints-strategy=dynamic", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsTrue) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) +} + +func TestKeyspace_UpdateSettingsCmd_NoFlagsUpdateNotPerformed(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} + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + 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}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + 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..b56b8b67 100644 --- a/internal/planetscale/keyspaces.go +++ b/internal/planetscale/keyspaces.go @@ -25,6 +25,15 @@ type Keyspace struct { VReplicationFlags *VReplicationFlags `json:"vreplication_flags"` ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints"` ReadOnlyRegions []*ReadOnlyRegionKeyspace `json:"read_only_regions"` + DiskAutoscaling *DiskAutoscaling `json:"disk_autoscaling"` +} + +// DiskAutoscaling configures how a keyspace's dedicated disks autoscale. +type DiskAutoscaling struct { + // Strategy is the disk autoscaling strategy: "grow" or "disable". + Strategy string `json:"strategy"` + // StorageLimitBytes is the maximum size in bytes disks may autoscale to. + StorageLimitBytes int64 `json:"storage_limit_bytes"` } type ReadOnlyRegionKeyspace struct { @@ -174,6 +183,14 @@ type UpdateKeyspaceSettingsRequest struct { Keyspace string `json:"-"` ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints,omitempty"` VReplicationFlags *VReplicationFlags `json:"vreplication_flags,omitempty"` + DiskAutoscaling *DiskAutoscalingUpdate `json:"disk_autoscaling,omitempty"` +} + +// DiskAutoscalingUpdate is the request body for changing a keyspace's disk +// autoscaling settings. Only the fields that are set are sent to the API. +type DiskAutoscalingUpdate struct { + Strategy *string `json:"strategy,omitempty"` + StorageLimitBytes *int64 `json:"storage_limit_bytes,omitempty"` } type ReplicationDurabilityConstraints struct { diff --git a/internal/planetscale/keyspaces_test.go b/internal/planetscale/keyspaces_test.go index 792ec397..45166dd5 100644 --- a/internal/planetscale/keyspaces_test.go +++ b/internal/planetscale/keyspaces_test.go @@ -3,6 +3,7 @@ package planetscale import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" "testing" @@ -479,3 +480,45 @@ func TestKeyspaces_UpdateSettings(t *testing.T) { c.Assert(keyspace.VReplicationFlags.VPlayerBatching, qt.Equals, true) c.Assert(keyspace.ReplicationDurabilityConstraints.Strategy, qt.Equals, "maximum") } + +func TestKeyspaces_UpdateSettingsDiskAutoscaling(t *testing.T) { + c := qt.New(t) + + var body []byte + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var err error + body, err = io.ReadAll(r.Body) + c.Assert(err, qt.IsNil) + + w.WriteHeader(200) + out := `{"type":"Keyspace","id":"thisisanid","name":"planetscale","shards":2,"sharded":true,"disk_autoscaling":{"strategy":"grow","storage_limit_bytes":8796093022208}}` + _, err = w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + c.Assert(r.Method, qt.Equals, http.MethodPatch) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + ctx := context.Background() + + strategy := "grow" + limit := int64(8796093022208) + + keyspace, err := client.Keyspaces.UpdateSettings(ctx, &UpdateKeyspaceSettingsRequest{ + Organization: "foo", + Database: "bar", + Branch: "baz", + Keyspace: "qux", + DiskAutoscaling: &DiskAutoscalingUpdate{ + Strategy: &strategy, + StorageLimitBytes: &limit, + }, + }) + + c.Assert(err, qt.IsNil) + c.Assert(string(body), qt.Contains, `"disk_autoscaling":{"strategy":"grow","storage_limit_bytes":8796093022208}`) + c.Assert(keyspace.DiskAutoscaling, qt.IsNotNil) + c.Assert(keyspace.DiskAutoscaling.Strategy, qt.Equals, "grow") + c.Assert(keyspace.DiskAutoscaling.StorageLimitBytes, qt.Equals, int64(8796093022208)) +}