feat: SELECT * REPLACE <Expr> AS <Identifier> for bigquery (#798)

* chore: add test for wildcard replace

* feat: define opt_replace for wildcard replace

* fix: modify replace option ast

* fix: add test cases

* chore: fmt

* redefine ast

* feat: parse select replace items

* ci

* Update src/ast/query.rs

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: togami2864

src/ast/mod.rs

@@ -36,8 +36,9 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, IdentWithAlias, Join, JoinConstraint,
     JoinOperator, LateralView, LockClause, LockType, NonBlock, Offset, OffsetRows, OrderByExpr,
-    Query, RenameSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier,
-    Table, TableAlias, TableFactor, TableWithJoins, Top, Values, WildcardAdditionalOptions, With,
+    Query, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto,
+    SelectItem, SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor,
+    TableWithJoins, Top, Values, WildcardAdditionalOptions, With,
 };
 pub use self::value::{
     escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,

src/ast/query.rs

@@ -394,6 +394,9 @@ pub struct WildcardAdditionalOptions {
     pub opt_except: Option<ExceptSelectItem>,
     /// `[RENAME ...]`.
     pub opt_rename: Option<RenameSelectItem>,
+    /// `[REPLACE]`
+    ///  BigQuery syntax: <https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#select_replace>
+    pub opt_replace: Option<ReplaceSelectItem>,
 }
 
 impl fmt::Display for WildcardAdditionalOptions {
@@ -407,6 +410,9 @@ impl fmt::Display for WildcardAdditionalOptions {
         if let Some(rename) = &self.opt_rename {
             write!(f, " {rename}")?;
         }
+        if let Some(replace) = &self.opt_replace {
+            write!(f, " {replace}")?;
+        }
         Ok(())
     }
 }
@@ -526,6 +532,51 @@ impl fmt::Display for ExceptSelectItem {
     }
 }
 
+/// Bigquery `REPLACE` information.
+///
+/// # Syntax
+/// ```plaintext
+/// REPLACE (<new_expr> [AS] <col_name>)
+/// REPLACE (<col_name> [AS] <col_alias>, <col_name> [AS] <col_alias>, ...)
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ReplaceSelectItem {
+    pub items: Vec<Box<ReplaceSelectElement>>,
+}
+
+impl fmt::Display for ReplaceSelectItem {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "REPLACE")?;
+        write!(f, " ({})", display_comma_separated(&self.items))?;
+        Ok(())
+    }
+}
+
+/// # Syntax
+/// ```plaintext
+/// <expr> [AS] <column_name>
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ReplaceSelectElement {
+    pub expr: Expr,
+    pub colum_name: Ident,
+    pub as_keyword: bool,
+}
+
+impl fmt::Display for ReplaceSelectElement {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if self.as_keyword {
+            write!(f, "{} AS {}", self.expr, self.colum_name)
+        } else {
+            write!(f, "{} {}", self.expr, self.colum_name)
+        }
+    }
+}
+
 impl fmt::Display for SelectItem {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match &self {

src/parser.rs

@@ -6164,10 +6164,17 @@ impl<'a> Parser<'a> {
             None
         };
 
+        let opt_replace = if dialect_of!(self is GenericDialect | BigQueryDialect) {
+            self.parse_optional_select_item_replace()?
+        } else {
+            None
+        };
+
         Ok(WildcardAdditionalOptions {
             opt_exclude,
             opt_except,
             opt_rename,
+            opt_replace,
         })
     }
 
@@ -6241,6 +6248,38 @@ impl<'a> Parser<'a> {
         Ok(opt_rename)
     }
 
+    /// Parse a [`Replace`](ReplaceSelectItem) information for wildcard select items.
+    pub fn parse_optional_select_item_replace(
+        &mut self,
+    ) -> Result<Option<ReplaceSelectItem>, ParserError> {
+        let opt_replace = if self.parse_keyword(Keyword::REPLACE) {
+            if self.consume_token(&Token::LParen) {
+                let items = self.parse_comma_separated(|parser| {
+                    Ok(Box::new(parser.parse_replace_elements()?))
+                })?;
+                self.expect_token(&Token::RParen)?;
+                Some(ReplaceSelectItem { items })
+            } else {
+                let tok = self.next_token();
+                return self.expected("( after REPLACE but", tok);
+            }
+        } else {
+            None
+        };
+
+        Ok(opt_replace)
+    }
+    pub fn parse_replace_elements(&mut self) -> Result<ReplaceSelectElement, ParserError> {
+        let expr = self.parse_expr()?;
+        let as_keyword = self.parse_keyword(Keyword::AS);
+        let ident = self.parse_identifier()?;
+        Ok(ReplaceSelectElement {
+            expr,
+            colum_name: ident,
+            as_keyword,
+        })
+    }
+
     /// Parse an expression, optionally followed by ASC or DESC (used in ORDER BY)
     pub fn parse_order_by_expr(&mut self) -> Result<OrderByExpr, ParserError> {
         let expr = self.parse_expr()?;

tests/sqlparser_bigquery.rs

@@ -17,6 +17,11 @@ use sqlparser::ast::*;
 use sqlparser::dialect::{BigQueryDialect, GenericDialect};
 use test_utils::*;
 
+#[cfg(feature = "bigdecimal")]
+use bigdecimal::*;
+#[cfg(feature = "bigdecimal")]
+use std::str::FromStr;
+
 #[test]
 fn parse_literal_string() {
     let sql = r#"SELECT 'single', "double""#;
@@ -313,6 +318,58 @@ fn test_select_wildcard_with_except() {
     );
 }
 
+#[test]
+fn test_select_wildcard_with_replace() {
+    let select = bigquery_and_generic()
+        .verified_only_select(r#"SELECT * REPLACE ('widget' AS item_name) FROM orders"#);
+    let expected = SelectItem::Wildcard(WildcardAdditionalOptions {
+        opt_replace: Some(ReplaceSelectItem {
+            items: vec![Box::new(ReplaceSelectElement {
+                expr: Expr::Value(Value::SingleQuotedString("widget".to_owned())),
+                colum_name: Ident::new("item_name"),
+                as_keyword: true,
+            })],
+        }),
+        ..Default::default()
+    });
+    assert_eq!(expected, select.projection[0]);
+
+    let select = bigquery_and_generic().verified_only_select(
+        r#"SELECT * REPLACE (quantity / 2 AS quantity, 3 AS order_id) FROM orders"#,
+    );
+    let expected = SelectItem::Wildcard(WildcardAdditionalOptions {
+        opt_replace: Some(ReplaceSelectItem {
+            items: vec![
+                Box::new(ReplaceSelectElement {
+                    expr: Expr::BinaryOp {
+                        left: Box::new(Expr::Identifier(Ident::new("quantity"))),
+                        op: BinaryOperator::Divide,
+                        #[cfg(not(feature = "bigdecimal"))]
+                        right: Box::new(Expr::Value(Value::Number("2".to_string(), false))),
+                        #[cfg(feature = "bigdecimal")]
+                        right: Box::new(Expr::Value(Value::Number(
+                            BigDecimal::from_str("2").unwrap(),
+                            false,
+                        ))),
+                    },
+                    colum_name: Ident::new("quantity"),
+                    as_keyword: true,
+                }),
+                Box::new(ReplaceSelectElement {
+                    #[cfg(not(feature = "bigdecimal"))]
+                    expr: Expr::Value(Value::Number("3".to_string(), false)),
+                    #[cfg(feature = "bigdecimal")]
+                    expr: Expr::Value(Value::Number(BigDecimal::from_str("3").unwrap(), false)),
+                    colum_name: Ident::new("order_id"),
+                    as_keyword: true,
+                }),
+            ],
+        }),
+        ..Default::default()
+    });
+    assert_eq!(expected, select.projection[0]);
+}
+
 fn bigquery() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(BigQueryDialect {})],