VReplication: Add --all-cells flag to create sub-commands (#14341)

Signed-off-by: Matt Lord <[email protected]>

go/cmd/vtctldclient/command/vreplication/common/utils.go

@@ -26,6 +26,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"vitess.io/vitess/go/cmd/vtctldclient/cli"
+	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
 	"vitess.io/vitess/go/vt/vtctl/workflow"
@@ -56,6 +57,7 @@ var (
 
 	CreateOptions = struct {
 		Cells                        []string
+		AllCells                     bool
 		TabletTypes                  []topodatapb.TabletType
 		TabletTypesInPreferenceOrder bool
 		OnDDL                        string
@@ -98,12 +100,28 @@ func GetCommandCtx() context.Context {
 	return commandCtx
 }
 
-func ParseCells(cmd *cobra.Command) {
-	if cmd.Flags().Lookup("cells").Changed { // Validate the provided value(s)
+func ParseCells(cmd *cobra.Command) error {
+	cf := cmd.Flags().Lookup("cells")
+	af := cmd.Flags().Lookup("all-cells")
+	if cf != nil && cf.Changed && af != nil && af.Changed {
+		return fmt.Errorf("cannot specify both --cells and --all-cells")
+	}
+	if cf.Changed { // Validate the provided value(s)
 		for i, cell := range CreateOptions.Cells { // Which only means trimming whitespace
 			CreateOptions.Cells[i] = strings.TrimSpace(cell)
 		}
 	}
+	if CreateOptions.AllCells { // Use all current cells
+		ctx, cancel := context.WithTimeout(commandCtx, topo.RemoteOperationTimeout)
+		defer cancel()
+		resp, err := client.GetCellInfoNames(ctx, &vtctldatapb.GetCellInfoNamesRequest{})
+		if err != nil {
+			return fmt.Errorf("failed to get current cells: %v", err)
+		}
+		CreateOptions.Cells = make([]string, len(resp.Names))
+		copy(CreateOptions.Cells, resp.Names)
+	}
+	return nil
 }
 
 func ParseTabletTypes(cmd *cobra.Command) error {
@@ -130,7 +148,9 @@ func ParseAndValidateCreateOptions(cmd *cobra.Command) error {
 	if err := validateOnDDL(cmd); err != nil {
 		return err
 	}
-	ParseCells(cmd)
+	if err := ParseCells(cmd); err != nil {
+		return err
+	}
 	if err := ParseTabletTypes(cmd); err != nil {
 		return err
 	}
@@ -192,6 +212,7 @@ func AddCommonFlags(cmd *cobra.Command) {
 
 func AddCommonCreateFlags(cmd *cobra.Command) {
 	cmd.Flags().StringSliceVarP(&CreateOptions.Cells, "cells", "c", nil, "Cells and/or CellAliases to copy table data from.")
+	cmd.Flags().BoolVarP(&CreateOptions.AllCells, "all-cells", "a", false, "Copy table data from any existing cell.")
 	cmd.Flags().Var((*topoproto.TabletTypeListFlag)(&CreateOptions.TabletTypes), "tablet-types", "Source tablet types to replicate table data from (e.g. PRIMARY,REPLICA,RDONLY).")
 	cmd.Flags().BoolVar(&CreateOptions.TabletTypesInPreferenceOrder, "tablet-types-in-preference-order", true, "When performing source tablet selection, look for candidates in the type order as they are listed in the tablet-types flag.")
 	cmd.Flags().StringVar(&CreateOptions.OnDDL, "on-ddl", onDDLDefault, "What to do when DDL is encountered in the VReplication stream. Possible values are IGNORE, STOP, EXEC, and EXEC_IGNORE.")

go/cmd/vtctldclient/command/vreplication/common/utils_test.go

@@ -14,19 +14,38 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package common
+package common_test
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/cmd/vtctldclient/command"
+	"vitess.io/vitess/go/cmd/vtctldclient/command/vreplication/common"
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
+	"vitess.io/vitess/go/vt/vtctl/grpcvtctldserver"
+	"vitess.io/vitess/go/vt/vtctl/localvtctldclient"
+	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
+	"vitess.io/vitess/go/vt/vttablet/tmclient"
 )
 
 func TestParseAndValidateCreateOptions(t *testing.T) {
+	common.SetCommandCtx(context.Background())
+	ctx, cancel := context.WithTimeout(common.GetCommandCtx(), 60*time.Second)
+	defer cancel()
+	cells := []string{"zone1", "zone2", "zone3"}
+	SetupLocalVtctldClient(t, ctx, cells...)
+
 	tests := []struct {
-		name    string
-		setFunc func(*cobra.Command) error
-		wantErr bool
+		name      string
+		setFunc   func(*cobra.Command) error
+		wantErr   bool
+		checkFunc func()
 	}{
 		{
 			name: "invalid tablet type",
@@ -58,25 +77,77 @@ func TestParseAndValidateCreateOptions(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "cells and all-cells",
+			setFunc: func(cmd *cobra.Command) error {
+				cellsFlag := cmd.Flags().Lookup("cells")
+				allCellsFlag := cmd.Flags().Lookup("all-cells")
+				if err := cellsFlag.Value.Set("cella"); err != nil {
+					return err
+				}
+				cellsFlag.Changed = true
+				if err := allCellsFlag.Value.Set("true"); err != nil {
+					return err
+				}
+				allCellsFlag.Changed = true
+				return nil
+			},
+			wantErr: true,
+		},
+		{
+			name: "all cells",
+			setFunc: func(cmd *cobra.Command) error {
+				allCellsFlag := cmd.Flags().Lookup("all-cells")
+				if err := allCellsFlag.Value.Set("true"); err != nil {
+					return err
+				}
+				allCellsFlag.Changed = true
+				return nil
+			},
+			wantErr: false,
+			checkFunc: func() {
+				require.Equal(t, cells, common.CreateOptions.Cells)
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			cmd := &cobra.Command{}
-			AddCommonCreateFlags(cmd)
+			common.AddCommonCreateFlags(cmd)
 			test := func() error {
 				if tt.setFunc != nil {
 					if err := tt.setFunc(cmd); err != nil {
 						return err
 					}
 				}
-				if err := ParseAndValidateCreateOptions(cmd); err != nil {
+				if err := common.ParseAndValidateCreateOptions(cmd); err != nil {
 					return err
 				}
 				return nil
 			}
 			if err := test(); (err != nil) != tt.wantErr {
 				t.Errorf("ParseAndValidateCreateOptions() error = %v, wantErr %t", err, tt.wantErr)
 			}
+			if tt.checkFunc != nil {
+				tt.checkFunc()
+			}
 		})
 	}
 }
+
+// SetupLocalVtctldClient sets up a local or internal VtctldServer and
+// VtctldClient for tests. It uses a memorytopo instance which contains
+// the cells provided.
+func SetupLocalVtctldClient(t *testing.T, ctx context.Context, cells ...string) {
+	ts, factory := memorytopo.NewServerAndFactory(ctx, cells...)
+	topo.RegisterFactory("test", factory)
+	tmclient.RegisterTabletManagerClientFactory("grpc", func() tmclient.TabletManagerClient {
+		return nil
+	})
+	vtctld := grpcvtctldserver.NewVtctldServer(ts)
+	localvtctldclient.SetServer(vtctld)
+	command.VtctldClientProtocol = "local"
+	client, err := vtctldclient.New(command.VtctldClientProtocol, "")
+	require.NoError(t, err, "failed to create local vtctld client which uses an internal vtctld server")
+	common.SetClient(client)
+}