fix: all query types should use cache if available (#454)

* fix: all query types should use cache if available
We were returning cache values in some cases and not others, generating some confusion
Refs #441

Author: cguardia

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

@@ -24,6 +24,7 @@
 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
 from google.cloud.ndb import exceptions
 from google.cloud.ndb import key as key_module
@@ -52,6 +53,8 @@
     ">=": query_pb2.PropertyFilter.GREATER_THAN_OR_EQUAL,
 }
 
+_KEY_NOT_IN_CACHE = object()
+
 
 def make_filter(name, op, value):
     """Make a property filter protocol buffer.
@@ -698,7 +701,7 @@ def _compare(self, other):
         return 0
 
     def entity(self):
-        """Get an entity for an entity result.
+        """Get an entity for an entity result. Use the cache if available.
 
         Args:
             projection (Optional[Sequence[str]]): Sequence of property names to
@@ -709,7 +712,21 @@ def entity(self):
         """
 
         if self.result_type == RESULT_TYPE_FULL:
-            entity = model._entity_from_protobuf(self.result_pb.entity)
+            # 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 _KEY_NOT_IN_CACHE:
+                # entity not in cache, create one.
+                entity = model._entity_from_protobuf(self.result_pb.entity)
             return entity
 
         elif self.result_type == RESULT_TYPE_PROJECTION:

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

@@ -279,9 +279,9 @@ def _clear_global_cache(self):
         if keys:
             yield [_cache.global_delete(key) for key in keys]
 
-    def _use_cache(self, key, options):
+    def _use_cache(self, key, options=None):
         """Return whether to use the context cache for this key."""
-        flag = options.use_cache
+        flag = options.use_cache if options else None
         if flag is None:
             flag = self.cache_policy(key)
         if flag is None:

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

@@ -1429,3 +1429,47 @@ class SomeKind(ndb.Model):
 
     retrieved = key.get()
     assert retrieved.foo == datetime.datetime(2020, 8, 8, 1, 2, 3)
+
+
+def test_cache_returns_entity_if_available(dispose_of, client_context):
+    """Regression test for #441
+
+    https://github.com/googleapis/python-ndb/issues/441
+    """
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    client_context.set_cache_policy(None)  # Use default
+
+    somekind = SomeKind(foo=1)
+    key = somekind.put()
+    dispose_of(key._key)
+
+    query = ndb.Query(kind="SomeKind")
+    ourkind = query.get()
+    ourkind.bar = "confusing"
+
+    assert somekind.bar == "confusing"
+
+
+def test_cache_off_new_entity_created(dispose_of, client_context):
+    """Regression test for #441
+
+    https://github.com/googleapis/python-ndb/issues/441
+    """
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    somekind = SomeKind(foo=1)
+    key = somekind.put()
+    dispose_of(key._key)
+
+    query = ndb.Query(kind="SomeKind")
+    ourkind = query.get()
+    ourkind.bar = "confusing"
+
+    assert somekind.bar is None

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

@@ -26,6 +26,7 @@
 from google.cloud.datastore_v1.proto import query_pb2
 
 from google.cloud.ndb import _datastore_query
+from google.cloud.ndb import context as context_module
 from google.cloud.ndb import exceptions
 from google.cloud.ndb import key as key_module
 from google.cloud.ndb import model
@@ -1052,16 +1053,64 @@ def test_entity_unsupported_result_type(model):
             result.entity()
 
     @staticmethod
+    @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb._datastore_query.model")
     def test_entity_full_entity(model):
-        model._entity_from_protobuf.return_value = "bar"
+        key_pb = entity_pb2.Key(
+            partition_id=entity_pb2.PartitionId(project_id="testing"),
+            path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+        )
+        entity = mock.Mock(key=key_pb)
+        model._entity_from_protobuf.return_value = entity
         result = _datastore_query._Result(
             _datastore_query.RESULT_TYPE_FULL,
-            mock.Mock(entity="foo", cursor=b"123", spec=("entity", "cursor")),
+            mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")),
         )
 
-        assert result.entity() == "bar"
-        model._entity_from_protobuf.assert_called_once_with("foo")
+        assert result.entity() is entity
+        model._entity_from_protobuf.assert_called_once_with(entity)
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query.model")
+    def test_entity_full_entity_cached(model):
+        key = key_module.Key("ThisKind", 42)
+        key_pb = entity_pb2.Key(
+            partition_id=entity_pb2.PartitionId(project_id="testing"),
+            path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+        )
+        entity = mock.Mock(key=key_pb)
+        cached_entity = mock.Mock(key=key_pb, _key=key)
+        context = context_module.get_context()
+        context.cache.data[key] = cached_entity
+        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")),
+        )
+
+        assert result.entity() is not entity
+        assert result.entity() is cached_entity
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query.model")
+    def test_entity_full_entity_no_cache(model):
+        context = context_module.get_context()
+        with context.new(cache_policy=False).use():
+            key_pb = entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)],
+            )
+            entity = mock.Mock(key=key_pb)
+            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")
+                ),
+            )
+            assert result.entity() is entity
 
     @staticmethod
     @pytest.mark.usefixtures("in_context")