fix(bigquery): Refactor EXPORT DATA statement (#4693)

Author: VaggelisD

sqlglot/dialects/bigquery.py

@@ -491,11 +491,6 @@ class Parser(parser.Parser):
         LOG_DEFAULTS_TO_LN = True
         SUPPORTS_IMPLICIT_UNNEST = True
 
-        ID_VAR_TOKENS = {
-            *parser.Parser.ID_VAR_TOKENS,
-            TokenType.EXPORT,
-        }
-
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
             "CONTAINS_SUBSTR": _build_contains_substring,
@@ -837,47 +832,13 @@ def _parse_features_at_time(self) -> exp.FeaturesAtTime:
             return expr
 
         def _parse_export_data(self) -> exp.Export:
-            # https://cloud.google.com/bigquery/docs/reference/standard-sql/export-statements
-            if not self._match_text_seq("DATA"):
-                self.raise_error("Expected 'DATA' after 'EXPORT'")
-
-            with_connection = None
-            options = None
-
-            if self._match_text_seq("WITH", "CONNECTION"):
-                parts = []
-                while True:
-                    part = self._parse_var()
-                    if not part:
-                        break
-                    parts.append(part.name)
-                    if not self._match(TokenType.DOT):
-                        break
-
-                if not parts:
-                    self.raise_error("Expected connection name after WITH CONNECTION")
-
-                with_connection = exp.Identifier(this=".".join(parts))
-
-            if self._match_text_seq("OPTIONS"):
-                self._match(TokenType.L_PAREN)
-                options = self._parse_properties()
-                self._match(TokenType.R_PAREN)
-            else:
-                self.raise_error("Expected 'OPTIONS' after 'EXPORT DATA'")
-
-            self._match_text_seq("AS")
-
-            # Parse the full SELECT statement
-            query = self._parse_statement()
-            if not isinstance(query, exp.Select):
-                self.raise_error("Expected SELECT statement in EXPORT DATA")
+            self._match_text_seq("DATA")
 
             return self.expression(
                 exp.Export,
-                this=query,
-                with_connection=with_connection,
-                options=options,
+                connection=self._match_text_seq("WITH", "CONNECTION") and self._parse_table_parts(),
+                options=self._parse_properties(),
+                this=self._match_text_seq("AS") and self._parse_select(),
             )
 
     class Generator(generator.Generator):
@@ -1288,11 +1249,3 @@ def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) ->
                 return f"{self.sql(expression, 'to')}{self.sql(this)}"
 
             return super().cast_sql(expression, safe_prefix=safe_prefix)
-
-        def export_sql(self, expression: exp.Export) -> str:
-            this = self.sql(expression, "this")
-            with_connection = self.sql(expression, "with_connection")
-            with_connection = f"WITH CONNECTION {with_connection} " if with_connection else ""
-            options = self.sql(expression, "options")
-            options = f"{options} " if options else ""
-            return f"EXPORT DATA {with_connection}{options}{this}"

sqlglot/expressions.py

@@ -2064,6 +2064,11 @@ def kind(self) -> t.Optional[str]:
         return kind and kind.upper()
 
 
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/export-statements
+class Export(Expression):
+    arg_types = {"this": True, "connection": False, "options": True}
+
+
 class Filter(Expression):
     arg_types = {"this": True, "expression": True}
 
@@ -8650,11 +8655,3 @@ def null() -> Null:
     Boolean,
     Null,
 )
-
-
-class Export(Expression):
-    arg_types = {
-        "this": True,
-        "with_connection": False,
-        "options": False,
-    }

sqlglot/generator.py

@@ -4752,3 +4752,10 @@ def xmltable_sql(self, expression: exp.XMLTable) -> str:
     def xmlnamespace_sql(self, expression: exp.XMLNamespace) -> str:
         this = self.sql(expression, "this")
         return this if isinstance(expression.this, exp.Alias) else f"DEFAULT {this}"
+
+    def export_sql(self, expression: exp.Export) -> str:
+        this = self.sql(expression, "this")
+        connection = self.sql(expression, "connection")
+        connection = f"WITH CONNECTION {connection} " if connection else ""
+        options = self.sql(expression, "options")
+        return f"EXPORT DATA {connection}{options} AS {this}"

sqlglot/parser.py

@@ -497,6 +497,7 @@ class Parser(metaclass=_Parser):
         TokenType.DIV,
         TokenType.END,
         TokenType.EXECUTE,
+        TokenType.EXPORT,
         TokenType.ESCAPE,
         TokenType.FALSE,
         TokenType.FIRST,

tests/dialects/test_bigquery.py

@@ -1662,10 +1662,10 @@ def test_bigquery(self):
         self.validate_identity("SELECT * FROM ML.FEATURES_AT_TIME((SELECT 1), num_rows => 1)")
 
         self.validate_identity(
-            "EXPORT DATA OPTIONS (URI='gs://path*.csv.gz', FORMAT='CSV') SELECT * FROM all_rows"
+            "EXPORT DATA OPTIONS (URI='gs://path*.csv.gz', FORMAT='CSV') AS SELECT * FROM all_rows"
         )
         self.validate_identity(
-            "EXPORT DATA WITH CONNECTION myproject.us.myconnection OPTIONS (URI='gs://path*.csv.gz', FORMAT='CSV') SELECT * FROM all_rows"
+            "EXPORT DATA WITH CONNECTION myproject.us.myconnection OPTIONS (URI='gs://path*.csv.gz', FORMAT='CSV') AS SELECT * FROM all_rows"
         )
 
     def test_errors(self):

tests/fixtures/identity.sql

@@ -888,3 +888,4 @@ SELECT 1 LIMIT 1
 CAST(x AS INT128)
 CAST(x AS UINT128)
 CAST(x AS UINT256)
+SELECT export