parse grouping sets, rollup, and cube

Author: jimexist

src/ast/mod.rs

@@ -278,6 +278,12 @@ pub enum Expr {
     Subquery(Box<Query>),
     /// The `LISTAGG` function `SELECT LISTAGG(...) WITHIN GROUP (ORDER BY ...)`
     ListAgg(ListAgg),
+    /// The `GROUPING SETS` expr.
+    GroupingSets(Vec<Vec<Expr>>),
+    /// The `CUBE` expr.
+    Cube(Vec<Vec<Expr>>),
+    /// The `ROLLUP` expr.
+    Rollup(Vec<Vec<Expr>>),
 }
 
 impl fmt::Display for Expr {

src/keywords.rs

@@ -401,6 +401,7 @@ define_keywords!(
     SESSION,
     SESSION_USER,
     SET,
+    SETS,
     SHOW,
     SIMILAR,
     SMALLINT,

src/parser.rs

@@ -336,7 +336,7 @@ impl<'a> Parser<'a> {
         return_ok_if_some!(self.maybe_parse(|parser| {
             match parser.parse_data_type()? {
                 DataType::Interval => parser.parse_literal_interval(),
-                // PosgreSQL allows almost any identifier to be used as custom data type name,
+                // PostgreSQL allows almost any identifier to be used as custom data type name,
                 // and we support that in `parse_data_type()`. But unlike Postgres we don't
                 // have a list of globally reserved keywords (since they vary across dialects),
                 // so given `NOT 'a' LIKE 'b'`, we'd accept `NOT` as a possible custom data type
@@ -557,6 +557,48 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// parse a group by expr. a group by expr can be one of group sets, roll up, cube, or simple
+    /// expr.
+    fn parse_group_by_expr(&mut self) -> Result<Expr, ParserError> {
+        if self.parse_keywords(&[Keyword::GROUPING, Keyword::SETS]) {
+            self.expect_token(&Token::LParen)?;
+            let result = self.parse_comma_separated(|p| p.parse_tuple(false))?;
+            self.expect_token(&Token::RParen)?;
+            Ok(Expr::GroupingSets(result))
+        } else if self.parse_keyword(Keyword::CUBE) {
+            self.expect_token(&Token::LParen)?;
+            let result = self.parse_comma_separated(|p| p.parse_tuple(true))?;
+            self.expect_token(&Token::RParen)?;
+            Ok(Expr::Cube(result))
+        } else if self.parse_keyword(Keyword::ROLLUP) {
+            self.expect_token(&Token::LParen)?;
+            let result = self.parse_comma_separated(|p| p.parse_tuple(true))?;
+            self.expect_token(&Token::RParen)?;
+            Ok(Expr::Rollup(result))
+        } else {
+            self.parse_expr()
+        }
+    }
+
+    /// parse a tuple with `(` and `)`, allowing for empty tuples. If lift_singleton is true,
+    /// then a singleton tuple is lifted to a tuple of length 1, otherwise it will fail.
+    fn parse_tuple(&mut self, lift_singleton: bool) -> Result<Vec<Expr>, ParserError> {
+        if lift_singleton {
+            if self.consume_token(&Token::LParen) {
+                let result = self.parse_comma_separated(Parser::parse_expr)?;
+                self.expect_token(&Token::RParen)?;
+                Ok(result)
+            } else {
+                Ok(vec![self.parse_expr()?])
+            }
+        } else {
+            self.expect_token(&Token::LParen)?;
+            let result = self.parse_comma_separated(Parser::parse_expr)?;
+            self.expect_token(&Token::RParen)?;
+            Ok(result)
+        }
+    }
+
     pub fn parse_case_expr(&mut self) -> Result<Expr, ParserError> {
         let mut operand = None;
         if !self.parse_keyword(Keyword::WHEN) {
@@ -2492,7 +2534,7 @@ impl<'a> Parser<'a> {
         };
 
         let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
-            self.parse_comma_separated(Parser::parse_expr)?
+            self.parse_comma_separated(Parser::parse_group_by_expr)?
         } else {
             vec![]
         };