Merge pull request #365 from DataDog/andor44-master

Bring distributed tracing to Flask

Author: LotharSee

ddtrace/contrib/flask/__init__.py

@@ -10,7 +10,7 @@
 
 and create a `TraceMiddleware` object::
 
-    traced_app = TraceMiddleware(app, tracer, service="my-flask-app")
+    traced_app = TraceMiddleware(app, tracer, service="my-flask-app", distributed_tracing=False)
 
 Here is the end result, in a sample app::
 
@@ -22,12 +22,14 @@
 
     app = Flask(__name__)
 
-    traced_app = TraceMiddleware(app, tracer, service="my-flask-app")
+    traced_app = TraceMiddleware(app, tracer, service="my-flask-app", distributed_tracing=False)
 
     @app.route("/")
     def home():
         return "hello world"
 
+Set `distributed_tracing=True` if this is called remotely from an instrumented application.
+We suggest to enable it only for internal services where headers are under your control.
 """
 
 from ..util import require_modules

ddtrace/contrib/flask/middleware.py

@@ -11,6 +11,7 @@
 # project
 from ... import compat
 from ...ext import http, errors, AppTypes
+from ...propagation.http import HTTPPropagator
 
 # 3p
 import flask.templating
@@ -22,13 +23,14 @@
 
 class TraceMiddleware(object):
 
-    def __init__(self, app, tracer, service="flask", use_signals=True):
+    def __init__(self, app, tracer, service="flask", use_signals=True, distributed_tracing=False):
         self.app = app
         self.app.logger.info("initializing trace middleware")
 
         # save our traces.
         self._tracer = tracer
         self._service = service
+        self._use_distributed_tracing = distributed_tracing
 
         self._tracer.set_service_info(
             service=service,
@@ -86,6 +88,12 @@ def _connect(self, signal_to_handler):
     # common methods
 
     def _start_span(self):
+        if self._use_distributed_tracing:
+            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)
         try:
             g.flask_datadog_span = self._tracer.trace(
                 "flask.request",
@@ -168,7 +176,6 @@ def _request_exception(self, *args, **kwargs):
         except Exception:
             self.app.logger.exception("error tracing error")
 
-
 def _patch_render(tracer):
     """ patch flask's render template methods with the given tracer. """
     # fall back to patching  global method

docs/index.rst

@@ -285,6 +285,7 @@ For that, refer to the configuration of the given integration.
 Supported web frameworks:
 
 - Django
+- Flask
 
 
 HTTP client/server

tests/contrib/flask/test_flask.py

@@ -11,6 +11,7 @@
 
 # project
 from ddtrace import Tracer
+from ddtrace.constants import SAMPLING_PRIORITY_KEY
 from ddtrace.contrib.flask import TraceMiddleware
 from ddtrace.ext import http, errors
 from ...test_tracer import DummyWriter
@@ -91,7 +92,7 @@ def handle_my_exception(e):
 # work)
 service = "test.flask.service"
 assert not writer.pop()  # should always be empty
-traced_app = TraceMiddleware(app, tracer, service=service)
+traced_app = TraceMiddleware(app, tracer, service=service, distributed_tracing=True)
 
 # make the app testable
 app.config['TESTING'] = True
@@ -341,3 +342,25 @@ def test_404(self):
         eq_(s.meta.get(http.STATUS_CODE), '404')
         eq_(s.meta.get(http.METHOD), 'GET')
         eq_(s.meta.get(http.URL), u'http://localhost/404/üŋïĉóđē')
+
+    def test_propagation(self):
+        rv = app.get('/', headers={
+            'x-datadog-trace-id': '1234',
+            'x-datadog-parent-id': '4567',
+            'x-datadog-sampling-priority': '2'
+        })
+
+        # ensure request worked
+        eq_(rv.status_code, 200)
+        eq_(rv.data, b'hello')
+
+        # ensure trace worked
+        assert not tracer.current_span(), tracer.current_span().pprint()
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+
+        # ensure the propagation worked well
+        eq_(s.trace_id, 1234)
+        eq_(s.parent_id, 4567)
+        eq_(s.get_metric(SAMPLING_PRIORITY_KEY), 2)