Revert changes that had crept in from other PR

Author: odeke-em

google/cloud/spanner_v1/database.py

@@ -53,10 +53,6 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
 )
-from google.cloud.spanner_v1._opentelemetry_tracing import (
-    get_current_span,
-    trace_call,
-)
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
 from google.cloud.spanner_v1.keyset import KeySet
@@ -71,6 +67,7 @@
     SpannerGrpcTransport,
 )
 from google.cloud.spanner_v1.table import Table
+from google.cloud.spanner_v1._opentelemetry_tracing import get_current_span
 
 
 SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
@@ -698,7 +695,7 @@ def execute_partitioned_dml(
             )
 
         def execute_pdml():
-            def do_execute_pdml(session):
+            with SessionCheckout(self._pool) as session:
                 txn = api.begin_transaction(
                     session=session.name, options=txn_options, metadata=metadata
                 )
@@ -730,15 +727,6 @@ def do_execute_pdml(session):
 
                 return result_set.stats.row_count_lower_bound
 
-            with SessionCheckout(self._pool) as session:
-                observability_options = getattr(self, "observability_options", None)
-                with trace_call(
-                    "CloudSpanner.execute_pdml",
-                    session,
-                    observability_options=observability_options,
-                ):
-                    return do_execute_pdml(session)
-
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
     def session(self, labels=None, database_role=None):
@@ -893,26 +881,17 @@ def run_in_transaction(self, func, *args, **kw):
         # Sanity check: Is there a transaction already running?
         # If there is, then raise a red flag. Otherwise, mark that this one
         # is running.
-        with SessionCheckout(self._pool) as session:
-            observability_options = getattr(self, "observability_options", None)
-            with trace_call(
-                "CloudSpanner.Database.run_in_transaction",
-                session,
-                observability_options=observability_options,
-            ):
-                # Sanity check: Is there a transaction already running?
-                # If there is, then raise a red flag. Otherwise, mark that this one
-                # is running.
-                if getattr(self._local, "transaction_running", False):
-                    raise RuntimeError("Spanner does not support nested transactions.")
-                self._local.transaction_running = True
-
-                # Check out a session and run the function in a transaction; once
-                # done, flip the sanity check bit back.
-                try:
-                    return session.run_in_transaction(func, *args, **kw)
-                finally:
-                    self._local.transaction_running = False
+        if getattr(self._local, "transaction_running", False):
+            raise RuntimeError("Spanner does not support nested transactions.")
+        self._local.transaction_running = True
+
+        # Check out a session and run the function in a transaction; once
+        # done, flip the sanity check bit back.
+        try:
+            with SessionCheckout(self._pool) as session:
+                return session.run_in_transaction(func, *args, **kw)
+        finally:
+            self._local.transaction_running = False
 
     def restore(self, source):
         """Restore from a backup to this database.
@@ -1189,7 +1168,7 @@ def __enter__(self):
         current_span = get_current_span()
         session = self._session = self._database._pool.get()
         if current_span:
-            current_span.add_event("Using session", {"id": self._session.session_id})
+            current_span.add_event("Using session", {"id": session.session_id})
         batch = self._batch = Batch(session)
         if self._request_options.transaction_tag:
             batch.transaction_tag = self._request_options.transaction_tag

google/cloud/spanner_v1/session.py

@@ -426,17 +426,23 @@ def run_in_transaction(self, func, *args, **kw):
         exclude_txn_from_change_streams = kw.pop(
             "exclude_txn_from_change_streams", None
         )
-
-        observability_options = getattr(self._database, "observability_options", None)
         attempts = 0
 
-        def __run_txn(txn, attempts):
+        while True:
+            if self._transaction is None:
+                txn = self.transaction()
+                txn.transaction_tag = transaction_tag
+                txn.exclude_txn_from_change_streams = exclude_txn_from_change_streams
+            else:
+                txn = self._transaction
+
             try:
+                attempts += 1
                 return_value = func(txn, *args, **kw)
             except Aborted as exc:
                 del self._transaction
                 _delay_until_retry(exc, deadline, attempts)
-                return None, False
+                continue
             except GoogleAPICallError:
                 del self._transaction
                 raise
@@ -462,28 +468,7 @@ def __run_txn(txn, attempts):
                         "CommitStats: {}".format(txn.commit_stats),
                         extra={"commit_stats": txn.commit_stats},
                     )
-                return return_value, True
-
-        while True:
-            if self._transaction is None:
-                with trace_call(
-                    "CloudSpanner.ReadWriteTransaction", self, observability_options
-                ):
-                    txn = self.transaction()
-                    txn.transaction_tag = transaction_tag
-                    txn.exclude_txn_from_change_streams = (
-                        exclude_txn_from_change_streams
-                    )
-                    return_value, completed = __run_txn(txn, attempts)
-                    if completed:
-                        return return_value
-            else:
-                txn = self._transaction
-                return_value, completed = __run_txn(txn, attempts)
-                if completed:
-                    return return_value
-
-            attempts += 1
+                return return_value
 
 
 # Rational:  this function factors out complex shared deadline / retry

google/cloud/spanner_v1/transaction.py

@@ -248,6 +248,31 @@ def commit(
             raise ValueError("Transaction is not begun")
 
         database = self._session._database
+        api = database.spanner_api
+        metadata = _metadata_with_prefix(database.name)
+        if database._route_to_leader_enabled:
+            metadata.append(
+                _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
+            )
+
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) is dict:
+            request_options = RequestOptions(request_options)
+        if self.transaction_tag is not None:
+            request_options.transaction_tag = self.transaction_tag
+        # Request tags are not supported for commit requests.
+        request_options.request_tag = None
+
+        request = CommitRequest(
+            session=self._session.name,
+            mutations=self._mutations,
+            transaction_id=self._transaction_id,
+            return_commit_stats=return_commit_stats,
+            max_commit_delay=max_commit_delay,
+            request_options=request_options,
+        )
+
         trace_attributes = {"num_mutations": len(self._mutations)}
         observability_options = getattr(database, "observability_options", None)
         with trace_call(
@@ -262,32 +287,6 @@ def commit(
             if span:
                 span.add_event("Starting Commit")
 
-            api = database.spanner_api
-            metadata = _metadata_with_prefix(database.name)
-            if database._route_to_leader_enabled:
-                metadata.append(
-                    _metadata_with_leader_aware_routing(
-                        database._route_to_leader_enabled
-                    )
-                )
-
-            if request_options is None:
-                request_options = RequestOptions()
-            elif type(request_options) is dict:
-                request_options = RequestOptions(request_options)
-            if self.transaction_tag is not None:
-                request_options.transaction_tag = self.transaction_tag
-            # Request tags are not supported for commit requests.
-            request_options.request_tag = None
-
-            request = CommitRequest(
-                session=self._session.name,
-                mutations=self._mutations,
-                transaction_id=self._transaction_id,
-                return_commit_stats=return_commit_stats,
-                max_commit_delay=max_commit_delay,
-                request_options=request_options,
-            )
             method = functools.partial(
                 api.commit,
                 request=request,