Fix comment reviews

Author: git-hulk

src/ast/mod.rs

@@ -43,14 +43,14 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
     ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
-    GroupByExpr, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator,
-    JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
+    GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint,
+    JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
     MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
     NonBlock, Offset, OffsetRows, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
     SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, SymbolDefinition, Table,
     TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode,
-    Values, WildcardAdditionalOptions, With, WithModifier,
+    Values, WildcardAdditionalOptions, With,
 };
 pub use self::value::{
     escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,

src/ast/query.rs

@@ -298,22 +298,7 @@ impl fmt::Display for Select {
         if let Some(ref selection) = self.selection {
             write!(f, " WHERE {selection}")?;
         }
-        match &self.group_by {
-            GroupByExpr::All(modifiers) => {
-                write!(f, " GROUP BY ALL")?;
-                if !modifiers.is_empty() {
-                    write!(f, " {}", display_separated(modifiers, " "))?;
-                }
-            }
-            GroupByExpr::Expressions(exprs, modifiers) => {
-                if !exprs.is_empty() {
-                    write!(f, " GROUP BY {}", display_comma_separated(exprs))?;
-                }
-                if !modifiers.is_empty() {
-                    write!(f, " {}", display_separated(modifiers, " "))?;
-                }
-            }
-        }
+        write!(f, "{}", self.group_by)?;
         if !self.cluster_by.is_empty() {
             write!(
                 f,
@@ -1876,18 +1861,22 @@ impl fmt::Display for SelectInto {
 #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub enum WithModifier {
+pub enum GroupByWithModifier {
+    /// ClickHouse supports GROUP BY WITH modifiers(includes ROLLUP|CUBE|TOTALS).
+    /// e.g. GROUP BY year WITH ROLLUP WITH TOTALS
+    ///
+    /// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
     Rollup,
     Cube,
     Totals,
 }
 
-impl fmt::Display for WithModifier {
+impl fmt::Display for GroupByWithModifier {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
-            WithModifier::Rollup => write!(f, "WITH ROLLUP"),
-            WithModifier::Cube => write!(f, "WITH CUBE"),
-            WithModifier::Totals => write!(f, "WITH TOTALS"),
+            GroupByWithModifier::Rollup => write!(f, "WITH ROLLUP"),
+            GroupByWithModifier::Cube => write!(f, "WITH CUBE"),
+            GroupByWithModifier::Totals => write!(f, "WITH TOTALS"),
         }
     }
 }
@@ -1905,24 +1894,31 @@ pub enum GroupByExpr {
     /// ClickHouse also supports WITH modifiers after GROUP BY ALL and expressions.
     ///
     /// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
-    All(Vec<WithModifier>),
+    All(Vec<GroupByWithModifier>),
 
     /// Expressions
-    Expressions(Vec<Expr>, Vec<WithModifier>),
+    Expressions(Vec<Expr>, Vec<GroupByWithModifier>),
 }
 
 impl fmt::Display for GroupByExpr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             GroupByExpr::All(modifiers) => {
-                write!(f, "GROUP BY ALL")?;
-                write!(f, " {}", display_separated(modifiers, " "))?;
+                write!(f, " GROUP BY ALL")?;
+                if !modifiers.is_empty() {
+                    write!(f, " {}", display_separated(modifiers, " "))?;
+                }
                 Ok(())
             }
             GroupByExpr::Expressions(col_names, modifiers) => {
+                if col_names.is_empty() {
+                    return Ok(());
+                }
                 let col_names = display_comma_separated(col_names);
-                write!(f, "GROUP BY ({col_names})")?;
-                write!(f, " {}", display_separated(modifiers, " "))?;
+                write!(f, " GROUP BY {col_names}")?;
+                if !modifiers.is_empty() {
+                    write!(f, " {}", display_separated(modifiers, " "))?;
+                }
                 Ok(())
             }
         }

src/parser/mod.rs

@@ -8294,21 +8294,22 @@ impl<'a> Parser<'a> {
             };
 
             let mut modifiers = vec![];
-            loop {
-                if dialect_of!(self is ClickHouseDialect | GenericDialect)
-                    && self.parse_keyword(Keyword::WITH)
-                {
-                    if self.parse_keyword(Keyword::ROLLUP) {
-                        modifiers.push(WithModifier::Rollup);
-                    } else if self.parse_keyword(Keyword::CUBE) {
-                        modifiers.push(WithModifier::Cube);
-                    } else if self.parse_keyword(Keyword::TOTALS) {
-                        modifiers.push(WithModifier::Totals);
-                    } else {
-                        self.expected("ROLLUP, CUBE, or TOTALS", self.peek_token())?;
+            if dialect_of!(self is ClickHouseDialect | GenericDialect) {
+                loop {
+                    if !self.parse_keyword(Keyword::WITH) {
+                        break;
                     }
-                } else {
-                    break;
+                    let keyword = self.expect_one_of_keywords(&[
+                        Keyword::ROLLUP,
+                        Keyword::CUBE,
+                        Keyword::TOTALS,
+                    ])?;
+                    modifiers.push(match keyword {
+                        Keyword::ROLLUP => GroupByWithModifier::Rollup,
+                        Keyword::CUBE => GroupByWithModifier::Cube,
+                        Keyword::TOTALS => GroupByWithModifier::Totals,
+                        _ => unreachable!(),
+                    });
                 }
             }
             match expressions {

tests/sqlparser_clickhouse.rs

@@ -628,19 +628,56 @@ fn parse_create_materialized_view() {
 
 #[test]
 fn parse_group_by_with_modifier() {
-    match clickhouse_and_generic().verified_stmt("SELECT * FROM t GROUP BY x WITH ROLLUP WITH CUBE")
-    {
-        Statement::Query(query) => {
-            let group_by = &query.body.as_select().unwrap().group_by;
-            assert_eq!(
-                group_by,
-                &GroupByExpr::Expressions(
-                    vec![Identifier(Ident::new("x"))],
-                    vec![WithModifier::Rollup, WithModifier::Cube]
-                )
-            );
+    let clauses = ["x", "a, b", "ALL"];
+    let modifiers = [
+        "WITH ROLLUP",
+        "WITH CUBE",
+        "WITH TOTALS",
+        "WITH ROLLUP WITH CUBE",
+    ];
+    let expected_modifiers = [
+        vec![GroupByWithModifier::Rollup],
+        vec![GroupByWithModifier::Cube],
+        vec![GroupByWithModifier::Totals],
+        vec![GroupByWithModifier::Rollup, GroupByWithModifier::Cube],
+    ];
+    for clause in &clauses {
+        for (modifier, expected_modifier) in modifiers.iter().zip(expected_modifiers.iter()) {
+            let sql = format!("SELECT * FROM t GROUP BY {clause} {modifier}");
+            match clickhouse_and_generic().verified_stmt(&sql) {
+                Statement::Query(query) => {
+                    let group_by = &query.body.as_select().unwrap().group_by;
+                    if clause == &"ALL" {
+                        assert_eq!(group_by, &GroupByExpr::All(expected_modifier.to_vec()));
+                    } else {
+                        assert_eq!(
+                            group_by,
+                            &GroupByExpr::Expressions(
+                                clause
+                                    .split(", ")
+                                    .map(|c| Identifier(Ident::new(c)))
+                                    .collect(),
+                                expected_modifier.to_vec()
+                            )
+                        );
+                    }
+                }
+                _ => unreachable!(),
+            }
         }
-        _ => unreachable!(),
+    }
+
+    // invalid cases
+    let invalid_cases = [
+        "SELECT * FROM t GROUP BY x WITH",
+        "SELECT * FROM t GROUP BY x WITH ROLLUP CUBE",
+        "SELECT * FROM t GROUP BY x WITH WITH ROLLUP",
+        "SELECT * FROM t GROUP BY WITH ROLLUP",
+    ];
+    for sql in invalid_cases {
+        clickhouse_and_generic()
+            .parse_sql_statements(sql)
+            .expect_err("Expected: one of ROLLUP or CUBE or TOTALS, found: WITH");
     }
 }