Add support for Flask 2.0

ddtrace/contrib/flask/patch.py

@@ -74,7 +74,10 @@ def patch():
     _w("flask", "Flask.preprocess_request", request_tracer("preprocess_request"))
     _w("flask", "Flask.add_url_rule", traced_add_url_rule)
     _w("flask", "Flask.endpoint", traced_endpoint)
-    _w("flask", "Flask._register_error_handler", traced_register_error_handler)
+    if flask_version >= (2, 0, 0):
+        _w("flask", "Flask.register_error_handler", traced_register_error_handler)
+    else:
+        _w("flask", "Flask._register_error_handler", traced__register_error_handler)
 
     # flask.blueprints.Blueprint methods that have custom tracing (add metadata, wrap functions, etc)
     _w("flask", "Blueprint.register", traced_blueprint_register)
@@ -175,7 +178,6 @@ def unpatch():
         "Flask.dispatch_request",
         "Flask.add_url_rule",
         "Flask.endpoint",
-        "Flask._register_error_handler",
         "Flask.preprocess_request",
         "Flask.process_response",
         "Flask.handle_exception",
@@ -218,6 +220,11 @@ def unpatch():
         "templating._render",
     ]
 
+    if flask_version >= (2, 0, 0):
+        props.append("Flask.register_error_handler")
+    else:
+        props.append("Flask._register_error_handler")
+
     # These were added in 0.11.0
     if flask_version >= (0, 11):
         props.append("before_render_template.receivers_for")
@@ -429,7 +436,7 @@ def _wrap(template, context, app):
     return _wrap(*args, **kwargs)
 
 
-def traced_register_error_handler(wrapped, instance, args, kwargs):
+def traced__register_error_handler(wrapped, instance, args, kwargs):
     """Wrapper to trace all functions registered with flask.app.register_error_handler"""
 
     def _wrap(key, code_or_exception, f):
@@ -438,6 +445,15 @@ def _wrap(key, code_or_exception, f):
     return _wrap(*args, **kwargs)
 
 
+def traced_register_error_handler(wrapped, instance, args, kwargs):
+    """Wrapper to trace all functions registered with flask.app.register_error_handler"""
+
+    def _wrap(code_or_exception, f):
+        return wrapped(code_or_exception, wrap_function(instance, f))
+
+    return _wrap(*args, **kwargs)
+
+
 def request_tracer(name):
     @with_instance_pin
     def _traced_request(pin, wrapped, instance, args, kwargs):

riotfile.py

@@ -585,13 +585,12 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                 "blinker": latest,
             },
             venvs=[
-                # Flask 0.10, 0.11
+                # Flask == 0.12.0
                 Venv(
                     pys=select_pys(),
                     pkgs={
-                        "flask": ["~=0.10.0", "~=0.11.0"],
+                        "flask": ["~=0.12.0"],
                         "pytest": "~=3.0",
-                        "Werkzeug": "<1.0",
                     },
                 ),
                 Venv(
@@ -602,17 +601,19 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                         "DATADOG_PATCH_MODULES": "jinja2:false",
                     },
                     pkgs={
-                        "flask": ["~=0.10.0", "~=0.11.0"],
+                        "flask": ["~=0.12.0"],
                         "pytest": "~=3.0",
-                        "Werkzeug": "<1.0",
                     },
                 ),
-                # Flask == 0.12.0
+                # Flask 1.x.x
                 Venv(
                     pys=select_pys(),
                     pkgs={
-                        "flask": ["~=0.12.0"],
-                        "pytest": "~=3.0",
+                        "flask": [
+                            "~=1.0.0",
+                            "==1.1.2",
+                            "~=1.0",  # latest 1.x
+                        ],
                     },
                 ),
                 Venv(
@@ -623,26 +624,37 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                         "DATADOG_PATCH_MODULES": "jinja2:false",
                     },
                     pkgs={
-                        "flask": ["~=0.12.0"],
-                        "pytest": "~=3.0",
+                        "flask": [
+                            "~=1.0.0",
+                            "==1.1.2",
+                            "~=1.0",  # latest 1.x
+                        ],
                     },
                 ),
-                # Flask >= 1.0.0
+                # Flask >= 2.0.0
                 Venv(
-                    pys=select_pys(),
+                    pys=select_pys(min_version="3.6"),
                     pkgs={
-                        "flask": ["~=1.0.0", "~=1.1.0", "<2.0.0"],
+                        "flask": [
+                            "~=2.0.0",
+                            "~=2.0",  # latest 2.x
+                            latest,
+                        ],
                     },
                 ),
                 Venv(
-                    pys=select_pys(),
+                    pys=select_pys(min_version="3.6"),
                     command="python tests/ddtrace_run.py pytest {cmdargs} tests/contrib/flask_autopatch",
                     env={
                         "DATADOG_SERVICE_NAME": "test.flask.service",
                         "DATADOG_PATCH_MODULES": "jinja2:false",
                     },
                     pkgs={
-                        "flask": ["~=1.0.0", "~=1.1.0", "<2.0.0"],
+                        "flask": [
+                            "~=2.0.0",
+                            "~=2.0",  # latest 2.x
+                            latest,
+                        ],
                     },
                 ),
             ],

tests/contrib/flask/__init__.py

@@ -18,9 +18,7 @@ def setUp(self):
         Pin.override(self.app, tracer=self.tracer)
 
     def tearDown(self):
-        # Remove any remaining spans
-        self.tracer.pop()
-
+        super(BaseFlaskTestCase, self).tearDown()
         # Unpatch Flask
         unpatch()
 

tests/contrib/flask/test_flask_helpers.py

@@ -1,8 +1,11 @@
+from io import BytesIO
+
 import flask
 
 from ddtrace import Pin
+from ddtrace.compat import StringIO
 from ddtrace.contrib.flask import unpatch
-from ddtrace.internal.compat import StringIO
+from ddtrace.contrib.flask.patch import flask_version
 
 from . import BaseFlaskTestCase
 
@@ -78,7 +81,10 @@ def test_send_file(self):
         When calling a patched ``flask.send_file``
             We create the expected spans
         """
-        fp = StringIO("static file")
+        if flask_version >= (2, 0, 0):
+            fp = BytesIO(b"static file")
+        else:
+            fp = StringIO("static file")
 
         with self.app.app_context():
             with self.app.test_request_context("/"):
@@ -110,7 +116,11 @@ def test_send_file_pin_disabled(self):
         pin = Pin.get_from(self.app)
         pin.tracer.enabled = False
 
-        fp = StringIO("static file")
+        if flask_version >= (2, 0, 0):
+            fp = BytesIO(b"static file")
+        else:
+            fp = StringIO("static file")
+
         with self.app.app_context():
             with self.app.test_request_context("/"):
                 # DEV: Flask >= (0, 12, 0) tries to infer mimetype, so set explicitly

tests/contrib/flask/test_hooks.py

@@ -14,7 +14,7 @@ def setUp(self):
         def index():
             return "Hello Flask", 200
 
-        self.bp = Blueprint(__name__, "bp")
+        self.bp = Blueprint("bp", __name__)
 
         @self.bp.route("/bp")
         def bp():