Wait for the schema snapshot after refresh-schema

internal/cmd/branch/refresh_schema.go

@@ -1,18 +1,28 @@
 package branch
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"time"
 
 	"github.com/planetscale/cli/internal/cmdutil"
 	"github.com/planetscale/cli/internal/planetscale"
 	"github.com/planetscale/cli/internal/printer"
 	"github.com/spf13/cobra"
 )
 
+// schemaRefreshPollInterval and schemaRefreshTimeout are vars so tests can shorten the wait.
+var (
+	schemaRefreshPollInterval = 2 * time.Second
+	schemaRefreshTimeout      = 2 * time.Minute
+)
+
 func RefreshSchemaCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "refresh-schema <database> <branch>",
 		Short: "Refresh the schema for a MySQL branch",
+		Long:  "Refresh the schema for a MySQL branch. Waits until the updated schema snapshot is ready so later commands, such as routing-rules, see it.",
 		Args:  cmdutil.RequiredArgs("database", "branch"),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
@@ -32,13 +42,11 @@ func RefreshSchemaCmd(ch *cmdutil.Helper) *cobra.Command {
 				Branch:       branch,
 			})
 			if err != nil {
-				switch cmdutil.ErrCode(err) {
-				case planetscale.ErrNotFound:
-					return fmt.Errorf("branch %s does not exist in database %s (organization: %s)",
-						printer.BoldBlue(branch), printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization))
-				default:
-					return cmdutil.HandleError(err)
-				}
+				return branchNotFoundOr(err, branch, database, ch.Config.Organization)
+			}
+
+			if err := waitUntilSchemaReady(ctx, client, ch.Config.Organization, database, branch); err != nil {
+				return branchNotFoundOr(err, branch, database, ch.Config.Organization)
 			}
 			end()
 
@@ -57,3 +65,53 @@ func RefreshSchemaCmd(ch *cmdutil.Helper) *cobra.Command {
 
 	return cmd
 }
+
+func branchNotFoundOr(err error, branch, database, org string) error {
+	if cmdutil.ErrCode(err) == planetscale.ErrNotFound {
+		return fmt.Errorf("branch %s does not exist in database %s (organization: %s)",
+			printer.BoldBlue(branch), printer.BoldBlue(database), printer.BoldBlue(org))
+	}
+	return cmdutil.HandleError(err)
+}
+
+// waitUntilSchemaReady polls until the branch schema snapshot is ready.
+// refresh-schema points the branch at a new snapshot before that snapshot
+// is ready, and commands such as routing-rules fail during that window.
+func waitUntilSchemaReady(ctx context.Context, client *planetscale.Client, org, database, branch string) error {
+	ctx, cancel := context.WithTimeout(ctx, schemaRefreshTimeout)
+	defer cancel()
+
+	ticker := time.NewTicker(schemaRefreshPollInterval)
+	defer ticker.Stop()
+
+	var pollErr error
+	for {
+		resp, err := client.DatabaseBranches.Get(ctx, &planetscale.GetDatabaseBranchRequest{
+			Organization: org,
+			Database:     database,
+			Branch:       branch,
+		})
+		switch {
+		case err != nil && cmdutil.ErrCode(err) == planetscale.ErrNotFound:
+			return err
+		case err != nil && ctx.Err() == nil:
+			pollErr = err
+		case err == nil && resp.SchemaReady:
+			return nil
+		case err == nil:
+			pollErr = nil
+		}
+
+		select {
+		case <-ctx.Done():
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) && pollErr != nil {
+				return pollErr
+			}
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				return errors.New("schema refreshing; please try again in a few moments")
+			}
+			return ctx.Err()
+		case <-ticker.C:
+		}
+	}
+}

internal/cmd/branch/refresh_schema_test.go

@@ -3,7 +3,9 @@ package branch
 import (
 	"bytes"
 	"context"
+	"errors"
 	"testing"
+	"time"
 
 	qt "github.com/frankban/quicktest"
 	"github.com/planetscale/cli/internal/cmdutil"
@@ -32,6 +34,12 @@ func TestSnapshot_CreateCmd(t *testing.T) {
 			c.Assert(req.Branch, qt.Equals, branch)
 			return nil
 		},
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Branch, qt.Equals, branch)
+			return &ps.DatabaseBranch{Name: branch, SchemaReady: true}, nil
+		},
 	}
 
 	ch := &cmdutil.Helper{
@@ -52,4 +60,106 @@ func TestSnapshot_CreateCmd(t *testing.T) {
 
 	c.Assert(err, qt.IsNil)
 	c.Assert(svc.RefreshSchemaFnInvoked, qt.IsTrue)
+	c.Assert(svc.GetFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.Contains, `"result": "schema refreshed"`)
+}
+
+func TestRefreshSchemaWaitsUntilSchemaReady(t *testing.T) {
+	c := qt.New(t)
+	restoreSchemaRefreshTiming(t)
+
+	schemaRefreshPollInterval = time.Millisecond
+	schemaRefreshTimeout = time.Second
+
+	var buf bytes.Buffer
+	var polls int
+	svc := &mock.DatabaseBranchesService{
+		RefreshSchemaFn: func(context.Context, *ps.RefreshSchemaRequest) error {
+			return nil
+		},
+		GetFn: func(context.Context, *ps.GetDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
+			polls++
+			return &ps.DatabaseBranch{SchemaReady: polls > 1}, nil
+		},
+	}
+
+	cmd := RefreshSchemaCmd(refreshSchemaHelper(svc, &buf))
+	cmd.SetArgs([]string{"planetscale", "development"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(polls, qt.Equals, 2)
+	c.Assert(buf.String(), qt.Contains, `"result": "schema refreshed"`)
+}
+
+func TestRefreshSchemaReturnsPollError(t *testing.T) {
+	c := qt.New(t)
+	restoreSchemaRefreshTiming(t)
+
+	schemaRefreshPollInterval = time.Millisecond
+	schemaRefreshTimeout = 50 * time.Millisecond
+
+	var buf bytes.Buffer
+	svc := &mock.DatabaseBranchesService{
+		RefreshSchemaFn: func(context.Context, *ps.RefreshSchemaRequest) error {
+			return nil
+		},
+		GetFn: func(context.Context, *ps.GetDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
+			return nil, errors.New("connection refused")
+		},
+	}
+
+	cmd := RefreshSchemaCmd(refreshSchemaHelper(svc, &buf))
+	cmd.SetArgs([]string{"planetscale", "development"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.ErrorMatches, "connection refused")
+	c.Assert(buf.String(), qt.Equals, "")
+}
+
+func TestRefreshSchemaTimesOutWhileSchemaIsRefreshing(t *testing.T) {
+	c := qt.New(t)
+	restoreSchemaRefreshTiming(t)
+
+	schemaRefreshPollInterval = time.Millisecond
+	schemaRefreshTimeout = 50 * time.Millisecond
+
+	var buf bytes.Buffer
+	svc := &mock.DatabaseBranchesService{
+		RefreshSchemaFn: func(context.Context, *ps.RefreshSchemaRequest) error {
+			return nil
+		},
+		GetFn: func(context.Context, *ps.GetDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
+			return &ps.DatabaseBranch{SchemaReady: false}, nil
+		},
+	}
+
+	cmd := RefreshSchemaCmd(refreshSchemaHelper(svc, &buf))
+	cmd.SetArgs([]string{"planetscale", "development"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.ErrorMatches, "schema refreshing; please try again in a few moments")
+	c.Assert(buf.String(), qt.Equals, "")
+}
+
+func refreshSchemaHelper(svc *mock.DatabaseBranchesService, buf *bytes.Buffer) *cmdutil.Helper {
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(buf)
+	return &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "planetscale"},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{DatabaseBranches: svc}, nil
+		},
+	}
+}
+
+func restoreSchemaRefreshTiming(t *testing.T) {
+	t.Helper()
+	interval, timeout := schemaRefreshPollInterval, schemaRefreshTimeout
+	t.Cleanup(func() {
+		schemaRefreshPollInterval = interval
+		schemaRefreshTimeout = timeout
+	})
 }

internal/cmdutil/json_error.go

@@ -177,6 +177,15 @@ func GlobalJSONError(err error) JSONErrorResponse {
 			AgentAuthCheckCmd(),
 		}
 
+	// refresh-schema timed out before the new snapshot was ready. The branch
+	// exists; retry the follow-up command once the snapshot finishes.
+	case strings.Contains(lower, "schema refreshing"):
+		status = "action_required"
+		code = "schema_snapshot_not_ready"
+		nextSteps = []string{
+			"Wait for the branch schema snapshot to become ready, then retry",
+		}
+
 	// Unclassified: auth state is the one thing worth ruling out first.
 	default:
 		nextSteps = []string{

internal/cmdutil/json_error_test.go

@@ -174,6 +174,20 @@ func TestGlobalJSONErrorSchemaMutationBlocked(t *testing.T) {
 	}
 }
 
+func TestGlobalJSONErrorSchemaRefreshing(t *testing.T) {
+	resp := GlobalJSONError(errors.New("schema refreshing; please try again in a few moments"))
+
+	if resp.Status != "action_required" {
+		t.Fatalf("status = %q", resp.Status)
+	}
+	if resp.Code() != "schema_snapshot_not_ready" {
+		t.Fatalf("code = %q", resp.Code())
+	}
+	if len(resp.NextSteps) != 1 || resp.NextSteps[0] != "Wait for the branch schema snapshot to become ready, then retry" {
+		t.Fatalf("next_steps = %#v", resp.NextSteps)
+	}
+}
+
 func TestGlobalJSONErrorSchemaSnapshotNotReady(t *testing.T) {
 	resp := GlobalJSONError(&planetscale.Error{
 		APICode: "schema_snapshot_not_ready",

internal/planetscale/branches.go

@@ -24,6 +24,7 @@ type DatabaseBranch struct {
 	Actor          Actor     `json:"actor"`
 	Region         Region    `json:"region"`
 	Ready          bool      `json:"ready"`
+	SchemaReady    bool      `json:"schema_ready"`
 	Production     bool      `json:"production"`
 	HtmlURL        string    `json:"html_url"`
 	CreatedAt      time.Time `json:"created_at"`