Fix(lineage): remove unnecessary optimization (tobymao#1663)

adrianisk · Jun 21, 2023 · 4aaeacb · 4aaeacb
1 parent a47de3f
commit 4aaeacb
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 11 deletions.
diff --git 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
 

diff --git 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, [])