bq: parse wildcard modifiers except/replace

Author: donhcd

src/ast/query.rs

@@ -207,19 +207,52 @@ pub enum SelectItem {
     UnnamedExpr(Expr),
     /// An expression, followed by `[ AS ] alias`
     ExprWithAlias { expr: Expr, alias: Ident },
-    /// `alias.*` or even `schema.table.*`
-    QualifiedWildcard(ObjectName),
-    /// An unqualified `*`
-    Wildcard,
+    /// if obj.is_some(), `alias.*` or even `schema.table.*`
+    /// else an unqualified `*`
+    /// except and replace are currently expected in bigquery
+    /// https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#modifiers_for_operator
+    Wildcard {
+        prefix: Option<ObjectName>,
+        except: Vec<Ident>,
+        replace: Vec<(Expr, Ident)>,
+    },
 }
 
 impl fmt::Display for SelectItem {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match &self {
             SelectItem::UnnamedExpr(expr) => write!(f, "{}", expr),
             SelectItem::ExprWithAlias { expr, alias } => write!(f, "{} AS {}", expr, alias),
-            SelectItem::QualifiedWildcard(prefix) => write!(f, "{}.*", prefix),
-            SelectItem::Wildcard => write!(f, "*"),
+            SelectItem::Wildcard {
+                prefix,
+                except,
+                replace,
+            } => {
+                if let Some(pre) = prefix {
+                    write!(f, "{}.*", pre)?;
+                } else {
+                    write!(f, "*")?;
+                }
+                let mut delim = "";
+                if !except.is_empty() {
+                    write!(f, " EXCEPT (")?;
+                    for col in except {
+                        write!(f, "{}{}", delim, col)?;
+                        delim = ", ";
+                    }
+                    write!(f, ")")?;
+                }
+                delim = "";
+                if !replace.is_empty() {
+                    write!(f, " REPLACE (")?;
+                    for &(ref expr, ref alias) in replace.iter() {
+                        write!(f, "{}{} AS {}", delim, expr, alias)?;
+                        delim = ", ";
+                    }
+                    write!(f, ")")?;
+                }
+                Ok(())
+            }
         }
     }
 }

src/parser.rs

@@ -2814,9 +2814,19 @@ impl<'a> Parser<'a> {
     pub fn parse_select_item(&mut self) -> Result<SelectItem, ParserError> {
         let expr = self.parse_expr()?;
         if let Expr::Wildcard = expr {
-            Ok(SelectItem::Wildcard)
+            let (except, replace) = self.parse_wildcard_modifiers()?;
+            Ok(SelectItem::Wildcard {
+                prefix: None,
+                except,
+                replace,
+            })
         } else if let Expr::QualifiedWildcard(prefix) = expr {
-            Ok(SelectItem::QualifiedWildcard(ObjectName(prefix)))
+            let (except, replace) = self.parse_wildcard_modifiers()?;
+            Ok(SelectItem::Wildcard {
+                prefix: Some(ObjectName(prefix)),
+                except,
+                replace,
+            })
         } else {
             // `expr` is a regular SQL expression and can be followed by an alias
             if let Some(alias) = self.parse_optional_alias(keywords::RESERVED_FOR_COLUMN_ALIAS)? {
@@ -2827,6 +2837,40 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse a comma-delimited list of projections after REPLACE
+    /// https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#select_replace
+    pub fn parse_replace_item(&mut self) -> Result<(Expr, Ident), ParserError> {
+        let expr = self.parse_expr()?;
+        // `expr` is a regular SQL expression and can be followed by an alias
+        if let Some(alias) = self.parse_optional_alias(keywords::RESERVED_FOR_COLUMN_ALIAS)? {
+            Ok((expr, alias))
+        } else {
+            parser_err!(self, "REPLACE expression must have alias")
+        }
+    }
+
+    pub fn parse_wildcard_modifiers(
+        &mut self,
+    ) -> Result<(Vec<Ident>, Vec<(Expr, Ident)>), ParserError> {
+        let except = if self.parse_keyword(Keyword::EXCEPT) {
+            self.expect_token(&Token::LParen)?;
+            let aliases = self.parse_comma_separated(Parser::parse_identifier)?;
+            self.expect_token(&Token::RParen)?;
+            aliases
+        } else {
+            vec![]
+        };
+        let replace = if self.parse_keyword(Keyword::EXCEPT) {
+            self.expect_token(&Token::LParen)?;
+            let replace = self.parse_comma_separated(Parser::parse_replace_item)?;
+            self.expect_token(&Token::RParen)?;
+            replace
+        } else {
+            vec![]
+        };
+        Ok((except, replace))
+    }
+
     /// 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()?;