Expose the --tablet_types_to_wait flag in vtcombo.

Signed-off-by: Arthur Schreiber <[email protected]>
github · Oct 15, 2024 · ce7168c · ce7168c
1 parent 2ffd196
commit ce7168c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/go/cmd/vtcombo/cli/main.go b/go/cmd/vtcombo/cli/main.go
@@ -40,6 +40,7 @@ import (
 	"vitess.io/vitess/go/vt/srvtopo"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
+	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/topotools"
 	"vitess.io/vitess/go/vt/vtcombo"
 	"vitess.io/vitess/go/vt/vtctld"
@@ -77,9 +78,10 @@ In particular, it contains:
 	plannerName           string
 	vschemaPersistenceDir string
 
-	tpb             vttestpb.VTTestTopology
-	ts              *topo.Server
-	resilientServer *srvtopo.ResilientServer
+	tpb               vttestpb.VTTestTopology
+	ts                *topo.Server
+	resilientServer   *srvtopo.ResilientServer
+	tabletTypesToWait []topodatapb.TabletType
 
 	env *vtenv.Environment
 )
@@ -113,6 +115,7 @@ func init() {
 
 	Main.Flags().Var(vttest.TextTopoData(&tpb), "proto_topo", "vttest proto definition of the topology, encoded in compact text format. See vttest.proto for more information.")
 	Main.Flags().Var(vttest.JSONTopoData(&tpb), "json_topo", "vttest proto definition of the topology, encoded in json format. See vttest.proto for more information.")
+	Main.Flags().Var((*topoproto.TabletTypeListFlag)(&tabletTypesToWait), "tablet_types_to_wait", "Wait till connected for specified tablet types during Gateway initialization. Should be provided as a comma-separated set of tablet types.")
 
 	// We're going to force the value later, so don't even bother letting the
 	// user know about this flag.
@@ -294,19 +297,30 @@ func run(cmd *cobra.Command, args []string) (err error) {
 
 	// vtgate configuration and init
 	resilientServer = srvtopo.NewResilientServer(context.Background(), ts, "ResilientSrvTopoServer")
-	tabletTypesToWait := []topodatapb.TabletType{
-		topodatapb.TabletType_PRIMARY,
-		topodatapb.TabletType_REPLICA,
-		topodatapb.TabletType_RDONLY,
+
+	tabletTypes := make([]topodatapb.TabletType, 0, 1)
+	if len(tabletTypesToWait) != 0 {
+		for _, tt := range tabletTypesToWait {
+			if topoproto.IsServingType(tt) {
+				tabletTypes = append(tabletTypes, tt)
+			}
+		}
+
+		if len(tabletTypes) == 0 {
+			log.Exitf("tablet_types_to_wait should contain at least one serving tablet type")
+		}
+	} else {
+		tabletTypes = append(tabletTypes, topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA, topodatapb.TabletType_RDONLY)
 	}
+
 	plannerVersion, _ := plancontext.PlannerNameToVersion(plannerName)
 
 	vtgate.QueryLogHandler = "/debug/vtgate/querylog"
 	vtgate.QueryLogzHandler = "/debug/vtgate/querylogz"
 	vtgate.QueryzHandler = "/debug/vtgate/queryz"
 
 	// pass nil for healthcheck, it will get created
-	vtg := vtgate.Init(context.Background(), env, nil, resilientServer, tpb.Cells[0], tabletTypesToWait, plannerVersion)
+	vtg := vtgate.Init(context.Background(), env, nil, resilientServer, tpb.Cells[0], tabletTypes, plannerVersion)
 
 	// vtctld configuration and init
 	err = vtctld.InitVtctld(env, ts)

diff --git a/go/flags/endtoend/vtcombo.txt b/go/flags/endtoend/vtcombo.txt
@@ -349,6 +349,7 @@ Flags:
       --tablet_manager_protocol string                                   Protocol to use to make tabletmanager RPCs to vttablets. (default "grpc")
       --tablet_refresh_interval duration                                 Tablet refresh interval. (default 1m0s)
       --tablet_refresh_known_tablets                                     Whether to reload the tablet's address/port map from topo in case they change. (default true)
+      --tablet_types_to_wait strings                                     Wait till connected for specified tablet types during Gateway initialization. Should be provided as a comma-separated set of tablet types.
       --tablet_url_template string                                       Format string describing debug tablet url formatting. See getTabletDebugURL() for how to customize this. (default "http://{{ "{{.GetTabletHostPort}}" }}")
       --throttle_tablet_types string                                     Comma separated VTTablet types to be considered by the throttler. default: 'replica'. example: 'replica,rdonly'. 'replica' always implicitly included (default "replica")
       --topo_consul_lock_delay duration                                  LockDelay for consul session. (default 15s)