From ec06262129280071caa38f41af7b20b6e53ed6a2 Mon Sep 17 00:00:00 2001 From: Hirad Pourtahmasbi Date: Thu, 3 Sep 2026 16:43:36 -0700 Subject: [PATCH] add replica selection to Postgres restores --- internal/cmd/backup/restore.go | 26 +++- internal/cmd/backup/restore_test.go | 35 ++++- internal/cmd/branch/create.go | 13 ++ internal/cmd/branch/create_test.go | 142 +++++++++++++++++- internal/planetscale/postgres_branches.go | 1 + .../planetscale/postgres_branches_test.go | 35 +++++ 6 files changed, 244 insertions(+), 8 deletions(-) diff --git a/internal/cmd/backup/restore.go b/internal/cmd/backup/restore.go index 61560fa7..219ec04e 100644 --- a/internal/cmd/backup/restore.go +++ b/internal/cmd/backup/restore.go @@ -13,7 +13,10 @@ import ( ) func RestoreCmd(ch *cmdutil.Helper) *cobra.Command { - var clusterSize string + var flags struct { + clusterSize string + replicas int + } cmd := &cobra.Command{ Use: "restore ", @@ -48,12 +51,16 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command { defer end() if db.Kind == "mysql" { + if cmd.Flags().Changed("replicas") { + return fmt.Errorf("--replicas is only supported for PostgreSQL backup restores") + } + newBranch, err := client.DatabaseBranches.Create(ctx, &planetscale.CreateDatabaseBranchRequest{ Organization: ch.Config.Organization, Database: database, Name: branchName, BackupID: backup, - ClusterSize: clusterSize, + ClusterSize: flags.clusterSize, }) if err != nil { return cmdutil.HandleError(err) @@ -62,13 +69,19 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command { end() return ch.Printer.PrintResource(branch.ToDatabaseBranch(newBranch)) } else { - newBranch, err := client.PostgresBranches.Create(ctx, &planetscale.CreatePostgresBranchRequest{ + createReq := &planetscale.CreatePostgresBranchRequest{ Organization: ch.Config.Organization, Database: database, Name: branchName, BackupID: backup, - ClusterName: clusterSize, - }) + ClusterName: flags.clusterSize, + } + if cmd.Flags().Changed("replicas") { + replicas := flags.replicas + createReq.Replicas = &replicas + } + + newBranch, err := client.PostgresBranches.Create(ctx, createReq) if err != nil { return cmdutil.HandleError(err) } @@ -79,7 +92,8 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command { }, } - cmd.Flags().StringVar(&clusterSize, "cluster-size", "PS-10", "Cluster size for restored backup branch. Use `pscale size cluster list` to see the valid sizes.") + cmd.Flags().StringVar(&flags.clusterSize, "cluster-size", "PS-10", "Cluster size for restored backup branch. Use `pscale size cluster list` to see the valid sizes.") + cmd.Flags().IntVar(&flags.replicas, "replicas", 0, "Number of additional replicas for a PostgreSQL restore. 0 creates a single-node branch; omit to use the target cluster size default.") cmd.MarkFlagRequired("cluster-size") cmd.RegisterFlagCompletionFunc("cluster-size", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { return cmdutil.ClusterSizesCompletionFunc(ch, cmd, args, toComplete) diff --git a/internal/cmd/backup/restore_test.go b/internal/cmd/backup/restore_test.go index 8574a76b..6b1506ed 100644 --- a/internal/cmd/backup/restore_test.go +++ b/internal/cmd/backup/restore_test.go @@ -92,6 +92,8 @@ func TestBackup_RestoreCmd_PostgreSQL(t *testing.T) { c.Assert(req.Name, qt.Equals, branch) c.Assert(req.BackupID, qt.Equals, backup) c.Assert(req.ClusterName, qt.Equals, "PS-20") + c.Assert(req.Replicas, qt.IsNotNil) + c.Assert(*req.Replicas, qt.Equals, 2) return res, nil }, } @@ -118,10 +120,41 @@ func TestBackup_RestoreCmd_PostgreSQL(t *testing.T) { } cmd := RestoreCmd(ch) - cmd.SetArgs([]string{db, branch, backup, "--cluster-size", "PS-20"}) + cmd.SetArgs([]string{db, branch, backup, "--cluster-size", "PS-20", "--replicas", "2"}) err := cmd.Execute() c.Assert(err, qt.IsNil) c.Assert(svc.CreateFnInvoked, qt.IsTrue) c.Assert(buf.String(), qt.JSONEquals, res) } + +func TestBackup_RestoreCmdRejectsReplicasForMySQL(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEngineMySQL}, 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 --replicas") + return nil, nil + }, + } + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: dbSvc, DatabaseBranches: svc}, nil + }, + } + + cmd := RestoreCmd(ch) + cmd.SetArgs([]string{"planetscale", "restored", "backup-id", "--cluster-size", "PS-20", "--replicas", "2"}) + + c.Assert(cmd.Execute(), qt.ErrorMatches, ".*--replicas is only supported for PostgreSQL.*") + c.Assert(svc.CreateFnInvoked, qt.IsFalse) +} diff --git a/internal/cmd/branch/create.go b/internal/cmd/branch/create.go index ef0ed11d..16231121 100644 --- a/internal/cmd/branch/create.go +++ b/internal/cmd/branch/create.go @@ -22,6 +22,7 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { backupID string restorePoint string majorVersion string + replicas int minStorage int64 maxStorage int64 } @@ -75,6 +76,9 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { if flags.backupID != "" && flags.parentBranch != "" && flags.restorePoint == "" { return fmt.Errorf("--from and --restore cannot be used together") } + if cmd.Flags().Changed("replicas") && flags.backupID == "" && flags.restorePoint == "" { + return fmt.Errorf("--replicas can only be used with a PostgreSQL backup restore or point-in-time recovery") + } client, err := ch.Client() if err != nil { @@ -118,6 +122,9 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { } if db.Kind == "mysql" { + if cmd.Flags().Changed("replicas") { + return fmt.Errorf("--replicas is only supported for PostgreSQL backup restores and point-in-time recovery") + } 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") } @@ -209,6 +216,11 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { MajorVersion: flags.majorVersion, } + if cmd.Flags().Changed("replicas") { + replicas := flags.replicas + createReq.Replicas = &replicas + } + if cmd.Flags().Changed("min-storage") || cmd.Flags().Changed("max-storage") { createReq.Storage = &ps.StorageConfig{} if cmd.Flags().Changed("min-storage") { @@ -274,6 +286,7 @@ 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().IntVar(&flags.replicas, "replicas", 0, "Number of additional replicas for a PostgreSQL restore. 0 creates a single-node branch; omit to use the target cluster size default.") 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") diff --git a/internal/cmd/branch/create_test.go b/internal/cmd/branch/create_test.go index 510b3835..2188cfae 100644 --- a/internal/cmd/branch/create_test.go +++ b/internal/cmd/branch/create_test.go @@ -3,6 +3,7 @@ package branch import ( "bytes" "context" + "errors" "testing" "time" @@ -779,6 +780,8 @@ func TestBranch_CreateCmdWithRestorePoint(t *testing.T) { c.Assert(req.ParentBranch, qt.Equals, parentBranch) c.Assert(req.BackupID, qt.Equals, backupID) c.Assert(req.ClusterName, qt.Equals, "PS-10") + c.Assert(req.Replicas, qt.IsNotNil) + c.Assert(*req.Replicas, qt.Equals, 3) return res, nil }, @@ -807,7 +810,7 @@ func TestBranch_CreateCmdWithRestorePoint(t *testing.T) { } cmd := CreateCmd(ch) - cmd.SetArgs([]string{db, branch, "--region", "us-east", "--from", parentBranch, "--restore-point", restorePoint}) + cmd.SetArgs([]string{db, branch, "--region", "us-east", "--from", parentBranch, "--restore-point", restorePoint, "--replicas", "3"}) err := cmd.Execute() c.Assert(err, qt.IsNil) @@ -1002,3 +1005,140 @@ func TestBranch_CreateCmdWithRestorePointMySQLError(t *testing.T) { c.Assert(err, qt.ErrorMatches, ".*only supported for PostgreSQL.*") c.Assert(svc.CreateFnInvoked, qt.IsFalse) } + +func TestBranch_CreateCmdWithPostgresRestoreReplicas(t *testing.T) { + zero := 0 + tests := []struct { + name string + args []string + wantReplicas *int + }{ + {name: "omitted", args: []string{"planetscale", "restored", "--restore", "backup-id"}}, + {name: "explicit zero", args: []string{"planetscale", "restored", "--restore", "backup-id", "--replicas", "0"}, wantReplicas: &zero}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + res := &ps.PostgresBranch{Name: "restored"} + svc := &mock.PostgresBranchesService{ + CreateFn: func(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) { + c.Assert(req.BackupID, qt.Equals, "backup-id") + if tt.wantReplicas == nil { + c.Assert(req.Replicas, qt.IsNil) + } else { + c.Assert(req.Replicas, qt.IsNotNil) + c.Assert(*req.Replicas, qt.Equals, *tt.wantReplicas) + } + return res, nil + }, + } + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEnginePostgres}, nil + }, + } + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: dbSvc, PostgresBranches: svc}, nil + }, + } + + cmd := CreateCmd(ch) + cmd.SetArgs(tt.args) + + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.CreateFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, res) + }) + } +} + +func TestBranch_CreateCmdRejectsReplicasWithoutRestore(t *testing.T) { + tests := []struct { + name string + args []string + }{ + {name: "ordinary branch", args: []string{"planetscale", "development", "--replicas", "2"}}, + {name: "data branching", args: []string{"planetscale", "development", "--seed-data", "--replicas", "2"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := qt.New(t) + cmd := CreateCmd(&cmdutil.Helper{}) + cmd.SetArgs(tt.args) + + c.Assert(cmd.Execute(), qt.ErrorMatches, ".*--replicas can only be used with a PostgreSQL backup restore or point-in-time recovery.*") + }) + } +} + +func TestBranch_CreateCmdRejectsRestoreReplicasForMySQL(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEngineMySQL}, 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 --replicas") + return nil, nil + }, + } + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: dbSvc, DatabaseBranches: svc}, nil + }, + } + + cmd := CreateCmd(ch) + cmd.SetArgs([]string{"planetscale", "restored", "--restore", "backup-id", "--replicas", "2"}) + + c.Assert(cmd.Execute(), qt.ErrorMatches, ".*--replicas is only supported for PostgreSQL.*") + c.Assert(svc.CreateFnInvoked, qt.IsFalse) +} + +func TestBranch_CreateCmdPropagatesReplicaValidationError(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + validationErr := errors.New("replica count is not valid for the selected cluster size") + svc := &mock.PostgresBranchesService{ + CreateFn: func(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) { + return nil, validationErr + }, + } + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEnginePostgres}, nil + }, + } + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: dbSvc, PostgresBranches: svc}, nil + }, + } + + cmd := CreateCmd(ch) + cmd.SetArgs([]string{"planetscale", "restored", "--restore", "backup-id", "--replicas", "1"}) + + c.Assert(cmd.Execute(), qt.Equals, validationErr) +} diff --git a/internal/planetscale/postgres_branches.go b/internal/planetscale/postgres_branches.go index de15148a..6fb4f16c 100644 --- a/internal/planetscale/postgres_branches.go +++ b/internal/planetscale/postgres_branches.go @@ -44,6 +44,7 @@ type CreatePostgresBranchRequest struct { RestorePoint string `json:"restore_point,omitempty"` ClusterName string `json:"cluster_name,omitempty"` MajorVersion string `json:"major_version,omitempty"` + Replicas *int `json:"replicas,omitempty"` Storage *StorageConfig `json:"storage,omitempty"` } diff --git a/internal/planetscale/postgres_branches_test.go b/internal/planetscale/postgres_branches_test.go index f07a01a4..ca5a0e56 100644 --- a/internal/planetscale/postgres_branches_test.go +++ b/internal/planetscale/postgres_branches_test.go @@ -18,6 +18,11 @@ func TestPostgresBranches_Create(t *testing.T) { c := qt.New(t) ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + c.Assert(json.NewDecoder(r.Body).Decode(&body), qt.IsNil) + _, hasReplicas := body["replicas"] + c.Assert(hasReplicas, qt.IsFalse) + w.WriteHeader(200) out := `{"id":"postgres-test-branch","name":"postgres-test-branch","created_at":"2021-01-14T10:19:23.000Z","updated_at":"2021-01-14T10:19:23.000Z", "region": {"slug": "us-west", "display_name": "US West"}}` _, err := w.Write([]byte(out)) @@ -54,6 +59,36 @@ func TestPostgresBranches_Create(t *testing.T) { c.Assert(branch, qt.DeepEquals, want) } +func TestCreatePostgresBranchRequestSerializesExplicitZeroReplicas(t *testing.T) { + c := qt.New(t) + replicas := 0 + + body, err := json.Marshal(&CreatePostgresBranchRequest{ + Name: testPostgresBranch, + Replicas: &replicas, + }) + c.Assert(err, qt.IsNil) + + var decoded map[string]any + c.Assert(json.Unmarshal(body, &decoded), qt.IsNil) + c.Assert(decoded["replicas"], qt.Equals, float64(0)) +} + +func TestCreatePostgresBranchRequestSerializesNonzeroReplicas(t *testing.T) { + c := qt.New(t) + replicas := 3 + + body, err := json.Marshal(&CreatePostgresBranchRequest{ + Name: testPostgresBranch, + Replicas: &replicas, + }) + c.Assert(err, qt.IsNil) + + var decoded map[string]any + c.Assert(json.Unmarshal(body, &decoded), qt.IsNil) + c.Assert(decoded["replicas"], qt.Equals, float64(3)) +} + func TestPostgresBranches_List(t *testing.T) { c := qt.New(t)