Accept human-readable sizes for keyspace --max-storage and --storage

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)"
+			}
 		}
 	}
 

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)")
 }

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
+}

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 {

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"`