fix: Fix bug with the _GlobalCacheGetBatch. (#305)

@VladNF found a bug (and the solution) that caused results from calls to
the global cache to get out of synch with the futures waiting on them.

Fixes #294.

Author: Chris Rossi

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

@@ -169,7 +169,7 @@ def done_callback(self, cache_call):
     def make_call(self):
         """Call :method:`GlobalCache.get`."""
         cache = context_module.get_context().global_cache
-        return cache.get(self.todo.keys())
+        return cache.get(self.keys)
 
     def future_info(self, key):
         """Generate info string for Future."""

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

@@ -8,6 +8,8 @@
 from google.cloud import datastore
 from google.cloud import ndb
 
+from google.cloud.ndb import global_cache as global_cache_module
+
 from . import KIND, OTHER_KIND, OTHER_NAMESPACE
 
 
@@ -110,8 +112,8 @@ def make_entity(*key_args, **entity_kwargs):
 
 @pytest.fixture
 def dispose_of(with_ds_client, to_delete):
-    def delete_entity(ds_key):
-        to_delete.append(ds_key)
+    def delete_entity(*ds_keys):
+        to_delete.extend(ds_keys)
 
     return delete_entity
 
@@ -126,3 +128,11 @@ def client_context(namespace):
     client = ndb.Client(namespace=namespace)
     with client.context(cache_policy=False, legacy_data=False) as the_context:
         yield the_context
+
+
+@pytest.fixture
+def redis_context(client_context):
+    global_cache = global_cache_module.RedisCache.from_environment()
+    with client_context.new(global_cache=global_cache).use() as context:
+        context.set_global_cache_policy(None)  # Use default
+        yield context

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

@@ -19,6 +19,7 @@
 import functools
 import operator
 import os
+import random
 import threading
 import zlib
 
@@ -118,7 +119,7 @@ class SomeKind(ndb.Model):
 
 
 @pytest.mark.skipif(not USE_REDIS_CACHE, reason="Redis is not configured")
-def test_retrieve_entity_with_redis_cache(ds_entity, client_context):
+def test_retrieve_entity_with_redis_cache(ds_entity, redis_context):
     entity_id = test_utils.system.unique_resource_id()
     ds_entity(KIND, entity_id, foo=42, bar="none", baz=b"night")
 
@@ -127,29 +128,25 @@ class SomeKind(ndb.Model):
         bar = ndb.StringProperty()
         baz = ndb.StringProperty()
 
-    global_cache = global_cache_module.RedisCache.from_environment()
-    with client_context.new(global_cache=global_cache).use() as context:
-        context.set_global_cache_policy(None)  # Use default
+    key = ndb.Key(KIND, entity_id)
+    entity = key.get()
+    assert isinstance(entity, SomeKind)
+    assert entity.foo == 42
+    assert entity.bar == "none"
+    assert entity.baz == "night"
 
-        key = ndb.Key(KIND, entity_id)
+    cache_key = _cache.global_cache_key(key._key)
+    assert redis_context.global_cache.redis.get(cache_key) is not None
+
+    patch = mock.patch("google.cloud.ndb._datastore_api._LookupBatch.add")
+    patch.side_effect = Exception("Shouldn't call this")
+    with patch:
         entity = key.get()
         assert isinstance(entity, SomeKind)
         assert entity.foo == 42
         assert entity.bar == "none"
         assert entity.baz == "night"
 
-        cache_key = _cache.global_cache_key(key._key)
-        assert global_cache.redis.get(cache_key) is not None
-
-        patch = mock.patch("google.cloud.ndb._datastore_api._LookupBatch.add")
-        patch.side_effect = Exception("Shouldn't call this")
-        with patch:
-            entity = key.get()
-            assert isinstance(entity, SomeKind)
-            assert entity.foo == 42
-            assert entity.bar == "none"
-            assert entity.baz == "night"
-
 
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity_not_found(ds_entity):
@@ -500,33 +497,29 @@ class SomeKind(ndb.Model):
 
 
 @pytest.mark.skipif(not USE_REDIS_CACHE, reason="Redis is not configured")
-def test_insert_entity_with_redis_cache(dispose_of, client_context):
+def test_insert_entity_with_redis_cache(dispose_of, redis_context):
     class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
 
-    global_cache = global_cache_module.RedisCache.from_environment()
-    with client_context.new(global_cache=global_cache).use() as context:
-        context.set_global_cache_policy(None)  # Use default
-
-        entity = SomeKind(foo=42, bar="none")
-        key = entity.put()
-        dispose_of(key._key)
-        cache_key = _cache.global_cache_key(key._key)
-        assert global_cache.redis.get(cache_key) is None
+    entity = SomeKind(foo=42, bar="none")
+    key = entity.put()
+    dispose_of(key._key)
+    cache_key = _cache.global_cache_key(key._key)
+    assert redis_context.global_cache.redis.get(cache_key) is None
 
-        retrieved = key.get()
-        assert retrieved.foo == 42
-        assert retrieved.bar == "none"
+    retrieved = key.get()
+    assert retrieved.foo == 42
+    assert retrieved.bar == "none"
 
-        assert global_cache.redis.get(cache_key) is not None
+    assert redis_context.global_cache.redis.get(cache_key) is not None
 
-        entity.foo = 43
-        entity.put()
+    entity.foo = 43
+    entity.put()
 
-        # This is py27 behavior. I can see a case being made for caching the
-        # entity on write rather than waiting for a subsequent lookup.
-        assert global_cache.redis.get(cache_key) is None
+    # This is py27 behavior. I can see a case being made for caching the
+    # entity on write rather than waiting for a subsequent lookup.
+    assert redis_context.global_cache.redis.get(cache_key) is None
 
 
 @pytest.mark.usefixtures("client_context")
@@ -671,7 +664,7 @@ class SomeKind(ndb.Model):
 
 
 @pytest.mark.skipif(not USE_REDIS_CACHE, reason="Redis is not configured")
-def test_delete_entity_with_redis_cache(ds_entity, client_context):
+def test_delete_entity_with_redis_cache(ds_entity, redis_context):
     entity_id = test_utils.system.unique_resource_id()
     ds_entity(KIND, entity_id, foo=42)
 
@@ -680,19 +673,17 @@ class SomeKind(ndb.Model):
 
     key = ndb.Key(KIND, entity_id)
     cache_key = _cache.global_cache_key(key._key)
-    global_cache = global_cache_module.RedisCache.from_environment()
 
-    with client_context.new(global_cache=global_cache).use():
-        assert key.get().foo == 42
-        assert global_cache.redis.get(cache_key) is not None
+    assert key.get().foo == 42
+    assert redis_context.global_cache.redis.get(cache_key) is not None
 
-        assert key.delete() is None
-        assert global_cache.redis.get(cache_key) is None
+    assert key.delete() is None
+    assert redis_context.global_cache.redis.get(cache_key) is None
 
-        # This is py27 behavior. Not entirely sold on leaving _LOCKED value for
-        # Datastore misses.
-        assert key.get() is None
-        assert global_cache.redis.get(cache_key) == b"0"
+    # This is py27 behavior. Not entirely sold on leaving _LOCKED value for
+    # Datastore misses.
+    assert key.get() is None
+    assert redis_context.global_cache.redis.get(cache_key) == b"0"
 
 
 @pytest.mark.usefixtures("client_context")
@@ -1113,3 +1104,30 @@ class SomeKind(ndb.Model):
 
     retreived = key.get()
     assert retreived.foo == ["", ""]
+
+
+@pytest.mark.usefixtures("redis_context")
+def test_multi_get_weirdness_with_redis(dispose_of):
+    """Regression test for issue #294.
+
+    https://github.com/googleapis/python-ndb/issues/294
+    """
+
+    class SomeKind(ndb.Model):
+        foo = ndb.StringProperty()
+
+    objects = [SomeKind(foo=str(i)) for i in range(10)]
+    keys = ndb.put_multi(objects)
+    for key in keys:
+        dispose_of(key._key)
+    ndb.get_multi(keys)
+
+    one_object = random.choice(keys).get()
+    one_object.foo = "CHANGED"
+    one_object.put()
+
+    objects_upd = ndb.get_multi(keys)
+    keys_upd = [obj.key for obj in objects_upd]
+    assert len(keys_upd) == len(keys)
+    assert len(set(keys_upd)) == len(set(keys))
+    assert set(keys_upd) == set(keys)

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

@@ -98,7 +98,7 @@ def test_add_and_idle_and_done_callbacks(in_context):
         with in_context.new(global_cache=cache).use():
             batch.idle_callback()
 
-        cache.get.assert_called_once_with(batch.todo.keys())
+        cache.get.assert_called_once_with(batch.keys)
         assert future1.result() == b"one"
         assert future2.result() == b"two"
         assert future3.result() == b"one"
@@ -118,7 +118,7 @@ def test_add_and_idle_and_done_callbacks_synchronous(in_context):
         with in_context.new(global_cache=cache).use():
             batch.idle_callback()
 
-        cache.get.assert_called_once_with(batch.todo.keys())
+        cache.get.assert_called_once_with(batch.keys)
         assert future1.result() == b"one"
         assert future2.result() == b"two"
 
@@ -139,7 +139,7 @@ def test_add_and_idle_and_done_callbacks_w_error(in_context):
         with in_context.new(global_cache=cache).use():
             batch.idle_callback()
 
-        cache.get.assert_called_once_with(batch.todo.keys())
+        cache.get.assert_called_once_with(batch.keys)
         assert future1.exception() is error
         assert future2.exception() is error