Skip to content

Commit

Permalink
Remove nullcontext from wsgi and asgi (#3983)
Browse files Browse the repository at this point in the history 
mypy was complaining too so just removed it
  • Loading branch information
@sl0thentr0py
sl0thentr0py authored Jan 21, 2025
1 parent ea06dab commit 3c7bb1e
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 71 deletions.
93 changes: 43 additions & 50 deletions sentry_sdk/integrations/asgi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

import asyncio
import inspect
from contextlib import nullcontext
from copy import deepcopy
from functools import partial

Expand Down Expand Up @@ -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:
Expand All @@ -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)

Expand Down
46 changes: 25 additions & 21 deletions sentry_sdk/integrations/wsgi.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import sys
from contextlib import nullcontext
from functools import partial

import sentry_sdk
Expand Down Expand Up @@ -123,50 +122,55 @@ 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,
origin=self.span_origin,
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]
):
# 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
Expand Down

0 comments on commit 3c7bb1e

Please sign in to comment.