From 98fbe07acc628375b3d113db7e547b5889437372 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Taylor?= <andres@planetscale.com>
Date: Thu, 13 Feb 2025 08:13:49 +0100
Subject: [PATCH] bugfix: support subqueries inside subqueries when merging
 (#17764)

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../operators/subquery_planning.go            | 50 ++++++++++--------
 .../planbuilder/testdata/select_cases.json    | 52 +++++++++++++++++++
 2 files changed, 81 insertions(+), 21 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/subquery_planning.go b/go/vt/vtgate/planbuilder/operators/subquery_planning.go
index fcae37ce1d1..b8e3c65b707 100644
--- a/go/vt/vtgate/planbuilder/operators/subquery_planning.go
+++ b/go/vt/vtgate/planbuilder/operators/subquery_planning.go
@@ -148,30 +148,38 @@ func mergeSubqueryExpr(ctx *plancontext.PlanningContext, pe *ProjExpr) {
 
 func rewriteMergedSubqueryExpr(ctx *plancontext.PlanningContext, se SubQueryExpression, expr sqlparser.Expr) (sqlparser.Expr, bool) {
 	rewritten := false
-	for _, sq := range se {
-		for _, sq2 := range ctx.MergedSubqueries {
-			if sq.originalSubquery == sq2 {
-				expr = sqlparser.Rewrite(expr, nil, func(cursor *sqlparser.Cursor) bool {
-					switch expr := cursor.Node().(type) {
-					case *sqlparser.ColName:
-						if expr.Name.String() != sq.ArgName { // TODO systay 2023.09.15 - This is not safe enough. We should figure out a better way.
+
+	merged := true
+	for merged {
+		// we need to keep rewriting the expression until we can't find any more subqueries to merge
+		// this is because we might have subqueries inside subqueries, and we need to merge them all
+		merged = false
+		for _, sq := range se {
+			for _, sq2 := range ctx.MergedSubqueries {
+				if sq.originalSubquery == sq2 {
+					expr = sqlparser.Rewrite(expr, nil, func(cursor *sqlparser.Cursor) bool {
+						switch expr := cursor.Node().(type) {
+						case *sqlparser.ColName:
+							if expr.Name.String() != sq.ArgName { // TODO systay 2023.09.15 - This is not safe enough. We should figure out a better way.
+								return true
+							}
+						case *sqlparser.Argument:
+							if expr.Name != sq.ArgName {
+								return true
+							}
+						default:
 							return true
 						}
-					case *sqlparser.Argument:
-						if expr.Name != sq.ArgName {
-							return true
+						rewritten = true
+						if sq.FilterType == opcode.PulloutExists {
+							cursor.Replace(&sqlparser.ExistsExpr{Subquery: sq.originalSubquery})
+						} else {
+							cursor.Replace(sq.originalSubquery)
 						}
-					default:
-						return true
-					}
-					rewritten = true
-					if sq.FilterType == opcode.PulloutExists {
-						cursor.Replace(&sqlparser.ExistsExpr{Subquery: sq.originalSubquery})
-					} else {
-						cursor.Replace(sq.originalSubquery)
-					}
-					return false
-				}).(sqlparser.Expr)
+						merged = true
+						return false
+					}).(sqlparser.Expr)
+				}
 			}
 		}
 	}
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
index e41d84adc69..edebb8a6119 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json
@@ -1228,6 +1228,58 @@
       ]
     }
   },
+  {
+    "comment": "Subquery inside subquery #1",
+    "query": "select (select (select col from `user` where id = 1) from `user` where id = 1) from `user` where id = 1",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select (select (select col from `user` where id = 1) from `user` where id = 1) from `user` where id = 1",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select (select (select col from `user` where 1 != 1) from `user` where 1 != 1) from `user` where 1 != 1",
+        "Query": "select (select (select col from `user` where id = 1) from `user` where id = 1) from `user` where id = 1",
+        "Table": "`user`",
+        "Values": [
+          "1"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "Subquery inside subquery #2",
+    "query": "select 1 from `user` where id = 1 and 1 = (select (select intcol from `user` where id = 1) from `user` where id = 1)",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select 1 from `user` where id = 1 and 1 = (select (select intcol from `user` where id = 1) from `user` where id = 1)",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select 1 from `user` where 1 != 1",
+        "Query": "select 1 from `user` where id = 1 and 1 = (select (select intcol from `user` where id = 1) from `user` where id = 1)",
+        "Table": "`user`",
+        "Values": [
+          "1"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "Multiple parenthesized expressions",
     "query": "select * from user where (id = 4 and name ='abc') limit 5",