Implementing Propagators API to use Context

ext/opentelemetry-ext-flask/src/opentelemetry/ext/flask/__init__.py

@@ -6,15 +6,17 @@
 from flask import request as flask_request
 
 import opentelemetry.ext.wsgi as otel_wsgi
-from opentelemetry import propagators, trace
+from opentelemetry import context, propagators, trace
 from opentelemetry.ext.flask.version import __version__
+from opentelemetry.trace.propagation import get_span_from_context
 from opentelemetry.util import time_ns
 
 logger = logging.getLogger(__name__)
 
 _ENVIRON_STARTTIME_KEY = "opentelemetry-flask.starttime_key"
 _ENVIRON_SPAN_KEY = "opentelemetry-flask.span_key"
 _ENVIRON_ACTIVATION_KEY = "opentelemetry-flask.activation_key"
+_ENVIRON_TOKEN = "opentelemetry-flask.token"
 
 
 def instrument_app(flask):
@@ -57,8 +59,8 @@ def _before_flask_request():
     span_name = flask_request.endpoint or otel_wsgi.get_default_span_name(
         environ
     )
-    parent_span = propagators.extract(
-        otel_wsgi.get_header_from_environ, environ
+    token = context.attach(
+        propagators.extract(otel_wsgi.get_header_from_environ, environ)
     )
 
     tracer = trace.get_tracer(__name__, __version__)
@@ -69,7 +71,6 @@ def _before_flask_request():
         attributes["http.route"] = flask_request.url_rule.rule
     span = tracer.start_span(
         span_name,
-        parent_span,
         kind=trace.SpanKind.SERVER,
         attributes=attributes,
         start_time=environ.get(_ENVIRON_STARTTIME_KEY),
@@ -78,6 +79,7 @@ def _before_flask_request():
     activation.__enter__()
     environ[_ENVIRON_ACTIVATION_KEY] = activation
     environ[_ENVIRON_SPAN_KEY] = span
+    environ[_ENVIRON_TOKEN] = token
 
 
 def _teardown_flask_request(exc):
@@ -95,3 +97,4 @@ def _teardown_flask_request(exc):
         activation.__exit__(
             type(exc), exc, getattr(exc, "__traceback__", None)
         )
+    context.detach(flask_request.environ.get(_ENVIRON_TOKEN))

ext/opentelemetry-ext-http-requests/src/opentelemetry/ext/http_requests/__init__.py

@@ -76,7 +76,7 @@ def instrumented_request(self, method, url, *args, **kwargs):
             # to access propagators.
 
             headers = kwargs.setdefault("headers", {})
-            propagators.inject(tracer, type(headers).__setitem__, headers)
+            propagators.inject(type(headers).__setitem__, headers)
             result = wrapped(self, method, url, *args, **kwargs)  # *** PROCEED
 
             span.set_attribute("http.status_code", result.status_code)

ext/opentelemetry-ext-http-requests/tests/test_requests_integration.py

@@ -41,6 +41,7 @@ def setUp(self):
         self.get_tracer = self.get_tracer_patcher.start()
         self.span_context_manager = mock.MagicMock()
         self.span = mock.create_autospec(trace.Span, spec_set=True)
+        self.span.get_context.return_value = trace.INVALID_SPAN_CONTEXT
         self.span_context_manager.__enter__.return_value = self.span
 
         def setspanattr(key, value):

ext/opentelemetry-ext-opentracing-shim/src/opentelemetry/ext/opentracing_shim/__init__.py

@@ -93,6 +93,10 @@
 from opentelemetry.ext.opentracing_shim import util
 from opentelemetry.ext.opentracing_shim.version import __version__
 from opentelemetry.trace import DefaultSpan
+from opentelemetry.trace.propagation import (
+    get_span_from_context,
+    set_span_in_context,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -677,11 +681,8 @@ def inject(self, span_context, format, carrier):
 
         propagator = propagators.get_global_httptextformat()
 
-        propagator.inject(
-            DefaultSpan(span_context.unwrap()),
-            type(carrier).__setitem__,
-            carrier,
-        )
+        ctx = set_span_in_context(DefaultSpan(span_context.unwrap()))
+        propagator.inject(type(carrier).__setitem__, carrier, context=ctx)
 
     def extract(self, format, carrier):
         """Implements the ``extract`` method from the base class."""
@@ -700,6 +701,7 @@ def get_as_list(dict_object, key):
             return [value] if value is not None else []
 
         propagator = propagators.get_global_httptextformat()
-        otel_context = propagator.extract(get_as_list, carrier)
+        ctx = propagator.extract(get_as_list, carrier)
+        otel_context = get_span_from_context(ctx).get_context()
 
         return SpanContextShim(otel_context)

ext/opentelemetry-ext-opentracing-shim/tests/test_shim.py

@@ -16,15 +16,26 @@
 # pylint:disable=no-member
 
 import time
+import typing
 from unittest import TestCase
 
 import opentracing
 
 import opentelemetry.ext.opentracing_shim as opentracingshim
 from opentelemetry import propagators, trace
-from opentelemetry.context.propagation.httptextformat import HTTPTextFormat
+from opentelemetry.context import Context
 from opentelemetry.ext.opentracing_shim import util
 from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace.propagation import (
+    get_span_from_context,
+    set_span_in_context,
+)
+from opentelemetry.trace.propagation.httptextformat import (
+    Getter,
+    HTTPTextFormat,
+    HTTPTextFormatT,
+    Setter,
+)
 
 
 class TestShim(TestCase):
@@ -49,7 +60,7 @@ def setUpClass(cls):
         cls._previous_propagator = propagators.get_global_httptextformat()
 
         # Set mock propagator for testing.
-        propagators.set_global_httptextformat(MockHTTPTextFormat)
+        propagators.set_global_httptextformat(MockHTTPTextFormat())
 
     @classmethod
     def tearDownClass(cls):
@@ -541,23 +552,37 @@ class MockHTTPTextFormat(HTTPTextFormat):
     TRACE_ID_KEY = "mock-traceid"
     SPAN_ID_KEY = "mock-spanid"
 
-    @classmethod
-    def extract(cls, get_from_carrier, carrier):
-        trace_id_list = get_from_carrier(carrier, cls.TRACE_ID_KEY)
-        span_id_list = get_from_carrier(carrier, cls.SPAN_ID_KEY)
+    def extract(
+        self,
+        get_from_carrier: Getter[HTTPTextFormatT],
+        carrier: HTTPTextFormatT,
+        context: typing.Optional[Context] = None,
+    ) -> Context:
+        trace_id_list = get_from_carrier(carrier, self.TRACE_ID_KEY)
+        span_id_list = get_from_carrier(carrier, self.SPAN_ID_KEY)
 
         if not trace_id_list or not span_id_list:
-            return trace.INVALID_SPAN_CONTEXT
+            return set_span_in_context(trace.INVALID_SPAN)
 
-        return trace.SpanContext(
-            trace_id=int(trace_id_list[0]), span_id=int(span_id_list[0])
+        return set_span_in_context(
+            trace.DefaultSpan(
+                trace.SpanContext(
+                    trace_id=int(trace_id_list[0]),
+                    span_id=int(span_id_list[0]),
+                )
+            )
         )
 
-    @classmethod
-    def inject(cls, span, set_in_carrier, carrier):
+    def inject(
+        self,
+        set_in_carrier: Setter[HTTPTextFormatT],
+        carrier: HTTPTextFormatT,
+        context: typing.Optional[Context] = None,
+    ) -> None:
+        span = get_span_from_context(context)
         set_in_carrier(
-            carrier, cls.TRACE_ID_KEY, str(span.get_context().trace_id)
+            carrier, self.TRACE_ID_KEY, str(span.get_context().trace_id)
         )
         set_in_carrier(
-            carrier, cls.SPAN_ID_KEY, str(span.get_context().span_id)
+            carrier, self.SPAN_ID_KEY, str(span.get_context().span_id)
         )

ext/opentelemetry-ext-wsgi/src/opentelemetry/ext/wsgi/__init__.py

@@ -22,8 +22,9 @@
 import typing
 import wsgiref.util as wsgiref_util
 
-from opentelemetry import propagators, trace
+from opentelemetry import context, propagators, trace
 from opentelemetry.ext.wsgi.version import __version__
+from opentelemetry.trace.propagation import get_span_from_context
 from opentelemetry.trace.status import Status, StatusCanonicalCode
 
 _HTTP_VERSION_PREFIX = "HTTP/"
@@ -181,12 +182,13 @@ def __call__(self, environ, start_response):
             start_response: The WSGI start_response callable.
         """
 
-        parent_span = propagators.extract(get_header_from_environ, environ)
+        token = context.attach(
+            propagators.extract(get_header_from_environ, environ)
+        )
         span_name = get_default_span_name(environ)
 
         span = self.tracer.start_span(
             span_name,
-            parent_span,
             kind=trace.SpanKind.SERVER,
             attributes=collect_request_attributes(environ),
         )
@@ -197,17 +199,20 @@ def __call__(self, environ, start_response):
                     span, start_response
                 )
                 iterable = self.wsgi(environ, start_response)
-                return _end_span_after_iterating(iterable, span, self.tracer)
+                return _end_span_after_iterating(
+                    iterable, span, self.tracer, token
+                )
         except:  # noqa
             # TODO Set span status (cf. https://github.com/open-telemetry/opentelemetry-python/issues/292)
             span.end()
+            context.detach(token)
             raise
 
 
 # Put this in a subfunction to not delay the call to the wrapped
 # WSGI application (instrumentation should change the application
 # behavior as little as possible).
-def _end_span_after_iterating(iterable, span, tracer):
+def _end_span_after_iterating(iterable, span, tracer, token):
     try:
         with tracer.use_span(span):
             for yielded in iterable:
@@ -217,3 +222,4 @@ def _end_span_after_iterating(iterable, span, tracer):
         if close:
             close()
         span.end()
+        context.detach(token)