Repository: tobymao/sqlglot

Issue:
Snowflake parsing/output error with Lateral Flatten
Snowflake SQL statement becomes invalid once parsed and output by SQLGlot. An extra `INPUT =>` is added to the query which then fails with `Syntax error: unexpected '=>'. ` The extra `INPUT =>` is visible in the last `CROSS JOIN` statement. 

**Fully reproducible code snippet**
```python
from sqlglot import parse_one, exp, parse
sql = """
  SELECT
    uc.user_id
    , uc.start_ts AS ts
    , CASE
        WHEN uc.start_ts::DATE >= '2023-01-01' AND uc.country_code IN ('US')
          AND uc.user_id NOT IN (SELECT DISTINCT _id FROM users, LATERAL FLATTEN(INPUT => PARSE_JSON(flags)) datasource WHERE datasource.value:name = 'something')
          THEN 'Sample1'
        ELSE
          'Sample2'
      END AS entity
  FROM user_countries AS uc
  LEFT JOIN (SELECT user_id, MAX(IFF(service_entity IS NULL,1,0)) AS le_null FROM accepted_user_agreements GROUP BY 1) AS aua
    ON uc.user_id = aua.user_id
"""
ast = parse_one(sql, dialect="snowflake")
print(ast.sql(dialect="snowflake", pretty=True))
``` 

Output: 
```
SELECT
  uc.user_id,
  uc.start_ts AS ts,
  CASE
    WHEN CAST(uc.start_ts AS DATE) >= '2023-01-01'
    AND uc.country_code IN ('US')
    AND uc.user_id <> ALL (
      SELECT DISTINCT
        _id
      FROM users, LATERAL IFF(_u.pos = _u_2.pos_2, _u_2.entity, NULL) AS datasource(SEQ, KEY, PATH, INDEX, VALUE, THIS)
      WHERE
        GET_PATH(datasource.value, 'name') = 'something'
    )
    THEN 'Sample1'
    ELSE 'Sample2'
  END AS entity
FROM user_countries AS uc
LEFT JOIN (
  SELECT
    user_id,
    MAX(IFF(service_entity IS NULL, 1, 0)) AS le_null
  FROM accepted_user_agreements
  GROUP BY
    1
) AS aua
  ON uc.user_id = aua.user_id
CROSS JOIN TABLE(FLATTEN(INPUT => ARRAY_GENERATE_RANGE(0, (
  GREATEST(ARRAY_SIZE(INPUT => PARSE_JSON(flags))) - 1
) + 1))) AS _u(seq, key, path, index, pos, this)
CROSS JOIN TABLE(FLATTEN(INPUT => INPUT => PARSE_JSON(flags))) AS _u_2(seq, key, path, pos_2, entity, this)
WHERE
  _u.pos = _u_2.pos_2
  OR (
    _u.pos > (
      ARRAY_SIZE(INPUT => PARSE_JSON(flags)) - 1
    )
    AND _u_2.pos_2 = (
      ARRAY_SIZE(INPUT => PARSE_JSON(flags)) - 1
    )
  )

``` 

Snowflake error:
![image](https://github.com/user-attachments/assets/78ee3848-2c27-4c04-b926-05176b5d97e0)



Discussion:

Patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 92a1671a..37842fd1 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -1097,7 +1097,11 @@ class Snowflake(Dialect):
             else:
                 unnest_alias = exp.TableAlias(this="_u", columns=columns)
 
-            explode = f"TABLE(FLATTEN(INPUT => {self.sql(expression.expressions[0])}))"
+            table_input = self.sql(expression.expressions[0])
+            if not table_input.startswith("INPUT =>"):
+                table_input = f"INPUT => {table_input}"
+
+            explode = f"TABLE(FLATTEN({table_input}))"
             alias = self.sql(unnest_alias)
             alias = f" AS {alias}" if alias else ""
             return f"{explode}{alias}"


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 5164b821..1d55f35d 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -21,27 +21,6 @@ class TestSnowflake(Validator):
         expr.selects[0].assert_is(exp.AggFunc)
         self.assertEqual(expr.sql(dialect="snowflake"), "SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
 
-        self.assertEqual(
-            exp.select(exp.Explode(this=exp.column("x")).as_("y", quoted=True)).sql(
-                "snowflake", pretty=True
-            ),
-            """SELECT
-  IFF(_u.pos = _u_2.pos_2, _u_2."y", NULL) AS "y"
-FROM TABLE(FLATTEN(INPUT => ARRAY_GENERATE_RANGE(0, (
-  GREATEST(ARRAY_SIZE(x)) - 1
-) + 1))) AS _u(seq, key, path, index, pos, this)
-CROSS JOIN TABLE(FLATTEN(INPUT => x)) AS _u_2(seq, key, path, pos_2, "y", this)
-WHERE
-  _u.pos = _u_2.pos_2
-  OR (
-    _u.pos > (
-      ARRAY_SIZE(x) - 1
-    ) AND _u_2.pos_2 = (
-      ARRAY_SIZE(x) - 1
-    )
-  )""",
-        )
-
         self.validate_identity("exclude := [foo]")
         self.validate_identity("SELECT CAST([1, 2, 3] AS VECTOR(FLOAT, 3))")
         self.validate_identity("SELECT CONNECT_BY_ROOT test AS test_column_alias")
@@ -1603,6 +1582,27 @@ WHERE
         )
 
     def test_flatten(self):
+        self.assertEqual(
+            exp.select(exp.Explode(this=exp.column("x")).as_("y", quoted=True)).sql(
+                "snowflake", pretty=True
+            ),
+            """SELECT
+  IFF(_u.pos = _u_2.pos_2, _u_2."y", NULL) AS "y"
+FROM TABLE(FLATTEN(INPUT => ARRAY_GENERATE_RANGE(0, (
+  GREATEST(ARRAY_SIZE(x)) - 1
+) + 1))) AS _u(seq, key, path, index, pos, this)
+CROSS JOIN TABLE(FLATTEN(INPUT => x)) AS _u_2(seq, key, path, pos_2, "y", this)
+WHERE
+  _u.pos = _u_2.pos_2
+  OR (
+    _u.pos > (
+      ARRAY_SIZE(x) - 1
+    ) AND _u_2.pos_2 = (
+      ARRAY_SIZE(x) - 1
+    )
+  )""",
+        )
+
         self.validate_all(
             """
             select
@@ -1627,6 +1627,75 @@ FROM cs.telescope.dag_report, TABLE(FLATTEN(input => SPLIT(operators, ','))) AS
             },
             pretty=True,
         )
+        self.validate_all(
+            """
+            SELECT
+              uc.user_id,
+              uc.start_ts AS ts,
+              CASE
+                WHEN uc.start_ts::DATE >= '2023-01-01' AND uc.country_code IN ('US') AND uc.user_id NOT IN (
+                  SELECT DISTINCT
+                    _id
+                  FROM
+                    users,
+                    LATERAL FLATTEN(INPUT => PARSE_JSON(flags)) datasource
+                  WHERE datasource.value:name = 'something'
+                )
+                  THEN 'Sample1'
+                  ELSE 'Sample2'
+              END AS entity
+            FROM user_countries AS uc
+            LEFT JOIN (
+              SELECT user_id, MAX(IFF(service_entity IS NULL,1,0)) AS le_null
+              FROM accepted_user_agreements
+              GROUP BY 1
+            ) AS aua
+              ON uc.user_id = aua.user_id
+            """,
+            write={
+                "snowflake": """SELECT
+  uc.user_id,
+  uc.start_ts AS ts,
+  CASE
+    WHEN CAST(uc.start_ts AS DATE) >= '2023-01-01'
+    AND uc.country_code IN ('US')
+    AND uc.user_id <> ALL (
+      SELECT DISTINCT
+        _id
+      FROM users, LATERAL IFF(_u.pos = _u_2.pos_2, _u_2.entity, NULL) AS datasource(SEQ, KEY, PATH, INDEX, VALUE, THIS)
+      WHERE
+        GET_PATH(datasource.value, 'name') = 'something'
+    )
+    THEN 'Sample1'
+    ELSE 'Sample2'
+  END AS entity
+FROM user_countries AS uc
+LEFT JOIN (
+  SELECT
+    user_id,
+    MAX(IFF(service_entity IS NULL, 1, 0)) AS le_null
+  FROM accepted_user_agreements
+  GROUP BY
+    1
+) AS aua
+  ON uc.user_id = aua.user_id
+CROSS JOIN TABLE(FLATTEN(INPUT => ARRAY_GENERATE_RANGE(0, (
+  GREATEST(ARRAY_SIZE(INPUT => PARSE_JSON(flags))) - 1
+) + 1))) AS _u(seq, key, path, index, pos, this)
+CROSS JOIN TABLE(FLATTEN(INPUT => PARSE_JSON(flags))) AS _u_2(seq, key, path, pos_2, entity, this)
+WHERE
+  _u.pos = _u_2.pos_2
+  OR (
+    _u.pos > (
+      ARRAY_SIZE(INPUT => PARSE_JSON(flags)) - 1
+    )
+    AND _u_2.pos_2 = (
+      ARRAY_SIZE(INPUT => PARSE_JSON(flags)) - 1
+    )
+  )""",
+            },
+            pretty=True,
+        )
 
         # All examples from https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax
         self.validate_all(