Include a couple of GRPC transient errors in retry. (#52)

This is a fix for #51.

Author: Chris Rossi

packages/google-cloud-ndb/src/google/cloud/ndb/_retry.py

@@ -15,6 +15,7 @@
 """Retry functions."""
 
 import functools
+import grpc
 import itertools
 
 from google.api_core import retry as core_retry
@@ -65,7 +66,7 @@ def retry_wrapper(*args, **kwargs):
             except Exception as e:
                 # `e` is removed from locals at end of block
                 error = e  # See: https://goo.gl/5J8BMK
-                if not core_retry.if_transient_error(error):
+                if not is_transient_error(error):
                     raise
 
             yield tasklets.sleep(sleep_time)
@@ -78,3 +79,23 @@ def retry_wrapper(*args, **kwargs):
         )
 
     return retry_wrapper
+
+
+TRANSIENT_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.INTERNAL)
+
+
+def is_transient_error(error):
+    """Determine whether an error is transient.
+
+    Returns:
+        bool: True if error is transient, else False.
+    """
+    if core_retry.if_transient_error(error):
+        return True
+
+    method = getattr(error, "code", None)
+    if method is not None:
+        code = method()
+        return code in TRANSIENT_CODES
+
+    return False

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

@@ -16,6 +16,7 @@
 
 from unittest import mock
 
+import grpc
 import pytest
 
 from google.api_core import exceptions as core_exceptions
@@ -125,3 +126,51 @@ def callback():
         assert error_context.value.cause is error
         assert sleep.call_count == 5
         assert sleep.call_args[0][0] == 4
+
+
+class Test_is_transient_error:
+    @staticmethod
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_core_says_yes(core_retry):
+        error = object()
+        core_retry.if_transient_error.return_value = True
+        assert _retry.is_transient_error(error) is True
+        core_retry.if_transient_error.assert_called_once_with(error)
+
+    @staticmethod
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_core_says_no_we_say_no(core_retry):
+        error = object()
+        core_retry.if_transient_error.return_value = False
+        assert _retry.is_transient_error(error) is False
+        core_retry.if_transient_error.assert_called_once_with(error)
+
+    @staticmethod
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_unavailable(core_retry):
+        error = mock.Mock(
+            code=mock.Mock(return_value=grpc.StatusCode.UNAVAILABLE)
+        )
+        core_retry.if_transient_error.return_value = False
+        assert _retry.is_transient_error(error) is True
+        core_retry.if_transient_error.assert_called_once_with(error)
+
+    @staticmethod
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_internal(core_retry):
+        error = mock.Mock(
+            code=mock.Mock(return_value=grpc.StatusCode.INTERNAL)
+        )
+        core_retry.if_transient_error.return_value = False
+        assert _retry.is_transient_error(error) is True
+        core_retry.if_transient_error.assert_called_once_with(error)
+
+    @staticmethod
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_unauthenticated(core_retry):
+        error = mock.Mock(
+            code=mock.Mock(return_value=grpc.StatusCode.UNAUTHENTICATED)
+        )
+        core_retry.if_transient_error.return_value = False
+        assert _retry.is_transient_error(error) is False
+        core_retry.if_transient_error.assert_called_once_with(error)