From 8f80c90ae2f11bc13f47fe1ae700cdcec1e8e6ca Mon Sep 17 00:00:00 2001 From: Daniel Szoke Date: Mon, 23 Sep 2024 14:58:59 +0200 Subject: [PATCH] reapply some changes --- sentry_sdk/integrations/bottle.py | 6 +++++- sentry_sdk/integrations/celery/__init__.py | 6 ++++-- sentry_sdk/integrations/cohere.py | 24 ++++++++++------------ sentry_sdk/integrations/django/__init__.py | 9 ++++---- 4 files changed, 25 insertions(+), 20 deletions(-) diff --git a/sentry_sdk/integrations/bottle.py b/sentry_sdk/integrations/bottle.py index b1800bd191..dc573eb958 100644 --- a/sentry_sdk/integrations/bottle.py +++ b/sentry_sdk/integrations/bottle.py @@ -1,3 +1,5 @@ +import functools + import sentry_sdk from sentry_sdk.tracing import SOURCE_FOR_STYLE from sentry_sdk.utils import ( @@ -81,10 +83,12 @@ def sentry_patched_wsgi_app(self, environ, start_response): old_handle = Bottle._handle - @ensure_integration_enabled(BottleIntegration, old_handle) + @functools.wraps(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 = sentry_sdk.get_isolation_scope() scope._name = "bottle" diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 28a44015aa..9a984de8c3 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -248,13 +248,15 @@ def __exit__(self, exc_type, exc_value, traceback): def _wrap_task_run(f): # type: (F) -> F @wraps(f) - @ensure_integration_enabled(CeleryIntegration, f) def apply_async(*args, **kwargs): # type: (*Any, **Any) -> Any # Note: kwargs can contain headers=None, so no setdefault! # Unsure which backend though. - kwarg_headers = kwargs.get("headers") or {} integration = sentry_sdk.get_client().get_integration(CeleryIntegration) + if integration is None: + return f(*args, **kwargs) + + kwarg_headers = kwargs.get("headers") or {} propagate_traces = kwarg_headers.pop( "sentry-propagate-traces", integration.propagate_traces ) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 4d6a4a244c..b4c2af91da 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -14,11 +14,7 @@ import sentry_sdk from sentry_sdk.scope import should_send_default_pii from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.utils import ( - capture_internal_exceptions, - event_from_exception, - ensure_integration_enabled, -) +from sentry_sdk.utils import capture_internal_exceptions, event_from_exception try: from cohere.client import Client @@ -134,13 +130,15 @@ def collect_chat_response_fields(span, res, include_pii): set_data_normalized(span, "ai.warnings", res.meta.warnings) @wraps(f) - @ensure_integration_enabled(CohereIntegration, f) def new_chat(*args, **kwargs): # type: (*Any, **Any) -> Any - if "message" not in kwargs: - return f(*args, **kwargs) + integration = sentry_sdk.get_client().get_integration(CohereIntegration) - if not isinstance(kwargs.get("message"), str): + if ( + integration is None + or "message" not in kwargs + or not isinstance(kwargs.get("message"), str) + ): return f(*args, **kwargs) message = kwargs.get("message") @@ -158,8 +156,6 @@ def new_chat(*args, **kwargs): span.__exit__(None, None, None) raise e from None - integration = sentry_sdk.get_client().get_integration(CohereIntegration) - with capture_internal_exceptions(): if should_send_default_pii() and integration.include_prompts: set_data_normalized( @@ -227,15 +223,17 @@ def _wrap_embed(f): # type: (Callable[..., Any]) -> Callable[..., Any] @wraps(f) - @ensure_integration_enabled(CohereIntegration, f) def new_embed(*args, **kwargs): # type: (*Any, **Any) -> Any + integration = sentry_sdk.get_client().get_integration(CohereIntegration) + if integration is None: + return f(*args, **kwargs) + with sentry_sdk.start_span( op=consts.OP.COHERE_EMBEDDINGS_CREATE, name="Cohere Embedding Creation", origin=CohereIntegration.origin, ) as span: - integration = sentry_sdk.get_client().get_integration(CohereIntegration) if "texts" in kwargs and ( should_send_default_pii() and integration.include_prompts ): diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py index fce93503e9..40d17b0507 100644 --- a/sentry_sdk/integrations/django/__init__.py +++ b/sentry_sdk/integrations/django/__init__.py @@ -411,10 +411,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() @@ -440,11 +441,10 @@ 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.transaction_style != "url": + if integration is None or integration.transaction_style != "url": return scope = sentry_sdk.get_current_scope() @@ -510,11 +510,12 @@ 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 None: + return if request is not None and integration.transaction_style == "url": scope = sentry_sdk.get_current_scope()