Update promote/demote cli

internal/cmd/branch/branch.go

@@ -59,8 +59,8 @@ func BranchCmd(ch *cmdutil.Helper) *cobra.Command {
 	add(branchGroupDatabase, UpdateCmd(ch))
 	add(branchGroupDatabase, SwitchCmd(ch))
 	add(branchGroupDatabase, SchemaCmd(ch))
-	add(branchGroupDatabase, PromoteCmd(ch))
-	add(branchGroupDatabase, DemoteCmd(ch))
+	add(branchGroupMySQLPostgres, PromoteCmd(ch))
+	add(branchGroupMySQLPostgres, DemoteCmd(ch))
 	add(branchGroupMySQLPostgres, ConnectionsCmd(ch))
 	add(branchGroupDatabase, InfraCmd(ch))
 	add(branchGroupVitess, DiffCmd(ch))

internal/cmd/branch/branch_test.go

@@ -30,6 +30,8 @@ func TestBranchCmdGroupsCommandsByEngine(t *testing.T) {
 	}
 
 	c.Assert(got["create"], qt.Equals, branchGroupDatabase)
+	c.Assert(got["promote"], qt.Equals, branchGroupMySQLPostgres)
+	c.Assert(got["demote"], qt.Equals, branchGroupMySQLPostgres)
 	c.Assert(got["connections"], qt.Equals, branchGroupMySQLPostgres)
 	c.Assert(got["lint"], qt.Equals, branchGroupVitess)
 	c.Assert(got["routing-rules"], qt.Equals, branchGroupVitess)

internal/cmd/branch/demote.go

@@ -11,15 +11,41 @@ import (
 )
 
 func DemoteCmd(ch *cmdutil.Helper) *cobra.Command {
-	return &cobra.Command{
-		Use:   "demote <database> <branch> [options]",
+	var flags struct {
+		clusterSize string
+		replicas    int
+	}
+
+	cmd := &cobra.Command{
+		Use:   "demote <database> <branch>",
 		Short: "Demote a production branch to development",
-		Args:  cmdutil.RequiredArgs("database", "branch"),
+		Long: `Demote a production branch to development.
+
+This command is supported for MySQL databases and for enabled Postgres branches. Neki databases are not supported.
+
+--cluster-size sets the Postgres cluster size to use. PS-DEV is allowed. Omit
+it to leave the size unchanged. Changing the size requires the write_database
+permission.
+
+--replicas sets the Postgres replica count, not including the primary. 0 is a
+single node. 2 to 8 is high availability. Omit it to leave the count
+unchanged. Changing the replica count requires the write_database permission.
+
+Vitess ignores both options.`,
+		Example: `  pscale branch demote mydb main
+  pscale branch demote mydb main --cluster-size PS-DEV
+  pscale branch demote mydb main --cluster-size PS-10 --replicas 0`,
+		Args: cmdutil.RequiredArgs("database", "branch"),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			req := &ps.DemoteRequest{
 				Database:     args[0],
 				Branch:       args[1],
 				Organization: ch.Config.Organization,
+				ClusterSize:  cmdutil.ToSizeSKUName(flags.clusterSize),
+			}
+			if cmd.Flags().Changed("replicas") {
+				replicas := flags.replicas
+				req.Replicas = &replicas
 			}
 
 			client, err := ch.Client()
@@ -54,4 +80,13 @@ func DemoteCmd(ch *cmdutil.Helper) *cobra.Command {
 			return ch.Printer.PrintResource(ToDatabaseBranch(b))
 		},
 	}
+
+	cmd.Flags().StringVar(&flags.clusterSize, "cluster-size", "", "Postgres cluster size to apply when demoting (for example, PS-DEV). Use 'pscale size cluster list' to see valid sizes.")
+	cmd.Flags().IntVar(&flags.replicas, "replicas", 0, "Postgres replica count, not including the primary. 0 is a single node. 2 to 8 is high availability.")
+
+	cmd.RegisterFlagCompletionFunc("cluster-size", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		return cmdutil.ClusterSizesIncludingDevelopmentCompletionFunc(ch, cmd, args, toComplete)
+	})
+
+	return cmd
 }

internal/cmd/branch/demote_test.go

@@ -36,6 +36,8 @@ func TestBranch_DemoteCmd(t *testing.T) {
 			c.Assert(req.Branch, qt.Equals, branch)
 			c.Assert(req.Database, qt.Equals, db)
 			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.ClusterSize, qt.Equals, "")
+			c.Assert(req.Replicas, qt.IsNil)
 
 			return res, nil
 		},
@@ -62,6 +64,48 @@ func TestBranch_DemoteCmd(t *testing.T) {
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
 
+func TestBranch_DemoteCmd_ClusterSizeAndReplicas(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 := "main"
+	replicas := 0
+
+	svc := &mock.DatabaseBranchesService{
+		DemoteFn: func(ctx context.Context, req *ps.DemoteRequest) (*ps.DatabaseBranch, error) {
+			c.Assert(req.ClusterSize, qt.Equals, "PS_DEV")
+			c.Assert(req.Replicas, qt.DeepEquals, &replicas)
+			return &ps.DatabaseBranch{Name: branch, Production: false}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{DatabaseBranches: svc}, nil
+		},
+	}
+
+	cmd := DemoteCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--cluster-size", "PS-DEV", "--replicas", "0"})
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(svc.DemoteFnInvoked, qt.IsTrue)
+}
+
+func TestBranch_DemoteCmd_HelpStatesEngineSupport(t *testing.T) {
+	c := qt.New(t)
+	cmd := DemoteCmd(&cmdutil.Helper{})
+
+	c.Assert(cmd.Long, qt.Contains, "This command is supported for MySQL databases and for enabled Postgres branches. Neki databases are not supported.")
+}
+
 func TestBranch_DemoteCmd_ServiceTokenPermissionError(t *testing.T) {
 	c := qt.New(t)
 

internal/cmd/branch/promote.go

@@ -11,12 +11,32 @@ import (
 )
 
 func PromoteCmd(ch *cmdutil.Helper) *cobra.Command {
+	var flags struct {
+		clusterSize string
+		replicas    int
+	}
+
 	promoteReq := &ps.PromoteRequest{}
 
 	cmd := &cobra.Command{
-		Use:   "promote <database> <branch> [options]",
-		Short: "Promote a new branch from a database",
-		Args:  cmdutil.RequiredArgs("database", "branch"),
+		Use:   "promote <database> <branch>",
+		Short: "Promote a branch to production",
+		Long: `Promote a branch to production.
+
+This command is supported for MySQL databases and for enabled Postgres branches. Neki databases are not supported.
+
+--cluster-size sets the cluster size to use. PS-DEV is not permitted. When
+omitted, a PS-DEV Postgres branch uses the smallest production size. For
+MySQL, the size applies to the default keyspace. Changing the size requires
+the write_database permission.
+
+--replicas sets the Postgres replica count, not including the primary. 0 is a
+single node. 2 to 8 is high availability. Omit it to leave the count
+unchanged. Changing the replica count requires the write_database permission.`,
+		Example: `  pscale branch promote mydb dev
+  pscale branch promote mydb dev --cluster-size PS-10
+  pscale branch promote mydb dev --cluster-size PS-10 --replicas 2`,
+		Args: cmdutil.RequiredArgs("database", "branch"),
 		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 			if len(args) != 0 {
 				return nil, cobra.ShellCompDirectiveNoFileComp
@@ -58,6 +78,11 @@ func PromoteCmd(ch *cmdutil.Helper) *cobra.Command {
 			promoteReq.Database = db
 			promoteReq.Organization = ch.Config.Organization
 			promoteReq.Branch = branch
+			promoteReq.ClusterSize = cmdutil.ToSizeSKUName(flags.clusterSize)
+			if cmd.Flags().Changed("replicas") {
+				replicas := flags.replicas
+				promoteReq.Replicas = &replicas
+			}
 
 			client, err := ch.Client()
 			if err != nil {
@@ -80,13 +105,23 @@ func PromoteCmd(ch *cmdutil.Helper) *cobra.Command {
 			}
 
 			if ch.Printer.Format() == printer.Human {
-				ch.Printer.Printf("Branch %s in %s was successfully promoted to production. Note: In order to use deploy requests to deploy schema changes to this branch, safe migrations must also be enabled.\n", printer.BoldBlue(branch), printer.BoldBlue(db))
+				ch.Printer.Printf("Branch %s in %s was successfully promoted to production.\n", printer.BoldBlue(branch), printer.BoldBlue(db))
+				if dbBranch.Kind == "" || dbBranch.Kind == ps.DatabaseEngineMySQL {
+					ch.Printer.Printf("Note: In order to use deploy requests to deploy schema changes to this branch, safe migrations must also be enabled.\n")
+				}
 				return nil
 			}
 
 			return ch.Printer.PrintResource(ToDatabaseBranch(dbBranch))
 		},
 	}
 
+	cmd.Flags().StringVar(&flags.clusterSize, "cluster-size", "", "Cluster size to apply when promoting (for example, PS-10). PS-DEV is not permitted. Use 'pscale size cluster list' to see valid sizes.")
+	cmd.Flags().IntVar(&flags.replicas, "replicas", 0, "Postgres replica count, not including the primary. 0 is a single node. 2 to 8 is high availability.")
+
+	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/promote_test.go

@@ -36,6 +36,8 @@ func TestBranch_PromoteCmd(t *testing.T) {
 			c.Assert(req.Branch, qt.Equals, branch)
 			c.Assert(req.Database, qt.Equals, db)
 			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.ClusterSize, qt.Equals, "")
+			c.Assert(req.Replicas, qt.IsNil)
 
 			return res, nil
 		},
@@ -62,6 +64,103 @@ func TestBranch_PromoteCmd(t *testing.T) {
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
 
+func TestBranch_PromoteCmd_ClusterSizeAndReplicas(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"
+	replicas := 0
+
+	svc := &mock.DatabaseBranchesService{
+		PromoteFn: func(ctx context.Context, req *ps.PromoteRequest) (*ps.DatabaseBranch, error) {
+			c.Assert(req.ClusterSize, qt.Equals, "PS_10")
+			c.Assert(req.Replicas, qt.DeepEquals, &replicas)
+			return &ps.DatabaseBranch{Name: branch, Production: true}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{DatabaseBranches: svc}, nil
+		},
+	}
+
+	cmd := PromoteCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--cluster-size", "PS-10", "--replicas", "0"})
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(svc.PromoteFnInvoked, qt.IsTrue)
+}
+
+func TestBranch_PromoteCmd_HelpStatesEngineSupport(t *testing.T) {
+	c := qt.New(t)
+	cmd := PromoteCmd(&cmdutil.Helper{})
+
+	c.Assert(cmd.Long, qt.Contains, "This command is supported for MySQL databases and for enabled Postgres branches. Neki databases are not supported.")
+}
+
+func TestBranch_PromoteCmd_MySQLSuccessText(t *testing.T) {
+	c := qt.New(t)
+	out := promoteHumanOutput(c, ps.DatabaseEngineMySQL)
+
+	c.Assert(out, qt.Contains, "successfully promoted to production")
+	c.Assert(out, qt.Contains, "safe migrations must also be enabled")
+}
+
+func TestBranch_PromoteCmd_PostgresSuccessText(t *testing.T) {
+	c := qt.New(t)
+	out := promoteHumanOutput(c, ps.DatabaseEnginePostgres)
+
+	c.Assert(out, qt.Contains, "successfully promoted to production")
+	c.Assert(out, qt.Not(qt.Contains), "safe migrations")
+}
+
+func TestBranch_PromoteCmd_NekiSuccessText(t *testing.T) {
+	c := qt.New(t)
+	out := promoteHumanOutput(c, ps.DatabaseEngineNeki)
+
+	c.Assert(out, qt.Contains, "successfully promoted to production")
+	c.Assert(out, qt.Not(qt.Contains), "safe migrations")
+}
+
+func promoteHumanOutput(c *qt.C, kind ps.DatabaseEngine) string {
+	var buf bytes.Buffer
+	format := printer.Human
+	p := printer.NewPrinter(&format)
+	p.SetHumanOutput(&buf)
+
+	org := "planetscale"
+	db := "planetscale"
+	branch := "development"
+
+	svc := &mock.DatabaseBranchesService{
+		PromoteFn: func(ctx context.Context, req *ps.PromoteRequest) (*ps.DatabaseBranch, error) {
+			return &ps.DatabaseBranch{Name: branch, Production: true, Kind: kind}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{DatabaseBranches: svc}, nil
+		},
+	}
+
+	cmd := PromoteCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	c.Assert(cmd.Execute(), qt.IsNil)
+
+	return buf.String()
+}
+
 func TestBranch_PromoteCmd_ServiceTokenPermissionError(t *testing.T) {
 	c := qt.New(t)
 

internal/cmdutil/completions.go

@@ -9,14 +9,20 @@ import (
 )
 
 func ClusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
-	return clusterSizesCompletionFunc(ch, cmd, toComplete)
+	return clusterSizesCompletionFunc(ch, cmd, toComplete, false)
+}
+
+// ClusterSizesIncludingDevelopmentCompletionFunc completes cluster sizes and
+// includes PS-DEV. ClusterSizesCompletionFunc omits that size.
+func ClusterSizesIncludingDevelopmentCompletionFunc(ch *Helper, cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
+	return clusterSizesCompletionFunc(ch, cmd, toComplete, true)
 }
 
 func ExternalClusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
-	return clusterSizesCompletionFunc(ch, cmd, toComplete, ps.WithExternal())
+	return clusterSizesCompletionFunc(ch, cmd, toComplete, false, ps.WithExternal())
 }
 
-func clusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, toComplete string, extraOpts ...ps.ListOption) ([]cobra.Completion, cobra.ShellCompDirective) {
+func clusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, toComplete string, includeDevelopment bool, extraOpts ...ps.ListOption) ([]cobra.Completion, cobra.ShellCompDirective) {
 	ctx := cmd.Context()
 
 	org := ch.Config.Organization // --org flag
@@ -52,7 +58,7 @@ func clusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, toComplete strin
 
 	clusterSizes := make([]cobra.Completion, 0)
 	for _, c := range clusterSKUs {
-		if c.Enabled && strings.Contains(c.Name, toComplete) && c.Rate != nil && c.Name != "PS_DEV" {
+		if c.Enabled && strings.Contains(c.Name, toComplete) && c.Rate != nil && (includeDevelopment || c.Name != "PS_DEV") {
 			var description strings.Builder
 			if c.Metal {
 				description.WriteString(c.DisplayName)