Add replica selection to Postgres restores

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 <database> <branch> <backup>",
@@ -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)

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)
+}

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")
 

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)
+}

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"`
 }
 

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)