Skip to content

Commit

Permalink
fix(integrations): Do not use convenience decorator (#3022)
Browse files Browse the repository at this point in the history
  • Loading branch information
sentrivana authored Apr 26, 2024
1 parent a626f01 commit 162773c
Show file tree
Hide file tree
Showing 13 changed files with 65 additions and 35 deletions.
9 changes: 6 additions & 3 deletions sentry_sdk/integrations/aiohttp.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
logger,
parse_url,
Expand Down Expand Up @@ -98,9 +97,11 @@ def setup_once():

old_handle = Application._handle

@ensure_integration_enabled_async(AioHttpIntegration, old_handle)
async def sentry_app_handle(self, request, *args, **kwargs):
# type: (Any, Request, *Any, **Any) -> Any
if sentry_sdk.get_client().get_integration(AioHttpIntegration) is None:
return await old_handle(self, request, *args, **kwargs)

weak_request = weakref.ref(request)

with sentry_sdk.isolation_scope() as scope:
Expand Down Expand Up @@ -190,9 +191,11 @@ def init(*args, **kwargs):
def create_trace_config():
# type: () -> TraceConfig

@ensure_integration_enabled_async(AioHttpIntegration)
async def on_request_start(session, trace_config_ctx, params):
# type: (ClientSession, SimpleNamespace, TraceRequestStartParams) -> None
if sentry_sdk.get_client().get_integration(AioHttpIntegration) is None:
return

method = params.method.upper()

parsed_url = None
Expand Down
16 changes: 12 additions & 4 deletions sentry_sdk/integrations/arq.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
SENSITIVE_DATA_SUBSTITUTE,
parse_version,
Expand Down Expand Up @@ -71,9 +70,12 @@ def patch_enqueue_job():
# type: () -> None
old_enqueue_job = ArqRedis.enqueue_job

@ensure_integration_enabled_async(ArqIntegration, old_enqueue_job)
async def _sentry_enqueue_job(self, function, *args, **kwargs):
# type: (ArqRedis, str, *Any, **Any) -> Optional[Job]
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await old_enqueue_job(self, function, *args, **kwargs)

with sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_ARQ, description=function):
return await old_enqueue_job(self, function, *args, **kwargs)

Expand All @@ -84,9 +86,12 @@ def patch_run_job():
# type: () -> None
old_run_job = Worker.run_job

@ensure_integration_enabled_async(ArqIntegration, old_run_job)
async def _sentry_run_job(self, job_id, score):
# type: (Worker, str, int) -> None
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await old_run_job(self, job_id, score)

with sentry_sdk.isolation_scope() as scope:
scope._name = "arq"
scope.clear_breadcrumbs()
Expand Down Expand Up @@ -157,9 +162,12 @@ def event_processor(event, hint):
def _wrap_coroutine(name, coroutine):
# type: (str, WorkerCoroutine) -> WorkerCoroutine

@ensure_integration_enabled_async(ArqIntegration, coroutine)
async def _sentry_coroutine(ctx, *args, **kwargs):
# type: (Dict[Any, Any], *Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await coroutine(ctx, *args, **kwargs)

Scope.get_isolation_scope().add_event_processor(
_make_event_processor({**ctx, "job_name": name}, *args, **kwargs)
)
Expand Down
12 changes: 8 additions & 4 deletions sentry_sdk/integrations/asyncpg.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
from sentry_sdk.utils import (
ensure_integration_enabled,
ensure_integration_enabled_async,
parse_version,
capture_internal_exceptions,
)
Expand Down Expand Up @@ -58,8 +57,10 @@ def setup_once() -> None:


def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
@ensure_integration_enabled_async(AsyncPGIntegration, f)
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)

# Avoid recording calls to _execute twice.
# Calls to Connection.execute with args also call
# Connection._execute, which is recorded separately
Expand Down Expand Up @@ -110,8 +111,9 @@ def _record(
def _wrap_connection_method(
f: Callable[..., Awaitable[T]], *, executemany: bool = False
) -> Callable[..., Awaitable[T]]:
@ensure_integration_enabled_async(AsyncPGIntegration, f)
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)
query = args[1]
params_list = args[2] if len(args) > 2 else None
with _record(None, query, params_list, executemany=executemany) as span:
Expand Down Expand Up @@ -145,8 +147,10 @@ def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807


def _wrap_connect_addr(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
@ensure_integration_enabled_async(AsyncPGIntegration, f)
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)

user = kwargs["params"].user
database = kwargs["params"].database

Expand Down
9 changes: 6 additions & 3 deletions sentry_sdk/integrations/django/asgi.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
)


Expand Down Expand Up @@ -72,9 +71,11 @@ def patch_django_asgi_handler_impl(cls):

old_app = cls.__call__

@ensure_integration_enabled_async(DjangoIntegration, old_app)
async def sentry_patched_asgi_handler(self, scope, receive, send):
# type: (Any, Any, Any, Any) -> Any
if sentry_sdk.get_client().get_integration(DjangoIntegration) is None:
return await old_app(self, scope, receive, send)

middleware = SentryAsgiMiddleware(
old_app.__get__(self, cls), unsafe_context_data=True
)._run_asgi3
Expand Down Expand Up @@ -120,9 +121,11 @@ def patch_channels_asgi_handler_impl(cls):
if channels.__version__ < "3.0.0":
old_app = cls.__call__

@ensure_integration_enabled_async(DjangoIntegration, old_app)
async def sentry_patched_asgi_handler(self, receive, send):
# type: (Any, Any, Any) -> Any
if sentry_sdk.get_client().get_integration(DjangoIntegration) is None:
return await old_app(self, receive, send)

middleware = SentryAsgiMiddleware(
lambda _scope: old_app.__get__(self, cls), unsafe_context_data=True
)
Expand Down
5 changes: 3 additions & 2 deletions sentry_sdk/integrations/fastapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
from sentry_sdk.utils import (
transaction_from_function,
logger,
ensure_integration_enabled_async,
)

if TYPE_CHECKING:
Expand Down Expand Up @@ -97,9 +96,11 @@ def _sentry_call(*args, **kwargs):

old_app = old_get_request_handler(*args, **kwargs)

@ensure_integration_enabled_async(FastApiIntegration, old_app)
async def _sentry_app(*args, **kwargs):
# type: (*Any, **Any) -> Any
if sentry_sdk.get_client().get_integration(FastApiIntegration) is None:
return await old_app(*args, **kwargs)

integration = sentry_sdk.get_client().get_integration(FastApiIntegration)
request = args[0]

Expand Down
5 changes: 3 additions & 2 deletions sentry_sdk/integrations/graphene.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
package_version,
)
Expand Down Expand Up @@ -69,9 +68,11 @@ def _sentry_patched_graphql_sync(schema, source, *args, **kwargs):

return result

@ensure_integration_enabled_async(GrapheneIntegration, old_graphql_async)
async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
# type: (GraphQLSchema, Union[str, Source], Any, Any) -> ExecutionResult
if sentry_sdk.get_client().get_integration(GrapheneIntegration) is None:
return await old_graphql_async(schema, source, *args, **kwargs)

scope = Scope.get_isolation_scope()
scope.add_event_processor(_event_processor)

Expand Down
5 changes: 3 additions & 2 deletions sentry_sdk/integrations/httpx.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
SENSITIVE_DATA_SUBSTITUTE,
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
logger,
parse_url,
)
Expand Down Expand Up @@ -98,9 +97,11 @@ def _install_httpx_async_client():
# type: () -> None
real_send = AsyncClient.send

@ensure_integration_enabled_async(HttpxIntegration, real_send)
async def send(self, request, **kwargs):
# type: (AsyncClient, Request, **Any) -> Response
if sentry_sdk.get_client().get_integration(HttpxIntegration) is None:
return await real_send(self, request, **kwargs)

parsed_url = None
with capture_internal_exceptions():
parsed_url = parse_url(str(request.url), sanitize=False)
Expand Down
9 changes: 6 additions & 3 deletions sentry_sdk/integrations/quart.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
)
from sentry_sdk._types import TYPE_CHECKING
Expand Down Expand Up @@ -150,10 +149,11 @@ def _set_transaction_name_and_source(scope, transaction_style, request):
pass


@ensure_integration_enabled_async(QuartIntegration)
async def _request_websocket_started(app, **kwargs):
# type: (Quart, **Any) -> None
integration = sentry_sdk.get_client().get_integration(QuartIntegration)
if integration is None:
return

if has_request_context():
request_websocket = request._get_current_object()
Expand Down Expand Up @@ -200,9 +200,12 @@ def inner(event, hint):
return inner


@ensure_integration_enabled_async(QuartIntegration)
async def _capture_exception(sender, exception, **kwargs):
# type: (Quart, Union[ValueError, BaseException], **Any) -> None
integration = sentry_sdk.get_client().get_integration(QuartIntegration)
if integration is None:
return

event, hint = event_from_exception(
exception,
client_options=sentry_sdk.get_client().options,
Expand Down
9 changes: 6 additions & 3 deletions sentry_sdk/integrations/redis/asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
from sentry_sdk.tracing import Span
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled_async,
)

if TYPE_CHECKING:
Expand All @@ -26,9 +25,11 @@ def patch_redis_async_pipeline(
# type: (Union[type[Pipeline[Any]], type[ClusterPipeline[Any]]], bool, Any, Callable[[Span, Any], None]) -> None
old_execute = pipeline_cls.execute

@ensure_integration_enabled_async(RedisIntegration, old_execute)
async def _sentry_execute(self, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
if sentry_sdk.get_client().get_integration(RedisIntegration) is None:
return await old_execute(self, *args, **kwargs)

with sentry_sdk.start_span(
op=OP.DB_REDIS, description="redis.pipeline.execute"
) as span:
Expand All @@ -51,9 +52,11 @@ def patch_redis_async_client(cls, is_cluster, set_db_data_fn):
# type: (Union[type[StrictRedis[Any]], type[RedisCluster[Any]]], bool, Callable[[Span, Any], None]) -> None
old_execute_command = cls.execute_command

@ensure_integration_enabled_async(RedisIntegration, old_execute_command) # type: ignore
async def _sentry_execute_command(self, name, *args, **kwargs):
# type: (Any, str, *Any, **Any) -> Any
if sentry_sdk.get_client().get_integration(RedisIntegration) is None:
return await old_execute_command(self, name, *args, **kwargs)

description = _get_span_description(name, *args)

with sentry_sdk.start_span(op=OP.DB_REDIS, description=description) as span:
Expand Down
5 changes: 3 additions & 2 deletions sentry_sdk/integrations/sanic.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
HAS_REAL_CONTEXTVARS,
CONTEXTVARS_ERROR_MESSAGE,
Expand Down Expand Up @@ -274,9 +273,11 @@ async def sentry_wrapped_error_handler(request, exception):
return sentry_wrapped_error_handler


@ensure_integration_enabled_async(SanicIntegration, old_handle_request)
async def _legacy_handle_request(self, request, *args, **kwargs):
# type: (Any, Request, *Any, **Any) -> Any
if sentry_sdk.get_client().get_integration(SanicIntegration) is None:
return await old_handle_request(self, request, *args, **kwargs)

weak_request = weakref.ref(request)

with sentry_sdk.isolation_scope() as scope:
Expand Down
7 changes: 4 additions & 3 deletions sentry_sdk/integrations/starlette.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
AnnotatedValue,
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
logger,
parse_version,
Expand Down Expand Up @@ -337,10 +336,11 @@ def patch_asgi_app():
"""
old_app = Starlette.__call__

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

middleware = SentryAsgiMiddleware(
lambda *a, **kw: old_app(self, *a, **kw),
Expand Down Expand Up @@ -377,12 +377,13 @@ 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
integration = sentry_sdk.get_client().get_integration(
StarletteIntegration
)
if integration is None:
return await old_func(*args, **kwargs)

request = args[0]

Expand Down
4 changes: 2 additions & 2 deletions sentry_sdk/integrations/starlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE
from sentry_sdk.utils import (
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
transaction_from_function,
)
Expand Down Expand Up @@ -179,10 +178,11 @@ async def _sentry_send(message: "Message") -> None:
def patch_http_route_handle() -> None:
old_handle = HTTPRoute.handle

@ensure_integration_enabled_async(StarliteIntegration, old_handle)
async def handle_wrapper(
self: "HTTPRoute", scope: "HTTPScope", receive: "Receive", send: "Send"
) -> None:
if sentry_sdk.get_client().get_integration(StarliteIntegration) is None:
return await old_handle(self, scope, receive, send)

sentry_scope = SentryScope.get_isolation_scope()
request: "Request[Any, Any]" = scope["app"].request_class(
Expand Down
5 changes: 3 additions & 2 deletions sentry_sdk/integrations/strawberry.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
ensure_integration_enabled_async,
event_from_exception,
logger,
package_version,
Expand Down Expand Up @@ -266,11 +265,13 @@ def _patch_execute():
old_execute_async = strawberry_schema.execute
old_execute_sync = strawberry_schema.execute_sync

@ensure_integration_enabled_async(StrawberryIntegration, old_execute_async)
async def _sentry_patched_execute_async(*args, **kwargs):
# type: (Any, Any) -> ExecutionResult
result = await old_execute_async(*args, **kwargs)

if sentry_sdk.get_client().get_integration(StrawberryIntegration) is None:
return result

if "execution_context" in kwargs and result.errors:
scope = Scope.get_isolation_scope()
event_processor = _make_request_event_processor(kwargs["execution_context"])
Expand Down

0 comments on commit 162773c

Please sign in to comment.