Multi-tenant migrations: add topo locking while updating keyspace routing rules (#15807)

Signed-off-by: Rohit Nayak <rohit@planetscale.com>
Signed-off-by: Matt Lord <mattalord@gmail.com>
Co-authored-by: Matt Lord <mattalord@gmail.com>

Author: rohit-nayak-ps

go/cmd/vtctldclient/command/keyspace_routing_rules.go

@@ -70,7 +70,6 @@ var applyKeyspaceRoutingRulesOptions = struct {
 func commandApplyKeyspaceRoutingRules(cmd *cobra.Command, args []string) error {
 	opts := applyKeyspaceRoutingRulesOptions
 	cli.FinishedParsing(cmd)
-
 	var rulesBytes []byte
 	if opts.RulesFilePath != "" {
 		data, err := os.ReadFile(opts.RulesFilePath)
@@ -86,13 +85,14 @@ func commandApplyKeyspaceRoutingRules(cmd *cobra.Command, args []string) error {
 	if err := json2.Unmarshal(rulesBytes, &krr); err != nil {
 		return err
 	}
-	// Round-trip so that when we display the result it's readable.
-	data, err := cli.MarshalJSON(krr)
-	if err != nil {
-		return err
-	}
 
 	if opts.DryRun {
+		// Round-trip so that when we display the result it's readable.
+		data, err := cli.MarshalJSON(krr)
+		if err != nil {
+			return err
+		}
+
 		fmt.Printf("[DRY RUN] Would have saved new KeyspaceRoutingRules object:\n%s\n", data)
 
 		if opts.SkipRebuild {
@@ -105,11 +105,10 @@ func commandApplyKeyspaceRoutingRules(cmd *cobra.Command, args []string) error {
 				fmt.Printf(" in the following cells: %s.\n", strings.Join(applyKeyspaceRoutingRulesOptions.Cells, ", "))
 			}
 		}
-
 		return nil
 	}
 
-	_, err = client.ApplyKeyspaceRoutingRules(commandCtx, &vtctldatapb.ApplyKeyspaceRoutingRulesRequest{
+	resp, err := client.ApplyKeyspaceRoutingRules(commandCtx, &vtctldatapb.ApplyKeyspaceRoutingRulesRequest{
 		KeyspaceRoutingRules: krr,
 		SkipRebuild:          opts.SkipRebuild,
 		RebuildCells:         opts.Cells,
@@ -118,12 +117,11 @@ func commandApplyKeyspaceRoutingRules(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	fmt.Printf("New KeyspaceRoutingRules object:\n%s\nIf this is not what you expected, check the input data (as JSON parsing will skip unexpected fields).\n", data)
-
-	if opts.SkipRebuild {
-		fmt.Println("Skipping rebuild of VSchema graph as requested, you will need to run RebuildVSchemaGraph for the changes to take effect.")
+	respJSON, err := cli.MarshalJSON(resp)
+	if err != nil {
+		return err
 	}
-
+	fmt.Printf("%s\n", respJSON)
 	return nil
 }
 

go/test/endtoend/vreplication/multi_tenant_test.go

@@ -37,14 +37,14 @@ import (
 	"testing"
 	"time"
 
-	"google.golang.org/protobuf/encoding/protojson"
-
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/encoding/protojson"
 
 	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/proto/vtctldata"
+
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
-	"vitess.io/vitess/go/vt/proto/vtctldata"
 )
 
 type tenantMigrationStatus int
@@ -71,6 +71,8 @@ var (
 	chNotSetup, chNotCreated, chInProgress, chSwitched, chCompleted chan int64
 	// counters to keep track of the number of tenants in each state
 	numSetup, numInProgress, numSwitched, numCompleted atomic.Int64
+
+	emptyKeyspaceRoutingRules = &vschemapb.KeyspaceRoutingRules{}
 )
 
 // multiTenantMigration manages the migration of multiple tenants to a single target keyspace.
@@ -193,6 +195,7 @@ func TestMultiTenantSimple(t *testing.T) {
 	}
 
 	require.Zero(t, len(getKeyspaceRoutingRules(t, vc).Rules))
+
 	mt.Create()
 	confirmKeyspacesRoutedTo(t, sourceKeyspace, "s1", "t1", nil)
 	validateKeyspaceRoutingRules(t, vc, initialRules)
@@ -223,6 +226,41 @@ func TestMultiTenantSimple(t *testing.T) {
 	log.Infof("Migration completed, total rows in target: %d", actualRowsInserted)
 	require.Equal(t, lastIndex, int64(actualRowsInserted))
 
+	t.Run("Test ApplyKeyspaceRoutingRules", func(t *testing.T) {
+		// First set of rules
+		applyKeyspaceRoutingRules(t, initialRules)
+
+		updatedRules := &vschemapb.KeyspaceRoutingRules{
+			Rules: []*vschemapb.KeyspaceRoutingRule{
+				{FromKeyspace: "s1", ToKeyspace: "mt"},
+				{FromKeyspace: "s1@rdonly", ToKeyspace: "mt"},
+				{FromKeyspace: "s1@replica", ToKeyspace: "mt"},
+			},
+		}
+		// Update the rules
+		applyKeyspaceRoutingRules(t, updatedRules)
+		// Update with the same rules
+		applyKeyspaceRoutingRules(t, updatedRules)
+		// Remove the rules
+		applyKeyspaceRoutingRules(t, emptyKeyspaceRoutingRules)
+		// Test setting empty rules again
+		applyKeyspaceRoutingRules(t, emptyKeyspaceRoutingRules)
+	})
+}
+
+func applyKeyspaceRoutingRules(t *testing.T, newRules *vschemapb.KeyspaceRoutingRules) {
+	var rulesJSON []byte
+	var err error
+	require.NotNil(t, newRules)
+	rulesJSON, err = json.Marshal(newRules)
+	require.NoError(t, err)
+	output, err := vc.VtctldClient.ExecuteCommandWithOutput("ApplyKeyspaceRoutingRules", "--rules", string(rulesJSON))
+	require.NoError(t, err)
+
+	response := &vtctldata.ApplyKeyspaceRoutingRulesResponse{}
+	err = json.Unmarshal([]byte(output), response)
+	require.NoError(t, err)
+	require.ElementsMatch(t, newRules.Rules, response.GetKeyspaceRoutingRules().Rules)
 }
 
 func confirmOnlyReadsSwitched(t *testing.T) {
@@ -251,10 +289,10 @@ func confirmOnlyWritesSwitched(t *testing.T) {
 	validateKeyspaceRoutingRules(t, vc, rules)
 }
 
-// TestMultiTenantSimpleSharded tests a single tenant migration to a sharded target. The aim is to test
+// TestMultiTenantSharded tests a single tenant migration to a sharded target. The aim is to test
 // the specification of the target shards in all the MoveTables subcommands, including creating only one stream
 // for a tenant on the shard to which this tenant id will be routed, using the specified Vindex.
-func TestMultiTenantSimpleSharded(t *testing.T) {
+func TestMultiTenantSharded(t *testing.T) {
 	setSidecarDBName("_vt")
 	// Don't create RDONLY tablets to reduce number of tablets created to reduce resource requirements for the test.
 	origDefaultRdonly := defaultRdonly