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:

Patch:
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index ef38934a..f2f076e0 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -109,10 +109,7 @@ def lineage(
             # 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
index b9289385..f33a2c2d 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -146,17 +146,43 @@ class TestLineage(unittest.TestCase):
         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, [])