with connection

Author: chenkovsky

src/ast/mod.rs

@@ -10146,16 +10146,26 @@ impl fmt::Display for MemberOf {
 pub struct ExportData {
     pub options: Vec<SqlOption>,
     pub query: Box<Query>,
+    pub connection: Option<ObjectName>,
 }
 
 impl fmt::Display for ExportData {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(
-            f,
-            "EXPORT DATA OPTIONS({}) AS {}",
-            display_comma_separated(&self.options),
-            self.query
-        )
+        if let Some(connection) = &self.connection {
+            write!(
+                f,
+                "EXPORT DATA WITH CONNECTION {connection} OPTIONS({}) AS {}",
+                display_comma_separated(&self.options),
+                self.query
+            )
+        } else {
+            write!(
+                f,
+                "EXPORT DATA OPTIONS({}) AS {}",
+                display_comma_separated(&self.options),
+                self.query
+            )
+        }
     }
 }
 /// Creates a user

src/ast/spans.rs

@@ -531,11 +531,16 @@ impl Spanned for Statement {
             Statement::Print { .. } => Span::empty(),
             Statement::Return { .. } => Span::empty(),
             Statement::List(..) | Statement::Remove(..) => Span::empty(),
-            Statement::ExportData(ExportData { options, query }) => union_spans(
+            Statement::ExportData(ExportData {
+                options,
+                query,
+                connection,
+            }) => union_spans(
                 options
                     .iter()
                     .map(|i| i.span())
-                    .chain(core::iter::once(query.span())),
+                    .chain(core::iter::once(query.span()))
+                    .chain(connection.iter().map(|i| i.span())),
             ),
             Statement::CreateUser(..) => Span::empty(),
         }

src/parser/mod.rs

@@ -16513,13 +16513,23 @@ impl<'a> Parser<'a> {
 
     fn parse_export(&mut self) -> Result<Statement, ParserError> {
         self.expect_keyword(Keyword::DATA)?;
+
+        let connection = if self.parse_keywords(&[Keyword::WITH, Keyword::CONNECTION]) {
+            Some(self.parse_object_name(false)?)
+        } else {
+            None
+        };
         self.expect_keyword(Keyword::OPTIONS)?;
         self.expect_token(&Token::LParen)?;
         let options = self.parse_comma_separated(|p| p.parse_sql_option())?;
         self.expect_token(&Token::RParen)?;
         self.expect_keyword(Keyword::AS)?;
         let query = self.parse_query()?;
-        Ok(Statement::ExportData(ExportData { options, query }))
+        Ok(Statement::ExportData(ExportData {
+            options,
+            query,
+            connection,
+        }))
     }
 
     /// Consume the parser and return its underlying token buffer

tests/sqlparser_bigquery.rs

@@ -2577,4 +2577,13 @@ fn test_export() {
         ") AS ",
         "SELECT field1, field2 FROM mydataset.table1 ORDER BY field1 LIMIT 10",
     ));
+
+    bigquery().verified_stmt(concat!(
+        "EXPORT DATA WITH CONNECTION myconnection.myproject.us OPTIONS(",
+        "uri = 'gs://bucket/folder/*', ",
+        "format = 'PARQUET', ",
+        "overwrite = true",
+        ") AS ",
+        "SELECT field1, field2 FROM mydataset.table1 ORDER BY field1 LIMIT 10",
+    ));
 }