added test case

Author: kiraksi

sqlalchemy_bigquery/base.py

@@ -343,13 +343,12 @@ def visit_label(self, *args, within_group_by=False, **kwargs):
                 for keyword in ("GROUPING SETS", "ROLLUP", "CUBE")
             ):
                 kwargs["render_label_as_label"] = args[0]
+
         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,
+            select, **kw, within_group_by=True
         )
 
     ############################################################################

tests/unit/test_compiler.py

@@ -26,6 +26,7 @@
     sqlalchemy_before_2_0,
 )
 from sqlalchemy.sql.functions import rollup, cube, grouping_sets
+from sqlalchemy import func
 
 
 def test_constraints_are_ignored(faux_conn, metadata):
@@ -343,3 +344,24 @@ def test_cube(faux_conn, metadata):
     )
     found_sql = q.compile(faux_conn).string
     assert found_sql == expected_sql
+
+
+def test_multiple_grouping_sets(faux_conn, metadata):
+    table = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.ARRAY(sqlalchemy.Integer)),
+    )
+
+    q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
+        grouping_sets(table.c.foo, table.c.bar), grouping_sets(table.c.foo)
+    )
+
+    expected_sql = (
+        "SELECT `table1`.`foo`, `table1`.`bar` \n"
+        "FROM `table1` GROUP BY GROUPING SETS(`table1`.`foo`, `table1`.`bar`), GROUPING SETS(`table1`.`foo`)"
+    )
+    found_sql = q.compile(faux_conn).string
+    assert found_sql == expected_sql