Add disk autoscaling settings to keyspace settings commands

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
 }

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
 }

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
 }
 

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)
 

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 {