Support polling a specific Vitess cell

timvaillancourt · timvaillancourt · commit 42d5227906a7 · 2020-07-20T16:54:44.000+02:00
diff --git a/pkg/config/mysql_config.go b/pkg/config/mysql_config.go
@@ -58,6 +58,7 @@ type MySQLConfigurationSettings struct {
 	HttpCheckPort        int      // port for HTTP check. -1 to disable.
 	HttpCheckPath        string   // If non-empty, requires HttpCheckPort
 	IgnoreHosts          []string // If non empty, substrings to indicate hosts to be ignored/skipped
+	VitessCell           string   // Name of the Vitess cell for polling tablet hosts
 
 	Clusters map[string](*MySQLClusterConfigurationSettings) // cluster name -> cluster config
 }
@@ -114,6 +115,9 @@ func (settings *MySQLConfigurationSettings) postReadAdjustments() error {
 		if len(clusterSettings.IgnoreHosts) == 0 {
 			clusterSettings.IgnoreHosts = settings.IgnoreHosts
 		}
+		if !clusterSettings.VitessSettings.IsEmpty() && clusterSettings.VitessSettings.Cell == "" {
+			clusterSettings.VitessSettings.Cell = settings.VitessCell
+		}
 	}
 	return nil
 }
diff --git a/pkg/config/vitess_config.go b/pkg/config/vitess_config.go
@@ -6,6 +6,7 @@ package config
 
 type VitessConfigurationSettings struct {
 	API         string
+	Cell        string
 	Keyspace    string
 	Shard       string
 	TimeoutSecs uint
diff --git a/pkg/vitess/api_client.go b/pkg/vitess/api_client.go
@@ -16,9 +16,10 @@ const defaultTimeout = time.Duration(5) * time.Second
 
 // Tablet represents information about a running instance of vttablet.
 type Tablet struct {
-	MysqlHostname string              `json:"mysql_hostname,omitempty"`
-	MysqlPort     int32               `json:"mysql_port,omitempty"`
-	Type          topodata.TabletType `json:"type,omitempty"`
+	Alias         *topodata.TabletAlias `json:"alias,omitempty"`
+	MysqlHostname string                `json:"mysql_hostname,omitempty"`
+	MysqlPort     int32                 `json:"mysql_port,omitempty"`
+	Type          topodata.TabletType   `json:"type,omitempty"`
 }
 
 // IsValidReplica returns a bool reflecting if a tablet type is REPLICA
@@ -41,8 +42,11 @@ func constructAPIURL(settings config.VitessConfigurationSettings) (url string) {
 }
 
 // filterReplicaTablets parses a list of tablets, returning replica tablets only
-func filterReplicaTablets(tablets []Tablet) (replicas []Tablet) {
+func filterReplicaTablets(settings config.VitessConfigurationSettings, tablets []Tablet) (replicas []Tablet) {
 	for _, tablet := range tablets {
+		if settings.Cell != "" && tablet.Alias.Cell != settings.Cell {
+			continue
+		}
 		if tablet.IsValidReplica() {
 			replicas = append(replicas, tablet)
 		}
@@ -72,5 +76,5 @@ func ParseTablets(settings config.VitessConfigurationSettings) (tablets []Tablet
 	}
 
 	err = json.Unmarshal(body, &tablets)
-	return filterReplicaTablets(tablets), err
+	return filterReplicaTablets(settings, tablets), err
 }
diff --git a/pkg/vitess/api_client_test.go b/pkg/vitess/api_client_test.go
@@ -18,27 +18,37 @@ func TestParseTablets(t *testing.T) {
 		case "/api/keyspace/test/tablets/00":
 			data, _ := json.Marshal([]Tablet{
 				{
+					Alias:         &topodata.TabletAlias{Cell: "ash1"},
 					MysqlHostname: "master",
 					Type:          topodata.TabletType_MASTER,
 				},
 				{
+					Alias:         &topodata.TabletAlias{Cell: "ac4"},
 					MysqlHostname: "replica",
 					Type:          topodata.TabletType_REPLICA,
 				},
 				{
+					Alias:         &topodata.TabletAlias{Cell: "va3"},
+					MysqlHostname: "replica",
+					Type:          topodata.TabletType_REPLICA,
+				},
+				{
+					Alias:         &topodata.TabletAlias{Cell: "ac4"},
 					MysqlHostname: "spare",
 					Type:          topodata.TabletType_SPARE,
 				},
 				{
+					Alias:         &topodata.TabletAlias{Cell: "va3"},
 					MysqlHostname: "batch",
 					Type:          topodata.TabletType_BATCH,
 				},
 				{
+					Alias:         &topodata.TabletAlias{Cell: "ac4"},
 					MysqlHostname: "backup",
 					Type:          topodata.TabletType_BACKUP,
 				},
 				{
-
+					Alias:         &topodata.TabletAlias{Cell: "ash1"},
 					MysqlHostname: "restore",
 					Type:          topodata.TabletType_RESTORE,
 				},
@@ -62,8 +72,8 @@ func TestParseTablets(t *testing.T) {
 			t.Fatalf("Expected no error, got %q", err)
 		}
 
-		if len(tablets) != 1 {
-			t.Fatalf("Expected 1 tablet, got %d", len(tablets))
+		if len(tablets) != 2 {
+			t.Fatalf("Expected 2 tablets, got %d", len(tablets))
 		}
 
 		if tablets[0].MysqlHostname != "replica" {
@@ -75,6 +85,26 @@ func TestParseTablets(t *testing.T) {
 		}
 	})
 
+	t.Run("with-cell", func(t *testing.T) {
+		tablets, err := ParseTablets(config.VitessConfigurationSettings{
+			API:      vitessApi.URL,
+			Cell:     "ac4",
+			Keyspace: "test",
+			Shard:    "00",
+		})
+		if err != nil {
+			t.Fatalf("Expected no error, got %q", err)
+		}
+
+		if len(tablets) > 1 {
+			t.Fatalf("Expected 1 tablet, got %d", len(tablets))
+		}
+
+		if tablets[0].Alias.Cell != "ac4" {
+			t.Fatalf("Expected vitess cell %s, got %s", "ac4", tablets[0].Alias.Cell)
+		}
+	})
+
 	t.Run("not-found", func(t *testing.T) {
 		tablets, err := ParseTablets(config.VitessConfigurationSettings{
 			API:      vitessApi.URL,
