feat(bigquery): add EXPORT DATA statement support (#4688)

* feat(bigquery): add EXPORT DATA statement support

* feat(bigquery): improve EXPORT DATA implementation with named function

* refactor(bigquery): improve EXPORT DATA implementation

* feat(bigquery): make code compatible with 3.7

* feat(bigquery): add test with connection

* fix(bigquery): improve WITH CONNECTION parsing logic

sqlglot/dialects/bigquery.py

@@ -474,6 +474,7 @@ class Tokenizer(tokens.Tokenizer):
             "DECLARE": TokenType.COMMAND,
             "ELSEIF": TokenType.COMMAND,
             "EXCEPTION": TokenType.COMMAND,
+            "EXPORT": TokenType.EXPORT,
             "FLOAT64": TokenType.DOUBLE,
             "FOR SYSTEM_TIME": TokenType.TIMESTAMP_SNAPSHOT,
             "MODEL": TokenType.MODEL,
@@ -490,6 +491,11 @@ 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,
@@ -596,6 +602,7 @@ class Parser(parser.Parser):
             TokenType.ELSE: lambda self: self._parse_as_command(self._prev),
             TokenType.END: lambda self: self._parse_as_command(self._prev),
             TokenType.FOR: lambda self: self._parse_for_in(),
+            TokenType.EXPORT: lambda self: self._parse_export_data(),
         }
 
         BRACKET_OFFSETS = {
@@ -829,6 +836,50 @@ 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")
+
+            return self.expression(
+                exp.Export,
+                this=query,
+                with_connection=with_connection,
+                options=options,
+            )
+
     class Generator(generator.Generator):
         INTERVAL_ALLOWS_PLURAL_FORM = False
         JOIN_HINTS = False
@@ -1237,3 +1288,11 @@ 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

@@ -8650,3 +8650,11 @@ def null() -> Null:
     Boolean,
     Null,
 )
+
+
+class Export(Expression):
+    arg_types = {
+        "this": True,
+        "with_connection": False,
+        "options": False,
+    }

sqlglot/tokens.py

@@ -414,6 +414,7 @@ class TokenType(AutoName):
     SOURCE = auto()
     ANALYZE = auto()
     NAMESPACE = auto()
+    EXPORT = auto()
 
 
 _ALL_TOKEN_TYPES = list(TokenType)

tests/dialects/test_bigquery.py

@@ -1661,6 +1661,13 @@ 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"
+        )
+        self.validate_identity(
+            "EXPORT DATA WITH CONNECTION myproject.us.myconnection OPTIONS (URI='gs://path*.csv.gz', FORMAT='CSV') SELECT * FROM all_rows"
+        )
+
     def test_errors(self):
         with self.assertRaises(TokenError):
             transpile("'\\'", read="bigquery")