Repository: tobymao/sqlglot

Issue:
Bug report : OptimizerError thrown by qualify when INTERSECT, UNION, or EXCEPT used in uncorrelated subquery
## Summary / Background

When using `qualify` with `'mysql'` dialect there is an unexpected error. The following code snippet throws an OptimizerError due to an unresolved column despite all columns being well defined. It appears as though quoting the identifier "B" resolves the issue but this should not be required since the query is unambiguous. 

## Steps to Reproduce

To reproduce, run the following code snippet with python3.11 and sqlglot==20.8.0.

```python3
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

schema = {
        '' : {
            'table_a' : {
                'col1' : 'INT'
            },
            'table_b' : {
                'col2' : 'INT'
            }
        }
    }

query = """
SELECT *
FROM 
"table_a" as "A"
WHERE "A"."col1" IN (
    SELECT MAX(B."col2") FROM "table_b" as "B"
    UNION ALL
    SELECT MAX("C"."col2") FROM "table_b" as "C"
);"""

dialect='mysql'
qualify(parse_one(query, dialect=dialect), schema=schema, dialect=dialect)
# throws error : OptimizeError: Column '"B".'col2'' could not be resolved for table: 'B'
# thrown by validate_qualify_columns
```

## Documentation Links
[qualify documentation](https://sqlglot.com/sqlglot/optimizer/qualify.html#qualify)
[MySQL dialect documentation](https://sqlglot.com/sqlglot/dialects/mysql.html#MySQL)
[validate_qualify_columns documentation](https://sqlglot.com/sqlglot/optimizer/qualify_columns.html#validate_qualify_columns)

## Additional Information 

When investigating the source of the exception, it appears as though the bug is in the `scope` code.  

From the source file `sqlglot/optimizer/qualify_columns.py` line 83.
```python3
            if scope.external_columns and not scope.is_correlated_subquery and not scope.pivots:
                column = scope.external_columns[0]
                for_table = f" for table: '{column.table}'" if column.table else ""
                raise OptimizeError(f"Column '{column}' could not be resolved{for_table}")
```
Based of the if statement, the scope has `B."col2"` as an external column (which I believe is incorrect), the subquery as uncorrelated (correct) and the scope containing no pivots (also correct).

Below are additional code snippets with the output which are possibly useful for debugging the issue. If you quote "B" or remove `UNION ALL` the validation passes and no error is thrown.

```python3
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

schema = {
        '' : {
            'table_a' : {
                'col1' : 'INT'
            },
            'table_b' : {
                'col2' : 'INT'
            }
        }
    }

query = """
SELECT *
FROM 
"table_a" as "A"
WHERE "A"."col1" IN (
    SELECT MAX("B"."col2") FROM "table_b" as "B"
    UNION ALL
    SELECT MAX("C"."col2") FROM "table_b" as "C"
);"""

dialect='mysql'
qualify(parse_one(query, dialect=dialect), schema=schema, dialect=dialect)
# Okay, no error thrown if we quote "B"
```
```python3
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

schema = {
        '' : {
            'table_a' : {
                'col1' : 'INT'
            },
            'table_b' : {
                'col2' : 'INT'
            }
        }
    }

query = """
SELECT *
FROM 
"table_a" as "A"
WHERE "A"."col1" IN (
    SELECT MAX(B."col2") FROM "table_b" as "B"
);"""

dialect='mysql'
qualify(parse_one(query, dialect=dialect), schema=schema, dialect=dialect)
# Okay, without UNION ALL there is no need to quote "B".
```

Thank you for the help, please let me know if there is any other information that I can provide.


Discussion:

Resolving patch:
diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
--- a/sqlglot/optimizer/scope.py
+++ b/sqlglot/optimizer/scope.py
@@ -323,9 +323,14 @@ def external_columns(self):
                 sources in the current scope.
         """
         if self._external_columns is None:
-            self._external_columns = [
-                c for c in self.columns if c.table not in self.selected_sources
-            ]
+            if isinstance(self.expression, exp.Union):
+                left, right = self.union_scopes
+                self._external_columns = left.external_columns + right.external_columns
+            else:
+                self._external_columns = [
+                    c for c in self.columns if c.table not in self.selected_sources
+                ]
+
         return self._external_columns
 
     @property


Test patch:
diff --git a/tests/fixtures/optimizer/qualify_columns.sql b/tests/fixtures/optimizer/qualify_columns.sql
--- a/tests/fixtures/optimizer/qualify_columns.sql
+++ b/tests/fixtures/optimizer/qualify_columns.sql
@@ -230,6 +230,11 @@ SELECT (SELECT y.c AS c FROM y AS y) AS _col_0 FROM x AS x;
 SELECT a FROM (SELECT a FROM x) WHERE a IN (SELECT b FROM (SELECT b FROM y));
 SELECT _q_1.a AS a FROM (SELECT x.a AS a FROM x AS x) AS _q_1 WHERE _q_1.a IN (SELECT _q_0.b AS b FROM (SELECT y.b AS b FROM y AS y) AS _q_0);
 
+# dialect: mysql
+# execute: false
+SELECT * FROM table_a as A WHERE A.col1 IN (SELECT MAX(B.col2) FROM table_b as B UNION ALL SELECT MAX(C.col2) FROM table_b as C);
+SELECT * FROM table_a AS `A` WHERE `A`.col1 IN (SELECT MAX(`B`.col2) AS _col_0 FROM table_b AS `B` UNION ALL SELECT MAX(`C`.col2) AS _col_0 FROM table_b AS `C`);
+
 --------------------------------------
 -- Correlated subqueries
 --------------------------------------