Implement propagator options which allows for custom keys to be specified (#5)

* add tracer init to readme

* added build and release features

* fix makefile typo

* can enable shared spans if desired

* add ability to provide custom propagation headers if needed. default to Span-ID, Trace-Id, Parent-ID

* fix lint issue

Author: rhilfers

examples/main.py

@@ -5,6 +5,8 @@
 from haystack import HaystackTracer
 from haystack import AsyncHttpRecorder
 from haystack import LoggerRecorder
+from haystack.text_propagator import TextPropagator
+from haystack.propagator import PropagatorOpts
 
 recorder = LoggerRecorder()
 logging.basicConfig(level=logging.DEBUG)
@@ -15,6 +17,10 @@ def act_as_remote_service(headers):
     with HaystackTracer("Service-B", recorder,) as tracer:
         opentracing.tracer = tracer
 
+        # ---ability to use custom propagation headers if needed---
+        # prop_opts = PropagatorOpts("X-Trace-ID", "X-Span-ID", "X-Parent-Span")
+        # opentracing.tracer.register_propagator(opentracing.Format.HTTP_HEADERS, TextPropagator(prop_opts))
+
         # simulate network transfer delay
         time.sleep(.25)
 

haystack/constants.py

@@ -1,16 +1,14 @@
 # Name of the HTTP header or Key used to encode Trace ID
-TRACE_ID_HEADER = "Trace-ID"
+TRACE_ID = "Trace-ID"
 
-# Name of the HTTP header or Key used to encode Span ID followed by any
-# alternate header names to extract as Span ID
-SPAN_ID_HEADERS = ("Span-ID", "Message-ID",)
+# Name of the HTTP header or Key used to encode Span ID
+SPAN_ID = "Span-ID"
 
-# Name of the HTTP header or Key used to encode Parent Span ID followed by any
-# alternate header names to extract as Parent Span ID
-PARENT_SPAN_ID_HEADERS = ("Parent-ID", "Parent-Message-ID",)
+# Name of the HTTP header or Key used to encode Parent Span ID
+PARENT_SPAN_ID = "Parent-ID"
 
 # Prefix of the HTTP header or Key used to encode Baggage items
-BAGGAGE_HEADER_PREFIX = "Baggage-"
+BAGGAGE_PREFIX = "Baggage-"
 
 # The number of microseconds in one second
 SECONDS_TO_MICRO = 1000000

haystack/propagator.py

@@ -1,4 +1,11 @@
 from abc import ABC, abstractmethod
+from collections import namedtuple
+from .constants import (
+    TRACE_ID,
+    SPAN_ID,
+    PARENT_SPAN_ID,
+    BAGGAGE_PREFIX,
+)
 
 
 class Propagator(ABC):
@@ -10,3 +17,11 @@ def inject(self, span_context, carrier):
     @abstractmethod
     def extract(self, carrier):
         raise NotImplementedError()
+
+
+PropagatorOpts = namedtuple("PropagatorOpts", ["trace_id_key",
+                                               "span_id_key",
+                                               "parent_id_key",
+                                               "baggage_key_prefix"])
+PropagatorOpts.__new__.__defaults__ = (TRACE_ID, SPAN_ID, PARENT_SPAN_ID,
+                                       BAGGAGE_PREFIX)

haystack/text_propagator.py

@@ -1,24 +1,26 @@
-from .propagator import Propagator
+from .propagator import Propagator, PropagatorOpts
 from .span import SpanContext
 from opentracing import SpanContextCorruptedException
-from .constants import (
-    SPAN_ID_HEADERS,
-    TRACE_ID_HEADER,
-    PARENT_SPAN_ID_HEADERS,
-    BAGGAGE_HEADER_PREFIX,
-)
 
 
 class TextPropagator(Propagator):
     """A propagator for Format.TEXT_MAP and Format.HTTP_HEADERS"""
 
+    def __init__(self, propagator_opts=None):
+        """
+        Initialize with optional custom keys for the carrier.
+        :param propagator_opts: PropagatorOps named tuple with custom keys
+        to inject and extract :class:SpanContext from the carrier.
+        """
+        self._propagator_opts = propagator_opts or PropagatorOpts()
+
     def inject(self, span_context, carrier):
-        carrier[TRACE_ID_HEADER] = span_context.trace_id
-        carrier[SPAN_ID_HEADERS[0]] = span_context.span_id
-        carrier[PARENT_SPAN_ID_HEADERS[0]] = span_context.parent_id
+        carrier[self._propagator_opts.trace_id_key] = span_context.trace_id
+        carrier[self._propagator_opts.span_id_key] = span_context.span_id
+        carrier[self._propagator_opts.parent_id_key] = span_context.parent_id
         if span_context.baggage is not None:
             for item in span_context.baggage:
-                carrier[BAGGAGE_HEADER_PREFIX + item] = \
+                carrier[self._propagator_opts.baggage_key_prefix + item] = \
                     span_context.baggage[item]
 
     def extract(self, carrier):
@@ -27,16 +29,20 @@ def extract(self, carrier):
         parent_id = None
 
         for key in carrier:
+            lc_key = key.lower()
             value = carrier[key]
-            if key in SPAN_ID_HEADERS:
+            if lc_key == self._propagator_opts.span_id_key.lower():
                 span_id = value
                 count += 1
-            elif key == TRACE_ID_HEADER:
+            elif lc_key == self._propagator_opts.trace_id_key.lower():
                 trace_id = value
                 count += 1
-            elif key.startswith(BAGGAGE_HEADER_PREFIX):
-                baggage[key[len(BAGGAGE_HEADER_PREFIX):]] = value
-            elif key in PARENT_SPAN_ID_HEADERS:
+            elif lc_key.startswith(
+                    self._propagator_opts.baggage_key_prefix.lower()):
+                baggage[
+                    key[len(self._propagator_opts.baggage_key_prefix):]] = \
+                    value
+            elif lc_key == self._propagator_opts.parent_id_key.lower():
                 parent_id = value
 
         if count == 0:

tests/unit/test_text_propagator.py

@@ -1,40 +1,76 @@
 import unittest
 from opentracing import SpanContextCorruptedException
 from haystack.text_propagator import TextPropagator
+from haystack.propagator import PropagatorOpts
+from haystack.constants import (
+    TRACE_ID,
+    SPAN_ID,
+    PARENT_SPAN_ID,
+    BAGGAGE_PREFIX,
+)
 
 
 class TextPropagatorTest(unittest.TestCase):
 
     def setUp(self):
         self.propagator = TextPropagator()
 
-    def test_message_id_in_carrier_should_be_extracted_as_span_id(self):
+    def test_span_id_without_trace_id_throws_exception(self):
+        carrier = {"Span-ID": "123456", "Foo-ID": "1234567"}
+
+        self.assertRaises(SpanContextCorruptedException,
+                          self.propagator.extract, carrier)
+
+    def test_trace_id_without_span_id_throws_exception(self):
+        carrier = {"Trace-ID": "123456", "Foo-ID": "1234567"}
+
+        self.assertRaises(SpanContextCorruptedException,
+                          self.propagator.extract, carrier)
+
+    def test_no_context_in_carrier_returns_none(self):
+        span_id = "1234"
+        carrier = {"Not-A-Message-ID": span_id, "Not-A-Trace-ID": "123456"}
+
+        ctx = self.propagator.extract(carrier)
+
+        self.assertIsNone(ctx)
+
+    def test_lowercase_headers_are_still_extracted(self):
+        pass
+
+    def test_default_propagator_options_injected_and_context_is_extracted(self):
         span_id = "1234"
         parent_id = "4321"
         trace_id = "1212"
-        carrier = {"Message-ID": span_id, "Parent-Message-ID": parent_id,
-                   "Trace-ID": trace_id}
+        baggage = {"Item1": "Value1", "Item2": "Value2"}
+        carrier = {TRACE_ID: trace_id, SPAN_ID: span_id, PARENT_SPAN_ID: parent_id}
+        for key, value in baggage.items():
+            carrier[BAGGAGE_PREFIX + key] = value
 
         ctx = self.propagator.extract(carrier)
 
         self.assertEqual(ctx.span_id, span_id)
         self.assertEqual(ctx.parent_id, parent_id)
         self.assertEqual(ctx.trace_id, trace_id)
+        self.assertDictEqual(ctx.baggage, baggage)
 
-    def test_corrupted_context_throws_exception(self):
+    def test_custom_propagator_options_are_injected_and_context_is_extracted(self):
+        propagator = TextPropagator(PropagatorOpts("X-Trace-ID", "X-Span-ID",
+                                                   "X-Parent-ID", "X-baggage-"))
         span_id = "1234"
-        carrier = {"Message-ID": span_id, "Span-ID": "123456",
-                   "Trace-ID": "1234567"}  # two span id's are invalid
-
-        self.assertRaises(SpanContextCorruptedException, self.propagator.extract, carrier)
-
-    def test_no_context_in_carrier_returns_none(self):
-        span_id = "1234"
-        carrier = {"Not-A-Message-ID": span_id, "Not-A-Trace-ID": "123456"}
+        parent_id = "4321"
+        trace_id = "1212"
+        carrier = {"X-Span-ID": span_id, "X-Parent-ID": parent_id,
+                   "X-Trace-ID": trace_id}
+        baggage = {"Item1": "Value1", "Item2": "Value2"}
+        for key, value in baggage.items():
+            carrier["X-baggage-" + key] = value
 
-        ctx = self.propagator.extract(carrier)
+        ctx = propagator.extract(carrier)
 
-        self.assertIsNone(ctx)
+        self.assertEqual(ctx.span_id, span_id)
+        self.assertEqual(ctx.parent_id, parent_id)
+        self.assertEqual(ctx.trace_id, trace_id)
 
 
 if __name__ == "__main__":