|
| 1 | +import sentry_sdk |
| 2 | +from sentry_sdk._types import TYPE_CHECKING |
| 3 | +from sentry_sdk.consts import OP |
| 4 | +from sentry_sdk.integrations import DidNotEnable, Integration |
| 5 | +from sentry_sdk.integrations.asgi import SentryAsgiMiddleware |
| 6 | +from sentry_sdk.integrations.logging import ignore_logger |
| 7 | +from sentry_sdk.scope import should_send_default_pii |
| 8 | +from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE |
| 9 | +from sentry_sdk.utils import ( |
| 10 | + ensure_integration_enabled, |
| 11 | + event_from_exception, |
| 12 | + transaction_from_function, |
| 13 | +) |
| 14 | + |
| 15 | +try: |
| 16 | + from litestar import Request, Litestar # type: ignore |
| 17 | + from litestar.handlers.base import BaseRouteHandler # type: ignore |
| 18 | + from litestar.middleware import DefineMiddleware # type: ignore |
| 19 | + from litestar.routes.http import HTTPRoute # type: ignore |
| 20 | + from litestar.data_extractors import ConnectionDataExtractor # type: ignore |
| 21 | +except ImportError: |
| 22 | + raise DidNotEnable("Litestar is not installed") |
| 23 | +if TYPE_CHECKING: |
| 24 | + from typing import Any, Optional, Union |
| 25 | + from litestar.types.asgi_types import ASGIApp # type: ignore |
| 26 | + from litestar.types import ( # type: ignore |
| 27 | + HTTPReceiveMessage, |
| 28 | + HTTPScope, |
| 29 | + Message, |
| 30 | + Middleware, |
| 31 | + Receive, |
| 32 | + Scope as LitestarScope, |
| 33 | + Send, |
| 34 | + WebSocketReceiveMessage, |
| 35 | + ) |
| 36 | + from litestar.middleware import MiddlewareProtocol |
| 37 | + from sentry_sdk._types import Event, Hint |
| 38 | + |
| 39 | +_DEFAULT_TRANSACTION_NAME = "generic Litestar request" |
| 40 | + |
| 41 | + |
| 42 | +class LitestarIntegration(Integration): |
| 43 | + identifier = "litestar" |
| 44 | + origin = f"auto.http.{identifier}" |
| 45 | + |
| 46 | + @staticmethod |
| 47 | + def setup_once(): |
| 48 | + # type: () -> None |
| 49 | + patch_app_init() |
| 50 | + patch_middlewares() |
| 51 | + patch_http_route_handle() |
| 52 | + |
| 53 | + # The following line follows the pattern found in other integrations such as `DjangoIntegration.setup_once`. |
| 54 | + # The Litestar `ExceptionHandlerMiddleware.__call__` catches exceptions and does the following |
| 55 | + # (among other things): |
| 56 | + # 1. Logs them, some at least (such as 500s) as errors |
| 57 | + # 2. Calls after_exception hooks |
| 58 | + # The `LitestarIntegration`` provides an after_exception hook (see `patch_app_init` below) to create a Sentry event |
| 59 | + # from an exception, which ends up being called during step 2 above. However, the Sentry `LoggingIntegration` will |
| 60 | + # by default create a Sentry event from error logs made in step 1 if we do not prevent it from doing so. |
| 61 | + ignore_logger("litestar") |
| 62 | + |
| 63 | + |
| 64 | +class SentryLitestarASGIMiddleware(SentryAsgiMiddleware): |
| 65 | + def __init__(self, app, span_origin=LitestarIntegration.origin): |
| 66 | + # type: (ASGIApp, str) -> None |
| 67 | + |
| 68 | + super().__init__( |
| 69 | + app=app, |
| 70 | + unsafe_context_data=False, |
| 71 | + transaction_style="endpoint", |
| 72 | + mechanism_type="asgi", |
| 73 | + span_origin=span_origin, |
| 74 | + ) |
| 75 | + |
| 76 | + |
| 77 | +def patch_app_init(): |
| 78 | + # type: () -> None |
| 79 | + """ |
| 80 | + Replaces the Litestar class's `__init__` function in order to inject `after_exception` handlers and set the |
| 81 | + `SentryLitestarASGIMiddleware` as the outmost middleware in the stack. |
| 82 | + See: |
| 83 | + - https://docs.litestar.dev/2/usage/applications.html#after-exception |
| 84 | + - https://docs.litestar.dev/2/usage/middleware/using-middleware.html |
| 85 | + """ |
| 86 | + old__init__ = Litestar.__init__ |
| 87 | + |
| 88 | + @ensure_integration_enabled(LitestarIntegration, old__init__) |
| 89 | + def injection_wrapper(self, *args, **kwargs): |
| 90 | + # type: (Litestar, *Any, **Any) -> None |
| 91 | + kwargs["after_exception"] = [ |
| 92 | + exception_handler, |
| 93 | + *(kwargs.get("after_exception") or []), |
| 94 | + ] |
| 95 | + |
| 96 | + SentryLitestarASGIMiddleware.__call__ = SentryLitestarASGIMiddleware._run_asgi3 # type: ignore |
| 97 | + middleware = kwargs.get("middleware") or [] |
| 98 | + kwargs["middleware"] = [SentryLitestarASGIMiddleware, *middleware] |
| 99 | + old__init__(self, *args, **kwargs) |
| 100 | + |
| 101 | + Litestar.__init__ = injection_wrapper |
| 102 | + |
| 103 | + |
| 104 | +def patch_middlewares(): |
| 105 | + # type: () -> None |
| 106 | + old_resolve_middleware_stack = BaseRouteHandler.resolve_middleware |
| 107 | + |
| 108 | + @ensure_integration_enabled(LitestarIntegration, old_resolve_middleware_stack) |
| 109 | + def resolve_middleware_wrapper(self): |
| 110 | + # type: (BaseRouteHandler) -> list[Middleware] |
| 111 | + return [ |
| 112 | + enable_span_for_middleware(middleware) |
| 113 | + for middleware in old_resolve_middleware_stack(self) |
| 114 | + ] |
| 115 | + |
| 116 | + BaseRouteHandler.resolve_middleware = resolve_middleware_wrapper |
| 117 | + |
| 118 | + |
| 119 | +def enable_span_for_middleware(middleware): |
| 120 | + # type: (Middleware) -> Middleware |
| 121 | + if ( |
| 122 | + not hasattr(middleware, "__call__") # noqa: B004 |
| 123 | + or middleware is SentryLitestarASGIMiddleware |
| 124 | + ): |
| 125 | + return middleware |
| 126 | + |
| 127 | + if isinstance(middleware, DefineMiddleware): |
| 128 | + old_call = middleware.middleware.__call__ # type: ASGIApp |
| 129 | + else: |
| 130 | + old_call = middleware.__call__ |
| 131 | + |
| 132 | + async def _create_span_call(self, scope, receive, send): |
| 133 | + # type: (MiddlewareProtocol, LitestarScope, Receive, Send) -> None |
| 134 | + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: |
| 135 | + return await old_call(self, scope, receive, send) |
| 136 | + |
| 137 | + middleware_name = self.__class__.__name__ |
| 138 | + with sentry_sdk.start_span( |
| 139 | + op=OP.MIDDLEWARE_LITESTAR, |
| 140 | + description=middleware_name, |
| 141 | + origin=LitestarIntegration.origin, |
| 142 | + ) as middleware_span: |
| 143 | + middleware_span.set_tag("litestar.middleware_name", middleware_name) |
| 144 | + |
| 145 | + # Creating spans for the "receive" callback |
| 146 | + async def _sentry_receive(*args, **kwargs): |
| 147 | + # type: (*Any, **Any) -> Union[HTTPReceiveMessage, WebSocketReceiveMessage] |
| 148 | + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: |
| 149 | + return await receive(*args, **kwargs) |
| 150 | + with sentry_sdk.start_span( |
| 151 | + op=OP.MIDDLEWARE_LITESTAR_RECEIVE, |
| 152 | + description=getattr(receive, "__qualname__", str(receive)), |
| 153 | + origin=LitestarIntegration.origin, |
| 154 | + ) as span: |
| 155 | + span.set_tag("litestar.middleware_name", middleware_name) |
| 156 | + return await receive(*args, **kwargs) |
| 157 | + |
| 158 | + receive_name = getattr(receive, "__name__", str(receive)) |
| 159 | + receive_patched = receive_name == "_sentry_receive" |
| 160 | + new_receive = _sentry_receive if not receive_patched else receive |
| 161 | + |
| 162 | + # Creating spans for the "send" callback |
| 163 | + async def _sentry_send(message): |
| 164 | + # type: (Message) -> None |
| 165 | + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: |
| 166 | + return await send(message) |
| 167 | + with sentry_sdk.start_span( |
| 168 | + op=OP.MIDDLEWARE_LITESTAR_SEND, |
| 169 | + description=getattr(send, "__qualname__", str(send)), |
| 170 | + origin=LitestarIntegration.origin, |
| 171 | + ) as span: |
| 172 | + span.set_tag("litestar.middleware_name", middleware_name) |
| 173 | + return await send(message) |
| 174 | + |
| 175 | + send_name = getattr(send, "__name__", str(send)) |
| 176 | + send_patched = send_name == "_sentry_send" |
| 177 | + new_send = _sentry_send if not send_patched else send |
| 178 | + |
| 179 | + return await old_call(self, scope, new_receive, new_send) |
| 180 | + |
| 181 | + not_yet_patched = old_call.__name__ not in ["_create_span_call"] |
| 182 | + |
| 183 | + if not_yet_patched: |
| 184 | + if isinstance(middleware, DefineMiddleware): |
| 185 | + middleware.middleware.__call__ = _create_span_call |
| 186 | + else: |
| 187 | + middleware.__call__ = _create_span_call |
| 188 | + |
| 189 | + return middleware |
| 190 | + |
| 191 | + |
| 192 | +def patch_http_route_handle(): |
| 193 | + # type: () -> None |
| 194 | + old_handle = HTTPRoute.handle |
| 195 | + |
| 196 | + async def handle_wrapper(self, scope, receive, send): |
| 197 | + # type: (HTTPRoute, HTTPScope, Receive, Send) -> None |
| 198 | + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: |
| 199 | + return await old_handle(self, scope, receive, send) |
| 200 | + |
| 201 | + sentry_scope = sentry_sdk.get_isolation_scope() |
| 202 | + request = scope["app"].request_class( |
| 203 | + scope=scope, receive=receive, send=send |
| 204 | + ) # type: Request[Any, Any] |
| 205 | + extracted_request_data = ConnectionDataExtractor( |
| 206 | + parse_body=True, parse_query=True |
| 207 | + )(request) |
| 208 | + body = extracted_request_data.pop("body") |
| 209 | + |
| 210 | + request_data = await body |
| 211 | + |
| 212 | + def event_processor(event, _): |
| 213 | + # type: (Event, Hint) -> Event |
| 214 | + route_handler = scope.get("route_handler") |
| 215 | + |
| 216 | + request_info = event.get("request", {}) |
| 217 | + request_info["content_length"] = len(scope.get("_body", b"")) |
| 218 | + if should_send_default_pii(): |
| 219 | + request_info["cookies"] = extracted_request_data["cookies"] |
| 220 | + if request_data is not None: |
| 221 | + request_info["data"] = request_data |
| 222 | + |
| 223 | + func = None |
| 224 | + if route_handler.name is not None: |
| 225 | + tx_name = route_handler.name |
| 226 | + # Accounts for use of type `Ref` in earlier versions of litestar without the need to reference it as a type |
| 227 | + elif hasattr(route_handler.fn, "value"): |
| 228 | + func = route_handler.fn.value |
| 229 | + else: |
| 230 | + func = route_handler.fn |
| 231 | + if func is not None: |
| 232 | + tx_name = transaction_from_function(func) |
| 233 | + |
| 234 | + tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]} |
| 235 | + |
| 236 | + if not tx_name: |
| 237 | + tx_name = _DEFAULT_TRANSACTION_NAME |
| 238 | + tx_info = {"source": TRANSACTION_SOURCE_ROUTE} |
| 239 | + |
| 240 | + event.update( |
| 241 | + { |
| 242 | + "request": request_info, |
| 243 | + "transaction": tx_name, |
| 244 | + "transaction_info": tx_info, |
| 245 | + } |
| 246 | + ) |
| 247 | + return event |
| 248 | + |
| 249 | + sentry_scope._name = LitestarIntegration.identifier |
| 250 | + sentry_scope.add_event_processor(event_processor) |
| 251 | + |
| 252 | + return await old_handle(self, scope, receive, send) |
| 253 | + |
| 254 | + HTTPRoute.handle = handle_wrapper |
| 255 | + |
| 256 | + |
| 257 | +def retrieve_user_from_scope(scope): |
| 258 | + # type: (LitestarScope) -> Optional[dict[str, Any]] |
| 259 | + scope_user = scope.get("user") |
| 260 | + if isinstance(scope_user, dict): |
| 261 | + return scope_user |
| 262 | + if hasattr(scope_user, "asdict"): # dataclasses |
| 263 | + return scope_user.asdict() |
| 264 | + |
| 265 | + return None |
| 266 | + |
| 267 | + |
| 268 | +@ensure_integration_enabled(LitestarIntegration) |
| 269 | +def exception_handler(exc, scope): |
| 270 | + # type: (Exception, LitestarScope) -> None |
| 271 | + user_info = None # type: Optional[dict[str, Any]] |
| 272 | + if should_send_default_pii(): |
| 273 | + user_info = retrieve_user_from_scope(scope) |
| 274 | + if user_info and isinstance(user_info, dict): |
| 275 | + sentry_scope = sentry_sdk.get_isolation_scope() |
| 276 | + sentry_scope.set_user(user_info) |
| 277 | + |
| 278 | + event, hint = event_from_exception( |
| 279 | + exc, |
| 280 | + client_options=sentry_sdk.get_client().options, |
| 281 | + mechanism={"type": LitestarIntegration.identifier, "handled": False}, |
| 282 | + ) |
| 283 | + |
| 284 | + sentry_sdk.capture_event(event, hint=hint) |
0 commit comments