Add ChangeTabletTags/ChangeTags RPCs

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>

go/cmd/vtctldclient/cli/tablets_test.go

@@ -1,3 +1,19 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package cli
 
 import (

go/vt/topotools/tablet.go

@@ -78,25 +78,33 @@ func isMapsEqual(a, b map[string]string) bool {
 // transitions need to be forced from time to time.
 //
 // If successful, the updated tablet record is returned.
-func ChangeTags(ctx context.Context, ts *topo.Server, tabletAlias *topodatapb.TabletAlias, tabletTags map[string]string, replace bool) (map[string]string, error) {
- var result map[string]string
+func ChangeTags(ctx context.Context, ts *topo.Server, tabletAlias *topodatapb.TabletAlias, tabletTags map[string]string, replace bool) (*topodatapb.Tablet, error) {
+ var result *topodatapb.Tablet
  _, err := ts.UpdateTabletFields(ctx, tabletAlias, func(tablet *topodatapb.Tablet) error {
  if replace && isMapsEqual(tablet.Tags, tabletTags) {
- result = tablet.Tags
+ result = tablet
  return topo.NewError(topo.NoUpdateNeeded, topoproto.TabletAliasString(tabletAlias))
  }
  if replace || tablet.Tags == nil {
  tablet.Tags = tabletTags
  } else {
+ var doUpdate bool
  for key, val := range tabletTags {
  if val == "" {
  delete(tablet.Tags, key)
+ doUpdate = true
  continue
  }
- tablet.Tags[key] = val
+ if tablet.Tags[key] != val {
+ tablet.Tags[key] = val
+ doUpdate = true
+ }
+ }
+ if !doUpdate {
+ return topo.NewError(topo.NoUpdateNeeded, topoproto.TabletAliasString(tabletAlias))
  }
  }
- result = tablet.Tags
+ result = tablet
  return nil
  })
  if err != nil {

go/vt/vtcombo/tablet_map.go

@@ -764,12 +764,18 @@ func (itmc *internalTabletManagerClient) SetReadWrite(ctx context.Context, table
  return fmt.Errorf("not implemented in vtcombo")
 }
 
-func (itmc *internalTabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (map[string]string, error) {
+func (itmc *internalTabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (*tabletmanagerdatapb.ChangeTagsResponse, error) {
  t, ok := tabletMap[tablet.Alias.Uid]
  if !ok {
  return nil, fmt.Errorf("tmclient: cannot find tablet %v", tablet.Alias.Uid)
  }
- return t.tm.ChangeTags(ctx, tabletTags, replace)
+ afterTags, err := t.tm.ChangeTags(ctx, tabletTags, replace)
+ if err != nil {
+ return nil, err
+ }
+ return &tabletmanagerdatapb.ChangeTagsResponse{
+ Tags: afterTags,
+ }, nil
 }
 
 func (itmc *internalTabletManagerClient) ChangeType(ctx context.Context, tablet *topodatapb.Tablet, dbType topodatapb.TabletType, semiSync bool) error {

go/vt/vtctl/grpcvtctldserver/server.go

@@ -518,16 +518,16 @@ func (s *VtctldServer) ChangeTabletTags(ctx context.Context, req *vtctldatapb.Ch
 
  span.Annotate("before_tablet_tags", tablet.Tags)
 
- afterTags, err := s.tmc.ChangeTags(ctx, tablet.Tablet, req.Tags, req.Replace)
+ changeTagsResp, err := s.tmc.ChangeTags(ctx, tablet.Tablet, req.Tags, req.Replace)
  if err != nil {
  return nil, err
  }
 
- span.Annotate("after_tablet_tags", afterTags)
+ span.Annotate("after_tablet_tags", changeTagsResp.Tags)
 
  return &vtctldatapb.ChangeTabletTagsResponse{
  BeforeTags: tablet.Tags,
- AfterTags: afterTags,
+ AfterTags: changeTagsResp.Tags,
  }, nil
 }
 

go/vt/vtctl/grpcvtctldserver/server_test.go

@@ -1213,6 +1213,206 @@ func TestChangeTabletTags(t *testing.T) {
  })
 }
 
+func TestChangeTabletTags(t *testing.T) {
+ t.Parallel()
+
+ tests := []struct {
+ name string
+ cells []string
+ tablet *topodatapb.Tablet
+ req *vtctldatapb.ChangeTabletTagsRequest
+ expected *vtctldatapb.ChangeTabletTagsResponse
+ shouldErr bool
+ }{
+ {
+ name: "success",
+ cells: []string{"zone1"},
+ tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Keyspace: "ks",
+ Shard: "0",
+ },
+ req: &vtctldatapb.ChangeTabletTagsRequest{
+ TabletAlias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Tags: map[string]string{
+ "test": t.Name(),
+ },
+ },
+ expected: &vtctldatapb.ChangeTabletTagsResponse{
+ AfterTags: map[string]string{
+ "test": t.Name(),
+ },
+ },
+ shouldErr: false,
+ },
+ {
+ name: "success with existing",
+ cells: []string{"zone1"},
+ tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Keyspace: "ks",
+ Shard: "0",
+ Tags: map[string]string{
+ "delete": "me",
+ "hello": "world!",
+ },
+ },
+ req: &vtctldatapb.ChangeTabletTagsRequest{
+ TabletAlias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Tags: map[string]string{
+ "delete": "",
+ "test": t.Name(),
+ },
+ },
+ expected: &vtctldatapb.ChangeTabletTagsResponse{
+ BeforeTags: map[string]string{
+ "delete": "me",
+ "hello": "world!",
+ },
+ AfterTags: map[string]string{
+ "hello": "world!",
+ "test": t.Name(),
+ },
+ },
+ shouldErr: false,
+ },
+ {
+ name: "success with replace",
+ cells: []string{"zone1"},
+ tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Keyspace: "ks",
+ Shard: "0",
+ Tags: map[string]string{
+ "hello": "world!",
+ },
+ },
+ req: &vtctldatapb.ChangeTabletTagsRequest{
+ TabletAlias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Tags: map[string]string{
+ "test": t.Name(),
+ },
+ Replace: true,
+ },
+ expected: &vtctldatapb.ChangeTabletTagsResponse{
+ BeforeTags: map[string]string{
+ "hello": "world!",
+ },
+ AfterTags: map[string]string{
+ "test": t.Name(),
+ },
+ },
+ shouldErr: false,
+ },
+ {
+ name: "tablet not found",
+ cells: []string{"zone1"},
+ tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 200,
+ },
+ Keyspace: "ks",
+ Shard: "0",
+ },
+ req: &vtctldatapb.ChangeTabletTagsRequest{
+ TabletAlias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Tags: map[string]string{
+ "test": t.Name(),
+ },
+ },
+ expected: nil,
+ shouldErr: true,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ t.Parallel()
+
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ ts := memorytopo.NewServer(ctx, tt.cells...)
+ vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &testutil.TabletManagerClient{
+ TopoServer: ts,
+ }, func(ts *topo.Server) vtctlservicepb.VtctldServer {
+ return NewVtctldServer(vtenv.NewTestEnv(), ts)
+ })
+
+ testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{
+ AlsoSetShardPrimary: true,
+ }, tt.tablet)
+
+ resp, err := vtctld.ChangeTabletTags(ctx, tt.req)
+ if tt.shouldErr {
+ assert.Error(t, err)
+ return
+ }
+ assert.NoError(t, err)
+ utils.MustMatch(t, tt.expected, resp)
+
+ tablet, err := ts.GetTablet(ctx, tt.req.TabletAlias)
+ assert.NoError(t, err)
+ utils.MustMatch(t, resp.AfterTags, tablet.Tags)
+ })
+ }
+
+ t.Run("tabletmanager failure", func(t *testing.T) {
+ t.Parallel()
+
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ ts := memorytopo.NewServer(ctx, "zone1")
+ vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &testutil.TabletManagerClient{
+ TopoServer: nil,
+ }, func(ts *topo.Server) vtctlservicepb.VtctldServer {
+ return NewVtctldServer(vtenv.NewTestEnv(), ts)
+ })
+
+ testutil.AddTablet(ctx, t, ts, &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Keyspace: "ks",
+ Shard: "0",
+ Type: topodatapb.TabletType_REPLICA,
+ }, nil)
+
+ _, err := vtctld.ChangeTabletTags(ctx, &vtctldatapb.ChangeTabletTagsRequest{
+ TabletAlias: &topodatapb.TabletAlias{
+ Cell: "zone1",
+ Uid: 100,
+ },
+ Tags: map[string]string{
+ "test": t.Name(),
+ },
+ })
+ assert.Error(t, err)
+ })
+}
+
 func TestChangeTabletType(t *testing.T) {
  t.Parallel()
 

go/vt/vtctl/grpcvtctldserver/testutil/test_tmclient.go

@@ -185,11 +185,11 @@ type TabletManagerClient struct {
  ErrorAfter time.Duration
  }
  // keyed by tablet alias.
- ChangeTabletTagsResult map[string]struct {
- Tags map[string]string
- Error error
+ ChangeTagsResult map[string]struct {
+ Response *tabletmanagerdatapb.ChangeTagsResponse
+ Error  error
  }
- ChangeTabletTagsDelays map[string]time.Duration
+ ChangeTagsDelays  map[string]time.Duration
  ChangeTabletTypeResult map[string]error
  // keyed by tablet alias.
  DemotePrimaryDelays map[string]time.Duration
@@ -449,11 +449,11 @@ func (fake *TabletManagerClient) Backup(ctx context.Context, tablet *topodatapb.
 }
 
 // ChangeTags is part of the tmclient.TabletManagerClient interface.
-func (fake *TabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (map[string]string, error) {
+func (fake *TabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (*tabletmanagerdatapb.ChangeTagsResponse, error) {
  key := topoproto.TabletAliasString(tablet.Alias)
 
- if fake.ChangeTabletTagsDelays != nil {
- if delay, ok := fake.ChangeTabletTagsDelays[key]; ok {
+ if fake.ChangeTagsDelays != nil {
+ if delay, ok := fake.ChangeTagsDelays[key]; ok {
  select {
  case <-ctx.Done():
  return nil, ctx.Err()
@@ -463,15 +463,22 @@ func (fake *TabletManagerClient) ChangeTags(ctx context.Context, tablet *topodat
  }
  }
 
- if result, ok := fake.ChangeTabletTagsResult[key]; ok {
- return result.Tags, result.Error
+ if result, ok := fake.ChangeTagsResult[key]; ok {
+ return result.Response, result.Error
  }
 
  if fake.TopoServer == nil {
  return nil, assert.AnError
  }
 
- return topotools.ChangeTags(ctx, fake.TopoServer, tablet.Alias, tabletTags, replace)
+ tablet, err := topotools.ChangeTags(ctx, fake.TopoServer, tablet.Alias, tabletTags, replace)
+ if err != nil {
+ return nil, err
+ }
+
+ return &tabletmanagerdatapb.ChangeTagsResponse{
+ Tags: tablet.Tags,
+ }, nil
 }
 
 // ChangeType is part of the tmclient.TabletManagerClient interface.

go/vt/vttablet/faketmclient/fake_client.go

@@ -35,6 +35,7 @@ import (
  logutilpb "vitess.io/vitess/go/vt/proto/logutil"
  querypb "vitess.io/vitess/go/vt/proto/query"
  replicationdatapb "vitess.io/vitess/go/vt/proto/replicationdata"
+ "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
  tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
  topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
@@ -121,8 +122,8 @@ func (client *FakeTabletManagerClient) SetReadWrite(ctx context.Context, tablet
 }
 
 // ChangeTags is part of the tmclient.TabletManagerClient interface.
-func (client *FakeTabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (map[string]string, error) {
- return map[string]string{}, nil
+func (client *FakeTabletManagerClient) ChangeTags(ctx context.Context, tablet *topodatapb.Tablet, tabletTags map[string]string, replace bool) (*tabletmanagerdata.ChangeTagsResponse, error) {
+ return &tabletmanagerdata.ChangeTagsResponse{}, nil
 }
 
 // ChangeType is part of the tmclient.TabletManagerClient interface.