fix(sqlalchemy-bigquery): wrap GEOGRAPHY bind arguments with ST_GeogFromText in _fixup_st_arguments

Author: ohmayr

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/__init__.py

@@ -20,6 +20,7 @@
 SQLAlchemy dialect for Google BigQuery
 """
 
+import sys
 import warnings
 
 from ._types import (
@@ -43,7 +44,6 @@
 )
 from .base import BigQueryDialect, dialect
 from .version import __version__
-import sys
 
 # Now that support for Python 3.7, 3.8 and 3.9 has been removed, we don't expect the
 # following check to succeed. The warning is only included for robustness.

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/_helpers.py

@@ -10,11 +10,11 @@
 import re
 from typing import Optional
 
-from google.api_core import client_info
 import google.auth
+import sqlalchemy
+from google.api_core import client_info
 from google.cloud import bigquery
 from google.oauth2 import service_account
-import sqlalchemy
 
 USER_AGENT_TEMPLATE = "sqlalchemy/{}"
 SCOPES = (

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/_struct.py

@@ -17,12 +17,12 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-from sqlalchemy.sql import operators
 import sqlalchemy.sql.coercions
 import sqlalchemy.sql.default_comparator
 import sqlalchemy.sql.roles
 import sqlalchemy.sql.sqltypes
 import sqlalchemy.types
+from sqlalchemy.sql import operators
 
 # from . import base  # Moved to _get_subtype_col_spec to break circular import
 

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/_types.py

@@ -17,9 +17,9 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-from google.cloud.bigquery.schema import SchemaField
 import sqlalchemy.types
 import sqlalchemy.util
+from google.cloud.bigquery.schema import SchemaField
 
 try:
     from .geography import GEOGRAPHY

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py

@@ -20,23 +20,28 @@
 """Integration between SQLAlchemy and BigQuery."""
 
 import datetime
-from decimal import Decimal
 import operator
 import random
 import re
 import uuid
+from decimal import Decimal
 
-from google import auth
 import google.api_core.exceptions
+import packaging.version
+import sqlalchemy
+import sqlalchemy.sql.expression
+import sqlalchemy.sql.functions
+import sqlalchemy.sql.sqltypes
+import sqlalchemy.sql.type_api
+import sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base as vendored_postgresql
+from google import auth
 from google.api_core.exceptions import NotFound
 from google.cloud.bigquery import ConnectionProperty, QueryJobConfig, dbapi
 from google.cloud.bigquery.table import (
     RangePartitioning,
     TableReference,
     TimePartitioning,
 )
-import packaging.version
-import sqlalchemy
 from sqlalchemy import util
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
@@ -49,14 +54,9 @@
     IdentifierPreparer,
     SQLCompiler,
 )
-import sqlalchemy.sql.expression
-import sqlalchemy.sql.functions
 from sqlalchemy.sql.schema import Column, Table
 from sqlalchemy.sql.selectable import CTE
-import sqlalchemy.sql.sqltypes
 from sqlalchemy.sql.sqltypes import Integer, NullType, Numeric, String
-import sqlalchemy.sql.type_api
-import sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base as vendored_postgresql
 
 from . import _helpers, _types
 from .parse_url import parse_url

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/geography.py

@@ -19,9 +19,9 @@
 
 import geoalchemy2
 import geoalchemy2.functions
+import sqlalchemy.ext.compiler
 from geoalchemy2.shape import to_shape
 from shapely import wkb, wkt
-import sqlalchemy.ext.compiler
 from sqlalchemy.sql.elements import BindParameter
 
 SRID = 4326  # WGS84, https://spatialreference.org/ref/epsg/wgs-84/
@@ -182,11 +182,19 @@ def _fixup_st_arguments(element, compiler, **kw):
     """
     argument_types = _argument_types.get(element.name.lower())
     if argument_types:
+        new_clauses = []
+        modified = False
         for argument_type, argument in zip(argument_types, element.clauses.clauses):
             if isinstance(argument, BindParameter) and not isinstance(
                 argument.type, argument_type
             ):
                 argument.type = argument_type()
+                if issubclass(argument_type, GEOGRAPHY):
+                    argument = geoalchemy2.functions.ST_GeogFromText(argument)
+                    modified = True
+            new_clauses.append(argument)
+        if modified:
+            element.clauses.clauses = new_clauses
 
     return compiler.visit_function(element, **kw)
 

packages/sqlalchemy-bigquery/tests/sqlalchemy_dialect_compliance/conftest.py

@@ -22,6 +22,7 @@
 import traceback
 
 import google.cloud.bigquery.dbapi.connection
+import test_utils.prefixer
 from sqlalchemy.testing import config
 from sqlalchemy.testing.plugin.pytestplugin import *  # noqa
 from sqlalchemy.testing.plugin.pytestplugin import (
@@ -30,7 +31,6 @@
 from sqlalchemy.testing.plugin.pytestplugin import (
     pytest_sessionstart as _pytest_sessionstart,
 )
-import test_utils.prefixer
 
 import sqlalchemy_bigquery.base
 

packages/sqlalchemy-bigquery/tests/sqlalchemy_dialect_compliance/test_dialect_compliance.py

@@ -25,33 +25,34 @@
 import pytest
 import pytz
 import sqlalchemy
-from sqlalchemy import and_
 import sqlalchemy.sql.sqltypes
+import sqlalchemy.testing.suite.test_types
+from sqlalchemy import and_
 from sqlalchemy.testing import config, util
 from sqlalchemy.testing.assertions import eq_
+from sqlalchemy.testing.suite import *  # noqa
 from sqlalchemy.testing.suite import (
+    Column,
+    DistinctOnTest,
+    HasIndexTest,
+    IdentityAutoincrementTest,
     Integer,
     LongNameBlowoutTest,
     PostCompileParamsTest,
     QuotedNameArgumentTest,
-)
-from sqlalchemy.testing.suite import (
+    String,
+    Table,
     WindowFunctionTest,
     bindparam,
     exists,
     select,
     testing,
 )
-from sqlalchemy.testing.suite import *  # noqa
 from sqlalchemy.testing.suite import CTETest as _CTETest
-from sqlalchemy.testing.suite import Column
-from sqlalchemy.testing.suite import DistinctOnTest
-from sqlalchemy.testing.suite import ExistsTest as _ExistsTest
-from sqlalchemy.testing.suite import HasIndexTest, IdentityAutoincrementTest
-from sqlalchemy.testing.suite import InsertBehaviorTest as _InsertBehaviorTest
-from sqlalchemy.testing.suite import String, Table
 from sqlalchemy.testing.suite import DifficultParametersTest as _DifficultParametersTest
+from sqlalchemy.testing.suite import ExistsTest as _ExistsTest
 from sqlalchemy.testing.suite import FetchLimitOffsetTest as _FetchLimitOffsetTest
+from sqlalchemy.testing.suite import InsertBehaviorTest as _InsertBehaviorTest
 from sqlalchemy.testing.suite import SimpleUpdateDeleteTest as _SimpleUpdateDeleteTest
 from sqlalchemy.testing.suite import (
     TimestampMicrosecondsTest as _TimestampMicrosecondsTest,
@@ -62,7 +63,6 @@
     ComponentReflectionTestExtra,
     HasTableTest,
 )
-import sqlalchemy.testing.suite.test_types
 from sqlalchemy.testing.suite.test_types import ArrayTest
 
 if packaging.version.parse(sqlalchemy.__version__) >= packaging.version.parse("2.0"):
@@ -637,8 +637,12 @@ def test_no_results_for_non_returning_insert(cls):
 del DistinctOnTest  # expects unquoted table names.
 del HasIndexTest  # BQ doesn't do the indexes that SQLA is loooking for.
 del IdentityAutoincrementTest  # BQ doesn't do autoincrement
-del LongNameBlowoutTest  # Requires features (indexes, primary keys, etc., that BigQuery doesn't have.
-del PostCompileParamsTest  # BQ adds backticks to bind parameters, causing failure of tests TODO: fix this?
+del (
+    LongNameBlowoutTest
+)  # Requires features (indexes, primary keys, etc., that BigQuery doesn't have.
+del (
+    PostCompileParamsTest
+)  # BQ adds backticks to bind parameters, causing failure of tests TODO: fix this?
 del QuotedNameArgumentTest  # Quotes aren't allowed in BigQuery table names.
 del (
     WindowFunctionTest.test_window_rows_between

packages/sqlalchemy-bigquery/tests/system/conftest.py

@@ -20,12 +20,12 @@
 import pathlib
 from typing import List
 
-from google.api_core import exceptions
-from google.cloud import bigquery
 import pytest
 import sqlalchemy
 import test_utils.prefixer
 import test_utils.retry
+from google.api_core import exceptions
+from google.cloud import bigquery
 
 from sqlalchemy_bigquery import BigQueryDialect
 

packages/sqlalchemy-bigquery/tests/system/test_alembic.py

@@ -20,8 +20,8 @@
 import contextlib
 
 import google.api_core.exceptions
-from google.cloud.bigquery import SchemaField, TimePartitioning
 import pytest
+from google.cloud.bigquery import SchemaField, TimePartitioning
 from sqlalchemy import Column, DateTime, Integer, Numeric, String
 
 alembic = pytest.importorskip("alembic")