feat(integrations): Update StarliteIntegration to be more in line with new LitestarIntegration (#3384)

The new LitestarIntegration was initially ported from the StarliteIntegration, but then had a thorough code review that resulted in use of type comments instead of type hints (the convention used throughout the repo), more concise code in several places, and additional/updated tests. This PR backports those improvements to the StarliteIntegration. See #3358. --------- Co-authored-by: Anton Pirker <anton.pirker@sentry.io>

Author: KellyWalker

sentry_sdk/integrations/starlite.py

@@ -1,6 +1,5 @@
-from typing import TYPE_CHECKING
-
 import sentry_sdk
+from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
@@ -20,26 +19,26 @@
  from starlite.routes.http import HTTPRoute # type: ignore
  from starlite.utils import ConnectionDataExtractor, is_async_callable, Ref # type: ignore
  from pydantic import BaseModel # type: ignore
-
- if TYPE_CHECKING:
- from typing import Any, Dict, List, Optional, Union
- from starlite.types import ( # type: ignore
- ASGIApp,
- Hint,
- HTTPReceiveMessage,
- HTTPScope,
- Message,
- Middleware,
- Receive,
- Scope as StarliteScope,
- Send,
- WebSocketReceiveMessage,
- )
- from starlite import MiddlewareProtocol
- from sentry_sdk._types import Event
 except ImportError:
  raise DidNotEnable("Starlite is not installed")
 
+if TYPE_CHECKING:
+ from typing import Any, Optional, Union
+ from starlite.types import ( # type: ignore
+ ASGIApp,
+ Hint,
+ HTTPReceiveMessage,
+ HTTPScope,
+ Message,
+ Middleware,
+ Receive,
+ Scope as StarliteScope,
+ Send,
+ WebSocketReceiveMessage,
+ )
+ from starlite import MiddlewareProtocol
+ from sentry_sdk._types import Event
+
 
 _DEFAULT_TRANSACTION_NAME = "generic Starlite request"
 
@@ -49,14 +48,16 @@ class StarliteIntegration(Integration):
  origin = f"auto.http.{identifier}"
 
  @staticmethod
- def setup_once() -> None:
+ def setup_once():
+ # type: () -> None
  patch_app_init()
  patch_middlewares()
  patch_http_route_handle()
 
 
 class SentryStarliteASGIMiddleware(SentryAsgiMiddleware):
- def __init__(self, app: "ASGIApp", span_origin: str = StarliteIntegration.origin):
+ def __init__(self, app, span_origin=StarliteIntegration.origin):
+ # type: (ASGIApp, str) -> None
  super().__init__(
  app=app,
  unsafe_context_data=False,
@@ -66,7 +67,8 @@ def __init__(self, app: "ASGIApp", span_origin: str = StarliteIntegration.origin
  )
 
 
-def patch_app_init() -> None:
+def patch_app_init():
+ # type: () -> None
  """
  Replaces the Starlite class's `__init__` function in order to inject `after_exception` handlers and set the
  `SentryStarliteASGIMiddleware` as the outmost middleware in the stack.
@@ -76,7 +78,9 @@ def patch_app_init() -> None:
  """
  old__init__ = Starlite.__init__
 
- def injection_wrapper(self: "Starlite", *args: "Any", **kwargs: "Any") -> None:
+ @ensure_integration_enabled(StarliteIntegration, old__init__)
+ def injection_wrapper(self, *args, **kwargs):
+ # type: (Starlite, *Any, **Any) -> None
  after_exception = kwargs.pop("after_exception", [])
  kwargs.update(
  after_exception=[
@@ -90,43 +94,43 @@ def injection_wrapper(self: "Starlite", *args: "Any", **kwargs: "Any") -> None:
  )
 
  SentryStarliteASGIMiddleware.__call__ = SentryStarliteASGIMiddleware._run_asgi3 # type: ignore
- middleware = kwargs.pop("middleware", None) or []
+ middleware = kwargs.get("middleware") or []
  kwargs["middleware"] = [SentryStarliteASGIMiddleware, *middleware]
  old__init__(self, *args, **kwargs)
 
  Starlite.__init__ = injection_wrapper
 
 
-def patch_middlewares() -> None:
- old__resolve_middleware_stack = BaseRouteHandler.resolve_middleware
+def patch_middlewares():
+ # type: () -> None
+ old_resolve_middleware_stack = BaseRouteHandler.resolve_middleware
 
- def resolve_middleware_wrapper(self: "Any") -> "List[Middleware]":
+ @ensure_integration_enabled(StarliteIntegration, old_resolve_middleware_stack)
+ def resolve_middleware_wrapper(self):
+ # type: (BaseRouteHandler) -> list[Middleware]
  return [
  enable_span_for_middleware(middleware)
- for middleware in old__resolve_middleware_stack(self)
+ for middleware in old_resolve_middleware_stack(self)
  ]
 
  BaseRouteHandler.resolve_middleware = resolve_middleware_wrapper
 
 
-def enable_span_for_middleware(middleware: "Middleware") -> "Middleware":
+def enable_span_for_middleware(middleware):
+ # type: (Middleware) -> Middleware
  if (
  not hasattr(middleware, "__call__") # noqa: B004
  or middleware is SentryStarliteASGIMiddleware
  ):
  return middleware
 
  if isinstance(middleware, DefineMiddleware):
- old_call: "ASGIApp" = middleware.middleware.__call__
+ old_call = middleware.middleware.__call__ # type: ASGIApp
  else:
  old_call = middleware.__call__
 
- async def _create_span_call(
- self: "MiddlewareProtocol",
- scope: "StarliteScope",
- receive: "Receive",
- send: "Send",
- ) -> None:
+ async def _create_span_call(self, scope, receive, send):
+ # type: (MiddlewareProtocol, StarliteScope, Receive, Send) -> None
  if sentry_sdk.get_client().get_integration(StarliteIntegration) is None:
  return await old_call(self, scope, receive, send)
 
@@ -139,9 +143,10 @@ async def _create_span_call(
  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]":
+ async def _sentry_receive(*args, **kwargs):
+ # type: (*Any, **Any) -> Union[HTTPReceiveMessage, WebSocketReceiveMessage]
+ if sentry_sdk.get_client().get_integration(StarliteIntegration) is None:
+ return await receive(*args, **kwargs)
  with sentry_sdk.start_span(
  op=OP.MIDDLEWARE_STARLITE_RECEIVE,
  description=getattr(receive, "__qualname__", str(receive)),
@@ -155,7 +160,10 @@ async def _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:
+ async def _sentry_send(message):
+ # type: (Message) -> None
+ if sentry_sdk.get_client().get_integration(StarliteIntegration) is None:
+ return await send(message)
  with sentry_sdk.start_span(
  op=OP.MIDDLEWARE_STARLITE_SEND,
  description=getattr(send, "__qualname__", str(send)),
@@ -181,27 +189,28 @@ async def _sentry_send(message: "Message") -> None:
  return middleware
 
 
-def patch_http_route_handle() -> None:
+def patch_http_route_handle():
+ # type: () -> None
  old_handle = HTTPRoute.handle
 
- async def handle_wrapper(
- self: "HTTPRoute", scope: "HTTPScope", receive: "Receive", send: "Send"
- ) -> None:
+ async def handle_wrapper(self, scope, receive, send):
+ # type: (HTTPRoute, HTTPScope, Receive, Send) -> None
  if sentry_sdk.get_client().get_integration(StarliteIntegration) is None:
  return await old_handle(self, scope, receive, send)
 
  sentry_scope = sentry_sdk.get_isolation_scope()
- request: "Request[Any, Any]" = scope["app"].request_class(
+ request = scope["app"].request_class(
  scope=scope, receive=receive, send=send
- )
+ ) # type: Request[Any, Any]
  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":
+ def event_processor(event, _):
+ # type: (Event, Hint) -> Event
  route_handler = scope.get("route_handler")
 
  request_info = event.get("request", {})
@@ -244,8 +253,9 @@ def event_processor(event: "Event", _: "Hint") -> "Event":
  HTTPRoute.handle = handle_wrapper
 
 
-def retrieve_user_from_scope(scope: "StarliteScope") -> "Optional[Dict[str, Any]]":
- scope_user = scope.get("user", {})
+def retrieve_user_from_scope(scope):
+ # type: (StarliteScope) -> Optional[dict[str, Any]]
+ scope_user = scope.get("user")
  if not scope_user:
  return None
  if isinstance(scope_user, dict):
@@ -263,8 +273,9 @@ def retrieve_user_from_scope(scope: "StarliteScope") -> "Optional[Dict[str, Any]
 
 
 @ensure_integration_enabled(StarliteIntegration)
-def exception_handler(exc: Exception, scope: "StarliteScope", _: "State") -> None:
- user_info: "Optional[Dict[str, Any]]" = None
+def exception_handler(exc, scope, _):
+ # type: (Exception, StarliteScope, State) -> None
+ user_info = None # type: Optional[dict[str, Any]]
  if should_send_default_pii():
  user_info = retrieve_user_from_scope(scope)
  if user_info and isinstance(user_info, dict):