Skip to content

Commit 9aaa88b

Browse files
nickvanwnickvanw
authored
Merge pull request #1247 from planetscale/nvanwiggeren/async-move-tables-create
Handle async MoveTables Create
2 parents 09b8752 + 7f61109 commit 9aaa88b

File tree

5 files changed

+59
-44
lines changed

5 files changed

+59
-44
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ require (
2525
github.com/mattn/go-shellwords v1.0.12
2626
github.com/mitchellh/go-homedir v1.1.0
2727
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
28-
github.com/planetscale/planetscale-go v0.161.0
28+
github.com/planetscale/planetscale-go v0.162.0
2929
github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
3030
github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
3131
github.com/spf13/cobra v1.10.2

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjL
176176
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
177177
github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2vvqGZLvoQfpaGg/j1fNDr4j03s3PRz4rVY=
178178
github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
179-
github.com/planetscale/planetscale-go v0.161.0 h1:YWpOJXJTJKX2AFHoKkO/IZdRxCmOQAS+O6A+jJuteL8=
180-
github.com/planetscale/planetscale-go v0.161.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
179+
github.com/planetscale/planetscale-go v0.162.0 h1:uOpqOeIikRXu+5AxwUBitvImhPtRca6Kfmt6ahmXzDs=
180+
github.com/planetscale/planetscale-go v0.162.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
181181
github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=
182182
github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs=
183183
github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8=

‎internal/cmd/branch/vtctld/move_tables.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,12 @@ func MoveTablesCreateCmd(ch *cmdutil.Helper) *cobra.Command {
105105
req.AtomicCopy = &flags.atomicCopy
106106
}
107107

108-
data, err := client.MoveTables.Create(ctx, req)
108+
operation, err := client.MoveTables.Create(ctx, req)
109+
if err != nil {
110+
return cmdutil.HandleError(err)
111+
}
112+
113+
data, err := waitForMoveTablesOperationResult(ctx, client, ch.Config.Organization, database, branch, operation.ID)
109114
if err != nil {
110115
return cmdutil.HandleError(err)
111116
}

‎internal/cmd/branch/vtctld/move_tables_test.go‎

Lines changed: 48 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -63,13 +63,14 @@ func setMoveTablesOperationTimeouts(t *testing.T, defaultTimeout, timeoutBuffer
6363

6464
func TestMoveTablesCreate(t *testing.T) {
6565
c := qt.New(t)
66+
setMoveTablesPollInterval(t, 0)
6667

6768
org := "my-org"
6869
db := "my-db"
6970
branch := "my-branch"
7071

7172
svc := &mock.MoveTablesService{
72-
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) {
73+
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) {
7374
c.Assert(req.Organization, qt.Equals, org)
7475
c.Assert(req.Database, qt.Equals, db)
7576
c.Assert(req.Branch, qt.Equals, branch)
@@ -82,25 +83,29 @@ func TestMoveTablesCreate(t *testing.T) {
8283
c.Assert(req.TabletTypes, qt.IsNil)
8384
c.Assert(req.ExcludeTables, qt.IsNil)
8485
c.Assert(req.AtomicCopy, qt.IsNil)
85-
return json.RawMessage(`{"summary":"created"}`), nil
86+
return &ps.VtctldOperationReference{ID: "create-op"}, nil
8687
},
8788
}
8889

89-
var buf bytes.Buffer
90-
format := printer.JSON
91-
p := printer.NewPrinter(&format)
92-
p.SetResourceOutput(&buf)
90+
vtctldSvc := &mock.VtctldService{
91+
GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) {
92+
c.Assert(req.Organization, qt.Equals, org)
93+
c.Assert(req.Database, qt.Equals, db)
94+
c.Assert(req.Branch, qt.Equals, branch)
95+
c.Assert(req.ID, qt.Equals, "create-op")
9396

94-
ch := &cmdutil.Helper{
95-
Printer: p,
96-
Config: &config.Config{Organization: org},
97-
Client: func() (*ps.Client, error) {
98-
return &ps.Client{
99-
MoveTables: svc,
97+
return &ps.VtctldOperation{
98+
ID: "create-op",
99+
State: "completed",
100+
Completed: true,
101+
Result: json.RawMessage(`{"summary":"created"}`),
100102
}, nil
101103
},
102104
}
103105

106+
var buf bytes.Buffer
107+
ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf)
108+
104109
cmd := MoveTablesCmd(ch)
105110
cmd.SetArgs([]string{"create", db, branch,
106111
"--workflow", "my-workflow",
@@ -110,41 +115,43 @@ func TestMoveTablesCreate(t *testing.T) {
110115
err := cmd.Execute()
111116
c.Assert(err, qt.IsNil)
112117
c.Assert(svc.CreateFnInvoked, qt.IsTrue)
118+
c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue)
119+
c.Assert(buf.String(), qt.JSONEquals, map[string]string{"summary": "created"})
113120
}
114121

115122
func TestMoveTablesCreateWithDeferSecondaryKeysFalse(t *testing.T) {
116123
c := qt.New(t)
124+
setMoveTablesPollInterval(t, 0)
117125

118126
org := "my-org"
119127
db := "my-db"
120128
branch := "my-branch"
121129

122130
svc := &mock.MoveTablesService{
123-
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) {
131+
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) {
124132
c.Assert(req.Workflow, qt.Equals, "my-workflow")
125133
c.Assert(req.TargetKeyspace, qt.Equals, "target-ks")
126134
c.Assert(req.SourceKeyspace, qt.Equals, "source-ks")
127135
c.Assert(req.DeferSecondaryKeys, qt.IsNotNil)
128136
c.Assert(*req.DeferSecondaryKeys, qt.IsFalse)
129-
return json.RawMessage(`{"summary":"created"}`), nil
137+
return &ps.VtctldOperationReference{ID: "create-op"}, nil
130138
},
131139
}
132140

133-
var buf bytes.Buffer
134-
format := printer.JSON
135-
p := printer.NewPrinter(&format)
136-
p.SetResourceOutput(&buf)
137-
138-
ch := &cmdutil.Helper{
139-
Printer: p,
140-
Config: &config.Config{Organization: org},
141-
Client: func() (*ps.Client, error) {
142-
return &ps.Client{
143-
MoveTables: svc,
141+
vtctldSvc := &mock.VtctldService{
142+
GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) {
143+
return &ps.VtctldOperation{
144+
ID: "create-op",
145+
State: "completed",
146+
Completed: true,
147+
Result: json.RawMessage(`{"summary":"created"}`),
144148
}, nil
145149
},
146150
}
147151

152+
var buf bytes.Buffer
153+
ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf)
154+
148155
cmd := MoveTablesCmd(ch)
149156
cmd.SetArgs([]string{"create", db, branch,
150157
"--workflow", "my-workflow",
@@ -155,17 +162,19 @@ func TestMoveTablesCreateWithDeferSecondaryKeysFalse(t *testing.T) {
155162
err := cmd.Execute()
156163
c.Assert(err, qt.IsNil)
157164
c.Assert(svc.CreateFnInvoked, qt.IsTrue)
165+
c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue)
158166
}
159167

160168
func TestMoveTablesCreateWithAllFlags(t *testing.T) {
161169
c := qt.New(t)
170+
setMoveTablesPollInterval(t, 0)
162171

163172
org := "my-org"
164173
db := "my-db"
165174
branch := "my-branch"
166175

167176
svc := &mock.MoveTablesService{
168-
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) {
177+
CreateFn: func(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) {
169178
c.Assert(req.Workflow, qt.Equals, "my-workflow")
170179
c.Assert(req.TargetKeyspace, qt.Equals, "target-ks")
171180
c.Assert(req.SourceKeyspace, qt.Equals, "source-ks")
@@ -176,25 +185,24 @@ func TestMoveTablesCreateWithAllFlags(t *testing.T) {
176185
c.Assert(*req.AtomicCopy, qt.IsTrue)
177186
c.Assert(req.AllTables, qt.IsNotNil)
178187
c.Assert(*req.AllTables, qt.IsTrue)
179-
return json.RawMessage(`{"summary":"created"}`), nil
188+
return &ps.VtctldOperationReference{ID: "create-op"}, nil
180189
},
181190
}
182191

183-
var buf bytes.Buffer
184-
format := printer.JSON
185-
p := printer.NewPrinter(&format)
186-
p.SetResourceOutput(&buf)
187-
188-
ch := &cmdutil.Helper{
189-
Printer: p,
190-
Config: &config.Config{Organization: org},
191-
Client: func() (*ps.Client, error) {
192-
return &ps.Client{
193-
MoveTables: svc,
192+
vtctldSvc := &mock.VtctldService{
193+
GetOperationFn: func(ctx context.Context, req *ps.GetVtctldOperationRequest) (*ps.VtctldOperation, error) {
194+
return &ps.VtctldOperation{
195+
ID: "create-op",
196+
State: "completed",
197+
Completed: true,
198+
Result: json.RawMessage(`{"summary":"created"}`),
194199
}, nil
195200
},
196201
}
197202

203+
var buf bytes.Buffer
204+
ch := moveTablesTestHelper(org, svc, vtctldSvc, &buf)
205+
198206
cmd := MoveTablesCmd(ch)
199207
cmd.SetArgs([]string{"create", db, branch,
200208
"--workflow", "my-workflow",
@@ -209,6 +217,8 @@ func TestMoveTablesCreateWithAllFlags(t *testing.T) {
209217
err := cmd.Execute()
210218
c.Assert(err, qt.IsNil)
211219
c.Assert(svc.CreateFnInvoked, qt.IsTrue)
220+
c.Assert(vtctldSvc.GetOperationFnInvoked, qt.IsTrue)
221+
c.Assert(buf.String(), qt.JSONEquals, map[string]string{"summary": "created"})
212222
}
213223

214224
func TestMoveTablesSwitchTrafficWithMaxLag(t *testing.T) {

‎internal/mock/vtctld_move_tables.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import (
88
)
99

1010
type MoveTablesService struct {
11-
CreateFn func(context.Context, *ps.MoveTablesCreateRequest) (json.RawMessage, error)
11+
CreateFn func(context.Context, *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error)
1212
CreateFnInvoked bool
1313

1414
ShowFn func(context.Context, *ps.MoveTablesShowRequest) (json.RawMessage, error)
@@ -30,7 +30,7 @@ type MoveTablesService struct {
3030
CompleteFnInvoked bool
3131
}
3232

33-
func (s *MoveTablesService) Create(ctx context.Context, req *ps.MoveTablesCreateRequest) (json.RawMessage, error) {
33+
func (s *MoveTablesService) Create(ctx context.Context, req *ps.MoveTablesCreateRequest) (*ps.VtctldOperationReference, error) {
3434
s.CreateFnInvoked = true
3535
return s.CreateFn(ctx, req)
3636
}

0 commit comments

Comments
 (0)