From cfa31fcaf86c19ec1e03e511c1b3aa714591a939 Mon Sep 17 00:00:00 2001 From: Nick Holden Date: Tue, 7 Apr 2026 10:44:24 -0700 Subject: [PATCH 1/2] Add --min-storage and --max-storage flags to database and branch creation Allow users to configure minimum and maximum storage size (in bytes) when creating PostgreSQL databases and branches. Returns a clear error if used with MySQL databases. Co-Authored-By: Claude Opus 4.6 (1M context) --- internal/cmd/branch/create.go | 19 +++++ internal/cmd/branch/create_test.go | 107 +++++++++++++++++++++++++++ internal/cmd/database/create.go | 19 +++++ internal/cmd/database/create_test.go | 97 ++++++++++++++++++++++++ 4 files changed, 242 insertions(+) diff --git a/internal/cmd/branch/create.go b/internal/cmd/branch/create.go index 1ad4d3ed..d650a84a 100644 --- a/internal/cmd/branch/create.go +++ b/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") diff --git a/internal/cmd/branch/create_test.go b/internal/cmd/branch/create_test.go index 0b185ba4..b41b7031 100644 --- a/internal/cmd/branch/create_test.go +++ b/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) diff --git a/internal/cmd/database/create.go b/internal/cmd/database/create.go index a9e88879..db3d059a 100644 --- a/internal/cmd/database/create.go +++ b/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 diff --git a/internal/cmd/database/create_test.go b/internal/cmd/database/create_test.go index 17c6d050..4ae4bdc5 100644 --- a/internal/cmd/database/create_test.go +++ b/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) From 1fc64babf5eb5b9ad18216fd23bc3c4cad13836a Mon Sep 17 00:00:00 2001 From: Nick Holden Date: Tue, 7 Apr 2026 10:53:00 -0700 Subject: [PATCH 2/2] Update planetscale-go to v0.161.0 Co-Authored-By: Claude Opus 4.6 (1M context) --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index b46a6f64..5ea32d3d 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index b54b7ced..2e3dd570 100644 --- a/go.sum +++ b/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=