vtctldclient: Fix Apply (Shard | Keyspace| Table) Routing Rules commands

go/cmd/vtctldclient/command/keyspace_routing_rules.go

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

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/marshal_test.go

@@ -36,6 +36,12 @@ func TestMarshalPB(t *testing.T) {
 	require.NoErrorf(t, err, "MarshalPB(%+v) error", col)
 	want := "{\"name\":\"c1\",\"type\":\"VARCHAR\"}"
 	assert.Equalf(t, want, string(b), "MarshalPB(%+v)", col)
+
+	// Test MarshalPB with nil proto.
+	b, err = MarshalPB(nil)
+	require.NoErrorf(t, err, "MarshalPB(nil) error")
+	want = "{}"
+	assert.Equalf(t, want, string(b), "MarshalPB(nil)")
 }
 
 func TestMarshalIndentPB(t *testing.T) {

go/json2/unmarshal.go

@@ -53,3 +53,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) {
+	protoData := &emptypb.Empty{}
+	protoJSONData, err := protojson.Marshal(protoData)
+	require.Nil(t, err)
+
+	var pb emptypb.Empty
+	err = UnmarshalPB(protoJSONData, &pb)
+	require.Nil(t, err)
+	require.Equal(t, protoData, &pb)
+}

go/test/endtoend/vreplication/vreplication_vtctldclient_cli_test.go

@@ -19,10 +19,13 @@ package vreplication
 import (
 	"encoding/json"
 	"fmt"
+	"os"
 	"slices"
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/json2"
+
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/maps"
 	"google.golang.org/protobuf/encoding/protojson"
@@ -61,6 +64,9 @@ func TestVtctldclientCLI(t *testing.T) {
 	workflowName := "wf1"
 	targetTabs := setupMinimalCustomerKeyspace(t)
 
+	t.Run("RoutingRulesCommands", func(t *testing.T) {
+		testAllRoutingRulesCommands(t)
+	})
 	t.Run("WorkflowList", func(t *testing.T) {
 		testWorkflowList(t, sourceKeyspaceName, targetKeyspaceName)
 	})
@@ -438,3 +444,137 @@ func validateMoveTablesWorkflow(t *testing.T, workflows []*vtctldatapb.Workflow)
 	require.Equal(t, binlogdatapb.OnDDLAction_STOP, bls.OnDdl)
 	require.True(t, bls.StopAfterCopy)
 }
+
+func testAllRoutingRulesCommands(t *testing.T) {
+	var rulesBytes []byte
+	var err error
+	var validateRules func(want, got string)
+	typs := []string{"RoutingRules"} // , "ShardRoutingRules", "KeyspaceRoutingRules"}
+	for _, typ := range typs {
+		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.Unmarshal([]byte(want), wantRules))
+				var gotRules = &vschemapb.RoutingRules{}
+				require.NoError(t, json2.Unmarshal([]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.Unmarshal([]byte(want), wantRules))
+				var gotRules = &vschemapb.ShardRoutingRules{}
+				require.NoError(t, json2.Unmarshal([]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.Unmarshal([]byte(want), wantRules))
+				var gotRules = &vschemapb.KeyspaceRoutingRules{}
+				require.NoError(t, json2.Unmarshal([]byte(got), gotRules))
+				require.EqualValues(t, wantRules, gotRules)
+			}
+		default:
+			t.Fatalf("Unknown type %s", typ)
+		}
+		testOneRoutingRulesCommand(t, typ, string(rulesBytes), validateRules)
+	}
+
+}
+
+func testOneRoutingRulesCommand(t *testing.T, typ string, rules string, validateRules func(want, got string)) {
+	type routingRulesTest struct {
+		name    string
+		rules   string
+		useFile bool
+	}
+	tests := []routingRulesTest{
+		{
+			name:  "inline",
+			rules: string(rules),
+		},
+		{
+			name:    "file",
+			rules:   string(rules),
+			useFile: true,
+		},
+		{
+			name:  "empty",
+			rules: "{}",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(typ+"/"+tt.name, func(t *testing.T) {
+			done := false
+			wantRules := tt.rules
+			typ2 := "camelCase"
+			for !done {
+				t.Run(typ2, 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)
+					if typ2 == "snake_case" {
+						done = true
+					} else {
+						wantRules = output
+						typ2 = "snake_case"
+					}
+				})
+			}
+		})
+	}
+
+}