Cherry-pick 8cf7f31 with conflicts

go/cmd/vtctldclient/command/keyspace_routing_rules.go

@@ -0,0 +1,154 @@
+/*
+Copyright 2024 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 command
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/cmd/vtctldclient/cli"
+	"vitess.io/vitess/go/json2"
+
+	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
+	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
+)
+
+var (
+	// ApplyKeyspaceRoutingRules makes an ApplyKeyspaceRoutingRules gRPC call to a vtctld.
+	ApplyKeyspaceRoutingRules = &cobra.Command{
+		Use:                   "ApplyKeyspaceRoutingRules {--rules RULES | --rules-file RULES_FILE} [--cells=c1,c2,...] [--skip-rebuild] [--dry-run]",
+		Short:                 "Applies the provided keyspace routing rules.",
+		DisableFlagsInUseLine: true,
+		Args:                  cobra.NoArgs,
+		PreRunE:               validateApplyKeyspaceRoutingRulesOptions,
+		RunE:                  commandApplyKeyspaceRoutingRules,
+	}
+	// GetKeyspaceRoutingRules makes a GetKeyspaceRoutingRules gRPC call to a vtctld.
+	GetKeyspaceRoutingRules = &cobra.Command{
+		Use:                   "GetKeyspaceRoutingRules",
+		Short:                 "Displays the currently active keyspace routing rules.",
+		DisableFlagsInUseLine: true,
+		Args:                  cobra.NoArgs,
+		RunE:                  commandGetKeyspaceRoutingRules,
+	}
+)
+
+func validateApplyKeyspaceRoutingRulesOptions(cmd *cobra.Command, args []string) error {
+	opts := applyKeyspaceRoutingRulesOptions
+	if (opts.Rules != "" && opts.RulesFilePath != "") || (opts.Rules == "" && opts.RulesFilePath == "") {
+		return errors.New("must pass exactly one of --rules or --rules-file")
+	}
+	return nil
+}
+
+var applyKeyspaceRoutingRulesOptions = struct {
+	Rules         string
+	RulesFilePath string
+	Cells         []string
+	SkipRebuild   bool
+	DryRun        bool
+}{}
+
+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)
+		if err != nil {
+			return err
+		}
+		rulesBytes = data
+	} else {
+		rulesBytes = []byte(opts.Rules)
+	}
+
+	krr := &vschemapb.KeyspaceRoutingRules{}
+	if err := json2.UnmarshalPB(rulesBytes, krr); 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 {
+			fmt.Println("[DRY RUN] Would not have rebuilt VSchema graph, would have required operator to run RebuildVSchemaGraph for changes to take effect.")
+		} else {
+			fmt.Print("[DRY RUN] Would have rebuilt the VSchema graph")
+			if len(opts.Cells) == 0 {
+				fmt.Print(" in all cells\n")
+			} else {
+				fmt.Printf(" in the following cells: %s.\n", strings.Join(applyKeyspaceRoutingRulesOptions.Cells, ", "))
+			}
+		}
+		return nil
+	}
+
+	resp, err := client.ApplyKeyspaceRoutingRules(commandCtx, &vtctldatapb.ApplyKeyspaceRoutingRulesRequest{
+		KeyspaceRoutingRules: krr,
+		SkipRebuild:          opts.SkipRebuild,
+		RebuildCells:         opts.Cells,
+	})
+	if err != nil {
+		return err
+	}
+
+	respJSON, err := cli.MarshalJSON(resp)
+	if err != nil {
+		return err
+	}
+	fmt.Printf("%s\n", respJSON)
+	return nil
+}
+
+func commandGetKeyspaceRoutingRules(cmd *cobra.Command, args []string) error {
+	cli.FinishedParsing(cmd)
+
+	resp, err := client.GetKeyspaceRoutingRules(commandCtx, &vtctldatapb.GetKeyspaceRoutingRulesRequest{})
+	if err != nil {
+		return err
+	}
+
+	data, err := cli.MarshalJSON(resp.KeyspaceRoutingRules)
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("%s\n", data)
+
+	return nil
+}
+
+func init() {
+	ApplyKeyspaceRoutingRules.Flags().StringVarP(&applyKeyspaceRoutingRulesOptions.Rules, "rules", "r", "", "Keyspace routing rules, specified as a string")
+	ApplyKeyspaceRoutingRules.Flags().StringVarP(&applyKeyspaceRoutingRulesOptions.RulesFilePath, "rules-file", "f", "", "Path to a file containing keyspace routing rules specified as JSON")
+	ApplyKeyspaceRoutingRules.Flags().StringSliceVarP(&applyKeyspaceRoutingRulesOptions.Cells, "cells", "c", nil, "Limit the VSchema graph rebuilding to the specified cells. Ignored if --skip-rebuild is specified.")
+	ApplyKeyspaceRoutingRules.Flags().BoolVar(&applyKeyspaceRoutingRulesOptions.SkipRebuild, "skip-rebuild", false, "Skip rebuilding the SrvVSchema objects.")
+	ApplyKeyspaceRoutingRules.Flags().BoolVarP(&applyKeyspaceRoutingRulesOptions.DryRun, "dry-run", "d", false, "Validate the specified keyspace routing rules and note actions that would be taken, but do not actually apply the rules to the topo.")
+	Root.AddCommand(ApplyKeyspaceRoutingRules)
+	Root.AddCommand(GetKeyspaceRoutingRules)
+}

go/cmd/vtctldclient/command/routing_rules.go

@@ -82,7 +82,7 @@ func commandApplyRoutingRules(cmd *cobra.Command, args []string) error {
 	}
 
 	rr := &vschemapb.RoutingRules{}
-	if err := json2.Unmarshal(rulesBytes, &rr); err != nil {
+	if err := json2.UnmarshalPB(rulesBytes, rr); err != nil {
 		return err
 	}
 

go/cmd/vtctldclient/command/shard_routing_rules.go

@@ -87,7 +87,7 @@ func commandApplyShardRoutingRules(cmd *cobra.Command, args []string) error {
 	}
 
 	srr := &vschemapb.ShardRoutingRules{}
-	if err := json2.Unmarshal(rulesBytes, &srr); err != nil {
+	if err := json2.UnmarshalPB(rulesBytes, srr); err != nil {
 		return err
 	}
 	// Round-trip so when we display the result it's readable.

go/json2/unmarshal.go

@@ -33,8 +33,7 @@ var carriageReturn = []byte("\n")
 // efficient and should not be used for high QPS operations.
 func Unmarshal(data []byte, v any) error {
 	if pb, ok := v.(proto.Message); ok {
-		opts := protojson.UnmarshalOptions{DiscardUnknown: true}
-		return annotate(data, opts.Unmarshal(data, pb))
+		return UnmarshalPB(data, pb)
 	}
 	return annotate(data, json.Unmarshal(data, v))
 }
@@ -53,3 +52,9 @@ func annotate(data []byte, err error) error {
 
 	return fmt.Errorf("line: %d, position %d: %v", line, pos, err)
 }
+
+// UnmarshalPB is similar to Unmarshal but specifically for proto.Message to add type safety.
+func UnmarshalPB(data []byte, pb proto.Message) error {
+	opts := protojson.UnmarshalOptions{DiscardUnknown: true}
+	return annotate(data, opts.Unmarshal(data, pb))
+}

go/json2/unmarshal_test.go

@@ -91,3 +91,14 @@ func TestAnnotate(t *testing.T) {
 		require.Equal(t, tcase.err, err, "annotate(%s, %v) error", string(tcase.data), tcase.err)
 	}
 }
+
+func TestUnmarshalPB(t *testing.T) {
+	want := &emptypb.Empty{}
+	json, err := protojson.Marshal(want)
+	require.NoError(t, err)
+
+	var got emptypb.Empty
+	err = UnmarshalPB(json, &got)
+	require.NoError(t, err)
+	require.Equal(t, want, &got)
+}

go/test/endtoend/vreplication/vreplication_vtctldclient_cli_test.go

@@ -18,6 +18,7 @@ package vreplication
 
 import (
 	"fmt"
+	"os"
 	"slices"
 	"strings"
 	"testing"
@@ -26,6 +27,7 @@ import (
 	"golang.org/x/exp/maps"
 	"google.golang.org/protobuf/encoding/protojson"
 
+	"vitess.io/vitess/go/json2"
 	"vitess.io/vitess/go/test/endtoend/cluster"
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
@@ -59,6 +61,15 @@ func TestVtctldclientCLI(t *testing.T) {
 	workflowName := "wf1"
 	targetTabs := setupMinimalCustomerKeyspace(t)
 
+<<<<<<< HEAD
+=======
+	t.Run("RoutingRulesApply", func(t *testing.T) {
+		testRoutingRulesApplyCommands(t)
+	})
+	t.Run("WorkflowList", func(t *testing.T) {
+		testWorkflowList(t, sourceKeyspaceName, targetKeyspaceName)
+	})
+>>>>>>> 8cf7f31d59 (vtctldclient: Apply (Shard | Keyspace| Table) Routing Rules commands don't work (#16096))
 	t.Run("MoveTablesCreateFlags1", func(t *testing.T) {
 		testMoveTablesFlags1(t, &mt, sourceKeyspaceName, targetKeyspaceName, workflowName, targetTabs)
 	})
@@ -389,3 +400,138 @@ func validateMoveTablesWorkflow(t *testing.T, workflows []*vtctldatapb.Workflow)
 	require.Equal(t, binlogdatapb.OnDDLAction_STOP, bls.OnDdl)
 	require.True(t, bls.StopAfterCopy)
 }
+
+// Test that routing rules can be applied using the vtctldclient CLI for all types of routing rules.
+func testRoutingRulesApplyCommands(t *testing.T) {
+	var rulesBytes []byte
+	var err error
+	var validateRules func(want, got string)
+
+	ruleTypes := []string{"RoutingRules", "ShardRoutingRules", "KeyspaceRoutingRules"}
+	for _, typ := range ruleTypes {
+		switch typ {
+		case "RoutingRules":
+			rr := &vschemapb.RoutingRules{
+				Rules: []*vschemapb.RoutingRule{
+					{
+						FromTable: "from1",
+						ToTables:  []string{"to1", "to2"},
+					},
+				},
+			}
+			rulesBytes, err = json2.MarshalPB(rr)
+			require.NoError(t, err)
+			validateRules = func(want, got string) {
+				var wantRules = &vschemapb.RoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(want), wantRules))
+				var gotRules = &vschemapb.RoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(got), gotRules))
+				require.EqualValues(t, wantRules, gotRules)
+			}
+		case "ShardRoutingRules":
+			srr := &vschemapb.ShardRoutingRules{
+				Rules: []*vschemapb.ShardRoutingRule{
+					{
+						FromKeyspace: "from1",
+						ToKeyspace:   "to1",
+						Shard:        "-80",
+					},
+				},
+			}
+			rulesBytes, err = json2.MarshalPB(srr)
+			require.NoError(t, err)
+			validateRules = func(want, got string) {
+				var wantRules = &vschemapb.ShardRoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(want), wantRules))
+				var gotRules = &vschemapb.ShardRoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(got), gotRules))
+				require.EqualValues(t, wantRules, gotRules)
+			}
+		case "KeyspaceRoutingRules":
+			krr := &vschemapb.KeyspaceRoutingRules{
+				Rules: []*vschemapb.KeyspaceRoutingRule{
+					{
+						FromKeyspace: "from1",
+						ToKeyspace:   "to1",
+					},
+				},
+			}
+			rulesBytes, err = json2.MarshalPB(krr)
+			require.NoError(t, err)
+			validateRules = func(want, got string) {
+				var wantRules = &vschemapb.KeyspaceRoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(want), wantRules))
+				var gotRules = &vschemapb.KeyspaceRoutingRules{}
+				require.NoError(t, json2.UnmarshalPB([]byte(got), gotRules))
+				require.EqualValues(t, wantRules, gotRules)
+			}
+		default:
+			require.FailNow(t, "Unknown type %s", typ)
+		}
+		testOneRoutingRulesCommand(t, typ, string(rulesBytes), validateRules)
+	}
+
+}
+
+// For a given routing rules type, test that the rules can be applied using the vtctldclient CLI.
+// We test both inline and file-based rules.
+// The test also validates that both camelCase and snake_case key names work correctly.
+func testOneRoutingRulesCommand(t *testing.T, typ string, rules string, validateRules func(want, got string)) {
+	type routingRulesTest struct {
+		name    string
+		rules   string
+		useFile bool // if true, use a file to pass the rules
+	}
+	tests := []routingRulesTest{
+		{
+			name:  "inline",
+			rules: rules,
+		},
+		{
+			name:    "file",
+			rules:   rules,
+			useFile: true,
+		},
+		{
+			name:  "empty", // finally, cleanup rules
+			rules: "{}",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(typ+"/"+tt.name, func(t *testing.T) {
+			wantRules := tt.rules
+			// The input rules are in camelCase, since they are the output of json2.MarshalPB
+			// The first iteration uses the output of routing rule Gets which are in snake_case.
+			for _, keyCase := range []string{"camelCase", "snake_case"} {
+				t.Run(keyCase, func(t *testing.T) {
+					var args []string
+					apply := fmt.Sprintf("Apply%s", typ)
+					get := fmt.Sprintf("Get%s", typ)
+					args = append(args, apply)
+					if tt.useFile {
+						tmpFile, err := os.CreateTemp("", fmt.Sprintf("%s_rules.json", tt.name))
+						require.NoError(t, err)
+						defer os.Remove(tmpFile.Name())
+						_, err = tmpFile.WriteString(wantRules)
+						require.NoError(t, err)
+						args = append(args, "--rules-file", tmpFile.Name())
+					} else {
+						args = append(args, "--rules", wantRules)
+					}
+					var output string
+					var err error
+					if output, err = vc.VtctldClient.ExecuteCommandWithOutput(args...); err != nil {
+						require.FailNowf(t, "failed action", apply, "%v: %s", err, output)
+					}
+					if output, err = vc.VtctldClient.ExecuteCommandWithOutput(get); err != nil {
+						require.FailNowf(t, "failed action", get, "%v: %s", err, output)
+					}
+					validateRules(wantRules, output)
+					// output of GetRoutingRules is in snake_case and we use it for the next iteration which
+					// tests applying rules with snake_case keys.
+					wantRules = output
+				})
+			}
+		})
+	}
+}