bigquery: parse EXECUTE IMMEDIATE in FOR...IN statements

sqlglot/parsers/bigquery.py

@@ -328,6 +328,7 @@ class BigQueryParser(parser.Parser):
         **parser.Parser.STATEMENT_PARSERS,
         TokenType.ELSE: lambda self: self._parse_as_command(self._prev),
         TokenType.END: lambda self: self._parse_as_command(self._prev),
+        TokenType.EXECUTE: lambda self: self._parse_execute_immediate(),
         TokenType.FOR: lambda self: self._parse_for_in(),
         TokenType.EXPORT: lambda self: self._parse_export_data(),
         TokenType.DECLARE: lambda self: self._parse_declare(),
@@ -349,6 +350,36 @@ def _parse_for_in(self) -> exp.ForIn | exp.Command:
             return self._parse_as_command(self._prev)
         return self.expression(exp.ForIn(this=this, expression=self._parse_statement()))
 
+    def _parse_execute_immediate(self) -> exp.Command:
+        start = self._prev
+
+        if not self._match_text_seq("IMMEDIATE"):
+            return self._parse_as_command(start)
+
+        if not self._parse_expression():
+            return self._parse_as_command(start)
+
+        if self._match_text_seq("INTO"):
+            self._parse_csv(lambda: self._parse_id_var(any_token=True) or self._parse_expression())
+
+        if self._match_text_seq("USING"):
+            self._parse_csv(self._parse_execute_immediate_using)
+
+        text = self._find_sql(start, self._prev)
+        size = len(start.text)
+        return exp.Command(this=text[:size], expression=text[size:])
+
+    def _parse_execute_immediate_using(self) -> exp.Expr | None:
+        this = self._parse_expression()
+
+        if not this:
+            return None
+
+        if self._match(TokenType.ALIAS):
+            self._parse_id_var(any_token=True)
+
+        return this
+
     def _parse_table_part(self, schema: bool = False) -> exp.Expr | None:
         this = super()._parse_table_part(schema=schema) or self._parse_number()
 

sqlglot/parsers/prql.py

@@ -8,8 +8,8 @@
 from collections.abc import Collection
 
 
-def _select_all(table: exp.Expr) -> exp.Select | None:
-    return exp.select("*").from_(table, copy=False) if table else None
+def _select_all(table: exp.Expr) -> exp.Select:
+    return exp.select("*").from_(table, copy=False)
 
 
 def _resolve_projection(s: exp.Expr, projections: dict[str, exp.Expr]) -> exp.Expr:
@@ -18,6 +18,10 @@ def _resolve_projection(s: exp.Expr, projections: dict[str, exp.Expr]) -> exp.Ex
     return s
 
 
+def _sum_with_zero(args: list[exp.Expr]) -> exp.Expr:
+    return exp.func("COALESCE", exp.Sum(this=seq_get(args, 0)), 0)
+
+
 class PRQLParser(parser.Parser):
     CONJUNCTION = {
         **parser.Parser.CONJUNCTION,
@@ -30,31 +34,53 @@ class PRQLParser(parser.Parser):
     }
 
     TRANSFORM_PARSERS = {
-        "DERIVE": lambda self, query: self._parse_selection(query),
-        "SELECT": lambda self, query: self._parse_selection(query, append=False),
-        "TAKE": lambda self, query: self._parse_take(query),
-        "FILTER": lambda self, query: query.where(self._parse_disjunction()),
-        "APPEND": lambda self, query: query.union(
-            _select_all(self._parse_table()), distinct=False, copy=False
-        ),
-        "REMOVE": lambda self, query: query.except_(
-            _select_all(self._parse_table()), distinct=False, copy=False
-        ),
-        "INTERSECT": lambda self, query: query.intersect(
-            _select_all(self._parse_table()), distinct=False, copy=False
-        ),
-        "SORT": lambda self, query: self._parse_order_by(query),
-        "AGGREGATE": lambda self, query: self._parse_selection(
-            query, parse_method=self._parse_aggregate, append=False
-        ),
+        "DERIVE": "_parse_transform_derive",
+        "SELECT": "_parse_transform_select",
+        "TAKE": "_parse_transform_take",
+        "FILTER": "_parse_transform_filter",
+        "APPEND": "_parse_transform_append",
+        "REMOVE": "_parse_transform_remove",
+        "INTERSECT": "_parse_transform_intersect",
+        "SORT": "_parse_transform_sort",
+        "AGGREGATE": "_parse_transform_aggregate",
     }
 
     FUNCTIONS = {
         **parser.Parser.FUNCTIONS,
         "AVERAGE": exp.Avg.from_arg_list,
-        "SUM": lambda args: exp.func("COALESCE", exp.Sum(this=seq_get(args, 0)), 0),
+        "SUM": _sum_with_zero,
     }
 
+    def _parse_transform_derive(self, query: exp.Query) -> exp.Query:
+        return self._parse_selection(query)
+
+    def _parse_transform_select(self, query: exp.Query) -> exp.Query:
+        return self._parse_selection(query, append=False)
+
+    def _parse_transform_take(self, query: exp.Query) -> exp.Query | None:
+        return self._parse_take(query)
+
+    def _parse_transform_filter(self, query: exp.Query) -> exp.Query:
+        return query.where(self._parse_disjunction())
+
+    def _parse_transform_append(self, query: exp.Query) -> exp.Query:
+        table = t.cast(exp.Expr, self._parse_table())
+        return query.union(_select_all(table), distinct=False, copy=False)
+
+    def _parse_transform_remove(self, query: exp.Query) -> exp.Query:
+        table = t.cast(exp.Expr, self._parse_table())
+        return query.except_(_select_all(table), distinct=False, copy=False)
+
+    def _parse_transform_intersect(self, query: exp.Query) -> exp.Query:
+        table = t.cast(exp.Expr, self._parse_table())
+        return query.intersect(_select_all(table), distinct=False, copy=False)
+
+    def _parse_transform_sort(self, query: exp.Query) -> exp.Query | None:
+        return self._parse_order_by(t.cast(exp.Select, query))
+
+    def _parse_transform_aggregate(self, query: exp.Query) -> exp.Query:
+        return self._parse_selection(query, parse_method=self._parse_aggregate, append=False)
+
     def _parse_equality(self) -> exp.Expr | None:
         eq = self._parse_comparison()
         while self._match_set(self.EQUALITY):
@@ -89,7 +115,8 @@ def _parse_query(self) -> exp.Query | None:
         query: exp.Query = exp.select("*").from_(from_, copy=False)
 
         while self._match_texts(self.TRANSFORM_PARSERS):
-            query = self.TRANSFORM_PARSERS[self._prev.text.upper()](self, query)
+            transform = self.TRANSFORM_PARSERS[self._prev.text.upper()]
+            query = getattr(self, transform)(query)
 
         return query
 

tests/dialects/test_bigquery.py

@@ -2262,6 +2262,19 @@ def test_warnings(self):
             self.assertIn("EXECUTE IMMEDIATE", cm.output[1])
             self.assertIn("END FOR", cm.output[2])
 
+        with self.assertLogs(parser_logger) as cm:
+            expression = parse_one(
+                'FOR tn IN (SELECT "101" AS number) DO EXECUTE IMMEDIATE FORMAT("SELECT %s", tn.number); END FOR',
+                read="bigquery",
+            )
+            self.assertEqual(
+                expression.sql(dialect="bigquery"),
+                "FOR tn IN (SELECT '101' AS number) DO EXECUTE IMMEDIATE FORMAT(\"SELECT %s\", tn.number); END FOR",
+            )
+            self.assertNotIn("FOR tn IN", "\n".join(cm.output))
+            self.assertEqual(len(cm.output), 1)
+            self.assertIn("'END FOR'", cm.output[0])
+
     def test_user_defined_functions(self):
         self.validate_identity(
             "CREATE TEMPORARY FUNCTION a(x FLOAT64, y FLOAT64) RETURNS FLOAT64 NOT DETERMINISTIC LANGUAGE js AS 'return x*y;'"