diff --git a/sqlglot/optimizer/eliminate_subqueries.py b/sqlglot/optimizer/eliminate_subqueries.py index 2245cc2761..c6bea5a849 100644 --- 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 index 5a3ed5ae9b..badbb87cd6 100644 --- 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 index c566657299..4fa63dd5d1 100644 --- 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);