Repository: tobymao/sqlglot

Issue:
Difficulty parsing snowflake query with `split()` nested in `flatten()`
Hello!

I've been unable to get sqlglot to parse the following query and I'm not sure if it's user error or a bug? Would love to get this resolved:

```python

import sqlglot
import sqlglot.expressions as exp

query = f"""
select
  dag_report.acct_id,
  dag_report.report_date,
  dag_report.report_uuid,
  dag_report.airflow_name,
  dag_report.dag_id,
  f.value::varchar as operator
from cs.telescope.dag_report,
table(flatten(input=>split(operators, ','))) f
"""

def get_query_columns(query):
    column_names = []
    for expression in sqlglot.parse_one(query, read="snowflake").find(exp.Select).args["expressions"]:
        if isinstance(expression, exp.Alias):
            column_names.append(expression.text("alias"))
        elif isinstance(expression, exp.Column):
            column_names.append(expression.text("this"))

    return column_names

if __name__ == '__main__':
    column_names = get_query_columns(query)
    print(column_names)

```

Error:

```text
sqlglot.errors.ParseError: Required keyword: 'this' missing for <class 'sqlglot.expressions.GT'>. Line 10, Col: 43.
```

Any help would be greatly appreciated. 🍻




Discussion:
Hello, thanks for the report.

This error happens because `=>` is not currently parsed properly to support Snowflake's `flatten` [syntax](https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax).

I'll take a look at it soon.

Resolving patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -58,7 +58,7 @@ def _snowflake_to_timestamp(args):
     return exp.UnixToTime.from_arg_list(args)
 
 
-def _unix_to_time(self, expression):
+def _unix_to_time_sql(self, expression):
     scale = expression.args.get("scale")
     timestamp = self.sql(expression, "this")
     if scale in [None, exp.UnixToTime.SECONDS]:
@@ -180,6 +180,7 @@ class Tokenizer(tokens.Tokenizer):
 
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
+            "=>": TokenType.ARROW,
             "QUALIFY": TokenType.QUALIFY,
             "DOUBLE PRECISION": TokenType.DOUBLE,
             "TIMESTAMP_LTZ": TokenType.TIMESTAMPLTZ,
@@ -197,7 +198,7 @@ class Generator(generator.Generator):
             exp.ArrayConcat: rename_func("ARRAY_CAT"),
             exp.If: rename_func("IFF"),
             exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
-            exp.UnixToTime: _unix_to_time,
+            exp.UnixToTime: _unix_to_time_sql,
             exp.TimeToUnix: lambda self, e: f"EXTRACT(epoch_second FROM {self.sql(e, 'this')})",
             exp.Array: inline_array_sql,
             exp.StrPosition: rename_func("POSITION"),
@@ -219,6 +220,9 @@ class Generator(generator.Generator):
             exp.VolatilityProperty,
         }
 
+        def lambda_sql(self, expression):
+            return super().lambda_sql(expression, arrow_sep="=>")
+
         def except_op(self, expression):
             if not expression.args.get("distinct", False):
                 self.unsupported("EXCEPT with All is not supported in Snowflake")
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -727,10 +727,10 @@ def join_sql(self, expression):
         this_sql = self.sql(expression, "this")
         return f"{expression_sql}{op_sql} {this_sql}{on_sql}"
 
-    def lambda_sql(self, expression):
+    def lambda_sql(self, expression, arrow_sep="->"):
         args = self.expressions(expression, flat=True)
         args = f"({args})" if len(args.split(",")) > 1 else args
-        return self.no_identify(lambda: f"{args} -> {self.sql(expression, 'this')}")
+        return self.no_identify(lambda: f"{args} {arrow_sep} {self.sql(expression, 'this')}")
 
     def lateral_sql(self, expression):
         this = self.sql(expression, "this")


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -371,3 +371,105 @@ def test_table_literal(self):
             r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10""",
             write={"snowflake": r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10"""},
         )
+
+    def test_flatten(self):
+        self.validate_all(
+            """
+            select
+              dag_report.acct_id,
+              dag_report.report_date,
+              dag_report.report_uuid,
+              dag_report.airflow_name,
+              dag_report.dag_id,
+              f.value::varchar as operator
+            from cs.telescope.dag_report,
+            table(flatten(input=>split(operators, ','))) f
+            """,
+            write={
+                "snowflake": """SELECT
+  dag_report.acct_id,
+  dag_report.report_date,
+  dag_report.report_uuid,
+  dag_report.airflow_name,
+  dag_report.dag_id,
+  CAST(f.value AS VARCHAR) AS operator
+FROM cs.telescope.dag_report, TABLE(FLATTEN(input => SPLIT(operators, ','))) AS f"""
+            },
+            pretty=True,
+        )
+
+        # All examples from https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax
+        self.validate_all(
+            "SELECT * FROM TABLE(FLATTEN(input => parse_json('[1, ,77]'))) f",
+            write={
+                "snowflake": "SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('[1, ,77]'))) AS f"
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), outer => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('{"a":1, "b":[77,88]}'), outer => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), path => 'b')) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('{"a":1, "b":[77,88]}'), path => 'b')) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('[]'))) f""",
+            write={"snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('[]'))) AS f"""},
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('[]'), outer => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('[]'), outer => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'))) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'))) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => true, mode => 'object')) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(input => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => TRUE, mode => 'object')) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """
+            SELECT id as "ID",
+              f.value AS "Contact",
+              f1.value:type AS "Type",
+              f1.value:content AS "Details"
+            FROM persons p,
+              lateral flatten(input => p.c, path => 'contact') f,
+              lateral flatten(input => f.value:business) f1
+            """,
+            write={
+                "snowflake": """SELECT
+  id AS "ID",
+  f.value AS "Contact",
+  f1.value['type'] AS "Type",
+  f1.value['content'] AS "Details"
+FROM persons AS p, LATERAL FLATTEN(input => p.c, path => 'contact') f, LATERAL FLATTEN(input => f.value['business']) f1""",
+            },
+            pretty=True,
+        )