Implement filters for queries. (#55)

Implement filters for queries. Does not include post-filter functionality.

Author: Chris Rossi

packages/google-cloud-ndb/src/google/cloud/ndb/_datastore_query.py

@@ -14,11 +14,13 @@
 
 """Translate NDB queries to Datastore calls."""
 
+import itertools
 import logging
 
 from google.cloud.datastore_v1.proto import datastore_pb2
 from google.cloud.datastore_v1.proto import entity_pb2
 from google.cloud.datastore_v1.proto import query_pb2
+from google.cloud.datastore import helpers
 
 from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _datastore_api
@@ -34,6 +36,50 @@
 RESULT_TYPE_FULL = ResultType.Value("FULL")
 RESULT_TYPE_PROJECTION = ResultType.Value("PROJECTION")
 
+FILTER_OPERATORS = {
+    "=": query_pb2.PropertyFilter.EQUAL,
+    "<": query_pb2.PropertyFilter.LESS_THAN,
+    "<=": query_pb2.PropertyFilter.LESS_THAN_OR_EQUAL,
+    ">": query_pb2.PropertyFilter.GREATER_THAN,
+    ">=": query_pb2.PropertyFilter.GREATER_THAN_OR_EQUAL,
+}
+
+
+def make_filter(name, op, value):
+    """Make a property filter protocol buffer.
+
+    Args:
+        name (str): The name of the property to filter by.
+        op (str): The operator to apply in the filter. Must be one of "=", "<",
+            "<=", ">", or ">=".
+        value (Any): The value for comparison.
+
+    Returns:
+        query_pb2.PropertyFilter: The filter protocol buffer.
+    """
+    filter_pb = query_pb2.PropertyFilter(
+        property=query_pb2.PropertyReference(name=name),
+        op=FILTER_OPERATORS[op],
+    )
+    helpers._set_protobuf_value(filter_pb.value, value)
+    return filter_pb
+
+
+def make_composite_and_filter(filter_pbs):
+    """Make a composite filter protocol buffer using AND.
+
+    Args:
+        List[Union[query_pb2.PropertyFilter, query_pb2.CompositeFilter]]: The
+            list of filters to be combined.
+
+    Returns:
+        query_pb2.CompositeFilter: The new composite filter.
+    """
+    return query_pb2.CompositeFilter(
+        op=query_pb2.CompositeFilter.AND,
+        filters=[_filter_pb(filter_pb) for filter_pb in filter_pbs],
+    )
+
 
 @tasklets.tasklet
 def fetch(query):
@@ -45,7 +91,7 @@ def fetch(query):
     Returns:
         tasklets.Future: Result is List[model.Model]: The query results.
     """
-    for name in ("filters", "orders", "default_options"):
+    for name in ("orders", "default_options"):
         if getattr(query, name, None):
             raise NotImplementedError(
                 "{} is not yet implemented for queries.".format(name)
@@ -61,14 +107,59 @@ def fetch(query):
     if not namespace:
         namespace = client.namespace
 
-    query_pb = _query_to_protobuf(query)
-    results = yield _run_query(project_id, namespace, query_pb)
+    filter_pbs = (None,)
+    if query.filters:
+        filter_pbs = query.filters._to_filter()
+        if not isinstance(filter_pbs, (tuple, list)):
+            filter_pbs = (filter_pbs,)
+
+    queries = [
+        _run_query(project_id, namespace, _query_to_protobuf(query, filter_pb))
+        for filter_pb in filter_pbs
+    ]
+    results = yield queries
+
+    if len(results) > 1:
+        results = _merge_results(results)
+    else:
+        results = results[0]
+
     return [
         _process_result(result_type, result, query.projection)
         for result_type, result in results
     ]
 
 
+def _merge_results(results):
+    """Merge the results of distinct queries.
+
+    Some queries that in NDB are logically a single query have to be broken
+    up into two or more Datastore queries, because Datastore doesn't have a
+    composite filter with a boolean OR. The `results` are the result sets from
+    two or more queries which logically form a composite query joined by OR.
+    The individual result sets are combined into a single result set,
+    consolidating any results which may be common to two or more result sets.
+
+    Args:
+        results (List[Tuple[query_pb2.EntityResult.ResultType,
+            query_pb2.EntityResult]]): List of individual result sets as
+            returned by :func:`_run_query`. These are merged into the final
+            result.
+
+    Returns:
+        List[Tuple[query_pb2.EntityResult.ResultType,
+            query_pb2.EntityResult]]: The merged result set.
+    """
+    seen_keys = set()
+    for result_type, result in itertools.chain(*results):
+        hash_key = result.entity.key.SerializeToString()
+        if hash_key in seen_keys:
+            continue
+
+        seen_keys.add(hash_key)
+        yield result_type, result
+
+
 def _process_result(result_type, result, projection):
     """Process a single entity result.
 
@@ -98,11 +189,13 @@ def _process_result(result_type, result, projection):
     )
 
 
-def _query_to_protobuf(query):
+def _query_to_protobuf(query, filter_pb=None):
     """Convert an NDB query to a Datastore protocol buffer.
 
     Args:
         query (query.Query): The query.
+        filter_pb (Optional[query_pb2.Filter]): The filter to apply for this
+        query.
 
     Returns:
         query_pb2.Query: The protocol buffer representation of the query.
@@ -125,22 +218,55 @@ def _query_to_protobuf(query):
             for name in query.distinct_on
         ]
 
-    filters = []
     if query.ancestor:
         ancestor_pb = query.ancestor._key.to_protobuf()
-        filter_pb = query_pb2.PropertyFilter(
+        ancestor_filter_pb = query_pb2.PropertyFilter(
             property=query_pb2.PropertyReference(name="__key__"),
             op=query_pb2.PropertyFilter.HAS_ANCESTOR,
         )
-        filter_pb.value.key_value.CopyFrom(ancestor_pb)
-        filters.append(filter_pb)
+        ancestor_filter_pb.value.key_value.CopyFrom(ancestor_pb)
+
+        if filter_pb is None:
+            filter_pb = ancestor_filter_pb
 
-    if len(filters) == 1:
-        query_args["filter"] = query_pb2.Filter(property_filter=filters[0])
+        elif isinstance(filter_pb, query_pb2.CompositeFilter):
+            filter_pb.filters.add(property_filter=ancestor_filter_pb)
+
+        else:
+            filter_pb = query_pb2.CompositeFilter(
+                op=query_pb2.CompositeFilter.AND,
+                filters=[
+                    _filter_pb(filter_pb),
+                    _filter_pb(ancestor_filter_pb),
+                ],
+            )
+
+    if filter_pb is not None:
+        query_args["filter"] = _filter_pb(filter_pb)
 
     return query_pb2.Query(**query_args)
 
 
+def _filter_pb(filter_pb):
+    """Convenience function to compose a filter protocol buffer.
+
+    The Datastore protocol uses a Filter message which has one of either a
+    PropertyFilter or CompositeFilter as a sole attribute.
+
+    Args:
+        filter_pb (Union[query_pb2.CompositeFilter, query_pb2.PropertyFilter]):
+            The actual filter.
+
+    Returns:
+        query_pb2.Filter: The filter at the higher level of abstraction
+            required to use it in a query.
+    """
+    if isinstance(filter_pb, query_pb2.CompositeFilter):
+        return query_pb2.Filter(composite_filter=filter_pb)
+
+    return query_pb2.Filter(property_filter=filter_pb)
+
+
 @tasklets.tasklet
 def _run_query(project_id, namespace, query_pb):
     """Run a query in Datastore.

packages/google-cloud-ndb/src/google/cloud/ndb/query.py

@@ -279,7 +279,7 @@ def __gt__(self, unused_other):
         raise TypeError("Nodes cannot be ordered")
 
     def _to_filter(self, post=False):
-        """Helper to convert to low-level filter, or :data:`None`.
+        """Helper to convert to low-level filter.
 
         Raises:
             NotImplementedError: Always. This method is virtual.
@@ -409,7 +409,7 @@ def __eq__(self, other):
         )
 
     def _to_filter(self, post=False):
-        """Helper to convert to low-level filter, or :data:`None`.
+        """Helper to convert to low-level filter.
 
         Args:
             post (bool): Indicates if this is a post-filter node.
@@ -540,22 +540,21 @@ def __eq__(self, other):
         )
 
     def _to_filter(self, post=False):
-        """Helper to convert to low-level filter, or :data:`None`.
+        """Helper to convert to low-level filter.
 
         Args:
             post (bool): Indicates if this is a post-filter node.
 
         Returns:
-            None: If this is a post-filter.
+            Optional[query_pb2.PropertyFilter]: Returns :data:`None`, if
+                this is a post-filter, otherwise returns the protocol buffer
+                representation of the filter.
 
         Raises:
             NotImplementedError: If the ``opsymbol`` is ``!=`` or ``in``, since
                 they should correspond to a composite filter. This should
                 never occur since the constructor will create ``OR`` nodes for
                 ``!=`` and ``in``
-            NotImplementedError: If not a post-filter and the ``opsymbol``
-                is a simple comparison. (For now) this is because the original
-                implementation relied on a low-level datastore query module.
         """
         if post:
             return None
@@ -566,7 +565,9 @@ def _to_filter(self, post=False):
                 "to a single filter ({!r})".format(self._opsymbol)
             )
 
-        raise NotImplementedError("Missing datastore_query.make_filter")
+        return _datastore_query.make_filter(
+            self._name, self._opsymbol, self._value
+        )
 
 
 class PostFilterNode(Node):
@@ -612,7 +613,7 @@ def __eq__(self, other):
         return self is other or self.predicate == other.predicate
 
     def _to_filter(self, post=False):
-        """Helper to convert to low-level filter, or :data:`None`.
+        """Helper to convert to low-level filter.
 
         Args:
             post (bool): Indicates if this is a post-filter node.
@@ -801,19 +802,14 @@ def __eq__(self, other):
         return self._nodes == other._nodes
 
     def _to_filter(self, post=False):
-        """Helper to convert to low-level filter, or :data:`None`.
+        """Helper to convert to low-level filter.
 
         Args:
             post (bool): Indicates if this is a post-filter node.
 
         Returns:
             Optional[Node]: The single or composite filter corresponding to
-            the pre- or post-filter nodes stored.
-
-        Raises:
-            NotImplementedError: If a composite filter must be returned. This
-                is because the original implementation relied on a low-level
-                datastore query module.
+                the pre- or post-filter nodes stored. May return :data:`None`.
         """
         filters = []
         for node in self._nodes:
@@ -827,7 +823,7 @@ def _to_filter(self, post=False):
         if len(filters) == 1:
             return filters[0]
 
-        raise NotImplementedError("Missing datastore_query.CompositeFilter")
+        return _datastore_query.make_composite_and_filter(filters)
 
     def _post_filters(self):
         """Helper to extract post-filter nodes, if any.
@@ -962,6 +958,23 @@ def resolve(self, bindings, used):
 
         return DisjunctionNode(*resolved_nodes)
 
+    def _to_filter(self, post=False):
+        """Helper to convert to low-level filters.
+
+        Args:
+            post (bool): Indicates if this is a post-filter node.
+
+        Returns:
+            Optional[List[Node]]: List of filter protocol buffers that should
+                be combined using OR. The code in `_datastore_query` will
+                recognize that a list has been returned and run multiple
+                queries.
+        """
+        if post:
+            raise NotImplementedError("No idea what I should do here, yet.")
+
+        return [node._to_filter(post=post) for node in self._nodes]
+
 
 # AND and OR are preferred aliases for these.
 AND = ConjunctionNode

packages/google-cloud-ndb/tests/system/test_query.py

@@ -162,3 +162,64 @@ class SomeKind(ndb.Model):
     assert results[0].foo == 1
     assert results[0].bar == "a"
     assert results[0].key.namespace() == OTHER_NAMESPACE
+
+
+@pytest.mark.usefixtures("client_context")
+def test_filter_equal(ds_entity):
+    for i in range(5):
+        entity_id = test_utils.system.unique_resource_id()
+        ds_entity(KIND, entity_id, foo=i)
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    # query = SomeKind.query()  # Not implemented yet
+    query = ndb.Query(kind=KIND).filter(SomeKind.foo == 2)
+    results = query.fetch()
+    assert len(results) == 1
+    assert results[0].foo == 2
+
+
+@pytest.mark.usefixtures("client_context")
+def test_filter_not_equal(ds_entity):
+    for i in range(5):
+        entity_id = test_utils.system.unique_resource_id()
+        ds_entity(KIND, entity_id, foo=i)
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    # query = SomeKind.query()  # Not implemented yet
+    query = ndb.Query(kind=KIND).filter(SomeKind.foo != 2)
+    results = query.fetch()
+    assert len(results) == 4
+
+    results = sorted(results, key=operator.attrgetter("foo"))
+    assert [entity.foo for entity in results] == [0, 1, 3, 4]
+
+
+@pytest.mark.usefixtures("client_context")
+def test_filter_or(dispose_of):
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    @ndb.tasklet
+    def make_entities():
+        keys = yield (
+            SomeKind(foo=1, bar="a").put_async(),
+            SomeKind(foo=2, bar="b").put_async(),
+            SomeKind(foo=1, bar="c").put_async(),
+        )
+        for key in keys:
+            dispose_of(key._key)
+
+    make_entities().check_success()
+    query = ndb.Query(kind=KIND).filter(
+        ndb.OR(SomeKind.foo == 1, SomeKind.bar == "c")
+    )
+    results = query.fetch()
+    assert len(results) == 2
+
+    results = sorted(results, key=operator.attrgetter("foo"))
+    assert [entity.bar for entity in results] == ["a", "c"]