fix: when merging routes with reference routing, we should also copy …

…the keyspace

Signed-off-by: Andres Taylor <[email protected]>

go/vt/vtgate/planbuilder/physical/route_planning.go

@@ -569,17 +569,22 @@ func leaves(op abstract.Operator) (sources []abstract.Operator) {
 }
 
 func tryMergeReferenceTable(aRoute, bRoute *Route, merger mergeFunc) (*Route, error) {
-	// if either side is a reference table, we can just merge it and use the opcode of the other side
-	var opCode engine.Opcode
-	var selected *VindexOption
+	var (
+		// if either side is a reference table, we can just merge it and use the opcode of the other side
+		opCode engine.Opcode
+		vindex *VindexOption
+		ks     *vindexes.Keyspace
+	)
 
 	switch {
 	case aRoute.RouteOpCode == engine.Reference:
-		selected = bRoute.Selected
+		vindex = bRoute.Selected
 		opCode = bRoute.RouteOpCode
+		ks = bRoute.Keyspace
 	case bRoute.RouteOpCode == engine.Reference:
-		selected = aRoute.Selected
+		vindex = aRoute.Selected
 		opCode = aRoute.RouteOpCode
+		ks = aRoute.Keyspace
 	default:
 		return nil, nil
 	}
@@ -589,7 +594,8 @@ func tryMergeReferenceTable(aRoute, bRoute *Route, merger mergeFunc) (*Route, er
 		return nil, err
 	}
 	r.RouteOpCode = opCode
-	r.Selected = selected
+	r.Selected = vindex
+	r.Keyspace = ks
 	return r, nil
 }
 

go/vt/vtgate/planbuilder/plan_test.go

@@ -270,6 +270,18 @@ func TestOneWithMainAsDefault(t *testing.T) {
 	testFile(t, "onecase.txt", "", vschema)
 }
 
+func TestOneWithSecondUserAsDefault(t *testing.T) {
+	vschema := &vschemaWrapper{
+		v: loadSchema(t, "schema_test.json", true),
+		keyspace: &vindexes.Keyspace{
+			Name:    "second_user",
+			Sharded: true,
+		},
+	}
+
+	testFile(t, "onecase.txt", "", vschema)
+}
+
 func TestRubyOnRailsQueries(t *testing.T) {
 	vschemaWrapper := &vschemaWrapper{
 		v:             loadSchema(t, "rails_schema_test.json", true),
@@ -384,6 +396,21 @@ func TestWithDefaultKeyspaceFromFile(t *testing.T) {
 	testFile(t, "call_cases.txt", testOutputTempDir, vschema)
 }
 
+func TestWithDefaultKeyspaceFromFileSharded(t *testing.T) {
+	// We are testing this separately so we can set a default keyspace
+	vschema := &vschemaWrapper{
+		v: loadSchema(t, "schema_test.json", true),
+		keyspace: &vindexes.Keyspace{
+			Name:    "second_user",
+			Sharded: true,
+		},
+		tabletType: topodatapb.TabletType_PRIMARY,
+	}
+
+	testOutputTempDir := makeTestOutput(t)
+	testFile(t, "select_cases_with_default.txt", testOutputTempDir, vschema)
+}
+
 func TestWithSystemSchemaAsDefaultKeyspace(t *testing.T) {
 	// We are testing this separately so we can set a default keyspace
 	vschema := &vschemaWrapper{

go/vt/vtgate/planbuilder/testdata/select_cases.txt

@@ -2195,8 +2195,8 @@ Gen4 plan same as above
     "OperatorType": "Route",
     "Variant": "Scatter",
     "Keyspace": {
-      "Name": "main",
-      "Sharded": false
+      "Name": "user",
+      "Sharded": true
     },
     "FieldQuery": "select 42, id from dual, `user` where 1 != 1",
     "Query": "select 42, id from dual, `user`",

go/vt/vtgate/planbuilder/testdata/select_cases_with_default.txt

@@ -0,0 +1,61 @@
+# EXISTS subquery when the default ks is different than the inner query
+"select exists(select * from user where id = 5)"
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select * from user where id = 5)",
+  "Instructions": {
+    "OperatorType": "Subquery",
+    "Variant": "PulloutExists",
+    "PulloutVars": [
+      "__sq_has_values1",
+      "__sq1"
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select * from `user` where 1 != 1",
+        "Query": "select * from `user` where id = 5",
+        "Table": "`user`",
+        "Values": [
+          "INT64(5)"
+        ],
+        "Vindex": "user_index"
+      },
+      {
+        "OperatorType": "Route",
+        "Variant": "Reference",
+        "Keyspace": {
+          "Name": "second_user",
+          "Sharded": true
+        },
+        "FieldQuery": "select :__sq_has_values1 from dual where 1 != 1",
+        "Query": "select :__sq_has_values1 from dual",
+        "Table": "dual"
+      }
+    ]
+  }
+}
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select * from user where id = 5)",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "EqualUnique",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select exists (select * from `user` where 1 != 1) from dual where 1 != 1",
+    "Query": "select exists (select * from `user` where id = 5) from dual",
+    "Table": "dual",
+    "Values": [
+      "INT64(5)"
+    ],
+    "Vindex": "user_index"
+  }
+}