From 617edc1371d892fe6e181e8e5a114fc9151ebc87 Mon Sep 17 00:00:00 2001 From: Inel Pandzic Date: Mon, 5 Oct 2026 17:12:33 +0200 Subject: [PATCH 1/2] Accept human-readable sizes for keyspace --max-storage and --storage --- internal/cmd/keyspace/update_settings.go | 60 +++++++-- internal/cmd/keyspace/update_settings_test.go | 120 ++++++++++++++++++ 2 files changed, 170 insertions(+), 10 deletions(-) diff --git a/internal/cmd/keyspace/update_settings.go b/internal/cmd/keyspace/update_settings.go index d3a11b29..e88cf4a6 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) or a higher limit set by PlanetScale staff.") + 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 { From 58c11951463b3943692b10b6911c4367305e2b16 Mon Sep 17 00:00:00 2001 From: Inel Pandzic Date: Tue, 6 Oct 2026 18:19:45 +0200 Subject: [PATCH 2/2] Show when a keyspace max storage limit is set by PlanetScale staff --- internal/cmd/keyspace/settings.go | 3 +++ internal/cmd/keyspace/settings_test.go | 13 +++++++++++++ internal/cmd/keyspace/update_settings.go | 2 +- internal/planetscale/keyspaces.go | 3 +++ 4 files changed, 20 insertions(+), 1 deletion(-) 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 e88cf4a6..855be348 100644 --- a/internal/cmd/keyspace/update_settings.go +++ b/internal/cmd/keyspace/update_settings.go @@ -214,7 +214,7 @@ 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().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) or a higher limit set by PlanetScale staff.") + 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") diff --git a/internal/planetscale/keyspaces.go b/internal/planetscale/keyspaces.go index 08a4c6a2..03c01b94 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"`