All add much deeper tests to check for span statuses plus retries+abort

Also while here, fixed #1246

Author: odeke-em

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -142,6 +142,9 @@ def trace_call_end_lazily(
     ctx_manager.__enter__()
 
     def discard(exc_type=None, exc_value=None, exc_traceback=None):
+        if not exc_type:
+            span.set_status(Status(StatusCode.OK))
+
         ctx_manager.__exit__(exc_type, exc_value, exc_traceback)
 
     return discard
@@ -175,7 +178,12 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
             span.record_exception(error)
             raise
         else:
-            span.set_status(Status(StatusCode.OK))
+            if span._status.status_code == StatusCode.UNSET:
+                # OpenTelemetry-Python only allows a status change
+                # if the current code is UNSET or ERROR. At the end
+                # of the generator's consumption, only set it to OK
+                # it wasn't previously set otherwise
+                span.set_status(Status(StatusCode.OK))
 
 
 def set_span_status_error(span, error):
@@ -204,5 +212,5 @@ def add_event_on_current_span(event_name, attributes=None, span=None):
 
 def record_span_exception_and_status(span, exc):
     if span:
-        span.set_status(Status(StatusCode.ERROR, "foo"))
+        span.set_status(Status(StatusCode.ERROR, str(exc)))
         span.record_exception(exc)

google/cloud/spanner_v1/batch.py

@@ -50,7 +50,7 @@ class _BatchBase(_SessionWrapper):
     def __init__(self, session):
         super(_BatchBase, self).__init__(session)
         self._mutations = []
-        self.__discard_span = trace_call_end_lazily(
+        self.__base_discard_span = trace_call_end_lazily(
             f"CloudSpanner.{type(self).__name__}",
             self._session,
             None,
@@ -157,6 +157,11 @@ def delete(self, table, keyset):
             dict(table=table),
         )
 
+    def _discard_on_end(self, exc_type=None, exc_val=None, exc_traceback=None):
+        if self.__base_discard_span:
+            self.__base_discard_span(exc_type, exc_val, exc_traceback)
+            self.__base_discard_span = None
+
 
 class Batch(_BatchBase):
     """Accumulate mutations for transmission during :meth:`commit`."""
@@ -253,6 +258,7 @@ def commit(
             )
         self.committed = response.commit_timestamp
         self.commit_stats = response.commit_stats
+        self._discard_on_end()
         return self.committed
 
     def __enter__(self):
@@ -276,6 +282,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if self.__discard_span:
             self.__discard_span(exc_type, exc_val, exc_tb)
             self.__discard_span = None
+        self._discard_on_end()
 
 
 class MutationGroup(_BatchBase):

google/cloud/spanner_v1/session.py

@@ -180,7 +180,6 @@ def exists(self):
             )
 
         observability_options = getattr(self._database, "observability_options", None)
-        print(f"obsopts {observability_options}")
         with trace_call(
             "CloudSpanner.GetSession", self, observability_options=observability_options
         ) as span:

google/cloud/spanner_v1/snapshot.py

@@ -90,6 +90,7 @@ def _restart_on_unavailable(
         span_name, session, attributes, observability_options=observability_options
     ):
         iterator = method(request=request)
+
     while True:
         try:
             for item in iterator:

google/cloud/spanner_v1/transaction.py

@@ -206,6 +206,7 @@ def rollback(self):
                 )
         self.rolled_back = True
         del self._session._transaction
+        self._discard_on_end()
 
     def commit(
         self, return_commit_stats=False, request_options=None, max_commit_delay=None
@@ -286,6 +287,7 @@ def commit(
         if return_commit_stats:
             self.commit_stats = response.commit_stats
         del self._session._transaction
+        self._discard_on_end()
         return self.committed
 
     @staticmethod

tests/system/test_observability_options.py

@@ -137,3 +137,48 @@ def _make_credentials():
     from google.auth.credentials import AnonymousCredentials
 
     return AnonymousCredentials()
+
+
+from tests import _helpers as ot_helpers
+
+
+@pytest.mark.skipif(
+    not ot_helpers.HAS_OPENTELEMETRY_INSTALLED,
+    reason="Tracing requires OpenTelemetry",
+)
+def test_trace_call_keeps_span_error_status():
+    # Verifies that after our span's status was set to ERROR
+    # that it doesn't unconditionally get changed to OK
+    # per https://github.com/googleapis/python-spanner/issues/1246
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+        InMemorySpanExporter,
+    )
+    from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
+    from opentelemetry.trace.status import Status, StatusCode
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.sampling import ALWAYS_ON
+    from opentelemetry import trace
+
+    tracer_provider = TracerProvider(sampler=ALWAYS_ON)
+    trace_exporter = InMemorySpanExporter()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(trace_exporter))
+    observability_options = dict(tracer_provider=tracer_provider)
+
+    with trace_call(
+        "VerifyBehavior", observability_options=observability_options
+    ) as span:
+        span.set_status(Status(StatusCode.ERROR, "Our error exhibit"))
+
+    span_list = trace_exporter.get_finished_spans()
+    got_statuses = []
+
+    for span in span_list:
+        got_statuses.append(
+            (span.name, span.status.status_code, span.status.description)
+        )
+
+    want_statuses = [
+        ("VerifyBehavior", StatusCode.ERROR, "Our error exhibit"),
+    ]
+    assert got_statuses == want_statuses