diff --git a/go/vt/topo/cell_info.go b/go/vt/topo/cell_info.go index ba59b656067..28bc65f86f3 100644 --- a/go/vt/topo/cell_info.go +++ b/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 } diff --git a/go/vt/topo/topotests/cell_info_test.go b/go/vt/topo/topotests/cell_info_test.go index 7854763a683..eaae838efab 100644 --- a/go/vt/topo/topotests/cell_info_test.go +++ b/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) + }) + } + }) }