fix: updating retry strategy to call _retry_on_aborted_exception for batch api calls

Author: aakashanandg

google/cloud/spanner_v1/_helpers.py

@@ -464,20 +464,32 @@ def _metadata_with_prefix(prefix, **kw):
     return [("google-cloud-resource-prefix", prefix)]
 
 
-def _retry_on_aborted_exception(exc, deadline, attempts, allowed_exceptions):
+def _retry_on_aborted_exception(
+    func,
+    deadline,
+    allowed_exceptions=None,
+):
     """
-    Handles the retry logic for Aborted exceptions, considering the deadline.
-    Returns True if the exception is retried, False otherwise.
+    Handles retry logic for Aborted exceptions, considering the deadline.
+    Retries the function in case of Aborted exceptions and other allowed exceptions.
     """
-    if isinstance(exc, Aborted) and deadline is not None:
-        # The logic for handling Aborted exceptions
-        if (
-            allowed_exceptions is not None
-            and allowed_exceptions.get(exc.__class__) is not None
-        ):
+    attempts = 0
+    while True:
+        try:
+            attempts += 1
+            return func()
+        except Aborted as exc:
             _delay_until_retry(exc, deadline=deadline, attempts=attempts)
-            return True
-    return False
+            continue
+        except Exception as exc:
+            try:
+                retry_result = _retry(func=func, allowed_exceptions=allowed_exceptions)
+                if retry_result is not None:
+                    return retry_result
+                else:
+                    raise exc
+            except Aborted:
+                continue
 
 
 def _retry(
@@ -486,14 +498,9 @@ def _retry(
     delay=2,
     allowed_exceptions=None,
     beforeNextRetry=None,
-    deadline=None,
 ):
     """
-    Retry a specified function with different logic based on the type of exception raised.
-
-    If the exception is of type google.api_core.exceptions.Aborted,
-    apply an alternate retry strategy that relies on the provided deadline value instead of a fixed number of retries.
-    For all other exceptions, retry the function up to a specified number of times.
+    Retry a function with a specified number of retries, delay between retries, and list of allowed exceptions.
 
     Args:
         func: The function to be retried.
@@ -507,19 +514,13 @@ def _retry(
         The result of the function if it is successful, or raises the last exception if all retries fail.
     """
     retries = 0
-    attempts = 0
-    while True:
-        if retries > retry_count:
-            raise Exception("Exceeded retry count.")
+    while retries <= retry_count:
         if retries > 0 and beforeNextRetry:
             beforeNextRetry(retries, delay)
 
         try:
-            attempts += 1
             return func()
         except Exception as exc:
-            if _retry_on_aborted_exception(exc, deadline, attempts, allowed_exceptions):
-                continue
             if (
                 allowed_exceptions is None or exc.__class__ in allowed_exceptions
             ) and retries < retry_count:
@@ -582,7 +583,6 @@ def _delay_until_retry(exc, deadline, attempts):
         raise
 
     delay = _get_retry_delay(cause, attempts)
-    print(now, delay, deadline)
     if delay is not None:
         if now + delay > deadline:
             raise

google/cloud/spanner_v1/batch.py

@@ -29,9 +29,9 @@
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1._helpers import _retry
+from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
 from google.cloud.spanner_v1._helpers import _check_rst_stream_error
 from google.api_core.exceptions import InternalServerError
-from google.api_core.exceptions import Aborted
 import time
 
 DEFAULT_RETRY_TIMEOUT_SECS = 30
@@ -235,11 +235,10 @@ def commit(
             deadline = time.time() + kwargs.get(
                 "timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
             )
-            response = _retry(
+            response = _retry_on_aborted_exception(
                 method,
                 allowed_exceptions={
                     InternalServerError: _check_rst_stream_error,
-                    Aborted: lambda exc: None,
                 },
                 deadline=deadline,
             )
@@ -388,4 +387,3 @@ def _make_write_pb(table, columns, values):
     return Mutation.Write(
         table=table, columns=columns, values=_make_list_value_pbs(values)
     )
-

tests/unit/test__helpers.py

@@ -882,6 +882,98 @@ def test_check_rst_stream_error(self):
 
         self.assertEqual(test_api.test_fxn.call_count, 3)
 
+    def test_retry_on_aborted_exception_with_success_after_first_aborted_retry(self):
+        from google.api_core.exceptions import Aborted
+        import time
+        from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
+        import functools
+
+        test_api = mock.create_autospec(self.test_class)
+        test_api.test_fxn.side_effect = [
+            Aborted("aborted exception", errors=("Aborted error")),
+            "true",
+        ]
+        deadline = time.time() + 30
+        result_after_retry = _retry_on_aborted_exception(
+            functools.partial(test_api.test_fxn), deadline
+        )
+
+        self.assertEqual(test_api.test_fxn.call_count, 2)
+        self.assertTrue(result_after_retry)
+
+    def test_retry_on_aborted_exception_with_success_after_three_retries(self):
+        from google.api_core.exceptions import Aborted
+        from google.api_core.exceptions import InternalServerError
+        import time
+        from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
+        import functools
+
+        test_api = mock.create_autospec(self.test_class)
+        # Case where aborted exception is thrown after other generic exceptions
+        test_api.test_fxn.side_effect = [
+            InternalServerError("testing"),
+            InternalServerError("testing"),
+            Aborted("aborted exception", errors=("Aborted error")),
+            "true",
+        ]
+        allowed_exceptions = {
+            InternalServerError: lambda exc: None,
+        }
+        deadline = time.time() + 30
+        _retry_on_aborted_exception(
+            functools.partial(test_api.test_fxn),
+            deadline=deadline,
+            allowed_exceptions=allowed_exceptions,
+        )
+
+        self.assertEqual(test_api.test_fxn.call_count, 4)
+
+    def test_retry_on_aborted_exception_raises_aborted_if_deadline_expires(self):
+        from google.api_core.exceptions import Aborted
+        import time
+        from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
+        import functools
+
+        test_api = mock.create_autospec(self.test_class)
+        test_api.test_fxn.side_effect = [
+            Aborted("aborted exception", errors=("Aborted error")),
+            "true",
+        ]
+        deadline = time.time() + 0.1
+        with self.assertRaises(Aborted):
+            _retry_on_aborted_exception(
+                functools.partial(test_api.test_fxn), deadline=deadline
+            )
+
+        self.assertEqual(test_api.test_fxn.call_count, 1)
+
+    def test_retry_on_aborted_exception_returns_response_after_internal_server_errors(
+        self,
+    ):
+        from google.api_core.exceptions import InternalServerError
+        import time
+        from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
+        import functools
+
+        test_api = mock.create_autospec(self.test_class)
+        test_api.test_fxn.side_effect = [
+            InternalServerError("testing"),
+            InternalServerError("testing"),
+            "true",
+        ]
+        allowed_exceptions = {
+            InternalServerError: lambda exc: None,
+        }
+        deadline = time.time() + 30
+        result_after_retries = _retry_on_aborted_exception(
+            functools.partial(test_api.test_fxn),
+            deadline=deadline,
+            allowed_exceptions=allowed_exceptions,
+        )
+
+        self.assertEqual(test_api.test_fxn.call_count, 3)
+        self.assertTrue(result_after_retries)
+
 
 class Test_metadata_with_leader_aware_routing(unittest.TestCase):
     def _call_fut(self, *args, **kw):

tests/unit/test_database.py

@@ -1931,12 +1931,13 @@ def test_context_mgr_w_commit_stats_error(self):
             return_commit_stats=True,
             request_options=RequestOptions(),
         )
-        api.commit.assert_called_once_with(
+        self.assertEqual(api.commit.call_count, 2)
+        api.commit.assert_any_call(
             request=request,
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
-            ],
+            ]
         )
 
         database.logger.info.assert_not_called()