Merge pull request #25 from getsentry/wsgi-middleware

ref: Share code from Flask/Django with WSGI middleware

Author: untitaker

sentry_sdk/_compat.py

@@ -16,6 +16,13 @@ def implements_str(cls):
         cls.__str__ = lambda x: unicode(x).encode("utf-8")  # noqa
         return cls
 
+    exec("def reraise(tp, value, tb=None):\n raise tp, value, tb")
+
+    def implements_iterator(cls):
+        cls.next = cls.__next__
+        del cls.__next__
+        return cls
+
 
 else:
     import urllib.parse as urlparse  # noqa
@@ -24,9 +31,19 @@ def implements_str(cls):
     text_type = str
     number_types = (int, float)
 
+    def _identity(x):
+        return x
+
+    implements_iterator = _identity
+
     def implements_str(x):
         return x
 
+    def reraise(tp, value, tb=None):
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+
 
 def with_metaclass(meta, *bases):
     class metaclass(type):

sentry_sdk/integrations/_wsgi.py

@@ -1,7 +1,10 @@
 import json
+import sys
 
-from sentry_sdk.hub import _should_send_default_pii
+import sentry_sdk
+from sentry_sdk.hub import _internal_exceptions, _should_send_default_pii
 from sentry_sdk.event import AnnotatedValue
+from sentry_sdk._compat import reraise, implements_iterator
 
 
 def get_environ(environ):
@@ -38,28 +41,11 @@ def __init__(self, request):
         self.request = request
 
     def extract_into_event(self, event, client_options):
-        if "request" in event:
-            return
-
-        # if the code below fails halfway through we at least have some data
-        event["request"] = request_info = {}
-
+        request_info = event.setdefault("request", {})
         request_info["url"] = self.url
-        request_info["query_string"] = self.query_string
-        request_info["method"] = self.method
-
-        request_info["env"] = dict(get_environ(self.env))
 
         if _should_send_default_pii():
-            request_info["headers"] = dict(self.headers)
             request_info["cookies"] = dict(self.cookies)
-        else:
-            request_info["headers"] = {
-                k: v
-                for k, v in dict(self.headers).items()
-                if k.lower().replace("_", "-")
-                not in ("set-cookie", "cookie", "authentication")
-            }
 
         bodies = client_options.get("request_bodies")
         if (
@@ -108,22 +94,6 @@ def content_length(self):
     def url(self):
         raise NotImplementedError()
 
-    @property
-    def query_string(self):
-        return self.env.get("QUERY_STRING")
-
-    @property
-    def method(self):
-        return self.env.get("REQUEST_METHOD")
-
-    @property
-    def headers(self):
-        return get_headers(self.env)
-
-    @property
-    def env(self):
-        raise NotImplementedError()
-
     @property
     def cookies(self):
         raise NotImplementedError()
@@ -136,10 +106,6 @@ def raw_data(self):
     def form(self):
         raise NotImplementedError()
 
-    @property
-    def form_is_multipart(self):
-        return self.env.get("CONTENT_TYPE").startswith("multipart/form-data")
-
     @property
     def is_json(self):
         mt = (self.env.get("CONTENT_TYPE") or "").split(";", 1)[0]
@@ -177,3 +143,89 @@ def get_client_ip(environ):
         return environ["HTTP_X_FORWARDED_FOR"].split(",")[0].strip()
     except (KeyError, IndexError):
         return environ.get("REMOTE_ADDR")
+
+
+def run_wsgi_app(app, environ, start_response):
+    hub = sentry_sdk.get_current_hub()
+    hub.push_scope()
+    with _internal_exceptions():
+        client_options = sentry_sdk.get_current_hub().client.options
+        sentry_sdk.get_current_hub().add_event_processor(
+            lambda: _make_wsgi_event_processor(environ, client_options)
+        )
+
+    try:
+        rv = app(environ, start_response)
+    except Exception:
+        einfo = sys.exc_info()
+        sentry_sdk.capture_exception(einfo)
+        hub.pop_scope_unsafe()
+        reraise(*einfo)
+
+    return _ScopePoppingResponse(hub, rv)
+
+
+@implements_iterator
+class _ScopePoppingResponse(object):
+    __slots__ = ("_response", "_hub")
+
+    def __init__(self, hub, response):
+        self._hub = hub
+        self._response = response
+
+    def __iter__(self):
+        try:
+            self._response = iter(self._response)
+        except Exception:
+            einfo = sys.exc_info()
+            sentry_sdk.capture_exception(einfo)
+            reraise(*einfo)
+        return self
+
+    def __next__(self):
+        try:
+            return next(self._response)
+        except StopIteration:
+            raise
+        except Exception:
+            einfo = sys.exc_info()
+            sentry_sdk.capture_exception(einfo)
+            reraise(*einfo)
+
+    def close(self):
+        self._hub.pop_scope_unsafe()
+        if hasattr(self._response, "close"):
+            try:
+                self._response.close()
+            except Exception:
+                einfo = sys.exc_info()
+                sentry_sdk.capture_exception(einfo)
+                reraise(*einfo)
+
+
+def _make_wsgi_event_processor(environ, client_options):
+    def event_processor(event):
+        with _internal_exceptions():
+            # if the code below fails halfway through we at least have some data
+            request_info = event.setdefault("request", {})
+
+            if "query_string" not in request_info:
+                request_info["query_string"] = environ.get("QUERY_STRING")
+
+            if "method" not in request_info:
+                request_info["method"] = environ.get("REQUEST_METHOD")
+
+            if "env" not in request_info:
+                request_info["env"] = dict(get_environ(environ))
+
+            if "headers" not in request_info:
+                request_info["headers"] = dict(get_headers(environ))
+                if not _should_send_default_pii():
+                    request_info["headers"] = {
+                        k: v
+                        for k, v in request_info["headers"].items()
+                        if k.lower().replace("_", "-")
+                        not in ("set-cookie", "cookie", "authorization")
+                    }
+
+    return event_processor

sentry_sdk/integrations/django.py

@@ -10,7 +10,7 @@
 
 from sentry_sdk import get_current_hub, capture_exception
 from sentry_sdk.hub import _internal_exceptions, _should_send_default_pii
-from ._wsgi import RequestExtractor, get_client_ip
+from ._wsgi import RequestExtractor, get_client_ip, run_wsgi_app
 from . import Integration
 
 
@@ -33,19 +33,31 @@ def __init__(self):
         pass
 
     def install(self, client):
-        from django.core.handlers.base import BaseHandler
+        # Patch in our custom middleware.
 
-        make_event_processor = self._make_event_processor
+        from django.core.handlers.wsgi import WSGIHandler
+
+        old_app = WSGIHandler.__call__
+
+        def sentry_patched_wsgi_handler(self, environ, start_response):
+            return run_wsgi_app(
+                lambda *a, **kw: old_app(self, *a, **kw), environ, start_response
+            )
+
+        WSGIHandler.__call__ = sentry_patched_wsgi_handler
+
+        # patch get_response, because at that point we have the Django request
+        # object
+
+        from django.core.handlers.base import BaseHandler
 
         old_get_response = BaseHandler.get_response
+        make_event_processor = self._make_event_processor
 
         def sentry_patched_get_response(self, request):
-            with get_current_hub().push_scope():
-                get_current_hub().add_event_processor(
-                    lambda: make_event_processor(request)
-                )
+            get_current_hub().add_event_processor(lambda: make_event_processor(request))
 
-                return old_get_response(self, request)
+            return old_get_response(self, request)
 
         BaseHandler.get_response = sentry_patched_get_response
 

sentry_sdk/integrations/flask.py

@@ -2,15 +2,15 @@
 
 from sentry_sdk import capture_exception, get_current_hub
 from sentry_sdk.hub import _internal_exceptions, _should_send_default_pii
-from ._wsgi import RequestExtractor
+from ._wsgi import RequestExtractor, run_wsgi_app
 from . import Integration
 
 try:
     from flask_login import current_user
 except ImportError:
     current_user = None
 
-from flask import _request_ctx_stack, _app_ctx_stack
+from flask import Flask, _request_ctx_stack, _app_ctx_stack
 from flask.signals import (
     appcontext_pushed,
     appcontext_tearing_down,
@@ -29,8 +29,19 @@ def install(self, client):
         appcontext_tearing_down.connect(_pop_appctx)
         got_request_exception.connect(_capture_exception)
 
+        old_app = Flask.__call__
+
+        def sentry_patched_wsgi_app(self, environ, start_response):
+            return run_wsgi_app(
+                lambda *a, **kw: old_app(self, *a, **kw), environ, start_response
+            )
+
+        Flask.__call__ = sentry_patched_wsgi_app
+
 
 def _push_appctx(*args, **kwargs):
+    # always want to push scope regardless of whether WSGI app might already
+    # have (not the case for CLI for example)
     get_current_hub().push_scope()
     get_current_hub().add_event_processor(_make_event_processor)
 

tests/integrations/conftest.py

@@ -6,12 +6,16 @@
 @pytest.fixture
 def capture_exceptions(monkeypatch):
     def inner():
-        errors = []
+        errors = set()
         old_capture_exception = sentry_sdk.Hub.current.capture_exception
 
         def capture_exception(error=None):
-            errors.append(error or sys.exc_info()[1])
-            return old_capture_exception(error)
+            given_error = error
+            error = error or sys.exc_info()[1]
+            if isinstance(error, tuple):
+                error = error[1]
+            errors.add(error)
+            return old_capture_exception(given_error)
 
         monkeypatch.setattr(
             sentry_sdk.Hub.current, "capture_exception", capture_exception

tests/integrations/django/myapp/settings.py

@@ -37,7 +37,7 @@
 # SECURITY WARNING: don't run with debug turned on in production!
 DEBUG = True
 
-ALLOWED_HOSTS = []
+ALLOWED_HOSTS = ["localhost"]
 
 
 # Application definition