diff --git a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/geography.py b/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/geography.py index 2c1b20d559aa..c7fb46fd72bc 100644 --- a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/geography.py +++ b/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/geography.py @@ -180,6 +180,17 @@ def _fixup_st_arguments(element, compiler, **kw): This is not a user-facing function. """ + # GeoAlchemy2 >= 0.21.0 wraps raw string arguments in an _UntypedString(NullType) + # bindparam for PostgreSQL compatibility. Restore String() for string bindparams + # so BigQueryCompiler.visit_bindparam emits typed %(...:STRING)s placeholders. + for argument in element.clauses.clauses: + if ( + isinstance(argument, BindParameter) + and isinstance(argument.type, sqlalchemy.types.NullType) + and isinstance(argument.value, str) + ): + argument.type = sqlalchemy.types.String() + argument_types = _argument_types.get(element.name.lower()) if argument_types: for argument_type, argument in zip(argument_types, element.clauses.clauses): @@ -191,6 +202,22 @@ def _fixup_st_arguments(element, compiler, **kw): return compiler.visit_function(element, **kw) +def _get_all_subclasses(cls): + subclasses = set(cls.__subclasses__()) + for sub in list(subclasses): + subclasses.update(_get_all_subclasses(sub)) + return subclasses + + +# GeoAlchemy2 registers dialect-specific @compiles handlers (e.g. for "mssql") +# on several GenericFunction subclasses such as ST_GeogFromText, which gives +# those subclasses their own _compiler_dispatcher that shadows GenericFunction's. +# Register _fixup_st_arguments on each subclass so the "bigquery" compiler +# hook is invoked for all spatial functions. +for _cls in _get_all_subclasses(geoalchemy2.functions.GenericFunction): + sqlalchemy.ext.compiler.compiles(_cls, "bigquery")(_fixup_st_arguments) + + _argument_types = dict( st_area=(GEOGRAPHY,), st_asbinary=(GEOGRAPHY,), diff --git a/packages/sqlalchemy-bigquery/tests/unit/test_geography.py b/packages/sqlalchemy-bigquery/tests/unit/test_geography.py index 332cb1c3a954..20d95ab07f23 100644 --- a/packages/sqlalchemy-bigquery/tests/unit/test_geography.py +++ b/packages/sqlalchemy-bigquery/tests/unit/test_geography.py @@ -184,6 +184,7 @@ def test_calling_st_functions_that_dont_take_geographies(faux_conn, last_query): def test_fixup_st_arguments(): + import sqlalchemy.types from geoalchemy2.functions import GenericFunction, ST_Area from sqlalchemy.sql.elements import BindParameter @@ -203,9 +204,46 @@ def visit_function(self, element, **kw): func_element2 = ST_Area(BindParameter("param", "point(0 0)", type_=GEOGRAPHY())) _fixup_st_arguments(func_element2, DummyCompiler()) - # Case 3: function without specified argument types + # Case 3: function without specified argument types coerces NullType str to String class ST_Unknown(GenericFunction): name = "ST_Unknown" - func_element3 = ST_Unknown(BindParameter("param", "point(0 0)")) + func_element3 = ST_Unknown( + BindParameter("param", "point(0 0)", type_=sqlalchemy.types.NullType()) + ) _fixup_st_arguments(func_element3, DummyCompiler()) + assert isinstance(func_element3.clauses.clauses[0].type, sqlalchemy.types.String) + + # Case 4: NullType with non-str value (e.g. None) remains NullType + func_element4 = ST_Unknown( + BindParameter("param", None, type_=sqlalchemy.types.NullType()) + ) + _fixup_st_arguments(func_element4, DummyCompiler()) + assert isinstance(func_element4.clauses.clauses[0].type, sqlalchemy.types.NullType) + + +def test_nested_generic_function_subclasses_register_compiler(): + """Verify that multi-tier GenericFunction subclasses register the bigquery compiler handler.""" + from geoalchemy2.functions import GenericFunction + + from sqlalchemy_bigquery.geography import _get_all_subclasses + + # Intermediate base class inheriting directly from GenericFunction + class ST_BaseFunction(GenericFunction): + name = "ST_BaseFunction" + + # Multi-tier / nested subclass inheriting through the intermediate base + class ST_NestedFunction(ST_BaseFunction): + name = "ST_NestedFunction" + + # Step 1: Ensure the recursive discovery function resolves transitive descendants + discovered_classes = _get_all_subclasses(GenericFunction) + assert ST_NestedFunction in discovered_classes, ( + "ST_NestedFunction must be discovered by recursive subclass collection." + ) + + # Step 2: Ensure the BigQuery compiler dispatcher hook is registered on the nested class + dispatcher = getattr(ST_NestedFunction, "_compiler_dispatcher", None) + assert dispatcher is not None and "bigquery" in dispatcher.specs, ( + "Nested GenericFunction subclasses must have the 'bigquery' compiler handler registered." + )