Set status for ended spans

Fixes #292

Author: ocelotl

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

@@ -122,10 +122,12 @@ class SpanKind(enum.Enum):
     https://github.com/open-telemetry/opentelemetry-specification/pull/226.
     """
 
-    #: Default value. Indicates that the span is used internally in the application.
+    #: Default value. Indicates that the span is used internally in the
+    # application.
     INTERNAL = 0
 
-    #: Indicates that the span describes an operation that handles a remote request.
+    #: Indicates that the span describes an operation that handles a remote
+    # request.
     SERVER = 1
 
     #: Indicates that the span describes a request to some remote service.
@@ -224,6 +226,7 @@ def __exit__(
         exc_tb: typing.Optional[python_types.TracebackType],
     ) -> None:
         """Ends context manager and calls `end` on the `Span`."""
+
         self.end()
 
 
@@ -392,6 +395,7 @@ def start_span(
         attributes: typing.Optional[types.Attributes] = None,
         links: typing.Sequence[Link] = (),
         start_time: typing.Optional[int] = None,
+        auto_update_status: bool = True,
     ) -> "Span":
         """Starts a span.
 
@@ -423,6 +427,8 @@ def start_span(
             attributes: The span's attributes.
             links: Links span to other spans
             start_time: Sets the start time of a span
+            auto_update_status: Defines if the status should be updated
+                automatically when the span finishes
 
         Returns:
             The newly-created span.

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

@@ -163,14 +163,14 @@ class Status:
 
     def __init__(
         self,
-        canonical_code: "StatusCanonicalCode" = StatusCanonicalCode.OK,
+        canonical_code: StatusCanonicalCode = StatusCanonicalCode.OK,
         description: typing.Optional[str] = None,
     ):
         self._canonical_code = canonical_code
         self._description = description
 
     @property
-    def canonical_code(self) -> "StatusCanonicalCode":
+    def canonical_code(self) -> StatusCanonicalCode:
         """Represents the canonical status code of a finished Span."""
         return self._canonical_code
 
@@ -183,3 +183,25 @@ def description(self) -> typing.Optional[str]:
     def is_ok(self) -> bool:
         """Returns false if this represents an error, true otherwise."""
         return self._canonical_code is StatusCanonicalCode.OK
+
+
+class StatusError(Exception):
+    def __init__(
+        self,
+        canonical_code: StatusCanonicalCode,
+        description: typing.Optional[str] = None,
+    ):
+        self._canonical_code = canonical_code
+        self._description = description
+
+        super(StatusError, self).__init__()
+
+    @property
+    def canonical_code(self) -> StatusCanonicalCode:
+        """Represents the canonical status code of a finished Span."""
+        return self._canonical_code
+
+    @property
+    def description(self) -> typing.Optional[str]:
+        """Status description"""
+        return self._description

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

@@ -18,13 +18,15 @@
 import random
 import threading
 from contextlib import contextmanager
-from typing import Iterator, Optional, Sequence, Tuple
+from types import TracebackType
+from typing import Iterator, Optional, Sequence, Tuple, Type
 
 from opentelemetry import trace as trace_api
 from opentelemetry.context import Context
 from opentelemetry.sdk import util
 from opentelemetry.sdk.util import BoundedDict, BoundedList
 from opentelemetry.trace import sampling
+from opentelemetry.trace.status import Status, StatusCanonicalCode, StatusError
 from opentelemetry.util import time_ns, types
 
 logger = logging.getLogger(__name__)
@@ -134,6 +136,7 @@ def __init__(
         links: Sequence[trace_api.Link] = (),
         kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
         span_processor: SpanProcessor = SpanProcessor(),
+        auto_update_status: bool = True,
     ) -> None:
 
         self.name = name
@@ -143,9 +146,10 @@ def __init__(
         self.trace_config = trace_config
         self.resource = resource
         self.kind = kind
+        self._auto_update_status = auto_update_status
 
         self.span_processor = span_processor
-        self.status = trace_api.Status()
+        self.status = None
         self._lock = threading.Lock()
 
         if attributes is None:
@@ -174,7 +178,10 @@ def __repr__(self):
         )
 
     def __str__(self):
-        return '{}(name="{}", context={}, kind={}, parent={}, start_time={}, end_time={})'.format(
+        return (
+            '{}(name="{}", context={}, kind={}, '
+            "parent={}, start_time={}, end_time={})"
+        ).format(
             type(self).__name__,
             self.name,
             self.context,
@@ -275,6 +282,30 @@ def set_status(self, status: trace_api.Status) -> None:
             return
         self.status = status
 
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Ends context manager and calls `end` on the `Span`."""
+
+        if self.status is None and self._auto_update_status:
+            if exc_val is not None:
+
+                if isinstance(exc_val, StatusError):
+                    self.set_status(
+                        Status(exc_val.canonical_code, exc_val.description)
+                    )
+
+                else:
+                    self.set_status(Status(StatusCanonicalCode.UNKNOWN))
+
+            else:
+                self.set_status(Status(StatusCanonicalCode.OK))
+
+        super().__exit__(exc_type, exc_val, exc_tb)
+
 
 def generate_span_id() -> int:
     """Get a new random span ID.
@@ -334,14 +365,15 @@ def start_as_current_span(
         span = self.start_span(name, parent, kind, attributes, links)
         return self.use_span(span, end_on_exit=True)
 
-    def start_span(
+    def start_span(  # pylint: disable=too-many-locals
         self,
         name: str,
         parent: trace_api.ParentSpan = trace_api.Tracer.CURRENT_SPAN,
         kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
         attributes: Optional[types.Attributes] = None,
         links: Sequence[trace_api.Link] = (),
         start_time: Optional[int] = None,
+        auto_update_status: bool = True,
     ) -> "Span":
         """See `opentelemetry.trace.Tracer.start_span`."""
 
@@ -401,6 +433,7 @@ def start_span(
                 span_processor=self._active_span_processor,
                 kind=kind,
                 links=links,
+                auto_update_status=auto_update_status,
             )
             span.start(start_time=start_time)
         else:

opentelemetry-sdk/tests/trace/test_trace.py

@@ -20,6 +20,7 @@
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk import trace
 from opentelemetry.trace import sampling
+from opentelemetry.trace.status import StatusCanonicalCode, StatusError
 from opentelemetry.util import time_ns
 
 
@@ -455,12 +456,7 @@ def test_start_span(self):
         span.start()
         self.assertEqual(start_time, span.start_time)
 
-        # default status
-        self.assertTrue(span.status.is_ok)
-        self.assertIs(
-            span.status.canonical_code, trace_api.status.StatusCanonicalCode.OK
-        )
-        self.assertIs(span.status.description, None)
+        self.assertIs(span.status, None)
 
         # status
         new_status = trace_api.status.Status(
@@ -515,13 +511,17 @@ def test_ended_span(self):
                 "Test description",
             )
             root.set_status(new_status)
-            # default status
-            self.assertTrue(root.status.is_ok)
-            self.assertEqual(
-                root.status.canonical_code,
-                trace_api.status.StatusCanonicalCode.OK,
-            )
-            self.assertIs(root.status.description, None)
+            self.assertIs(root.status, None)
+
+    def test_error_status(self):
+        with self.assertRaises(StatusError):
+            with trace.Tracer("test_error_status").start_span("root") as root:
+                raise StatusError(StatusCanonicalCode.CANCELLED, "cancelled")
+
+        self.assertIs(
+            root.status.canonical_code, StatusCanonicalCode.CANCELLED
+        )
+        self.assertEqual(root.status.description, "cancelled")
 
 
 def span_event_start_fmt(span_processor_name, span_name):

tox.ini

@@ -90,7 +90,7 @@ commands_pre =
   mypyinstalled: pip install file://{toxinidir}/opentelemetry-api/
 
 commands =
-  test: pytest
+  test: pytest {posargs}
   coverage: {toxinidir}/scripts/coverage.sh
 
   mypy: mypy --namespace-packages opentelemetry-api/src/opentelemetry/