diff --git a/go/cmd/vtctldclient/command/vreplication/common/utils.go b/go/cmd/vtctldclient/command/vreplication/common/utils.go index be3e52f9ab9..da6e3329579 100644 --- a/go/cmd/vtctldclient/command/vreplication/common/utils.go +++ b/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.") diff --git a/go/cmd/vtctldclient/command/vreplication/common/utils_test.go b/go/cmd/vtctldclient/command/vreplication/common/utils_test.go index 77672d82350..0dc179060d6 100644 --- a/go/cmd/vtctldclient/command/vreplication/common/utils_test.go +++ b/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,18 +77,50 @@ 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 @@ -77,6 +128,26 @@ func TestParseAndValidateCreateOptions(t *testing.T) { 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) +}