Add cluster-size to backup restoration commands (#1023)

* Add cluster size to restore commands

* Add cluster size to backup restore command

* Update tests

* Update planetscale-go

* Require restore and cluster-size flags to be present together

* Make PS-20 flag required for backup restoration

* Remove deprecated flag

Author: iheanyi

go.mod

@@ -23,7 +23,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.131.0
+	github.com/planetscale/planetscale-go v0.132.0
 	github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7
 	github.com/planetscale/psdbproxy v0.0.0-20250117221522-0c8e2b0e36e6
 	github.com/spf13/cobra v1.9.1

go.sum

@@ -142,8 +142,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.131.0 h1:ojr8A1yrkJMRS3oUcJH58mv5efRQwCV4MFFQYJnLs9s=
-github.com/planetscale/planetscale-go v0.131.0/go.mod h1:VE4Vn8+qWV8KlfcasKurXByncyqaL4NIA68DGep+AFg=
+github.com/planetscale/planetscale-go v0.132.0 h1:SoQczeRfMUwpNSlZwsht+T+x4BlvOCJZA3+xW52In6c=
+github.com/planetscale/planetscale-go v0.132.0/go.mod h1:Ht/CUQOhz22XoXd5+ouiQAFfPWSyiGy5dATcyyoXcq0=
 github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7 h1:dxdoFKWVDlV1gq8UQC8NWCofLjCEjEHw47gfeojgs28=
 github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7/go.mod h1:WZmi4gw3rOK+ryd1inGxgfKwoFV04O7xBCqzWzv0/0U=
 github.com/planetscale/psdbproxy v0.0.0-20250117221522-0c8e2b0e36e6 h1:/Ox1ZTAdk+soSngzzKoJh5voOzptrpPrux11o30rIaw=

internal/cmd/backup/restore.go

@@ -13,6 +13,8 @@ import (
 )
 
 func RestoreCmd(ch *cmdutil.Helper) *cobra.Command {
+	var clusterSize string
+
 	cmd := &cobra.Command{
 		Use:   "restore <database> <branch> <backup>",
 		Short: "Restore a backup to a new branch",
@@ -35,6 +37,7 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command {
 				Database:     database,
 				Name:         branchName,
 				BackupID:     backup,
+				ClusterSize:  clusterSize,
 			})
 			if err != nil {
 				return cmdutil.HandleError(err)
@@ -45,5 +48,11 @@ 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.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)
+	})
+
 	return cmd
 }

internal/cmd/backup/restore_test.go

@@ -35,6 +35,7 @@ func TestBackup_RestoreCmd(t *testing.T) {
 			c.Assert(req.Database, qt.Equals, db)
 			c.Assert(req.Name, qt.Equals, branch)
 			c.Assert(req.BackupID, qt.Equals, backup)
+			c.Assert(req.ClusterSize, qt.Equals, "PS-20")
 			return res, nil
 		},
 	}
@@ -52,7 +53,7 @@ func TestBackup_RestoreCmd(t *testing.T) {
 	}
 
 	cmd := RestoreCmd(ch)
-	cmd.SetArgs([]string{db, branch, backup})
+	cmd.SetArgs([]string{db, branch, backup, "--cluster-size", "PS-20"})
 	err := cmd.Execute()
 
 	c.Assert(err, qt.IsNil)

internal/cmd/branch/create.go

@@ -119,18 +119,23 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	}
 
 	cmd.Flags().StringVar(&createReq.ParentBranch, "from", "", "Parent branch to create the new branch from. Cannot be used with --restore")
-	cmd.Flags().StringVar(&createReq.Region, "region", "", "Region for the branch. Can not be used with --restore")
-	cmd.Flags().StringVar(&createReq.BackupID, "restore", "", "Backup to restore into the branch. Backup can only be restored to its original region")
+	cmd.Flags().StringVar(&createReq.Region, "region", "", "Region for the branch to be created in.")
+	cmd.Flags().StringVar(&createReq.BackupID, "restore", "", "ID of Backup to restore into branch.")
+	cmd.Flags().StringVar(&createReq.ClusterSize, "cluster-size", "PS-10", "Cluster size for branches being created from a backup or seeded with data. Use `pscale size cluster list` to see the valid sizes.")
 	cmd.Flags().BoolVar(&flags.dataBranching, "seed-data", false, "Add seed data using the Data Branching™ feature")
 	cmd.Flags().BoolVar(&flags.wait, "wait", false, "Wait until the branch is ready")
 	cmd.MarkFlagsMutuallyExclusive("from", "restore")
-	cmd.MarkFlagsMutuallyExclusive("region", "restore")
 	cmd.MarkFlagsMutuallyExclusive("restore", "seed-data")
+	cmd.MarkFlagsRequiredTogether("restore", "cluster-size")
 
 	cmd.RegisterFlagCompletionFunc("region", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		return cmdutil.RegionsCompletionFunc(ch, cmd, args, toComplete)
 	})
 
+	cmd.RegisterFlagCompletionFunc("cluster-size", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		return cmdutil.ClusterSizesCompletionFunc(ch, cmd, args, toComplete)
+	})
+
 	return cmd
 }
 

internal/cmd/branch/create_test.go

@@ -59,3 +59,98 @@ func TestBranch_CreateCmd(t *testing.T) {
 	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
+
+func TestBranch_CreateCmdWithRestore(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.DatabaseBranch{Name: branch}
+
+	svc := &mock.DatabaseBranchesService{
+		CreateFn: func(ctx context.Context, req *ps.CreateDatabaseBranchRequest) (*ps.DatabaseBranch, 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.BackupID, qt.Equals, "somebackupid")
+			c.Assert(req.ClusterSize, qt.Equals, "PS-20")
+
+			return res, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				DatabaseBranches: svc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--region", "us-east", "--restore", "somebackupid", "--cluster-size", "PS-20"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}
+
+func TestBranch_CreateCmdWithSeedData(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.DatabaseBranch{Name: branch}
+
+	svc := &mock.DatabaseBranchesService{
+		CreateFn: func(ctx context.Context, req *ps.CreateDatabaseBranchRequest) (*ps.DatabaseBranch, 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.SeedData, qt.Equals, "last_successful_backup")
+
+			return res, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				DatabaseBranches: svc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--region", "us-east", "--seed-data"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}

internal/cmd/database/create.go

@@ -60,8 +60,6 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 		},
 	}
 
-	cmd.Flags().StringVar(&createReq.Notes, "notes", "", "notes for the database")
-	cmd.Flags().MarkDeprecated("notes", "is no longer available.")
 	cmd.Flags().StringVar(&createReq.Region, "region", "", "region for the database")
 
 	cmd.Flags().String("cluster-size", "PS-10", "cluster size for Scaler Pro databases. Use `pscale size cluster list` to see the valid sizes.")