[core] Add WSGI http header support to HTTP propagator (#522)

* [core] add WSGI compatible http header support to propagator

* [core] clean up old TODOs, unused imports

* [core] use string literals in tests instead of constants

Author: Kyle-Verhoog

ddtrace/propagation/http.py

@@ -2,13 +2,28 @@
 
 from ..context import Context
 
+from .utils import get_wsgi_header
+
 log = logging.getLogger(__name__)
 
 # HTTP headers one should set for distributed tracing.
 # These are cross-language (eg: Python, Go and other implementations should honor these)
-HTTP_HEADER_TRACE_ID = 'x-datadog-trace-id'
-HTTP_HEADER_PARENT_ID = 'x-datadog-parent-id'
-HTTP_HEADER_SAMPLING_PRIORITY = 'x-datadog-sampling-priority'
+HTTP_HEADER_TRACE_ID = "x-datadog-trace-id"
+HTTP_HEADER_PARENT_ID = "x-datadog-parent-id"
+HTTP_HEADER_SAMPLING_PRIORITY = "x-datadog-sampling-priority"
+
+
+# Note that due to WSGI spec we have to also check for uppercased and prefixed
+# versions of these headers
+POSSIBLE_HTTP_HEADER_TRACE_IDS = frozenset(
+    [HTTP_HEADER_TRACE_ID, get_wsgi_header(HTTP_HEADER_TRACE_ID)]
+)
+POSSIBLE_HTTP_HEADER_PARENT_IDS = frozenset(
+    [HTTP_HEADER_PARENT_ID, get_wsgi_header(HTTP_HEADER_PARENT_ID)]
+)
+POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES = frozenset(
+    [HTTP_HEADER_SAMPLING_PRIORITY, get_wsgi_header(HTTP_HEADER_SAMPLING_PRIORITY)]
+)
 
 
 class HTTPPropagator(object):
@@ -39,6 +54,36 @@ def parent_call():
         if sampling_priority is not None:
             headers[HTTP_HEADER_SAMPLING_PRIORITY] = str(span_context.sampling_priority)
 
+    @staticmethod
+    def extract_trace_id(headers):
+        trace_id = 0
+
+        for key in POSSIBLE_HTTP_HEADER_TRACE_IDS:
+            if key in headers:
+                trace_id = headers.get(key)
+
+        return int(trace_id)
+
+    @staticmethod
+    def extract_parent_span_id(headers):
+        parent_span_id = 0
+
+        for key in POSSIBLE_HTTP_HEADER_PARENT_IDS:
+            if key in headers:
+                parent_span_id = headers.get(key)
+
+        return int(parent_span_id)
+
+    @staticmethod
+    def extract_sampling_priority(headers):
+        sampling_priority = None
+
+        for key in POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES:
+            if key in headers:
+                sampling_priority = headers.get(key)
+
+        return sampling_priority
+
     def extract(self, headers):
         """Extract a Context from HTTP headers into a new Context.
 
@@ -60,9 +105,10 @@ def my_controller(url, headers):
             return Context()
 
         try:
-            trace_id = int(headers.get(HTTP_HEADER_TRACE_ID, 0))
-            parent_span_id = int(headers.get(HTTP_HEADER_PARENT_ID, 0))
-            sampling_priority = headers.get(HTTP_HEADER_SAMPLING_PRIORITY)
+            trace_id = HTTPPropagator.extract_trace_id(headers)
+            parent_span_id = HTTPPropagator.extract_parent_span_id(headers)
+            sampling_priority = HTTPPropagator.extract_sampling_priority(headers)
+
             if sampling_priority is not None:
                 sampling_priority = int(sampling_priority)
 

ddtrace/propagation/utils.py

@@ -0,0 +1,6 @@
+def get_wsgi_header(header):
+    """Returns a WSGI compliant HTTP header.
+    See https://www.python.org/dev/peps/pep-3333/#environ-variables for
+    information from the spec.
+    """
+    return "HTTP_{}".format(header.upper().replace("-", "_"))

tests/propagation/test_http.py

@@ -1,42 +1,63 @@
 from unittest import TestCase
-from nose.tools import eq_, ok_
+from nose.tools import eq_
 from tests.test_tracer import get_dummy_tracer
 
-from ddtrace.span import Span
-from ddtrace.context import Context, ThreadLocalContext
-
 from ddtrace.propagation.http import (
     HTTPPropagator,
     HTTP_HEADER_TRACE_ID,
     HTTP_HEADER_PARENT_ID,
     HTTP_HEADER_SAMPLING_PRIORITY,
 )
 
+
 class TestHttpPropagation(TestCase):
     """
     Tests related to the ``Context`` class that hosts the trace for the
     current execution flow.
     """
+
     def test_inject(self):
         tracer = get_dummy_tracer()
 
         with tracer.trace("global_root_span") as span:
+            span.context.sampling_priority = 2
             headers = {}
             propagator = HTTPPropagator()
             propagator.inject(span.context, headers)
 
             eq_(int(headers[HTTP_HEADER_TRACE_ID]), span.trace_id)
             eq_(int(headers[HTTP_HEADER_PARENT_ID]), span.span_id)
-            # TODO: do it for priority too
-
+            eq_(
+                int(headers[HTTP_HEADER_SAMPLING_PRIORITY]),
+                span.context.sampling_priority,
+            )
 
     def test_extract(self):
         tracer = get_dummy_tracer()
 
         headers = {
-            HTTP_HEADER_TRACE_ID: '1234',
-            HTTP_HEADER_PARENT_ID: '5678',
-            HTTP_HEADER_SAMPLING_PRIORITY: '1',
+            "x-datadog-trace-id": "1234",
+            "x-datadog-parent-id": "5678",
+            "x-datadog-sampling-priority": "1",
+        }
+
+        propagator = HTTPPropagator()
+        context = propagator.extract(headers)
+        tracer.context_provider.activate(context)
+
+        with tracer.trace("local_root_span") as span:
+            eq_(span.trace_id, 1234)
+            eq_(span.parent_id, 5678)
+            eq_(span.context.sampling_priority, 1)
+
+    def test_WSGI_extract(self):
+        """Ensure we support the WSGI formatted headers as well."""
+        tracer = get_dummy_tracer()
+
+        headers = {
+            "HTTP_X_DATADOG_TRACE_ID": "1234",
+            "HTTP_X_DATADOG_PARENT_ID": "5678",
+            "HTTP_X_DATADOG_SAMPLING_PRIORITY": "1",
         }
 
         propagator = HTTPPropagator()
@@ -46,4 +67,4 @@ def test_extract(self):
         with tracer.trace("local_root_span") as span:
             eq_(span.trace_id, 1234)
             eq_(span.parent_id, 5678)
-            # TODO: do it for priority too
+            eq_(span.context.sampling_priority, 1)

tests/propagation/test_utils.py

@@ -0,0 +1,6 @@
+from ddtrace.propagation.utils import get_wsgi_header
+
+
+class TestPropagationUtils(object):
+    def test_get_wsgi_header(self):
+        assert get_wsgi_header("x-datadog-trace-id") == "HTTP_X_DATADOG_TRACE_ID"