From 1425acf76e780ad49058741fd61641d1a6212911 Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Mon, 20 Jan 2025 20:18:56 +0100 Subject: [PATCH] Remove nullcontext from wsgi and asgi --- sentry_sdk/integrations/asgi.py | 93 +++++++++++++++------------------ sentry_sdk/integrations/wsgi.py | 46 ++++++++-------- 2 files changed, 68 insertions(+), 71 deletions(-) diff --git a/sentry_sdk/integrations/asgi.py b/sentry_sdk/integrations/asgi.py index 4a3fe830eb..2a8bbe5091 100644 --- a/sentry_sdk/integrations/asgi.py +++ b/sentry_sdk/integrations/asgi.py @@ -6,7 +6,6 @@ import asyncio import inspect -from contextlib import nullcontext from copy import deepcopy from functools import partial @@ -169,20 +168,24 @@ async def _run_asgi3(self, scope, receive, send): # type: (Any, Any, Any) -> Any return await self._run_app(scope, receive, send, asgi_version=3) + async def _run_original_app(self, scope, receive, send, asgi_version): + # type: (Any, Any, Any, Any, int) -> Any + try: + if asgi_version == 2: + return await self.app(scope)(receive, send) + else: + return await self.app(scope, receive, send) + + except Exception as exc: + _capture_exception(exc, mechanism_type=self.mechanism_type) + raise exc from None + async def _run_app(self, scope, receive, send, asgi_version): # type: (Any, Any, Any, Any, int) -> Any is_recursive_asgi_middleware = _asgi_middleware_applied.get(False) is_lifespan = scope["type"] == "lifespan" if is_recursive_asgi_middleware or is_lifespan: - try: - if asgi_version == 2: - return await self.app(scope)(receive, send) - else: - return await self.app(scope, receive, send) - - except Exception as exc: - _capture_exception(exc, mechanism_type=self.mechanism_type) - raise exc from None + return await self._run_original_app(scope, receive, send, asgi_version) _asgi_middleware_applied.set(True) try: @@ -209,52 +212,42 @@ async def _run_app(self, scope, receive, send, asgi_version): method = scope.get("method", "").upper() should_trace = method in self.http_methods_to_capture + if not should_trace: + return await self._run_original_app( + scope, receive, send, asgi_version + ) + with sentry_sdk.continue_trace(_get_headers(scope)): - with ( - sentry_sdk.start_span( - op=( - OP.WEBSOCKET_SERVER - if ty == "websocket" - else OP.HTTP_SERVER - ), - name=transaction_name, - source=transaction_source, - origin=self.span_origin, - attributes=_prepopulate_attributes(scope), - ) - if should_trace - else nullcontext() + with sentry_sdk.start_span( + op=( + OP.WEBSOCKET_SERVER + if ty == "websocket" + else OP.HTTP_SERVER + ), + name=transaction_name, + source=transaction_source, + origin=self.span_origin, + attributes=_prepopulate_attributes(scope), ) as span: if span is not None: logger.debug("[ASGI] Started transaction: %s", span) span.set_tag("asgi.type", ty) - try: - - async def _sentry_wrapped_send(event): - # type: (Dict[str, Any]) -> Any - is_http_response = ( - event.get("type") == "http.response.start" - and span is not None - and "status" in event - ) - if is_http_response: - span.set_http_status(event["status"]) - - return await send(event) - - if asgi_version == 2: - return await self.app(scope)( - receive, _sentry_wrapped_send - ) - else: - return await self.app( - scope, receive, _sentry_wrapped_send - ) - except Exception as exc: - _capture_exception( - exc, mechanism_type=self.mechanism_type + + async def _sentry_wrapped_send(event): + # type: (Dict[str, Any]) -> Any + is_http_response = ( + event.get("type") == "http.response.start" + and span is not None + and "status" in event ) - raise exc from None + if is_http_response: + span.set_http_status(event["status"]) + + return await send(event) + + return await self._run_original_app( + scope, receive, _sentry_wrapped_send, asgi_version + ) finally: _asgi_middleware_applied.set(False) diff --git a/sentry_sdk/integrations/wsgi.py b/sentry_sdk/integrations/wsgi.py index e9cc65d716..74051df0db 100644 --- a/sentry_sdk/integrations/wsgi.py +++ b/sentry_sdk/integrations/wsgi.py @@ -1,5 +1,4 @@ import sys -from contextlib import nullcontext from functools import partial import sentry_sdk @@ -123,9 +122,9 @@ def __call__(self, environ, start_response): ) method = environ.get("REQUEST_METHOD", "").upper() should_trace = method in self.http_methods_to_capture - with sentry_sdk.continue_trace(environ): - with ( - sentry_sdk.start_span( + if should_trace: + with sentry_sdk.continue_trace(environ): + with sentry_sdk.start_span( op=OP.HTTP_SERVER, name=DEFAULT_TRANSACTION_NAME, source=TRANSACTION_SOURCE_ROUTE, @@ -133,31 +132,36 @@ def __call__(self, environ, start_response): attributes=_prepopulate_attributes( environ, self.use_x_forwarded_for ), - ) - if should_trace - else nullcontext() - ) as transaction: - try: - response = self.app( - environ, - partial( - _sentry_start_response, - start_response, - transaction, - ), + ) as span: + response = self._run_original_app( + environ, start_response, span ) - except BaseException: - reraise(*_capture_exception()) + else: + response = self._run_original_app(environ, start_response, None) finally: _wsgi_middleware_applied.set(False) return _ScopedResponse(scope, response) + def _run_original_app(self, environ, start_response, span): + # type: (dict[str, str], StartResponse, Optional[Span]) -> Any + try: + return self.app( + environ, + partial( + _sentry_start_response, + start_response, + span, + ), + ) + except BaseException: + reraise(*_capture_exception()) + def _sentry_start_response( # type: ignore old_start_response, # type: StartResponse - transaction, # type: Optional[Span] + span, # type: Optional[Span] status, # type: str response_headers, # type: WsgiResponseHeaders exc_info=None, # type: Optional[WsgiExcInfo] @@ -165,8 +169,8 @@ def _sentry_start_response( # type: ignore # type: (...) -> WsgiResponseIter with capture_internal_exceptions(): status_int = int(status.split(" ", 1)[0]) - if transaction is not None: - transaction.set_http_status(status_int) + if span is not None: + span.set_http_status(status_int) if exc_info is None: # The Django Rest Framework WSGI test client, and likely other