Merge pull request vitessio#8291 from tinyspeck/am_expandcells_handle…

…_aliases

[topo] Refactor `ExpandCells` to not error on valid aliases

Signed-off-by: Andrew Mason <[email protected]>

go/vt/topo/cell_info.go

@@ -17,12 +17,12 @@ limitations under the License.
 package topo
 
 import (
+	"context"
 	"path"
 	"strings"
 
-	"context"
-
 	"github.com/golang/protobuf/proto"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	"vitess.io/vitess/go/vt/vterrors"
 
@@ -176,37 +176,52 @@ func (ts *Server) GetKnownCells(ctx context.Context) ([]string, error) {
 // ExpandCells takes a comma-separated list of cells and returns an array of cell names
 // Aliases are expanded and an empty string returns all cells
 func (ts *Server) ExpandCells(ctx context.Context, cells string) ([]string, error) {
-	var err error
-	var outputCells []string
-	inputCells := strings.Split(cells, ",")
+	var (
+		err         error
+		inputCells  []string
+		outputCells = sets.NewString() // Use a set to dedupe if the input cells list includes an alias and a cell in that alias.
+	)
+
 	if cells == "" {
 		inputCells, err = ts.GetCellInfoNames(ctx)
 		if err != nil {
 			return nil, err
 		}
+	} else {
+		inputCells = strings.Split(cells, ",")
 	}
 
-	for _, cell := range inputCells {
-		cell2 := strings.TrimSpace(cell)
+	expandCell := func(ctx context.Context, cell string) error {
 		shortCtx, cancel := context.WithTimeout(ctx, *RemoteOperationTimeout)
 		defer cancel()
-		_, err := ts.GetCellInfo(shortCtx, cell2, false)
+
+		_, err := ts.GetCellInfo(shortCtx, cell, false /* strongRead */)
 		if err != nil {
-			// not a valid cell, check whether it is a cell alias
+			// Not a valid cell name. Check whether it is an alias.
 			shortCtx, cancel := context.WithTimeout(ctx, *RemoteOperationTimeout)
 			defer cancel()
-			alias, err2 := ts.GetCellsAlias(shortCtx, cell2, false)
-			// if we get an error, either cellAlias doesn't exist or it isn't a cell alias at all. Ignore and continue
-			if err2 == nil {
-				outputCells = append(outputCells, alias.Cells...)
-			}
-			if err != nil {
-				return nil, err
+
+			alias, err2 := ts.GetCellsAlias(shortCtx, cell, false /* strongRead */)
+			if err2 != nil {
+				return err // return the original err to indicate the cell does not exist
 			}
-		} else {
-			// valid cell, add it to our list
-			outputCells = append(outputCells, cell2)
+
+			// Expand the alias cells list into the final set.
+			outputCells.Insert(alias.Cells...)
+			return nil
 		}
+
+		// Valid cell.
+		outputCells.Insert(cell)
+		return nil
 	}
-	return outputCells, nil
+
+	for _, cell := range inputCells {
+		cell2 := strings.TrimSpace(cell)
+		if err := expandCell(ctx, cell2); err != nil {
+			return nil, err
+		}
+	}
+
+	return outputCells.List(), nil
 }

go/vt/topo/topotests/cell_info_test.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"context"
@@ -176,4 +177,52 @@ func TestExpandCells(t *testing.T) {
 		})
 	}
 
+	t.Run("aliases", func(t *testing.T) {
+		cells := []string{"cell1", "cell2", "cell3"}
+		ts := memorytopo.NewServer(cells...)
+		err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: cells})
+		require.NoError(t, err)
+
+		tests := []struct {
+			name      string
+			in        string
+			out       []string
+			shouldErr bool
+		}{
+			{
+				name: "alias only",
+				in:   "alias",
+				out:  []string{"cell1", "cell2", "cell3"},
+			},
+			{
+				name: "alias and cell in alias", // test deduping logic
+				in:   "alias,cell1",
+				out:  []string{"cell1", "cell2", "cell3"},
+			},
+			{
+				name: "just cells",
+				in:   "cell1",
+				out:  []string{"cell1"},
+			},
+			{
+				name:      "missing alias",
+				in:        "not_an_alias",
+				shouldErr: true,
+			},
+		}
+
+		for _, tt := range tests {
+			tt := tt
+			t.Run(tt.name, func(t *testing.T) {
+				expanded, err := ts.ExpandCells(ctx, tt.in)
+				if tt.shouldErr {
+					assert.Error(t, err)
+					return
+				}
+
+				require.NoError(t, err)
+				assert.ElementsMatch(t, expanded, tt.out)
+			})
+		}
+	})
 }