8080from logging import getLogger
8181from typing import Iterator , Optional , Sequence , cast
8282
83+ from opentelemetry import context as context_api
8384from opentelemetry .context .context import Context
8485from opentelemetry .environment_variables import OTEL_PYTHON_TRACER_PROVIDER
8586from opentelemetry .trace .propagation import (
87+ SPAN_KEY ,
8688 get_current_span ,
8789 set_span_in_context ,
8890)
101103 format_span_id ,
102104 format_trace_id ,
103105)
104- from opentelemetry .trace .status import Status
106+ from opentelemetry .trace .status import Status , StatusCode
105107from opentelemetry .util import types
106108from opentelemetry .util ._providers import _load_provider
107109
@@ -324,7 +326,7 @@ def start_as_current_span(
324326 is equivalent to::
325327
326328 span = tracer.start_span(name)
327- with tracer .use_span(span, end_on_exit=True):
329+ with opentelemetry.trace .use_span(span, end_on_exit=True):
328330 do_work()
329331
330332 Args:
@@ -350,28 +352,6 @@ def start_as_current_span(
350352 The newly-created span.
351353 """
352354
353- @contextmanager # type: ignore
354- @abstractmethod
355- def use_span (
356- self , span : "Span" , end_on_exit : bool = False ,
357- ) -> Iterator [None ]:
358- """Context manager for setting the passed span as the
359- current span in the context, as well as resetting the
360- context back upon exiting the context manager.
361-
362- Set the given span as the current span in this tracer's context.
363-
364- On exiting the context manager set the span that was previously active
365- as the current span (this is usually but not necessarily the parent of
366- the given span). If ``end_on_exit`` is ``True``, then the span is also
367- ended when exiting the context manager.
368-
369- Args:
370- span: The span to start and make current.
371- end_on_exit: Whether to end the span automatically when leaving the
372- context manager.
373- """
374-
375355
376356class DefaultTracer (Tracer ):
377357 """The default Tracer, used when no Tracer implementation is available.
@@ -409,13 +389,6 @@ def start_as_current_span(
409389 # pylint: disable=unused-argument,no-self-use
410390 yield INVALID_SPAN
411391
412- @contextmanager # type: ignore
413- def use_span (
414- self , span : "Span" , end_on_exit : bool = False ,
415- ) -> Iterator [None ]:
416- # pylint: disable=unused-argument,no-self-use
417- yield
418-
419392
420393_TRACER_PROVIDER = None
421394
@@ -466,6 +439,44 @@ def get_tracer_provider() -> TracerProvider:
466439 return _TRACER_PROVIDER
467440
468441
442+ @contextmanager # type: ignore
443+ def use_span (
444+ span : Span ,
445+ end_on_exit : bool = False ,
446+ record_exception : bool = True ,
447+ set_status_on_exception : bool = True ,
448+ ) -> Iterator [Span ]:
449+ try :
450+ token = context_api .attach (context_api .set_value (SPAN_KEY , span ))
451+ try :
452+ yield span
453+ finally :
454+ context_api .detach (token )
455+
456+ except Exception as exc : # pylint: disable=broad-except
457+ # Record the exception as an event
458+ if isinstance (span , Span ) and span .is_recording ():
459+ if record_exception :
460+ span .record_exception (exc )
461+
462+ if set_status_on_exception :
463+ status = getattr (span , "_status" , None ) # type: Status
464+ if not status or status .status_code is StatusCode .UNSET :
465+ span .set_status (
466+ Status (
467+ status_code = StatusCode .ERROR ,
468+ description = "{}: {}" .format (
469+ type (exc ).__name__ , exc
470+ ),
471+ )
472+ )
473+ raise
474+
475+ finally :
476+ if end_on_exit :
477+ span .end ()
478+
479+
469480__all__ = [
470481 "DEFAULT_TRACE_OPTIONS" ,
471482 "DEFAULT_TRACE_STATE" ,
@@ -492,5 +503,6 @@ def get_tracer_provider() -> TracerProvider:
492503 "get_tracer_provider" ,
493504 "set_tracer_provider" ,
494505 "set_span_in_context" ,
506+ "use_span" ,
495507 "Status" ,
496508]
0 commit comments