Fix: clickhouse group by with totals closes #1609

tobymao · May 12, 2023 · 2f34d6f · 2f34d6f
1 parent 72f1984
commit 2f34d6f
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 3 deletions.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -223,6 +223,7 @@ class Generator(generator.Generator):
         JOIN_HINTS = False
         TABLE_HINTS = False
         EXPLICIT_UNION = True
+        GROUPINGS_SEP = ""
 
         def _param_args_sql(
             self,

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1518,6 +1518,7 @@ class Group(Expression):
         "grouping_sets": False,
         "cube": False,
         "rollup": False,
+        "totals": False,
     }
 
 

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -137,6 +137,9 @@ class Generator:
     # Whether a table is allowed to be renamed with a db
     RENAME_TABLE_WITH_DB = True
 
+    # The separator for grouping sets and rollups
+    GROUPINGS_SEP = ","
+
     TYPE_MAPPING = {
         exp.DataType.Type.NCHAR: "CHAR",
         exp.DataType.Type.NVARCHAR: "VARCHAR",
@@ -1259,10 +1262,16 @@ def group_sql(self, expression: exp.Group) -> str:
             rollup_sql = self.expressions(expression, key="rollup", indent=False)
             rollup_sql = f"{self.seg('ROLLUP')} {self.wrap(rollup_sql)}" if rollup_sql else ""
 
-        groupings = csv(grouping_sets, cube_sql, rollup_sql, sep=",")
+        groupings = csv(
+            grouping_sets,
+            cube_sql,
+            rollup_sql,
+            self.seg("WITH TOTALS") if expression.args.get("totals") else "",
+            sep=self.GROUPINGS_SEP,
+        )
 
         if expression.args.get("expressions") and groupings:
-            group_by = f"{group_by},"
+            group_by = f"{group_by}{self.GROUPINGS_SEP}"
 
         return f"{group_by}{groupings}"
 

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -2522,6 +2522,7 @@ def _parse_group(self, skip_group_by_token: bool = False) -> t.Optional[exp.Expr
 
             rollup = None
             cube = None
+            totals = None
 
             with_ = self._match(TokenType.WITH)
             if self._match(TokenType.ROLLUP):
@@ -2532,7 +2533,11 @@ def _parse_group(self, skip_group_by_token: bool = False) -> t.Optional[exp.Expr
                 cube = with_ or self._parse_wrapped_csv(self._parse_column)
                 elements["cube"].extend(ensure_list(cube))
 
-            if not (grouping_sets or rollup or cube):
+            if self._match_text_seq("TOTALS"):
+                totals = True
+                elements["totals"] = True  # type: ignore
+
+            if not (grouping_sets or rollup or cube or totals):
                 break
 
         return self.expression(exp.Group, **elements)  # type: ignore

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -8,6 +8,7 @@ class TestClickhouse(Validator):
     def test_clickhouse(self):
         self.validate_identity("ATTACH DATABASE DEFAULT ENGINE = ORDINARY")
         self.validate_identity("CAST(['hello'], 'Array(Enum8(''hello'' = 1))')")
+        self.validate_identity("SELECT x, COUNT() FROM y GROUP BY x WITH TOTALS")
         self.validate_identity("SELECT INTERVAL t.days day")
         self.validate_identity("SELECT match('abc', '([a-z]+)')")
         self.validate_identity("dictGet(x, 'y')")