diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py index 2ae35239d1..352da79093 100644 --- a/sqlglot/dialects/bigquery.py +++ b/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}" diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py index ecb516a930..6a9bf50b08 100644 --- a/sqlglot/expressions.py +++ b/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, - } diff --git a/sqlglot/generator.py b/sqlglot/generator.py index 65337639bb..2b6ef97fd3 100644 --- a/sqlglot/generator.py +++ b/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}" diff --git a/sqlglot/parser.py b/sqlglot/parser.py index 295319996a..34027d6b91 100644 --- a/sqlglot/parser.py +++ b/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, diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py index 08c949b185..b4e4793249 100644 --- a/tests/dialects/test_bigquery.py +++ b/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): diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql index aa0ed403e3..34dd90da1d 100644 --- a/tests/fixtures/identity.sql +++ b/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 \ No newline at end of file