Move use_span() method from Tracer to opentelemetry.trace.use_span()

Fixes: open-telemetry#1630

Author: owais

.github/workflows/test.yml

@@ -10,7 +10,7 @@ env:
   # Otherwise, set variable to the commit of your branch on
   # opentelemetry-python-contrib which is compatible with these Core repo
   # changes.
-  CONTRIB_REPO_SHA: 16ae58b341b960df52c1cf4541695164821b3638
+  CONTRIB_REPO_SHA: 8783e0ff97ad123006ff1ff2c2cf3f52161a406f
 
 jobs:
   build:

CHANGELOG.md

@@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add `max_attr_value_length` support to Jaeger exporter
   ([#1633])(https://github.com/open-telemetry/opentelemetry-python/pull/1633)
+- Moved `use_span` from Tracer to `opentelemetry.trace.use_span`.
+  ([#1668](https://github.com/open-telemetry/opentelemetry-python/pull/1668))
+- `opentelemetry.trace.use_span()` will now overwrite previously set status on span in case an
+  exception is raised inside the context manager and `set_status_on_exception` is set to `True`.
+  ([#1668](https://github.com/open-telemetry/opentelemetry-python/pull/1668))
 
 ### Changed
 - Rename `IdsGenerator` to `IdGenerator`

docs/examples/basic_context/async_context.py

@@ -24,7 +24,7 @@
 
 
 async def async_span(span):
-    with tracer.use_span(span):
+    with trace.use_span(span):
         ctx = baggage.set_baggage("foo", "bar")
     return ctx
 

opentelemetry-api/src/opentelemetry/trace/__init__.py

@@ -80,9 +80,11 @@
 from logging import getLogger
 from typing import Iterator, Optional, Sequence, cast
 
+from opentelemetry import context as context_api
 from opentelemetry.context.context import Context
 from opentelemetry.environment_variables import OTEL_PYTHON_TRACER_PROVIDER
 from opentelemetry.trace.propagation import (
+    SPAN_KEY,
     get_current_span,
     set_span_in_context,
 )
@@ -101,7 +103,7 @@
     format_span_id,
     format_trace_id,
 )
-from opentelemetry.trace.status import Status
+from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util import types
 from opentelemetry.util._providers import _load_provider
 
@@ -324,7 +326,7 @@ def start_as_current_span(
         is equivalent to::
 
             span = tracer.start_span(name)
-            with tracer.use_span(span, end_on_exit=True):
+            with opentelemetry.trace.use_span(span, end_on_exit=True):
                 do_work()
 
         Args:
@@ -350,28 +352,6 @@ def start_as_current_span(
             The newly-created span.
         """
 
-    @contextmanager  # type: ignore
-    @abstractmethod
-    def use_span(
-        self, span: "Span", end_on_exit: bool = False,
-    ) -> Iterator[None]:
-        """Context manager for setting the passed span as the
-        current span in the context, as well as resetting the
-        context back upon exiting the context manager.
-
-        Set the given span as the current span in this tracer's context.
-
-        On exiting the context manager set the span that was previously active
-        as the current span (this is usually but not necessarily the parent of
-        the given span). If ``end_on_exit`` is ``True``, then the span is also
-        ended when exiting the context manager.
-
-        Args:
-            span: The span to start and make current.
-            end_on_exit: Whether to end the span automatically when leaving the
-                context manager.
-        """
-
 
 class DefaultTracer(Tracer):
     """The default Tracer, used when no Tracer implementation is available.
@@ -409,13 +389,6 @@ def start_as_current_span(
         # pylint: disable=unused-argument,no-self-use
         yield INVALID_SPAN
 
-    @contextmanager  # type: ignore
-    def use_span(
-        self, span: "Span", end_on_exit: bool = False,
-    ) -> Iterator[None]:
-        # pylint: disable=unused-argument,no-self-use
-        yield
-
 
 _TRACER_PROVIDER = None
 
@@ -466,6 +439,55 @@ def get_tracer_provider() -> TracerProvider:
     return _TRACER_PROVIDER
 
 
+@contextmanager  # type: ignore
+def use_span(
+    span: Span,
+    end_on_exit: bool = False,
+    record_exception: bool = True,
+    set_status_on_exception: bool = True,
+) -> Iterator[Span]:
+    """Takes a non-active span and activates it in the current context.
+
+    Args:
+        span: The span that should be activated in the current context.
+        end_on_exit: Whether to end the span automatically when leaving the
+            context manager.
+        record_exception: Whether to record any exceptions raised within the
+            context as error event on the span.
+        set_status_on_exception: Only relevant if the returned span is used
+            in a with/context manager. Defines wether the span status will
+            be automatically set to ERROR when an uncaught exception is
+            raised in the span with block. The span status won't be set by
+            this mechanism if it was previously set manually.
+    """
+    try:
+        token = context_api.attach(context_api.set_value(SPAN_KEY, span))
+        try:
+            yield span
+        finally:
+            context_api.detach(token)
+
+    except Exception as exc:  # pylint: disable=broad-except
+        if isinstance(span, Span) and span.is_recording():
+            # Record the exception as an event
+            if record_exception:
+                span.record_exception(exc)
+
+            # Set status in case exception was raised
+            if set_status_on_exception:
+                span.set_status(
+                    Status(
+                        status_code=StatusCode.ERROR,
+                        description="{}: {}".format(type(exc).__name__, exc),
+                    )
+                )
+        raise
+
+    finally:
+        if end_on_exit:
+            span.end()
+
+
 __all__ = [
     "DEFAULT_TRACE_OPTIONS",
     "DEFAULT_TRACE_STATE",
@@ -492,5 +514,6 @@ def get_tracer_provider() -> TracerProvider:
     "get_tracer_provider",
     "set_tracer_provider",
     "set_span_in_context",
+    "use_span",
     "Status",
 ]

opentelemetry-api/src/opentelemetry/trace/propagation/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from typing import Optional
 
-from opentelemetry.context import get_value, set_value
+from opentelemetry.context import attach, get_value, set_value
 from opentelemetry.context.context import Context
 from opentelemetry.trace.span import INVALID_SPAN, Span
 

opentelemetry-api/tests/trace/test_globals.py

@@ -2,6 +2,27 @@
 from unittest.mock import patch
 
 from opentelemetry import context, trace
+from opentelemetry.trace.status import Status, StatusCode
+
+
+class TestSpan(trace.NonRecordingSpan):
+    has_ended = False
+    recorded_exception = None
+    recorded_status = Status(status_code=StatusCode.UNSET)
+
+    def set_status(self, status):
+        self.recorded_status = status
+
+    def end(self, end_time=None):
+        self.has_ended = True
+
+    def is_recording(self):
+        return not self.has_ended
+
+    def record_exception(
+        self, exception, attributes=None, timestamp=None, escaped=False
+    ):
+        self.recorded_exception = exception
 
 
 class TestGlobals(unittest.TestCase):
@@ -38,3 +59,53 @@ def test_get_current_span(self):
         finally:
             context.detach(token)
         self.assertEqual(trace.get_current_span(), trace.INVALID_SPAN)
+
+
+class TestUseTracer(unittest.TestCase):
+    def test_use_span(self):
+        self.assertEqual(trace.get_current_span(), trace.INVALID_SPAN)
+        span = trace.NonRecordingSpan(trace.INVALID_SPAN_CONTEXT)
+        with trace.use_span(span):
+            self.assertIs(trace.get_current_span(), span)
+        self.assertEqual(trace.get_current_span(), trace.INVALID_SPAN)
+
+    def test_use_span_end_on_exit(self):
+
+        test_span = TestSpan(trace.INVALID_SPAN_CONTEXT)
+
+        with trace.use_span(test_span):
+            pass
+        self.assertFalse(test_span.has_ended)
+
+        with trace.use_span(test_span, end_on_exit=True):
+            pass
+        self.assertTrue(test_span.has_ended)
+
+    def test_use_span_exception(self):
+        class TestUseSpanException(Exception):
+            pass
+
+        test_span = TestSpan(trace.INVALID_SPAN_CONTEXT)
+        exception = TestUseSpanException("test exception")
+        with self.assertRaises(TestUseSpanException):
+            with trace.use_span(test_span):
+                raise exception
+
+        self.assertEqual(test_span.recorded_exception, exception)
+
+    def test_use_span_set_status(self):
+        class TestUseSpanException(Exception):
+            pass
+
+        test_span = TestSpan(trace.INVALID_SPAN_CONTEXT)
+        with self.assertRaises(TestUseSpanException):
+            with trace.use_span(test_span):
+                raise TestUseSpanException("test error")
+
+        self.assertEqual(
+            test_span.recorded_status.status_code, StatusCode.ERROR
+        )
+        self.assertEqual(
+            test_span.recorded_status.description,
+            "TestUseSpanException: test error",
+        )

opentelemetry-api/tests/trace/test_tracer.py

@@ -29,11 +29,6 @@ def test_start_as_current_span(self):
         with self.tracer.start_as_current_span("") as span:
             self.assertIsInstance(span, trace.Span)
 
-    def test_use_span(self):
-        span = trace.NonRecordingSpan(trace.INVALID_SPAN_CONTEXT)
-        with self.tracer.use_span(span):
-            pass
-
     def test_get_current_span(self):
         with self.tracer.start_as_current_span("test") as span:
             trace.get_current_span().set_attribute("test", "test")

opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

@@ -776,10 +776,7 @@ def __exit__(
                 self.record_exception(exception=exc_val, escaped=True)
             # Records status if span is used as context manager
             # i.e. with tracer.start_span() as span:
-            if (
-                self._status.status_code is StatusCode.UNSET
-                and self._set_status_on_exception
-            ):
+            if self._set_status_on_exception:
                 self.set_status(
                     Status(
                         status_code=StatusCode.ERROR,
@@ -869,7 +866,12 @@ def start_as_current_span(
             record_exception=record_exception,
             set_status_on_exception=set_status_on_exception,
         )
-        with self.use_span(span, end_on_exit=end_on_exit) as span_context:
+        with trace_api.use_span(
+            span,
+            end_on_exit=end_on_exit,
+            record_exception=record_exception,
+            set_status_on_exception=set_status_on_exception,
+        ) as span_context:
             yield span_context
 
     def start_span(  # pylint: disable=too-many-locals
@@ -950,44 +952,6 @@ def start_span(  # pylint: disable=too-many-locals
             span = trace_api.NonRecordingSpan(context=span_context)
         return span
 
-    @contextmanager
-    def use_span(
-        self, span: trace_api.Span, end_on_exit: bool = False,
-    ) -> Iterator[trace_api.Span]:
-        try:
-            token = context_api.attach(context_api.set_value(SPAN_KEY, span))
-            try:
-                yield span
-            finally:
-                context_api.detach(token)
-
-        except Exception as exc:  # pylint: disable=broad-except
-            # Record the exception as an event
-            if isinstance(span, Span) and span.is_recording():
-                # pylint:disable=protected-access
-                if span._record_exception:
-                    span.record_exception(exc)
-
-                # Records status if use_span is used
-                # i.e. with tracer.start_as_current_span() as span:
-                if (
-                    span._status.status_code is StatusCode.UNSET
-                    and span._set_status_on_exception
-                ):
-                    span.set_status(
-                        Status(
-                            status_code=StatusCode.ERROR,
-                            description="{}: {}".format(
-                                type(exc).__name__, exc
-                            ),
-                        )
-                    )
-            raise
-
-        finally:
-            if end_on_exit:
-                span.end()
-
 
 class TracerProvider(trace_api.TracerProvider):
     """See `opentelemetry.trace.TracerProvider`."""

opentelemetry-sdk/tests/trace/test_trace.py

@@ -122,18 +122,6 @@ def run_general_code(shutdown_on_exit, explicit_shutdown):
         out = run_general_code(False, False)
         self.assertTrue(out.startswith(b"0"))
 
-    def test_use_span_exception(self):
-        class TestUseSpanException(Exception):
-            pass
-
-        default_span = trace_api.NonRecordingSpan(
-            trace_api.INVALID_SPAN_CONTEXT
-        )
-        tracer = new_tracer()
-        with self.assertRaises(TestUseSpanException):
-            with tracer.use_span(default_span):
-                raise TestUseSpanException()
-
     def test_tracer_provider_accepts_concurrent_multi_span_processor(self):
         span_processor = trace.ConcurrentMultiSpanProcessor(2)
         tracer_provider = trace.TracerProvider(
@@ -307,7 +295,7 @@ def test_start_span_implicit(self):
         self.assertIsNone(root.end_time)
         self.assertEqual(root.kind, trace_api.SpanKind.INTERNAL)
 
-        with tracer.use_span(root, True):
+        with trace_api.use_span(root, True):
             self.assertIs(trace_api.get_current_span(), root)
 
             with tracer.start_span(
@@ -364,7 +352,7 @@ def test_start_span_explicit(self):
         self.assertIsNone(root.end_time)
 
         # Test with the implicit root span
-        with tracer.use_span(root, True):
+        with trace_api.use_span(root, True):
             self.assertIs(trace_api.get_current_span(), root)
 
             with tracer.start_span("stepchild", other_parent_context) as child:
@@ -947,7 +935,7 @@ def error_status_test(context):
             .start_as_current_span("root")
         )
 
-    def test_override_error_status(self):
+    def test_last_status_wins(self):
         def error_status_test(context):
             with self.assertRaises(AssertionError):
                 with context as root:
@@ -956,8 +944,10 @@ def error_status_test(context):
                     )
                     raise AssertionError("unknown")
 
-            self.assertIs(root.status.status_code, StatusCode.OK)
-            self.assertEqual(root.status.description, "OK")
+            self.assertIs(root.status.status_code, StatusCode.ERROR)
+            self.assertEqual(
+                root.status.description, "AssertionError: unknown"
+            )
 
         error_status_test(
             trace.TracerProvider().get_tracer(__name__).start_span("root")