tracing: spans are unconditionally set to a status of OK even after setting to ERROR after invoking `with trace_call` and that fails due to unconditional overriding because OpenTelemetry-Python can change the status after ERROR to OK

[odeke-em]

Noticed while writing extensive tests for modernized tracing functionality that we've got this code

python-spanner/google/cloud/spanner_v1/_opentelemetry_tracing.py

Lines 98 to 108 in ccae6e0

	with tracer.start_as_current_span(
	name, kind=trace.SpanKind.CLIENT, attributes=attributes
	) as span:
	try:
	yield span
	except Exception as error:
	span.set_status(Status(StatusCode.ERROR, str(error)))
	span.record_exception(error)
	raise
	else:
	span.set_status(Status(StatusCode.OK))

which doesn't do exactly what you expect with OpenTelemetry

Once yield span is invoked, if no exception was encountered, we unconditionally set the span's value to OK even if it previously were set to ERROR and here is a test

Reproduction

from tests import _helpers as ot_helpers
        
@pytest.mark.skipif(
    not ot_helpers.HAS_OPENTELEMETRY_INSTALLED,
    reason="Tracing requires OpenTelemetry",
)               
def test_trace_call_status():
    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

which would fail with

$ nox -s system-3.8 -- -k test_trace_call_status
>       assert got_statuses == want_statuses
E       AssertionError: assert [('VerifyBeha...OK: 1>, None)] == [('VerifyBeha...ror exhibit')]
E         
E         At index 0 diff: ('VerifyBehavior', <StatusCode.OK: 1>, None) != ('VerifyBehavior', <StatusCode.ERROR: 2>, 'Our error exhibit')
E         Use -v to get more diff

tests/system/test_observability_options.py:180: AssertionError

Fix

diff --git a/google/cloud/spanner_v1/_opentelemetry_tracing.py b/google/cloud/spanner_v1/_opentelemetry_tracing.py
index 9472251..65b3441 100644
--- a/google/cloud/spanner_v1/_opentelemetry_tracing.py
+++ b/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,8 +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))

/cc @harshachinta

[odeke-em]

I have a fix for it.