[pyramid] add distributed tracing

Author: Emanuele Palazzetti

ddtrace/contrib/pyramid/constants.py

@@ -1,3 +1,4 @@
 SETTINGS_SERVICE = 'datadog_trace_service'
 SETTINGS_TRACER = 'datadog_tracer'
 SETTINGS_TRACE_ENABLED = 'datadog_trace_enabled'
+SETTINGS_DISTRIBUTED_TRACING = 'datadog_distributed_tracing'

ddtrace/contrib/pyramid/trace.py

@@ -9,13 +9,21 @@
 # project
 import ddtrace
 from ...ext import http, AppTypes
-from .constants import SETTINGS_SERVICE, SETTINGS_TRACE_ENABLED, SETTINGS_TRACER
+from ...propagation.http import HTTPPropagator
+from .constants import (
+    SETTINGS_TRACER,
+    SETTINGS_SERVICE,
+    SETTINGS_TRACE_ENABLED,
+    SETTINGS_DISTRIBUTED_TRACING,
+)
+
 
 log = logging.getLogger(__name__)
 
 DD_TWEEN_NAME = 'ddtrace.contrib.pyramid:trace_tween_factory'
 DD_SPAN = '_datadog_span'
 
+
 def trace_pyramid(config):
     config.include('ddtrace.contrib.pyramid')
 
@@ -49,6 +57,7 @@ def trace_tween_factory(handler, registry):
     service = settings.get(SETTINGS_SERVICE) or 'pyramid'
     tracer = settings.get(SETTINGS_TRACER) or ddtrace.tracer
     enabled = asbool(settings.get(SETTINGS_TRACE_ENABLED, tracer.enabled))
+    distributed_tracing = asbool(settings.get(SETTINGS_DISTRIBUTED_TRACING, False))
 
     # set the service info
     tracer.set_service_info(
@@ -59,6 +68,12 @@ def trace_tween_factory(handler, registry):
     if enabled:
         # make a request tracing function
         def trace_tween(request):
+            if distributed_tracing:
+                propagator = HTTPPropagator()
+                context = propagator.extract(request.headers)
+                # only need to active the new context if something was propagated
+                if context.trace_id:
+                    tracer.context_provider.activate(context)
             with tracer.trace('pyramid.request', service=service, resource='404') as span:
                 setattr(request, DD_SPAN, span)  # used to find the tracer in templates
                 response = None

tests/contrib/pyramid/test_pyramid.py

@@ -3,7 +3,6 @@
 
 from nose.tools import eq_, assert_raises
 
-# project
 from ddtrace import compat
 from ddtrace.contrib.pyramid.patch import insert_tween_if_needed
 
@@ -267,3 +266,30 @@ def test_tween_overridden(self):
         self.app.get('/json', status=200)
         spans = self.tracer.writer.pop()
         eq_(len(spans), 0)
+
+
+class TestPyramidDistributedTracing(PyramidBase):
+    instrument = True
+
+    def get_settings(self):
+        return {
+            'datadog_distributed_tracing': True,
+        }
+
+    def test_distributed_tracing(self):
+        # ensure the Context is properly created
+        # if distributed tracing is enabled
+        headers = {
+            'x-datadog-trace-id': '100',
+            'x-datadog-parent-id': '42',
+            'x-datadog-sampling-priority': '2',
+        }
+        res = self.app.get('/', headers=headers, status=200)
+        writer = self.tracer.writer
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        # check the propagated Context
+        span = spans[0]
+        eq_(span.trace_id, 100)
+        eq_(span.parent_id, 42)
+        eq_(span.get_metric('_sampling_priority_v1'), 2)