Skip to content

Commit 245195e

Browse files
sl0thentr0pysl0thentr0py
authored
Update integrations with new continue_trace callback usage (#3486)
* Update integrations with new continue_trace callback usage * aiohttp * asgi * aws_lambda * celery * gcp * huey * rq * sanic - this uses manual enter/exit on the contextmanager so we'll have to test properly * tornado * ray
1 parent 7c4a3a0 commit 245195e

File tree

12 files changed

+200
-256
lines changed

12 files changed

+200
-256
lines changed

‎sentry_sdk/integrations/_asgi_common.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def _get_headers(asgi_scope):
2121
Extract headers from the ASGI scope, in the format that the Sentry protocol expects.
2222
"""
2323
headers = {} # type: Dict[str, str]
24-
for raw_key, raw_value in asgi_scope["headers"]:
24+
for raw_key, raw_value in asgi_scope.get("headers", {}):
2525
key = raw_key.decode("latin-1")
2626
value = raw_value.decode("latin-1")
2727
if key in headers:

‎sentry_sdk/integrations/aiohttp.py‎

Lines changed: 25 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
import weakref
33

44
import sentry_sdk
5-
from sentry_sdk.api import continue_trace
65
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
76
from sentry_sdk.integrations import Integration, DidNotEnable
87
from sentry_sdk.integrations.logging import ignore_logger
@@ -113,34 +112,31 @@ async def sentry_app_handle(self, request, *args, **kwargs):
113112
scope.add_event_processor(_make_request_processor(weak_request))
114113

115114
headers = dict(request.headers)
116-
transaction = continue_trace(
117-
headers,
118-
op=OP.HTTP_SERVER,
119-
# If this transaction name makes it to the UI, AIOHTTP's
120-
# URL resolver did not find a route or died trying.
121-
name="generic AIOHTTP request",
122-
source=TRANSACTION_SOURCE_ROUTE,
123-
origin=AioHttpIntegration.origin,
124-
)
125-
with sentry_sdk.start_transaction(
126-
transaction,
127-
custom_sampling_context={"aiohttp_request": request},
128-
):
129-
try:
130-
response = await old_handle(self, request)
131-
except HTTPException as e:
132-
transaction.set_http_status(e.status_code)
133-
raise
134-
except (asyncio.CancelledError, ConnectionResetError):
135-
transaction.set_status(SPANSTATUS.CANCELLED)
136-
raise
137-
except Exception:
138-
# This will probably map to a 500 but seems like we
139-
# have no way to tell. Do not set span status.
140-
reraise(*_capture_exception())
141-
142-
transaction.set_http_status(response.status)
143-
return response
115+
with sentry_sdk.continue_trace(headers):
116+
with sentry_sdk.start_transaction(
117+
op=OP.HTTP_SERVER,
118+
# If this transaction name makes it to the UI, AIOHTTP's
119+
# URL resolver did not find a route or died trying.
120+
name="generic AIOHTTP request",
121+
source=TRANSACTION_SOURCE_ROUTE,
122+
origin=AioHttpIntegration.origin,
123+
custom_sampling_context={"aiohttp_request": request},
124+
) as transaction:
125+
try:
126+
response = await old_handle(self, request)
127+
except HTTPException as e:
128+
transaction.set_http_status(e.status_code)
129+
raise
130+
except (asyncio.CancelledError, ConnectionResetError):
131+
transaction.set_status(SPANSTATUS.CANCELLED)
132+
raise
133+
except Exception:
134+
# This will probably map to a 500 but seems like we
135+
# have no way to tell. Do not set span status.
136+
reraise(*_capture_exception())
137+
138+
transaction.set_http_status(response.status)
139+
return response
144140

145141
Application._handle = sentry_app_handle
146142

‎sentry_sdk/integrations/asgi.py‎

Lines changed: 37 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
from functools import partial
1111

1212
import sentry_sdk
13-
from sentry_sdk.api import continue_trace
1413
from sentry_sdk.consts import OP
1514

1615
from sentry_sdk.integrations._asgi_common import (
@@ -34,7 +33,6 @@
3433
transaction_from_function,
3534
_get_installed_modules,
3635
)
37-
from sentry_sdk.tracing import Transaction
3836

3937
from typing import TYPE_CHECKING
4038

@@ -185,66 +183,47 @@ async def _run_app(self, scope, receive, send, asgi_version):
185183
scope,
186184
)
187185

188-
if ty in ("http", "websocket"):
189-
transaction = continue_trace(
190-
_get_headers(scope),
191-
op="{}.server".format(ty),
186+
with sentry_sdk.continue_trace(_get_headers(scope)):
187+
with sentry_sdk.start_transaction(
188+
op=(
189+
OP.WEBSOCKET_SERVER
190+
if ty == "websocket"
191+
else OP.HTTP_SERVER
192+
),
192193
name=transaction_name,
193194
source=transaction_source,
194195
origin=self.span_origin,
195-
)
196-
logger.debug(
197-
"[ASGI] Created transaction (continuing trace): %s",
198-
transaction,
199-
)
200-
else:
201-
transaction = Transaction(
202-
op=OP.HTTP_SERVER,
203-
name=transaction_name,
204-
source=transaction_source,
205-
origin=self.span_origin,
206-
)
207-
logger.debug(
208-
"[ASGI] Created transaction (new): %s", transaction
209-
)
210-
211-
transaction.set_tag("asgi.type", ty)
212-
logger.debug(
213-
"[ASGI] Set transaction name and source on transaction: '%s' / '%s'",
214-
transaction.name,
215-
transaction.source,
216-
)
217-
218-
with sentry_sdk.start_transaction(
219-
transaction,
220-
custom_sampling_context={"asgi_scope": scope},
221-
):
222-
logger.debug("[ASGI] Started transaction: %s", transaction)
223-
try:
224-
225-
async def _sentry_wrapped_send(event):
226-
# type: (Dict[str, Any]) -> Any
227-
is_http_response = (
228-
event.get("type") == "http.response.start"
229-
and transaction is not None
230-
and "status" in event
231-
)
232-
if is_http_response:
233-
transaction.set_http_status(event["status"])
234-
235-
return await send(event)
236-
237-
if asgi_version == 2:
238-
return await self.app(scope)(
239-
receive, _sentry_wrapped_send
240-
)
241-
else:
242-
return await self.app(
243-
scope, receive, _sentry_wrapped_send
196+
custom_sampling_context={"asgi_scope": scope},
197+
) as transaction:
198+
logger.debug("[ASGI] Started transaction: %s", transaction)
199+
transaction.set_tag("asgi.type", ty)
200+
try:
201+
202+
async def _sentry_wrapped_send(event):
203+
# type: (Dict[str, Any]) -> Any
204+
is_http_response = (
205+
event.get("type") == "http.response.start"
206+
and transaction is not None
207+
and "status" in event
208+
)
209+
if is_http_response:
210+
transaction.set_http_status(event["status"])
211+
212+
return await send(event)
213+
214+
if asgi_version == 2:
215+
return await self.app(scope)(
216+
receive, _sentry_wrapped_send
217+
)
218+
else:
219+
return await self.app(
220+
scope, receive, _sentry_wrapped_send
221+
)
222+
except Exception as exc:
223+
_capture_exception(
224+
exc, mechanism_type=self.mechanism_type
244225
)
245-
except Exception as exc:
246-
_capture_exception(exc, mechanism_type=self.mechanism_type)
247-
raise exc from None
226+
raise exc from None
248227
finally:
249228
_asgi_middleware_applied.set(False)
250229

‎sentry_sdk/integrations/aws_lambda.py‎

Lines changed: 25 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
from os import environ
55

66
import sentry_sdk
7-
from sentry_sdk.api import continue_trace
87
from sentry_sdk.consts import OP
98
from sentry_sdk.scope import should_send_default_pii
109
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
@@ -135,34 +134,31 @@ def sentry_handler(aws_event, aws_context, *args, **kwargs):
135134
if not isinstance(headers, dict):
136135
headers = {}
137136

138-
transaction = continue_trace(
139-
headers,
140-
op=OP.FUNCTION_AWS,
141-
name=aws_context.function_name,
142-
source=TRANSACTION_SOURCE_COMPONENT,
143-
origin=AwsLambdaIntegration.origin,
144-
)
145-
with sentry_sdk.start_transaction(
146-
transaction,
147-
custom_sampling_context={
148-
"aws_event": aws_event,
149-
"aws_context": aws_context,
150-
},
151-
):
152-
try:
153-
return handler(aws_event, aws_context, *args, **kwargs)
154-
except Exception:
155-
exc_info = sys.exc_info()
156-
sentry_event, hint = event_from_exception(
157-
exc_info,
158-
client_options=client.options,
159-
mechanism={"type": "aws_lambda", "handled": False},
160-
)
161-
sentry_sdk.capture_event(sentry_event, hint=hint)
162-
reraise(*exc_info)
163-
finally:
164-
if timeout_thread:
165-
timeout_thread.stop()
137+
with sentry_sdk.continue_trace(headers):
138+
with sentry_sdk.start_transaction(
139+
op=OP.FUNCTION_AWS,
140+
name=aws_context.function_name,
141+
source=TRANSACTION_SOURCE_COMPONENT,
142+
origin=AwsLambdaIntegration.origin,
143+
custom_sampling_context={
144+
"aws_event": aws_event,
145+
"aws_context": aws_context,
146+
},
147+
):
148+
try:
149+
return handler(aws_event, aws_context, *args, **kwargs)
150+
except Exception:
151+
exc_info = sys.exc_info()
152+
sentry_event, hint = event_from_exception(
153+
exc_info,
154+
client_options=client.options,
155+
mechanism={"type": "aws_lambda", "handled": False},
156+
)
157+
sentry_sdk.capture_event(sentry_event, hint=hint)
158+
reraise(*exc_info)
159+
finally:
160+
if timeout_thread:
161+
timeout_thread.stop()
166162

167163
return sentry_handler # type: ignore
168164

‎sentry_sdk/integrations/celery/__init__.py‎

Lines changed: 16 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
import sentry_sdk
66
from sentry_sdk import isolation_scope
7-
from sentry_sdk.api import continue_trace
87
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
98
from sentry_sdk.integrations import Integration, DidNotEnable
109
from sentry_sdk.integrations.celery.beat import (
@@ -301,38 +300,27 @@ def _inner(*args, **kwargs):
301300
scope.clear_breadcrumbs()
302301
scope.add_event_processor(_make_event_processor(task, *args, **kwargs))
303302

304-
transaction = None
305-
306303
# Celery task objects are not a thing to be trusted. Even
307304
# something such as attribute access can fail.
308-
with capture_internal_exceptions():
309-
headers = args[3].get("headers") or {}
310-
transaction = continue_trace(
311-
headers,
305+
headers = args[3].get("headers") or {}
306+
with sentry_sdk.continue_trace(headers):
307+
with sentry_sdk.start_transaction(
312308
op=OP.QUEUE_TASK_CELERY,
313-
name="unknown celery task",
309+
name=task.name,
314310
source=TRANSACTION_SOURCE_TASK,
315311
origin=CeleryIntegration.origin,
316-
)
317-
transaction.name = task.name
318-
transaction.set_status(SPANSTATUS.OK)
319-
320-
if transaction is None:
321-
return f(*args, **kwargs)
322-
323-
with sentry_sdk.start_transaction(
324-
transaction,
325-
custom_sampling_context={
326-
"celery_job": {
327-
"task": task.name,
328-
# for some reason, args[1] is a list if non-empty but a
329-
# tuple if empty
330-
"args": list(args[1]),
331-
"kwargs": args[2],
332-
}
333-
},
334-
):
335-
return f(*args, **kwargs)
312+
custom_sampling_context={
313+
"celery_job": {
314+
"task": task.name,
315+
# for some reason, args[1] is a list if non-empty but a
316+
# tuple if empty
317+
"args": list(args[1]),
318+
"kwargs": args[2],
319+
}
320+
},
321+
) as transaction:
322+
transaction.set_status(SPANSTATUS.OK)
323+
return f(*args, **kwargs)
336324

337325
return _inner # type: ignore
338326

0 commit comments

Comments
 (0)