[flask] rewrite Flask integration

.circleci/config.yml

@@ -326,10 +326,12 @@ jobs:
       - *restore_cache_step
       - run: tox -e 'flask_contrib-{py27,py34,py35,py36}-flask{010,011,012,10}-blinker' --result-json /tmp/flask.1.results
       - run: TOX_SKIP_DIST=False tox -e 'flask_contrib_autopatch-{py27,py34,py35,py36}-flask{010,011,012,10}-blinker' --result-json /tmp/flask.2.results
-      - run: tox -e 'flask_cache_contrib-{py27,py34,py35,py36}-flask{010,011,012}-flaskcache{013}-memcached-redis{210}-blinker' --result-json /tmp/flask.3.results
-      - run: TOX_SKIP_DIST=False tox -e 'flask_cache_contrib_autopatch-{py27,py34,py35,py36}-flask{010,011,012}-flaskcache{013}-memcached-redis{210}-blinker' --result-json /tmp/flask.4.results
-      - run: tox -e 'flask_cache_contrib-{py27}-flask{010,011}-flaskcache{012}-memcached-redis{210}-blinker' --result-json /tmp/flask.5.results
-      - run: TOX_SKIP_DIST=False tox -e 'flask_cache_contrib_autopatch-{py27}-flask{010,011}-flaskcache{012}-memcached-redis{210}-blinker' --result-json /tmp/flask.6.results
+      - run: tox -e 'flask_contrib-{py27}-flask{09}-blinker' --result-json /tmp/flask.3.results
+      - run: TOX_SKIP_DIST=False tox -e 'flask_contrib_autopatch-{py27}-flask{09}-blinker' --result-json /tmp/flask.4.results
+      - run: tox -e 'flask_cache_contrib-{py27,py34,py35,py36}-flask{010,011,012}-flaskcache{013}-memcached-redis{210}-blinker' --result-json /tmp/flask.5.results
+      - run: TOX_SKIP_DIST=False tox -e 'flask_cache_contrib_autopatch-{py27,py34,py35,py36}-flask{010,011,012}-flaskcache{013}-memcached-redis{210}-blinker' --result-json /tmp/flask.6.results
+      - run: tox -e 'flask_cache_contrib-{py27}-flask{010,011}-flaskcache{012}-memcached-redis{210}-blinker' --result-json /tmp/flask.7.results
+      - run: TOX_SKIP_DIST=False tox -e 'flask_cache_contrib_autopatch-{py27}-flask{010,011}-flaskcache{012}-memcached-redis{210}-blinker' --result-json /tmp/flask.8.results
       - persist_to_workspace:
           root: /tmp
           paths:
@@ -339,6 +341,8 @@ jobs:
             - flask.4.results
             - flask.5.results
             - flask.6.results
+            - flask.7.results
+            - flask.8.results
       - *save_cache_step
 
   gevent:

ddtrace/contrib/flask/__init__.py

@@ -1,35 +1,32 @@
 """
-The Flask trace middleware will track request timings and templates. It
-requires the `Blinker <https://pythonhosted.org/blinker/>`_ library, which
-Flask uses for signalling.
+The `Flask <http://flask.pocoo.org/>`_ integration will add tracing to all requests to your Flask application.
 
-To install the middleware, add::
+This integration will track the entire Flask lifecycle including user-defined endpoints, hooks,
+signals, and templating rendering.
 
-    from ddtrace import tracer
-    from ddtrace.contrib.flask import TraceMiddleware
+To configure tracing manually::
 
-and create a `TraceMiddleware` object::
+    from ddtrace import patch_all
+    patch_all()
 
-    traced_app = TraceMiddleware(app, tracer, service="my-flask-app", distributed_tracing=False)
+    from flask import Flask
 
-Here is the end result, in a sample app::
+    app = Flask(__name__)
 
-    from flask import Flask
-    import blinker as _
 
-    from ddtrace import tracer
-    from ddtrace.contrib.flask import TraceMiddleware
+    @app.route('/')
+    def index():
+        return 'hello world'
 
-    app = Flask(__name__)
 
-    traced_app = TraceMiddleware(app, tracer, service="my-flask-app", distributed_tracing=False)
+    if __name__ == '__main__':
+        app.run()
 
-    @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.
+You may also enable Flask tracing automatically via ddtrace-run::
+
+    ddtrace-run python app.py
+
 """
 
 from ...utils.importlib import require_modules
@@ -39,7 +36,11 @@ def home():
 
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
+        # DEV: We do this so we can `@mock.patch('ddtrace.contrib.flask._patch.<func>')` in tests
+        from . import patch as _patch
         from .middleware import TraceMiddleware
-        from .patch import patch
 
-        __all__ = ['TraceMiddleware', 'patch']
+        patch = _patch.patch
+        unpatch = _patch.unpatch
+
+        __all__ = ['TraceMiddleware', 'patch', 'unpatch']

ddtrace/contrib/flask/helpers.py

@@ -0,0 +1,44 @@
+from ddtrace import Pin
+import flask
+
+
+def get_current_app():
+    """Helper to get the flask.app.Flask from the current app context"""
+    appctx = flask._app_ctx_stack.top
+    if appctx:
+        return appctx.app
+    return None
+
+
+def with_instance_pin(func):
+    """Helper to wrap a function wrapper and ensure an enabled pin is available for the `instance`"""
+    def wrapper(wrapped, instance, args, kwargs):
+        pin = Pin._find(wrapped, instance, get_current_app())
+        if not pin or not pin.enabled():
+            return wrapped(*args, **kwargs)
+
+        return func(pin, wrapped, instance, args, kwargs)
+    return wrapper
+
+
+def simple_tracer(name, span_type=None):
+    """Generate a simple tracer that wraps the function call with `with tracer.trace()`"""
+    @with_instance_pin
+    def wrapper(pin, wrapped, instance, args, kwargs):
+        with pin.tracer.trace(name, service=pin.service, span_type=span_type):
+            return wrapped(*args, **kwargs)
+    return wrapper
+
+
+def get_current_span(pin, root=False):
+    """Helper to get the current span from the provided pins current call context"""
+    if not pin or not pin.enabled():
+        return None
+
+    ctx = pin.tracer.get_call_context()
+    if not ctx:
+        return None
+
+    if root:
+        return ctx.get_current_root_span()
+    return ctx.get_current_span()

ddtrace/contrib/flask/middleware.py

@@ -3,6 +3,7 @@
 from ... import compat
 from ...ext import http, errors, AppTypes
 from ...propagation.http import HTTPPropagator
+from ...utils.deprecation import deprecated
 
 import flask.templating
 from flask import g, request, signals
@@ -16,6 +17,7 @@
 
 class TraceMiddleware(object):
 
+    @deprecated(message='Use patching instead (see the docs).', version='1.0.0')
     def __init__(self, app, tracer, service="flask", use_signals=True, distributed_tracing=False):
         self.app = app
         log.debug('flask: initializing trace middleware')