Add Vtctld.SetShardTabletControl client

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: maxenglander

planetscale/vtctld_general.go

@@ -17,6 +17,7 @@ type VtctldService interface {
 	ListKeyspaces(context.Context, *VtctldListKeyspacesRequest) (json.RawMessage, error)
 	GetRoutingRules(context.Context, *VtctldGetRoutingRulesRequest) (json.RawMessage, error)
 	GetShard(context.Context, *VtctldGetShardRequest) (json.RawMessage, error)
+	SetShardTabletControl(context.Context, *VtctldSetShardTabletControlRequest) (json.RawMessage, error)
 	ListTablets(context.Context, *ListBranchTabletsRequest) ([]*TabletGroup, error)
 	StartWorkflow(context.Context, *VtctldStartWorkflowRequest) (json.RawMessage, error)
 	StopWorkflow(context.Context, *VtctldStopWorkflowRequest) (json.RawMessage, error)
@@ -60,6 +61,22 @@ type VtctldGetShardRequest struct {
 	Shard        string `json:"-"`
 }
 
+// VtctldSetShardTabletControlRequest is a request for updating shard tablet
+// controls via vtctld.
+type VtctldSetShardTabletControlRequest struct {
+	Organization string `json:"-"`
+	Database     string `json:"-"`
+	Branch       string `json:"-"`
+
+	Keyspace            string   `json:"keyspace"`
+	Shard               string   `json:"shard"`
+	TabletType          string   `json:"tablet_type"`
+	Cells               []string `json:"cells,omitempty"`
+	DeniedTables        []string `json:"denied_tables,omitempty"`
+	Remove              *bool    `json:"remove,omitempty"`
+	DisableQueryService *bool    `json:"disable_query_service,omitempty"`
+}
+
 // VtctldStartWorkflowRequest is a request for starting a workflow.
 type VtctldStartWorkflowRequest struct {
 	Organization string `json:"-"`
@@ -166,6 +183,10 @@ func vtctldShardAPIPath(org, db, branch string) string {
 	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "shard")
 }
 
+func vtctldShardTabletControlAPIPath(org, db, branch string) string {
+	return path.Join(databaseBranchAPIPath(org, db, branch), "vtctld", "shard", "tablet-control")
+}
+
 func (s *vtctldService) ListWorkflows(ctx context.Context, req *VtctldListWorkflowsRequest) (json.RawMessage, error) {
 	p := vtctldWorkflowsAPIPath(req.Organization, req.Database, req.Branch)
 	v := url.Values{}
@@ -283,6 +304,20 @@ func (s *vtctldService) GetThrottlerStatus(ctx context.Context, req *VtctldGetTh
 	return resp.Data, nil
 }
 
+// SetShardTabletControl updates tablet controls on a shard via vtctld.
+func (s *vtctldService) SetShardTabletControl(ctx context.Context, req *VtctldSetShardTabletControlRequest) (json.RawMessage, error) {
+	p := vtctldShardTabletControlAPIPath(req.Organization, req.Database, req.Branch)
+	httpReq, err := s.client.newRequest(http.MethodPut, 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

@@ -176,6 +176,49 @@ func TestVtctld_GetShard(t *testing.T) {
 	c.Assert(string(data), qt.Equals, `{"keyspace":"commerce","name":"-"}`)
 }
 
+func TestVtctld_SetShardTabletControl(t *testing.T) {
+	c := qt.New(t)
+
+	remove := true
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, qt.Equals, http.MethodPut)
+		c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/vtctld/shard/tablet-control")
+
+		var body VtctldSetShardTabletControlRequest
+		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.TabletType, qt.Equals, "rdonly")
+		c.Assert(body.DeniedTables, qt.DeepEquals, []string{"customers"})
+		c.Assert(body.Remove, qt.Not(qt.IsNil))
+		c.Assert(*body.Remove, qt.Equals, true)
+
+		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.SetShardTabletControl(ctx, &VtctldSetShardTabletControlRequest{
+		Organization: "my-org",
+		Database:     "my-db",
+		Branch:       "my-branch",
+		Keyspace:     "commerce",
+		Shard:        "-",
+		TabletType:   "rdonly",
+		DeniedTables: []string{"customers"},
+		Remove:       &remove,
+	})
+	c.Assert(err, qt.IsNil)
+	c.Assert(string(data), qt.Equals, `{}`)
+}
+
 func TestVtctld_ListKeyspaces(t *testing.T) {
 	c := qt.New(t)