feat: grouping sets, rollup and cube compatibility

Author: kiraksi

sqlalchemy_bigquery/base.py

@@ -38,6 +38,7 @@
 import sqlalchemy
 import sqlalchemy.sql.expression
 import sqlalchemy.sql.functions
+from sqlalchemy.sql.functions import rollup, cube, grouping_sets
 import sqlalchemy.sql.sqltypes
 import sqlalchemy.sql.type_api
 from sqlalchemy.exc import NoSuchTableError, NoSuchColumnError
@@ -340,9 +341,36 @@ def visit_label(self, *args, within_group_by=False, **kwargs):
         return super(BigQueryCompiler, self).visit_label(*args, **kwargs)
 
     def group_by_clause(self, select, **kw):
-        return super(BigQueryCompiler, self).group_by_clause(
-            select, **kw, within_group_by=True
-        )
+        grouping_sets_exprs = []
+        rollup_exprs = []
+        cube_exprs = []
+
+        # Traverse select statement to extract grouping sets, rollup, and cube expressions
+        for expr in select._group_by_clause:
+            if isinstance(expr, grouping_sets):
+                grouping_sets_exprs.append(
+                    self.process(expr.clauses)
+                )  # Assuming SQLAlchemy syntax
+            elif isinstance(expr, rollup):  # Assuming SQLAlchemy syntax
+                rollup_exprs.append(self.process(expr.clauses))
+            elif isinstance(expr, cube):  # Assuming SQLAlchemy syntax
+                cube_exprs.append(self.process(expr.clauses))
+            else:
+                # Handle regular group by expressions
+                pass
+
+        clause = super(BigQueryCompiler, self).group_by_clause(select, **kw)
+
+        if grouping_sets_exprs:
+            clause = (
+                f"GROUP BY {clause} GROUPING SETS ({', '.join(grouping_sets_exprs)})"
+            )
+        if rollup_exprs:
+            clause = f"GROUP BY {clause} ROLLUP ({', '.join(rollup_exprs)})"
+        if cube_exprs:
+            clause = f"GROUP BY {clause} CUBE ({', '.join(cube_exprs)})"
+
+        return clause
 
     ############################################################################
     # Handle parameters in in

tests/unit/test_compiler.py

@@ -22,6 +22,7 @@
 
 from .conftest import setup_table
 from .conftest import sqlalchemy_1_4_or_higher, sqlalchemy_before_1_4
+from sqlalchemy.sql.functions import rollup, cube, grouping_sets
 
 
 def test_constraints_are_ignored(faux_conn, metadata):
@@ -278,3 +279,66 @@ def test_no_implicit_join_for_inner_unnest_no_table2_column(faux_conn, metadata)
     )
     found_outer_sql = q.compile(faux_conn).string
     assert found_outer_sql == expected_outer_sql
+
+
+def test_grouping_sets(faux_conn, metadata):
+    table = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+
+    q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
+        grouping_sets(table.c.foo, table.c.bar)
+    )
+
+    expected_sql = (
+        "SELECT `table1`.`foo`, `table1`.`bar` \n"
+        "FROM `table1` GROUP BY GROUPING SETS ((`table1`.`foo`), (`table1`.`bar`))"
+    )
+    found_sql = q.compile(faux_conn).string
+    assert found_sql == expected_sql
+
+
+def test_rollup(faux_conn, metadata):
+    table = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+
+    q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
+        rollup(table.c.foo, table.c.bar)
+    )
+
+    expected_sql = (
+        "SELECT `table1`.`foo`, `table1`.`bar` \n"
+        "FROM `table1` GROUP BY ROLLUP(`table1`.`foo`, `table1`.`bar`)"
+    )
+    found_sql = q.compile(faux_conn).string
+    assert found_sql == expected_sql
+
+
+def test_cube(faux_conn, metadata):
+    table = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+
+    q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
+        cube(table.c.foo, table.c.bar)
+    )
+
+    expected_sql = (
+        "SELECT `table1`.`foo`, `table1`.`bar` \n"
+        "FROM `table1` GROUP BY CUBE(`table1`.`foo`, `table1`.`bar`)"
+    )
+    found_sql = q.compile(faux_conn).string
+    assert found_sql == expected_sql