From 3333386f94dead214df0db9562d087e99a2bbf75 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Thu, 21 Mar 2024 14:15:21 +0100 Subject: [PATCH 1/2] ref(starlite): Use new scopes API --- sentry_sdk/integrations/starlite.py | 207 ++++++++++++++-------------- 1 file changed, 104 insertions(+), 103 deletions(-) diff --git a/sentry_sdk/integrations/starlite.py b/sentry_sdk/integrations/starlite.py index 47a91d495d..d51fcc4bba 100644 --- a/sentry_sdk/integrations/starlite.py +++ b/sentry_sdk/integrations/starlite.py @@ -1,9 +1,10 @@ from typing import TYPE_CHECKING +import sentry_sdk from sentry_sdk.consts import OP -from sentry_sdk.hub import Hub, _should_send_default_pii from sentry_sdk.integrations import DidNotEnable, Integration from sentry_sdk.integrations.asgi import SentryAsgiMiddleware +from sentry_sdk.scope import Scope as SentryScope, should_send_default_pii from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE from sentry_sdk.utils import event_from_exception, transaction_from_function @@ -25,7 +26,7 @@ Message, Middleware, Receive, - Scope, + Scope as StarliteScope, Send, WebSocketReceiveMessage, ) @@ -114,51 +115,50 @@ def enable_span_for_middleware(middleware: "Middleware") -> "Middleware": old_call = middleware.__call__ async def _create_span_call( - self: "MiddlewareProtocol", scope: "Scope", receive: "Receive", send: "Send" + self: "MiddlewareProtocol", + scope: "StarliteScope", + receive: "Receive", + send: "Send", ) -> None: - hub = Hub.current - integration = hub.get_integration(StarliteIntegration) - if integration is not None: - middleware_name = self.__class__.__name__ - with hub.start_span( - op=OP.MIDDLEWARE_STARLITE, description=middleware_name - ) as middleware_span: - middleware_span.set_tag("starlite.middleware_name", middleware_name) - - # Creating spans for the "receive" callback - async def _sentry_receive( - *args: "Any", **kwargs: "Any" - ) -> "Union[HTTPReceiveMessage, WebSocketReceiveMessage]": - hub = Hub.current - with hub.start_span( - op=OP.MIDDLEWARE_STARLITE_RECEIVE, - description=getattr(receive, "__qualname__", str(receive)), - ) as span: - span.set_tag("starlite.middleware_name", middleware_name) - return await receive(*args, **kwargs) - - receive_name = getattr(receive, "__name__", str(receive)) - receive_patched = receive_name == "_sentry_receive" - new_receive = _sentry_receive if not receive_patched else receive - - # Creating spans for the "send" callback - async def _sentry_send(message: "Message") -> None: - hub = Hub.current - with hub.start_span( - op=OP.MIDDLEWARE_STARLITE_SEND, - description=getattr(send, "__qualname__", str(send)), - ) as span: - span.set_tag("starlite.middleware_name", middleware_name) - return await send(message) - - send_name = getattr(send, "__name__", str(send)) - send_patched = send_name == "_sentry_send" - new_send = _sentry_send if not send_patched else send - - return await old_call(self, scope, new_receive, new_send) - else: + if sentry_sdk.get_client().get_integration(StarliteIntegration) is None: return await old_call(self, scope, receive, send) + middleware_name = self.__class__.__name__ + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE, description=middleware_name + ) as middleware_span: + middleware_span.set_tag("starlite.middleware_name", middleware_name) + + # Creating spans for the "receive" callback + async def _sentry_receive( + *args: "Any", **kwargs: "Any" + ) -> "Union[HTTPReceiveMessage, WebSocketReceiveMessage]": + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE_RECEIVE, + description=getattr(receive, "__qualname__", str(receive)), + ) as span: + span.set_tag("starlite.middleware_name", middleware_name) + return await receive(*args, **kwargs) + + receive_name = getattr(receive, "__name__", str(receive)) + receive_patched = receive_name == "_sentry_receive" + new_receive = _sentry_receive if not receive_patched else receive + + # Creating spans for the "send" callback + async def _sentry_send(message: "Message") -> None: + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE_SEND, + description=getattr(send, "__qualname__", str(send)), + ) as span: + span.set_tag("starlite.middleware_name", middleware_name) + return await send(message) + + send_name = getattr(send, "__name__", str(send)) + send_patched = send_name == "_sentry_send" + new_send = _sentry_send if not send_patched else send + + return await old_call(self, scope, new_receive, new_send) + not_yet_patched = old_call.__name__ not in ["_create_span_call"] if not_yet_patched: @@ -176,66 +176,67 @@ def patch_http_route_handle() -> None: async def handle_wrapper( self: "HTTPRoute", scope: "HTTPScope", receive: "Receive", send: "Send" ) -> None: - hub = Hub.current - integration: StarliteIntegration = hub.get_integration(StarliteIntegration) + integration: StarliteIntegration = sentry_sdk.get_client().get_integration( + StarliteIntegration + ) if integration is None: return await old_handle(self, scope, receive, send) - with hub.configure_scope() as sentry_scope: - request: "Request[Any, Any]" = scope["app"].request_class( - scope=scope, receive=receive, send=send + sentry_scope = SentryScope.get_isolation_scope() + request: "Request[Any, Any]" = scope["app"].request_class( + scope=scope, receive=receive, send=send + ) + extracted_request_data = ConnectionDataExtractor( + parse_body=True, parse_query=True + )(request) + body = extracted_request_data.pop("body") + + request_data = await body + + def event_processor(event: "Event", _: "Dict[str, Any]") -> "Event": + route_handler = scope.get("route_handler") + + request_info = event.get("request", {}) + request_info["content_length"] = len(scope.get("_body", b"")) + if should_send_default_pii(): + request_info["cookies"] = extracted_request_data["cookies"] + if request_data is not None: + request_info["data"] = request_data + + func = None + if route_handler.name is not None: + tx_name = route_handler.name + elif isinstance(route_handler.fn, Ref): + func = route_handler.fn.value + else: + func = route_handler.fn + if func is not None: + tx_name = transaction_from_function(func) + + tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]} + + if not tx_name: + tx_name = _DEFAULT_TRANSACTION_NAME + tx_info = {"source": TRANSACTION_SOURCE_ROUTE} + + event.update( + { + "request": request_info, + "transaction": tx_name, + "transaction_info": tx_info, + } ) - extracted_request_data = ConnectionDataExtractor( - parse_body=True, parse_query=True - )(request) - body = extracted_request_data.pop("body") - - request_data = await body - - def event_processor(event: "Event", _: "Dict[str, Any]") -> "Event": - route_handler = scope.get("route_handler") - - request_info = event.get("request", {}) - request_info["content_length"] = len(scope.get("_body", b"")) - if _should_send_default_pii(): - request_info["cookies"] = extracted_request_data["cookies"] - if request_data is not None: - request_info["data"] = request_data - - func = None - if route_handler.name is not None: - tx_name = route_handler.name - elif isinstance(route_handler.fn, Ref): - func = route_handler.fn.value - else: - func = route_handler.fn - if func is not None: - tx_name = transaction_from_function(func) - - tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]} - - if not tx_name: - tx_name = _DEFAULT_TRANSACTION_NAME - tx_info = {"source": TRANSACTION_SOURCE_ROUTE} - - event.update( - { - "request": request_info, - "transaction": tx_name, - "transaction_info": tx_info, - } - ) - return event - - sentry_scope._name = StarliteIntegration.identifier - sentry_scope.add_event_processor(event_processor) + return event - return await old_handle(self, scope, receive, send) + sentry_scope._name = StarliteIntegration.identifier + sentry_scope.add_event_processor(event_processor) + + return await old_handle(self, scope, receive, send) HTTPRoute.handle = handle_wrapper -def retrieve_user_from_scope(scope: "Scope") -> "Optional[Dict[str, Any]]": +def retrieve_user_from_scope(scope: "StarliteScope") -> "Optional[Dict[str, Any]]": scope_user = scope.get("user", {}) if not scope_user: return None @@ -253,22 +254,22 @@ def retrieve_user_from_scope(scope: "Scope") -> "Optional[Dict[str, Any]]": return None -def exception_handler(exc: Exception, scope: "Scope", _: "State") -> None: - hub = Hub.current - if hub.get_integration(StarliteIntegration) is None: +def exception_handler(exc: Exception, scope: "StarliteScope", _: "State") -> None: + client = sentry_sdk.get_client() + if client.get_integration(StarliteIntegration) is None: return user_info: "Optional[Dict[str, Any]]" = None - if _should_send_default_pii(): + if should_send_default_pii(): user_info = retrieve_user_from_scope(scope) if user_info and isinstance(user_info, dict): - with hub.configure_scope() as sentry_scope: - sentry_scope.set_user(user_info) + sentry_scope = SentryScope.get_isolation_scope() + sentry_scope.set_user(user_info) event, hint = event_from_exception( exc, - client_options=hub.client.options if hub.client else None, + client_options=client.options, mechanism={"type": StarliteIntegration.identifier, "handled": False}, ) - hub.capture_event(event, hint=hint) + sentry_sdk.capture_event(event, hint=hint) From 0b41b5b305946c2ab222bc22fab81265981e5ff3 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Fri, 22 Mar 2024 11:05:26 +0100 Subject: [PATCH 2/2] nicer type --- sentry_sdk/integrations/starlite.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/starlite.py b/sentry_sdk/integrations/starlite.py index d51fcc4bba..1ee2e479ea 100644 --- a/sentry_sdk/integrations/starlite.py +++ b/sentry_sdk/integrations/starlite.py @@ -21,6 +21,7 @@ from typing import Any, Dict, List, Optional, Union from starlite.types import ( # type: ignore ASGIApp, + Hint, HTTPReceiveMessage, HTTPScope, Message, @@ -193,7 +194,7 @@ async def handle_wrapper( request_data = await body - def event_processor(event: "Event", _: "Dict[str, Any]") -> "Event": + def event_processor(event: "Event", _: "Hint") -> "Event": route_handler = scope.get("route_handler") request_info = event.get("request", {})