diff --git a/internal/cmd/keyspace/keyspace.go b/internal/cmd/keyspace/keyspace.go index 7a59d599..027f6f70 100644 --- a/internal/cmd/keyspace/keyspace.go +++ b/internal/cmd/keyspace/keyspace.go @@ -57,6 +57,7 @@ type KeyspaceSettings struct { VReplicationFlags VReplicationFlags `header:"inline" json:"vreplication_flags"` MaxRollout string `header:"max rollout" json:"max_rollout"` Throttler Throttler `header:"inline" json:"throttler"` + Storage Storage `header:"inline" json:"storage"` orig *ps.Keyspace } @@ -80,6 +81,12 @@ type Throttler struct { Threshold string `header:"throttler threshold" json:"threshold"` } +type Storage struct { + DiskScalingStrategy string `header:"disk scaling strategy" json:"disk_scaling_strategy"` + StorageBytes string `header:"storage" json:"storage_bytes"` + MaxStorageBytes string `header:"max storage" json:"max_storage_bytes"` +} + 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 233fe6d6..2b820cfa 100644 --- a/internal/cmd/keyspace/settings.go +++ b/internal/cmd/keyspace/settings.go @@ -4,6 +4,7 @@ import ( "fmt" "strconv" + "github.com/dustin/go-humanize" "github.com/planetscale/cli/internal/cmdutil" ps "github.com/planetscale/cli/internal/planetscale" "github.com/planetscale/cli/internal/printer" @@ -55,7 +56,12 @@ func SettingsCmd(ch *cmdutil.Helper) *cobra.Command { func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings { settings := &KeyspaceSettings{ MaxRollout: "not set", - orig: ks, + Storage: Storage{ + DiskScalingStrategy: "not set", + StorageBytes: "not set", + MaxStorageBytes: "not set", + }, + orig: ks, } if ks.MaxRollout != nil { @@ -100,5 +106,18 @@ func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings { } } + // Set the disk storage settings if available + if ks.Storage != nil { + if ks.Storage.DiskScalingStrategy != "" { + settings.Storage.DiskScalingStrategy = ks.Storage.DiskScalingStrategy + } + if ks.Storage.StorageBytes > 0 { + settings.Storage.StorageBytes = humanize.IBytes(uint64(ks.Storage.StorageBytes)) + } + if ks.Storage.MaxStorageBytes > 0 { + settings.Storage.MaxStorageBytes = humanize.IBytes(uint64(ks.Storage.MaxStorageBytes)) + } + } + return settings } diff --git a/internal/cmd/keyspace/settings_test.go b/internal/cmd/keyspace/settings_test.go index 07cb5f44..658b2798 100644 --- a/internal/cmd/keyspace/settings_test.go +++ b/internal/cmd/keyspace/settings_test.go @@ -202,6 +202,11 @@ func TestBuildKeyspaceSettings(t *testing.T) { maxRollout := 8 fullKs.MaxRollout = &maxRollout + fullKs.Storage = &ps.KeyspaceStorage{ + StorageBytes: 107374182400, + MaxStorageBytes: 4398046511104, + DiskScalingStrategy: "grow", + } settings := toKeyspaceSettings(fullKs) c.Assert(settings.ReplicationDurabilityConstraintStrategy, qt.Equals, "maximum") // Should be translated @@ -209,6 +214,9 @@ func TestBuildKeyspaceSettings(t *testing.T) { c.Assert(settings.VReplicationFlags.OptimizeInserts, qt.Equals, true) c.Assert(settings.VReplicationFlags.AllowNoBlobBinlogRowImage, qt.Equals, true) c.Assert(settings.VReplicationFlags.VPlayerBatching, qt.Equals, false) + c.Assert(settings.Storage.DiskScalingStrategy, qt.Equals, "grow") + c.Assert(settings.Storage.StorageBytes, qt.Equals, "100 GiB") + c.Assert(settings.Storage.MaxStorageBytes, qt.Equals, "4.0 TiB") // Test with nil settings nilKs := &ps.Keyspace{ @@ -226,4 +234,22 @@ func TestBuildKeyspaceSettings(t *testing.T) { 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) + c.Assert(nilSettings.Storage.DiskScalingStrategy, qt.Equals, "not set") + c.Assert(nilSettings.Storage.StorageBytes, qt.Equals, "not set") + c.Assert(nilSettings.Storage.MaxStorageBytes, qt.Equals, "not set") + + // Test with a storage object that only carries a strategy, which is what + // the API returns once autoscaling is disabled. + disabledKs := &ps.Keyspace{ + ID: "ks1", + Name: "test", + Storage: &ps.KeyspaceStorage{ + DiskScalingStrategy: "disable", + }, + } + + disabledSettings := toKeyspaceSettings(disabledKs) + c.Assert(disabledSettings.Storage.DiskScalingStrategy, qt.Equals, "disable") + c.Assert(disabledSettings.Storage.StorageBytes, qt.Equals, "not set") + c.Assert(disabledSettings.Storage.MaxStorageBytes, qt.Equals, "not set") } diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index 953cb27e..f6a0e034 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/internal/cmd/keyspace/update_settings.go @@ -4,15 +4,26 @@ import ( "context" "errors" "fmt" + "slices" "strconv" + "strings" "github.com/charmbracelet/huh" + "github.com/dustin/go-humanize" "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 scaling strategies accepted by the +// --disk-scaling-strategy flag. +var diskScalingStrategies = []string{"grow", "disable", "shrink"} + +// shrinkStrategy recreates disks at the requested size and then disables +// autoscaling. It is the only strategy that accepts --storage. +const shrinkStrategy = "shrink" + func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { updateReq := &ps.UpdateKeyspaceSettingsRequest{} @@ -22,6 +33,9 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { throttlerEnabled bool throttlerThreshold float64 maxRollout int + diskScalingStrategy string + maxStorage int64 + storage int64 interactive bool } @@ -56,7 +70,12 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { maxRolloutChanged := cmd.Flags().Changed("max-rollout") - if !rdcChanged && !vrfChanged && !throttlerChanged && !maxRolloutChanged { + strategyChanged := cmd.Flags().Changed("disk-scaling-strategy") + maxStorageChanged := cmd.Flags().Changed("max-storage") + storageChanged := cmd.Flags().Changed("storage") + diskStorageChanged := strategyChanged || maxStorageChanged || storageChanged + + if !rdcChanged && !vrfChanged && !throttlerChanged && !maxRolloutChanged && !diskStorageChanged { ch.Printer.Println("No changes were requested. No update performed.") return nil } @@ -69,6 +88,30 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { return errors.New("--max-rollout must be between 1 and 32") } + if strategyChanged && !slices.Contains(diskScalingStrategies, flags.diskScalingStrategy) { + return fmt.Errorf("invalid --disk-scaling-strategy %q, must be one of: %s", flags.diskScalingStrategy, strings.Join(diskScalingStrategies, ", ")) + } + + if maxStorageChanged && flags.maxStorage <= 0 { + return errors.New("--max-storage must be greater than 0") + } + + if storageChanged { + // The API only recreates disks at a new size when shrinking, so + // the strategy has to be part of the same request. + if !strategyChanged || flags.diskScalingStrategy != shrinkStrategy { + return fmt.Errorf("--storage can only be set when --disk-scaling-strategy is %s", shrinkStrategy) + } + + if flags.storage <= 0 { + return errors.New("--storage must be greater than 0") + } + + if flags.storage%humanize.GiByte != 0 { + return errors.New("--storage must be a multiple of 1 GiB") + } + } + client, err := ch.Client() if err != nil { return err @@ -127,6 +170,23 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { updateReq.MaxRollout = &flags.maxRollout } + // Disk storage fields that are left out keep their current values. + if diskStorageChanged { + updateReq.Storage = &ps.KeyspaceStorageUpdate{} + + if strategyChanged { + updateReq.Storage.DiskScalingStrategy = &flags.diskScalingStrategy + } + + if maxStorageChanged { + updateReq.Storage.MaxStorageBytes = &flags.maxStorage + } + + if storageChanged { + updateReq.Storage.StorageBytes = &flags.storage + } + } + k, err := updateKeyspaceSettings(ctx, client, updateReq) if err != nil { return err @@ -145,8 +205,13 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { 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().StringVar(&flags.diskScalingStrategy, "disk-scaling-strategy", "grow", fmt.Sprintf("The disk scaling strategy (%s). 'grow' lets dedicated disks grow automatically up to --max-storage; 'disable' turns autoscaling off; 'shrink' recreates disks at --storage 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.") + cmd.Flags().Int64Var(&flags.storage, "storage", 0, fmt.Sprintf("The disk size in bytes to recreate disks at. Must be a multiple of 1 GiB and requires --disk-scaling-strategy %s.", shrinkStrategy)) 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 612cb268..5cd7071a 100644 --- a/internal/cmd/keyspace/update_settings_test.go +++ b/internal/cmd/keyspace/update_settings_test.go @@ -1142,6 +1142,383 @@ func TestKeyspace_UpdateSettingsCmd_MaxRolloutAndVReplicationFetchesOnce(t *test c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) } +func TestKeyspace_UpdateSettingsCmd_Storage(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() + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + CreatedAt: ts, + UpdatedAt: ts, + Storage: &ps.KeyspaceStorage{ + StorageBytes: 107374182400, + MaxStorageBytes: 4398046511104, + DiskScalingStrategy: "grow", + }, + } + + 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.Storage, qt.Not(qt.IsNil)) + c.Assert(req.Storage.DiskScalingStrategy, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.DiskScalingStrategy, qt.Equals, "grow") + c.Assert(req.Storage.MaxStorageBytes, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.MaxStorageBytes, qt.Equals, int64(4398046511104)) + c.Assert(req.Storage.StorageBytes, qt.IsNil) + c.Assert(req.ReplicationDurabilityConstraints, qt.IsNil) + c.Assert(req.VReplicationFlags, qt.IsNil) + c.Assert(req.Throttler, qt.IsNil) + c.Assert(req.MaxRollout, 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, "--disk-scaling-strategy=grow", "--max-storage=4398046511104"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, updatedKs) +} + +func TestKeyspace_UpdateSettingsCmd_ShrinkStorage(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" + + updatedKs := &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + Storage: &ps.KeyspaceStorage{ + StorageBytes: 214748364800, + DiskScalingStrategy: "shrink", + }, + } + + 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.Storage, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.DiskScalingStrategy, qt.Equals, "shrink") + c.Assert(req.Storage.StorageBytes, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.StorageBytes, qt.Equals, int64(214748364800)) + c.Assert(req.Storage.MaxStorageBytes, 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, "--disk-scaling-strategy=shrink", "--storage=214748364800"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, updatedKs) +} + +func TestKeyspace_UpdateSettingsCmd_DisableDiskScaling(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.Storage, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.DiskScalingStrategy, qt.Equals, "disable") + c.Assert(req.Storage.MaxStorageBytes, qt.IsNil) + c.Assert(req.Storage.StorageBytes, qt.IsNil) + + return &ps.Keyspace{ + ID: "ks1", + Name: keyspace, + Storage: &ps.KeyspaceStorage{ + DiskScalingStrategy: "disable", + }, + }, 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=disable"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) +} + +func TestKeyspace_UpdateSettingsCmd_RejectsInvalidDiskScalingStrategy(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Keyspaces: svc}, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{"planetscale", "main", "sharded", "--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) +} + +func TestKeyspace_UpdateSettingsCmd_RejectsStorageWithoutShrink(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + args := [][]string{ + {"--storage=214748364800"}, + {"--disk-scaling-strategy=grow", "--storage=214748364800"}, + } + + for _, extra := range args { + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Keyspaces: svc}, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs(append([]string{"planetscale", "main", "sharded"}, extra...)) + err := cmd.Execute() + c.Assert(err, qt.ErrorMatches, "--storage can only be set when --disk-scaling-strategy is shrink") + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse) + } +} + +func TestKeyspace_UpdateSettingsCmd_RejectsInvalidStorageSizes(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + tests := []struct { + args []string + err string + }{ + { + args: []string{"--disk-scaling-strategy=shrink", "--storage=53687091201"}, + err: "--storage must be a multiple of 1 GiB", + }, + { + args: []string{"--disk-scaling-strategy=shrink", "--storage=0"}, + err: "--storage must be greater than 0", + }, + { + args: []string{"--max-storage=0"}, + err: "--max-storage must be greater than 0", + }, + } + + for _, tt := range tests { + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Keyspaces: svc}, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs(append([]string{"planetscale", "main", "sharded"}, tt.args...)) + err := cmd.Execute() + c.Assert(err, qt.ErrorMatches, tt.err) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsFalse) + } +} + +func TestKeyspace_UpdateSettingsCmd_StorageAndThrottlerSkipGet(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Storage, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.DiskScalingStrategy, qt.Equals, "grow") + c.Assert(req.Throttler, qt.Not(qt.IsNil)) + c.Assert(*req.Throttler.Enabled, qt.IsFalse) + c.Assert(req.VReplicationFlags, qt.IsNil) + + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Keyspaces: svc}, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{"planetscale", "main", "sharded", "--disk-scaling-strategy=grow", "--throttler-enabled=false"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.GetFnInvoked, qt.IsFalse) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) +} + +func TestKeyspace_UpdateSettingsCmd_OmittedStorageIsNotSent(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + svc := &mock.KeyspacesService{ + GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) { + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Storage, qt.IsNil) + + return &ps.Keyspace{ID: "ks1", Name: "sharded"}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Keyspaces: svc}, nil + }, + } + + cmd := UpdateSettingsCmd(ch) + cmd.SetArgs([]string{"planetscale", "main", "sharded", "--max-rollout=4"}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) +} + func TestKeyspace_UpdateSettingsCmd_NoFlagsMakesNoRequests(t *testing.T) { c := qt.New(t) diff --git a/internal/planetscale/keyspaces.go b/internal/planetscale/keyspaces.go index 3d60c78b..64bbfb81 100644 --- a/internal/planetscale/keyspaces.go +++ b/internal/planetscale/keyspaces.go @@ -26,9 +26,23 @@ type Keyspace struct { ReplicationDurabilityConstraints *ReplicationDurabilityConstraints `json:"replication_durability_constraints"` MaxRollout *int `json:"max_rollout"` Throttler *KeyspaceThrottler `json:"throttler"` + Storage *KeyspaceStorage `json:"storage"` ReadOnlyRegions []*ReadOnlyRegionKeyspace `json:"read_only_regions"` } +// KeyspaceStorage describes the disk storage configuration for a keyspace's +// dedicated disks. +type KeyspaceStorage struct { + // StorageBytes is the provisioned disk size in bytes. Disks grow from and + // shrink to this size. + StorageBytes int64 `json:"storage_bytes"` + // MaxStorageBytes is the maximum size in bytes disks may autoscale to. + MaxStorageBytes int64 `json:"max_storage_bytes"` + // DiskScalingStrategy is the disk scaling strategy: "grow", "disable" or + // "shrink". + DiskScalingStrategy string `json:"disk_scaling_strategy"` +} + type ReadOnlyRegionKeyspace struct { Region string `json:"region"` ClusterName string `json:"cluster_name"` @@ -178,6 +192,7 @@ type UpdateKeyspaceSettingsRequest struct { VReplicationFlags *VReplicationFlags `json:"vreplication_flags,omitempty"` Throttler *KeyspaceThrottler `json:"throttler,omitempty"` MaxRollout *int `json:"max_rollout,omitempty"` + Storage *KeyspaceStorageUpdate `json:"storage,omitempty"` } type ReplicationDurabilityConstraints struct { @@ -195,6 +210,14 @@ type KeyspaceThrottler struct { Threshold *float64 `json:"threshold,omitempty"` } +// KeyspaceStorageUpdate changes a keyspace's disk storage settings. Only the +// fields that are set are sent to the API. +type KeyspaceStorageUpdate struct { + DiskScalingStrategy *string `json:"disk_scaling_strategy,omitempty"` + MaxStorageBytes *int64 `json:"max_storage_bytes,omitempty"` + StorageBytes *int64 `json:"storage_bytes,omitempty"` +} + // KeyspacesService is an interface for interacting with the keyspace endpoints of the PlanetScale API type KeyspacesService interface { Create(context.Context, *CreateKeyspaceRequest) (*Keyspace, error) diff --git a/internal/planetscale/keyspaces_test.go b/internal/planetscale/keyspaces_test.go index 993e983c..94826714 100644 --- a/internal/planetscale/keyspaces_test.go +++ b/internal/planetscale/keyspaces_test.go @@ -519,3 +519,47 @@ func TestKeyspaces_UpdateSettingsMaxRollout(t *testing.T) { c.Assert(keyspace.MaxRollout, qt.Not(qt.IsNil)) c.Assert(*keyspace.MaxRollout, qt.Equals, 8) } + +func TestKeyspaces_UpdateSettingsStorage(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","storage":{"storage_bytes":214748364800,"max_storage_bytes":4398046511104,"disk_scaling_strategy":"shrink"}}` + _, err = w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + ctx := context.Background() + strategy := "shrink" + storageBytes := int64(214748364800) + + keyspace, err := client.Keyspaces.UpdateSettings(ctx, &UpdateKeyspaceSettingsRequest{ + Organization: "foo", + Database: "bar", + Branch: "baz", + Keyspace: "qux", + Storage: &KeyspaceStorageUpdate{ + DiskScalingStrategy: &strategy, + StorageBytes: &storageBytes, + }, + }) + + c.Assert(err, qt.IsNil) + c.Assert(strings.TrimSpace(body), qt.Equals, `{"storage":{"disk_scaling_strategy":"shrink","storage_bytes":214748364800}}`) + c.Assert(keyspace.Storage, qt.Not(qt.IsNil)) + c.Assert(keyspace.Storage.DiskScalingStrategy, qt.Equals, "shrink") + c.Assert(keyspace.Storage.StorageBytes, qt.Equals, int64(214748364800)) + c.Assert(keyspace.Storage.MaxStorageBytes, qt.Equals, int64(4398046511104)) +}