From eb1c0f8e50bc80ed9250b959f05690957c5308f6 Mon Sep 17 00:00:00 2001
From: Chris Rossi <chris@archimedeanco.com>
Date: Mon, 26 Apr 2021 13:18:14 -0400
Subject: [PATCH] fix: properly handle error when clearing cache

Fixes #633
---
 google/cloud/ndb/_cache.py | 21 +++++++++++----------
 tests/unit/test__cache.py  | 30 ++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/google/cloud/ndb/_cache.py b/google/cloud/ndb/_cache.py
index a7ef0cc2..c5f7f095 100644
--- a/google/cloud/ndb/_cache.py
+++ b/google/cloud/ndb/_cache.py
@@ -165,22 +165,23 @@ def retry_wrapper(*args, **kwargs):
         @tasklets.tasklet
         def wrapper(*args, **kwargs):
             cache = _global_cache()
+
+            is_read = read
+            if not is_read:
+                is_read = kwargs.get("read", False)
+
+            strict = cache.strict_read if is_read else cache.strict_write
+            if strict:
+                function = retry(wrapped, cache.transient_errors)
+            else:
+                function = wrapped
+
             try:
                 if cache.clear_cache_soon:
                     warnings.warn("Clearing global cache...", RuntimeWarning)
                     cache.clear()
                     cache.clear_cache_soon = False
 
-                is_read = read
-                if not is_read:
-                    is_read = kwargs.get("read", False)
-
-                strict = cache.strict_read if is_read else cache.strict_write
-                if strict:
-                    function = retry(wrapped, cache.transient_errors)
-                else:
-                    function = wrapped
-
                 result = yield function(*args, **kwargs)
                 raise tasklets.Return(result)
 
diff --git a/tests/unit/test__cache.py b/tests/unit/test__cache.py
index 914ce6c2..54600555 100644
--- a/tests/unit/test__cache.py
+++ b/tests/unit/test__cache.py
@@ -137,6 +137,36 @@ def test_global_get_clear_cache_soon(_batch, _global_cache):
     _global_cache.return_value.clear.assert_called_once_with()
 
 
+@pytest.mark.usefixtures("in_context")
+@mock.patch("google.cloud.ndb._cache._global_cache")
+@mock.patch("google.cloud.ndb._cache._batch")
+def test_global_get_clear_cache_soon_with_error(_batch, _global_cache):
+    """Regression test for #633
+
+    https://github.com/googleapis/python-ndb/issues/633
+    """
+
+    class TransientError(Exception):
+        pass
+
+    batch = _batch.get_batch.return_value
+    future = _future_result("hi mom!")
+    batch.add.return_value = future
+    _global_cache.return_value = mock.Mock(
+        transient_errors=(TransientError),
+        clear_cache_soon=True,
+        strict_read=False,
+        clear=mock.Mock(side_effect=TransientError("oops!"), spec=()),
+        spec=("transient_errors", "clear_cache_soon", "clear", "strict_read"),
+    )
+
+    with warnings.catch_warnings(record=True) as logged:
+        assert _cache.global_get(b"foo").result() is None
+        assert len(logged) == 2
+
+    _global_cache.return_value.clear.assert_called_once_with()
+
+
 @pytest.mark.usefixtures("in_context")
 @mock.patch("google.cloud.ndb.tasklets.sleep")
 @mock.patch("google.cloud.ndb._cache._global_cache")