fix(optimizer): unions on nested subqueries (#4603)

sqlglot/optimizer/scope.py

@@ -146,9 +146,7 @@ def _collect(self):
                 self._udtfs.append(node)
             elif isinstance(node, exp.CTE):
                 self._ctes.append(node)
-            elif _is_derived_table(node) and isinstance(
-                node.parent, (exp.From, exp.Join, exp.Subquery)
-            ):
+            elif _is_derived_table(node) and _is_from_or_join(node):
                 self._derived_tables.append(node)
             elif isinstance(node, exp.UNWRAPPED_QUERIES):
                 self._subqueries.append(node)
@@ -661,6 +659,19 @@ def _is_derived_table(expression: exp.Subquery) -> bool:
     )
 
 
+def _is_from_or_join(expression: exp.Expression) -> bool:
+    """
+    Determine if `expression` is the FROM or JOIN clause of a SELECT statement.
+    """
+    parent = expression.parent
+
+    # Subqueries can be arbitrarily nested
+    while isinstance(parent, exp.Subquery):
+        parent = parent.parent
+
+    return isinstance(parent, (exp.From, exp.Join))
+
+
 def _traverse_tables(scope):
     sources = {}
 

tests/fixtures/optimizer/qualify_columns.sql

@@ -277,6 +277,10 @@ SELECT x.a AS a FROM x AS x UNION SELECT x.a AS a FROM x AS x UNION SELECT x.a A
 SELECT a FROM (SELECT a FROM x UNION SELECT a FROM x) ORDER BY a;
 SELECT _q_0.a AS a FROM (SELECT x.a AS a FROM x AS x UNION SELECT x.a AS a FROM x AS x) AS _q_0 ORDER BY a;
 
+# title: nested subqueries in union
+((select a from x where a < 1)) UNION ((select a from x where a > 2));
+((SELECT x.a AS a FROM x AS x WHERE x.a < 1)) UNION ((SELECT x.a AS a FROM x AS x WHERE x.a > 2));
+
 --------------------------------------
 -- Subqueries
 --------------------------------------