diff --git a/internal/cmd/keyspace/settings.go b/internal/cmd/keyspace/settings.go index 2b820cfa..abbf3a21 100644 --- a/internal/cmd/keyspace/settings.go +++ b/internal/cmd/keyspace/settings.go @@ -116,6 +116,9 @@ func toKeyspaceSettings(ks *ps.Keyspace) *KeyspaceSettings { } if ks.Storage.MaxStorageBytes > 0 { settings.Storage.MaxStorageBytes = humanize.IBytes(uint64(ks.Storage.MaxStorageBytes)) + if ks.Storage.MaxStorageBytesManagedByStaff { + settings.Storage.MaxStorageBytes += " (set by PlanetScale staff)" + } } } diff --git a/internal/cmd/keyspace/settings_test.go b/internal/cmd/keyspace/settings_test.go index 658b2798..48ce4f44 100644 --- a/internal/cmd/keyspace/settings_test.go +++ b/internal/cmd/keyspace/settings_test.go @@ -252,4 +252,17 @@ func TestBuildKeyspaceSettings(t *testing.T) { 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") + + // A limit raised by PlanetScale staff is read-only, so the table says so. + staffKs := &ps.Keyspace{ + ID: "ks1", + Name: "test", + Storage: &ps.KeyspaceStorage{ + MaxStorageBytes: 8796093022208, + MaxStorageBytesManagedByStaff: true, + }, + } + + staffSettings := toKeyspaceSettings(staffKs) + c.Assert(staffSettings.Storage.MaxStorageBytes, qt.Equals, "8.0 TiB (set by PlanetScale staff)") } diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index d3a11b29..855be348 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/internal/cmd/keyspace/update_settings.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "math" "slices" "strconv" "strings" @@ -30,8 +31,8 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { throttlerThreshold float64 maxRollout int diskScalingStrategy string - maxStorage int64 - storage int64 + maxStorage string + storage string interactive bool } @@ -88,16 +89,33 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { 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") + var ( + maxStorageBytes int64 + err error + ) + if maxStorageChanged { + maxStorageBytes, err = parseStorageSize(flags.maxStorage) + if err != nil { + return fmt.Errorf("invalid --max-storage %q: %w", flags.maxStorage, err) + } + + if maxStorageBytes <= 0 { + return errors.New("--max-storage must be greater than 0") + } } + var storageBytes int64 if storageChanged { - if flags.storage <= 0 { + storageBytes, err = parseStorageSize(flags.storage) + if err != nil { + return fmt.Errorf("invalid --storage %q: %w", flags.storage, err) + } + + if storageBytes <= 0 { return errors.New("--storage must be greater than 0") } - if flags.storage%humanize.GiByte != 0 { + if storageBytes%humanize.GiByte != 0 { return errors.New("--storage must be a multiple of 1 GiB") } } @@ -169,11 +187,11 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { } if maxStorageChanged { - updateReq.Storage.MaxStorageBytes = &flags.maxStorage + updateReq.Storage.MaxStorageBytes = &maxStorageBytes } if storageChanged { - updateReq.Storage.StorageBytes = &flags.storage + updateReq.Storage.StorageBytes = &storageBytes } } @@ -196,8 +214,8 @@ func UpdateSettingsCmd(ch *cmdutil.Helper) *cobra.Command { 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, "The disk size in bytes to recreate disks at. Must be a multiple of 1 GiB. Only accepted when the disk scaling strategy is shrink, either passed with --disk-scaling-strategy or already set on the keyspace.") + cmd.Flags().StringVar(&flags.maxStorage, "max-storage", "", "The maximum size that dedicated disks may autoscale to, in bytes or with a unit such as 200GiB or 1TiB. Must be at least 12 GiB and no smaller than the current disk size. Cannot be raised above the organization default (4 TiB, or 16 TiB for managed tenancy). Read-only while PlanetScale staff have raised it above that default.") + cmd.Flags().StringVar(&flags.storage, "storage", "", "The disk size to recreate disks at, in bytes or with a unit such as 200GiB or 1TiB. Must be a multiple of 1 GiB. Only accepted when the disk scaling strategy is shrink, either passed with --disk-scaling-strategy or already set on the keyspace.") cmd.Flags().BoolVarP(&flags.interactive, "interactive", "i", false, "Run the command in interactive mode") _ = cmd.RegisterFlagCompletionFunc("disk-scaling-strategy", cobra.FixedCompletions(diskScalingStrategies, cobra.ShellCompDirectiveNoFileComp)) @@ -390,3 +408,25 @@ func constraintsToStrategy(constraint string) string { return constraint } } + +// parseStorageSize parses a storage size given either as a plain number of +// bytes or as a size with a unit, such as 200GiB or 1TiB. It only parses; the +// caller is responsible for validating the returned value. +func parseStorageSize(value string) (int64, error) { + // A plain integer is parsed exactly. humanize.ParseBytes goes through a + // float64 and would silently round byte counts above 2^53. + if n, err := strconv.ParseInt(strings.TrimSpace(value), 10, 64); err == nil { + return n, nil + } + + parsed, err := humanize.ParseBytes(value) + if err != nil { + return 0, errors.New("must be a size in bytes or with a unit, such as 200GiB") + } + + if parsed > math.MaxInt64 { + return 0, errors.New("value is too large") + } + + return int64(parsed), nil +} diff --git a/internal/cmd/keyspace/update_settings_test.go b/internal/cmd/keyspace/update_settings_test.go index ea20c932..a4fbb130 100644 --- a/internal/cmd/keyspace/update_settings_test.go +++ b/internal/cmd/keyspace/update_settings_test.go @@ -1211,6 +1211,101 @@ func TestKeyspace_UpdateSettingsCmd_Storage(t *testing.T) { c.Assert(buf.String(), qt.JSONEquals, updatedKs) } +func TestKeyspace_UpdateSettingsCmd_MaxStorageAcceptsHumanReadableSize(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + tests := []struct { + arg string + bytes int64 + }{ + {arg: "200GiB", bytes: 214748364800}, + {arg: "1TiB", bytes: 1099511627776}, + {arg: "1.5 TiB", bytes: 1649267441664}, + {arg: "107374182400", bytes: 107374182400}, + // Plain integers are parsed exactly, including above 2^53 where a + // float64 would round. + {arg: "9007199254740993", bytes: 9007199254740993}, + } + + for _, tt := range tests { + svc := &mock.KeyspacesService{ + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Storage, qt.Not(qt.IsNil)) + c.Assert(req.Storage.MaxStorageBytes, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.MaxStorageBytes, qt.Equals, tt.bytes) + + 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-storage=" + tt.arg}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + } +} + +func TestKeyspace_UpdateSettingsCmd_StorageAcceptsHumanReadableSize(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + tests := []struct { + arg string + bytes int64 + }{ + {arg: "200GiB", bytes: 214748364800}, + {arg: "10TiB", bytes: 10995116277760}, + {arg: "1.5 TiB", bytes: 1649267441664}, + {arg: "214748364800", bytes: 214748364800}, + } + + for _, tt := range tests { + svc := &mock.KeyspacesService{ + UpdateSettingsFn: func(ctx context.Context, req *ps.UpdateKeyspaceSettingsRequest) (*ps.Keyspace, error) { + c.Assert(req.Storage, qt.Not(qt.IsNil)) + c.Assert(req.Storage.StorageBytes, qt.Not(qt.IsNil)) + c.Assert(*req.Storage.StorageBytes, qt.Equals, tt.bytes) + + 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=shrink", "--storage=" + tt.arg}) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.UpdateSettingsFnInvoked, qt.IsTrue) + } +} + func TestKeyspace_UpdateSettingsCmd_ShrinkStorage(t *testing.T) { c := qt.New(t) @@ -1423,10 +1518,35 @@ func TestKeyspace_UpdateSettingsCmd_RejectsInvalidStorageSizes(t *testing.T) { args: []string{"--disk-scaling-strategy=shrink", "--storage=0"}, err: "--storage must be greater than 0", }, + { + args: []string{"--disk-scaling-strategy=shrink", "--storage=lots"}, + err: `invalid --storage "lots": must be a size in bytes or with a unit, such as 200GiB`, + }, + { + // Decimal units are not GiB-aligned. + args: []string{"--disk-scaling-strategy=shrink", "--storage=10TB"}, + err: "--storage must be a multiple of 1 GiB", + }, + { + args: []string{"--disk-scaling-strategy=shrink", "--storage=-1"}, + err: "--storage must be greater than 0", + }, { args: []string{"--max-storage=0"}, err: "--max-storage must be greater than 0", }, + { + args: []string{"--max-storage=-1073741824"}, + err: "--max-storage must be greater than 0", + }, + { + args: []string{"--max-storage=9EiB"}, + err: `invalid --max-storage "9EiB": value is too large`, + }, + { + args: []string{"--max-storage=lots"}, + err: `invalid --max-storage "lots": must be a size in bytes or with a unit, such as 200GiB`, + }, } for _, tt := range tests { diff --git a/internal/planetscale/keyspaces.go b/internal/planetscale/keyspaces.go index 78ddaedb..d85d1877 100644 --- a/internal/planetscale/keyspaces.go +++ b/internal/planetscale/keyspaces.go @@ -39,6 +39,9 @@ type KeyspaceStorage struct { StorageBytes int64 `json:"storage_bytes"` // MaxStorageBytes is the maximum size in bytes disks may autoscale to. MaxStorageBytes int64 `json:"max_storage_bytes"` + // MaxStorageBytesManagedByStaff is true when PlanetScale staff raised the + // limit above the organization default. The limit is read-only while set. + MaxStorageBytesManagedByStaff bool `json:"max_storage_bytes_managed_by_staff"` // DiskScalingStrategy is the disk scaling strategy: "grow", "disable" or // "shrink". DiskScalingStrategy string `json:"disk_scaling_strategy"`