Additional vitess region commands

internal/cmd/keyspace/read_only_regions.go

@@ -54,6 +54,12 @@ func ReadOnlyRegionsCmd(ch *cmdutil.Helper) *cobra.Command {
 		},
 	}
 
+	cmd.AddCommand(
+		ReadOnlyRegionsAddCmd(ch),
+		ReadOnlyRegionsUpdateCmd(ch),
+		ReadOnlyRegionsRemoveCmd(ch),
+	)
+
 	return cmd
 }
 
@@ -85,3 +91,73 @@ func (r *ReadOnlyRegion) MarshalJSON() ([]byte, error) {
 func (r *ReadOnlyRegion) MarshalCSVValue() interface{} {
 	return []*ReadOnlyRegion{r}
 }
+
+func readOnlyRegionConfigs(regions []*ps.ReadOnlyRegionKeyspace) []*ps.ReadOnlyRegionKeyspaceConfig {
+	configs := make([]*ps.ReadOnlyRegionKeyspaceConfig, 0, len(regions))
+	for _, region := range regions {
+		clusterSize := region.ClusterName
+		replicas := region.Replicas
+		configs = append(configs, &ps.ReadOnlyRegionKeyspaceConfig{
+			Region:      region.Region,
+			ClusterSize: &clusterSize,
+			Replicas:    &replicas,
+		})
+	}
+	return configs
+}
+
+func readOnlyRegionIndex(regions []*ps.ReadOnlyRegionKeyspace, name string) int {
+	for i, region := range regions {
+		if region.Region == name {
+			return i
+		}
+	}
+	return -1
+}
+
+func getReadOnlyRegions(cmd *cobra.Command, ch *cmdutil.Helper, database, branch, keyspace string) (*ps.Client, []*ps.ReadOnlyRegionKeyspace, error) {
+	client, err := ch.Client()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	k, err := client.Keyspaces.Get(cmd.Context(), &ps.GetKeyspaceRequest{
+		Organization: ch.Config.Organization,
+		Database:     database,
+		Branch:       branch,
+		Keyspace:     keyspace,
+		Full:         true,
+	})
+	if err != nil {
+		switch cmdutil.ErrCode(err) {
+		case ps.ErrNotFound:
+			return nil, nil, fmt.Errorf("keyspace %s does not exist in branch %s (database: %s, organization: %s)", printer.BoldBlue(keyspace), printer.BoldBlue(branch), printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization))
+		default:
+			return nil, nil, cmdutil.HandleError(err)
+		}
+	}
+
+	return client, k.ReadOnlyRegions, nil
+}
+
+func updateReadOnlyRegions(cmd *cobra.Command, ch *cmdutil.Helper, client *ps.Client, database, branch, keyspace string, configs []*ps.ReadOnlyRegionKeyspaceConfig) ([]*ps.ReadOnlyRegionKeyspace, error) {
+	regions, err := client.Keyspaces.UpdateReadOnlyRegions(cmd.Context(), &ps.UpdateReadOnlyRegionsRequest{
+		Organization:    ch.Config.Organization,
+		Database:        database,
+		Branch:          branch,
+		Keyspace:        keyspace,
+		ReadOnlyRegions: configs,
+	})
+	if err != nil {
+		return nil, cmdutil.HandleError(err)
+	}
+	return regions, nil
+}
+
+func printReadOnlyRegionMutation(ch *cmdutil.Helper, message string, regions []*ps.ReadOnlyRegionKeyspace) error {
+	if ch.Printer.Format() == printer.Human {
+		ch.Printer.Println(message)
+		return nil
+	}
+	return ch.Printer.PrintResource(toReadOnlyRegions(regions))
+}

internal/cmd/keyspace/read_only_regions_add.go

@@ -0,0 +1,64 @@
+package keyspace
+
+import (
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/planetscale/cli/internal/printer"
+	"github.com/spf13/cobra"
+)
+
+func ReadOnlyRegionsAddCmd(ch *cmdutil.Helper) *cobra.Command {
+	var clusterSize string
+	var replicas int
+
+	cmd := &cobra.Command{
+		Use:   "add <database> <branch> <keyspace> <region>",
+		Short: "Add a read-only region to a keyspace",
+		Long: "Add a read-only region to a Vitess keyspace.\n\n" +
+			"<region> is a PlanetScale region slug. List available slugs with: pscale region list.",
+		Args: cmdutil.RequiredArgs("database", "branch", "keyspace", "region"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			database, branch, keyspace, region := args[0], args[1], args[2], args[3]
+
+			client, current, err := getReadOnlyRegions(cmd, ch, database, branch, keyspace)
+			if err != nil {
+				return err
+			}
+			if readOnlyRegionIndex(current, region) >= 0 {
+				return fmt.Errorf("read-only region %s is already configured for keyspace %s", printer.BoldBlue(region), printer.BoldBlue(keyspace))
+			}
+
+			config := &ps.ReadOnlyRegionKeyspaceConfig{Region: region}
+			if cmd.Flags().Changed("cluster-size") {
+				config.ClusterSize = &clusterSize
+			}
+			if cmd.Flags().Changed("replicas") {
+				if replicas < 1 {
+					return fmt.Errorf("--replicas must be greater than 0")
+				}
+				config.Replicas = &replicas
+			}
+
+			end := ch.Printer.PrintProgress(fmt.Sprintf("Adding read-only region %s to keyspace %s", printer.BoldBlue(region), printer.BoldBlue(keyspace)))
+			defer end()
+
+			regions, err := updateReadOnlyRegions(cmd, ch, client, database, branch, keyspace, append(readOnlyRegionConfigs(current), config))
+			if err != nil {
+				return err
+			}
+			end()
+
+			return printReadOnlyRegionMutation(ch, fmt.Sprintf("Added read-only region %s to keyspace %s.", printer.BoldBlue(region), printer.BoldBlue(keyspace)), regions)
+		},
+	}
+
+	cmd.Flags().StringVar(&clusterSize, "cluster-size", "", "cluster size for the keyspace in this read-only region. Use `pscale size cluster list` to get a list of valid sizes.")
+	cmd.Flags().IntVar(&replicas, "replicas", 0, "number of replicas per shard in this read-only region")
+	cmd.RegisterFlagCompletionFunc("cluster-size", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
+		return cmdutil.BranchClusterSizesCompletionFunc(ch, cmd, args, toComplete)
+	})
+
+	return cmd
+}

internal/cmd/keyspace/read_only_regions_mutation_test.go

@@ -0,0 +1,199 @@
+package keyspace
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	qt "github.com/frankban/quicktest"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/mock"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/planetscale/cli/internal/printer"
+)
+
+func readOnlyRegionsTestHelper(format printer.Format, svc *mock.KeyspacesService, out *bytes.Buffer) *cmdutil.Helper {
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(out)
+	if format == printer.Human {
+		p.SetHumanOutput(out)
+	}
+	return &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "planetscale"},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{Keyspaces: svc}, nil
+		},
+	}
+}
+
+func getReadOnlyRegionsFn(regions []*ps.ReadOnlyRegionKeyspace) func(context.Context, *ps.GetKeyspaceRequest) (*ps.Keyspace, error) {
+	return func(_ context.Context, _ *ps.GetKeyspaceRequest) (*ps.Keyspace, error) {
+		return &ps.Keyspace{ReadOnlyRegions: regions}, nil
+	}
+}
+
+func TestKeyspace_ReadOnlyRegionsAddCmd(t *testing.T) {
+	c := qt.New(t)
+	var out bytes.Buffer
+	current := []*ps.ReadOnlyRegionKeyspace{{
+		Region:      "eu-west",
+		ClusterName: "PS_10",
+		Replicas:    1,
+	}}
+	result := append(current, &ps.ReadOnlyRegionKeyspace{
+		Region:             "us-west",
+		ClusterName:        "PS_20",
+		ClusterDisplayName: "PS-20",
+		Replicas:           2,
+	})
+
+	svc := &mock.KeyspacesService{
+		GetFn: getReadOnlyRegionsFn(current),
+		UpdateReadOnlyRegionsFn: func(_ context.Context, req *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+			c.Assert(req.Organization, qt.Equals, "planetscale")
+			c.Assert(req.Database, qt.Equals, "analytics")
+			c.Assert(req.Branch, qt.Equals, "main")
+			c.Assert(req.Keyspace, qt.Equals, "events")
+			c.Assert(req.ReadOnlyRegions, qt.HasLen, 2)
+			c.Assert(req.ReadOnlyRegions[0].Region, qt.Equals, "eu-west")
+			c.Assert(*req.ReadOnlyRegions[0].ClusterSize, qt.Equals, "PS_10")
+			c.Assert(*req.ReadOnlyRegions[0].Replicas, qt.Equals, 1)
+			c.Assert(req.ReadOnlyRegions[1], qt.DeepEquals, &ps.ReadOnlyRegionKeyspaceConfig{Region: "us-west"})
+			return result, nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.JSON, svc, &out))
+	cmd.SetArgs([]string{"add", "analytics", "main", "events", "us-west"})
+
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(out.String(), qt.JSONEquals, result)
+}
+
+func TestKeyspace_ReadOnlyRegionsAddCmdRejectsDuplicate(t *testing.T) {
+	c := qt.New(t)
+	current := []*ps.ReadOnlyRegionKeyspace{{Region: "us-west", ClusterName: "PS_10", Replicas: 1}}
+	svc := &mock.KeyspacesService{
+		GetFn: getReadOnlyRegionsFn(current),
+		UpdateReadOnlyRegionsFn: func(context.Context, *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+			c.Fatal("UpdateReadOnlyRegions should not be called")
+			return nil, nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.JSON, svc, &bytes.Buffer{}))
+	cmd.SetArgs([]string{"add", "analytics", "main", "events", "us-west"})
+
+	c.Assert(cmd.Execute(), qt.ErrorMatches, ".*already configured.*")
+}
+
+func TestKeyspace_ReadOnlyRegionsUpdateCmd(t *testing.T) {
+	c := qt.New(t)
+	var out bytes.Buffer
+	current := []*ps.ReadOnlyRegionKeyspace{
+		{Region: "us-west", ClusterName: "PS_10", Replicas: 2},
+		{Region: "eu-west", ClusterName: "PS_20", Replicas: 1},
+	}
+	result := []*ps.ReadOnlyRegionKeyspace{
+		{Region: "us-west", ClusterName: "PS_30", ClusterDisplayName: "PS-30", Replicas: 2},
+		current[1],
+	}
+	svc := &mock.KeyspacesService{
+		GetFn: getReadOnlyRegionsFn(current),
+		UpdateReadOnlyRegionsFn: func(_ context.Context, req *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+			c.Assert(req.ReadOnlyRegions, qt.HasLen, 2)
+			c.Assert(*req.ReadOnlyRegions[0].ClusterSize, qt.Equals, "PS_30")
+			c.Assert(*req.ReadOnlyRegions[0].Replicas, qt.Equals, 2)
+			c.Assert(*req.ReadOnlyRegions[1].ClusterSize, qt.Equals, "PS_20")
+			c.Assert(*req.ReadOnlyRegions[1].Replicas, qt.Equals, 1)
+			return result, nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.JSON, svc, &out))
+	cmd.SetArgs([]string{"update", "analytics", "main", "events", "us-west", "--cluster-size", "PS_30"})
+
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(out.String(), qt.JSONEquals, result)
+}
+
+func TestKeyspace_ReadOnlyRegionsUpdateCmdValidation(t *testing.T) {
+	tests := []struct {
+		name    string
+		args    []string
+		regions []*ps.ReadOnlyRegionKeyspace
+		want    string
+	}{
+		{
+			name:    "requires sizing flag",
+			args:    []string{"update", "analytics", "main", "events", "us-west"},
+			regions: nil,
+			want:    ".*at least one of --cluster-size or --replicas is required.*",
+		},
+		{
+			name:    "requires configured region",
+			args:    []string{"update", "analytics", "main", "events", "us-west", "--replicas", "2"},
+			regions: []*ps.ReadOnlyRegionKeyspace{{Region: "eu-west", ClusterName: "PS_10", Replicas: 1}},
+			want:    ".*is not configured.*",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := qt.New(t)
+			svc := &mock.KeyspacesService{
+				GetFn: getReadOnlyRegionsFn(tt.regions),
+				UpdateReadOnlyRegionsFn: func(context.Context, *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+					c.Fatal("UpdateReadOnlyRegions should not be called")
+					return nil, nil
+				},
+			}
+			cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.JSON, svc, &bytes.Buffer{}))
+			cmd.SetArgs(tt.args)
+			c.Assert(cmd.Execute(), qt.ErrorMatches, tt.want)
+		})
+	}
+}
+
+func TestKeyspace_ReadOnlyRegionsRemoveCmd(t *testing.T) {
+	c := qt.New(t)
+	var out bytes.Buffer
+	current := []*ps.ReadOnlyRegionKeyspace{
+		{Region: "us-west", ClusterName: "PS_10", Replicas: 2},
+		{Region: "eu-west", ClusterName: "PS_20", Replicas: 1},
+	}
+	svc := &mock.KeyspacesService{
+		GetFn: getReadOnlyRegionsFn(current),
+		UpdateReadOnlyRegionsFn: func(_ context.Context, req *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+			c.Assert(req.ReadOnlyRegions, qt.HasLen, 1)
+			c.Assert(req.ReadOnlyRegions[0].Region, qt.Equals, "eu-west")
+			c.Assert(*req.ReadOnlyRegions[0].ClusterSize, qt.Equals, "PS_20")
+			c.Assert(*req.ReadOnlyRegions[0].Replicas, qt.Equals, 1)
+			return current[1:], nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.Human, svc, &out))
+	cmd.SetArgs([]string{"remove", "analytics", "main", "events", "us-west"})
+
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(out.String(), qt.Contains, "Removed read-only region us-west from keyspace events.")
+}
+
+func TestKeyspace_ReadOnlyRegionsRemoveCmdRejectsMissingRegion(t *testing.T) {
+	c := qt.New(t)
+	svc := &mock.KeyspacesService{
+		GetFn: getReadOnlyRegionsFn(nil),
+		UpdateReadOnlyRegionsFn: func(context.Context, *ps.UpdateReadOnlyRegionsRequest) ([]*ps.ReadOnlyRegionKeyspace, error) {
+			c.Fatal("UpdateReadOnlyRegions should not be called")
+			return nil, nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(readOnlyRegionsTestHelper(printer.JSON, svc, &bytes.Buffer{}))
+	cmd.SetArgs([]string{"remove", "analytics", "main", "events", "us-west"})
+
+	c.Assert(cmd.Execute(), qt.ErrorMatches, ".*is not configured.*")
+}

internal/cmd/keyspace/read_only_regions_remove.go

@@ -0,0 +1,47 @@
+package keyspace
+
+import (
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/printer"
+	"github.com/spf13/cobra"
+)
+
+func ReadOnlyRegionsRemoveCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "remove <database> <branch> <keyspace> <region>",
+		Short: "Remove a read-only region from a keyspace",
+		Long: "Remove a read-only region from a Vitess keyspace.\n\n" +
+			"<region> is a PlanetScale region slug already configured on the keyspace. List configured regions with: pscale keyspace read-only-regions <database> <branch> <keyspace>.",
+		Args: cmdutil.RequiredArgs("database", "branch", "keyspace", "region"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			database, branch, keyspace, region := args[0], args[1], args[2], args[3]
+
+			client, current, err := getReadOnlyRegions(cmd, ch, database, branch, keyspace)
+			if err != nil {
+				return err
+			}
+			index := readOnlyRegionIndex(current, region)
+			if index < 0 {
+				return fmt.Errorf("read-only region %s is not configured for keyspace %s", printer.BoldBlue(region), printer.BoldBlue(keyspace))
+			}
+
+			configs := readOnlyRegionConfigs(current)
+			configs = append(configs[:index], configs[index+1:]...)
+
+			end := ch.Printer.PrintProgress(fmt.Sprintf("Removing read-only region %s from keyspace %s", printer.BoldBlue(region), printer.BoldBlue(keyspace)))
+			defer end()
+
+			regions, err := updateReadOnlyRegions(cmd, ch, client, database, branch, keyspace, configs)
+			if err != nil {
+				return err
+			}
+			end()
+
+			return printReadOnlyRegionMutation(ch, fmt.Sprintf("Removed read-only region %s from keyspace %s.", printer.BoldBlue(region), printer.BoldBlue(keyspace)), regions)
+		},
+	}
+
+	return cmd
+}