feat(datastore): support "IN" query filter (#760)

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',

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'
 
 

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

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(