fix:Refactoring existing retry logic for aborted transactions and clean up redundant code

Author: aakashanandg

google/cloud/spanner_v1/_helpers.py

@@ -464,20 +464,43 @@ def _metadata_with_prefix(prefix, **kw):
     return [("google-cloud-resource-prefix", prefix)]
 
 
+def _retry_on_aborted_exception(
+    func,
+    deadline,
+    allowed_exceptions=None,
+):
+    """
+    Handles retry logic for Aborted exceptions, considering the deadline.
+    Retries the function in case of Aborted exceptions and other allowed exceptions.
+    """
+    attempts = 0
+    while True:
+        try:
+            attempts += 1
+            return func()
+        except Aborted as exc:
+            _delay_until_retry(exc, deadline=deadline, attempts=attempts)
+            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(
     func,
     retry_count=5,
     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.
@@ -491,21 +514,13 @@ def _retry(
         The result of the function if it is successful, or raises the last exception if all retries fail.
     """
     retries = 0
-    while True:
+    while retries <= retry_count:
         if retries > 0 and beforeNextRetry:
             beforeNextRetry(retries, delay)
 
         try:
             return func()
         except Exception as exc:
-            if isinstance(exc, Aborted) and deadline is not None:
-                if (
-                    allowed_exceptions is not None
-                    and allowed_exceptions.get(exc.__class__) is not None
-                ):
-                    retries += 1
-                    _delay_until_retry(exc, deadline=deadline, attempts=retries)
-                    continue
             if (
                 allowed_exceptions is None or exc.__class__ in allowed_exceptions
             ) and retries < retry_count:
@@ -568,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: no_op_handler,
                 },
                 deadline=deadline,
             )
@@ -360,16 +359,11 @@ def batch_write(
                 request=request,
                 metadata=metadata,
             )
-            deadline = time.time() + kwargs.get(
-                "timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
-            )
             response = _retry(
                 method,
                 allowed_exceptions={
                     InternalServerError: _check_rst_stream_error,
-                    Aborted: no_op_handler,
                 },
-                deadline=deadline,
             )
         self.committed = True
         return response
@@ -393,8 +387,3 @@ def _make_write_pb(table, columns, values):
     return Mutation.Write(
         table=table, columns=columns, values=_make_list_value_pbs(values)
     )
-
-
-def no_op_handler(exc):
-    # No-op (does nothing)
-    pass

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_batch.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 
-import time
 import unittest
 from unittest.mock import MagicMock
 from tests._helpers import (
@@ -651,36 +650,6 @@ def __init__(self, database=None, name=TestBatch.SESSION_NAME):
     def session_id(self):
         return self.name
 
-    def run_in_transaction(self, fnc):
-        """
-        Runs a function in a transaction, retrying if an exception occurs.
-        :param fnc: The function to run in the transaction.
-        :param max_retries: Maximum number of retry attempts.
-        :param delay: Delay (in seconds) between retries.
-        :return: The result of the function, or raises the exception after max retries.
-        """
-        from google.api_core.exceptions import Aborted
-
-        attempt = 0
-        max_retries = 3
-        delay = 1
-        while attempt < max_retries:
-            try:
-                result = fnc()
-                return result
-            except Aborted as exc:
-                attempt += 1
-                if attempt < max_retries:
-                    print(
-                        f"Attempt {attempt} failed with Aborted. Retrying in {delay} seconds..."
-                    )
-                    time.sleep(delay)  # Wait before retrying
-                else:
-                    raise exc  # After max retries, raise the exception
-            except Exception as exc:
-                print(f"Unexpected exception occurred: {exc}")
-                raise  # Raise any other unexpected exception immediately
-
 
 class _Database(object):
     name = "testing"

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()