Fix eliminate_subqueries on unions (#1098)

tobymao · Feb 6, 2023 · b5df65e · b5df65e
1 parent 6e182cf
commit b5df65e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/sqlglot/optimizer/eliminate_subqueries.py b/sqlglot/optimizer/eliminate_subqueries.py
@@ -114,7 +114,7 @@ def _eliminate_union(scope, existing_ctes, taken):
     taken[alias] = scope
 
     # Try to maintain the selections
-    expressions = scope.expression.args.get("expressions")
+    expressions = scope.selects
     selects = [
         exp.alias_(exp.column(e.alias_or_name, table=alias), alias=e.alias_or_name)
         for e in expressions

diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
@@ -300,7 +300,7 @@ def selects(self):
             list[exp.Expression]: expressions
         """
         if isinstance(self.expression, exp.Union):
-            return []
+            return self.expression.unnest().selects
         return self.expression.selects
 
     @property

diff --git a/tests/fixtures/optimizer/eliminate_subqueries.sql b/tests/fixtures/optimizer/eliminate_subqueries.sql
@@ -50,6 +50,10 @@ WITH cte AS (SELECT 1 AS x, 2 AS y) SELECT cte.x AS x, cte.y AS y FROM cte AS ct
 (SELECT a FROM (SELECT b FROM x)) UNION (SELECT a FROM (SELECT b FROM y));
 WITH cte AS (SELECT b FROM x), cte_2 AS (SELECT a FROM cte AS cte), cte_3 AS (SELECT b FROM y), cte_4 AS (SELECT a FROM cte_3 AS cte_3) (SELECT cte_2.a AS a FROM cte_2 AS cte_2) UNION (SELECT cte_4.a AS a FROM cte_4 AS cte_4);
 
+-- Three unions
+SELECT a FROM x UNION ALL SELECT a FROM y UNION ALL SELECT a FROM z;
+WITH cte AS (SELECT a FROM x), cte_2 AS (SELECT a FROM y), cte_3 AS (SELECT a FROM z), cte_4 AS (SELECT cte_2.a AS a FROM cte_2 AS cte_2 UNION ALL SELECT cte_3.a AS a FROM cte_3 AS cte_3) SELECT cte.a AS a FROM cte AS cte UNION ALL SELECT cte_4.a AS a FROM cte_4 AS cte_4;
+
 -- Subquery
 SELECT a FROM x WHERE b = (SELECT y.c FROM y);
 SELECT a FROM x WHERE b = (SELECT y.c FROM y);