Handle subquery merging with references correctly

Fixes vitessio#8638 Signed-off-by: Andres Taylor <[email protected]>
planetscale · Aug 23, 2021 · e9aa716 · e9aa716
1 parent 56a2288
commit e9aa716
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 9 deletions.
diff --git a/go/vt/vtgate/planbuilder/route.go b/go/vt/vtgate/planbuilder/route.go
@@ -458,9 +458,15 @@ func (rb *route) SubqueryCanMerge(pb *primitiveBuilder, inner *route) bool {
 	if rb.eroute.Keyspace.Name != inner.eroute.Keyspace.Name {
 		return false
 	}
+
+	// if either side is a reference table, we can just merge it and use the opcode of the other side
+	if rb.eroute.Opcode == engine.SelectReference || inner.eroute.Opcode == engine.SelectReference {
+		return true
+	}
+
 	switch rb.eroute.Opcode {
-	case engine.SelectUnsharded, engine.SelectDBA, engine.SelectReference:
-		return rb.eroute.Opcode == inner.eroute.Opcode || inner.eroute.Opcode == engine.SelectReference
+	case engine.SelectUnsharded, engine.SelectDBA:
+		return rb.eroute.Opcode == inner.eroute.Opcode
 	case engine.SelectEqualUnique:
 		// Check if they target the same shard.
 		if inner.eroute.Opcode == engine.SelectEqualUnique && rb.eroute.Vindex == inner.eroute.Vindex && valEqual(rb.condition, inner.condition) {
@@ -469,11 +475,7 @@ func (rb *route) SubqueryCanMerge(pb *primitiveBuilder, inner *route) bool {
 	case engine.SelectNext:
 		return false
 	}
-	// Any sharded plan (including SelectEqualUnique) can merge on a reference table subquery.
-	// This excludes the case of SelectReference with a sharded subquery.
-	if inner.eroute.Opcode == engine.SelectReference {
-		return true
-	}
+
 	switch vals := inner.condition.(type) {
 	case *sqlparser.ColName:
 		if pb.st.Vindex(vals, rb) == inner.eroute.Vindex {

diff --git a/go/vt/vtgate/planbuilder/route_test.go b/go/vt/vtgate/planbuilder/route_test.go
@@ -87,9 +87,9 @@ func TestSubqueryCanMerge(t *testing.T) {
 		{false, false, false, false, false, false, false, false, true, false},
 		{false, false, false, false, false, false, false, false, true, false},
 		{false, false, false, false, false, false, false, false, true, false},
-		{false, false, false, false, false, false, false, false, false, false},
-		{false, false, false, false, false, false, false, true, true, false},
 		{false, false, false, false, false, false, false, false, true, false},
+		{false, false, false, false, false, false, false, true, true, false},
+		{true, true, true, true, true, true, true, true, true, true},
 		{false, false, false, false, false, false, false, false, true, false},
 	}
 

diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.txt b/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
@@ -2427,3 +2427,22 @@ Gen4 plan same as above
 "delete user from (select * from user) music where id = 1"
 "Unknown table 'user' in MULTI DELETE"
 Gen4 plan same as above
+
+"INSERT INTO main.user_privacy_consents (user_id, accepted_at)  SELECT user_id, accepted_at FROM (SELECT 1 as user_id, 1629194864 as accepted_at) AS tmp WHERE NOT EXISTS (SELECT user_id FROM main.user_privacy_consents WHERE user_id = 1)"
+{
+  "QueryType": "INSERT",
+  "Original": "INSERT INTO main.user_privacy_consents (user_id, accepted_at)  SELECT user_id, accepted_at FROM (SELECT 1 as user_id, 1629194864 as accepted_at) AS tmp WHERE NOT EXISTS (SELECT user_id FROM main.user_privacy_consents WHERE user_id = 1)",
+  "Instructions": {
+    "OperatorType": "Insert",
+    "Variant": "Unsharded",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    },
+    "TargetTabletType": "PRIMARY",
+    "MultiShardAutocommit": false,
+    "Query": "insert into user_privacy_consents(user_id, accepted_at) select user_id, accepted_at from (select 1 as user_id, 1629194864 as accepted_at from dual) as tmp where not exists (select user_id from user_privacy_consents where user_id = 1)",
+    "TableName": "user_privacy_consents"
+  }
+}
+Gen4 plan same as above