Make code a bit nicer

Author: antonpirker

sentry_sdk/integrations/bottle.py

@@ -2,6 +2,7 @@
 
 import sentry_sdk
 from sentry_sdk.tracing import SOURCE_FOR_STYLE
+from sentry_sdk.tracing_utils import finish_running_transaction
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -139,7 +140,7 @@ def wrapped_callback(*args, **kwargs):
                     _capture_exception(res, handled=True)
 
                 # Manually close the transaction because Bottle does not call `close()` on the WSGI response
-                sentry_sdk.get_current_scope().transaction.__exit__(None, None, None)
+                finish_running_transaction()
 
                 return res
 

sentry_sdk/integrations/flask.py

@@ -7,6 +7,7 @@
 from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.tracing import SOURCE_FOR_STYLE
+from sentry_sdk.tracing_utils import finish_running_transaction
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -156,7 +157,7 @@ def _request_started(app, **kwargs):
 
 def _request_finished(sender, response, **kwargs):
     # Manually close the transaction because Bottle does not call `close()` on the WSGI response
-    sentry_sdk.get_current_scope().transaction.__exit__(None, None, None)
+    finish_running_transaction()
 
 
 class FlaskRequestExtractor(RequestExtractor):

sentry_sdk/integrations/wsgi.py

@@ -9,10 +9,10 @@
 from sentry_sdk.integrations._wsgi_common import (
     DEFAULT_HTTP_METHODS_TO_CAPTURE,
     _filter_headers,
-    nullcontext,
 )
 from sentry_sdk.sessions import track_session
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_ROUTE
+from sentry_sdk.tracing_utils import finish_running_transaction
 from sentry_sdk.utils import (
     ContextVar,
     capture_internal_exceptions,
@@ -262,18 +262,16 @@ def __iter__(self):
                     try:
                         chunk = next(iterator)
                     except StopIteration:
+                        # Close the Sentry transaction (it could be that response.close() is never called by the framework)
+                        # This is done here to make sure the Transaction stays
+                        # open until all streaming responses are done.
+                        finish_running_transaction()
                         break
                     except BaseException:
                         reraise(*_capture_exception())
 
             yield chunk
 
-        # Close the Sentry transaction (it could be that response.close() is never called by the framework)
-        # This is done here to make sure the Transaction stays
-        # open until all streaming responses are done.
-        if self._transaction is not None and hasattr(self._transaction, "_context_manager_state"):
-            self._transaction.__exit__(None, None, None)
-
     def close(self):
         # type: () -> None
         with use_isolation_scope(self._isolation_scope):
@@ -284,8 +282,7 @@ def close(self):
                     # Close the Sentry transaction
                     # This is done here to make sure the Transaction stays
                     # open until all streaming responses are done.
-                    if self._transaction is not None and hasattr(self._transaction, "_context_manager_state"):
-                        self._transaction.__exit__(None, None, None)
+                    finish_running_transaction()
                 except AttributeError:
                     pass
                 except BaseException:

sentry_sdk/tracing_utils.py

@@ -724,3 +724,10 @@ def get_current_span(scope=None):
 
 if TYPE_CHECKING:
     from sentry_sdk.tracing import Span
+
+
+def finish_running_transaction():
+    # type: () -> None
+    current_scope = sentry_sdk.get_current_scope()
+    if current_scope._transaction is not None and hasattr(current_scope._transaction, "_context_manager_state"):
+        current_scope._transaction.__exit__(None, None, None)