feat: fail correlated subqueries if we are unable to merge them in a …

…route

Signed-off-by: Manan Gupta <manan@planetscale.com>

go/vt/vtgate/planbuilder/operators/ast_to_op.go

@@ -108,23 +108,24 @@ func findTablesContained(ctx *plancontext.PlanningContext, node sqlparser.SQLNod
 	return
 }
 
-func rewriteRemainingColumns(
+func checkForCorrelatedSubqueries(
 	ctx *plancontext.PlanningContext,
 	stmt sqlparser.SelectStatement,
 	subqID semantics.TableSet,
-) sqlparser.SelectStatement {
-	return sqlparser.CopyOnRewrite(stmt, nil, func(cursor *sqlparser.CopyOnWriteCursor) {
-		colname, isColname := cursor.Node().(*sqlparser.ColName)
+) (correlated bool) {
+	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		colname, isColname := node.(*sqlparser.ColName)
 		if !isColname {
-			return
+			return true, nil
 		}
 		deps := ctx.SemTable.RecursiveDeps(colname)
 		if deps.IsSolvedBy(subqID) {
-			return
+			return true, nil
 		}
-		rsv := ctx.GetReservedArgumentFor(colname)
-		cursor.Replace(sqlparser.NewArgument(rsv))
-	}, nil).(sqlparser.SelectStatement)
+		correlated = true
+		return false, nil
+	}, stmt)
+	return correlated
 }
 
 // joinPredicateCollector is used to inspect the predicates inside the subquery, looking for any

go/vt/vtgate/planbuilder/operators/subquery.go

@@ -49,6 +49,9 @@ type SubQuery struct {
 	// Fields related to correlated subqueries:
 	Vars    map[string]int // Arguments copied from outer to inner, set during offset planning.
 	outerID semantics.TableSet
+	// correlated stores whether this subquery is correlated or not.
+	// We use this information to fail the planning if we are unable to merge the subquery with a route.
+	correlated bool
 
 	IsProjection bool
 }
@@ -200,6 +203,9 @@ func (sq *SubQuery) settle(ctx *plancontext.PlanningContext, outer Operator) (Op
 	if !sq.TopLevel {
 		return nil, subqueryNotAtTopErr
 	}
+	if sq.correlated {
+		return nil, correlatedSubqueryErr
+	}
 	if sq.IsProjection {
 		if len(sq.GetMergePredicates()) > 0 {
 			// this means that we have a correlated subquery on our hands

go/vt/vtgate/planbuilder/operators/subquery_builder.go

@@ -169,12 +169,9 @@ func createSubquery(
 	sqc := &SubQueryBuilder{totalID: totalID, subqID: subqID, outerID: outerID}
 
 	predicates, joinCols := sqc.inspectStatement(ctx, subq.Select)
-	stmt := rewriteRemainingColumns(ctx, subq.Select, subqID)
+	correlated := checkForCorrelatedSubqueries(ctx, subq.Select, subqID)
 
-	// TODO: this should not be needed. We are using CopyOnRewrite above, but somehow this is not getting copied
-	ctx.SemTable.CopySemanticInfo(subq.Select, stmt)
-
-	opInner := translateQueryToOp(ctx, stmt)
+	opInner := translateQueryToOp(ctx, subq.Select)
 
 	opInner = sqc.getRootOperator(opInner, nil)
 	return &SubQuery{
@@ -187,6 +184,7 @@ func createSubquery(
 		IsProjection:     isProjection,
 		TopLevel:         topLevel,
 		JoinColumns:      joinCols,
+		correlated:       correlated,
 	}
 }
 

go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json

@@ -1037,10 +1037,10 @@
   },
   {
     "comment": "merge even one side have schema name in subquery",
-    "query": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `COLUMN_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
+    "query": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `TABLE_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
     "plan": {
       "QueryType": "SELECT",
-      "Original": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `COLUMN_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
+      "Original": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `TABLE_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
       "Instructions": {
         "OperatorType": "UncorrelatedSubquery",
         "Variant": "PulloutIn",
@@ -1066,8 +1066,8 @@
                       "Name": "main",
                       "Sharded": false
                     },
-                    "FieldQuery": "select :COLUMN_NAME from information_schema.`tables` as t where 1 != 1",
-                    "Query": "select distinct :COLUMN_NAME from information_schema.`tables` as t where t.TABLE_SCHEMA = :__vtschemaname /* VARCHAR */",
+                    "FieldQuery": "select TABLE_NAME from information_schema.`tables` as t where 1 != 1",
+                    "Query": "select distinct TABLE_NAME from information_schema.`tables` as t where t.TABLE_SCHEMA = :__vtschemaname /* VARCHAR */",
                     "SysTableTableSchema": "['a']",
                     "Table": "information_schema.`tables`"
                   },

go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json

@@ -1102,10 +1102,10 @@
   },
   {
     "comment": "merge even one side have schema name in subquery",
-    "query": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `COLUMN_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
+    "query": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `TABLE_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
     "plan": {
       "QueryType": "SELECT",
-      "Original": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `COLUMN_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
+      "Original": "select `COLLATION_NAME` from information_schema.`COLUMNS` t where `COLUMN_NAME` in (select `TABLE_NAME` from information_schema.tables t where t.TABLE_SCHEMA = 'a' union select `COLUMN_NAME` from information_schema.columns)",
       "Instructions": {
         "OperatorType": "UncorrelatedSubquery",
         "Variant": "PulloutIn",
@@ -1131,8 +1131,8 @@
                       "Name": "main",
                       "Sharded": false
                     },
-                    "FieldQuery": "select :COLUMN_NAME from information_schema.`tables` as t where 1 != 1",
-                    "Query": "select distinct :COLUMN_NAME from information_schema.`tables` as t where t.TABLE_SCHEMA = :__vtschemaname /* VARCHAR */",
+                    "FieldQuery": "select TABLE_NAME from information_schema.`tables` as t where 1 != 1",
+                    "Query": "select distinct TABLE_NAME from information_schema.`tables` as t where t.TABLE_SCHEMA = :__vtschemaname /* VARCHAR */",
                     "SysTableTableSchema": "['a']",
                     "Table": "information_schema.`tables`"
                   },

go/vt/vtgate/planbuilder/testdata/unsupported_cases.json

@@ -378,5 +378,10 @@
     "comment": "Alias cannot clash with base tables",
     "query": "WITH user AS (SELECT col FROM user) SELECT * FROM user",
     "plan": "VT12001: unsupported: do not support CTE that use the CTE alias inside the CTE query"
+  },
+  {
+    "comment": "correlated subqueries in select expressions are unsupported",
+    "query": "SELECT (SELECT sum(user.name) FROM music LIMIT 1) FROM user",
+    "plan": "VT12001: unsupported: correlated subquery is only supported for EXISTS"
   }
 ]