Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ref(starlite): Use new scopes API #2876

Merged
merged 2 commits into from
Mar 25, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 105 additions & 103 deletions sentry_sdk/integrations/starlite.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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,
)
Expand Down Expand Up @@ -114,51 +116,50 @@
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:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can use the @encure_... decorator here, no?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd do this after we've fixed it and we're sure it works as intended, then we can go over the async integrations and add it.

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:
Expand All @@ -176,66 +177,67 @@
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)
Comment on lines +180 to 184
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also use the @ensure_... decorator here.


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"]

Check warning on line 203 in sentry_sdk/integrations/starlite.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlite.py#L203

Added line #L203 was not covered by tests
if request_data is not None:
request_info["data"] = request_data

Check warning on line 205 in sentry_sdk/integrations/starlite.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlite.py#L205

Added line #L205 was not covered by tests

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

Check warning on line 213 in sentry_sdk/integrations/starlite.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlite.py#L213

Added line #L213 was not covered by tests
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}

Check warning on line 221 in sentry_sdk/integrations/starlite.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlite.py#L220-L221

Added lines #L220 - L221 were not covered by tests

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
Expand All @@ -253,22 +255,22 @@
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)

Check warning on line 268 in sentry_sdk/integrations/starlite.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlite.py#L267-L268

Added lines #L267 - L268 were not covered by tests

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)
Loading