Add vtctld get-vschema command

internal/cmd/branch/vtctld/vschema.go

@@ -0,0 +1,55 @@
+package vtctld
+
+import (
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/spf13/cobra"
+)
+
+// GetVSchemaCmd reads the live VSchema for a keyspace from the cluster via vtctld.
+func GetVSchemaCmd(ch *cmdutil.Helper) *cobra.Command {
+	var keyspace string
+
+	cmd := &cobra.Command{
+		Use:   "get-vschema <database> <branch>",
+		Short: "Get the live VSchema for a keyspace",
+		Long: "Get the live VSchema for a keyspace from the cluster via vtctld. " +
+			"This reads the current cluster state, unlike `pscale keyspace vschema show`, " +
+			"which reads from the schema snapshot.",
+		Args: cmdutil.RequiredArgs("database", "branch"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			database, branch := args[0], args[1]
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			end := ch.Printer.PrintProgress(
+				fmt.Sprintf("Fetching VSchema for keyspace %s on %s\u2026",
+					keyspace, progressTarget(ch.Config.Organization, database, branch)))
+			defer end()
+
+			data, err := client.Vtctld.GetVSchema(ctx, &ps.VtctldGetVSchemaRequest{
+				Organization: ch.Config.Organization,
+				Database:     database,
+				Branch:       branch,
+				Keyspace:     keyspace,
+			})
+			if err != nil {
+				return cmdutil.HandleError(err)
+			}
+
+			end()
+			return ch.Printer.PrettyPrintJSON(data)
+		},
+	}
+
+	cmd.Flags().StringVar(&keyspace, "keyspace", "", "Keyspace name")
+	cmd.MarkFlagRequired("keyspace")
+
+	return cmd
+}

internal/cmd/branch/vtctld/vschema_test.go

@@ -0,0 +1,62 @@
+package vtctld
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"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 TestGetVSchema(t *testing.T) {
+	c := qt.New(t)
+
+	const (
+		org      = "my-org"
+		db       = "my-db"
+		branch   = "my-branch"
+		keyspace = "commerce"
+	)
+
+	svc := &mock.VtctldService{
+		GetVSchemaFn: func(ctx context.Context, req *ps.VtctldGetVSchemaRequest) (json.RawMessage, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Branch, qt.Equals, branch)
+			c.Assert(req.Keyspace, qt.Equals, keyspace)
+			return json.RawMessage(`{"multi_tenant_spec":{"tenant_id_column_name":"source_shard_id","tenant_id_column_type":"INT64"}}`), nil
+		},
+	}
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{Vtctld: svc}, nil
+		},
+	}
+
+	cmd := GetVSchemaCmd(ch)
+	cmd.SetArgs([]string{db, branch, "--keyspace", keyspace})
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.GetVSchemaFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, map[string]any{
+		"multi_tenant_spec": map[string]any{
+			"tenant_id_column_name": "source_shard_id",
+			"tenant_id_column_type": "INT64",
+		},
+	})
+}

internal/cmd/branch/vtctld/vtctld.go

@@ -20,6 +20,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(PlannedReparentShardCmd(ch))
 	cmd.AddCommand(ListWorkflowsCmd(ch))
 	cmd.AddCommand(ListKeyspacesCmd(ch))
+	cmd.AddCommand(GetVSchemaCmd(ch))
 	cmd.AddCommand(GetRoutingRulesCmd(ch))
 	cmd.AddCommand(GetShardCmd(ch))
 	cmd.AddCommand(SetShardTabletControlCmd(ch))

internal/mock/vtctld_general.go

@@ -14,6 +14,9 @@ type VtctldService struct {
 	ListKeyspacesFn        func(context.Context, *ps.VtctldListKeyspacesRequest) (json.RawMessage, error)
 	ListKeyspacesFnInvoked bool
 
+	GetVSchemaFn        func(context.Context, *ps.VtctldGetVSchemaRequest) (json.RawMessage, error)
+	GetVSchemaFnInvoked bool
+
 	GetRoutingRulesFn        func(context.Context, *ps.VtctldGetRoutingRulesRequest) (json.RawMessage, error)
 	GetRoutingRulesFnInvoked bool
 
@@ -58,6 +61,11 @@ func (s *VtctldService) ListKeyspaces(ctx context.Context, req *ps.VtctldListKey
 	return s.ListKeyspacesFn(ctx, req)
 }
 
+func (s *VtctldService) GetVSchema(ctx context.Context, req *ps.VtctldGetVSchemaRequest) (json.RawMessage, error) {
+	s.GetVSchemaFnInvoked = true
+	return s.GetVSchemaFn(ctx, req)
+}
+
 func (s *VtctldService) GetRoutingRules(ctx context.Context, req *ps.VtctldGetRoutingRulesRequest) (json.RawMessage, error) {
 	s.GetRoutingRulesFnInvoked = true
 	return s.GetRoutingRulesFn(ctx, req)

internal/planetscale/vtctld_general.go

@@ -15,6 +15,7 @@ import (
 type VtctldService interface {
 	ListWorkflows(context.Context, *VtctldListWorkflowsRequest) (json.RawMessage, error)
 	ListKeyspaces(context.Context, *VtctldListKeyspacesRequest) (json.RawMessage, error)
+	GetVSchema(context.Context, *VtctldGetVSchemaRequest) (json.RawMessage, error)
 	GetRoutingRules(context.Context, *VtctldGetRoutingRulesRequest) (json.RawMessage, error)
 	GetShard(context.Context, *VtctldGetShardRequest) (json.RawMessage, error)
 	SetShardTabletControl(context.Context, *VtctldSetShardTabletControlRequest) (json.RawMessage, error)
@@ -44,6 +45,15 @@ type VtctldListKeyspacesRequest struct {
 	Name         string `json:"-"`
 }
 
+// VtctldGetVSchemaRequest is a request for reading the live VSchema for a
+// keyspace from the cluster via vtctld.
+type VtctldGetVSchemaRequest struct {
+	Organization string `json:"-"`
+	Database     string `json:"-"`
+	Branch       string `json:"-"`
+	Keyspace     string `json:"-"`
+}
+
 // VtctldGetRoutingRulesRequest is a request for reading live routing rules
 // from the cluster via vtctld.
 type VtctldGetRoutingRulesRequest struct {
@@ -188,6 +198,10 @@ func vtctldKeyspacesAPIPath(org, db, branch string) string {
 	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "keyspaces")
 }
 
+func vtctldVSchemaAPIPath(org, db, branch string) string {
+	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "vschema")
+}
+
 func vtctldRoutingRulesAPIPath(org, db, branch string) string {
 	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "routing-rules")
 }
@@ -242,6 +256,22 @@ func (s *vtctldService) ListKeyspaces(ctx context.Context, req *VtctldListKeyspa
 	return resp.Data, nil
 }
 
+// GetVSchema reads the live VSchema for a keyspace from the cluster via vtctld.
+func (s *vtctldService) GetVSchema(ctx context.Context, req *VtctldGetVSchemaRequest) (json.RawMessage, error) {
+	p := vtctldVSchemaAPIPath(req.Organization, req.Database, req.Branch)
+	v := url.Values{}
+	v.Set("keyspace", req.Keyspace)
+	httpReq, err := s.client.newRequest(http.MethodGet, p, nil, WithQueryParams(v))
+	if err != nil {
+		return nil, fmt.Errorf("error creating http request: %w", err)
+	}
+	resp := &vtctldDataResponse{}
+	if err := s.client.do(ctx, httpReq, resp); err != nil {
+		return nil, err
+	}
+	return resp.Data, nil
+}
+
 func (s *vtctldService) GetRoutingRules(ctx context.Context, req *VtctldGetRoutingRulesRequest) (json.RawMessage, error) {
 	p := vtctldRoutingRulesAPIPath(req.Organization, req.Database, req.Branch)
 	httpReq, err := s.client.newRequest(http.MethodGet, p, nil)

internal/planetscale/vtctld_general_test.go

@@ -283,6 +283,34 @@ func TestVtctld_ListKeyspaces(t *testing.T) {
 	c.Assert(string(data), qt.Equals, `{"result":"ok"}`)
 }
 
+func TestVtctld_GetVSchema(t *testing.T) {
+	c := qt.New(t)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, qt.Equals, http.MethodGet)
+		c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/vtctld/vschema")
+		c.Assert(r.URL.Query().Get("keyspace"), qt.Equals, "commerce")
+
+		w.WriteHeader(200)
+		_, err := w.Write([]byte(`{"data":{"multi_tenant_spec":{"tenant_id_column_name":"source_shard_id","tenant_id_column_type":"INT64"}}}`))
+		c.Assert(err, qt.IsNil)
+	}))
+	defer ts.Close()
+
+	client, err := NewClient(WithBaseURL(ts.URL))
+	c.Assert(err, qt.IsNil)
+
+	ctx := context.Background()
+	data, err := client.Vtctld.GetVSchema(ctx, &VtctldGetVSchemaRequest{
+		Organization: "my-org",
+		Database:     "my-db",
+		Branch:       "my-branch",
+		Keyspace:     "commerce",
+	})
+	c.Assert(err, qt.IsNil)
+	c.Assert(string(data), qt.Equals, `{"multi_tenant_spec":{"tenant_id_column_name":"source_shard_id","tenant_id_column_type":"INT64"}}`)
+}
+
 func TestVtctld_StartWorkflow(t *testing.T) {
 	c := qt.New(t)