From 0c9bff3d7f0daa973c8bf8edcf421f6d83185644 Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Wed, 20 Nov 2024 02:28:50 -0800
Subject: [PATCH] Database.run_in_transaction: add span_events for using
 Transaction and aborted

---
 google/cloud/spanner_v1/session.py | 41 ++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py
index f4b548f3b1..31792b141f 100644
--- a/google/cloud/spanner_v1/session.py
+++ b/google/cloud/spanner_v1/session.py
@@ -424,9 +424,11 @@ def run_in_transaction(self, func, *args, **kw):
         observability_options = getattr(self._database, "observability_options", None)
         with trace_call(
             "CloudSpanner.ReadWriteTransaction", self, observability_options
-        ):
+        ) as span:
             while True:
                 if self._transaction is None:
+                    if span:
+                        span.add_event("Creating Transaction")
                     txn = self.transaction()
                     txn.transaction_tag = transaction_tag
                     txn.exclude_txn_from_change_streams = (
@@ -435,17 +437,38 @@ def run_in_transaction(self, func, *args, **kw):
                 else:
                     txn = self._transaction
 
+                attempts += 1
+
+                span_attributes = {"transaction.id": txn.id, "attempt": attempts}
+                if span:
+                    span.add_event("Using Transaction", span_attributes)
+
                 try:
-                    attempts += 1
                     return_value = func(txn, *args, **kw)
                 except Aborted as exc:
                     del self._transaction
+                    if span:
+                        delay_seconds = _get_retry_delay(exc, attempts)
+                        attributes = dict(delay_seconds=delay_seconds)
+                        attributes.update(span_attributes)
+                        span.add_event("Transaction was aborted, retrying", attributes)
+
                     _delay_until_retry(exc, deadline, attempts)
                     continue
                 except GoogleAPICallError:
                     del self._transaction
+                    if span:
+                        span.add_event(
+                            "Transaction.commit failed due to GoogleAPICallError, not retrying",
+                            span_attributes,
+                        )
                     raise
                 except Exception:
+                    if span:
+                        span.add_event(
+                            "Invoking Transaction.rollback(), not retrying",
+                            span_attributes,
+                        )
                     txn.rollback()
                     raise
 
@@ -457,9 +480,23 @@ def run_in_transaction(self, func, *args, **kw):
                     )
                 except Aborted as exc:
                     del self._transaction
+                    if span:
+                        delay_seconds = _get_retry_delay(exc, attempts)
+                        attributes = dict(delay_seconds=delay_seconds)
+                        attributes.update(span_attributes)
+                        span.add_event(
+                            "Transaction.commit was aborted, retrying afresh",
+                            attributes,
+                        )
+
                     _delay_until_retry(exc, deadline, attempts)
                 except GoogleAPICallError:
                     del self._transaction
+                    if span:
+                        span.add_event(
+                            "Transaction.commit failed due to GoogleAPICallError, not retrying",
+                            span_attributes,
+                        )
                     raise
                 else:
                     if self._database.log_commit_stats and txn.commit_stats: