[contrib:bottle] get distributed tracing context

ddtrace/contrib/bottle/__init__.py

@@ -9,6 +9,10 @@
     app = bottle.Bottle()
     plugin = TracePlugin(service="my-web-app")
     app.install(plugin)
+
+To enable distributed tracing::
+
+    plugin = TracePlugin(service="my-web-app", distributed_tracing=True)
 """
 
 from ..util import require_modules

ddtrace/contrib/bottle/trace.py

@@ -6,19 +6,22 @@
 import ddtrace
 from ddtrace.ext import http, AppTypes
 
+# project
+from ...propagation.http import HTTPPropagator
 
 class TracePlugin(object):
 
     name = 'trace'
     api = 2
 
-    def __init__(self, service="bottle", tracer=None):
+    def __init__(self, service="bottle", tracer=None, distributed_tracing=None):
         self.service = service
         self.tracer = tracer or ddtrace.tracer
         self.tracer.set_service_info(
             service=service,
             app="bottle",
             app_type=AppTypes.web)
+        self.distributed_tracing = distributed_tracing
 
     def apply(self, callback, route):
 
@@ -28,6 +31,13 @@ def wrapped(*args, **kwargs):
 
             resource = "%s %s" % (request.method, request.route.rule)
 
+            # Propagate headers such as x-datadog-trace-id.
+            if self.distributed_tracing:
+                propagator = HTTPPropagator()
+                context = propagator.extract(request.headers)
+                if context.trace_id:
+                    self.tracer.context_provider.activate(context)
+
             with self.tracer.trace("bottle.request", service=self.service, resource=resource) as s:
                 code = 0
                 try:
@@ -43,5 +53,3 @@ def wrapped(*args, **kwargs):
                     s.set_tag(http.METHOD, request.method)
 
         return wrapped
-
-

tests/contrib/bottle/test_distributed.py

@@ -0,0 +1,92 @@
+import bottle
+import ddtrace
+import webtest
+
+from unittest import TestCase
+from nose.tools import eq_, assert_not_equal
+from tests.test_tracer import get_dummy_tracer
+
+from ddtrace import compat
+from ddtrace.contrib.bottle import TracePlugin
+
+
+SERVICE = 'bottle-app'
+
+
+class TraceBottleDistributedTest(TestCase):
+    """
+    Ensures that Bottle is properly traced.
+    """
+    def setUp(self):
+        # provide a dummy tracer
+        self.tracer = get_dummy_tracer()
+        self._original_tracer = ddtrace.tracer
+        ddtrace.tracer = self.tracer
+        # provide a Bottle app
+        self.app = bottle.Bottle()
+
+    def tearDown(self):
+        # restore the tracer
+        ddtrace.tracer = self._original_tracer
+
+    def _trace_app_distributed(self, tracer=None):
+        self.app.install(TracePlugin(service=SERVICE, tracer=tracer, distributed_tracing=True))
+        self.app = webtest.TestApp(self.app)
+
+    def _trace_app_not_distributed(self, tracer=None):
+        self.app.install(TracePlugin(service=SERVICE, tracer=tracer, distributed_tracing=False))
+        self.app = webtest.TestApp(self.app)
+
+    def test_distributed(self):
+        # setup our test app
+        @self.app.route('/hi/<name>')
+        def hi(name):
+            return 'hi %s' % name
+        self._trace_app_distributed(self.tracer)
+
+        # make a request
+        headers = {'x-datadog-trace-id': '123',
+                   'x-datadog-parent-id': '456'}
+        resp = self.app.get('/hi/dougie', headers=headers)
+        eq_(resp.status_int, 200)
+        eq_(compat.to_unicode(resp.body), u'hi dougie')
+
+        # validate it's traced
+        spans = self.tracer.writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.name, 'bottle.request')
+        eq_(s.service, 'bottle-app')
+        eq_(s.resource, 'GET /hi/<name>')
+        eq_(s.get_tag('http.status_code'), '200')
+        eq_(s.get_tag('http.method'), 'GET')
+        # check distributed headers
+        eq_(123, s.trace_id)
+        eq_(456, s.parent_id)
+
+    def test_not_distributed(self):
+        # setup our test app
+        @self.app.route('/hi/<name>')
+        def hi(name):
+            return 'hi %s' % name
+        self._trace_app_not_distributed(self.tracer)
+
+        # make a request
+        headers = {'x-datadog-trace-id': '123',
+                   'x-datadog-parent-id': '456'}
+        resp = self.app.get('/hi/dougie', headers=headers)
+        eq_(resp.status_int, 200)
+        eq_(compat.to_unicode(resp.body), u'hi dougie')
+
+        # validate it's traced
+        spans = self.tracer.writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.name, 'bottle.request')
+        eq_(s.service, 'bottle-app')
+        eq_(s.resource, 'GET /hi/<name>')
+        eq_(s.get_tag('http.status_code'), '200')
+        eq_(s.get_tag('http.method'), 'GET')
+        # check distributed headers
+        assert_not_equal(123, s.trace_id)
+        assert_not_equal(456, s.parent_id)