Read-only region support

internal/cmd/database/dump.go

@@ -26,19 +26,20 @@ import (
 )
 
 type dumpFlags struct {
-	localAddr    string
-	remoteAddr   string
-	keyspace     string
-	shard        string
-	replica      bool
-	rdonly       bool
-	tables       string
-	wheres       string
-	columns      []string
-	output       string
-	threads      int
-	schemaOnly   bool
-	outputFormat string
+	localAddr      string
+	remoteAddr     string
+	keyspace       string
+	shard          string
+	replica        bool
+	rdonly         bool
+	readOnlyRegion string
+	tables         string
+	wheres         string
+	columns        []string
+	output         string
+	threads        int
+	schemaOnly     bool
+	outputFormat   string
 }
 
 // DumpCmd encapsulates the commands for dumping a database
@@ -59,8 +60,10 @@ func DumpCmd(ch *cmdutil.Helper) *cobra.Command {
 		"", "Local address to bind and listen for connections. By default the proxy binds to 127.0.0.1 with a random port.")
 	cmd.PersistentFlags().StringVar(&f.remoteAddr, "remote-addr", "",
 		"PlanetScale Database remote network address. By default the remote address is populated automatically from the PlanetScale API. (format: `hostname:port`)")
-	cmd.PersistentFlags().BoolVar(&f.replica, "replica", false, "Dump from a replica (if available; will fail if not).")
-	cmd.PersistentFlags().BoolVar(&f.rdonly, "rdonly", false, "Dump from a rdonly tablet (if available; will fail if not).")
+	cmd.PersistentFlags().BoolVar(&f.replica, "replica", false, "Dump from a replica tablet in the primary region (if available; will fail if not).")
+	cmd.PersistentFlags().BoolVar(&f.rdonly, "rdonly", false, "Dump from a rdonly tablet in the primary region (if available; will fail if not). Not for separate read-only regions — use --read-only-region instead.")
+	cmd.PersistentFlags().StringVar(&f.readOnlyRegion, "read-only-region", "",
+		"Dump from a Vitess read-only region (region slug, display name, or id). List regions with: pscale keyspace read-only-regions <database> <branch> <keyspace>.")
 	cmd.PersistentFlags().StringVar(&f.tables, "tables", "",
 		"Comma separated string of tables to dump. By default all tables are dumped.")
 	cmd.PersistentFlags().StringVar(&f.wheres, "wheres", "",
@@ -93,6 +96,10 @@ func dump(ch *cmdutil.Helper, cmd *cobra.Command, flags *dumpFlags, args []strin
 		return fmt.Errorf("to target a single shard, please pass the --keyspace flag")
 	}
 
+	if flags.readOnlyRegion != "" && (flags.rdonly || flags.replica) {
+		return fmt.Errorf("--read-only-region cannot be combined with --rdonly or --replica")
+	}
+
 	validFormats := map[string]bool{"sql": true, "json": true, "csv": true}
 	if !validFormats[flags.outputFormat] {
 		return fmt.Errorf("invalid output format: %s. Valid options are: sql, json, csv", flags.outputFormat)
@@ -144,13 +151,37 @@ func dump(ch *cmdutil.Helper, cmd *cobra.Command, flags *dumpFlags, args []strin
 		return errors.New("database branch is not ready yet, please try again in a few minutes")
 	}
 
+	role := cmdutil.AdministratorRole
+	var readOnlyRegionID string
+	if flags.readOnlyRegion != "" {
+		regions, err := client.ReadOnlyRegions.List(ctx, &ps.ListReadOnlyRegionsRequest{
+			Organization: ch.Config.Organization,
+			Database:     database,
+		})
+		if err != nil {
+			return cmdutil.HandleError(err)
+		}
+
+		ror, err := ps.FindReadOnlyRegion(regions, flags.readOnlyRegion)
+		if err != nil {
+			return err
+		}
+		if !ror.Ready {
+			return fmt.Errorf("read-only region %s is not ready yet", printer.BoldBlue(flags.readOnlyRegion))
+		}
+
+		readOnlyRegionID = ror.ID
+		role = cmdutil.ReaderRole
+	}
+
 	pw, err := passwordutil.New(ctx, client, passwordutil.Options{
-		Organization: ch.Config.Organization,
-		Database:     database,
-		Branch:       branch,
-		Role:         cmdutil.AdministratorRole,
-		Name:         passwordutil.GenerateName("pscale-cli-dump"),
-		TTL:          5 * time.Minute,
+		Organization:     ch.Config.Organization,
+		Database:         database,
+		Branch:           branch,
+		Role:             role,
+		Name:             passwordutil.GenerateName("pscale-cli-dump"),
+		TTL:              5 * time.Minute,
+		ReadOnlyRegionID: readOnlyRegionID,
 	})
 	if err != nil {
 		return cmdutil.HandleError(err)
@@ -253,16 +284,8 @@ func dump(ch *cmdutil.Helper, cmd *cobra.Command, flags *dumpFlags, args []strin
 	cfg.OutputFormat = flags.outputFormat
 
 	if flags.shard != "" {
-		if flags.replica {
-			useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
-			cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
-		} else if flags.rdonly {
-			useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
-			cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
-		} else {
-			useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
-			cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
-		}
+		useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
+		cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
 	}
 
 	if flags.replica && flags.shard == "" {

internal/cmd/database/dump_test.go

@@ -4,6 +4,10 @@ import (
 	"testing"
 
 	qt "github.com/frankban/quicktest"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/planetscale/cli/internal/printer"
 )
 
 func TestParseColumnIncludes(t *testing.T) {
@@ -101,6 +105,34 @@ func TestParseColumnIncludes(t *testing.T) {
 	}
 }
 
+func TestDump_ReadOnlyRegionFlagConflicts(t *testing.T) {
+	c := qt.New(t)
+
+	format := printer.Human
+	p := printer.NewPrinter(&format)
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: "planetscale",
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{}, nil
+		},
+	}
+
+	cmd := DumpCmd(ch)
+	cmd.SetArgs([]string{"db", "main", "--read-only-region", "eu-west", "--rdonly"})
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "cannot be combined")
+
+	cmd = DumpCmd(ch)
+	cmd.SetArgs([]string{"db", "main", "--read-only-region", "eu-west", "--replica"})
+	err = cmd.Execute()
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "cannot be combined")
+}
+
 func TestShardUseCommand(t *testing.T) {
 	c := qt.New(t)
 

internal/cmd/keyspace/keyspace.go

@@ -27,6 +27,7 @@ func KeyspaceCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(CreateCmd(ch))
 	cmd.AddCommand(ResizeCmd(ch))
 	cmd.AddCommand(RolloutStatusCmd(ch))
+	cmd.AddCommand(ReadOnlyRegionsCmd(ch))
 	cmd.AddCommand(UpdateSettingsCmd(ch))
 	cmd.AddCommand(SettingsCmd(ch))
 

internal/cmd/keyspace/read_only_regions.go

@@ -0,0 +1,87 @@
+package keyspace
+
+import (
+	"encoding/json"
+	"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 ReadOnlyRegionsCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "read-only-regions <database> <branch> <keyspace>",
+		Short: "List read-only regions for a keyspace",
+		Long: "List read-only regions configured for a Vitess keyspace.\n\n" +
+			"This command is only supported for Vitess databases.",
+		Args: cmdutil.RequiredArgs("database", "branch", "keyspace"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			database, branch, keyspace := args[0], args[1], args[2]
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching read-only regions for keyspace %s in %s/%s", printer.BoldBlue(keyspace), printer.BoldBlue(database), printer.BoldBlue(branch)))
+			defer end()
+
+			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 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 cmdutil.HandleError(err)
+				}
+			}
+			end()
+
+			if len(k.ReadOnlyRegions) == 0 && ch.Printer.Format() == printer.Human {
+				ch.Printer.Println("No read-only regions have been configured for this keyspace.")
+				return nil
+			}
+
+			return ch.Printer.PrintResource(toReadOnlyRegions(k.ReadOnlyRegions))
+		},
+	}
+
+	return cmd
+}
+
+type ReadOnlyRegion struct {
+	Region      string `header:"region" json:"region"`
+	ClusterSize string `header:"cluster_size" json:"cluster_name"`
+	Replicas    int    `header:"replicas" json:"replicas"`
+
+	orig *ps.ReadOnlyRegionKeyspace
+}
+
+func toReadOnlyRegions(regions []*ps.ReadOnlyRegionKeyspace) []*ReadOnlyRegion {
+	out := make([]*ReadOnlyRegion, 0, len(regions))
+	for _, region := range regions {
+		out = append(out, &ReadOnlyRegion{
+			Region:      region.Region,
+			ClusterSize: region.ClusterDisplayName,
+			Replicas:    region.Replicas,
+			orig:        region,
+		})
+	}
+	return out
+}
+
+func (r *ReadOnlyRegion) MarshalJSON() ([]byte, error) {
+	return json.MarshalIndent(r.orig, "", "  ")
+}
+
+func (r *ReadOnlyRegion) MarshalCSVValue() interface{} {
+	return []*ReadOnlyRegion{r}
+}

internal/cmd/keyspace/read_only_regions_test.go

@@ -0,0 +1,55 @@
+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 TestKeyspace_ReadOnlyRegionsCmd(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	regions := []*ps.ReadOnlyRegionKeyspace{{
+		Region:             "us-west",
+		ClusterName:        "PS_20",
+		ClusterDisplayName: "PS-20",
+		Replicas:           2,
+	}}
+	svc := &mock.KeyspacesService{
+		GetFn: func(ctx context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, 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.Full, qt.IsTrue)
+			return &ps.Keyspace{ReadOnlyRegions: regions}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "planetscale"},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{Keyspaces: svc}, nil
+		},
+	}
+
+	cmd := ReadOnlyRegionsCmd(ch)
+	cmd.SetArgs([]string{"analytics", "main", "events"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(buf.String(), qt.JSONEquals, regions)
+}