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

Author: aakashanandg

google/cloud/spanner_v1/_helpers.py

@@ -464,6 +464,22 @@ def _metadata_with_prefix(prefix, **kw):
     return [("google-cloud-resource-prefix", prefix)]
 
 
+def _retry_on_aborted_exception(exc, deadline, attempts, allowed_exceptions):
+    """
+    Handles the retry logic for Aborted exceptions, considering the deadline.
+    Returns True if the exception is retried, False otherwise.
+    """
+    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
+        ):
+            _delay_until_retry(exc, deadline=deadline, attempts=attempts)
+            return True
+    return False
+
+
 def _retry(
     func,
     retry_count=5,
@@ -491,21 +507,19 @@ 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.")
         if retries > 0 and beforeNextRetry:
             beforeNextRetry(retries, delay)
 
         try:
+            attempts += 1
             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 _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:

google/cloud/spanner_v1/batch.py

@@ -239,7 +239,7 @@ def commit(
                 method,
                 allowed_exceptions={
                     InternalServerError: _check_rst_stream_error,
-                    Aborted: no_op_handler,
+                    Aborted: lambda exc: None,
                 },
                 deadline=deadline,
             )
@@ -360,16 +360,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
@@ -394,7 +389,3 @@ def _make_write_pb(table, columns, values):
         table=table, columns=columns, values=_make_list_value_pbs(values)
     )
 
-
-def no_op_handler(exc):
-    # No-op (does nothing)
-    pass

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"