fix: Push scope for each response chunk (#218)

* fix: Push scope for each response chunk

* fix: Remove implements_iterator

* fix: Just use a hub everywhere

Author: untitaker

sentry_sdk/_compat.py

@@ -21,11 +21,6 @@ def implements_str(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
@@ -40,8 +35,6 @@ def implements_iterator(cls):
     def _identity(x):
         return x
 
-    implements_iterator = _identity
-
     def implements_str(x):
         return x
 

sentry_sdk/integrations/wsgi.py

@@ -2,7 +2,7 @@
 
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
-from sentry_sdk._compat import PY2, reraise, implements_iterator
+from sentry_sdk._compat import PY2, reraise
 from sentry_sdk.integrations._wsgi_common import _filter_headers
 
 
@@ -57,21 +57,20 @@ def __init__(self, app):
         self.app = app
 
     def __call__(self, environ, start_response):
-        hub = Hub.current
-        hub.push_scope()
-        with capture_internal_exceptions():
-            with hub.configure_scope() as scope:
-                scope._name = "wsgi"
-                scope.add_event_processor(_make_wsgi_event_processor(environ))
+        hub = Hub(Hub.current)
+
+        with hub:
+            with capture_internal_exceptions():
+                with hub.configure_scope() as scope:
+                    scope._name = "wsgi"
+                    scope.add_event_processor(_make_wsgi_event_processor(environ))
 
-        try:
-            rv = self.app(environ, start_response)
-        except Exception:
-            einfo = _capture_exception(hub)
-            hub.pop_scope_unsafe()
-            reraise(*einfo)
+            try:
+                rv = self.app(environ, start_response)
+            except Exception:
+                reraise(*_capture_exception(hub))
 
-        return _ScopePoppingResponse(hub, rv)
+        return _ScopedResponse(hub, rv)
 
 
 def _get_environ(environ):
@@ -132,45 +131,35 @@ def _capture_exception(hub):
     return exc_info
 
 
-@implements_iterator
-class _ScopePoppingResponse(object):
-    __slots__ = ("_response", "_iterator", "_hub", "_popped")
+class _ScopedResponse(object):
+    __slots__ = ("_response", "_hub")
 
     def __init__(self, hub, response):
         self._hub = hub
         self._response = response
-        self._iterator = None
-        self._popped = False
 
     def __iter__(self):
-        try:
-            self._iterator = iter(self._response)
-        except Exception:
-            reraise(*_capture_exception(self._hub))
-        return self
-
-    def __next__(self):
-        if self._iterator is None:
-            self.__iter__()
-
-        try:
-            return next(self._iterator)
-        except StopIteration:
-            raise
-        except Exception:
-            reraise(*_capture_exception(self._hub))
+        iterator = iter(self._response)
+
+        while True:
+            with self._hub:
+                try:
+                    chunk = next(iterator)
+                except StopIteration:
+                    break
+                except Exception:
+                    reraise(*_capture_exception(self._hub))
+
+            yield chunk
 
     def close(self):
-        if not self._popped:
-            self._hub.pop_scope_unsafe()
-            self._popped = True
-
-        try:
-            self._response.close()
-        except AttributeError:
-            pass
-        except Exception:
-            reraise(*_capture_exception(self._hub))
+        with self._hub:
+            try:
+                self._response.close()
+            except AttributeError:
+                pass
+            except Exception:
+                reraise(*_capture_exception(self._hub))
 
 
 def _make_wsgi_event_processor(environ):

tests/integrations/conftest.py

@@ -6,16 +6,16 @@
 def capture_exceptions(monkeypatch):
     def inner():
         errors = set()
-        old_capture_event = sentry_sdk.Hub.current.capture_event
+        old_capture_event = sentry_sdk.Hub.capture_event
 
-        def capture_event(event, hint=None):
+        def capture_event(self, event, hint=None):
             if hint:
                 if "exc_info" in hint:
                     error = hint["exc_info"][1]
                     errors.add(error)
-            return old_capture_event(event, hint=hint)
+            return old_capture_event(self, event, hint=hint)
 
-        monkeypatch.setattr(sentry_sdk.Hub.current, "capture_event", capture_event)
+        monkeypatch.setattr(sentry_sdk.Hub, "capture_event", capture_event)
         return errors
 
     return inner

tests/integrations/django/test_basic.py

@@ -12,7 +12,7 @@
 except ImportError:
     from django.core.urlresolvers import reverse
 
-from sentry_sdk import last_event_id, capture_message
+from sentry_sdk import capture_message
 from sentry_sdk.integrations.django import DjangoIntegration
 
 from tests.integrations.django.myapp.wsgi import application
@@ -120,16 +120,17 @@ def test_custom_error_handler_request_context(sentry_init, client, capture_event
     }
 
 
-def test_500(sentry_init, client):
+def test_500(sentry_init, client, capture_events):
     sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
-    old_event_id = last_event_id()
+    events = capture_events()
+
     content, status, headers = client.get("/view-exc")
     assert status.lower() == "500 internal server error"
     content = b"".join(content).decode("utf-8")
-    event_id = last_event_id()
+
+    event, = events
+    event_id = event["event_id"]
     assert content == "Sentry error: %s" % event_id
-    assert event_id is not None
-    assert old_event_id != event_id
 
 
 def test_management_command_raises():

tests/integrations/flask/test_flask.py

@@ -5,11 +5,16 @@
 
 flask = pytest.importorskip("flask")
 
-from flask import Flask, request, abort
+from flask import Flask, Response, request, abort, stream_with_context
 
 from flask_login import LoginManager, login_user
 
-from sentry_sdk import capture_message, capture_exception, last_event_id
+from sentry_sdk import (
+    configure_scope,
+    capture_message,
+    capture_exception,
+    last_event_id,
+)
 from sentry_sdk.integrations.logging import LoggingIntegration
 import sentry_sdk.integrations.flask as flask_sentry
 
@@ -71,7 +76,7 @@ def test_transaction_style(
 @pytest.mark.parametrize("debug", (True, False))
 @pytest.mark.parametrize("testing", (True, False))
 def test_errors(sentry_init, capture_exceptions, capture_events, app, debug, testing):
-    sentry_init(integrations=[flask_sentry.FlaskIntegration()])
+    sentry_init(integrations=[flask_sentry.FlaskIntegration()], debug=True)
 
     app.debug = debug
     app.testing = testing
@@ -460,3 +465,33 @@ def index():
     client.get("/")
 
     assert not events
+
+
+def test_does_not_leak_scope(sentry_init, capture_events, app):
+    sentry_init(integrations=[flask_sentry.FlaskIntegration()])
+    events = capture_events()
+
+    with configure_scope() as scope:
+        scope.set_tag("request_data", False)
+
+    @app.route("/")
+    def index():
+        with configure_scope() as scope:
+            scope.set_tag("request_data", True)
+
+        def generate():
+            for row in range(1000):
+                with configure_scope() as scope:
+                    assert scope._tags["request_data"]
+
+                yield str(row) + "\n"
+
+        return Response(stream_with_context(generate()), mimetype="text/csv")
+
+    client = app.test_client()
+    response = client.get("/")
+    assert response.data.decode() == "".join(str(row) + "\n" for row in range(1000))
+    assert not events
+
+    with configure_scope() as scope:
+        assert not scope._tags["request_data"]

tests/integrations/wsgi/test_wsgi.py

@@ -1,9 +1,7 @@
 from werkzeug.test import Client
 import pytest
 
-from sentry_sdk import Hub
-
-from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware, _ScopePoppingResponse
+from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
 
 
 @pytest.fixture
@@ -32,30 +30,3 @@ def test_basic(sentry_init, crashing_app, capture_events):
         "query_string": "",
         "url": "http://localhost/",
     }
-
-
-def test_calls_close():
-    hub = Hub.current
-    stack_size = len(hub._stack)
-    closes = []
-
-    hub.push_scope()
-
-    class Foo(object):
-        def __iter__(self):
-            yield 1
-            yield 2
-            yield 3
-
-        def close(self):
-            closes.append(1)
-
-    response = _ScopePoppingResponse(hub, Foo())
-    list(response)
-    response.close()
-    response.close()
-    response.close()
-
-    # multiple close calls are just forwarded, but the scope is only popped once
-    assert len(closes) == 3
-    assert len(hub._stack) == stack_size