Skip to content

Update integrations with new continue_trace callback usage #3486

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.

Sign up for GitHub

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

Merged
merged 2 commits into from
Sep 2, 2024
Merged

Update integrations with new continue_trace callback usage #3486

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
609e21a
Update integrations with new continue_trace callback usage
sl0thentr0py Aug 30, 2024
02efbe6
Make asgi_scope headers access safe
sl0thentr0py Sep 2, 2024
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
2 changes: 1 addition & 1 deletion sentry_sdk/integrations/_asgi_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def _get_headers(asgi_scope):
Extract headers from the ASGI scope, in the format that the Sentry protocol expects.
"""
headers = {} # type: Dict[str, str]
for raw_key, raw_value in asgi_scope["headers"]:
for raw_key, raw_value in asgi_scope.get("headers", {}):
key = raw_key.decode("latin-1")
value = raw_value.decode("latin-1")
if key in headers:
Expand Down
54 changes: 25 additions & 29 deletions sentry_sdk/integrations/aiohttp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import weakref

import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.integrations.logging import ignore_logger
Expand Down Expand Up @@ -113,34 +112,31 @@ async def sentry_app_handle(self, request, *args, **kwargs):
scope.add_event_processor(_make_request_processor(weak_request))

headers = dict(request.headers)
transaction = continue_trace(
headers,
op=OP.HTTP_SERVER,
# If this transaction name makes it to the UI, AIOHTTP's
# URL resolver did not find a route or died trying.
name="generic AIOHTTP request",
source=TRANSACTION_SOURCE_ROUTE,
origin=AioHttpIntegration.origin,
)
with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={"aiohttp_request": request},
):
try:
response = await old_handle(self, request)
except HTTPException as e:
transaction.set_http_status(e.status_code)
raise
except (asyncio.CancelledError, ConnectionResetError):
transaction.set_status(SPANSTATUS.CANCELLED)
raise
except Exception:
# This will probably map to a 500 but seems like we
# have no way to tell. Do not set span status.
reraise(*_capture_exception())

transaction.set_http_status(response.status)
return response
with sentry_sdk.continue_trace(headers):
with sentry_sdk.start_transaction(
Comment on lines +115 to +116
Copy link
Member

Choose a reason for hiding this comment

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

I wonder whether it might make sense to introduce a new convenience function that would continue the trace and start a span in one call, so that these two with context managers could be combined into a single with continue_trace_start_span (perhaps with better naming) statement.

op=OP.HTTP_SERVER,
# If this transaction name makes it to the UI, AIOHTTP's
# URL resolver did not find a route or died trying.
name="generic AIOHTTP request",
source=TRANSACTION_SOURCE_ROUTE,
origin=AioHttpIntegration.origin,
custom_sampling_context={"aiohttp_request": request},
) as transaction:
try:
response = await old_handle(self, request)
except HTTPException as e:
transaction.set_http_status(e.status_code)
raise
except (asyncio.CancelledError, ConnectionResetError):
transaction.set_status(SPANSTATUS.CANCELLED)
raise
except Exception:
# This will probably map to a 500 but seems like we
# have no way to tell. Do not set span status.
reraise(*_capture_exception())

transaction.set_http_status(response.status)
return response

Application._handle = sentry_app_handle

Expand Down
95 changes: 37 additions & 58 deletions sentry_sdk/integrations/asgi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
from functools import partial

import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP

from sentry_sdk.integrations._asgi_common import (
Expand All @@ -34,7 +33,6 @@
transaction_from_function,
_get_installed_modules,
)
from sentry_sdk.tracing import Transaction

from typing import TYPE_CHECKING

Expand Down Expand Up @@ -185,66 +183,47 @@ async def _run_app(self, scope, receive, send, asgi_version):
scope,
)

if ty in ("http", "websocket"):
transaction = continue_trace(
_get_headers(scope),
op="{}.server".format(ty),
with sentry_sdk.continue_trace(_get_headers(scope)):
with sentry_sdk.start_transaction(
op=(
OP.WEBSOCKET_SERVER
if ty == "websocket"
else OP.HTTP_SERVER
),
name=transaction_name,
source=transaction_source,
origin=self.span_origin,
)
logger.debug(
"[ASGI] Created transaction (continuing trace): %s",
transaction,
)
else:
Copy link
Contributor

Choose a reason for hiding this comment

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

We're removing this else logic for non-http and non-websocket scope types, was this intentional? I'm not sure why it was there in the first place (maybe future proofing for new asgi scope types?), but it seems deliberate

Copy link
Contributor

Choose a reason for hiding this comment

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

Ok this is why: 44c43a8 Apparently fetching headers might be problematic for other scope types. So we should probably keep this

Copy link
Member Author

Choose a reason for hiding this comment

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

thanks for digging, made the accessor safe to avoid the if
02efbe6

transaction = Transaction(
op=OP.HTTP_SERVER,
name=transaction_name,
source=transaction_source,
origin=self.span_origin,
)
logger.debug(
"[ASGI] Created transaction (new): %s", transaction
)

transaction.set_tag("asgi.type", ty)
logger.debug(
"[ASGI] Set transaction name and source on transaction: '%s' / '%s'",
transaction.name,
transaction.source,
)

with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={"asgi_scope": scope},
):
logger.debug("[ASGI] Started transaction: %s", transaction)
try:

async def _sentry_wrapped_send(event):
# type: (Dict[str, Any]) -> Any
is_http_response = (
event.get("type") == "http.response.start"
and transaction is not None
and "status" in event
)
if is_http_response:
transaction.set_http_status(event["status"])

return await send(event)

if asgi_version == 2:
return await self.app(scope)(
receive, _sentry_wrapped_send
)
else:
return await self.app(
scope, receive, _sentry_wrapped_send
custom_sampling_context={"asgi_scope": scope},
) as transaction:
logger.debug("[ASGI] Started transaction: %s", transaction)
transaction.set_tag("asgi.type", ty)
try:

async def _sentry_wrapped_send(event):
# type: (Dict[str, Any]) -> Any
is_http_response = (
event.get("type") == "http.response.start"
and transaction is not None
and "status" in event
)
if is_http_response:
transaction.set_http_status(event["status"])

return await send(event)

if asgi_version == 2:
return await self.app(scope)(
receive, _sentry_wrapped_send
)
else:
return await self.app(
scope, receive, _sentry_wrapped_send
)
except Exception as exc:
_capture_exception(
exc, mechanism_type=self.mechanism_type
)
except Exception as exc:
_capture_exception(exc, mechanism_type=self.mechanism_type)
raise exc from None
raise exc from None
finally:
_asgi_middleware_applied.set(False)

Expand Down
54 changes: 25 additions & 29 deletions sentry_sdk/integrations/aws_lambda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from os import environ

import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
Expand Down Expand Up @@ -135,34 +134,31 @@ def sentry_handler(aws_event, aws_context, *args, **kwargs):
if not isinstance(headers, dict):
headers = {}

transaction = continue_trace(
headers,
op=OP.FUNCTION_AWS,
name=aws_context.function_name,
source=TRANSACTION_SOURCE_COMPONENT,
origin=AwsLambdaIntegration.origin,
)
with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={
"aws_event": aws_event,
"aws_context": aws_context,
},
):
try:
return handler(aws_event, aws_context, *args, **kwargs)
except Exception:
exc_info = sys.exc_info()
sentry_event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "aws_lambda", "handled": False},
)
sentry_sdk.capture_event(sentry_event, hint=hint)
reraise(*exc_info)
finally:
if timeout_thread:
timeout_thread.stop()
with sentry_sdk.continue_trace(headers):
with sentry_sdk.start_transaction(
op=OP.FUNCTION_AWS,
name=aws_context.function_name,
source=TRANSACTION_SOURCE_COMPONENT,
origin=AwsLambdaIntegration.origin,
custom_sampling_context={
"aws_event": aws_event,
"aws_context": aws_context,
},
):
try:
return handler(aws_event, aws_context, *args, **kwargs)
except Exception:
exc_info = sys.exc_info()
sentry_event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "aws_lambda", "handled": False},
)
sentry_sdk.capture_event(sentry_event, hint=hint)
reraise(*exc_info)
finally:
if timeout_thread:
timeout_thread.stop()

return sentry_handler # type: ignore

Expand Down
44 changes: 16 additions & 28 deletions sentry_sdk/integrations/celery/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import sentry_sdk
from sentry_sdk import isolation_scope
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.integrations.celery.beat import (
Expand Down Expand Up @@ -301,38 +300,27 @@ def _inner(*args, **kwargs):
scope.clear_breadcrumbs()
scope.add_event_processor(_make_event_processor(task, *args, **kwargs))

transaction = None

# Celery task objects are not a thing to be trusted. Even
# something such as attribute access can fail.
with capture_internal_exceptions():
headers = args[3].get("headers") or {}
transaction = continue_trace(
headers,
headers = args[3].get("headers") or {}
with sentry_sdk.continue_trace(headers):
with sentry_sdk.start_transaction(
op=OP.QUEUE_TASK_CELERY,
name="unknown celery task",
name=task.name,
source=TRANSACTION_SOURCE_TASK,
origin=CeleryIntegration.origin,
)
transaction.name = task.name
transaction.set_status(SPANSTATUS.OK)

if transaction is None:
return f(*args, **kwargs)

with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={
"celery_job": {
"task": task.name,
# for some reason, args[1] is a list if non-empty but a
# tuple if empty
"args": list(args[1]),
"kwargs": args[2],
}
},
):
return f(*args, **kwargs)
custom_sampling_context={
"celery_job": {
"task": task.name,
# for some reason, args[1] is a list if non-empty but a
# tuple if empty
"args": list(args[1]),
"kwargs": args[2],
}
},
) as transaction:
transaction.set_status(SPANSTATUS.OK)
return f(*args, **kwargs)

return _inner # type: ignore

Expand Down
Loading
Loading