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:

Patch:
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
index 71339b88..2e53a675 100644
--- a/sqlglot/dialects/postgres.py
+++ b/sqlglot/dialects/postgres.py
@@ -518,6 +518,7 @@ class Postgres(Dialect):
             exp.Variance: rename_func("VAR_SAMP"),
             exp.Xor: bool_xor_sql,
         }
+        TRANSFORMS.pop(exp.CommentColumnConstraint)
 
         PROPERTIES_LOCATION = {
             **generator.Generator.PROPERTIES_LOCATION,
@@ -526,6 +527,10 @@ class Postgres(Dialect):
             exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
         }
 
+        def commentcolumnconstraint_sql(self, expression: exp.CommentColumnConstraint) -> str:
+            self.unsupported("Column comments are not supported in the CREATE statement")
+            return ""
+
         def unnest_sql(self, expression: exp.Unnest) -> str:
             if len(expression.expressions) == 1:
                 from sqlglot.optimizer.annotate_types import annotate_types
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index c91bb36e..f4a3dec5 100644
--- 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
diff --git a/sqlglot/optimizer/__init__.py b/sqlglot/optimizer/__init__.py
index 34ea6cb1..050f246c 100644
--- a/sqlglot/optimizer/__init__.py
+++ b/sqlglot/optimizer/__init__.py
@@ -1,11 +1,11 @@
 # ruff: noqa: F401
 
-from sqlglot.optimizer.optimizer import RULES, optimize
+from sqlglot.optimizer.optimizer import RULES as RULES, optimize as optimize
 from sqlglot.optimizer.scope import (
-    Scope,
-    build_scope,
-    find_all_in_scope,
-    find_in_scope,
-    traverse_scope,
-    walk_in_scope,
+    Scope as Scope,
+    build_scope as build_scope,
+    find_all_in_scope as find_all_in_scope,
+    find_in_scope as find_in_scope,
+    traverse_scope as traverse_scope,
+    walk_in_scope as walk_in_scope,
 )


Test patch:
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
index 5a55a7d6..1ed7d82f 100644
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -314,6 +314,12 @@ class TestPostgres(Validator):
         )
         self.validate_identity("SELECT * FROM t1*", "SELECT * FROM t1")
 
+        self.validate_all(
+            "CREATE TABLE t (c INT)",
+            read={
+                "mysql": "CREATE TABLE t (c INT COMMENT 'comment')",
+            },
+        )
         self.validate_all(
             'SELECT * FROM "test_table" ORDER BY RANDOM() LIMIT 5',
             write={
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index c782d9ae..cbadf7bb 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -224,16 +224,50 @@ class TestLineage(unittest.TestCase):
             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",