Repository: tobymao/sqlglot

Issue:
Snowflake join table(flatten(.)) parsed as anonymous table
**Fully reproducible code snippet**
_Please include a fully reproducible code snippet or the input sql, dialect, and expected output._

In Snowflake SQL it is possible to flatten objects in (at least) the following two ways: using a lateral flatten and using a join table flatten. If we then compute the lineage of a column taken from the flattened source, the catalog and db of the source is accounted for in the former but not the lateral. I think this is because the lateral is parsed as an Anonymous table.

```python
sql_lateral = """
select f.value:external_id::string as external_id
from database_name.schema_name.table_name raw,
lateral flatten(events) f
"""
sql_join = """
select f.value:external_id::string as external_id
from database_name.schema_name.table_name raw
join table(flatten(events)) f
"""

lat_expression = parse_one(sql_lateral, dialect="snowflake")
join_expression = parse_one(sql_join, dialect="snowflake")

lat_root_node = lineage("external_id", lat_expression, dialect="snowflake")
join_root_node = lineage("external_id", join_expression, dialect="snowflake")

def print_source(root_node):
    for node in root_node.walk():
        source = node.source
        if isinstance(source, exp.Table):
            print(f"{source.catalog}.{source.db}.{source.this}")

print_source(lat_root_node) # DATABASE_NAME.SCHEMA_NAME.TABLE_NAME
print_source(join_root_node) # ..TABLE(EXPLODE(EVENTS))
```

**Official Documentation**
_Please include links to official SQL documentation related to your issue._

Discussion:

Solution patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 8473eaf4dc..228167c599 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -450,6 +450,7 @@ class Parser(parser.Parser):
             "REGEXP_SUBSTR_ALL": _build_regexp_extract(exp.RegexpExtractAll),
             "RLIKE": exp.RegexpLike.from_arg_list,
             "SQUARE": lambda args: exp.Pow(this=seq_get(args, 0), expression=exp.Literal.number(2)),
+            "TABLE": lambda args: exp.TableFromRows(this=seq_get(args, 0)),
             "TIMEADD": _build_date_time_add(exp.TimeAdd),
             "TIMEDIFF": _build_datediff,
             "TIMESTAMPADD": _build_date_time_add(exp.DateAdd),
@@ -747,6 +748,33 @@ def _parse_table_parts(
 
             return table
 
+        def _parse_table(
+            self,
+            schema: bool = False,
+            joins: bool = False,
+            alias_tokens: t.Optional[t.Collection[TokenType]] = None,
+            parse_bracket: bool = False,
+            is_db_reference: bool = False,
+            parse_partition: bool = False,
+        ) -> t.Optional[exp.Expression]:
+            table = super()._parse_table(
+                schema=schema,
+                joins=joins,
+                alias_tokens=alias_tokens,
+                parse_bracket=parse_bracket,
+                is_db_reference=is_db_reference,
+                parse_partition=parse_partition,
+            )
+            if isinstance(table, exp.Table) and isinstance(table.this, exp.TableFromRows):
+                table_from_rows = table.this
+                for arg in exp.TableFromRows.arg_types:
+                    if arg != "this":
+                        table_from_rows.set(arg, table.args.get(arg))
+
+                table = table_from_rows
+
+            return table
+
         def _parse_id_var(
             self,
             any_token: bool = True,
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 0f549e3147..4d72be9144 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2587,6 +2587,18 @@ class Lateral(UDTF):
     }
 
 
+# https://docs.snowflake.com/sql-reference/literals-table
+# https://docs.snowflake.com/en/sql-reference/functions-table#using-a-table-function
+class TableFromRows(UDTF):
+    arg_types = {
+        "this": True,
+        "alias": False,
+        "joins": False,
+        "pivots": False,
+        "sample": False,
+    }
+
+
 class MatchRecognizeMeasure(Expression):
     arg_types = {
         "this": True,
@@ -8520,19 +8532,21 @@ def expand(
         The transformed expression.
     """
     normalized_sources = {normalize_table_name(k, dialect=dialect): v for k, v in sources.items()}
-    # Create a query provider based on the sources parameter
 
     def _expand(node: Expression):
         if isinstance(node, Table):
             name = normalize_table_name(node, dialect=dialect)
             source = normalized_sources.get(name)
+
             if source:
                 # Create a subquery with the same alias (or table name if no alias)
                 parsed_source = source() if callable(source) else source
                 subquery = parsed_source.subquery(node.alias or name)
                 subquery.comments = [f"source: {name}"]
+
                 # Continue expanding within the subquery
                 return subquery.transform(_expand, copy=False)
+
         return node
 
     return expression.transform(_expand, copy=copy)
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index b318826cea..926654a7f2 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -191,8 +191,8 @@ class Generator(metaclass=_Generator):
         exp.StreamingTableProperty: lambda *_: "STREAMING",
         exp.StrictProperty: lambda *_: "STRICT",
         exp.SwapTable: lambda self, e: f"SWAP WITH {self.sql(e, 'this')}",
-        exp.TemporaryProperty: lambda *_: "TEMPORARY",
         exp.Tags: lambda self, e: f"TAG ({self.expressions(e, flat=True)})",
+        exp.TemporaryProperty: lambda *_: "TEMPORARY",
         exp.TitleColumnConstraint: lambda self, e: f"TITLE {self.sql(e, 'this')}",
         exp.ToMap: lambda self, e: f"MAP {self.sql(e, 'this')}",
         exp.ToTableProperty: lambda self, e: f"TO {self.sql(e.this)}",
@@ -1999,6 +1999,17 @@ def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
 
         return f"{only}{table}{changes}{partition}{version}{file_format}{sample_pre_alias}{alias}{hints}{pivots}{sample_post_alias}{joins}{laterals}{ordinality}"
 
+    def tablefromrows_sql(self, expression: exp.TableFromRows) -> str:
+        table = self.func("TABLE", expression.this)
+        alias = self.sql(expression, "alias")
+        alias = f" AS {alias}" if alias else ""
+        sample = self.sql(expression, "sample")
+        pivots = self.expressions(expression, key="pivots", sep="", flat=True)
+        joins = self.indent(
+            self.expressions(expression, key="joins", sep="", flat=True), skip_first=True
+        )
+        return f"{table}{alias}{pivots}{sample}{joins}"
+
     def tablesample_sql(
         self,
         expression: exp.TableSample,
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 4dd6a8d92a..fc844c5179 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1387,6 +1387,8 @@ class Parser(metaclass=_Parser):
 
     RECURSIVE_CTE_SEARCH_KIND = {"BREADTH", "DEPTH", "CYCLE"}
 
+    MODIFIABLES = (exp.Query, exp.Table, exp.TableFromRows)
+
     STRICT_CAST = True
 
     PREFIXED_PIVOT_COLUMNS = False
@@ -3351,7 +3353,7 @@ def _implicit_unnests_to_explicit(self, this: E) -> E:
     def _parse_query_modifiers(
         self, this: t.Optional[exp.Expression]
     ) -> t.Optional[exp.Expression]:
-        if isinstance(this, (exp.Query, exp.Table)):
+        if isinstance(this, self.MODIFIABLES):
             for join in self._parse_joins():
                 this.append("joins", join)
             for lateral in iter(self._parse_lateral, None):


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index c5da2085f2..cc4e6257da 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -1627,44 +1627,21 @@ def test_stored_procedures(self):
             "CREATE PROCEDURE a.b.c(x INT, y VARIANT) RETURNS OBJECT EXECUTE AS CALLER AS 'BEGIN SELECT 1; END;'"
         )
 
-    def test_table_literal(self):
-        # All examples from https://docs.snowflake.com/en/sql-reference/literals-table.html
-        self.validate_all(
-            r"""SELECT * FROM TABLE('MYTABLE')""",
-            write={"snowflake": r"""SELECT * FROM TABLE('MYTABLE')"""},
-        )
-
-        self.validate_all(
-            r"""SELECT * FROM TABLE('MYDB."MYSCHEMA"."MYTABLE"')""",
-            write={"snowflake": r"""SELECT * FROM TABLE('MYDB."MYSCHEMA"."MYTABLE"')"""},
-        )
-
-        # Per Snowflake documentation at https://docs.snowflake.com/en/sql-reference/literals-table.html
-        # one can use either a  " ' " or " $$ " to enclose the object identifier.
-        # Capturing the single tokens seems like lot of work. Hence adjusting tests to use these interchangeably,
-        self.validate_all(
-            r"""SELECT * FROM TABLE($$MYDB. "MYSCHEMA"."MYTABLE"$$)""",
-            write={"snowflake": r"""SELECT * FROM TABLE('MYDB. "MYSCHEMA"."MYTABLE"')"""},
-        )
-
-        self.validate_all(
-            r"""SELECT * FROM TABLE($MYVAR)""",
-            write={"snowflake": r"""SELECT * FROM TABLE($MYVAR)"""},
-        )
-
-        self.validate_all(
-            r"""SELECT * FROM TABLE(?)""",
-            write={"snowflake": r"""SELECT * FROM TABLE(?)"""},
-        )
-
-        self.validate_all(
-            r"""SELECT * FROM TABLE(:BINDING)""",
-            write={"snowflake": r"""SELECT * FROM TABLE(:BINDING)"""},
-        )
-
-        self.validate_all(
-            r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10""",
-            write={"snowflake": r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10"""},
+    def test_table_function(self):
+        self.validate_identity("SELECT * FROM TABLE('MYTABLE')")
+        self.validate_identity("SELECT * FROM TABLE($MYVAR)")
+        self.validate_identity("SELECT * FROM TABLE(?)")
+        self.validate_identity("SELECT * FROM TABLE(:BINDING)")
+        self.validate_identity("SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10")
+        self.validate_identity("SELECT * FROM TABLE('t1') AS f")
+        self.validate_identity("SELECT * FROM (TABLE('t1') CROSS JOIN TABLE('t2'))")
+        self.validate_identity("SELECT * FROM TABLE('t1'), LATERAL (SELECT * FROM t2)")
+        self.validate_identity("SELECT * FROM TABLE('t1') UNION ALL SELECT * FROM TABLE('t2')")
+        self.validate_identity("SELECT * FROM TABLE('t1') TABLESAMPLE BERNOULLI (20.3)")
+        self.validate_identity("""SELECT * FROM TABLE('MYDB."MYSCHEMA"."MYTABLE"')""")
+        self.validate_identity(
+            'SELECT * FROM TABLE($$MYDB. "MYSCHEMA"."MYTABLE"$$)',
+            """SELECT * FROM TABLE('MYDB. "MYSCHEMA"."MYTABLE"')""",
         )
 
     def test_flatten(self):
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index 095ee80e47..36380dd665 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -576,3 +576,48 @@ def test_pivot_with_implicit_column_of_pivoted_source_and_cte(self) -> None:
         self.assertEqual(node.downstream[0].name, "t.empid")
         self.assertEqual(node.downstream[0].reference_node_name, "t")
         self.assertEqual(node.downstream[0].downstream[0].name, "quarterly_sales.empid")
+
+    def test_table_udtf_snowflake(self) -> None:
+        lateral_flatten = """
+        SELECT f.value:external_id::string AS external_id
+        FROM database_name.schema_name.table_name AS raw,
+        LATERAL FLATTEN(events) AS f
+        """
+        table_flatten = """
+        SELECT f.value:external_id::string AS external_id
+        FROM database_name.schema_name.table_name AS raw
+        JOIN TABLE(FLATTEN(events)) AS f
+        """
+
+        lateral_node = lineage("external_id", lateral_flatten, dialect="snowflake")
+        table_node = lineage("external_id", table_flatten, dialect="snowflake")
+
+        self.assertEqual(lateral_node.name, "EXTERNAL_ID")
+        self.assertEqual(table_node.name, "EXTERNAL_ID")
+
+        lateral_node = lateral_node.downstream[0]
+        table_node = table_node.downstream[0]
+
+        self.assertEqual(lateral_node.name, "F.VALUE")
+        self.assertEqual(
+            lateral_node.source.sql("snowflake"),
+            "LATERAL FLATTEN(RAW.EVENTS) AS F(SEQ, KEY, PATH, INDEX, VALUE, THIS)",
+        )
+
+        self.assertEqual(table_node.name, "F.VALUE")
+        self.assertEqual(table_node.source.sql("snowflake"), "TABLE(FLATTEN(RAW.EVENTS)) AS F")
+
+        lateral_node = lateral_node.downstream[0]
+        table_node = table_node.downstream[0]
+
+        self.assertEqual(lateral_node.name, "RAW.EVENTS")
+        self.assertEqual(
+            lateral_node.source.sql("snowflake"),
+            "DATABASE_NAME.SCHEMA_NAME.TABLE_NAME AS RAW",
+        )
+
+        self.assertEqual(table_node.name, "RAW.EVENTS")
+        self.assertEqual(
+            table_node.source.sql("snowflake"),
+            "DATABASE_NAME.SCHEMA_NAME.TABLE_NAME AS RAW",
+        )