Repository: googleapis/python-spanner-django

Issue:
dbapi: when retrying or discarding transactions, maintain a full stack of all statements
Raised by @bvandiver in https://github.com/orijtech/django-spanner/commit/4f32b239486f464b64b69904aefc4ba00e6278d7#r37546752 after my retry logic PR https://github.com/orijtech/django-spanner/pull/292.

The issue here is that on encountering the Aborted exception, we are erroneously retrying the currently executing statement ONLY, but within an aborted Transaction, there could multiple prior statements.

The fix for this is to maintain a FIFO stack of prior (sql, args,) until commit() or rollback() are executed, or any exception or until we retry.


Discussion:

Resolving patch:
diff --git a/spanner/dbapi/connection.py b/spanner/dbapi/connection.py
--- a/spanner/dbapi/connection.py
+++ b/spanner/dbapi/connection.py
@@ -16,6 +16,7 @@ def __init__(self, db_handle, session, discard_session):
         self.__txn = None
         self.__dbhandle = db_handle
         self.__closed = False
+        self.__on_transaction_clean_up = None
         self.__ddl_statements = []
 
     def __raise_if_already_closed(self):
@@ -56,7 +57,13 @@ def __can_commit_or_rollback(self):
         #   https://github.com/googleapis/python-spanner/issues/13
         return self.__txn and not (self.__txn.committed or self.__txn._rolled_back)
 
+    def __clean_up_transaction_state(self):
+        if self.__on_transaction_clean_up:
+            self.__on_transaction_clean_up()
+
     def commit(self):
+        self.__clean_up_transaction_state()
+
         if self.__can_commit_or_rollback():
             res = self.__txn.commit()
             self.__txn = None
@@ -65,6 +72,8 @@ def commit(self):
             self.__txn.stop()
 
     def rollback(self):
+        self.__clean_up_transaction_state()
+
         if self.__can_commit_or_rollback():
             res = self.__txn.rollback()
             self.__txn = None
@@ -75,7 +84,9 @@ def rollback(self):
     def cursor(self):
         self.__raise_if_already_closed()
 
-        return Cursor(self)
+        cur = Cursor(self)
+        self.__on_transaction_clean_up = cur._clear_transaction_state
+        return cur
 
     def discard_aborted_txn(self):
         # Discard the prior, now bad transaction.
diff --git a/spanner/dbapi/cursor.py b/spanner/dbapi/cursor.py
--- a/spanner/dbapi/cursor.py
+++ b/spanner/dbapi/cursor.py
@@ -4,6 +4,8 @@
 # license that can be found in the LICENSE file or at
 # https://developers.google.com/open-source/licenses/bsd
 
+import time
+
 import google.api_core.exceptions as grpc_exceptions
 
 from .exceptions import (
@@ -26,6 +28,7 @@ def __init__(self, db_handle):
         self.__connection = db_handle
         self.__last_op = None
         self.__closed = False
+        self.__sql_in_same_txn = []
 
         # arraysize is a readable and writable property mandated
         # by PEP-0249 https://www.python.org/dev/peps/pep-0249/#arraysize
@@ -64,7 +67,7 @@ def __discard_aborted_txn(self):
     def __get_txn(self):
         return self.__connection.get_txn()
 
-    def execute(self, sql, args=None, current_retry=0):
+    def execute(self, sql, args=None, already_in_retry=False):
         """
         Abstracts and implements execute SQL statements on Cloud Spanner.
         If it encounters grpc_exceptions.Aborted error, it optimistically retries
@@ -102,21 +105,79 @@ def execute(self, sql, args=None, current_retry=0):
                 self.__handle_insert(self.__get_txn(), sql, args or None)
             else:
                 self.__handle_update(self.__get_txn(), sql, args or None)
+
+        except grpc_exceptions.InvalidArgument as e:  # We can't retry a syntax issue, fail fast.
+            self.__discard_aborted_txn()
+            raise ProgrammingError(e.details if hasattr(e, 'details') else e)
+
         except (grpc_exceptions.AlreadyExists, grpc_exceptions.FailedPrecondition) as e:
+            # We can't retry an integrity error within the same transaction regardless.
+            self.__discard_aborted_txn()
             raise IntegrityError(e.details if hasattr(e, 'details') else e)
+
+        except Exception as e:
+            # Firstly discard the aborted transaction.
+            self.__discard_aborted_txn()
+
+            if already_in_retry:  # It is already being retried, so return immediately.
+                raise e
+
+            # Attempt to replay all the prior sql within the same transaction.
+            sql_args_tuples = self.__sql_in_same_txn[:]
+            sql_args_tuples.append((sql, args,))
+
+            return self.__replay_all_prior_statements_in_transaction(sql_args_tuples)
+        else:  # No error here
+            self.__sql_in_same_txn.append((sql, args,))
+
+    def _clear_transaction_state(self):
+        """
+        Invoked on every Connection.commit() or Connection.rollback()
+        """
+        if self.__sql_in_same_txn:
+            self.__sql_in_same_txn.clear()
+
+    def __replay_all_prior_statements_in_transaction(self, sql_args_tuples):
+        if not sql_args_tuples:
+            return
+
+        lastException = None
+
+        for i in range(5):
+            # Clean up before attempting the replay.
+            self.__sql_in_same_txn.clear()
+
+            print("\033[31mAttempting transaction replay #%d with elements:\n%s\033[00m" % (i, sql_args_tuples))
+
+            for sql, args in sql_args_tuples:
+                try:
+                    self.execute(sql, args, already_in_retry=True)
+                except grpc_exceptions.InvalidArgument as e:  # We can't retry a syntax issue, fail fast.
+                    raise ProgrammingError(e.details if hasattr(e, 'details') else e)
+                except (grpc_exceptions.AlreadyExists, grpc_exceptions.FailedPrecondition) as e:
+                    raise IntegrityError(e.details if hasattr(e, 'details') else e)
+                except Exception as e:
+                    lastException = e
+                    # TODO: Use exponential backoff with jitter, before retrying.
+                    time.sleep(0.57)
+                    break
+            else:
+                # All the elements in sql_args_tuples were executed,
+                # thus we can now break out of the retry loop.
+                # But first, reset all the executed (sql, args) for future replay.
+                self.__sql_in_same_txn = sql_args_tuples[:]
+                break
+
+        try:
+            if lastException:
+                self.__discard_aborted_txn()
+                raise lastException
         except grpc_exceptions.InvalidArgument as e:
             raise ProgrammingError(e.details if hasattr(e, 'details') else e)
         except grpc_exceptions.InternalServerError as e:
             raise OperationalError(e.details if hasattr(e, 'details') else e)
         except grpc_exceptions.Aborted as e:
-            if current_retry > 2:  # Arbitrary limit that should probably be a setting.
-                raise InternalError(e.details if hasattr(e, 'details') else e)
-
-            self.__discard_aborted_txn()
-            # Otherwise retry it.
-            print('\033[31mRetrying execution #%d of:\nSQL: %s\nArgs: %s\033[00m' % (
-                current_retry, sql, args))
-            return self.execute(sql, args, current_retry=current_retry+1)
+            raise InternalError(e.details if hasattr(e, 'details') else e)
 
     def __handle_update(self, txn, sql, params, param_types=None):
         sql = ensure_where_clause(sql)


Test patch: