Skip to content

Commit c018f92

Browse files
authored
Add vtctld refresh-state-by-shard command (#1269)
1 parent 8b6f660 commit c018f92

6 files changed

Lines changed: 144 additions & 3 deletions

File tree

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ require (
3030
github.com/mitchellh/go-homedir v1.1.0
3131
github.com/muesli/termenv v0.16.0
3232
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
33-
github.com/planetscale/planetscale-go v0.171.0
33+
github.com/planetscale/planetscale-go v0.174.0
3434
github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
3535
github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
3636
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.171.0 h1:ZBltEV1SANoQDUMc0VVUmGD4urjquZqh/y9XgGxG3Pw=
180-
github.com/planetscale/planetscale-go v0.171.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
179+
github.com/planetscale/planetscale-go v0.174.0 h1:u7UeL+XApeY2ZvettqecRIoemsVMBptAVV4i+2Dq3AA=
180+
github.com/planetscale/planetscale-go v0.174.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=
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package vtctld
2+
3+
import (
4+
"fmt"
5+
6+
"github.com/planetscale/cli/internal/cmdutil"
7+
ps "github.com/planetscale/planetscale-go/planetscale"
8+
"github.com/spf13/cobra"
9+
)
10+
11+
// RefreshStateByShardCmd reloads tablet records for all tablets in a shard via vtctld.
12+
func RefreshStateByShardCmd(ch *cmdutil.Helper) *cobra.Command {
13+
var flags struct {
14+
keyspace string
15+
shard string
16+
cells []string
17+
}
18+
19+
cmd := &cobra.Command{
20+
Use: "refresh-state-by-shard <database> <branch>",
21+
Short: "Reload tablet records for all tablets in a shard",
22+
Long: "Reload the tablet record for all tablets in a shard via vtctld, " +
23+
"optionally limited to the specified cells.",
24+
Args: cmdutil.RequiredArgs("database", "branch"),
25+
RunE: func(cmd *cobra.Command, args []string) error {
26+
ctx := cmd.Context()
27+
database, branch := args[0], args[1]
28+
29+
if flags.keyspace == "" {
30+
return fmt.Errorf("keyspace is required")
31+
}
32+
if flags.shard == "" {
33+
return fmt.Errorf("shard is required")
34+
}
35+
36+
client, err := ch.Client()
37+
if err != nil {
38+
return err
39+
}
40+
41+
end := ch.Printer.PrintProgress(
42+
fmt.Sprintf("Refreshing tablet state for %s/%s on %s…",
43+
flags.keyspace, flags.shard,
44+
progressTarget(ch.Config.Organization, database, branch)))
45+
defer end()
46+
47+
data, err := client.Vtctld.RefreshStateByShard(ctx, &ps.VtctldRefreshStateByShardRequest{
48+
Organization: ch.Config.Organization,
49+
Database: database,
50+
Branch: branch,
51+
Keyspace: flags.keyspace,
52+
Shard: flags.shard,
53+
Cells: flags.cells,
54+
})
55+
if err != nil {
56+
return cmdutil.HandleError(err)
57+
}
58+
59+
end()
60+
return ch.Printer.PrettyPrintJSON(data)
61+
},
62+
}
63+
64+
cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Keyspace name")
65+
cmd.Flags().StringVar(&flags.shard, "shard", "", "Shard name (e.g. \"-\" for unsharded)")
66+
cmd.Flags().StringSliceVar(&flags.cells, "cells", nil, "Cells to refresh (comma-separated)")
67+
cmd.MarkFlagRequired("keyspace")
68+
cmd.MarkFlagRequired("shard")
69+
70+
return cmd
71+
}
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package vtctld
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/json"
7+
"testing"
8+
9+
qt "github.com/frankban/quicktest"
10+
11+
"github.com/planetscale/cli/internal/cmdutil"
12+
"github.com/planetscale/cli/internal/config"
13+
"github.com/planetscale/cli/internal/mock"
14+
"github.com/planetscale/cli/internal/printer"
15+
ps "github.com/planetscale/planetscale-go/planetscale"
16+
)
17+
18+
func TestRefreshStateByShard(t *testing.T) {
19+
c := qt.New(t)
20+
21+
org := "my-org"
22+
db := "my-db"
23+
branch := "my-branch"
24+
25+
svc := &mock.VtctldService{
26+
RefreshStateByShardFn: func(ctx context.Context, req *ps.VtctldRefreshStateByShardRequest) (json.RawMessage, error) {
27+
c.Assert(req.Organization, qt.Equals, org)
28+
c.Assert(req.Database, qt.Equals, db)
29+
c.Assert(req.Branch, qt.Equals, branch)
30+
c.Assert(req.Keyspace, qt.Equals, "commerce")
31+
c.Assert(req.Shard, qt.Equals, "-")
32+
c.Assert(req.Cells, qt.DeepEquals, []string{"zone1"})
33+
return json.RawMessage(`{}`), nil
34+
},
35+
}
36+
37+
var buf bytes.Buffer
38+
format := printer.JSON
39+
p := printer.NewPrinter(&format)
40+
p.SetResourceOutput(&buf)
41+
42+
ch := &cmdutil.Helper{
43+
Printer: p,
44+
Config: &config.Config{Organization: org},
45+
Client: func() (*ps.Client, error) {
46+
return &ps.Client{
47+
Vtctld: svc,
48+
}, nil
49+
},
50+
}
51+
52+
cmd := RefreshStateByShardCmd(ch)
53+
cmd.SetArgs([]string{db, branch})
54+
cmd.Flags().Set("keyspace", "commerce")
55+
cmd.Flags().Set("shard", "-")
56+
cmd.Flags().Set("cells", "zone1")
57+
58+
err := cmd.Execute()
59+
c.Assert(err, qt.IsNil)
60+
c.Assert(svc.RefreshStateByShardFnInvoked, qt.IsTrue)
61+
}

internal/cmd/branch/vtctld/vtctld.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ func VtctldCmd(ch *cmdutil.Helper) *cobra.Command {
2323
cmd.AddCommand(GetRoutingRulesCmd(ch))
2424
cmd.AddCommand(GetShardCmd(ch))
2525
cmd.AddCommand(SetShardTabletControlCmd(ch))
26+
cmd.AddCommand(RefreshStateByShardCmd(ch))
2627
cmd.AddCommand(ListTabletsCmd(ch))
2728
cmd.AddCommand(StartWorkflowCmd(ch))
2829
cmd.AddCommand(StopWorkflowCmd(ch))

internal/mock/vtctld_general.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ type VtctldService struct {
2323
SetShardTabletControlFn func(context.Context, *ps.VtctldSetShardTabletControlRequest) (json.RawMessage, error)
2424
SetShardTabletControlFnInvoked bool
2525

26+
RefreshStateByShardFn func(context.Context, *ps.VtctldRefreshStateByShardRequest) (json.RawMessage, error)
27+
RefreshStateByShardFnInvoked bool
28+
2629
ListTabletsFn func(context.Context, *ps.ListBranchTabletsRequest) ([]*ps.TabletGroup, error)
2730
ListTabletsFnInvoked bool
2831

@@ -70,6 +73,11 @@ func (s *VtctldService) SetShardTabletControl(ctx context.Context, req *ps.Vtctl
7073
return s.SetShardTabletControlFn(ctx, req)
7174
}
7275

76+
func (s *VtctldService) RefreshStateByShard(ctx context.Context, req *ps.VtctldRefreshStateByShardRequest) (json.RawMessage, error) {
77+
s.RefreshStateByShardFnInvoked = true
78+
return s.RefreshStateByShardFn(ctx, req)
79+
}
80+
7381
func (s *VtctldService) ListTablets(ctx context.Context, req *ps.ListBranchTabletsRequest) ([]*ps.TabletGroup, error) {
7482
s.ListTabletsFnInvoked = true
7583
return s.ListTabletsFn(ctx, req)

0 commit comments

Comments
 (0)