Merge pull request #1246 from planetscale/nick/storage-flags

Add --min-storage and --max-storage flags to database and branch creation

Author: nholden

go.mod

@@ -25,7 +25,7 @@ require (
 	github.com/mattn/go-shellwords v1.0.12
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
-	github.com/planetscale/planetscale-go v0.160.0
+	github.com/planetscale/planetscale-go v0.161.0
 	github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
 	github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
 	github.com/spf13/cobra v1.10.2

go.sum

@@ -176,8 +176,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjL
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2vvqGZLvoQfpaGg/j1fNDr4j03s3PRz4rVY=
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
-github.com/planetscale/planetscale-go v0.160.0 h1:nWgTrMJPk+FzV71OV+wKU17F6DCnSX0lHxHSUH6Yhn0=
-github.com/planetscale/planetscale-go v0.160.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
+github.com/planetscale/planetscale-go v0.161.0 h1:YWpOJXJTJKX2AFHoKkO/IZdRxCmOQAS+O6A+jJuteL8=
+github.com/planetscale/planetscale-go v0.161.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs=
 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8=

internal/cmd/branch/create.go

@@ -21,6 +21,8 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 		clusterSize   string
 		backupID      string
 		majorVersion  string
+		minStorage    int64
+		maxStorage    int64
 	}
 
 	cmd := &cobra.Command{
@@ -103,6 +105,10 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 			}
 
 			if db.Kind == "mysql" {
+				if cmd.Flags().Changed("min-storage") || cmd.Flags().Changed("max-storage") {
+					return fmt.Errorf("--min-storage and --max-storage are only supported for PostgreSQL databases")
+				}
+
 				createReq := &ps.CreateDatabaseBranchRequest{
 					Organization: ch.Config.Organization,
 					Database:     source,
@@ -171,6 +177,16 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 					MajorVersion: flags.majorVersion,
 				}
 
+				if cmd.Flags().Changed("min-storage") || cmd.Flags().Changed("max-storage") {
+					createReq.Storage = &ps.StorageConfig{}
+					if cmd.Flags().Changed("min-storage") {
+						createReq.Storage.MinimumStorageBytes = &flags.minStorage
+					}
+					if cmd.Flags().Changed("max-storage") {
+						createReq.Storage.MaximumStorageBytes = &flags.maxStorage
+					}
+				}
+
 				dbBranch, err := client.PostgresBranches.Create(cmd.Context(), createReq)
 				if err != nil {
 					switch cmdutil.ErrCode(err) {
@@ -224,6 +240,9 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.Flags().BoolVar(&flags.dataBranching, "seed-data", false, "Add seed data using the Data Branching™ feature. This branch will be created with the same resources as the base branch.")
 	cmd.Flags().BoolVar(&flags.wait, "wait", false, "Wait until the branch is ready")
 	cmd.Flags().StringVar(&flags.majorVersion, "major-version", "", "For PostgreSQL databases, the PostgreSQL major version to use for the branch. Defaults to the major version of the parent branch if it exists or the database's default branch major version. Ignored for branches restored from backups.")
+	cmd.Flags().Int64Var(&flags.minStorage, "min-storage", 0, "Minimum storage size in bytes")
+	cmd.Flags().Int64Var(&flags.maxStorage, "max-storage", 0, "Maximum storage size in bytes for autoscaling")
+
 	cmd.MarkFlagsMutuallyExclusive("from", "restore")
 	cmd.MarkFlagsMutuallyExclusive("restore", "seed-data")
 

internal/cmd/branch/create_test.go

@@ -298,6 +298,113 @@ func TestBranch_CreateCmdWithMajorVersion(t *testing.T) {
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
 
+func TestBranch_CreateCmdWithStorageMySQLError(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 := "development"
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			return &ps.Database{Kind: "mysql"}, nil
+		},
+	}
+
+	svc := &mock.DatabaseBranchesService{
+		CreateFn: func(ctx context.Context, req *ps.CreateDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
+			c.Fatal("CreateFn should not be called for MySQL with storage flags")
+			return nil, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				DatabaseBranches: svc,
+				Databases:        dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--region", "us-east", "--min-storage", "10737418240"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err, qt.ErrorMatches, ".*only supported for PostgreSQL.*")
+	c.Assert(svc.CreateFnInvoked, qt.IsFalse)
+}
+
+func TestBranch_CreateCmdWithStorage(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 := "development"
+
+	res := &ps.PostgresBranch{Name: branch}
+
+	svc := &mock.PostgresBranchesService{
+		CreateFn: func(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) {
+			c.Assert(req.Name, qt.Equals, branch)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Region, qt.Equals, "us-east")
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Storage, qt.IsNotNil)
+			c.Assert(req.Storage.MinimumStorageBytes, qt.IsNotNil)
+			c.Assert(*req.Storage.MinimumStorageBytes, qt.Equals, int64(10737418240))
+			c.Assert(req.Storage.MaximumStorageBytes, qt.IsNotNil)
+			c.Assert(*req.Storage.MaximumStorageBytes, qt.Equals, int64(107374182400))
+
+			return res, nil
+		},
+	}
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Organization, qt.Equals, org)
+			return &ps.Database{Kind: "postgresql"}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				PostgresBranches: svc,
+				Databases:        dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--region", "us-east", "--min-storage", "10737418240", "--max-storage", "107374182400"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}
+
 func TestBranch_CreateCmd_ServiceTokenAuthError(t *testing.T) {
 	c := qt.New(t)
 

internal/cmd/database/create.go

@@ -24,6 +24,8 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 		wait         bool
 		replicas     *int
 		majorVersion string
+		minStorage   int64
+		maxStorage   int64
 	}
 
 	cmd := &cobra.Command{
@@ -51,6 +53,20 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 				createReq.MajorVersion = flags.majorVersion
 			}
 
+			if (cmd.Flags().Changed("min-storage") || cmd.Flags().Changed("max-storage")) && engine != ps.DatabaseEnginePostgres {
+				return fmt.Errorf("--min-storage and --max-storage are only supported for PostgreSQL databases")
+			}
+
+			if cmd.Flags().Changed("min-storage") || cmd.Flags().Changed("max-storage") {
+				createReq.Storage = &ps.StorageConfig{}
+				if cmd.Flags().Changed("min-storage") {
+					createReq.Storage.MinimumStorageBytes = &flags.minStorage
+				}
+				if cmd.Flags().Changed("max-storage") {
+					createReq.Storage.MaximumStorageBytes = &flags.maxStorage
+				}
+			}
+
 			client, err := ch.Client()
 			if err != nil {
 				return err
@@ -120,6 +136,9 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.RegisterFlagCompletionFunc("cluster-size", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
 		return cmdutil.ClusterSizesCompletionFunc(ch, cmd, args, toComplete)
 	})
+	cmd.Flags().Int64Var(&flags.minStorage, "min-storage", 0, "Minimum storage size in bytes")
+	cmd.Flags().Int64Var(&flags.maxStorage, "max-storage", 0, "Maximum storage size in bytes for autoscaling")
+
 	cmd.Flags().BoolVar(&flags.wait, "wait", false, "Wait until the database is ready")
 
 	return cmd

internal/cmd/database/create_test.go

@@ -228,6 +228,103 @@ func TestDatabase_CreateCmdWithReplicas(t *testing.T) {
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
 
+func TestDatabase_CreateCmdWithStorageMySQLError(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"
+
+	svc := &mock.DatabaseService{
+		CreateFn: func(ctx context.Context, req *ps.CreateDatabaseRequest) (*ps.Database, error) {
+			c.Fatal("CreateFn should not be called for MySQL with storage flags")
+			return nil, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Databases: svc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, "--region", "us-east", "--min-storage", "10737418240"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err, qt.ErrorMatches, ".*only supported for PostgreSQL.*")
+	c.Assert(svc.CreateFnInvoked, qt.IsFalse)
+}
+
+func TestDatabase_CreateCmdWithStorage(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"
+
+	res := &ps.Database{Name: "foo"}
+
+	svc := &mock.DatabaseService{
+		CreateFn: func(ctx context.Context, req *ps.CreateDatabaseRequest) (*ps.Database, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Name, qt.Equals, db)
+			c.Assert(req.Region, qt.Equals, "us-east")
+			c.Assert(req.Kind, qt.Equals, ps.DatabaseEnginePostgres)
+			c.Assert(req.Storage, qt.IsNotNil)
+			c.Assert(req.Storage.MinimumStorageBytes, qt.IsNotNil)
+			c.Assert(*req.Storage.MinimumStorageBytes, qt.Equals, int64(10737418240))
+			c.Assert(req.Storage.MaximumStorageBytes, qt.IsNotNil)
+			c.Assert(*req.Storage.MaximumStorageBytes, qt.Equals, int64(107374182400))
+
+			return res, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Databases: svc,
+				Organizations: &mock.OrganizationsService{
+					GetFn: func(ctx context.Context, request *ps.GetOrganizationRequest) (*ps.Organization, error) {
+						return &ps.Organization{
+							RemainingFreeDatabases: 1,
+							Name:                   request.Organization,
+						}, nil
+					},
+				},
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, "--region", "us-east", "--engine", "postgresql", "--min-storage", "10737418240", "--max-storage", "107374182400"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}
+
 func TestDatabase_CreateCmdPostgresWithMajorVersion(t *testing.T) {
 	c := qt.New(t)