Add RefreshStateByShard vtctld client method

planetscale/vtctld_general.go

@@ -18,6 +18,7 @@ type VtctldService interface {
 	GetRoutingRules(context.Context, *VtctldGetRoutingRulesRequest) (json.RawMessage, error)
 	GetShard(context.Context, *VtctldGetShardRequest) (json.RawMessage, error)
 	SetShardTabletControl(context.Context, *VtctldSetShardTabletControlRequest) (json.RawMessage, error)
+	RefreshStateByShard(context.Context, *VtctldRefreshStateByShardRequest) (json.RawMessage, error)
 	ListTablets(context.Context, *ListBranchTabletsRequest) ([]*TabletGroup, error)
 	StartWorkflow(context.Context, *VtctldStartWorkflowRequest) (json.RawMessage, error)
 	StopWorkflow(context.Context, *VtctldStopWorkflowRequest) (json.RawMessage, error)
@@ -77,6 +78,18 @@ type VtctldSetShardTabletControlRequest struct {
 	DisableQueryService *bool    `json:"disable_query_service,omitempty"`
 }
 
+// VtctldRefreshStateByShardRequest reloads tablet records for all tablets in
+// a shard via vtctld.
+type VtctldRefreshStateByShardRequest struct {
+	Organization string `json:"-"`
+	Database     string `json:"-"`
+	Branch       string `json:"-"`
+
+	Keyspace string   `json:"keyspace"`
+	Shard    string   `json:"shard"`
+	Cells    []string `json:"cells,omitempty"`
+}
+
 // VtctldStartWorkflowRequest is a request for starting a workflow.
 type VtctldStartWorkflowRequest struct {
 	Organization string `json:"-"`
@@ -187,6 +200,10 @@ func vtctldShardTabletControlAPIPath(org, db, branch string) string {
 	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "shard", "tablet-control")
 }
 
+func vtctldShardRefreshStateAPIPath(org, db, branch string) string {
+	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "shard", "refresh-state")
+}
+
 func (s *vtctldService) ListWorkflows(ctx context.Context, req *VtctldListWorkflowsRequest) (json.RawMessage, error) {
 	p := vtctldWorkflowsAPIPath(req.Organization, req.Database, req.Branch)
 	v := url.Values{}
@@ -318,6 +335,20 @@ func (s *vtctldService) SetShardTabletControl(ctx context.Context, req *VtctldSe
 	return resp.Data, nil
 }
 
+// RefreshStateByShard reloads tablet records for all tablets in a shard via vtctld.
+func (s *vtctldService) RefreshStateByShard(ctx context.Context, req *VtctldRefreshStateByShardRequest) (json.RawMessage, error) {
+	p := vtctldShardRefreshStateAPIPath(req.Organization, req.Database, req.Branch)
+	httpReq, err := s.client.newRequest(http.MethodPost, p, req)
+	if err != nil {
+		return nil, fmt.Errorf("error creating http request: %w", err)
+	}
+	resp := &vtctldDataResponse{}
+	if err := s.client.do(ctx, httpReq, resp); err != nil {
+		return nil, err
+	}
+	return resp.Data, nil
+}
+
 // CheckThrottler issues a throttler check against a single tablet.
 func (s *vtctldService) CheckThrottler(ctx context.Context, req *VtctldCheckThrottlerRequest) (json.RawMessage, error) {
 	p := path.Join(vtctldThrottlerAPIPath(req.Organization, req.Database, req.Branch), "check")

planetscale/vtctld_general_test.go

@@ -219,6 +219,42 @@ func TestVtctld_SetShardTabletControl(t *testing.T) {
 	c.Assert(string(data), qt.Equals, `{}`)
 }
 
+func TestVtctld_RefreshStateByShard(t *testing.T) {
+	c := qt.New(t)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, qt.Equals, http.MethodPost)
+		c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/vtctld/shard/refresh-state")
+
+		var body VtctldRefreshStateByShardRequest
+		err := json.NewDecoder(r.Body).Decode(&body)
+		c.Assert(err, qt.IsNil)
+		c.Assert(body.Keyspace, qt.Equals, "commerce")
+		c.Assert(body.Shard, qt.Equals, "-")
+		c.Assert(body.Cells, qt.DeepEquals, []string{"zone1"})
+
+		w.WriteHeader(200)
+		_, err = w.Write([]byte(`{"data":{}}`))
+		c.Assert(err, qt.IsNil)
+	}))
+	defer ts.Close()
+
+	client, err := NewClient(WithBaseURL(ts.URL))
+	c.Assert(err, qt.IsNil)
+
+	ctx := context.Background()
+	data, err := client.Vtctld.RefreshStateByShard(ctx, &VtctldRefreshStateByShardRequest{
+		Organization: "my-org",
+		Database:     "my-db",
+		Branch:       "my-branch",
+		Keyspace:     "commerce",
+		Shard:        "-",
+		Cells:        []string{"zone1"},
+	})
+	c.Assert(err, qt.IsNil)
+	c.Assert(string(data), qt.Equals, `{}`)
+}
+
 func TestVtctld_ListKeyspaces(t *testing.T) {
 	c := qt.New(t)