Skip to content

Commit 7d004f0

Browse files
antonpirkerantonpirker
authored
Have instrumentation for ASGI middleware receive/send callbacks. (#1673)
* Have instrumentation for ASGI middleware receive/send callbacks. * Added tests for new callback spans.
1 parent 9e1e760 commit 7d004f0

File tree

4 files changed

+136
-6
lines changed

4 files changed

+136
-6
lines changed

‎sentry_sdk/consts.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ class OP:
118118
HTTP_SERVER = "http.server"
119119
MIDDLEWARE_DJANGO = "middleware.django"
120120
MIDDLEWARE_STARLETTE = "middleware.starlette"
121+
MIDDLEWARE_STARLETTE_RECEIVE = "middleware.starlette.receive"
122+
MIDDLEWARE_STARLETTE_SEND = "middleware.starlette.send"
121123
QUEUE_SUBMIT_CELERY = "queue.submit.celery"
122124
QUEUE_TASK_CELERY = "queue.task.celery"
123125
QUEUE_TASK_RQ = "queue.task.rq"

‎sentry_sdk/integrations/starlette.py‎

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,21 +85,49 @@ def _enable_span_for_middleware(middleware_class):
8585
# type: (Any) -> type
8686
old_call = middleware_class.__call__
8787

88-
async def _create_span_call(*args, **kwargs):
89-
# type: (Any, Any) -> None
88+
async def _create_span_call(app, scope, receive, send, **kwargs):
89+
# type: (Any, Dict[str, Any], Callable[[], Awaitable[Dict[str, Any]]], Callable[[Dict[str, Any]], Awaitable[None]], Any) -> None
9090
hub = Hub.current
9191
integration = hub.get_integration(StarletteIntegration)
9292
if integration is not None:
93-
middleware_name = args[0].__class__.__name__
93+
middleware_name = app.__class__.__name__
94+
9495
with hub.start_span(
9596
op=OP.MIDDLEWARE_STARLETTE, description=middleware_name
9697
) as middleware_span:
9798
middleware_span.set_tag("starlette.middleware_name", middleware_name)
9899

99-
await old_call(*args, **kwargs)
100+
# Creating spans for the "receive" callback
101+
async def _sentry_receive(*args, **kwargs):
102+
# type: (*Any, **Any) -> Any
103+
hub = Hub.current
104+
with hub.start_span(
105+
op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
106+
description=receive.__qualname__,
107+
) as span:
108+
span.set_tag("starlette.middleware_name", middleware_name)
109+
await receive(*args, **kwargs)
110+
111+
receive_patched = receive.__name__ == "_sentry_receive"
112+
new_receive = _sentry_receive if not receive_patched else receive
113+
114+
# Creating spans for the "send" callback
115+
async def _sentry_send(*args, **kwargs):
116+
# type: (*Any, **Any) -> Any
117+
hub = Hub.current
118+
with hub.start_span(
119+
op=OP.MIDDLEWARE_STARLETTE_SEND, description=send.__qualname__
120+
) as span:
121+
span.set_tag("starlette.middleware_name", middleware_name)
122+
await send(*args, **kwargs)
123+
124+
send_patched = send.__name__ == "_sentry_send"
125+
new_send = _sentry_send if not send_patched else send
126+
127+
await old_call(app, scope, new_receive, new_send, **kwargs)
100128

101129
else:
102-
await old_call(*args, **kwargs)
130+
await old_call(app, scope, receive, send, **kwargs)
103131

104132
not_yet_patched = old_call.__name__ not in [
105133
"_create_span_call",

‎tests/integrations/starlette/test_starlette.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
from starlette.middleware.authentication import AuthenticationMiddleware
3232
from starlette.testclient import TestClient
3333

34+
STARLETTE_VERSION = tuple([int(x) for x in starlette.__version__.split(".")])
35+
3436
PICTURE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "photo.jpg")
3537

3638
BODY_JSON = {"some": "json", "for": "testing", "nested": {"numbers": 123}}
@@ -152,6 +154,26 @@ async def __anext__(self):
152154
raise StopAsyncIteration
153155

154156

157+
class SampleMiddleware:
158+
def __init__(self, app):
159+
self.app = app
160+
161+
async def __call__(self, scope, receive, send):
162+
# only handle http requests
163+
if scope["type"] != "http":
164+
await self.app(scope, receive, send)
165+
return
166+
167+
async def do_stuff(message):
168+
if message["type"] == "http.response.start":
169+
# do something here.
170+
pass
171+
172+
await send(message)
173+
174+
await self.app(scope, receive, do_stuff)
175+
176+
155177
@pytest.mark.asyncio
156178
async def test_starlettrequestextractor_content_length(sentry_init):
157179
with mock.patch(
@@ -546,6 +568,82 @@ def test_middleware_spans(sentry_init, capture_events):
546568
idx += 1
547569

548570

571+
def test_middleware_callback_spans(sentry_init, capture_events):
572+
sentry_init(
573+
traces_sample_rate=1.0,
574+
integrations=[StarletteIntegration()],
575+
)
576+
starlette_app = starlette_app_factory(middleware=[Middleware(SampleMiddleware)])
577+
events = capture_events()
578+
579+
client = TestClient(starlette_app, raise_server_exceptions=False)
580+
try:
581+
client.get("/message", auth=("Gabriela", "hello123"))
582+
except Exception:
583+
pass
584+
585+
(_, transaction_event) = events
586+
587+
expected = [
588+
{
589+
"op": "middleware.starlette",
590+
"description": "ServerErrorMiddleware",
591+
"tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
592+
},
593+
{
594+
"op": "middleware.starlette",
595+
"description": "SampleMiddleware",
596+
"tags": {"starlette.middleware_name": "SampleMiddleware"},
597+
},
598+
{
599+
"op": "middleware.starlette",
600+
"description": "ExceptionMiddleware",
601+
"tags": {"starlette.middleware_name": "ExceptionMiddleware"},
602+
},
603+
{
604+
"op": "middleware.starlette.send",
605+
"description": "SampleMiddleware.__call__.<locals>.do_stuff",
606+
"tags": {"starlette.middleware_name": "ExceptionMiddleware"},
607+
},
608+
{
609+
"op": "middleware.starlette.send",
610+
"description": "ServerErrorMiddleware.__call__.<locals>._send",
611+
"tags": {"starlette.middleware_name": "SampleMiddleware"},
612+
},
613+
{
614+
"op": "middleware.starlette.send",
615+
"description": "_ASGIAdapter.send.<locals>.send"
616+
if STARLETTE_VERSION < (0, 21)
617+
else "_TestClientTransport.handle_request.<locals>.send",
618+
"tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
619+
},
620+
{
621+
"op": "middleware.starlette.send",
622+
"description": "SampleMiddleware.__call__.<locals>.do_stuff",
623+
"tags": {"starlette.middleware_name": "ExceptionMiddleware"},
624+
},
625+
{
626+
"op": "middleware.starlette.send",
627+
"description": "ServerErrorMiddleware.__call__.<locals>._send",
628+
"tags": {"starlette.middleware_name": "SampleMiddleware"},
629+
},
630+
{
631+
"op": "middleware.starlette.send",
632+
"description": "_ASGIAdapter.send.<locals>.send"
633+
if STARLETTE_VERSION < (0, 21)
634+
else "_TestClientTransport.handle_request.<locals>.send",
635+
"tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
636+
},
637+
]
638+
639+
idx = 0
640+
for span in transaction_event["spans"]:
641+
assert span["op"] == expected[idx]["op"]
642+
assert span["description"] == expected[idx]["description"]
643+
assert span["tags"] == expected[idx]["tags"]
644+
idx += 1
645+
646+
549647
def test_last_event_id(sentry_init, capture_events):
550648
sentry_init(
551649
integrations=[StarletteIntegration()],

‎tox.ini‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ envlist =
3636

3737
{py3.7,py3.8,py3.9,py3.10}-asgi
3838

39-
{py3.7,py3.8,py3.9,py3.10}-starlette-{0.19.1,0.20}
39+
{py3.7,py3.8,py3.9,py3.10}-starlette-{0.19.1,0.20,0.21}
4040

4141
{py3.7,py3.8,py3.9,py3.10}-fastapi
4242

@@ -152,8 +152,10 @@ deps =
152152
starlette: pytest-asyncio
153153
starlette: python-multipart
154154
starlette: requests
155+
starlette-0.21: httpx
155156
starlette-0.19.1: starlette==0.19.1
156157
starlette-0.20: starlette>=0.20.0,<0.21.0
158+
starlette-0.21: starlette>=0.21.0,<0.22.0
157159

158160
fastapi: fastapi
159161
fastapi: pytest-asyncio

0 commit comments

Comments
 (0)