fix: replicate legacy behavior for using cache with queries (#613)

Corrects the issue pointed out in #586, that we weren't doing the right
thing for deleted entities. Also corrects an issue noticed while fixing
that, where the cache wasn't being updated with entities from queries.
Behavior should now match legacy in both these regards.

Fixes #586

Author: Chris Rossi

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

@@ -379,6 +379,18 @@ def _next_batch(self):
             for result_pb in response.batch.entity_results
         ]
 
+        if result_type == RESULT_TYPE_FULL:
+            # If we cached a delete, remove it from the result set. This may come cause
+            # some queries to return less than their limit even if there are more
+            # results. As far as I can tell, that was also a possibility with the legacy
+            # version.
+            context = context_module.get_context()
+            self._batch = [
+                result
+                for result in self._batch
+                if result.check_cache(context) is not None
+            ]
+
         self._has_next_batch = more_results = batch.more_results == NOT_FINISHED
 
         self._more_results_after_limit = batch.more_results == MORE_RESULTS_AFTER_LIMIT
@@ -745,6 +757,8 @@ class _Result(object):
             order.
     """
 
+    _key = None
+
     def __init__(self, result_type, result_pb, order_by=None):
         self.result_type = result_type
         self.result_pb = result_pb
@@ -802,8 +816,38 @@ def _compare(self, other):
 
         return 0
 
+    def key(self):
+        """Construct the key for this result.
+
+        Returns:
+            key.Key: The key.
+        """
+        if self._key is None:
+            key_pb = self.result_pb.entity.key
+            ds_key = helpers.key_from_protobuf(key_pb)
+            self._key = key_module.Key._from_ds_key(ds_key)
+
+        return self._key
+
+    def check_cache(self, context):
+        """Check local context cache for entity.
+
+        Returns:
+            Any: The NDB entity for this result, if it is cached, otherwise
+                `_KEY_NOT_IN_CACHE`. May also return `None` if entity was deleted which
+                will cause `None` to be recorded in the cache.
+        """
+        key = self.key()
+        if context._use_cache(key):
+            try:
+                return context.cache.get_and_validate(key)
+            except KeyError:
+                pass
+
+        return _KEY_NOT_IN_CACHE
+
     def entity(self):
-        """Get an entity for an entity result. Use the cache if available.
+        """Get an entity for an entity result. Use or update the cache if available.
 
         Args:
             projection (Optional[Sequence[str]]): Sequence of property names to
@@ -816,19 +860,12 @@ def entity(self):
         if self.result_type == RESULT_TYPE_FULL:
             # First check the cache.
             context = context_module.get_context()
-            key_pb = self.result_pb.entity.key
-            ds_key = helpers.key_from_protobuf(key_pb)
-            key = key_module.Key._from_ds_key(ds_key)
-            entity = _KEY_NOT_IN_CACHE
-            use_cache = context._use_cache(key)
-            if use_cache:
-                try:
-                    entity = context.cache.get_and_validate(key)
-                except KeyError:
-                    pass
-            if entity is None or entity is _KEY_NOT_IN_CACHE:
-                # entity not in cache, create one.
+            entity = self.check_cache(context)
+            if entity is _KEY_NOT_IN_CACHE:
+                # entity not in cache, create one, and then add it to cache
                 entity = model._entity_from_protobuf(self.result_pb.entity)
+                if context._use_cache(entity.key):
+                    context.cache[entity.key] = entity
             return entity
 
         elif self.result_type == RESULT_TYPE_PROJECTION:
@@ -838,10 +875,7 @@ def entity(self):
             return entity
 
         elif self.result_type == RESULT_TYPE_KEY_ONLY:
-            key_pb = self.result_pb.entity.key
-            ds_key = helpers.key_from_protobuf(key_pb)
-            key = key_module.Key._from_ds_key(ds_key)
-            return key
+            return self.key()
 
         raise NotImplementedError("Got unexpected entity result type for query.")
 

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

@@ -1859,7 +1859,7 @@ class SomeKind(ndb.Model):
     assert result.foo == index
 
 
-def test_uncomitted_deletes(dispose_of, client_context):
+def test_uncommitted_deletes(dispose_of, client_context):
     """Regression test for Issue #586
 
     https://github.com/googleapis/python-ndb/issues/586
@@ -1878,7 +1878,24 @@ class SomeKind(ndb.Model):
     @ndb.transactional()
     def do_the_thing(key):
         key.delete()  # Will be cached but not committed when query runs
-        return SomeKind.query(SomeKind.foo == 42, ancestor=parent_key).get()
+        return SomeKind.query(SomeKind.foo == 42, ancestor=parent_key).fetch()
 
     with client_context.new(cache_policy=None).use():  # Use default cache policy
-        assert do_the_thing(key).foo == 42
+        assert len(do_the_thing(key)) == 0
+
+
+def test_query_updates_cache(dispose_of, client_context):
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    entity = SomeKind(foo=42)
+    key = entity.put()
+    dispose_of(key._key)
+    eventually(SomeKind.query().fetch, length_equals(1))
+
+    with client_context.new(cache_policy=None).use():  # Use default cache policy
+        retrieved = SomeKind.query().get()
+        assert retrieved.foo == 42
+
+        # If there is a cache hit, we'll get back the same object, not just a copy
+        assert key.get() is retrieved

packages/google-cloud-ndb/tests/unit/test__datastore_query.py

@@ -590,10 +590,28 @@ def test_probably_has_next_finished():
     @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
     def test__next_batch(_datastore_run_query):
+        entity1 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+            )
+        )
+        entity2 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=43)],
+            )
+        )
+        entity3 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=44)],
+            )
+        )
         entity_results = [
-            mock.Mock(entity="entity1", cursor=b"a"),
-            mock.Mock(entity="entity2", cursor=b"b"),
-            mock.Mock(entity="entity3", cursor=b"c"),
+            mock.Mock(entity=entity1, cursor=b"a"),
+            mock.Mock(entity=entity2, cursor=b"b"),
+            mock.Mock(entity=entity3, cursor=b"c"),
         ]
         _datastore_run_query.return_value = utils.future_result(
             mock.Mock(
@@ -611,24 +629,91 @@ def test__next_batch(_datastore_run_query):
         assert iterator._next_batch().result() is None
         assert iterator._index == 0
         assert len(iterator._batch) == 3
-        assert iterator._batch[0].result_pb.entity == "entity1"
+        assert iterator._batch[0].result_pb.entity == entity1
         assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
         assert iterator._batch[0].order_by is None
         assert not iterator._has_next_batch
 
+    @staticmethod
+    @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
+    def test__next_batch_cached_delete(_datastore_run_query, in_context):
+        entity1 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+            )
+        )
+        entity2 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=43)],
+            )
+        )
+        entity3 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=44)],
+            )
+        )
+        entity_results = [
+            mock.Mock(entity=entity1, cursor=b"a"),
+            mock.Mock(entity=entity2, cursor=b"b"),
+            mock.Mock(entity=entity3, cursor=b"c"),
+        ]
+        in_context.cache[key_module.Key("ThisKind", 43)] = None
+        _datastore_run_query.return_value = utils.future_result(
+            mock.Mock(
+                batch=mock.Mock(
+                    entity_result_type=query_pb2.EntityResult.FULL,
+                    entity_results=entity_results,
+                    end_cursor=b"abc",
+                    more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
+                )
+            )
+        )
+
+        query = query_module.QueryOptions()
+        iterator = _datastore_query._QueryIteratorImpl(query)
+        assert iterator._next_batch().result() is None
+        assert iterator._index == 0
+        assert len(iterator._batch) == 2
+        assert iterator._batch[0].result_pb.entity == entity1
+        assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
+        assert iterator._batch[0].order_by is None
+        assert iterator._batch[1].result_pb.entity == entity3
+        assert not iterator._has_next_batch
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
     def test__next_batch_has_more(_datastore_run_query):
+        entity1 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+            )
+        )
+        entity2 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=43)],
+            )
+        )
+        entity3 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=44)],
+            )
+        )
         entity_results = [
-            mock.Mock(entity="entity1", cursor=b"a"),
-            mock.Mock(entity="entity2", cursor=b"b"),
-            mock.Mock(entity="entity3", cursor=b"c"),
+            mock.Mock(entity=entity1, cursor=b"a"),
+            mock.Mock(entity=entity2, cursor=b"b"),
+            mock.Mock(entity=entity3, cursor=b"c"),
         ]
         _datastore_run_query.return_value = utils.future_result(
             mock.Mock(
                 batch=mock.Mock(
-                    entity_result_type=query_pb2.EntityResult.FULL,
+                    entity_result_type=query_pb2.EntityResult.PROJECTION,
                     entity_results=entity_results,
                     end_cursor=b"abc",
                     more_results=query_pb2.QueryResultBatch.NOT_FINISHED,
@@ -641,8 +726,8 @@ def test__next_batch_has_more(_datastore_run_query):
         assert iterator._next_batch().result() is None
         assert iterator._index == 0
         assert len(iterator._batch) == 3
-        assert iterator._batch[0].result_pb.entity == "entity1"
-        assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
+        assert iterator._batch[0].result_pb.entity == entity1
+        assert iterator._batch[0].result_type == query_pb2.EntityResult.PROJECTION
         assert iterator._batch[0].order_by is None
         assert iterator._has_next_batch
         assert iterator._query.start_cursor.cursor == b"abc"
@@ -655,10 +740,28 @@ def test__next_batch_has_more_w_offset_and_limit(_datastore_run_query):
 
         https://github.com/googleapis/python-ndb/issues/236
         """
+        entity1 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+            )
+        )
+        entity2 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=43)],
+            )
+        )
+        entity3 = mock.Mock(
+            key=entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=44)],
+            )
+        )
         entity_results = [
-            mock.Mock(entity="entity1", cursor=b"a"),
-            mock.Mock(entity="entity2", cursor=b"b"),
-            mock.Mock(entity="entity3", cursor=b"c"),
+            mock.Mock(entity=entity1, cursor=b"a"),
+            mock.Mock(entity=entity2, cursor=b"b"),
+            mock.Mock(entity=entity3, cursor=b"c"),
         ]
         _datastore_run_query.return_value = utils.future_result(
             mock.Mock(
@@ -677,7 +780,7 @@ def test__next_batch_has_more_w_offset_and_limit(_datastore_run_query):
         assert iterator._next_batch().result() is None
         assert iterator._index == 0
         assert len(iterator._batch) == 3
-        assert iterator._batch[0].result_pb.entity == "entity1"
+        assert iterator._batch[0].result_pb.entity == entity1
         assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
         assert iterator._batch[0].order_by is None
         assert iterator._has_next_batch
@@ -1466,15 +1569,16 @@ def test_entity_full_entity(model):
             partition_id=entity_pb2.PartitionId(project_id="testing"),
             path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
         )
-        entity = mock.Mock(key=key_pb)
+        entity_pb = mock.Mock(key=key_pb)
+        entity = mock.Mock(key=key_module.Key("ThisKind", 42))
         model._entity_from_protobuf.return_value = entity
         result = _datastore_query._Result(
             _datastore_query.RESULT_TYPE_FULL,
-            mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")),
+            mock.Mock(entity=entity_pb, cursor=b"123", spec=("entity", "cursor")),
         )
 
         assert result.entity() is entity
-        model._entity_from_protobuf.assert_called_once_with(entity)
+        model._entity_from_protobuf.assert_called_once_with(entity_pb)
 
     @staticmethod
     @pytest.mark.usefixtures("in_context")