diff --git a/AGENTS.md b/AGENTS.md index a8d7b684..01f988ef 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -492,6 +492,8 @@ pscale branch vtctld move-tables create --org --format --workflow --source-keyspace --target-keyspace --tables --auto-start=false pscale branch vtctld move-tables start --org --format json \ --workflow --target-keyspace +pscale branch vtctld move-tables stop --org --format json \ + --workflow --target-keyspace pscale branch vtctld move-tables status --org --format json \ --workflow --target-keyspace pscale branch vtctld move-tables switch-traffic --org --format json \ diff --git a/internal/cmd/branch/vtctld/move_tables.go b/internal/cmd/branch/vtctld/move_tables.go index ec56b9dc..79e7d2b0 100644 --- a/internal/cmd/branch/vtctld/move_tables.go +++ b/internal/cmd/branch/vtctld/move_tables.go @@ -30,6 +30,7 @@ func MoveTablesCmd(ch *cmdutil.Helper) *cobra.Command { cmd.AddCommand(MoveTablesShowCmd(ch)) cmd.AddCommand(MoveTablesStatusCmd(ch)) cmd.AddCommand(MoveTablesStartCmd(ch)) + cmd.AddCommand(MoveTablesStopCmd(ch)) cmd.AddCommand(MoveTablesSwitchTrafficCmd(ch)) cmd.AddCommand(MoveTablesReverseTrafficCmd(ch)) cmd.AddCommand(MoveTablesCancelCmd(ch)) @@ -351,6 +352,56 @@ func MoveTablesStartCmd(ch *cmdutil.Helper) *cobra.Command { return cmd } +func MoveTablesStopCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + workflow string + targetKeyspace string + } + + cmd := &cobra.Command{ + Use: "stop ", + Short: "Stop a MoveTables workflow", + 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("Stopping MoveTables workflow %s on %s\u2026", + printer.BoldBlue(flags.workflow), progressTarget(ch.Config.Organization, database, branch))) + defer end() + + data, err := client.MoveTables.Stop(ctx, &ps.MoveTablesStopRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Workflow: flags.workflow, + TargetKeyspace: flags.targetKeyspace, + }) + if err != nil { + return cmdutil.HandleError(err) + } + + end() + return printWorkflowJSON(ch.Printer, data, []workflowNextStep{ + moveTablesStatusStep(ch.Config.Organization, database, branch, flags.workflow, flags.targetKeyspace, "Confirm streams are stopped"), + }) + }, + } + + cmd.Flags().StringVar(&flags.workflow, "workflow", "", "Name of the workflow") + cmd.Flags().StringVar(&flags.targetKeyspace, "target-keyspace", "", "Target keyspace") + cmd.MarkFlagRequired("workflow") // nolint:errcheck + cmd.MarkFlagRequired("target-keyspace") // nolint:errcheck + + return cmd +} + func MoveTablesSwitchTrafficCmd(ch *cmdutil.Helper) *cobra.Command { var flags struct { workflow string diff --git a/internal/cmd/branch/vtctld/move_tables_test.go b/internal/cmd/branch/vtctld/move_tables_test.go index 78189d57..65d7bedd 100644 --- a/internal/cmd/branch/vtctld/move_tables_test.go +++ b/internal/cmd/branch/vtctld/move_tables_test.go @@ -758,6 +758,44 @@ func TestMoveTablesStart(t *testing.T) { }) } +func TestMoveTablesStop(t *testing.T) { + c := qt.New(t) + + org := "my-org" + db := "my-db" + branch := "my-branch" + + svc := &mock.MoveTablesService{ + StopFn: func(ctx context.Context, req *ps.MoveTablesStopRequest) (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.Workflow, qt.Equals, "my-workflow") + c.Assert(req.TargetKeyspace, qt.Equals, "target-ks") + return json.RawMessage(`{"summary":"Streams stopped"}`), nil + }, + } + + var buf bytes.Buffer + ch := moveTablesTestHelper(org, svc, nil, &buf) + + cmd := MoveTablesCmd(ch) + cmd.SetArgs([]string{"stop", db, branch, + "--workflow", "my-workflow", + "--target-keyspace", "target-ks", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.StopFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, map[string]any{ + "summary": "Streams stopped", + "next_steps": []any{map[string]any{ + "command": "pscale branch vtctld move-tables status my-db my-branch --org my-org --workflow my-workflow --target-keyspace target-ks --format json", + "reason": "Confirm streams are stopped", + }}, + }) +} + func TestMoveTablesStatusAddsNextSteps(t *testing.T) { c := qt.New(t) diff --git a/internal/mock/vtctld_move_tables.go b/internal/mock/vtctld_move_tables.go index a1f054d9..dd0cfe83 100644 --- a/internal/mock/vtctld_move_tables.go +++ b/internal/mock/vtctld_move_tables.go @@ -23,6 +23,9 @@ type MoveTablesService struct { StartFn func(context.Context, *ps.MoveTablesStartRequest) (json.RawMessage, error) StartFnInvoked bool + StopFn func(context.Context, *ps.MoveTablesStopRequest) (json.RawMessage, error) + StopFnInvoked bool + SwitchTrafficFn func(context.Context, *ps.MoveTablesSwitchTrafficRequest) (*ps.VtctldOperationReference, error) SwitchTrafficFnInvoked bool @@ -61,6 +64,11 @@ func (s *MoveTablesService) Start(ctx context.Context, req *ps.MoveTablesStartRe return s.StartFn(ctx, req) } +func (s *MoveTablesService) Stop(ctx context.Context, req *ps.MoveTablesStopRequest) (json.RawMessage, error) { + s.StopFnInvoked = true + return s.StopFn(ctx, req) +} + func (s *MoveTablesService) SwitchTraffic(ctx context.Context, req *ps.MoveTablesSwitchTrafficRequest) (*ps.VtctldOperationReference, error) { s.SwitchTrafficFnInvoked = true return s.SwitchTrafficFn(ctx, req) diff --git a/internal/planetscale/vtctld_move_tables.go b/internal/planetscale/vtctld_move_tables.go index 03859ed2..c39f5b12 100644 --- a/internal/planetscale/vtctld_move_tables.go +++ b/internal/planetscale/vtctld_move_tables.go @@ -17,6 +17,7 @@ type MoveTablesService interface { Show(context.Context, *MoveTablesShowRequest) (json.RawMessage, error) Status(context.Context, *MoveTablesStatusRequest) (json.RawMessage, error) Start(context.Context, *MoveTablesStartRequest) (json.RawMessage, error) + Stop(context.Context, *MoveTablesStopRequest) (json.RawMessage, error) SwitchTraffic(context.Context, *MoveTablesSwitchTrafficRequest) (*VtctldOperationReference, error) ReverseTraffic(context.Context, *MoveTablesReverseTrafficRequest) (*VtctldOperationReference, error) Cancel(context.Context, *MoveTablesCancelRequest) (*VtctldOperationReference, error) @@ -83,6 +84,15 @@ type MoveTablesStartRequest struct { TargetKeyspace string `json:"target_keyspace"` } +// MoveTablesStopRequest is a request for stopping a MoveTables workflow. +type MoveTablesStopRequest struct { + Organization string `json:"-"` + Database string `json:"-"` + Branch string `json:"-"` + Workflow string `json:"-"` + TargetKeyspace string `json:"target_keyspace"` +} + // MoveTablesSwitchTrafficRequest is a request for switching traffic for a MoveTables workflow. type MoveTablesSwitchTrafficRequest struct { Organization string `json:"-"` @@ -215,6 +225,19 @@ func (s *moveTablesService) Start(ctx context.Context, req *MoveTablesStartReque return resp.Data, nil } +func (s *moveTablesService) Stop(ctx context.Context, req *MoveTablesStopRequest) (json.RawMessage, error) { + p := path.Join(moveTablesWorkflowAPIPath(req.Organization, req.Database, req.Branch, req.Workflow), "stop") + httpReq, err := s.client.newRequest(http.MethodPost, p, req) + 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 *moveTablesService) SwitchTraffic(ctx context.Context, req *MoveTablesSwitchTrafficRequest) (*VtctldOperationReference, error) { p := path.Join(moveTablesWorkflowAPIPath(req.Organization, req.Database, req.Branch, req.Workflow), "switch-traffic") return s.enqueueOperation(ctx, p, req) diff --git a/internal/planetscale/vtctld_move_tables_test.go b/internal/planetscale/vtctld_move_tables_test.go index fc617ffa..5d530b71 100644 --- a/internal/planetscale/vtctld_move_tables_test.go +++ b/internal/planetscale/vtctld_move_tables_test.go @@ -354,6 +354,39 @@ func TestMoveTables_Start(t *testing.T) { c.Assert(string(data), qt.Equals, `{"summary":"Streams started"}`) } +func TestMoveTables_Stop(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.MethodPost) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/move-tables/workflows/my-workflow/stop") + + var body map[string]interface{} + err := json.NewDecoder(r.Body).Decode(&body) + c.Assert(err, qt.IsNil) + c.Assert(body["target_keyspace"], qt.Equals, "target") + + w.WriteHeader(http.StatusOK) + _, err = w.Write([]byte(`{"data":{"summary":"Streams stopped"}}`)) + 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.MoveTables.Stop(ctx, &MoveTablesStopRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + Workflow: "my-workflow", + TargetKeyspace: "target", + }) + c.Assert(err, qt.IsNil) + c.Assert(string(data), qt.Equals, `{"summary":"Streams stopped"}`) +} + func TestMoveTables_SwitchTraffic(t *testing.T) { c := qt.New(t)