diff --git a/datastore/gcloud/aio/datastore/__init__.py b/datastore/gcloud/aio/datastore/__init__.py index 14cf87ba6..fe7e558ae 100644 --- a/datastore/gcloud/aio/datastore/__init__.py +++ b/datastore/gcloud/aio/datastore/__init__.py @@ -181,6 +181,7 @@ class MyGQLQuery(gcloud.aio.datastore.GQLQuery): """ import importlib.metadata +from .array import Array from .constants import CompositeFilterOperator from .constants import Consistency from .constants import Direction @@ -215,6 +216,7 @@ class MyGQLQuery(gcloud.aio.datastore.GQLQuery): __version__ = importlib.metadata.version('gcloud-aio-datastore') __all__ = [ + 'Array', 'CompositeFilter', 'CompositeFilterOperator', 'Consistency', diff --git a/datastore/gcloud/aio/datastore/constants.py b/datastore/gcloud/aio/datastore/constants.py index 1b2aeb4a9..729a8201a 100644 --- a/datastore/gcloud/aio/datastore/constants.py +++ b/datastore/gcloud/aio/datastore/constants.py @@ -44,14 +44,15 @@ class Operation(enum.Enum): class PropertyFilterOperator(enum.Enum): - # TODO: support IN / NOT_IN (requires rhs to be ArrayValue) EQUAL = 'EQUAL' GREATER_THAN = 'GREATER_THAN' GREATER_THAN_OR_EQUAL = 'GREATER_THAN_OR_EQUAL' HAS_ANCESTOR = 'HAS_ANCESTOR' + IN = 'IN' LESS_THAN = 'LESS_THAN' LESS_THAN_OR_EQUAL = 'LESS_THAN_OR_EQUAL' NOT_EQUAL = 'NOT_EQUAL' + NOT_IN = 'NOT_IN' UNSPECIFIED = 'OPERATOR_UNSPECIFIED' diff --git a/datastore/gcloud/aio/datastore/filter.py b/datastore/gcloud/aio/datastore/filter.py index ec66dc1ea..af2819847 100644 --- a/datastore/gcloud/aio/datastore/filter.py +++ b/datastore/gcloud/aio/datastore/filter.py @@ -1,7 +1,9 @@ from typing import Any from typing import Dict from typing import List +from typing import Union +from .array import Array from .constants import CompositeFilterOperator from .constants import PropertyFilterOperator from .value import Value @@ -89,7 +91,7 @@ class PropertyFilter(BaseFilter): def __init__( self, prop: str, operator: PropertyFilterOperator, - value: Value, + value: Union[Value, Array], ) -> None: self.prop = prop self.operator = operator @@ -113,8 +115,13 @@ def from_repr(cls, data: Dict[str, Any]) -> 'PropertyFilter': return cls(prop=prop, operator=operator, value=value) def to_repr(self) -> Dict[str, Any]: - return { + rep: Dict[str, Any] = { 'op': self.operator.value, 'property': {'name': self.prop}, - 'value': self.value.to_repr(), } + # TODO: consider refactoring to look more like Value.to_repr() + if isinstance(self.value, Array): + rep['value'] = {'arrayValue': self.value.to_repr()} + else: + rep['value'] = self.value.to_repr() + return rep diff --git a/datastore/tests/integration/smoke_test.py b/datastore/tests/integration/smoke_test.py index ce8e2ddc8..959c3441e 100644 --- a/datastore/tests/integration/smoke_test.py +++ b/datastore/tests/integration/smoke_test.py @@ -2,6 +2,7 @@ import pytest from gcloud.aio.auth import BUILD_GCLOUD_REST # pylint: disable=no-name-in-module +from gcloud.aio.datastore import Array from gcloud.aio.datastore import Datastore from gcloud.aio.datastore import Filter from gcloud.aio.datastore import GQLCursor @@ -272,6 +273,86 @@ async def test_query(creds: str, kind: str, project: str) -> None: assert len(after.entity_results) == num_results + 2 +@pytest.mark.asyncio +@pytest.mark.xfail(strict=False) +async def test_query_with_in_filter( + creds: str, kind: str, project: str) -> None: + async with Session() as s: + ds = Datastore(project=project, service_file=creds, session=s) + + property_filter = PropertyFilter( + prop='value', operator=PropertyFilterOperator.IN, + value=Array([Value(99), Value(100)]), + ) + query = Query(kind=kind, query_filter=Filter(property_filter)) + + before = await ds.runQuery(query, session=s) + num_results = len(before.entity_results) + + transaction = await ds.beginTransaction(session=s) + mutations = [ + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 99}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 100}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 101}, + ), + ] + await ds.commit(mutations, transaction=transaction, session=s) + + after = await ds.runQuery(query, session=s) + assert len(after.entity_results) == num_results + 2 + + +@pytest.mark.asyncio +@pytest.mark.xfail(strict=False) +async def test_query_with_not_in_filter( + creds: str, kind: str, project: str) -> None: + async with Session() as s: + ds = Datastore(project=project, service_file=creds, session=s) + + property_filter = PropertyFilter( + prop='value', operator=PropertyFilterOperator.NOT_IN, + value=Array([Value(99), Value(100), Value(30), Value(42)]), + ) + query = Query(kind=kind, query_filter=Filter(property_filter)) + + before = await ds.runQuery(query, session=s) + num_results = len(before.entity_results) + + transaction = await ds.beginTransaction(session=s) + mutations = [ + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 99}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 100}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 999}, + ), + ] + await ds.commit(mutations, transaction=transaction, session=s) + + after = await ds.runQuery(query, session=s) + assert len(after.entity_results) == num_results + 1 + + @pytest.mark.asyncio @pytest.mark.xfail(strict=False) async def test_gql_query(creds: str, kind: str, project: str) -> None: @@ -310,6 +391,85 @@ async def test_gql_query(creds: str, kind: str, project: str) -> None: assert len(after.entity_results) == num_results + 3 +@pytest.mark.asyncio +@pytest.mark.xfail(strict=False) +async def test_gql_query_with_in_filter( + creds: str, kind: str, project: str) -> None: + async with Session() as s: + ds = Datastore(project=project, service_file=creds, session=s) + + query = GQLQuery( + f'SELECT * FROM {kind} WHERE value IN @values', + named_bindings={'values': Array([Value(99), Value(100)])}, + ) + + before = await ds.runQuery(query, session=s) + num_results = len(before.entity_results) + + transaction = await ds.beginTransaction(session=s) + mutations = [ + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 99}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 100}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 101}, + ), + ] + await ds.commit(mutations, transaction=transaction, session=s) + + after = await ds.runQuery(query, session=s) + assert len(after.entity_results) == num_results + 2 + + +@pytest.mark.asyncio +@pytest.mark.xfail(strict=False) +async def test_gql_query_with_not_in_filter( + creds: str, kind: str, project: str) -> None: + async with Session() as s: + ds = Datastore(project=project, service_file=creds, session=s) + + query = GQLQuery( + f'SELECT * FROM {kind} WHERE value NOT IN @values', + named_bindings={'values': Array( + [Value(30), Value(42), Value(99), Value(100)])}, + ) + + before = await ds.runQuery(query, session=s) + num_results = len(before.entity_results) + + transaction = await ds.beginTransaction(session=s) + mutations = [ + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 99}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 100}, + ), + ds.make_mutation( + Operation.INSERT, + Key(project, [PathElement(kind)]), + properties={'value': 999}, + ), + ] + await ds.commit(mutations, transaction=transaction, session=s) + + after = await ds.runQuery(query, session=s) + assert len(after.entity_results) == num_results + 1 + + @pytest.mark.asyncio @pytest.mark.xfail(strict=False) async def test_gql_query_pagination(