[pylons] add distributed tracing via kwarg and environment variable

Author: Emanuele Palazzetti

ddtrace/contrib/pylons/__init__.py

@@ -10,9 +10,12 @@
 
     app = PylonsApp(...)
 
-    traced_app = PylonsTraceMiddleware(app, tracer, service="my-pylons-app")
+    traced_app = PylonsTraceMiddleware(app, tracer, service='my-pylons-app')
 
-Then you can define your routes and views as usual.
+Then you can define your routes and views as usual. To enable distributed tracing,
+set the following keyword argument::
+
+    traced_app = PylonsTraceMiddleware(app, tracer, service='my-pylons-app', distributed_tracing=True)
 """
 
 from ..util import require_modules

ddtrace/contrib/pylons/middleware.py

@@ -1,22 +1,24 @@
 import logging
 import sys
 
+from webob import Request
 from pylons import config
 
 from .renderer import trace_rendering
 from .constants import CONFIG_MIDDLEWARE
 
-from ...ext import http
-from ...ext import AppTypes
+from ...ext import http, AppTypes
+from ...propagation.http import HTTPPropagator
 
 log = logging.getLogger(__name__)
 
 
 class PylonsTraceMiddleware(object):
 
-    def __init__(self, app, tracer, service="pylons"):
+    def __init__(self, app, tracer, service='pylons', distributed_tracing=False):
         self.app = app
         self._service = service
+        self._distributed_tracing = distributed_tracing
         self._tracer = tracer
 
         # register middleware reference
@@ -32,6 +34,15 @@ def __init__(self, app, tracer, service="pylons"):
         )
 
     def __call__(self, environ, start_response):
+        if self._distributed_tracing:
+            # retrieve distributed tracing headers
+            request = Request(environ)
+            propagator = HTTPPropagator()
+            context = propagator.extract(request.headers)
+            # only need to active the new context if something was propagated
+            if context.trace_id:
+                self._tracer.context_provider.activate(context)
+
         with self._tracer.trace("pylons.request", service=self._service) as span:
             # Set the service in tracer.trace() as priority sampling requires it to be
             # set as early as possible when different services share one single agent.

ddtrace/contrib/pylons/patch.py

@@ -1,6 +1,5 @@
 import os
 import wrapt
-
 import pylons.wsgiapp
 
 from ddtrace import tracer, Pin
@@ -31,9 +30,10 @@ def traced_init(wrapped, instance, args, kwargs):
     wrapped(*args, **kwargs)
 
     # set tracing options and create the TraceMiddleware
-    service = os.environ.get('DATADOG_SERVICE_NAME') or 'pylons'
+    service = os.environ.get('DATADOG_SERVICE_NAME', 'pylons')
+    distributed_tracing = os.environ.get('DATADOG_PYLONS_DISTRIBUTED_TRACING', False)
     Pin(service=service, tracer=tracer).onto(instance)
-    traced_app = PylonsTraceMiddleware(instance, tracer, service=service)
+    traced_app = PylonsTraceMiddleware(instance, tracer, service=service, distributed_tracing=distributed_tracing)
 
     # re-order the middleware stack so that the first middleware is ours
     traced_app.app = instance.app

tests/contrib/pylons/test_pylons.py

@@ -8,6 +8,7 @@
 from paste.deploy import loadapp
 
 from ddtrace.ext import http
+from ddtrace.constants import SAMPLING_PRIORITY_KEY
 from ddtrace.contrib.pylons import PylonsTraceMiddleware
 
 from ...test_tracer import get_dummy_tracer
@@ -145,3 +146,44 @@ def test_failure_500_with_code_method(self):
         eq_(span.error, 1)
         eq_(span.get_tag('http.status_code'), '500')
         eq_(span.get_tag('error.msg'), 'Ouch!')
+
+    def test_distributed_tracing_default(self):
+        # ensure by default, distributed tracing is not enabled
+        headers = {
+            'x-datadog-trace-id': '100',
+            'x-datadog-parent-id': '42',
+            'x-datadog-sampling-priority': '2',
+        }
+        res = self.app.get(url_for(controller='root', action='index'), headers=headers)
+        eq_(res.status, 200)
+
+        spans = self.tracer.writer.pop()
+        ok_(spans, spans)
+        eq_(len(spans), 1)
+        span = spans[0]
+
+        ok_(span.trace_id != 100)
+        ok_(span.parent_id != 42)
+        ok_(span.get_metric(SAMPLING_PRIORITY_KEY) is None)
+
+    def test_distributed_tracing_enabled(self):
+        # ensure distributed tracing propagator is working
+        middleware = self.app.app
+        middleware._distributed_tracing = True
+        headers = {
+            'x-datadog-trace-id': '100',
+            'x-datadog-parent-id': '42',
+            'x-datadog-sampling-priority': '2',
+        }
+
+        res = self.app.get(url_for(controller='root', action='index'), headers=headers)
+        eq_(res.status, 200)
+
+        spans = self.tracer.writer.pop()
+        ok_(spans, spans)
+        eq_(len(spans), 1)
+        span = spans[0]
+
+        eq_(span.trace_id, 100)
+        eq_(span.parent_id, 42)
+        eq_(span.get_metric(SAMPLING_PRIORITY_KEY), 2)