vtctldclient: Apply tablet type filtering for keyspace+shard in GetTa…

…blets (#14467) Signed-off-by: Matt Lord <mattalord@gmail.com>
vitessio · Nov 5, 2023 · cfe88e6 · cfe88e6
1 parent 3ad0171
commit cfe88e6
Show file tree

Hide file tree

Showing 2 changed files with 150 additions and 2 deletions.
diff --git a/go/vt/vtctl/grpcvtctldserver/server.go b/go/vt/vtctl/grpcvtctldserver/server.go
@@ -1945,6 +1945,9 @@ func (s *VtctldServer) GetTablets(ctx context.Context, req *vtctldatapb.GetTable
 	defer panicHandler(&err)
 
 	span.Annotate("cells", strings.Join(req.Cells, ","))
+	if req.Keyspace != "" {
+		span.Annotate("keyspace", req.Keyspace)
+	}
 	if req.TabletType != topodatapb.TabletType_UNKNOWN {
 		span.Annotate("tablet_type", topodatapb.TabletType_name[int32(req.TabletType)])
 	}
@@ -1978,7 +1981,6 @@ func (s *VtctldServer) GetTablets(ctx context.Context, req *vtctldatapb.GetTable
 			err = fmt.Errorf("GetTabletMap(%v) failed: %w", req.TabletAliases, err)
 		}
 	case req.Keyspace != "" && req.Shard != "":
-		span.Annotate("keyspace", req.Keyspace)
 		span.Annotate("shard", req.Shard)
 
 		tabletMap, err = s.ts.GetTabletMapForShard(ctx, req.Keyspace, req.Shard)
@@ -2016,6 +2018,9 @@ func (s *VtctldServer) GetTablets(ctx context.Context, req *vtctldatapb.GetTable
 
 		tablets := make([]*topodatapb.Tablet, 0, len(tabletMap))
 		for _, ti := range tabletMap {
+			if req.TabletType != 0 && ti.Type != req.TabletType {
+				continue
+			}
 			adjustTypeForStalePrimary(ti, truePrimaryTimestamp)
 			tablets = append(tablets, ti.Tablet)
 		}

diff --git a/go/vt/vtctl/grpcvtctldserver/server_test.go b/go/vt/vtctl/grpcvtctldserver/server_test.go
@@ -39,6 +39,7 @@ import (
 	"vitess.io/vitess/go/test/utils"
 	hk "vitess.io/vitess/go/vt/hook"
 	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
+	"vitess.io/vitess/go/vt/proto/vttime"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
@@ -57,7 +58,6 @@ import (
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
 	vtctlservicepb "vitess.io/vitess/go/vt/proto/vtctlservice"
-	"vitess.io/vitess/go/vt/proto/vttime"
 )
 
 func init() {
@@ -6822,6 +6822,149 @@ func TestGetTablets(t *testing.T) {
 			expected:  []*topodatapb.Tablet{},
 			shouldErr: false,
 		},
+		{
+			name:  "tablet type filter",
+			cells: []string{"cell1"},
+			tablets: []*topodatapb.Tablet{
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  100,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  101,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_REPLICA,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  200,
+					},
+					Keyspace: "ks1",
+					Shard:    "80-",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  201,
+					},
+					Keyspace: "ks1",
+					Shard:    "80-",
+					Type:     topodatapb.TabletType_REPLICA,
+				},
+			},
+			req: &vtctldatapb.GetTabletsRequest{
+				TabletType: topodatapb.TabletType_PRIMARY,
+			},
+			expected: []*topodatapb.Tablet{
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  100,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  200,
+					},
+					Keyspace: "ks1",
+					Shard:    "80-",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+			},
+			shouldErr: false,
+		},
+		{
+			name:  "keyspace, shard, and tablet type filter",
+			cells: []string{"cell1"},
+			tablets: []*topodatapb.Tablet{
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  100,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  101,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_REPLICA,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  200,
+					},
+					Keyspace: "ks1",
+					Shard:    "80-",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  201,
+					},
+					Keyspace: "ks1",
+					Shard:    "80-",
+					Type:     topodatapb.TabletType_REPLICA,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  300,
+					},
+					Keyspace: "ks2",
+					Shard:    "-",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  301,
+					},
+					Keyspace: "ks2",
+					Shard:    "-",
+					Type:     topodatapb.TabletType_REPLICA,
+				},
+			},
+			req: &vtctldatapb.GetTabletsRequest{
+				Keyspace:   "ks1",
+				Shard:      "-80",
+				TabletType: topodatapb.TabletType_PRIMARY,
+			},
+			expected: []*topodatapb.Tablet{
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "cell1",
+						Uid:  100,
+					},
+					Keyspace: "ks1",
+					Shard:    "-80",
+					Type:     topodatapb.TabletType_PRIMARY,
+				},
+			},
+			shouldErr: false,
+		},
 	}
 
 	for _, tt := range tests {