Repository: tobymao/sqlglot

Issue:
Issue when calculating lineage when a CTE has the same name as a table from the schema
**Fully reproducible code snippet**

The following code 
```sql
from sqlglot.lineage import lineage

selected_column = "col_a"
sql = """
with 

my_cte_name_also_a_table_name as (
    select * from raw.schema.my_table
)

, inter as (
    select * from my_cte_name_also_a_table_name
)

select * from inter
"""

schema = {"raw": {"schema" : {"my_cte_name_also_a_table_name": {"col_1": "int"}, "my_table": {"col_a": "int"}}}}
schema_without_table = {"my_table": {"col_a": "int"}}

l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
```

returns the error
```
Traceback (most recent call last):
  File "/xxx/short_issue_lineage.py", line 21, in <module>
    l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 148, in lineage
    return to_node(column if isinstance(column, str) else column.name, scope)
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 136, in to_node
    to_node(
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 112, in to_node
    source = optimize(
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/optimizer.py", line 89, in optimize
    expression = rule(expression, **rule_kwargs)
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py", line 49, in qualify_columns
    _qualify_columns(scope, resolver)
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py", line 250, in _qualify_columns
    raise OptimizeError(f"Unknown column: {column_name}")
sqlglot.errors.OptimizeError: Unknown column: col_a
```


It looks like there is an issue with the logic being confused between `my_cte_name_also_a_table_name` the CTE (with a column called `col_a`) and `my_cte_name_also_a_table_name` the table from the schema with a column called `col_1`.

With the example above, if I remove the table from the schema and calculate the lineage, the error goes away.
```
l = lineage(column=selected_column, sql=sql, schema=schema_without_table, dialect="snowflake")
```


Interestingly, the code without the intermediate cte `inter` calculates the lineage correctly when the table with the CTE name is provided in the schema.
```
sql = """
with 

my_cte_name_also_a_table_name as (
    select * from raw.schema.my_table
)

select * from my_cte_name_also_a_table_name
"""
l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
```


Discussion:
You can't give a CTE a schema, schema is only used for "external sources". Remove the cte from the schema. Happy to discuss more, feel free to join our slack. 
Thanks for the answer!
What you describe is what I see when I used `schema_without_table` instead of `schema`, which works as you mention.

The overall issue I was raising was when we provide the schema for a table that has the same name as the CTE (even though the table is not part of what the `sql` analyzed reads from). My intent was not to provide the schema for the CTE but for another table with the same name.
@b-per can you provide an example of this?
Sure, let me try to clarify with a less abstract example

### Situation

Let's say that I have 2 tables in my warehouse:

1. a table called `customers` coming from some EL tool. It has 2 columns, `id` and `sales`.
1. a "staging" table/view called `stg_customers`, reading from `customers` and renaming `id` to `customer_id` and name it 

![image](https://github.com/tobymao/sqlglot/assets/8754100/025deb16-64ed-4875-a2fd-d8341a7de0b5)

I now want to build another table/view that reads from `stg_customers`.


### Issue

The following code raises the issue I showed earlier. This is because I used the CTE name `customers` which is also the name of a table in my `schema` variable.

```python
from sqlglot.lineage import lineage

selected_column = "customer_id"
sql = """
with 

customers as (
    select * from stg_customers
)

, final as (
    select * from customers
)

select customer_id from final
"""


schema = {
    "my_db1": {
        "my_schema1" : {
            "customers": {
                "id": "int", 
                "sales": "int"
            }
        },
    },
    "my_db2" : {
        "my_schema2": { 
            "stg_customers": {
                "customer_id": "int",
                "sales": "int"
            }
        }
    }
}

l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
```

Error: `sqlglot.errors.OptimizeError: Unknown column: customer_id`

### Investigation

From debugging a bit sqlglot, it looks like at some point in the code `resolver.get_source_columns(column_table)` looks for the column in `my_db1.my_schema1.customers` rather than at the columns from the CTE called `customers`.


### Workaround

As you described and as I saw, a workaround to this particular issue is to remove the `customers` table from the `schema` variable, but in that case I need to have a `schema` variable that only represents the table being read by the `sql` I am analyzing. 

My intent was to use a `schema` variable that would contain all the tables from my warehouse, without the need to know which ones are used or not in the `sql` I  am analyzing.
I'm looking into this, thanks for the report @b-per

Resolving patch:
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -109,10 +109,7 @@ def to_node(
             # a version that has only the column we care about.
             #   "x", SELECT x, y FROM foo
             #     => "x", SELECT x FROM foo
-            source = optimize(
-                scope.expression.select(select, append=False), schema=schema, rules=rules
-            )
-            select = source.selects[0]
+            source = t.cast(exp.Expression, scope.expression.select(select, append=False))
         else:
             source = scope.expression
 


Test patch:
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -146,17 +146,43 @@ def test_lineage_values(self) -> None:
         self.assertEqual(node.alias, "")
 
         downstream = node.downstream[0]
-        self.assertEqual(
-            downstream.source.sql(),
-            "SELECT t.a AS a FROM (VALUES (1), (2)) AS t(a)",
-        )
+        self.assertEqual(downstream.source.sql(), "SELECT t.a AS a FROM (VALUES (1), (2)) AS t(a)")
         self.assertEqual(downstream.expression.sql(), "t.a AS a")
         self.assertEqual(downstream.alias, "y")
 
         downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "(VALUES (1), (2)) AS t(a)")
+        self.assertEqual(downstream.expression.sql(), "a")
+        self.assertEqual(downstream.alias, "")
+
+    def test_lineage_cte_name_appears_in_schema(self) -> None:
+        schema = {"a": {"b": {"t1": {"c1": "int"}, "t2": {"c2": "int"}}}}
+
+        node = lineage(
+            "c2",
+            "WITH t1 AS (SELECT * FROM a.b.t2), inter AS (SELECT * FROM t1) SELECT * FROM inter",
+            schema=schema,
+        )
+
         self.assertEqual(
-            downstream.source.sql(),
-            "(VALUES (1), (2)) AS t(a)",
+            node.source.sql(),
+            "WITH t1 AS (SELECT t2.c2 AS c2 FROM a.b.t2 AS t2), inter AS (SELECT t1.c2 AS c2 FROM t1) SELECT inter.c2 AS c2 FROM inter",
         )
-        self.assertEqual(downstream.expression.sql(), "a")
+        self.assertEqual(node.alias, "")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.source.sql(), "SELECT t1.c2 AS c2 FROM t1")
+        self.assertEqual(downstream.expression.sql(), "t1.c2 AS c2")
         self.assertEqual(downstream.alias, "")
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "SELECT t2.c2 AS c2 FROM a.b.t2 AS t2")
+        self.assertEqual(downstream.expression.sql(), "t2.c2 AS c2")
+        self.assertEqual(downstream.alias, "")
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "a.b.t2 AS t2")
+        self.assertEqual(downstream.expression.sql(), "a.b.t2 AS t2")
+        self.assertEqual(downstream.alias, "")
+
+        self.assertEqual(downstream.downstream, [])