Repository: tobymao/sqlglot
Issue category: Bug Report

Issue:
`sqlglot.optimizer.qualify.qualify` fails when using `CONNECT BY` + `level` from a CTE
## Fully reproducible code snippet
In Snowflake, the `SELECT FROM <data_source> START WITH ... CONNECT BY` construct allows a `<level_expression>` by basically allowing to select a `level` column in addition to the existing columns:

```sql
SELECT
  employee_ID,
  manager_ID,
  title,
  level --<-- This isn't a column from "employees", but a column allowed by Snowflake in `CONNECT BY` constructs
FROM employees
START WITH title = 'President'
CONNECT BY manager_ID = PRIOR employee_id
ORDER BY
  employee_ID NULLS LAST
```

`sqlglot.optimizer.qualify.qualify` handles this fine ✅

However, if the `SELECT` is `FROM` a CTE, it fails!

Taking the same example as before **but** putting all the data in a CTE (instead of a table) breaks `qualify`:
```py
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

SQL_DIALECT = "snowflake"

print(qualify(parse_one("""
WITH employees AS (
    SELECT *
    FROM (VALUES ('President', 1, NULL),
                 ('Vice President Engineering', 10, 1),
                 ('Programmer', 100, 10),
                 ('QA Engineer', 101, 10),
                 ('Vice President HR', 20, 1),
                 ('Health Insurance Analyst', 200, 20)
    ) AS t(title, employee_ID, manager_ID)
)
SELECT
  employee_ID,
  manager_ID,
  title,
  level
FROM employees
START WITH title = 'President'
CONNECT BY manager_ID = PRIOR employee_id
ORDER BY
  employee_ID NULLS LAST
""", dialect="snowflake"), schema=None, dialect="snowflake").sql(pretty=True, dialect=SQL_DIALECT))
```

Output:
```
Traceback (most recent call last):
  File "test_sqlglot.py", line 7, in <module>
    print(qualify(parse_one("""
          ^^^^^^^^^^^^^^^^^^^^^
  File "/sqlglot/optimizer/qualify.py", line 102, in qualify
    validate_qualify_columns_func(expression)
  File "sqlglot/optimizer/qualify_columns.py", line 109, in validate_qualify_columns
    raise OptimizeError(f"Column '{column}' could not be resolved{for_table}")
sqlglot.errors.OptimizeError: Column '"LEVEL"' could not be resolved
```


## Official Documentation
(see [documentation](https://docs.snowflake.com/en/sql-reference/constructs/connect-by#examples)).


Hints:

Patch:
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
index ceb9ceb44f..871bb365ba 100644
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -898,10 +898,22 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequenc
                     for (name, alias) in itertools.zip_longest(columns, column_aliases)
                 ]
 
+            pseudocolumns = self._get_source_pseudocolumns(name)
+            if pseudocolumns:
+                columns = list(columns)
+                columns.extend(c for c in pseudocolumns if c not in columns)
+
             self._get_source_columns_cache[cache_key] = columns
 
         return self._get_source_columns_cache[cache_key]
 
+    def _get_source_pseudocolumns(self, name: str) -> t.Sequence[str]:
+        if self.schema.dialect == "snowflake" and self.scope.expression.args.get("connect"):
+            # When there is a CONNECT BY clause, there is only one table being scanned
+            # See: https://docs.snowflake.com/en/sql-reference/constructs/connect-by
+            return ["LEVEL"]
+        return []
+
     def _get_all_source_columns(self) -> t.Dict[str, t.Sequence[str]]:
         if self._source_columns is None:
             self._source_columns = {


Test patch:
diff --git a/tests/fixtures/optimizer/qualify_columns.sql b/tests/fixtures/optimizer/qualify_columns.sql
index eeaf8b3555..ecb6eee5ce 100644
--- a/tests/fixtures/optimizer/qualify_columns.sql
+++ b/tests/fixtures/optimizer/qualify_columns.sql
@@ -205,6 +205,30 @@ SELECT x.a + 1 AS i, missing_column AS missing_column FROM x AS x;
 SELECT s, arr1, arr2 FROM arrays_test LEFT ARRAY JOIN arr1, arrays_test.arr2;
 SELECT arrays_test.s AS s, arrays_test.arr1 AS arr1, arrays_test.arr2 AS arr2 FROM arrays_test AS arrays_test LEFT ARRAY JOIN arrays_test.arr1, arrays_test.arr2;
 
+# execute: false
+# dialect: snowflake
+WITH employees AS (
+    SELECT *
+    FROM (VALUES ('President', 1, NULL),
+                 ('Vice President Engineering', 10, 1),
+                 ('Programmer', 100, 10),
+                 ('QA Engineer', 101, 10),
+                 ('Vice President HR', 20, 1),
+                 ('Health Insurance Analyst', 200, 20)
+    ) AS t(title, employee_ID, manager_ID)
+)
+SELECT
+  employee_ID,
+  manager_ID,
+  title,
+  level
+FROM employees
+START WITH title = 'President'
+CONNECT BY manager_ID = PRIOR employee_id
+ORDER BY
+  employee_ID NULLS LAST;
+WITH EMPLOYEES AS (SELECT T.TITLE AS TITLE, T.EMPLOYEE_ID AS EMPLOYEE_ID, T.MANAGER_ID AS MANAGER_ID FROM (VALUES ('President', 1, NULL), ('Vice President Engineering', 10, 1), ('Programmer', 100, 10), ('QA Engineer', 101, 10), ('Vice President HR', 20, 1), ('Health Insurance Analyst', 200, 20)) AS T(TITLE, EMPLOYEE_ID, MANAGER_ID)) SELECT EMPLOYEES.EMPLOYEE_ID AS EMPLOYEE_ID, EMPLOYEES.MANAGER_ID AS MANAGER_ID, EMPLOYEES.TITLE AS TITLE, EMPLOYEES.LEVEL AS LEVEL FROM EMPLOYEES AS EMPLOYEES START WITH EMPLOYEES.TITLE = 'President' CONNECT BY EMPLOYEES.MANAGER_ID = PRIOR EMPLOYEES.EMPLOYEE_ID ORDER BY EMPLOYEE_ID;
+
 --------------------------------------
 -- Derived tables
 --------------------------------------