parse grouping sets, rollup, and cube for postgresql (#366)

* parse grouping sets, rollup, and cube

* add postgresql flag

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 {
@@ -376,6 +382,44 @@ impl fmt::Display for Expr {
             Expr::Exists(s) => write!(f, "EXISTS ({})", s),
             Expr::Subquery(s) => write!(f, "({})", s),
             Expr::ListAgg(listagg) => write!(f, "{}", listagg),
+            Expr::GroupingSets(sets) => {
+                write!(f, "GROUPING SETS (")?;
+                let mut sep = "";
+                for set in sets {
+                    write!(f, "{}", sep)?;
+                    sep = ", ";
+                    write!(f, "({})", display_comma_separated(set))?;
+                }
+                write!(f, ")")
+            }
+            Expr::Cube(sets) => {
+                write!(f, "CUBE (")?;
+                let mut sep = "";
+                for set in sets {
+                    write!(f, "{}", sep)?;
+                    sep = ", ";
+                    if set.len() == 1 {
+                        write!(f, "{}", set[0])?;
+                    } else {
+                        write!(f, "({})", display_comma_separated(set))?;
+                    }
+                }
+                write!(f, ")")
+            }
+            Expr::Rollup(sets) => {
+                write!(f, "ROLLUP (")?;
+                let mut sep = "";
+                for set in sets {
+                    write!(f, "{}", sep)?;
+                    sep = ", ";
+                    if set.len() == 1 {
+                        write!(f, "{}", set[0])?;
+                    } else {
+                        write!(f, "({})", display_comma_separated(set))?;
+                    }
+                }
+                write!(f, ")")
+            }
             Expr::Substring {
                 expr,
                 substring_from,
@@ -1903,4 +1947,93 @@ mod tests {
         let window_frame = WindowFrame::default();
         assert_eq!(WindowFrameBound::Preceding(None), window_frame.start_bound);
     }
+
+    #[test]
+    fn test_grouping_sets_display() {
+        // a and b in different group
+        let grouping_sets = Expr::GroupingSets(vec![
+            vec![Expr::Identifier(Ident::new("a"))],
+            vec![Expr::Identifier(Ident::new("b"))],
+        ]);
+        assert_eq!("GROUPING SETS ((a), (b))", format!("{}", grouping_sets));
+
+        // a and b in the same group
+        let grouping_sets = Expr::GroupingSets(vec![vec![
+            Expr::Identifier(Ident::new("a")),
+            Expr::Identifier(Ident::new("b")),
+        ]]);
+        assert_eq!("GROUPING SETS ((a, b))", format!("{}", grouping_sets));
+
+        // (a, b) and (c, d) in different group
+        let grouping_sets = Expr::GroupingSets(vec![
+            vec![
+                Expr::Identifier(Ident::new("a")),
+                Expr::Identifier(Ident::new("b")),
+            ],
+            vec![
+                Expr::Identifier(Ident::new("c")),
+                Expr::Identifier(Ident::new("d")),
+            ],
+        ]);
+        assert_eq!(
+            "GROUPING SETS ((a, b), (c, d))",
+            format!("{}", grouping_sets)
+        );
+    }
+
+    #[test]
+    fn test_rollup_display() {
+        let rollup = Expr::Rollup(vec![vec![Expr::Identifier(Ident::new("a"))]]);
+        assert_eq!("ROLLUP (a)", format!("{}", rollup));
+
+        let rollup = Expr::Rollup(vec![vec![
+            Expr::Identifier(Ident::new("a")),
+            Expr::Identifier(Ident::new("b")),
+        ]]);
+        assert_eq!("ROLLUP ((a, b))", format!("{}", rollup));
+
+        let rollup = Expr::Rollup(vec![
+            vec![Expr::Identifier(Ident::new("a"))],
+            vec![Expr::Identifier(Ident::new("b"))],
+        ]);
+        assert_eq!("ROLLUP (a, b)", format!("{}", rollup));
+
+        let rollup = Expr::Rollup(vec![
+            vec![Expr::Identifier(Ident::new("a"))],
+            vec![
+                Expr::Identifier(Ident::new("b")),
+                Expr::Identifier(Ident::new("c")),
+            ],
+            vec![Expr::Identifier(Ident::new("d"))],
+        ]);
+        assert_eq!("ROLLUP (a, (b, c), d)", format!("{}", rollup));
+    }
+
+    #[test]
+    fn test_cube_display() {
+        let cube = Expr::Cube(vec![vec![Expr::Identifier(Ident::new("a"))]]);
+        assert_eq!("CUBE (a)", format!("{}", cube));
+
+        let cube = Expr::Cube(vec![vec![
+            Expr::Identifier(Ident::new("a")),
+            Expr::Identifier(Ident::new("b")),
+        ]]);
+        assert_eq!("CUBE ((a, b))", format!("{}", cube));
+
+        let cube = Expr::Cube(vec![
+            vec![Expr::Identifier(Ident::new("a"))],
+            vec![Expr::Identifier(Ident::new("b"))],
+        ]);
+        assert_eq!("CUBE (a, b)", format!("{}", cube));
+
+        let cube = Expr::Cube(vec![
+            vec![Expr::Identifier(Ident::new("a"))],
+            vec![
+                Expr::Identifier(Ident::new("b")),
+                Expr::Identifier(Ident::new("c")),
+            ],
+            vec![Expr::Identifier(Ident::new("d"))],
+        ]);
+        assert_eq!("CUBE (a, (b, c), d)", format!("{}", cube));
+    }
 }

src/keywords.rs

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

src/parser.rs

@@ -338,7 +338,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
@@ -559,6 +559,68 @@ 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 dialect_of!(self is PostgreSqlDialect) {
+            if self.parse_keywords(&[Keyword::GROUPING, Keyword::SETS]) {
+                self.expect_token(&Token::LParen)?;
+                let result = self.parse_comma_separated(|p| p.parse_tuple(false, true))?;
+                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, 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, true))?;
+                self.expect_token(&Token::RParen)?;
+                Ok(Expr::Rollup(result))
+            } else {
+                self.parse_expr()
+            }
+        } else {
+            // TODO parse rollup for other dialects
+            self.parse_expr()
+        }
+    }
+
+    /// parse a tuple with `(` and `)`.
+    /// If `lift_singleton` is true, then a singleton tuple is lifted to a tuple of length 1, otherwise it will fail.
+    /// If `allow_empty` is true, then an empty tuple is allowed.
+    fn parse_tuple(
+        &mut self,
+        lift_singleton: bool,
+        allow_empty: bool,
+    ) -> Result<Vec<Expr>, ParserError> {
+        if lift_singleton {
+            if self.consume_token(&Token::LParen) {
+                let result = if allow_empty && self.consume_token(&Token::RParen) {
+                    vec![]
+                } else {
+                    let result = self.parse_comma_separated(Parser::parse_expr)?;
+                    self.expect_token(&Token::RParen)?;
+                    result
+                };
+                Ok(result)
+            } else {
+                Ok(vec![self.parse_expr()?])
+            }
+        } else {
+            self.expect_token(&Token::LParen)?;
+            let result = if allow_empty && self.consume_token(&Token::RParen) {
+                vec![]
+            } else {
+                let result = self.parse_comma_separated(Parser::parse_expr)?;
+                self.expect_token(&Token::RParen)?;
+                result
+            };
+            Ok(result)
+        }
+    }
+
     pub fn parse_case_expr(&mut self) -> Result<Expr, ParserError> {
         let mut operand = None;
         if !self.parse_keyword(Keyword::WHEN) {
@@ -2494,7 +2556,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![]
         };

tests/sqlparser_common.rs

@@ -20,13 +20,14 @@
 
 #[macro_use]
 mod test_utils;
-use test_utils::{all_dialects, expr_from_projection, join, number, only, table, table_alias};
-
 use matches::assert_matches;
 use sqlparser::ast::*;
-use sqlparser::dialect::{GenericDialect, SQLiteDialect};
+use sqlparser::dialect::{GenericDialect, PostgreSqlDialect, SQLiteDialect};
 use sqlparser::keywords::ALL_KEYWORDS;
 use sqlparser::parser::{Parser, ParserError};
+use test_utils::{
+    all_dialects, expr_from_projection, join, number, only, table, table_alias, TestedDialects,
+};
 
 #[test]
 fn parse_insert_values() {
@@ -1039,6 +1040,65 @@ fn parse_select_group_by() {
     );
 }
 
+#[test]
+fn parse_select_group_by_grouping_sets() {
+    let dialects = TestedDialects {
+        dialects: vec![Box::new(PostgreSqlDialect {})],
+    };
+    let sql =
+        "SELECT brand, size, sum(sales) FROM items_sold GROUP BY size, GROUPING SETS ((brand), (size), ())";
+    let select = dialects.verified_only_select(sql);
+    assert_eq!(
+        vec![
+            Expr::Identifier(Ident::new("size")),
+            Expr::GroupingSets(vec![
+                vec![Expr::Identifier(Ident::new("brand"))],
+                vec![Expr::Identifier(Ident::new("size"))],
+                vec![],
+            ])
+        ],
+        select.group_by
+    );
+}
+
+#[test]
+fn parse_select_group_by_rollup() {
+    let dialects = TestedDialects {
+        dialects: vec![Box::new(PostgreSqlDialect {})],
+    };
+    let sql = "SELECT brand, size, sum(sales) FROM items_sold GROUP BY size, ROLLUP (brand, size)";
+    let select = dialects.verified_only_select(sql);
+    assert_eq!(
+        vec![
+            Expr::Identifier(Ident::new("size")),
+            Expr::Rollup(vec![
+                vec![Expr::Identifier(Ident::new("brand"))],
+                vec![Expr::Identifier(Ident::new("size"))],
+            ])
+        ],
+        select.group_by
+    );
+}
+
+#[test]
+fn parse_select_group_by_cube() {
+    let dialects = TestedDialects {
+        dialects: vec![Box::new(PostgreSqlDialect {})],
+    };
+    let sql = "SELECT brand, size, sum(sales) FROM items_sold GROUP BY size, CUBE (brand, size)";
+    let select = dialects.verified_only_select(sql);
+    assert_eq!(
+        vec![
+            Expr::Identifier(Ident::new("size")),
+            Expr::Cube(vec![
+                vec![Expr::Identifier(Ident::new("brand"))],
+                vec![Expr::Identifier(Ident::new("size"))],
+            ])
+        ],
+        select.group_by
+    );
+}
+
 #[test]
 fn parse_select_having() {
     let sql = "SELECT foo FROM bar GROUP BY foo HAVING COUNT(*) > 1";