fix: Show a non-None error for core_exception.Unknown errors. (#968)

* fix: Show a non-None error for core_exception.Unknown errors. * Add test for Unknown api errors.
googleapis · Mar 9, 2024 · 66e61cc · 66e61cc
1 parent 98fb176
commit 66e61cc
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/google/cloud/ndb/_retry.py b/google/cloud/ndb/_retry.py
@@ -82,9 +82,10 @@ def retry_wrapper(*args, **kwargs):
                     result = yield result
             except exceptions.NestedRetryException as e:
                 error = e
-            except Exception as e:
+            except BaseException as e:
                 # `e` is removed from locals at end of block
                 error = e  # See: https://goo.gl/5J8BMK
+
                 if not is_transient_error(error):
                     # If we are in an inner retry block, use special nested
                     # retry exception to bubble up to outer retry. Else, raise
@@ -104,6 +105,10 @@ def retry_wrapper(*args, **kwargs):
 
             yield tasklets.sleep(sleep_time)
 
+        # Unknown errors really want to show up as None, so manually set the error.
+        if isinstance(error, core_exceptions.Unknown):
+            error = "google.api_core.exceptions.Unknown"
+
         raise core_exceptions.RetryError(
             "Maximum number of {} retries exceeded while calling {}".format(
                 retries, callback

diff --git a/tests/unit/test__retry.py b/tests/unit/test__retry.py
@@ -98,6 +98,18 @@ def callback():
         retry = _retry.retry_async(callback)
         assert retry().exception() is error
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    def test_api_core_unknown():
+        def callback():
+            raise core_exceptions.Unknown("Unknown")
+
+        with pytest.raises(core_exceptions.RetryError) as e:
+            retry = _retry.retry_async(callback, retries=1)
+            retry().result()
+
+        assert e.value.cause == "google.api_core.exceptions.Unknown"
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb.tasklets.sleep")