VSCopy: Enable to copy from all shards in either a specified keyspace or all keyspaces

go/vt/vtgate/endtoend/main_test.go

@@ -52,6 +52,12 @@ create table t1_copy_basic(
 	primary key(id1)
 ) Engine=InnoDB;
 
+create table t1_copy_all(
+	id1 bigint,
+	id2 bigint,
+	primary key(id1)
+) Engine=InnoDB;
+
 create table t1_copy_resume(
 	id1 bigint,
 	id2 bigint,
@@ -152,6 +158,12 @@ create table t1_sharded(
 					Name:   "hash",
 				}},
 			},
+			"t1_copy_all": {
+				ColumnVindexes: []*vschemapb.ColumnVindex{{
+					Column: "id1",
+					Name:   "hash",
+				}},
+			},
 			"t1_copy_resume": {
 				ColumnVindexes: []*vschemapb.ColumnVindex{{
 					Column: "id1",
@@ -219,6 +231,31 @@ create table t1_sharded(
 			},
 		},
 	}
+
+	schema2 = `
+create table t1_copy_all_ks2(
+	id1 bigint,
+	id2 bigint,
+	primary key(id1)
+) Engine=InnoDB;
+`
+
+	vschema2 = &vschemapb.Keyspace{
+		Sharded: true,
+		Vindexes: map[string]*vschemapb.Vindex{
+			"hash": {
+				Type: "hash",
+			},
+		},
+		Tables: map[string]*vschemapb.Table{
+			"t1_copy_all_ks2": {
+				ColumnVindexes: []*vschemapb.ColumnVindex{{
+					Column: "id1",
+					Name:   "hash",
+				}},
+			},
+		},
+	}
 )
 
 func TestMain(m *testing.M) {
@@ -227,21 +264,36 @@ func TestMain(m *testing.M) {
 	exitCode := func() int {
 		var cfg vttest.Config
 		cfg.Topology = &vttestpb.VTTestTopology{
-			Keyspaces: []*vttestpb.Keyspace{{
-				Name: "ks",
-				Shards: []*vttestpb.Shard{{
-					Name: "-80",
-				}, {
-					Name: "80-",
-				}},
-			}},
+			Keyspaces: []*vttestpb.Keyspace{
+				{
+					Name: "ks",
+					Shards: []*vttestpb.Shard{{
+						Name: "-80",
+					}, {
+						Name: "80-",
+					}},
+				},
+				{
+					Name: "ks2",
+					Shards: []*vttestpb.Shard{{
+						Name: "-80",
+					}, {
+						Name: "80-",
+					}},
+				},
+			},
 		}
 		if err := cfg.InitSchemas("ks", schema, vschema); err != nil {
 			fmt.Fprintf(os.Stderr, "%v\n", err)
 			os.RemoveAll(cfg.SchemaDir)
 			return 1
 		}
 		defer os.RemoveAll(cfg.SchemaDir)
+		if err := cfg.InitSchemas("ks2", schema2, vschema2); err != nil {
+			fmt.Fprintf(os.Stderr, "%v\n", err)
+			os.RemoveAll(cfg.SchemaDir)
+			return 1
+		}
 
 		cfg.TabletHostName = *tabletHostName
 

go/vt/vtgate/endtoend/misc_test.go

@@ -19,6 +19,7 @@ package endtoend
 import (
 	"context"
 	"fmt"
+	osExec "os/exec"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -55,6 +56,16 @@ func TestCreateAndDropDatabase(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 
+	// cleanup the keyspace from the topology.
+	defer func() {
+		// the corresponding database needs to be created in advance.
+		// a subsequent DeleteKeyspace command returns the error of 'node doesn't exist' without it.
+		_ = exec(t, conn, "create database testitest")
+
+		_, err := osExec.Command("vtctldclient", "--server", grpcAddress, "DeleteKeyspace", "--recursive", "--force", "testitest").CombinedOutput()
+		require.NoError(t, err)
+	}()
+
 	// run it 3 times.
 	for count := 0; count < 3; count++ {
 		t.Run(fmt.Sprintf("exec:%d", count), func(t *testing.T) {

go/vt/vtgate/endtoend/row_count_test.go

@@ -24,13 +24,15 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/utils"
 )
 
 func TestRowCount(t *testing.T) {
 	ctx := context.Background()
 	conn, err := mysql.Connect(ctx, &vtParams)
 	require.NoError(t, err)
 	defer conn.Close()
+	utils.Exec(t, conn, "use ks")
 	type tc struct {
 		query    string
 		expected int

go/vt/vtgate/endtoend/vstream_test.go

@@ -258,6 +258,130 @@ func TestVStreamCopyBasic(t *testing.T) {
 	}
 }
 
+func TestVStreamCopyWithoutKeyspaceShard(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, err := mysql.Connect(ctx, &vtParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	_, err = conn.ExecuteFetch("insert into t1_copy_all(id1,id2) values(1,1), (2,2), (3,3), (4,4), (5,5), (6,6), (7,7), (8,8)", 1, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = conn.ExecuteFetch("insert into t1_copy_all_ks2(id1,id2) values(10,10), (20,20)", 1, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	filter := &binlogdatapb.Filter{
+		Rules: []*binlogdatapb.Rule{{
+			Match: "/t1_copy_all.*/",
+		}},
+	}
+	flags := &vtgatepb.VStreamFlags{}
+
+	expectedKs1EventNum := 2 /* num shards */ * (9 /* begin/field/vgtid:pos/4 rowevents avg/vgitd: lastpk/commit) */ + 3 /* begin/vgtid/commit for completed table */ + 1 /* copy operation completed */)
+	expectedKs2EventNum := 2 /* num shards */ * (6 /* begin/field/vgtid:pos/1 rowevents avg/vgitd: lastpk/commit) */ + 3 /* begin/vgtid/commit for completed table */ + 1 /* copy operation completed */)
+	expectedFullyCopyCompletedNum := 1
+
+	cases := []struct {
+		name                    string
+		shardGtid               *binlogdatapb.ShardGtid
+		expectedEventNum        int
+		expectedCompletedEvents []string
+	}{
+		{
+			name: "copy from all keyspaces",
+			shardGtid: &binlogdatapb.ShardGtid{
+				Gtid: "",
+			},
+			expectedEventNum: expectedKs1EventNum + expectedKs2EventNum + expectedFullyCopyCompletedNum,
+			expectedCompletedEvents: []string{
+				`type:COPY_COMPLETED keyspace:"ks" shard:"-80"`,
+				`type:COPY_COMPLETED keyspace:"ks" shard:"80-"`,
+				`type:COPY_COMPLETED keyspace:"ks2" shard:"-80"`,
+				`type:COPY_COMPLETED keyspace:"ks2" shard:"80-"`,
+				`type:COPY_COMPLETED`,
+			},
+		},
+		{
+			name: "copy from all shards in the keyspace",
+			shardGtid: &binlogdatapb.ShardGtid{
+				Keyspace: "ks",
+				Gtid:     "",
+			},
+			expectedEventNum: expectedKs1EventNum + expectedFullyCopyCompletedNum,
+			expectedCompletedEvents: []string{
+				`type:COPY_COMPLETED keyspace:"ks" shard:"-80"`,
+				`type:COPY_COMPLETED keyspace:"ks" shard:"80-"`,
+				`type:COPY_COMPLETED`,
+			},
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			gconn, conn, mconn, closeConnections := initialize(ctx, t)
+			defer closeConnections()
+
+			var vgtid = &binlogdatapb.VGtid{}
+			vgtid.ShardGtids = []*binlogdatapb.ShardGtid{c.shardGtid}
+			reader, err := gconn.VStream(ctx, topodatapb.TabletType_PRIMARY, vgtid, filter, flags)
+			_, _ = conn, mconn
+			if err != nil {
+				t.Fatal(err)
+			}
+			require.NotNil(t, reader)
+			var evs []*binlogdatapb.VEvent
+			var completedEvs []*binlogdatapb.VEvent
+			for {
+				e, err := reader.Recv()
+				switch err {
+				case nil:
+					evs = append(evs, e...)
+
+					for _, ev := range e {
+						if ev.Type == binlogdatapb.VEventType_COPY_COMPLETED {
+							completedEvs = append(completedEvs, ev)
+						}
+					}
+
+					printEvents(evs) // for debugging ci failures
+
+					if len(evs) == c.expectedEventNum {
+						// The arrival order of COPY_COMPLETED events with keyspace/shard is not constant.
+						// On the other hand, the last event should always be a fully COPY_COMPLETED event.
+						// That's why the sort.Slice doesn't have to handle the last element in completedEvs.
+						sort.Slice(completedEvs[:len(completedEvs)-1], func(i, j int) bool {
+							if completedEvs[i].GetKeyspace() != completedEvs[j].GetKeyspace() {
+								return completedEvs[i].GetKeyspace() < completedEvs[j].GetKeyspace()
+							}
+							return completedEvs[i].GetShard() < completedEvs[j].GetShard()
+						})
+						for i, ev := range completedEvs {
+							require.Regexp(t, c.expectedCompletedEvents[i], ev.String())
+						}
+						t.Logf("TestVStreamCopyWithoutKeyspaceShard was successful")
+						return
+					} else if c.expectedEventNum < len(evs) {
+						t.Fatalf("len(events)=%v are not expected\n", len(evs))
+					}
+				case io.EOF:
+					log.Infof("stream ended\n")
+					cancel()
+				default:
+					log.Errorf("Returned err %v", err)
+					t.Fatalf("remote error: %v\n", err)
+				}
+			}
+		})
+	}
+}
+
 func TestVStreamCopyResume(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

go/vt/vtgate/vstream_manager.go

@@ -179,13 +179,8 @@ func (vsm *vstreamManager) resolveParams(ctx context.Context, tabletType topodat
 		return nil, nil, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "vgtid must have at least one value with a starting position")
 	}
 	// To fetch from all keyspaces, the input must contain a single ShardGtid
-	// that has an empty keyspace, and the Gtid must be "current". In the
-	// future, we'll allow the Gtid to be empty which will also support
-	// copying of existing data.
+	// that has an empty keyspace.
 	if len(vgtid.ShardGtids) == 1 && vgtid.ShardGtids[0].Keyspace == "" {
-		if vgtid.ShardGtids[0].Gtid != "current" {
-			return nil, nil, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "for an empty keyspace, the Gtid value must be 'current': %v", vgtid)
-		}
 		keyspaces, err := vsm.toposerv.GetSrvKeyspaceNames(ctx, vsm.cell, false)
 		if err != nil {
 			return nil, nil, nil, err
@@ -194,17 +189,14 @@ func (vsm *vstreamManager) resolveParams(ctx context.Context, tabletType topodat
 		for _, keyspace := range keyspaces {
 			newvgtid.ShardGtids = append(newvgtid.ShardGtids, &binlogdatapb.ShardGtid{
 				Keyspace: keyspace,
-				Gtid:     "current",
+				Gtid:     vgtid.ShardGtids[0].Gtid,
 			})
 		}
 		vgtid = newvgtid
 	}
 	newvgtid := &binlogdatapb.VGtid{}
 	for _, sgtid := range vgtid.ShardGtids {
 		if sgtid.Shard == "" {
-			if sgtid.Gtid != "current" {
-				return nil, nil, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "if shards are unspecified, the Gtid value must be 'current': %v", vgtid)
-			}
 			// TODO(sougou): this should work with the new Migrate workflow
 			_, _, allShards, err := vsm.resolver.GetKeyspaceShards(ctx, sgtid.Keyspace, tabletType)
 			if err != nil {

go/vt/vtgate/vstream_manager_test.go

@@ -886,16 +886,37 @@ func TestResolveVStreamParams(t *testing.T) {
 		err:   "vgtid must have at least one value with a starting position",
 	}, {
 		input: &binlogdatapb.VGtid{
-			ShardGtids: []*binlogdatapb.ShardGtid{{}},
+			ShardGtids: []*binlogdatapb.ShardGtid{{
+				Keyspace: "TestVStream",
+			}},
 		},
-		err: "for an empty keyspace, the Gtid value must be 'current'",
-	}, {
-		input: &binlogdatapb.VGtid{
+		output: &binlogdatapb.VGtid{
 			ShardGtids: []*binlogdatapb.ShardGtid{{
 				Keyspace: "TestVStream",
+				Shard:    "-20",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "20-40",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "40-60",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "60-80",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "80-a0",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "a0-c0",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "c0-e0",
+			}, {
+				Keyspace: "TestVStream",
+				Shard:    "e0-",
 			}},
 		},
-		err: "if shards are unspecified, the Gtid value must be 'current'",
 	}, {
 		input: &binlogdatapb.VGtid{
 			ShardGtids: []*binlogdatapb.ShardGtid{{
@@ -987,17 +1008,34 @@ func TestResolveVStreamParams(t *testing.T) {
 		assert.Equal(t, wantFilter, filter, tcase.input)
 		require.False(t, flags.MinimizeSkew)
 	}
+
 	// Special-case: empty keyspace because output is too big.
-	input := &binlogdatapb.VGtid{
-		ShardGtids: []*binlogdatapb.ShardGtid{{
-			Gtid: "current",
-		}},
+	specialCases := []struct {
+		input  string
+		output string
+	}{
+		{
+			input:  "current",
+			output: "current",
+		},
+		{},
 	}
-	vgtid, _, _, err := vsm.resolveParams(context.Background(), topodatapb.TabletType_REPLICA, input, nil, nil)
-	require.NoError(t, err, input)
-	if got, want := len(vgtid.ShardGtids), 8; want >= got {
-		t.Errorf("len(vgtid.ShardGtids): %v, must be >%d", got, want)
+	for _, tcase := range specialCases {
+		input := &binlogdatapb.VGtid{
+			ShardGtids: []*binlogdatapb.ShardGtid{{
+				Gtid: tcase.input,
+			}},
+		}
+		vgtid, _, _, err := vsm.resolveParams(context.Background(), topodatapb.TabletType_REPLICA, input, nil, nil)
+		require.NoError(t, err, tcase.input)
+		if got, want := len(vgtid.ShardGtids), 8; want >= got {
+			t.Errorf("len(vgtid.ShardGtids): %v, must be >%d", got, want)
+		}
+		for _, s := range vgtid.ShardGtids {
+			require.Equal(t, tcase.output, s.Gtid)
+		}
 	}
+
 	for _, minimizeSkew := range []bool{true, false} {
 		t.Run(fmt.Sprintf("resolveParams MinimizeSkew %t", minimizeSkew), func(t *testing.T) {
 			flags := &vtgatepb.VStreamFlags{MinimizeSkew: minimizeSkew}