Merge pull request #409 from clutchski/flask-errs

flask: don't override code of already handled errors

Author: Emanuele Palazzetti

ddtrace/contrib/flask/middleware.py

@@ -1,8 +1,8 @@
 """
-Datadog trace code for flask.
+Datadog tracing code for flask.
 
-Requires a modern version of flask and the `blinker` library (which is a
-dependency of flask signals).
+Installing the blinker library will allow the tracing middleware to collect
+more exception info.
 """
 
 # stdlib
@@ -21,54 +21,46 @@
 log = logging.getLogger(__name__)
 
 
+SPAN_NAME = 'flask.request'
+
+
 class TraceMiddleware(object):
 
     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.use_signals = use_signals
 
         self._tracer.set_service_info(
             service=service,
             app="flask",
             app_type=AppTypes.web,
         )
 
-        # warn the user if signals are unavailable (because blinker isn't
-        # installed) if they are asking to use them.
+        # Install hooks which time requests.
+        self.app.before_request(self._before_request)
+        self.app.after_request(self._after_request)
+        self.app.teardown_request(self._teardown_request)
+
+        # Add exception handling signals. This will annotate exceptions that
+        # are caught and handled in custom user code.
+        # See https://github.com/DataDog/dd-trace-py/issues/390
         if use_signals and not signals.signals_available:
             self.app.logger.info(_blinker_not_installed_msg)
         self.use_signals = use_signals and signals.signals_available
-
-        # our signal receivers
-        self._receivers = []
-
-        # instrument request timings
         timing_signals = {
-            'request_started': self._request_started,
-            'request_finished': self._request_finished,
             'got_request_exception': self._request_exception,
         }
+        self._receivers = []
         if self.use_signals and _signals_exist(timing_signals):
             self._connect(timing_signals)
-        else:
-            # Fallback to request hooks. Won't catch exceptions.
-            # handle exceptions.
-            self.app.before_request(self._before_request)
-            self.app.after_request(self._after_request)
 
         _patch_render(tracer)
 
-    def _flask_signals_exist(self, names):
-        """ Return true if the current version of flask has all of the given
-            signals.
-        """
-        return all(getattr(signals, n, None) for n in names)
-
     def _connect(self, signal_to_handler):
         connected = True
         for name, handler in signal_to_handler.items():
@@ -85,7 +77,34 @@ def _connect(self, signal_to_handler):
             self._receivers.append(handler)
         return connected
 
-    # common methods
+    def _before_request(self):
+        """ Starts tracing the current request and stores it in the global
+            request object.
+        """
+        self._start_span()
+
+    def _after_request(self, response):
+        """ Runs after the server can process a response. """
+        try:
+            self._process_response(response)
+        except Exception:
+            self.app.logger.exception("error tracing response")
+        return response
+
+    def _teardown_request(self, exception):
+        """ Runs at the end of a request. If there's an unhandled exception, it
+            will be passed in.
+        """
+        # when we teardown the span, ensure we have a clean slate.
+        span = getattr(g, 'flask_datadog_span', None)
+        setattr(g, 'flask_datadog_span', None)
+        if not span:
+            return
+
+        try:
+            self._finish_span(span, exception=exception)
+        except Exception:
+            self.app.logger.exception("error finishing span")
 
     def _start_span(self):
         if self._use_distributed_tracing:
@@ -96,85 +115,76 @@ def _start_span(self):
                 self._tracer.context_provider.activate(context)
         try:
             g.flask_datadog_span = self._tracer.trace(
-                "flask.request",
+                SPAN_NAME,
                 service=self._service,
                 span_type=http.TYPE,
             )
         except Exception:
             self.app.logger.exception("error tracing request")
 
-    def _finish_span(self, response=None, exception=None):
-        """ Close and finish the active span if it exists. """
+    def _process_response(self, response):
         span = getattr(g, 'flask_datadog_span', None)
-        if span:
-            if span.sampled:
-                error = 0
-                code = response.status_code if response else None
-                method = request.method if request else None
-
-                # if we didn't get a response, but we did get an exception, set
-                # codes accordingly.
-                if not response and exception:
-                    code = 500
-                    # The 3 next lines might not be strictly required, since `set_traceback`
-                    # also get the exception from the sys.exc_info (and fill the error meta).
-                    # Since we aren't sure it always work/for insuring no BC break, keep
-                    # these lines which get overridden anyway.
-                    error = 1
-                    span.set_tag(errors.ERROR_TYPE, type(exception))
-                    span.set_tag(errors.ERROR_MSG, exception)
-                    # The provided `exception` object doesn't have a stack trace attached,
-                    # so attach the stack trace with `set_traceback`.
-                    span.set_traceback()
-
-                # the endpoint that matched the request is None if an exception
-                # happened so we fallback to a common resource
-                resource = code if not request.endpoint else request.endpoint
-                span.resource = compat.to_unicode(resource).lower()
-                span.set_tag(http.URL, compat.to_unicode(request.base_url or ''))
-                span.set_tag(http.STATUS_CODE, code)
-                span.set_tag(http.METHOD, method)
-                span.error = error
-            span.finish()
-            # Clear our span just in case.
-            g.flask_datadog_span = None
-
-    # Request hook methods
+        if not (span and span.sampled):
+            return
 
-    def _before_request(self):
-        """ Starts tracing the current request and stores it in the global
-            request object.
-        """
-        self._start_span()
+        code = response.status_code if response else ''
+        span.set_tag(http.STATUS_CODE, code)
 
-    def _after_request(self, response):
-        """ handles a successful response. """
-        try:
-            self._finish_span(response=response)
-        except Exception:
-            self.app.logger.exception("error finishing trace")
-        finally:
-            return response
-
-    # signal handling methods
-
-    def _request_started(self, sender):
-        self._start_span()
+    def _request_exception(self, *args, **kwargs):
+        exception = kwargs.get("exception", None)
+        span = getattr(g, 'flask_datadog_span', None)
+        if span and exception:
+            _set_error_on_span(span, exception)
 
-    def _request_finished(self, sender, response, **kwargs):
-        try:
-            self._finish_span(response=response)
-        except Exception:
-            self.app.logger.exception("error finishing trace")
-        return response
+    def _finish_span(self, span, exception=None):
+        if not span or not span.sampled:
+            return
 
-    def _request_exception(self, *args, **kwargs):
-        """ handles an error response. """
-        exception = kwargs.pop("exception", None)
+        code = span.get_tag(http.STATUS_CODE) or 0
         try:
-            self._finish_span(exception=exception)
+            code = int(code)
         except Exception:
-            self.app.logger.exception("error tracing error")
+            code = 0
+
+        if exception:
+            # if the request has already had a code set, don't override it.
+            code = code or 500
+            _set_error_on_span(span, exception)
+
+        # the endpoint that matched the request is None if an exception
+        # happened so we fallback to a common resource
+        span.error = 0 if code < 500 else 1
+
+        # the request isn't guaranteed to exist here, so only use it carefully.
+        method = ''
+        endpoint = ''
+        url = ''
+        if request:
+            method = request.method
+            endpoint = request.endpoint or code
+            url = request.base_url or ''
+
+        # Let users specify their own resource in middleware if they so desire.
+        # See case https://github.com/DataDog/dd-trace-py/issues/353
+        if span.resource == SPAN_NAME:
+            resource = endpoint or code
+            span.resource = compat.to_unicode(resource).lower()
+
+        span.set_tag(http.URL, compat.to_unicode(url))
+        span.set_tag(http.STATUS_CODE, code)
+        span.set_tag(http.METHOD, method)
+        span.finish()
+
+def _set_error_on_span(span, exception):
+    # The 3 next lines might not be strictly required, since `set_traceback`
+    # also get the exception from the sys.exc_info (and fill the error meta).
+    # Since we aren't sure it always work/for insuring no BC break, keep
+    # these lines which get overridden anyway.
+    span.set_tag(errors.ERROR_TYPE, type(exception))
+    span.set_tag(errors.ERROR_MSG, exception)
+    # The provided `exception` object doesn't have a stack trace attached,
+    # so attach the stack trace with `set_traceback`.
+    span.set_traceback()
 
 def _patch_render(tracer):
     """ patch flask's render template methods with the given tracer. """

tests/contrib/flask/test_flask.py

@@ -9,6 +9,7 @@
 from flask import Flask, render_template
 from nose.tools import eq_
 
+
 # project
 from ddtrace import Tracer
 from ddtrace.constants import SAMPLING_PRIORITY_KEY
@@ -25,8 +26,8 @@
 tracer.writer = writer
 
 
-class TestError(Exception):
-    pass
+class TestError(Exception): pass
+class HandleMe(Exception): pass
 
 
 # define a toy flask app.
@@ -45,6 +46,9 @@ def index():
 def error():
     raise TestError()
 
+@app.route('/handleme')
+def handle_me():
+    raise HandleMe()
 
 @app.route('/fatal')
 def fatal():
@@ -70,6 +74,13 @@ def child():
         span.set_tag('a', 'b')
         return 'child'
 
+@app.route("/custom_span")
+def custom_span():
+    span = tracer.current_span()
+    assert span
+    span.resource = "overridden"
+    return 'hiya'
+
 
 def unicode_view():
     return u'üŋïĉóđē'
@@ -87,6 +98,11 @@ def handle_my_exception(e):
     assert isinstance(e, TestError)
     return 'error', 500
 
+@app.errorhandler(HandleMe)
+def err_to_202(e):
+    assert isinstance(e, HandleMe)
+    return 'handled', 202
+
 
 # add tracing to the app (we use a global app to help ensure multiple requests
 # work)
@@ -195,6 +211,28 @@ def test_template(self):
         eq_(t.trace_id, s.trace_id)
         assert s.start < t.start < t.start + t.duration < end
 
+    def test_handleme(self):
+        start = time.time()
+        rv = app.get('/handleme')
+        end = time.time()
+
+        # ensure request worked
+        eq_(rv.status_code, 202)
+        eq_(rv.data, b'handled')
+
+        # ensure trace worked
+        assert not tracer.current_span(), tracer.current_span().pprint()
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.service, service)
+        eq_(s.resource, "handle_me")
+        assert s.start >= start
+        assert s.duration <= end - start
+        eq_(s.error, 0)
+        eq_(s.meta.get(http.STATUS_CODE), '202')
+        eq_(s.meta.get(http.METHOD), 'GET')
+
     def test_template_err(self):
         start = time.time()
         try:
@@ -294,7 +332,7 @@ def test_fatal(self):
         assert s.duration <= end - start
         eq_(s.meta.get(http.STATUS_CODE), '500')
         eq_(s.meta.get(http.METHOD), 'GET')
-        assert "ZeroDivisionError" in s.meta.get(errors.ERROR_TYPE)
+        assert "ZeroDivisionError" in s.meta.get(errors.ERROR_TYPE), s.meta
         assert "by zero" in s.meta.get(errors.ERROR_MSG)
         assert re.search('File ".*/contrib/flask/test_flask.py", line [0-9]+, in fatal', s.meta.get(errors.ERROR_STACK))
 
@@ -364,3 +402,20 @@ def test_propagation(self):
         eq_(s.trace_id, 1234)
         eq_(s.parent_id, 4567)
         eq_(s.get_metric(SAMPLING_PRIORITY_KEY), 2)
+
+    def test_custom_span(self):
+        rv = app.get('/custom_span')
+        eq_(rv.status_code, 200)
+        # ensure trace worked
+        assert not tracer.current_span(), tracer.current_span().pprint()
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.service, service)
+        eq_(s.resource, "overridden")
+        eq_(s.error, 0)
+        eq_(s.meta.get(http.STATUS_CODE), '200')
+        eq_(s.meta.get(http.METHOD), 'GET')
+
+
+