Add --replicas flag to database create command

Author: no-itsbackpack

internal/cmd/database/create.go

@@ -22,6 +22,7 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 		clusterSize string
 		engine      string
 		wait        bool
+		replicas    *int
 	}
 
 	cmd := &cobra.Command{
@@ -41,6 +42,10 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 
 			createReq.Kind = engine
 
+			if cmd.Flags().Changed("replicas") {
+				createReq.Replicas = flags.replicas
+			}
+
 			client, err := ch.Client()
 			if err != nil {
 				return err
@@ -85,6 +90,9 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.Flags().StringVar(&createReq.Region, "region", "", "region for the database")
 
 	cmd.Flags().StringVar(&createReq.ClusterSize, "cluster-size", "", "cluster size for Scaler Pro databases. Use `pscale size cluster list` to see the valid sizes.")
+
+	flags.replicas = cmd.Flags().Int("replicas", 0, "The number of replicas for the database. 0 for single non-HA node, 2+ for HA.")
+
 	cmd.Flags().StringVar(&flags.engine, "engine", string(ps.DatabaseEngineMySQL), "The database engine for the database. Supported values: mysql, postgresql. Defaults to mysql.")
 	cmd.RegisterFlagCompletionFunc("engine", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
 		return []cobra.Completion{

internal/cmd/database/create_test.go

@@ -32,6 +32,7 @@ func TestDatabase_CreateCmd(t *testing.T) {
 			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.Replicas, qt.IsNil) // Should be nil when not set
 
 			return res, nil
 		},
@@ -66,6 +67,60 @@ func TestDatabase_CreateCmd(t *testing.T) {
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
 
+func TestDatabase_CreateCmdWithReplicasZero(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.Replicas, qt.IsNotNil)   // Should be set when explicitly passed
+			c.Assert(*req.Replicas, qt.Equals, 0) // Should be 0
+
+			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", "--replicas", "0"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}
+
 func TestDatabase_CreateCmdPostgres(t *testing.T) {
 	c := qt.New(t)
 
@@ -118,3 +173,57 @@ func TestDatabase_CreateCmdPostgres(t *testing.T) {
 	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
+
+func TestDatabase_CreateCmdWithReplicas(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.Replicas, qt.IsNotNil)
+			c.Assert(*req.Replicas, qt.Equals, 3)
+
+			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", "--replicas", "3"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}