ext/opentracing: Implement get_attribute and set_attribute (#912)

Author: ocelotl

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

@@ -90,11 +90,18 @@
 import opentracing
 from deprecated import deprecated
 
-import opentelemetry.trace as trace_api
 from opentelemetry import propagators
+from opentelemetry.context import Context
+from opentelemetry.correlationcontext import get_correlation, set_correlation
 from opentelemetry.ext.opentracing_shim import util
 from opentelemetry.ext.opentracing_shim.version import __version__
-from opentelemetry.trace import DefaultSpan, set_span_in_context
+from opentelemetry.trace import (
+    INVALID_SPAN_CONTEXT,
+    DefaultSpan,
+    Link,
+    get_current_span,
+    set_span_in_context,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -130,6 +137,8 @@ class SpanContextShim(opentracing.SpanContext):
 
     def __init__(self, otel_context):
         self._otel_context = otel_context
+        # Context is being used here since it must be immutable.
+        self._baggage = Context()
 
     def unwrap(self):
         """Returns the wrapped :class:`opentelemetry.trace.SpanContext`
@@ -144,17 +153,9 @@ def unwrap(self):
 
     @property
     def baggage(self):
-        """Implements the ``baggage`` property from the base class.
+        """Implements the ``baggage`` property from the base class."""
 
-        Warning:
-            Not implemented yet.
-        """
-
-        logger.warning(
-            "Using unimplemented property baggage on class %s.",
-            self.__class__.__name__,
-        )
-        # TODO: Implement.
+        return self._baggage
 
 
 class SpanShim(opentracing.Span):
@@ -270,31 +271,17 @@ def log(self, **kwargs):
     def log_event(self, event, payload=None):
         super().log_event(event, payload=payload)
 
-    def set_baggage_item(self, key, value):  # pylint:disable=unused-argument
-        """Implements the ``set_baggage_item()`` method from the base class.
-
-        Warning:
-            Not implemented yet.
-        """
-
-        logger.warning(
-            "Calling unimplemented method set_baggage_item() on class %s",
-            self.__class__.__name__,
+    def set_baggage_item(self, key, value):
+        """Implements the ``set_baggage_item`` method from the base class."""
+        # pylint: disable=protected-access
+        self._context._baggage = set_correlation(
+            key, value, context=self._context._baggage
         )
-        # TODO: Implement.
-
-    def get_baggage_item(self, key):  # pylint:disable=unused-argument
-        """Implements the ``get_baggage_item()`` method from the base class.
-
-        Warning:
-            Not implemented yet.
-        """
 
-        logger.warning(
-            "Calling unimplemented method get_baggage_item() on class %s",
-            self.__class__.__name__,
-        )
-        # TODO: Implement.
+    def get_baggage_item(self, key):
+        """Implements the ``get_baggage_item`` method from the base class."""
+        # pylint: disable=protected-access
+        return get_correlation(key, context=self._context._baggage)
 
 
 class ScopeShim(opentracing.Scope):
@@ -469,8 +456,8 @@ def active(self):
             shim and is likely to be handled in future versions.
         """
 
-        span = trace_api.get_current_span()
-        if span.get_context() == trace_api.INVALID_SPAN_CONTEXT:
+        span = get_current_span()
+        if span.get_context() == INVALID_SPAN_CONTEXT:
             return None
 
         span_context = SpanContextShim(span.get_context())
@@ -643,7 +630,7 @@ def start_span(
         links = []
         if references:
             for ref in references:
-                links.append(trace_api.Link(ref.referenced_context.unwrap()))
+                links.append(Link(ref.referenced_context.unwrap()))
 
         # The OpenTracing API expects time values to be `float` values which
         # represent the number of seconds since the epoch. OpenTelemetry
@@ -699,10 +686,10 @@ def get_as_list(dict_object, key):
 
         propagator = propagators.get_global_httptextformat()
         ctx = propagator.extract(get_as_list, carrier)
-        span = trace_api.get_current_span(ctx)
+        span = get_current_span(ctx)
         if span is not None:
             otel_context = span.get_context()
         else:
-            otel_context = trace_api.INVALID_SPAN_CONTEXT
+            otel_context = INVALID_SPAN_CONTEXT
 
         return SpanContextShim(otel_context)

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

@@ -17,12 +17,17 @@
 
 import time
 from unittest import TestCase
+from unittest.mock import Mock
 
 import opentracing
 
-import opentelemetry.ext.opentracing_shim as opentracingshim
 from opentelemetry import propagators, trace
-from opentelemetry.ext.opentracing_shim import util
+from opentelemetry.ext.opentracing_shim import (
+    SpanContextShim,
+    SpanShim,
+    create_tracer,
+    util,
+)
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.test.mock_httptextformat import (
     MockHTTPTextFormat,
@@ -36,7 +41,7 @@ class TestShim(TestCase):
     def setUp(self):
         """Create an OpenTelemetry tracer and a shim before every test case."""
         trace.set_tracer_provider(TracerProvider())
-        self.shim = opentracingshim.create_tracer(trace.get_tracer_provider())
+        self.shim = create_tracer(trace.get_tracer_provider())
 
     @classmethod
     def setUpClass(cls):
@@ -448,7 +453,7 @@ def test_span_context(self):
         """Test construction of `SpanContextShim` objects."""
 
         otel_context = trace.SpanContext(1234, 5678, is_remote=False)
-        context = opentracingshim.SpanContextShim(otel_context)
+        context = SpanContextShim(otel_context)
 
         self.assertIsInstance(context, opentracing.SpanContext)
         self.assertEqual(context.unwrap().trace_id, 1234)
@@ -473,7 +478,7 @@ def test_inject_http_headers(self):
         otel_context = trace.SpanContext(
             trace_id=1220, span_id=7478, is_remote=False
         )
-        context = opentracingshim.SpanContextShim(otel_context)
+        context = SpanContextShim(otel_context)
 
         headers = {}
         self.shim.inject(context, opentracing.Format.HTTP_HEADERS, headers)
@@ -486,7 +491,7 @@ def test_inject_text_map(self):
         otel_context = trace.SpanContext(
             trace_id=1220, span_id=7478, is_remote=False
         )
-        context = opentracingshim.SpanContextShim(otel_context)
+        context = SpanContextShim(otel_context)
 
         # Verify Format.TEXT_MAP
         text_map = {}
@@ -500,7 +505,7 @@ def test_inject_binary(self):
         otel_context = trace.SpanContext(
             trace_id=1220, span_id=7478, is_remote=False
         )
-        context = opentracingshim.SpanContextShim(otel_context)
+        context = SpanContextShim(otel_context)
 
         # Verify exception for non supported binary format.
         with self.assertRaises(opentracing.UnsupportedFormatException):
@@ -550,3 +555,20 @@ def test_extract_binary(self):
         # Verify exception for non supported binary format.
         with self.assertRaises(opentracing.UnsupportedFormatException):
             self.shim.extract(opentracing.Format.BINARY, bytearray())
+
+    def test_baggage(self):
+
+        span_context_shim = SpanContextShim(
+            trace.SpanContext(1234, 5678, is_remote=False)
+        )
+
+        baggage = span_context_shim.baggage
+
+        with self.assertRaises(ValueError):
+            baggage[1] = 3
+
+        span_shim = SpanShim(Mock(), span_context_shim, Mock())
+
+        span_shim.set_baggage_item(1, 2)
+
+        self.assertTrue(span_shim.get_baggage_item(1), 2)