Repository: tobymao/sqlglot

Issue:
Hive Dialect for "WITH ROLLUP" error
when parsing hive sql, there is a Parse Error

![image](https://user-images.githubusercontent.com/34465781/215717746-cfa1fca6-ca43-4ac9-8a7d-d12958b981a6.png)
![image](https://user-images.githubusercontent.com/34465781/215717796-da538963-9eb5-4a8c-86ba-86f0af9afb3c.png)
AND **the error** is 
![image](https://user-images.githubusercontent.com/34465781/215718058-87b636f2-7422-4637-a0bf-d528deb23129.png)

This SQL case is just copied from the official documentation https://cwiki.apache.org/confluence/display/Hive/Enhanced+Aggregation%2C+Cube%2C+Grouping+and+Rollup

Discussion:
GeorgeSittas: An alternative for this one is to add a flag in the base generator to switch between the previous supported syntax and Hive's `WITH CUBE/ROLLUP`, so that we duplicate less code over to `hive.py`. Happy to go either way.
GeorgeSittas: > does hive support the standard way?

The CUBE/ROLLUP docs didn't have an example, so I would guess not.
tobymao: can we do a flag on parser and gen?
GeorgeSittas: Yeah I'll fix it
GeorgeSittas: @tobymao should we rather add the "WITH CUBE/ROLLUP" logic in the base Parser, Generator without using a flag approach? That way we can capture both syntax forms without having to exclude either with a flag.
tobymao: yes that's what i was suggesting
GeorgeSittas: Ok I made the changes. It feels like it can be simplified more, let me know if I missed a cleaner way to implement this.

Patch:
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index c2547822..a3b18513 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -832,10 +832,21 @@ class Generator:
         grouping_sets = (
             f"{self.seg('GROUPING SETS')} {self.wrap(grouping_sets)}" if grouping_sets else ""
         )
-        cube = self.expressions(expression, key="cube", indent=False)
-        cube = f"{self.seg('CUBE')} {self.wrap(cube)}" if cube else ""
-        rollup = self.expressions(expression, key="rollup", indent=False)
-        rollup = f"{self.seg('ROLLUP')} {self.wrap(rollup)}" if rollup else ""
+
+        cube = expression.args.get("cube")
+        if cube is True:
+            cube = self.seg("WITH CUBE")
+        else:
+            cube = self.expressions(expression, key="cube", indent=False)
+            cube = f"{self.seg('CUBE')} {self.wrap(cube)}" if cube else ""
+
+        rollup = expression.args.get("rollup")
+        if rollup is True:
+            rollup = self.seg("WITH ROLLUP")
+        else:
+            rollup = self.expressions(expression, key="rollup", indent=False)
+            rollup = f"{self.seg('ROLLUP')} {self.wrap(rollup)}" if rollup else ""
+
         return f"{group_by}{grouping_sets}{cube}{rollup}"
 
     def having_sql(self, expression: exp.Having) -> str:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index f7af5a86..6df1659d 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1864,12 +1864,19 @@ class Parser(metaclass=_Parser):
         if not skip_group_by_token and not self._match(TokenType.GROUP_BY):
             return None
 
+        expressions = self._parse_csv(self._parse_conjunction)
+        grouping_sets = self._parse_grouping_sets()
+
+        with_ = self._match(TokenType.WITH)
+        cube = self._match(TokenType.CUBE) and (with_ or self._parse_wrapped_id_vars())
+        rollup = self._match(TokenType.ROLLUP) and (with_ or self._parse_wrapped_id_vars())
+
         return self.expression(
             exp.Group,
-            expressions=self._parse_csv(self._parse_conjunction),
-            grouping_sets=self._parse_grouping_sets(),
-            cube=self._match(TokenType.CUBE) and self._parse_wrapped_id_vars(),
-            rollup=self._match(TokenType.ROLLUP) and self._parse_wrapped_id_vars(),
+            expressions=expressions,
+            grouping_sets=grouping_sets,
+            cube=cube,
+            rollup=rollup,
         )
 
     def _parse_grouping_sets(self) -> t.Optional[t.List[t.Optional[exp.Expression]]]:


Test patch:
diff --git a/tests/dialects/test_hive.py b/tests/dialects/test_hive.py
index 893cac9e..94479af7 100644
--- a/tests/dialects/test_hive.py
+++ b/tests/dialects/test_hive.py
@@ -353,6 +353,12 @@ class TestHive(Validator):
         self.validate_identity(
             "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, FOO(b) GROUPING SETS ((a, FOO(b)), a)"
         )
+        self.validate_identity(
+            "SELECT key, value, GROUPING__ID, COUNT(*) FROM T1 GROUP BY key, value WITH CUBE"
+        )
+        self.validate_identity(
+            "SELECT key, value, GROUPING__ID, COUNT(*) FROM T1 GROUP BY key, value WITH ROLLUP"
+        )
         self.validate_all(
             "SELECT A.1a AS b FROM test_a AS A",
             write={