From b5df65e3fb5ee1ebc3cbab64b6d89598cf47a10b Mon Sep 17 00:00:00 2001
From: barak <barak.alon@airbnb.com>
Date: Sun, 5 Feb 2023 18:54:34 -0800
Subject: [PATCH] Fix eliminate_subqueries on unions (#1098)

---
 sqlglot/optimizer/eliminate_subqueries.py         | 2 +-
 sqlglot/optimizer/scope.py                        | 2 +-
 tests/fixtures/optimizer/eliminate_subqueries.sql | 4 ++++
 3 files changed, 6 insertions(+), 2 deletions(-)

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);