ref(starlite): Use new scopes API (#2876)

Author: sentrivana

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
 
@@ -20,12 +21,13 @@
         from typing import Any, Dict, List, Optional, Union
         from starlite.types import (  # type: ignore
             ASGIApp,
+            Hint,
             HTTPReceiveMessage,
             HTTPScope,
             Message,
             Middleware,
             Receive,
-            Scope,
+            Scope as StarliteScope,
             Send,
             WebSocketReceiveMessage,
         )
@@ -114,51 +116,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 +177,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", _: "Hint") -> "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 +255,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)