From a53cbea7ce6720c6b3589b5a8a79ac09f287dbf1 Mon Sep 17 00:00:00 2001 From: Elom Gomez Date: Tue, 26 May 2026 16:58:20 -0500 Subject: [PATCH] Update move table cancel to use async operation --- go.mod | 2 +- go.sum | 4 +- internal/cmd/branch/vtctld/move_tables.go | 7 ++- .../cmd/branch/vtctld/move_tables_test.go | 50 +++++++++++++++++++ internal/mock/vtctld_move_tables.go | 4 +- 5 files changed, 61 insertions(+), 6 deletions(-) diff --git a/go.mod b/go.mod index 5802ea1a..52b1671a 100644 --- a/go.mod +++ b/go.mod @@ -25,7 +25,7 @@ require ( github.com/mattn/go-shellwords v1.0.12 github.com/mitchellh/go-homedir v1.1.0 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c - github.com/planetscale/planetscale-go v0.163.0 + github.com/planetscale/planetscale-go v0.164.0 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 github.com/spf13/cobra v1.10.2 diff --git a/go.sum b/go.sum index cac8c69b..12a75e5e 100644 --- a/go.sum +++ b/go.sum @@ -176,8 +176,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjL github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2vvqGZLvoQfpaGg/j1fNDr4j03s3PRz4rVY= github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q= -github.com/planetscale/planetscale-go v0.163.0 h1:LsMWsnrDUncbstRBMdYLCTyu5nQzNyQc60ytXaHuezQ= -github.com/planetscale/planetscale-go v0.163.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= +github.com/planetscale/planetscale-go v0.164.0 h1:k/Jw1robfLhcGS3f4FtqxQXEDuuddj0bnTE/jbAo4/o= +github.com/planetscale/planetscale-go v0.164.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs= github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs= github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8= diff --git a/internal/cmd/branch/vtctld/move_tables.go b/internal/cmd/branch/vtctld/move_tables.go index 2cedf6a5..8c93a11a 100644 --- a/internal/cmd/branch/vtctld/move_tables.go +++ b/internal/cmd/branch/vtctld/move_tables.go @@ -429,7 +429,12 @@ func MoveTablesCancelCmd(ch *cmdutil.Helper) *cobra.Command { req.KeepRoutingRules = &flags.keepRoutingRules } - data, err := client.MoveTables.Cancel(ctx, req) + operation, err := client.MoveTables.Cancel(ctx, req) + if err != nil { + return cmdutil.HandleError(err) + } + + data, err := waitForMoveTablesOperationResult(ctx, client, ch.Config.Organization, database, branch, operation.ID) if err != nil { return cmdutil.HandleError(err) } diff --git a/internal/cmd/branch/vtctld/move_tables_test.go b/internal/cmd/branch/vtctld/move_tables_test.go index c1ad301e..026a8b31 100644 --- a/internal/cmd/branch/vtctld/move_tables_test.go +++ b/internal/cmd/branch/vtctld/move_tables_test.go @@ -497,6 +497,56 @@ func TestMoveTablesSwitchTrafficOperationTimeout(t *testing.T) { c.Assert(buf.String(), qt.Equals, "") } +func TestMoveTablesCancelWithFlags(t *testing.T) { + c := qt.New(t) + setMoveTablesPollInterval(t, 0) + + org := "my-org" + db := "my-db" + branch := "my-branch" + + keepData := false + svc := &mock.MoveTablesService{ + CancelFn: func(ctx context.Context, req *ps.MoveTablesCancelRequest) (*ps.VtctldOperationReference, error) { + c.Assert(req.Workflow, qt.Equals, "my-workflow") + c.Assert(req.TargetKeyspace, qt.Equals, "target-ks") + c.Assert(req.KeepData, qt.IsNotNil) + c.Assert(*req.KeepData, qt.Equals, keepData) + return &ps.VtctldOperationReference{ID: "cancel-op"}, nil + }, + } + vtctldSvc := &mock.VtctldService{ + GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, 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.ID, qt.Equals, "cancel-op") + + return &ps.VtctldOperation{ + ID: "cancel-op", + State: "completed", + Completed: true, + Result: json.RawMessage(`{"summary":"cancelled"}`), + }, nil + }, + } + + var buf bytes.Buffer + ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf) + + cmd := MoveTablesCmd(ch) + cmd.SetArgs([]string{"cancel", db, branch, + "--workflow", "my-workflow", + "--target-keyspace", "target-ks", + "--keep-data=false", + }) + err := cmd.Execute() + c.Assert(err, qt.IsNil) + c.Assert(svc.CancelFnInvoked, qt.IsTrue) + c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, map[string]string{"summary": "cancelled"}) +} + func TestMoveTablesShow(t *testing.T) { c := qt.New(t) diff --git a/internal/mock/vtctld_move_tables.go b/internal/mock/vtctld_move_tables.go index 6e630a28..9497b229 100644 --- a/internal/mock/vtctld_move_tables.go +++ b/internal/mock/vtctld_move_tables.go @@ -23,7 +23,7 @@ type MoveTablesService struct { ReverseTrafficFn func(context.Context, *ps.MoveTablesReverseTrafficRequest) (*ps.VtctldOperationReference, error) ReverseTrafficFnInvoked bool - CancelFn func(context.Context, *ps.MoveTablesCancelRequest) (json.RawMessage, error) + CancelFn func(context.Context, *ps.MoveTablesCancelRequest) (*ps.VtctldOperationReference, error) CancelFnInvoked bool CompleteFn func(context.Context, *ps.MoveTablesCompleteRequest) (*ps.VtctldOperationReference, error) @@ -55,7 +55,7 @@ func (s *MoveTablesService) ReverseTraffic(ctx context.Context, req *ps.MoveTabl return s.ReverseTrafficFn(ctx, req) } -func (s *MoveTablesService) Cancel(ctx context.Context, req *ps.MoveTablesCancelRequest) (json.RawMessage, error) { +func (s *MoveTablesService) Cancel(ctx context.Context, req *ps.MoveTablesCancelRequest) (*ps.VtctldOperationReference, error) { s.CancelFnInvoked = true return s.CancelFn(ctx, req) }