feat: support WITH with UNION

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

go/vt/sqlparser/ast.go

@@ -293,11 +293,11 @@ type (
 
 	// Union represents a UNION statement.
 	Union struct {
+		With     *With
 		Left     SelectStatement
 		Right    SelectStatement
 		Distinct bool
 		OrderBy  OrderBy
-		With     *With
 		Limit    *Limit
 		Lock     Lock
 		Into     *SelectInto

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

@@ -4050,5 +4050,59 @@
     "comment": "select with a target destination",
     "query": "select * from `user[-]`.user_metadata",
     "plan": "VT09017: SELECT with a target destination is not allowed"
+  },
+  {
+    "comment": "simple WITH query",
+    "query": "with x as (select * from user) select * from x",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "with x as (select * from user) select * from x",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select * from (select * from `user` where 1 != 1) as x where 1 != 1",
+        "Query": "select * from (select * from `user`) as x",
+        "Table": "`user`"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "UNION with WITH clause",
+    "query": "with x as (select id, foo from user) select * from x union select * from x",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "with x as (select id, foo from user) select * from x union select * from x",
+      "Instructions": {
+        "OperatorType": "Distinct",
+        "Collations": [
+          "(0:2)",
+          "(1:3)"
+        ],
+        "ResultColumns": 2,
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "Scatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select id, foo, weight_string(id), weight_string(foo) from (select id, foo from (select id, foo from `user` where 1 != 1) as x where 1 != 1 union select id, foo from (select id, foo from `user` where 1 != 1) as x where 1 != 1) as dt where 1 != 1",
+            "Query": "select id, foo, weight_string(id), weight_string(foo) from (select id, foo from (select id, foo from `user`) as x union select id, foo from (select id, foo from `user`) as x) as dt",
+            "Table": "`user`"
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
   }
 ]

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

@@ -254,16 +254,6 @@
     "query": "with x as (select * from user) update x set name = 'f'",
     "plan": "VT12001: unsupported: WITH expression in UPDATE statement"
   },
-  {
-    "comment": "unsupported with clause in select statement",
-    "query": "with x as (select * from user) select * from x",
-    "plan": "VT12001: unsupported: WITH expression in SELECT statement"
-  },
-  {
-    "comment": "unsupported with clause in union statement",
-    "query": "with x as (select * from user) select * from x union select * from x",
-    "plan": "VT12001: unsupported: WITH expression in UNION statement"
-  },
   {
     "comment": "insert having subquery in row values",
     "query": "insert into user(id, name) values ((select 1 from user where id = 1), 'A')",

go/vt/vtgate/semantics/early_rewriter.go

@@ -60,24 +60,29 @@ func (r *earlyRewriter) down(cursor *sqlparser.Cursor) error {
 	case *sqlparser.With:
 		return r.handleWith(node)
 	case *sqlparser.AliasedTableExpr:
-		tbl, ok := node.Expr.(sqlparser.TableName)
-		if !ok || !tbl.Qualifier.IsEmpty() {
-			return nil
-		}
-		scope := r.scoper.currentScope()
-		cte := scope.findCTE(tbl.Name.String())
-		if cte == nil {
-			return nil
-		}
-		if node.As.IsEmpty() {
-			node.As = tbl.Name
-		}
-		node.Expr = &sqlparser.DerivedTable{
-			Select: cte.Subquery.Select,
-		}
-		if len(cte.Columns) > 0 {
-			node.Columns = cte.Columns
-		}
+		return r.handleAliasedTable(node)
+	}
+	return nil
+}
+
+func (r *earlyRewriter) handleAliasedTable(node *sqlparser.AliasedTableExpr) error {
+	tbl, ok := node.Expr.(sqlparser.TableName)
+	if !ok || !tbl.Qualifier.IsEmpty() {
+		return nil
+	}
+	scope := r.scoper.currentScope()
+	cte := scope.findCTE(tbl.Name.String())
+	if cte == nil {
+		return nil
+	}
+	if node.As.IsEmpty() {
+		node.As = tbl.Name
+	}
+	node.Expr = &sqlparser.DerivedTable{
+		Select: cte.Subquery.Select,
+	}
+	if len(cte.Columns) > 0 {
+		node.Columns = cte.Columns
 	}
 	return nil
 }

go/vt/vtgate/semantics/scoper.go

@@ -66,6 +66,8 @@ func (s *scoper) down(cursor *sqlparser.Cursor) error {
 		s.pushDMLScope(node)
 	case *sqlparser.Select:
 		s.pushSelectScope(node)
+	case *sqlparser.Union:
+		s.pushUnionScope(node)
 	case sqlparser.TableExpr:
 		s.enterJoinScope(cursor)
 	case sqlparser.SelectExprs:
@@ -75,14 +77,21 @@ func (s *scoper) down(cursor *sqlparser.Cursor) error {
 	case sqlparser.GroupBy:
 		return s.addColumnInfoForGroupBy(cursor, node)
 	case *sqlparser.Where:
-		if node.Type != sqlparser.HavingClause {
-			break
+		if node.Type == sqlparser.HavingClause {
+			return s.createSpecialScopePostProjection(cursor.Parent())
 		}
-		return s.createSpecialScopePostProjection(cursor.Parent())
 	}
 	return nil
 }
 
+func (s *scoper) pushUnionScope(union *sqlparser.Union) {
+	currentScope := s.currentScope()
+	currScope := newScope(currentScope)
+	currScope.stmtScope = true
+	currScope.stmt = union
+	s.push(currScope)
+}
+
 func (s *scoper) addColumnInfoForGroupBy(cursor *sqlparser.Cursor, node sqlparser.GroupBy) error {
 	err := s.createSpecialScopePostProjection(cursor.Parent())
 	if err != nil {