feat: Multiplexed sessions - Add retry logic to run_in_transaction with previous transaction ID.

Signed-off-by: Taylor Curran <taylor.curran@improving.com>

Author: currantw

google/cloud/spanner_v1/session.py

@@ -464,7 +464,6 @@ def batch(self):
 
         return Batch(self)
 
-    # TODO multiplexed - deprecate
     def transaction(self) -> Transaction:
         """Create a transaction to perform a set of reads with shared staleness.
 
@@ -531,6 +530,12 @@ def run_in_transaction(self, func, *args, **kw):
         ) as span, MetricsCapture():
             attempts: int = 0
 
+            # If the transaction is retried after an aborted user operation, it should include the previous transaction ID
+            # in the transaction options used to begin the transaction. This allows the backend to recognize the transaction
+            # and increase the lock order for the new transaction ID that is created.
+            # See :attr:`~google.cloud.spanner_v1.types.TransactionOptions.ReadWrite.multiplexed_session_previous_transaction_id`
+            previous_transaction_id: Optional[bytes] = None
+
             while True:
                 # [A] Build transaction
                 # ---------------------
@@ -539,6 +544,7 @@ def run_in_transaction(self, func, *args, **kw):
                 txn.transaction_tag = transaction_tag
                 txn.exclude_txn_from_change_streams = exclude_txn_from_change_streams
                 txn.isolation_level = isolation_level
+                txn._previous_transaction_id = previous_transaction_id
 
                 # [B] Run user operation
                 # ----------------------
@@ -549,8 +555,8 @@ def run_in_transaction(self, func, *args, **kw):
                 try:
                     return_value = func(txn, *args, **kw)
 
-                # TODO multiplexed: store previous transaction ID.
                 except Aborted as exc:
+                    previous_transaction_id = txn._transaction_id
                     if span:
                         delay_seconds = _get_retry_delay(
                             exc.errors[0],
@@ -598,6 +604,7 @@ def run_in_transaction(self, func, *args, **kw):
                     )
 
                 except Aborted as exc:
+                    previous_transaction_id = txn._transaction_id
                     if span:
                         delay_seconds = _get_retry_delay(
                             exc.errors[0],
@@ -615,13 +622,15 @@ def run_in_transaction(self, func, *args, **kw):
                     _delay_until_retry(
                         exc, deadline, attempts, default_retry_delay=default_retry_delay
                     )
+
                 except GoogleAPICallError:
                     add_span_event(
                         span,
                         "Transaction.commit failed due to GoogleAPICallError, not retrying",
                         span_attributes,
                     )
                     raise
+
                 else:
                     if log_commit_stats and txn.commit_stats:
                         database.logger.info(

google/cloud/spanner_v1/transaction.py

@@ -70,6 +70,12 @@ def __init__(self, session):
         super(Transaction, self).__init__(session)
         self.rolled_back: bool = False
 
+        # If this transaction is used to retry a previous aborted transaction, the
+        # identifier for that transaction is used to increase the lock order of the new
+        # transaction (see :meth:`_build_transaction_options_pb`). This attribute should
+        # only be set by :meth:`~google.cloud.spanner_v1.session.Session.run_in_transaction`.
+        self._previous_transaction_id: Optional[bytes] = None
+
     def _build_transaction_options_pb(self) -> TransactionOptions:
         """Builds and returns transaction options for this transaction.
 
@@ -82,7 +88,9 @@ def _build_transaction_options_pb(self) -> TransactionOptions:
         )
 
         merge_transaction_options = TransactionOptions(
-            read_write=TransactionOptions.ReadWrite(),
+            read_write=TransactionOptions.ReadWrite(
+                multiplexed_session_previous_transaction_id=self._previous_transaction_id
+            ),
             exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
             isolation_level=self.isolation_level,
         )