Skip to content

Commit

Permalink
ref(starlette): Use new scopes API
Browse files Browse the repository at this point in the history
  • Loading branch information
sentrivana committed Mar 21, 2024
1 parent 7659554 commit 3a6f564
Showing 1 changed file with 119 additions and 132 deletions.
251 changes: 119 additions & 132 deletions sentry_sdk/integrations/starlette.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@
import functools
from copy import deepcopy

import sentry_sdk
from sentry_sdk._types import TYPE_CHECKING
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._wsgi_common import (
_is_json_content_type,
request_body_within_bounds,
)
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
from sentry_sdk.scope import Scope
from sentry_sdk.scope import Scope, should_send_default_pii
from sentry_sdk.tracing import (
SOURCE_FOR_STYLE,
TRANSACTION_SOURCE_COMPONENT,
Expand All @@ -20,6 +20,8 @@
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
logger,
parse_version,
Expand All @@ -29,7 +31,6 @@
if TYPE_CHECKING:
from typing import Any, Awaitable, Callable, Dict, Optional, Tuple

from sentry_sdk.scope import Scope as SentryScope
from sentry_sdk._types import Event

try:
Expand Down Expand Up @@ -102,60 +103,54 @@ def _enable_span_for_middleware(middleware_class):
# type: (Any) -> type
old_call = middleware_class.__call__

@ensure_integration_enabled_async(StarletteIntegration, old_call)
async def _create_span_call(app, scope, receive, send, **kwargs):
# type: (Any, Dict[str, Any], Callable[[], Awaitable[Dict[str, Any]]], Callable[[Dict[str, Any]], Awaitable[None]], Any) -> None
hub = Hub.current
integration = hub.get_integration(StarletteIntegration)
if integration is not None:
middleware_name = app.__class__.__name__

# Update transaction name with middleware name
name, source = _get_transaction_from_middleware(app, scope, integration)
if name is not None:
Scope.get_current_scope().set_transaction_name(
name,
source=source,
)
integration = sentry_sdk.get_client().get_integration(StarletteIntegration)
middleware_name = app.__class__.__name__

# Update transaction name with middleware name
name, source = _get_transaction_from_middleware(app, scope, integration)
if name is not None:
Scope.get_current_scope().set_transaction_name(

Check warning on line 115 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L115

Added line #L115 was not covered by tests
name,
source=source,
)

with hub.start_span(
op=OP.MIDDLEWARE_STARLETTE, description=middleware_name
) as middleware_span:
middleware_span.set_tag("starlette.middleware_name", middleware_name)

# Creating spans for the "receive" callback
async def _sentry_receive(*args, **kwargs):
# type: (*Any, **Any) -> Any
hub = Hub.current
with hub.start_span(
op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
description=getattr(receive, "__qualname__", str(receive)),
) as span:
span.set_tag("starlette.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(*args, **kwargs):
# type: (*Any, **Any) -> Any
hub = Hub.current
with hub.start_span(
op=OP.MIDDLEWARE_STARLETTE_SEND,
description=getattr(send, "__qualname__", str(send)),
) as span:
span.set_tag("starlette.middleware_name", middleware_name)
return await send(*args, **kwargs)

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(app, scope, new_receive, new_send, **kwargs)
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_STARLETTE, description=middleware_name
) as middleware_span:
middleware_span.set_tag("starlette.middleware_name", middleware_name)

else:
return await old_call(app, scope, receive, send, **kwargs)
# Creating spans for the "receive" callback
async def _sentry_receive(*args, **kwargs):
# type: (*Any, **Any) -> Any
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
description=getattr(receive, "__qualname__", str(receive)),
) as span:
span.set_tag("starlette.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(*args, **kwargs):
# type: (*Any, **Any) -> Any
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_STARLETTE_SEND,
description=getattr(send, "__qualname__", str(send)),
) as span:
span.set_tag("starlette.middleware_name", middleware_name)
return await send(*args, **kwargs)

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(app, scope, new_receive, new_send, **kwargs)

not_yet_patched = old_call.__name__ not in [
"_create_span_call",
Expand All @@ -171,17 +166,17 @@ async def _sentry_send(*args, **kwargs):

def _capture_exception(exception, handled=False):
# type: (BaseException, **Any) -> None
hub = Hub.current
if hub.get_integration(StarletteIntegration) is None:
client = sentry_sdk.get_client()

Check warning on line 169 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L169

Added line #L169 was not covered by tests
if client.get_integration(StarletteIntegration) is None:
return

event, hint = event_from_exception(
exception,
client_options=hub.client.options if hub.client else None,
client_options=client.options,
mechanism={"type": StarletteIntegration.identifier, "handled": handled},
)

hub.capture_event(event, hint=hint)
sentry_sdk.capture_event(event, hint=hint)

Check warning on line 179 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L179

Added line #L179 was not covered by tests


def patch_exception_middleware(middleware_class):
Expand Down Expand Up @@ -265,30 +260,29 @@ def _add_user_to_sentry_scope(scope):
if "user" not in scope:
return

if not _should_send_default_pii():
if not should_send_default_pii():
return

hub = Hub.current
if hub.get_integration(StarletteIntegration) is None:
if sentry_sdk.get_client().get_integration(StarletteIntegration) is None:
return

with hub.configure_scope() as sentry_scope:
user_info = {} # type: Dict[str, Any]
starlette_user = scope["user"]
user_info = {} # type: Dict[str, Any]
starlette_user = scope["user"]

Check warning on line 270 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L269-L270

Added lines #L269 - L270 were not covered by tests

username = getattr(starlette_user, "username", None)
if username:
user_info.setdefault("username", starlette_user.username)
username = getattr(starlette_user, "username", None)

Check warning on line 272 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L272

Added line #L272 was not covered by tests
if username:
user_info.setdefault("username", starlette_user.username)

Check warning on line 274 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L274

Added line #L274 was not covered by tests

user_id = getattr(starlette_user, "id", None)
if user_id:
user_info.setdefault("id", starlette_user.id)
user_id = getattr(starlette_user, "id", None)

Check warning on line 276 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L276

Added line #L276 was not covered by tests
if user_id:
user_info.setdefault("id", starlette_user.id)

Check warning on line 278 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L278

Added line #L278 was not covered by tests

email = getattr(starlette_user, "email", None)
if email:
user_info.setdefault("email", starlette_user.email)
email = getattr(starlette_user, "email", None)

Check warning on line 280 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L280

Added line #L280 was not covered by tests
if email:
user_info.setdefault("email", starlette_user.email)

Check warning on line 282 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L282

Added line #L282 was not covered by tests

sentry_scope.user = user_info
sentry_scope = Scope.get_isolation_scope()
sentry_scope.user = user_info

Check warning on line 285 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L284-L285

Added lines #L284 - L285 were not covered by tests


def patch_authentication_middleware(middleware_class):
Expand Down Expand Up @@ -348,7 +342,7 @@ def patch_asgi_app():

async def _sentry_patched_asgi_app(self, scope, receive, send):
# type: (Starlette, StarletteScope, Receive, Send) -> None
integration = Hub.current.get_integration(StarletteIntegration)
integration = sentry_sdk.get_client().get_integration(StarletteIntegration)
if integration is None:
return await old_app(self, scope, receive, send)

Expand Down Expand Up @@ -387,40 +381,39 @@ def _sentry_request_response(func):
is_coroutine = _is_async_callable(old_func)
if is_coroutine:

@ensure_integration_enabled_async(StarletteIntegration, old_func)
async def _sentry_async_func(*args, **kwargs):
# type: (*Any, **Any) -> Any
hub = Hub.current
integration = hub.get_integration(StarletteIntegration)
if integration is None:
return await old_func(*args, **kwargs)

integration = sentry_sdk.get_client().get_integration(

Check warning on line 387 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L387

Added line #L387 was not covered by tests
StarletteIntegration
)
request = args[0]

_set_transaction_name_and_source(
Scope.get_current_scope(), integration.transaction_style, request
)

with hub.configure_scope() as sentry_scope:
extractor = StarletteRequestExtractor(request)
info = await extractor.extract_request_info()
sentry_scope = Scope.get_isolation_scope()
extractor = StarletteRequestExtractor(request)
info = await extractor.extract_request_info()

Check warning on line 398 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L396-L398

Added lines #L396 - L398 were not covered by tests

def _make_request_event_processor(req, integration):
# type: (Any, Any) -> Callable[[Event, dict[str, Any]], Event]
def event_processor(event, hint):
# type: (Event, Dict[str, Any]) -> Event
def _make_request_event_processor(req, integration):

Check warning on line 400 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L400

Added line #L400 was not covered by tests
# type: (Any, Any) -> Callable[[Event, dict[str, Any]], Event]
def event_processor(event, hint):

Check warning on line 402 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L402

Added line #L402 was not covered by tests
# type: (Event, Dict[str, Any]) -> Event

# Add info from request to event
request_info = event.get("request", {})
if info:
if "cookies" in info:
request_info["cookies"] = info["cookies"]
if "data" in info:
request_info["data"] = info["data"]
event["request"] = deepcopy(request_info)
# Add info from request to event
request_info = event.get("request", {})

Check warning on line 406 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L406

Added line #L406 was not covered by tests
if info:
if "cookies" in info:
request_info["cookies"] = info["cookies"]

Check warning on line 409 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L409

Added line #L409 was not covered by tests
if "data" in info:
request_info["data"] = info["data"]
event["request"] = deepcopy(request_info)

Check warning on line 412 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L411-L412

Added lines #L411 - L412 were not covered by tests

return event
return event

Check warning on line 414 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L414

Added line #L414 was not covered by tests

return event_processor
return event_processor

Check warning on line 416 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L416

Added line #L416 was not covered by tests

sentry_scope._name = StarletteIntegration.identifier
sentry_scope.add_event_processor(
Expand All @@ -430,43 +423,44 @@ def event_processor(event, hint):
return await old_func(*args, **kwargs)

func = _sentry_async_func

else:

@ensure_integration_enabled(StarletteIntegration, old_func)

Check warning on line 429 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L429

Added line #L429 was not covered by tests
def _sentry_sync_func(*args, **kwargs):
# type: (*Any, **Any) -> Any
hub = Hub.current
integration = hub.get_integration(StarletteIntegration)
if integration is None:
return old_func(*args, **kwargs)
integration = sentry_sdk.get_client().get_integration(

Check warning on line 432 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L432

Added line #L432 was not covered by tests
StarletteIntegration
)
sentry_scope = Scope.get_isolation_scope()

Check warning on line 435 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L435

Added line #L435 was not covered by tests

with hub.configure_scope() as sentry_scope:
if sentry_scope.profile is not None:
sentry_scope.profile.update_active_thread_id()
if sentry_scope.profile is not None:
sentry_scope.profile.update_active_thread_id()

Check warning on line 438 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L438

Added line #L438 was not covered by tests

request = args[0]
request = args[0]

Check warning on line 440 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L440

Added line #L440 was not covered by tests

_set_transaction_name_and_source(
sentry_scope, integration.transaction_style, request
)
_set_transaction_name_and_source(

Check warning on line 442 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L442

Added line #L442 was not covered by tests
sentry_scope, integration.transaction_style, request
)

extractor = StarletteRequestExtractor(request)
cookies = extractor.extract_cookies_from_request()
extractor = StarletteRequestExtractor(request)
cookies = extractor.extract_cookies_from_request()

Check warning on line 447 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L446-L447

Added lines #L446 - L447 were not covered by tests

def _make_request_event_processor(req, integration):
# type: (Any, Any) -> Callable[[Event, dict[str, Any]], Event]
def event_processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
def _make_request_event_processor(req, integration):

Check warning on line 449 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L449

Added line #L449 was not covered by tests
# type: (Any, Any) -> Callable[[Event, dict[str, Any]], Event]
def event_processor(event, hint):

Check warning on line 451 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L451

Added line #L451 was not covered by tests
# type: (Event, dict[str, Any]) -> Event

# Extract information from request
request_info = event.get("request", {})
if cookies:
request_info["cookies"] = cookies
# Extract information from request
request_info = event.get("request", {})

Check warning on line 455 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L455

Added line #L455 was not covered by tests
if cookies:
request_info["cookies"] = cookies

Check warning on line 457 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L457

Added line #L457 was not covered by tests

event["request"] = deepcopy(request_info)
event["request"] = deepcopy(request_info)

Check warning on line 459 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L459

Added line #L459 was not covered by tests

return event
return event

Check warning on line 461 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L461

Added line #L461 was not covered by tests

return event_processor
return event_processor

Check warning on line 463 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L463

Added line #L463 was not covered by tests

sentry_scope._name = StarletteIntegration.identifier
sentry_scope.add_event_processor(
Expand Down Expand Up @@ -507,8 +501,7 @@ def _sentry_jinja2templates_init(self, *args, **kwargs):
# type: (Jinja2Templates, *Any, **Any) -> None
def add_sentry_trace_meta(request):
# type: (Request) -> Dict[str, Any]
hub = Hub.current
trace_meta = Markup(hub.trace_propagation_meta())
trace_meta = Markup(Scope.get_current_scope().trace_propagation_meta())

Check warning on line 504 in sentry_sdk/integrations/starlette.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/starlette.py#L504

Added line #L504 was not covered by tests
return {
"sentry_trace_meta": trace_meta,
}
Expand Down Expand Up @@ -537,27 +530,21 @@ def __init__(self, request):

def extract_cookies_from_request(self):
# type: (StarletteRequestExtractor) -> Optional[Dict[str, Any]]
client = Hub.current.client
if client is None:
return None

cookies = None # type: Optional[Dict[str, Any]]
if _should_send_default_pii():
if should_send_default_pii():
cookies = self.cookies()

return cookies

async def extract_request_info(self):
# type: (StarletteRequestExtractor) -> Optional[Dict[str, Any]]
client = Hub.current.client
if client is None:
return None
client = sentry_sdk.get_client()

request_info = {} # type: Dict[str, Any]

with capture_internal_exceptions():
# Add cookies
if _should_send_default_pii():
if should_send_default_pii():
request_info["cookies"] = self.cookies()

# If there is no body, just return the cookies
Expand Down Expand Up @@ -648,7 +635,7 @@ def _transaction_name_from_router(scope):


def _set_transaction_name_and_source(scope, transaction_style, request):
# type: (SentryScope, str, Any) -> None
# type: (Scope, str, Any) -> None
name = None
source = SOURCE_FOR_STYLE[transaction_style]

Expand Down

0 comments on commit 3a6f564

Please sign in to comment.