Merge pull request #957 from planetscale/add-rollouts

Add rollout status

go.mod

@@ -23,7 +23,7 @@ require (
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
 	github.com/pkg/errors v0.9.1
-	github.com/planetscale/planetscale-go v0.116.0
+	github.com/planetscale/planetscale-go v0.117.0
 	github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7
 	github.com/planetscale/psdbproxy v0.0.0-20250117221522-0c8e2b0e36e6
 	github.com/spf13/cobra v1.8.1

go.sum

@@ -116,6 +116,8 @@ github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2v
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
 github.com/planetscale/planetscale-go v0.116.0 h1:jXR32sTfQOWzoVBuwHLjqh0XktH6ATXyGbvt0u95ALc=
 github.com/planetscale/planetscale-go v0.116.0/go.mod h1:ldGffCLckkR8fjGDjDFs4WcjlDr8uqg2qRUZhRYBEMI=
+github.com/planetscale/planetscale-go v0.117.0 h1:MfmmDwEqkEcO2y31zSZbQONV1WpVFZAWV+EyI4Ky/wY=
+github.com/planetscale/planetscale-go v0.117.0/go.mod h1:ldGffCLckkR8fjGDjDFs4WcjlDr8uqg2qRUZhRYBEMI=
 github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7 h1:dxdoFKWVDlV1gq8UQC8NWCofLjCEjEHw47gfeojgs28=
 github.com/planetscale/psdb v0.0.0-20240109164348-6848e728f6e7/go.mod h1:WZmi4gw3rOK+ryd1inGxgfKwoFV04O7xBCqzWzv0/0U=
 github.com/planetscale/psdbproxy v0.0.0-20250117221522-0c8e2b0e36e6 h1:/Ox1ZTAdk+soSngzzKoJh5voOzptrpPrux11o30rIaw=

internal/cmd/keyspace/keyspace.go

@@ -25,6 +25,7 @@ func KeyspaceCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(VSchemaCmd(ch))
 	cmd.AddCommand(CreateCmd(ch))
 	cmd.AddCommand(ResizeCmd(ch))
+	cmd.AddCommand(RolloutStatusCmd(ch))
 
 	return cmd
 }

internal/cmd/keyspace/rollout_status.go

@@ -0,0 +1,97 @@
+package keyspace
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+	"github.com/spf13/cobra"
+)
+
+func RolloutStatusCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "rollout-status <database> <branch> <keyspace>",
+		Short: "Show keyspace rollout status per shard",
+		Args:  cmdutil.RequiredArgs("database", "branch", "keyspace"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			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 rollout status for keyspace %s in %s/%s", printer.BoldBlue(keyspace), printer.BoldBlue(database), printer.BoldBlue(branch)))
+			defer end()
+
+			k, err := client.Keyspaces.RolloutStatus(ctx, &ps.KeyspaceRolloutStatusRequest{
+				Organization: ch.Config.Organization,
+				Database:     database,
+				Branch:       branch,
+				Keyspace:     keyspace,
+			})
+			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()
+
+			return ch.Printer.PrintResource(toShardRollouts(k))
+		},
+	}
+
+	return cmd
+}
+
+func toShardRollouts(keyspaceRollout *ps.KeyspaceRollout) []*ShardRollout {
+	shards := make([]*ShardRollout, 0, len(keyspaceRollout.Shards))
+
+	for _, s := range keyspaceRollout.Shards {
+		shards = append(shards, toShardRollout(s))
+	}
+
+	return shards
+}
+
+type ShardRollout struct {
+	Name  string `header:"shard" json:"name"`
+	State string `header:"state" json:"state"`
+
+	LastRolloutStartedAt  *int64 `header:"started,timestamp(ms|utc|human)" json:"last_rollout_started_at"`
+	LastRolloutFinishedAt *int64 `header:"finished,timestamp(ms|utc|human)" json:"last_rollout_finished_at"`
+
+	orig *ps.ShardRollout
+}
+
+func toShardRollout(sr ps.ShardRollout) *ShardRollout {
+	var startedAt, finishedAt *time.Time
+	if !sr.LastRolloutStartedAt.IsZero() {
+		startedAt = &sr.LastRolloutStartedAt
+	}
+	if !sr.LastRolloutFinishedAt.IsZero() {
+		finishedAt = &sr.LastRolloutFinishedAt
+	}
+	return &ShardRollout{
+		Name:                  sr.Name,
+		State:                 cmdutil.SnakeToSentenceCase(sr.State),
+		LastRolloutStartedAt:  printer.GetMillisecondsIfExists(startedAt),
+		LastRolloutFinishedAt: printer.GetMillisecondsIfExists(finishedAt),
+		orig:                  &sr,
+	}
+}
+
+func (s *ShardRollout) MarshalJSON() ([]byte, error) {
+	return json.MarshalIndent(s.orig, "", "  ")
+}
+
+func (s *ShardRollout) MarshalCSVValue() interface{} {
+	return []*ShardRollout{s}
+}

internal/cmd/keyspace/rollout_status_test.go

@@ -0,0 +1,98 @@
+package keyspace
+
+import (
+	"bytes"
+	"context"
+	"testing"
+	"time"
+
+	qt "github.com/frankban/quicktest"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/mock"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+)
+
+func TestKeyspace_RolloutStatusCmd(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"
+	keyspace := "sharded"
+
+	ts := time.Now()
+
+	kr := &ps.KeyspaceRollout{
+		Name:  "sharded",
+		State: "resizing",
+		Shards: []ps.ShardRollout{
+			{
+				Name:                  "-80",
+				State:                 "resizing",
+				LastRolloutFinishedAt: time.Time{},
+				LastRolloutStartedAt:  ts,
+			},
+			{
+				Name:                  "80-",
+				State:                 "resizing",
+				LastRolloutFinishedAt: time.Time{},
+				LastRolloutStartedAt:  ts,
+			},
+		},
+	}
+
+	expected := []ShardRollout{
+		{
+			Name:                  "-80",
+			State:                 cmdutil.SnakeToSentenceCase("resizing"),
+			LastRolloutFinishedAt: nil,
+			LastRolloutStartedAt:  printer.GetMillisecondsIfExists(&ts),
+			orig:                  &kr.Shards[0],
+		},
+		{
+			Name:                  "80-",
+			State:                 cmdutil.SnakeToSentenceCase("resizing"),
+			LastRolloutFinishedAt: nil,
+			LastRolloutStartedAt:  printer.GetMillisecondsIfExists(&ts),
+			orig:                  &kr.Shards[1],
+		},
+	}
+
+	svc := &mock.KeyspacesService{
+		RolloutStatusFn: func(ctx context.Context, req *ps.KeyspaceRolloutStatusRequest) (*ps.KeyspaceRollout, error) {
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Branch, qt.Equals, branch)
+			c.Assert(req.Keyspace, qt.Equals, keyspace)
+
+			return kr, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Keyspaces: svc,
+			}, nil
+		},
+	}
+
+	cmd := RolloutStatusCmd(ch)
+	cmd.SetArgs([]string{db, branch, keyspace})
+	err := cmd.Execute()
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.RolloutStatusFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, expected)
+}

internal/cmdutil/cmdutil.go

@@ -20,6 +20,8 @@ import (
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	exec "golang.org/x/sys/execabs"
+	"golang.org/x/text/cases"
+	"golang.org/x/text/language"
 
 	"vitess.io/vitess/go/mysql"
 )
@@ -297,3 +299,14 @@ func DeployRequestBranchToNumber(ctx context.Context, client *ps.Client, organiz
 func ToClusterSizeSlug(c ps.ClusterSize) string {
 	return strings.ReplaceAll(string(c), "_", "-")
 }
+
+func SnakeToSentenceCase(s string) string {
+	titleCaser := cases.Title(language.English)
+
+	words := strings.Split(s, "_")
+	if len(words) > 0 {
+		words[0] = titleCaser.String(words[0])
+	}
+
+	return strings.Join(words, " ")
+}

internal/mock/keyspace.go

@@ -30,6 +30,9 @@ type KeyspacesService struct {
 
 	ResizeStatusFn        func(context.Context, *ps.KeyspaceResizeStatusRequest) (*ps.KeyspaceResizeRequest, error)
 	ResizeStatusFnInvoked bool
+
+	RolloutStatusFn        func(context.Context, *ps.KeyspaceRolloutStatusRequest) (*ps.KeyspaceRollout, error)
+	RolloutStatusFnInvoked bool
 }
 
 func (s *KeyspacesService) List(ctx context.Context, req *ps.ListKeyspacesRequest) ([]*ps.Keyspace, error) {
@@ -72,3 +75,8 @@ func (s *KeyspacesService) ResizeStatus(ctx context.Context, req *ps.KeyspaceRes
 	s.ResizeStatusFnInvoked = true
 	return s.ResizeStatusFn(ctx, req)
 }
+
+func (s *KeyspacesService) RolloutStatus(ctx context.Context, req *ps.KeyspaceRolloutStatusRequest) (*ps.KeyspaceRollout, error) {
+	s.RolloutStatusFnInvoked = true
+	return s.RolloutStatusFn(ctx, req)
+}