Repository: tobymao/sqlglot

Issue:
 BigQuery dialect parse error on "FOR... IN(...) DO"
version 19.0.1
steps to reproduce:
```python
sqlglot.parse(
  "FOR record IN (SELECT word FROM shakespeare) DO SELECT record.word; END FOR;",
  sqlglot.dialects.BigQuery
)
```

<details>
<summary>traceback</summary>

>>> sqlglot.parse("FOR record IN (SELECT word FROM shakespeare) DO SELECT record.word; END FOR;", sqlglot.dialects.BigQuery)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/__init__.py", line 86, in parse
    return dialect.parse(sql, **opts)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/dialects/dialect.py", line 311, in parse
    return self.parser(**opts).parse(self.tokenize(sql), sql)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 987, in parse
    return self._parse(
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 1053, in _parse
    expressions.append(parse_method(self))
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 1242, in _parse_statement
    expression = self._parse_expression()
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 3232, in _parse_expression
    return self._parse_alias(self._parse_conjunction())
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 3235, in _parse_conjunction
    return self._parse_tokens(self._parse_equality, self.CONJUNCTION)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 4905, in _parse_tokens
    this = parse_method()
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 3238, in _parse_equality
    return self._parse_tokens(self._parse_comparison, self.EQUALITY)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 4905, in _parse_tokens
    this = parse_method()
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 3241, in _parse_comparison
    return self._parse_tokens(self._parse_range, self.COMPARISON)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 4905, in _parse_tokens
    this = parse_method()
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 3248, in _parse_range
    expression = self.RANGE_PARSERS[self._prev.token_type](self, this)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 626, in <lambda>
    TokenType.FOR: lambda self, this: self._parse_comprehension(this),
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 5303, in _parse_comprehension
    return self.expression(
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 1117, in expression
    return self.validate_expression(instance)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 1137, in validate_expression
    self.raise_error(error_message)
  File "/.pyenv/versions/3.10.6/lib/python3.10/site-packages/sqlglot/parser.py", line 1097, in raise_error
    raise error
sqlglot.errors.ParseError: Required keyword: 'this' missing for <class 'sqlglot.expressions.Comprehension'>. Line 1, Col: 47.
  FOR record IN (SELECT word FROM shakespeare) DO SELECT record.word; END FOR;
</details>

The syntax is described in the BigQuery docs [here](https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#for-in).


Discussion:

Resolving patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -358,6 +358,17 @@ class Parser(parser.Parser):
 
         NULL_TOKENS = {TokenType.NULL, TokenType.UNKNOWN}
 
+        STATEMENT_PARSERS = {
+            **parser.Parser.STATEMENT_PARSERS,
+            TokenType.END: lambda self: self._parse_as_command(self._prev),
+            TokenType.FOR: lambda self: self._parse_for_in(),
+        }
+
+        def _parse_for_in(self) -> exp.ForIn:
+            this = self._parse_range()
+            self._match_text_seq("DO")
+            return self.expression(exp.ForIn, this=this, expression=self._parse_statement())
+
         def _parse_table_part(self, schema: bool = False) -> t.Optional[exp.Expression]:
             this = super()._parse_table_part(schema=schema) or self._parse_number()
 
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -4090,6 +4090,11 @@ class In(Predicate):
     }
 
 
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#for-in
+class ForIn(Expression):
+    arg_types = {"this": True, "expression": True}
+
+
 class TimeUnit(Expression):
     """Automatically converts unit arg into a var."""
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -2993,6 +2993,11 @@ def predict_sql(self, expression: exp.Predict) -> str:
         parameters = self.sql(expression, "params_struct")
         return self.func("PREDICT", model, table, parameters or None)
 
+    def forin_sql(self, expression: exp.ForIn) -> str:
+        this = self.sql(expression, "this")
+        expression_sql = self.sql(expression, "expression")
+        return f"FOR {this} DO {expression_sql}"
+
     def _simplify_unless_literal(self, expression: E) -> E:
         if not isinstance(expression, exp.Literal):
             from sqlglot.optimizer.simplify import simplify


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1,6 +1,13 @@
 from unittest import mock
 
-from sqlglot import ErrorLevel, ParseError, TokenError, UnsupportedError, transpile
+from sqlglot import (
+    ErrorLevel,
+    ParseError,
+    TokenError,
+    UnsupportedError,
+    parse,
+    transpile,
+)
 from tests.dialects.test_dialect import Validator
 
 
@@ -37,6 +44,15 @@ def test_bigquery(self):
         with self.assertRaises(ParseError):
             transpile("DATE_ADD(x, day)", read="bigquery")
 
+        for_in_stmts = parse(
+            "FOR record IN (SELECT word FROM shakespeare) DO SELECT record.word; END FOR;",
+            read="bigquery",
+        )
+        self.assertEqual(
+            [s.sql(dialect="bigquery") for s in for_in_stmts],
+            ["FOR record IN (SELECT word FROM shakespeare) DO SELECT record.word", "END FOR"],
+        )
+
         self.validate_identity("SELECT test.Unknown FROM test")
         self.validate_identity(r"SELECT '\n\r\a\v\f\t'")
         self.validate_identity("SELECT * FROM tbl FOR SYSTEM_TIME AS OF z")
@@ -89,6 +105,9 @@ def test_bigquery(self):
         self.validate_identity("ROLLBACK TRANSACTION")
         self.validate_identity("CAST(x AS BIGNUMERIC)")
         self.validate_identity("SELECT y + 1 FROM x GROUP BY y + 1 ORDER BY 1")
+        self.validate_identity(
+            "FOR record IN (SELECT word, word_count FROM bigquery-public-data.samples.shakespeare LIMIT 5) DO SELECT record.word, record.word_count"
+        )
         self.validate_identity(
             "DATE(CAST('2016-12-25 05:30:00+07' AS DATETIME), 'America/Los_Angeles')"
         )