Fix(snowflake)!: treat TABLE(...) as a UDTF (#4875)

Author: georgesittas

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,

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,

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,

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):

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):

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",
+        )