Skip to content

ref(integrations): Use ensure_integration_enabled decorator #2906

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.

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 18 commits into from
Mar 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
7 changes: 3 additions & 4 deletions sentry_sdk/integrations/aiohttp.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,12 +190,9 @@ 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
client = sentry_sdk.get_client()
if client.get_integration(AioHttpIntegration) is None:
return

method = params.method.upper()

parsed_url = None
Expand All @@ -213,6 +210,8 @@ async def on_request_start(session, trace_config_ctx, params):
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)

client = sentry_sdk.get_client()

if should_propagate_trace(client, str(params.url)):
for key, value in Scope.get_current_scope().iter_trace_propagation_headers(
span=span
Expand Down
5 changes: 2 additions & 3 deletions sentry_sdk/integrations/asyncpg.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,14 +58,13 @@ 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:
integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)

# Avoid recording calls to _execute twice.
# Calls to Connection.execute with args also call
# Connection._execute, which is recorded separately
# args[0] = the connection object, args[1] is the query
if integration is None or len(args) > 2:
if len(args) > 2:
return await f(*args, **kwargs)

query = args[1]
Expand Down
10 changes: 5 additions & 5 deletions sentry_sdk/integrations/atexit.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from sentry_sdk import Scope
from sentry_sdk.utils import logger
from sentry_sdk.integrations import Integration

from sentry_sdk.utils import ensure_integration_enabled
from sentry_sdk._types import TYPE_CHECKING

if TYPE_CHECKING:
Expand Down Expand Up @@ -44,13 +44,13 @@ def __init__(self, callback=None):
def setup_once():
# type: () -> None
@atexit.register
@ensure_integration_enabled(AtexitIntegration)
def _shutdown():
# type: () -> None
logger.debug("atexit: got shutdown signal")
client = sentry_sdk.get_client()
integration = client.get_integration(AtexitIntegration)
if integration is not None:
logger.debug("atexit: shutting down client")

Scope.get_isolation_scope().end_session()
client.close(callback=integration.callback)
logger.debug("atexit: shutting down client")
Scope.get_isolation_scope().end_session()
client.close(callback=integration.callback)
8 changes: 3 additions & 5 deletions sentry_sdk/integrations/aws_lambda.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
logger,
TimeoutThread,
Expand All @@ -37,12 +38,10 @@

def _wrap_init_error(init_error):
# type: (F) -> F
@ensure_integration_enabled(AwsLambdaIntegration, init_error)
def sentry_init_error(*args, **kwargs):
# type: (*Any, **Any) -> Any
client = sentry_sdk.get_client()
integration = client.get_integration(AwsLambdaIntegration)
if integration is None:
return init_error(*args, **kwargs)

with capture_internal_exceptions():
Scope.get_isolation_scope().clear_breadcrumbs()
Expand All @@ -63,6 +62,7 @@ def sentry_init_error(*args, **kwargs):

def _wrap_handler(handler):
# type: (F) -> F
@ensure_integration_enabled(AwsLambdaIntegration, handler)
def sentry_handler(aws_event, aws_context, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> Any

Expand Down Expand Up @@ -91,8 +91,6 @@ def sentry_handler(aws_event, aws_context, *args, **kwargs):

client = sentry_sdk.get_client()
integration = client.get_integration(AwsLambdaIntegration)
if integration is None:
return handler(aws_event, aws_context, *args, **kwargs)

configured_time = aws_context.get_remaining_time_in_millis()

Expand Down
11 changes: 7 additions & 4 deletions sentry_sdk/integrations/boto3.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,12 @@
from sentry_sdk.tracing import Span

from sentry_sdk._types import TYPE_CHECKING
from sentry_sdk.utils import capture_internal_exceptions, parse_url, parse_version
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
parse_url,
parse_version,
)

if TYPE_CHECKING:
from typing import Any
Expand Down Expand Up @@ -57,11 +62,9 @@ def sentry_patched_init(self, *args, **kwargs):
BaseClient.__init__ = sentry_patched_init


@ensure_integration_enabled(Boto3Integration)
def _sentry_request_created(service_id, request, operation_name, **kwargs):
# type: (str, AWSRequest, str, **Any) -> None
if sentry_sdk.get_client().get_integration(Boto3Integration) is None:
return

description = "aws.%s.%s" % (service_id, operation_name)
span = sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
Expand Down
7 changes: 2 additions & 5 deletions sentry_sdk/integrations/bottle.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,10 @@ def sentry_patched_wsgi_app(self, environ, start_response):

old_handle = Bottle._handle

@ensure_integration_enabled(BottleIntegration, old_handle)
def _patched_handle(self, environ):
# type: (Bottle, Dict[str, Any]) -> Any
integration = sentry_sdk.get_client().get_integration(BottleIntegration)
if integration is None:
return old_handle(self, environ)

scope = Scope.get_isolation_scope()
scope._name = "bottle"
Expand All @@ -96,13 +95,11 @@ def _patched_handle(self, environ):

old_make_callback = Route._make_callback

@ensure_integration_enabled(BottleIntegration, old_make_callback)
def patched_make_callback(self, *args, **kwargs):
# type: (Route, *object, **object) -> Any
client = sentry_sdk.get_client()
integration = client.get_integration(BottleIntegration)
prepared_callback = old_make_callback(self, *args, **kwargs)
if integration is None:
return prepared_callback

def wrapped_callback(*args, **kwargs):
# type: (*object, **object) -> Any
Expand Down
5 changes: 2 additions & 3 deletions sentry_sdk/integrations/clickhouse_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from sentry_sdk.tracing import Span
from sentry_sdk._types import TYPE_CHECKING
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import capture_internal_exceptions
from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled

from typing import TypeVar

Expand Down Expand Up @@ -74,9 +74,8 @@ def setup_once() -> None:


def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
@ensure_integration_enabled(ClickhouseDriverIntegration, f)
def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
if sentry_sdk.get_client().get_integration(ClickhouseDriverIntegration) is None:
return f(*args, **kwargs)
connection = args[0]
query = args[1]
query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
Expand Down
30 changes: 15 additions & 15 deletions sentry_sdk/integrations/django/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -392,13 +392,11 @@ def _set_transaction_name_and_source(scope, transaction_style, request):
pass


@ensure_integration_enabled(DjangoIntegration)
def _before_get_response(request):
# type: (WSGIRequest) -> None
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)

if integration is None:
return

_patch_drf()

scope = Scope.get_current_scope()
Expand All @@ -423,10 +421,11 @@ def _attempt_resolve_again(request, scope, transaction_style):
_set_transaction_name_and_source(scope, transaction_style, request)


@ensure_integration_enabled(DjangoIntegration)
def _after_get_response(request):
# type: (WSGIRequest) -> None
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None or integration.transaction_style != "url":
if integration.transaction_style != "url":
return

scope = Scope.get_current_scope()
Expand Down Expand Up @@ -492,21 +491,22 @@ def wsgi_request_event_processor(event, hint):
return wsgi_request_event_processor


@ensure_integration_enabled(DjangoIntegration)
def _got_request_exception(request=None, **kwargs):
# type: (WSGIRequest, **Any) -> None
client = sentry_sdk.get_client()
integration = client.get_integration(DjangoIntegration)
if integration is not None:
if request is not None and integration.transaction_style == "url":
scope = Scope.get_current_scope()
_attempt_resolve_again(request, scope, integration.transaction_style)

event, hint = event_from_exception(
sys.exc_info(),
client_options=client.options,
mechanism={"type": "django", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)

if request is not None and integration.transaction_style == "url":
scope = Scope.get_current_scope()
_attempt_resolve_again(request, scope, integration.transaction_style)

event, hint = event_from_exception(
sys.exc_info(),
client_options=client.options,
mechanism={"type": "django", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)


class DjangoRequestExtractor(RequestExtractor):
Expand Down
4 changes: 1 addition & 3 deletions sentry_sdk/integrations/django/templates.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,9 @@ def patch_templates():
real_rendered_content = SimpleTemplateResponse.rendered_content

@property # type: ignore
@ensure_integration_enabled(DjangoIntegration, real_rendered_content.fget)
def rendered_content(self):
# type: (SimpleTemplateResponse) -> str
if sentry_sdk.get_client().get_integration(DjangoIntegration) is None:
return real_rendered_content.fget(self)

with sentry_sdk.start_span(
op=OP.TEMPLATE_RENDER,
description=_get_template_name_description(self.template_name),
Expand Down
9 changes: 7 additions & 2 deletions sentry_sdk/integrations/excepthook.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
import sys

import sentry_sdk
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
)
from sentry_sdk.integrations import Integration

from sentry_sdk._types import TYPE_CHECKING
Expand Down Expand Up @@ -43,11 +47,12 @@ def setup_once():

def _make_excepthook(old_excepthook):
# type: (Excepthook) -> Excepthook
@ensure_integration_enabled(ExcepthookIntegration, old_excepthook)
def sentry_sdk_excepthook(type_, value, traceback):
# type: (Type[BaseException], BaseException, Optional[TracebackType]) -> None
integration = sentry_sdk.get_client().get_integration(ExcepthookIntegration)

if integration is not None and _should_send(integration.always_run):
if _should_send(integration.always_run):
with capture_internal_exceptions():
event, hint = event_from_exception(
(type_, value, traceback),
Expand Down
10 changes: 4 additions & 6 deletions sentry_sdk/integrations/falcon.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from sentry_sdk.tracing import SOURCE_FOR_STYLE
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
parse_version,
)
Expand Down Expand Up @@ -167,6 +168,7 @@ def _patch_handle_exception():
# type: () -> None
original_handle_exception = falcon_app_class._handle_exception

@ensure_integration_enabled(FalconIntegration, original_handle_exception)
def sentry_patched_handle_exception(self, *args):
# type: (falcon.API, *Any) -> Any
# NOTE(jmagnusson): falcon 2.0 changed falcon.API._handle_exception
Expand All @@ -187,14 +189,10 @@ def sentry_patched_handle_exception(self, *args):
# capture_internal_exceptions block above.
return was_handled

client = sentry_sdk.get_client()
integration = client.get_integration(FalconIntegration)

if integration is not None and _exception_leads_to_http_5xx(ex, response):
# If an integration is there, a client has to be there.
if _exception_leads_to_http_5xx(ex, response):
event, hint = event_from_exception(
ex,
client_options=client.options,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "falcon", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
Expand Down
15 changes: 6 additions & 9 deletions sentry_sdk/integrations/flask.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,11 @@ def setup_once():

old_app = Flask.__call__

@ensure_integration_enabled(FlaskIntegration, old_app)
def sentry_patched_wsgi_app(self, environ, start_response):
# type: (Any, Dict[str, str], Callable[..., Any]) -> _ScopedResponse
if sentry_sdk.get_client().get_integration(FlaskIntegration) is None:
return old_app(self, environ, start_response)

return SentryWsgiMiddleware(lambda *a, **kw: old_app(self, *a, **kw))(
environ, start_response
)
Expand Down Expand Up @@ -112,12 +114,10 @@ def _set_transaction_name_and_source(scope, transaction_style, request):
pass


@ensure_integration_enabled(FlaskIntegration)
def _request_started(app, **kwargs):
# type: (Flask, **Any) -> None
integration = sentry_sdk.get_client().get_integration(FlaskIntegration)
if integration is None:
return

request = flask_request._get_current_object()

# Set the transaction name and source here,
Expand Down Expand Up @@ -192,15 +192,12 @@ def inner(event, hint):
return inner


@ensure_integration_enabled(FlaskIntegration)
def _capture_exception(sender, exception, **kwargs):
# type: (Flask, Union[ValueError, BaseException], **Any) -> None
client = sentry_sdk.get_client()
if client.get_integration(FlaskIntegration) is None:
return

event, hint = event_from_exception(
exception,
client_options=client.options,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "flask", "handled": False},
)

Expand Down
4 changes: 2 additions & 2 deletions sentry_sdk/integrations/gcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
logger,
TimeoutThread,
Expand All @@ -38,13 +39,12 @@

def _wrap_func(func):
# type: (F) -> F
@ensure_integration_enabled(GcpIntegration, func)
def sentry_func(functionhandler, gcp_event, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> Any
client = sentry_sdk.get_client()

integration = client.get_integration(GcpIntegration)
if integration is None:
return func(functionhandler, gcp_event, *args, **kwargs)

configured_time = environ.get("FUNCTION_TIMEOUT_SEC")
if not configured_time:
Expand Down
Loading