MoveTables: allow copying all tables in a single atomic copy phase cy…

…cle (#13137)

Signed-off-by: Rohit Nayak <rohit@planetscale.com>

Makefile

@@ -267,6 +267,7 @@ $(PROTO_GO_OUTS): minimaltools install_protoc-gen-go proto/*.proto
 		--go-vtproto_opt=features=marshal+unmarshal+size+pool+clone \
 		--go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/query.Row \
 		--go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/binlogdata.VStreamRowsResponse \
+		--go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/binlogdata.VStreamTablesResponse \
 		-I${PWD}/dist/vt-protoc-21.3/include:proto $(PROTO_SRCS)
 	cp -Rf vitess.io/vitess/go/vt/proto/* go/vt/proto
 	rm -rf vitess.io/vitess/go/vt/proto/

go/cmd/vtctldclient/command/movetables.go

@@ -89,6 +89,7 @@ See the --help output for each command for more details.`,
 		Aliases:               []string{"Create"},
 		Args:                  cobra.NoArgs,
 		PreRunE: func(cmd *cobra.Command, args []string) error {
+
 			// Either specific tables or the all tables flags are required.
 			if !cmd.Flags().Lookup("tables").Changed && !cmd.Flags().Lookup("all-tables").Changed {
 				return fmt.Errorf("tables or all-tables are required to specify which tables to move")
@@ -104,6 +105,27 @@ See the --help output for each command for more details.`,
 			if _, ok := binlogdatapb.OnDDLAction_value[strings.ToUpper(moveTablesCreateOptions.OnDDL)]; !ok {
 				return fmt.Errorf("invalid on-ddl value: %s", moveTablesCreateOptions.OnDDL)
 			}
+
+			checkAtomicCopyOptions := func() error {
+				var errors []string
+				if !moveTablesCreateOptions.AtomicCopy {
+					return nil
+				}
+				if !moveTablesCreateOptions.AllTables {
+					errors = append(errors, "atomic copy requires --all-tables.")
+				}
+				if len(moveTablesCreateOptions.IncludeTables) > 0 || len(moveTablesCreateOptions.ExcludeTables) > 0 {
+					errors = append(errors, "atomic copy does not support specifying tables.")
+				}
+				if len(errors) > 0 {
+					errors = append(errors, "Found options incompatible with atomic copy:")
+					return fmt.Errorf(strings.Join(errors, " "))
+				}
+				return nil
+			}
+			if err := checkAtomicCopyOptions(); err != nil {
+				return err
+			}
 			return nil
 		},
 		RunE: commandMoveTablesCreate,
@@ -235,6 +257,7 @@ var (
 		AutoStart                    bool
 		StopAfterCopy                bool
 		NoRoutingRules               bool
+		AtomicCopy                   bool
 	}{}
 	moveTablesSwitchTrafficOptions = struct {
 		Cells                     []string
@@ -285,6 +308,7 @@ func commandMoveTablesCreate(cmd *cobra.Command, args []string) error {
 		AutoStart:                 moveTablesCreateOptions.AutoStart,
 		StopAfterCopy:             moveTablesCreateOptions.StopAfterCopy,
 		NoRoutingRules:            moveTablesCreateOptions.NoRoutingRules,
+		AtomicCopy:                moveTablesCreateOptions.AtomicCopy,
 	}
 
 	resp, err := client.MoveTablesCreate(commandCtx, req)
@@ -527,6 +551,7 @@ func init() {
 	MoveTablesCreate.Flags().BoolVar(&moveTablesCreateOptions.AutoStart, "auto-start", true, "Start the MoveTables workflow after creating it")
 	MoveTablesCreate.Flags().BoolVar(&moveTablesCreateOptions.StopAfterCopy, "stop-after-copy", false, "Stop the MoveTables workflow after it's finished copying the existing rows and before it starts replicating changes")
 	MoveTablesCreate.Flags().BoolVar(&moveTablesCreateOptions.NoRoutingRules, "no-routing-rules", false, "(Advanced) Do not create routing rules while creating the workflow. See the reference documentation for limitations if you use this flag.")
+	MoveTablesCreate.Flags().BoolVar(&moveTablesCreateOptions.AtomicCopy, "atomic-copy", false, "(EXPERIMENTAL) A single copy phase is run for all tables from the source. Use this, for example, if your source keyspace has tables which use foreign key constraints.")
 	MoveTables.AddCommand(MoveTablesCreate)
 
 	MoveTables.AddCommand(MoveTablesShow)

go/cmd/vtctldclient/command/workflows.go

@@ -21,6 +21,9 @@ import (
 	"sort"
 	"strings"
 
+	"vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
+
 	"github.com/spf13/cobra"
 
 	"vitess.io/vitess/go/cmd/vtctldclient/cli"
@@ -223,6 +226,45 @@ func commandWorkflowDelete(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+func getWorkflow(keyspace, workflow string) (*vtctldatapb.GetWorkflowsResponse, error) {
+	resp, err := client.GetWorkflows(commandCtx, &vtctldatapb.GetWorkflowsRequest{
+		Keyspace: keyspace,
+		Workflow: workflow,
+	})
+	if err != nil {
+		return &vtctldatapb.GetWorkflowsResponse{}, err
+	}
+	return resp, nil
+}
+
+// canRestartWorkflow validates that, for an atomic copy workflow, none of the streams are still in the copy phase.
+// Since we copy all tables in a single snapshot, we cannot restart a workflow which broke before all tables were copied.
+func canRestartWorkflow(keyspace, workflow string) error {
+	resp, err := getWorkflow(keyspace, workflow)
+	if err != nil {
+		return err
+	}
+	if len(resp.Workflows) == 0 {
+		return fmt.Errorf("workflow %s not found", workflow)
+	}
+	if len(resp.Workflows) > 1 {
+		return vterrors.Errorf(vtrpc.Code_INTERNAL, "multiple results found for workflow %s", workflow)
+	}
+	wf := resp.Workflows[0]
+	if wf.WorkflowSubType != binlogdatapb.VReplicationWorkflowSubType_AtomicCopy.String() {
+		return nil
+	}
+	// If we're here, we have an atomic copy workflow.
+	for _, shardStream := range wf.ShardStreams {
+		for _, stream := range shardStream.Streams {
+			if len(stream.CopyStates) > 0 {
+				return fmt.Errorf("stream %d is still in the copy phase: can only start workflow %s if all streams have completed the copy phase.", stream.Id, workflow)
+			}
+		}
+	}
+	return nil
+}
+
 func commandWorkflowShow(cmd *cobra.Command, args []string) error {
 	cli.FinishedParsing(cmd)
 
@@ -312,6 +354,10 @@ func commandWorkflowUpdateState(cmd *cobra.Command, args []string) error {
 	var state binlogdatapb.VReplicationWorkflowState
 	switch strings.ToLower(cmd.Name()) {
 	case "start":
+		if err := canRestartWorkflow(workflowUpdateOptions.Workflow, workflowOptions.Keyspace); err != nil {
+			return err
+		}
+
 		state = binlogdatapb.VReplicationWorkflowState_Running
 	case "stop":
 		state = binlogdatapb.VReplicationWorkflowState_Stopped

go/test/endtoend/vreplication/cluster_test.go

@@ -89,6 +89,7 @@ type ClusterConfig struct {
 
 // VitessCluster represents all components within the test cluster
 type VitessCluster struct {
+	t             *testing.T
 	ClusterConfig *ClusterConfig
 	Name          string
 	Cells         map[string]*Cell
@@ -320,7 +321,7 @@ func init() {
 
 // NewVitessCluster starts a basic cluster with vtgate, vtctld and the topo
 func NewVitessCluster(t *testing.T, name string, cellNames []string, clusterConfig *ClusterConfig) *VitessCluster {
-	vc := &VitessCluster{Name: name, Cells: make(map[string]*Cell), ClusterConfig: clusterConfig}
+	vc := &VitessCluster{t: t, Name: name, Cells: make(map[string]*Cell), ClusterConfig: clusterConfig}
 	require.NotNil(t, vc)
 
 	vc.CleanupDataroot(t, true)

go/test/endtoend/vreplication/fk_config_test.go

@@ -0,0 +1,65 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package vreplication
+
+var (
+	initialFKSchema = `
+create table parent(id int, name varchar(128), primary key(id)) engine=innodb;
+create table child(id int, parent_id int, name varchar(128), primary key(id), foreign key(parent_id) references parent(id) on delete cascade) engine=innodb;
+`
+	initialFKData = `
+insert into parent values(1, 'parent1'), (2, 'parent2');
+insert into child values(1, 1, 'child11'), (2, 1, 'child21'), (3, 2, 'child32');`
+
+	initialFKSourceVSchema = `
+{
+  "tables": {
+	"parent": {},
+	"child": {}
+  }
+}
+`
+
+	initialFKTargetVSchema = `
+{
+  "sharded": true,
+  "vindexes": {
+    "reverse_bits": {
+      "type": "reverse_bits"
+    }
+  },
+  "tables": {
+    "parent": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "reverse_bits"
+        }
+      ]
+    },
+    "child": {
+      "column_vindexes": [
+        {
+          "column": "parent_id",
+          "name": "reverse_bits"
+        }
+      ]
+    }
+  }
+}
+`
+)