Attach observability_options directly before trace_call

odeke-em · odeke-em · commit 3c3ed5a3da37 · 2024-11-11T23:32:35.000-08:00
diff --git a/google/cloud/spanner_v1/_opentelemetry_tracing.py b/google/cloud/spanner_v1/_opentelemetry_tracing.py
@@ -59,10 +59,6 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
 
     tracer_provider = None
     enable_extended_tracing = False
-    if observability_options is None and getattr(session, "_database", None):
-        observability_options = getattr(
-            session._database, "observability_options", None
-        )
 
     if observability_options:
         tracer_provider = observability_options.tracer_provider
@@ -85,7 +81,9 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
 
     # TODO(@odeke-em) enable after discussion with team and agreement
     # over extended tracing changes as the legacy default is always to
-    # record SQL statements on spans.
+    # record SQL statements on spans, because the prior behavior was
+    # to always record the SQL statement and changing it is 's considered
+    # a breaking change by the Python-Spanner team.
     if False and not enable_extended_tracing:
         attributes.pop("db.statement", False)
         attributes.pop("sql", False)
diff --git a/google/cloud/spanner_v1/batch.py b/google/cloud/spanner_v1/batch.py
@@ -205,7 +205,13 @@ def commit(
             max_commit_delay=max_commit_delay,
             request_options=request_options,
         )
-        with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
+        observability_options = getattr(database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.Commit",
+            self._session,
+            trace_attributes,
+            observability_options=observability_options,
+        ):
             method = functools.partial(
                 api.commit,
                 request=request,
@@ -318,7 +324,13 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             request_options=request_options,
             exclude_txn_from_change_streams=exclude_txn_from_change_streams,
         )
-        with trace_call("CloudSpanner.BatchWrite", self._session, trace_attributes):
+        observability_options = getattr(database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.BatchWrite",
+            self._session,
+            trace_attributes,
+            observability_options=observability_options,
+        ):
             method = functools.partial(
                 api.batch_write,
                 request=request,
diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py
@@ -142,7 +142,13 @@ def create(self):
         if self._labels:
             request.session.labels = self._labels
 
-        with trace_call("CloudSpanner.CreateSession", self, self._labels):
+        observability_options = getattr(self._database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.CreateSession",
+            self,
+            self._labels,
+            observability_options=observability_options,
+        ):
             session_pb = api.create_session(
                 request=request,
                 metadata=metadata,
@@ -169,7 +175,10 @@ def exists(self):
                 )
             )
 
-        with trace_call("CloudSpanner.GetSession", self) as span:
+        observability_options = getattr(self._database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.GetSession", self, observability_options=observability_options
+        ) as span:
             try:
                 api.get_session(name=self.name, metadata=metadata)
                 if span:
@@ -194,7 +203,12 @@ def delete(self):
             raise ValueError("Session ID not set by back-end")
         api = self._database.spanner_api
         metadata = _metadata_with_prefix(self._database.name)
-        with trace_call("CloudSpanner.DeleteSession", self):
+        observability_options = getattr(self._database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.DeleteSession",
+            self,
+            observability_options=observability_options,
+        ):
             api.delete_session(name=self.name, metadata=metadata)
 
     def ping(self):
diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py
@@ -84,7 +84,14 @@ def _restart_on_unavailable(
         )
 
     request.transaction = transaction_selector
-    with trace_call(trace_name, session, attributes):
+    observability_options = None
+    if session and session._database:
+        observability_options = getattr(
+            session._database, "observability_options", None
+        )
+    with trace_call(
+        trace_name, session, attributes, observability_options=observability_options
+    ):
         iterator = method(request=request)
     while True:
         try:
diff --git a/google/cloud/spanner_v1/transaction.py b/google/cloud/spanner_v1/transaction.py
@@ -110,7 +110,12 @@ def _execute_request(
         """
         transaction = self._make_txn_selector()
         request.transaction = transaction
-        with trace_call(trace_name, session, attributes):
+        observability_options = getattr(
+            self._session._database, "observability_options", None
+        )
+        with trace_call(
+            trace_name, session, attributes, observability_options=observability_options
+        ):
             method = functools.partial(method, request=request)
             response = _retry(
                 method,
@@ -147,7 +152,12 @@ def begin(self):
             read_write=TransactionOptions.ReadWrite(),
             exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
         )
-        with trace_call("CloudSpanner.BeginTransaction", self._session):
+        observability_options = getattr(database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.BeginTransaction",
+            self._session,
+            observability_options=observability_options,
+        ):
             method = functools.partial(
                 api.begin_transaction,
                 session=self._session.name,
@@ -175,7 +185,12 @@ def rollback(self):
                         database._route_to_leader_enabled
                     )
                 )
-            with trace_call("CloudSpanner.Rollback", self._session):
+            observability_options = getattr(database, "observability_options", None)
+            with trace_call(
+                "CloudSpanner.Rollback",
+                self._session,
+                observability_options=observability_options,
+            ):
                 method = functools.partial(
                     api.rollback,
                     session=self._session.name,
@@ -248,7 +263,13 @@ def commit(
             max_commit_delay=max_commit_delay,
             request_options=request_options,
         )
-        with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
+        observability_options = getattr(database, "observability_options", None)
+        with trace_call(
+            "CloudSpanner.Commit",
+            self._session,
+            trace_attributes,
+            observability_options=observability_options,
+        ):
             method = functools.partial(
                 api.commit,
                 request=request,
