1 files changed, 78 insertions, 0 deletions
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index ed1a448..c782d9a 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -269,6 +269,41 @@ class TestLineage(unittest.TestCase):
         node = node.downstream[0]
         self.assertEqual(node.name, "z.a")
 
+        node = lineage(
+            "a",
+            """
+            WITH foo AS (
+              SELECT
+                1 AS a
+            ), bar AS (
+              (
+                SELECT
+                  a + 1 AS a
+                FROM foo
+              )
+            )
+            (
+              SELECT
+                a + b AS a
+              FROM bar
+              CROSS JOIN (
+                SELECT
+                  2 AS b
+              ) AS baz
+            )
+            """,
+        )
+        self.assertEqual(node.name, "a")
+        self.assertEqual(len(node.downstream), 2)
+        a, b = sorted(node.downstream, key=lambda n: n.name)
+        self.assertEqual(a.name, "bar.a")
+        self.assertEqual(len(a.downstream), 1)
+        self.assertEqual(b.name, "baz.b")
+        self.assertEqual(b.downstream, [])
+
+        node = a.downstream[0]
+        self.assertEqual(node.name, "foo.a")
+
     def test_lineage_cte_union(self) -> None:
         query = """
         WITH dataset AS (
@@ -353,3 +388,46 @@ class TestLineage(unittest.TestCase):
 
         with self.assertRaises(sqlglot.errors.SqlglotError):
             lineage('"a"', "WITH x AS (SELECT 1 a) SELECT a FROM x", dialect="snowflake")
+
+    def test_ddl_lineage(self) -> None:
+        sql = """
+        INSERT /*+ HINT1 */
+        INTO target (x, y)
+        SELECT subq.x, subq.y
+        FROM (
+          SELECT /*+ HINT2 */
+            t.x AS x,
+            TO_DATE('2023-12-19', 'YYYY-MM-DD') AS y
+          FROM s.t t
+          WHERE 1 = 1 AND y = TO_DATE('2023-12-19', 'YYYY-MM-DD')
+        ) subq
+        """
+
+        node = lineage("y", sql, dialect="oracle")
+
+        self.assertEqual(node.name, "Y")
+        self.assertEqual(node.expression.sql(dialect="oracle"), "SUBQ.Y AS Y")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.name, "SUBQ.Y")
+        self.assertEqual(
+            downstream.expression.sql(dialect="oracle"), "TO_DATE('2023-12-19', 'YYYY-MM-DD') AS Y"
+        )
+
+    def test_trim(self) -> None:
+        sql = """
+            SELECT a, b, c
+            FROM (select a, b, c from y) z
+        """
+
+        node = lineage("a", sql, trim_selects=False)
+
+        self.assertEqual(node.name, "a")
+        self.assertEqual(
+            node.source.sql(),
+            "SELECT z.a AS a, z.b AS b, z.c AS c FROM (SELECT y.a AS a, y.b AS b, y.c AS c FROM y AS y) AS z",
+        )
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.name, "z.a")
+        self.assertEqual(downstream.source.sql(), "SELECT y.a AS a, y.b AS b, y.c AS c FROM y AS y")