Repository: tobymao/sqlglot

Issue:
No `source_name` in column lineage with `LATERAL FLATTEN`
**Fully reproducible code snippet**

```python
from sqlglot.lineage import lineage

query = """SELECT
    FLATTENED.VALUE:field::text AS FIELD
FROM
    SNOWFLAKE.SCHEMA.MODEL AS MODEL_ALIAS,
    LATERAL FLATTEN(INPUT =>
        MODEL_ALIAS.A
    ) AS FLATTENED
"""
sources = {"SNOWFLAKE.SCHEMA.MODEL": "SELECT A FROM SNOWFLAKE.SCHEMA.TABLE"}
schemas = {
    "SCHEMA": {
        "TABLE": {"A": "integer"},
    }
}

result = lineage("FIELD", query, schemas, sources, dialect="snowflake")
for node in result.walk():
    print(f"Name: {node.name}, Source: {node.source_name}")
```

The output is:

```
Name: FIELD, Source:
Name: FLATTENED.VALUE, Source:
Name: MODEL_ALIAS.A, Source:
Name: TABLE.A, Source:
```

I would expect the `MODEL_ALIAS.A` node to have `source_name` equal to `SNOWFLAKE.SCHEMA.MODEL` because that's the source (`sources` argument to the `lineage` function) where that column appears. That's how that field seems to work in other queries.

For example, changing the query to:

```sql
SELECT
    MODEL_ALIAS.A AS FIELD
FROM
    SNOWFLAKE.SCHEMA.MODEL AS MODEL_ALIAS
```

gives:

```
Name: FIELD, Source:
Name: MODEL_ALIAS.A, Source: SNOWFLAKE.SCHEMA.MODEL
Name: TABLE.A, Source:
```

I believe the root cause is that the expanded-and-qualified query has a `source: ` comment after the the first element in the `FROM` clause but not for the `LATERAL FLATTEN`?


Discussion:

Resolving patch:
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -129,12 +129,6 @@ def to_node(
     reference_node_name: t.Optional[str] = None,
     trim_selects: bool = True,
 ) -> Node:
-    source_names = {
-        dt.alias: dt.comments[0].split()[1]
-        for dt in scope.derived_tables
-        if dt.comments and dt.comments[0].startswith("source: ")
-    }
-
     # Find the specific select clause that is the source of the column we want.
     # This can either be a specific, named select or a generic `*` clause.
     select = (
@@ -242,6 +236,19 @@ def to_node(
     # If the source is a UDTF find columns used in the UTDF to generate the table
     if isinstance(source, exp.UDTF):
         source_columns |= set(source.find_all(exp.Column))
+        derived_tables = [
+            source.expression.parent
+            for source in scope.sources.values()
+            if isinstance(source, Scope) and source.is_derived_table
+        ]
+    else:
+        derived_tables = scope.derived_tables
+
+    source_names = {
+        dt.alias: dt.comments[0].split()[1]
+        for dt in derived_tables
+        if dt.comments and dt.comments[0].startswith("source: ")
+    }
 
     for c in source_columns:
         table = c.table


Test patch:
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -224,16 +224,50 @@ def test_lineage_lateral_flatten(self) -> None:
             downstream.source.sql(dialect="snowflake"),
             "LATERAL FLATTEN(INPUT => TEST_TABLE.RESULT, OUTER => TRUE) AS FLATTENED(SEQ, KEY, PATH, INDEX, VALUE, THIS)",
         )
-        self.assertEqual(
-            downstream.expression.sql(dialect="snowflake"),
-            "VALUE",
-        )
+        self.assertEqual(downstream.expression.sql(dialect="snowflake"), "VALUE")
         self.assertEqual(len(downstream.downstream), 1)
 
         downstream = downstream.downstream[0]
         self.assertEqual(downstream.name, "TEST_TABLE.RESULT")
         self.assertEqual(downstream.source.sql(dialect="snowflake"), "TEST_TABLE AS TEST_TABLE")
 
+        node = lineage(
+            "FIELD",
+            "SELECT FLATTENED.VALUE:field::text AS FIELD FROM SNOWFLAKE.SCHEMA.MODEL AS MODEL_ALIAS, LATERAL FLATTEN(INPUT => MODEL_ALIAS.A) AS FLATTENED",
+            schema={"SNOWFLAKE": {"SCHEMA": {"TABLE": {"A": "integer"}}}},
+            sources={"SNOWFLAKE.SCHEMA.MODEL": "SELECT A FROM SNOWFLAKE.SCHEMA.TABLE"},
+            dialect="snowflake",
+        )
+        self.assertEqual(node.name, "FIELD")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.name, "FLATTENED.VALUE")
+        self.assertEqual(
+            downstream.source.sql(dialect="snowflake"),
+            "LATERAL FLATTEN(INPUT => MODEL_ALIAS.A) AS FLATTENED(SEQ, KEY, PATH, INDEX, VALUE, THIS)",
+        )
+        self.assertEqual(downstream.expression.sql(dialect="snowflake"), "VALUE")
+        self.assertEqual(len(downstream.downstream), 1)
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.name, "MODEL_ALIAS.A")
+        self.assertEqual(downstream.source_name, "SNOWFLAKE.SCHEMA.MODEL")
+        self.assertEqual(
+            downstream.source.sql(dialect="snowflake"),
+            "SELECT TABLE.A AS A FROM SNOWFLAKE.SCHEMA.TABLE AS TABLE",
+        )
+        self.assertEqual(downstream.expression.sql(dialect="snowflake"), "TABLE.A AS A")
+        self.assertEqual(len(downstream.downstream), 1)
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.name, "TABLE.A")
+        self.assertEqual(
+            downstream.source.sql(dialect="snowflake"), "SNOWFLAKE.SCHEMA.TABLE AS TABLE"
+        )
+        self.assertEqual(
+            downstream.expression.sql(dialect="snowflake"), "SNOWFLAKE.SCHEMA.TABLE AS TABLE"
+        )
+
     def test_subquery(self) -> None:
         node = lineage(
             "output",