Merge pull request #1257 from planetscale/update-move-table-cancel-to-async

Update move table cancel to use async operation

Author: no-itsbackpack

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

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=

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)
 			}

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)
 

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)
 }