From 7f611093330d6e84f7594d1c9cfe15c0061f956d Mon Sep 17 00:00:00 2001 From: Nick Van Wiggeren Date: Tue, 7 Apr 2026 21:25:56 +0000 Subject: [PATCH] Handle async MoveTables Create MoveTables Create is now an async operation that returns an operation reference instead of the result directly. Poll GetOperation for the result, matching the pattern already used by SwitchTraffic, ReverseTraffic, and Complete. Bumps planetscale-go to v0.162.0. --- go.mod | 2 +- go.sum | 4 +- internal/cmd/branch/vtctld/move_tables.go | 7 +- .../cmd/branch/vtctld/move_tables_test.go | 86 +++++++++++-------- internal/mock/vtctld_move_tables.go | 4 +- 5 files changed, 59 insertions(+), 44 deletions(-) diff --git a/go.mod b/go.mod index 5ea32d3d..d59354a5 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.161.0 + github.com/planetscale/planetscale-go v0.162.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 2e3dd570..598e8c23 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.161.0 h1:YWpOJXJTJKX2AFHoKkO/IZdRxCmOQAS+O6A+jJuteL8= -github.com/planetscale/planetscale-go v0.161.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0= +github.com/planetscale/planetscale-go v0.162.0 h1:uOpqOeIikRXu+5AxwUBitvImhPtRca6Kfmt6ahmXzDs= +github.com/planetscale/planetscale-go v0.162.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 b81b76da..64a44182 100644 --- a/internal/cmd/branch/vtctld/move_tables.go +++ b/internal/cmd/branch/vtctld/move_tables.go @@ -105,7 +105,12 @@ func MoveTablesCreateCmd(ch *cmdutil.Helper) *cobra.Command { req.AtomicCopy = &flags.atomicCopy } - data, err := client.MoveTables.Create(ctx, req) + operation, err := client.MoveTables.Create(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 bfa69c11..6355192e 100644 --- a/internal/cmd/branch/vtctld/move_tables_test.go +++ b/internal/cmd/branch/vtctld/move_tables_test.go @@ -63,13 +63,14 @@ func setMoveTablesOperationTimeouts(t *testing.T, defaultTimeout, timeoutBuffer func TestMoveTablesCreate(t *testing.T) { c := qt.New(t) + setMoveTablesPollInterval(t, 0) org := "my-org" db := "my-db" branch := "my-branch" svc := &mock.MoveTablesService{ - CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) { + CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) { c.Assert(req.Organization, qt.Equals, org) c.Assert(req.Database, qt.Equals, db) c.Assert(req.Branch, qt.Equals, branch) @@ -82,25 +83,29 @@ func TestMoveTablesCreate(t *testing.T) { c.Assert(req.TabletTypes, qt.IsNil) c.Assert(req.ExcludeTables, qt.IsNil) c.Assert(req.AtomicCopy, qt.IsNil) - return json.RawMessage(`{"summary":"created"}`), nil + return &ps.VtctldOperationReference{ID: "create-op"}, nil }, } - var buf bytes.Buffer - format := printer.JSON - p := printer.NewPrinter(&format) - p.SetResourceOutput(&buf) + 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, "create-op") - ch := &cmdutil.Helper{ - Printer: p, - Config: &config.Config{Organization: org}, - Client: func() (*ps.Client, error) { - return &ps.Client{ - MoveTables: svc, + return &ps.VtctldOperation{ + ID: "create-op", + State: "completed", + Completed: true, + Result: json.RawMessage(`{"summary":"created"}`), }, nil }, } + var buf bytes.Buffer + ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf) + cmd := MoveTablesCmd(ch) cmd.SetArgs([]string{"create", db, branch, "--workflow", "my-workflow", @@ -110,41 +115,43 @@ func TestMoveTablesCreate(t *testing.T) { err := cmd.Execute() c.Assert(err, qt.IsNil) c.Assert(svc.CreateFnInvoked, qt.IsTrue) + c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, map[string]string{"summary": "created"}) } func TestMoveTablesCreateWithDeferSecondaryKeysFalse(t *testing.T) { c := qt.New(t) + setMoveTablesPollInterval(t, 0) org := "my-org" db := "my-db" branch := "my-branch" svc := &mock.MoveTablesService{ - CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) { + CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) { c.Assert(req.Workflow, qt.Equals, "my-workflow") c.Assert(req.TargetKeyspace, qt.Equals, "target-ks") c.Assert(req.SourceKeyspace, qt.Equals, "source-ks") c.Assert(req.DeferSecondaryKeys, qt.IsNotNil) c.Assert(*req.DeferSecondaryKeys, qt.IsFalse) - return json.RawMessage(`{"summary":"created"}`), nil + return &ps.VtctldOperationReference{ID: "create-op"}, nil }, } - var buf bytes.Buffer - format := printer.JSON - p := printer.NewPrinter(&format) - p.SetResourceOutput(&buf) - - ch := &cmdutil.Helper{ - Printer: p, - Config: &config.Config{Organization: org}, - Client: func() (*ps.Client, error) { - return &ps.Client{ - MoveTables: svc, + vtctldSvc := &mock.VtctldService{ + GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) { + return &ps.VtctldOperation{ + ID: "create-op", + State: "completed", + Completed: true, + Result: json.RawMessage(`{"summary":"created"}`), }, nil }, } + var buf bytes.Buffer + ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf) + cmd := MoveTablesCmd(ch) cmd.SetArgs([]string{"create", db, branch, "--workflow", "my-workflow", @@ -155,17 +162,19 @@ func TestMoveTablesCreateWithDeferSecondaryKeysFalse(t *testing.T) { err := cmd.Execute() c.Assert(err, qt.IsNil) c.Assert(svc.CreateFnInvoked, qt.IsTrue) + c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue) } func TestMoveTablesCreateWithAllFlags(t *testing.T) { c := qt.New(t) + setMoveTablesPollInterval(t, 0) org := "my-org" db := "my-db" branch := "my-branch" svc := &mock.MoveTablesService{ - CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) { + CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) { c.Assert(req.Workflow, qt.Equals, "my-workflow") c.Assert(req.TargetKeyspace, qt.Equals, "target-ks") c.Assert(req.SourceKeyspace, qt.Equals, "source-ks") @@ -176,25 +185,24 @@ func TestMoveTablesCreateWithAllFlags(t *testing.T) { c.Assert(*req.AtomicCopy, qt.IsTrue) c.Assert(req.AllTables, qt.IsNotNil) c.Assert(*req.AllTables, qt.IsTrue) - return json.RawMessage(`{"summary":"created"}`), nil + return &ps.VtctldOperationReference{ID: "create-op"}, nil }, } - var buf bytes.Buffer - format := printer.JSON - p := printer.NewPrinter(&format) - p.SetResourceOutput(&buf) - - ch := &cmdutil.Helper{ - Printer: p, - Config: &config.Config{Organization: org}, - Client: func() (*ps.Client, error) { - return &ps.Client{ - MoveTables: svc, + vtctldSvc := &mock.VtctldService{ + GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) { + return &ps.VtctldOperation{ + ID: "create-op", + State: "completed", + Completed: true, + Result: json.RawMessage(`{"summary":"created"}`), }, nil }, } + var buf bytes.Buffer + ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf) + cmd := MoveTablesCmd(ch) cmd.SetArgs([]string{"create", db, branch, "--workflow", "my-workflow", @@ -209,6 +217,8 @@ func TestMoveTablesCreateWithAllFlags(t *testing.T) { err := cmd.Execute() c.Assert(err, qt.IsNil) c.Assert(svc.CreateFnInvoked, qt.IsTrue) + c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, map[string]string{"summary": "created"}) } func TestMoveTablesSwitchTrafficWithMaxLag(t *testing.T) { diff --git a/internal/mock/vtctld_move_tables.go b/internal/mock/vtctld_move_tables.go index 9a653fe7..6e630a28 100644 --- a/internal/mock/vtctld_move_tables.go +++ b/internal/mock/vtctld_move_tables.go @@ -8,7 +8,7 @@ import ( ) type MoveTablesService struct { - CreateFn func(context.Context, *ps.MoveTablesCreateRequest) (json.RawMessage, error) + CreateFn func(context.Context, *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) CreateFnInvoked bool ShowFn func(context.Context, *ps.MoveTablesShowRequest) (json.RawMessage, error) @@ -30,7 +30,7 @@ type MoveTablesService struct { CompleteFnInvoked bool } -func (s *MoveTablesService) Create(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) { +func (s *MoveTablesService) Create(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) { s.CreateFnInvoked = true return s.CreateFn(ctx, req) }