Support many cells

timvaillancourt · timvaillancourt · commit 977db88869d6 · 2020-07-21T14:58:26.000+02:00
diff --git a/pkg/config/mysql_config.go b/pkg/config/mysql_config.go
@@ -58,7 +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
+	VitessCells          []string // Name of the Vitess cells for polling tablet hosts
 
 	Clusters map[string](*MySQLClusterConfigurationSettings) // cluster name -> cluster config
 }
@@ -115,8 +115,8 @@ 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
+		if !clusterSettings.VitessSettings.IsEmpty() && len(clusterSettings.VitessSettings.Cells) < 1 {
+			clusterSettings.VitessSettings.Cells = settings.VitessCells
 		}
 	}
 	return nil
diff --git a/pkg/config/vitess_config.go b/pkg/config/vitess_config.go
@@ -6,7 +6,7 @@ package config
 
 type VitessConfigurationSettings struct {
 	API         string
-	Cell        string
+	Cells       []string
 	Keyspace    string
 	Shard       string
 	TimeoutSecs uint
diff --git a/pkg/vitess/api_client.go b/pkg/vitess/api_client.go
@@ -22,6 +22,19 @@ type Tablet struct {
 	Type          topodata.TabletType   `json:"type,omitempty"`
 }
 
+// IsValidCell returns a bool reflecting if a tablet type is in a valid cell
+func (t Tablet) HasValidCell(validCells []string) bool {
+	if len(validCells) == 0 {
+		return true
+	}
+	for _, cell := range validCells {
+		if t.Alias.GetCell() == strings.TrimSpace(cell) {
+			return true
+		}
+	}
+	return false
+}
+
 // IsValidReplica returns a bool reflecting if a tablet type is REPLICA
 func (t Tablet) IsValidReplica() bool {
 	return t.Type == topodata.TabletType_REPLICA
@@ -44,11 +57,7 @@ func constructAPIURL(settings config.VitessConfigurationSettings) (url string) {
 // filterReplicaTablets parses a list of tablets, returning replica tablets only
 func filterReplicaTablets(settings config.VitessConfigurationSettings, tablets []Tablet) (replicas []Tablet) {
 	for _, tablet := range tablets {
-		expectedCell := strings.TrimSpace(settings.Cell)
-		if expectedCell != "" && tablet.Alias.GetCell() != expectedCell {
-			continue
-		}
-		if tablet.IsValidReplica() {
+		if tablet.HasValidCell(settings.Cells) && tablet.IsValidReplica() {
 			replicas = append(replicas, tablet)
 		}
 	}
diff --git a/pkg/vitess/api_client_test.go b/pkg/vitess/api_client_test.go
@@ -91,7 +91,7 @@ func TestParseTablets(t *testing.T) {
 	t.Run("with-cell", func(t *testing.T) {
 		tablets, err := ParseTablets(config.VitessConfigurationSettings{
 			API:      vitessApi.URL,
-			Cell:     "cell2",
+			Cells:    []string{"cell2"},
 			Keyspace: "test",
 			Shard:    "00",
 		})
@@ -109,7 +109,6 @@ func TestParseTablets(t *testing.T) {
 		if tablets[0].Alias.Cell != "cell2" {
 			t.Fatalf("Expected vitess cell %s, got %s", "cell2", tablets[0].Alias.Cell)
 		}
-
 	})
 
 	t.Run("not-found", func(t *testing.T) {
