fix: compilation of a labeled custom FunctionElement when used in grouping (#1155)

* fix: fix grouped labels using custom functions

If we have a labeled custom function that we are grouping by, and the
function does not support the `default` dialect, we can not compile our
query.

* Update tests/unit/test_compiler.py

---------

Co-authored-by: Tim Sweña (Swast) <[email protected]>

Author: snapiri

sqlalchemy_bigquery/base.py

@@ -343,10 +343,8 @@ def visit_label(self, *args, within_group_by=False, **kwargs):
         if within_group_by:
             column_label = args[0]
             sql_keywords = {"GROUPING SETS", "ROLLUP", "CUBE"}
-            for keyword in sql_keywords:
-                if keyword in str(column_label):
-                    break
-            else:  # for/else always happens unless break gets called
+            label_str = column_label.compile(dialect=self.dialect).string
+            if not any(keyword in label_str for keyword in sql_keywords):
                 kwargs["render_label_as_label"] = column_label
 
         return super(BigQueryCompiler, self).visit_label(*args, **kwargs)

tests/unit/test_compiler.py

@@ -417,3 +417,42 @@ def test_complex_grouping_ops_vs_nested_grouping_ops(
     )
 
     assert found_sql == expected_sql
+
+
+def test_label_compiler(faux_conn, metadata):
+    class CustomLower(sqlalchemy.sql.functions.FunctionElement):
+        name = "custom_lower"
+
+    @sqlalchemy.ext.compiler.compiles(CustomLower)
+    def compile_custom_intersect(element, compiler, **kwargs):
+        if compiler.dialect.name != "bigquery":
+            # We only test with the BigQuery dialect, so this should never happen.
+            raise sqlalchemy.exc.CompileError(  # pragma: NO COVER
+                f"custom_lower is not supported for dialect {compiler.dialect.name}"
+            )
+
+        clauses = list(element.clauses)
+        field = compiler.process(clauses[0], **kwargs)
+        return f"LOWER({field})"
+
+    table1 = setup_table(
+        faux_conn,
+        "table1",
+        metadata,
+        sqlalchemy.Column("foo", sqlalchemy.String),
+        sqlalchemy.Column("bar", sqlalchemy.Integer),
+    )
+
+    lower_foo = CustomLower(table1.c.foo).label("some_label")
+    q = (
+        sqlalchemy.select(lower_foo, sqlalchemy.func.max(table1.c.bar))
+        .select_from(table1)
+        .group_by(lower_foo)
+    )
+    expected_sql = (
+        "SELECT LOWER(`table1`.`foo`) AS `some_label`, max(`table1`.`bar`) AS `max_1` \n"
+        "FROM `table1` GROUP BY `some_label`"
+    )
+
+    found_sql = q.compile(faux_conn).string
+    assert found_sql == expected_sql