diff --git a/sqlalchemy_datastore/base.py b/sqlalchemy_datastore/base.py index 90d4368..be9d5b8 100644 --- a/sqlalchemy_datastore/base.py +++ b/sqlalchemy_datastore/base.py @@ -20,6 +20,7 @@ from typing import Any, List, Optional from google.cloud import firestore_admin_v1 +from google.cloud.datastore.query import PropertyFilter from google.oauth2 import service_account from sqlalchemy import exc from sqlalchemy.engine import Connection, default @@ -54,6 +55,7 @@ class CloudDatastoreDialect(default.DefaultDialect): supports_unicode_binds = True returns_unicode_strings = True description_encoding = None + supports_statement_cache = False # JSON support - required for SQLAlchemy JSON type _json_serializer = None @@ -88,10 +90,15 @@ def __init__( self._client = None @classmethod - def dbapi(cls): + def import_dbapi(cls): """Return the DBAPI 2.0 driver.""" return datastore_dbapi + @classmethod + def dbapi(cls): + """Return the DBAPI 2.0 driver.""" + return cls.import_dbapi() + def do_ping(self, dbapi_connection): """Performs a simple operation to check if the connection is still alive.""" try: @@ -212,7 +219,7 @@ def get_columns( """Retrieve column information from the database.""" client = self._client query = client.query(kind="__Stat_PropertyType_PropertyName_Kind__") - query.add_filter("kind_name", "=", table_name) + query.add_filter(filter=PropertyFilter("kind_name", "=", table_name)) properties = list(query.fetch()) return [ diff --git a/tests/test_integration_queries.py b/tests/test_integration_queries.py index ef19d0b..66641e8 100644 --- a/tests/test_integration_queries.py +++ b/tests/test_integration_queries.py @@ -33,8 +33,10 @@ def test_update_user(conn, datastore_client): # Find the inserted entity by querying the datastore client directly query = datastore_client.query(kind="users") - query.add_filter("name", "=", "UpdateTestUser") - entities = list(query.fetch()) + entities = [ + entity for entity in query.fetch() + if entity.get("name") == "UpdateTestUser" + ] assert len(entities) >= 1 entity_id = entities[0].key.id @@ -62,8 +64,10 @@ def test_delete_user(conn, datastore_client): # Find the inserted entity query = datastore_client.query(kind="users") - query.add_filter("name", "=", "DeleteTestUser") - entities = list(query.fetch()) + entities = [ + entity for entity in query.fetch() + if entity.get("name") == "DeleteTestUser" + ] assert len(entities) >= 1 entity_id = entities[0].key.id @@ -332,8 +336,10 @@ def test_task_insert_update_delete_raw_sql(conn, datastore_client): # Find the inserted entity query = datastore_client.query(kind="tasks") - query.add_filter("task", "=", "Coverage Test Task") - entities = list(query.fetch()) + entities = [ + entity for entity in query.fetch() + if entity.get("task") == "Coverage Test Task" + ] assert len(entities) >= 1 entity_id = entities[0].key.id diff --git a/tests/test_unit_types.py b/tests/test_unit_types.py index 09cdd17..c07f707 100644 --- a/tests/test_unit_types.py +++ b/tests/test_unit_types.py @@ -17,6 +17,7 @@ # 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. """Unit tests for _types module (no emulator required).""" +import pytest import sqlalchemy.types from google.cloud.bigquery.schema import SchemaField @@ -190,7 +191,13 @@ def test_get_sqla_column_type_repeated(): def test_get_sqla_column_type_unknown(): field = SchemaField("mystery", "UNKNOWN_TYPE_XYZ") - coltype = _get_sqla_column_type(field) + + with pytest.warns( + sqlalchemy.exc.SAWarning, + match="Did not recognize type 'UNKNOWN_TYPE_XYZ' of column 'mystery'", + ): + coltype = _get_sqla_column_type(field) + assert coltype is sqlalchemy.types.NullType