Merge branch 'master' into potel-base

Author: sentrivana

sentry_sdk/integrations/celery/__init__.py

@@ -40,6 +40,7 @@
 
 try:
     from celery import VERSION as CELERY_VERSION  # type: ignore
+    from celery.app.task import Task  # type: ignore
     from celery.app.trace import task_has_custom
     from celery.exceptions import (  # type: ignore
         Ignore,
@@ -82,6 +83,7 @@ def setup_once():
 
         _patch_build_tracer()
         _patch_task_apply_async()
+        _patch_celery_send_task()
         _patch_worker_exit()
         _patch_producer_publish()
 
@@ -242,7 +244,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         return None
 
 
-def _wrap_apply_async(f):
+def _wrap_task_run(f):
     # type: (F) -> F
     @wraps(f)
     @ensure_integration_enabled(CeleryIntegration, f)
@@ -259,14 +261,19 @@ def apply_async(*args, **kwargs):
         if not propagate_traces:
             return f(*args, **kwargs)
 
-        task = args[0]
+        if isinstance(args[0], Task):
+            task_name = args[0].name  # type: str
+        elif len(args) > 1 and isinstance(args[1], str):
+            task_name = args[1]
+        else:
+            task_name = "<unknown Celery task>"
 
         task_started_from_beat = sentry_sdk.get_isolation_scope()._name == "celery-beat"
 
         span_mgr = (
             sentry_sdk.start_span(
                 op=OP.QUEUE_SUBMIT_CELERY,
-                description=task.name,
+                description=task_name,
                 origin=CeleryIntegration.origin,
             )
             if not task_started_from_beat
@@ -425,9 +432,14 @@ def sentry_build_tracer(name, task, *args, **kwargs):
 
 def _patch_task_apply_async():
     # type: () -> None
-    from celery.app.task import Task  # type: ignore
+    Task.apply_async = _wrap_task_run(Task.apply_async)
+
+
+def _patch_celery_send_task():
+    # type: () -> None
+    from celery import Celery
 
-    Task.apply_async = _wrap_apply_async(Task.apply_async)
+    Celery.send_task = _wrap_task_run(Celery.send_task)
 
 
 def _patch_worker_exit():

sentry_sdk/integrations/django/middleware.py

@@ -30,7 +30,9 @@
     "import_string_should_wrap_middleware"
 )
 
-if DJANGO_VERSION < (3, 1):
+DJANGO_SUPPORTS_ASYNC_MIDDLEWARE = DJANGO_VERSION >= (3, 1)
+
+if not DJANGO_SUPPORTS_ASYNC_MIDDLEWARE:
     _asgi_middleware_mixin_factory = lambda _: object
 else:
     from .asgi import _asgi_middleware_mixin_factory
@@ -123,7 +125,9 @@ def sentry_wrapped_method(*args, **kwargs):
     class SentryWrappingMiddleware(
         _asgi_middleware_mixin_factory(_check_middleware_span)  # type: ignore
     ):
-        async_capable = getattr(middleware, "async_capable", False)
+        async_capable = DJANGO_SUPPORTS_ASYNC_MIDDLEWARE and getattr(
+            middleware, "async_capable", False
+        )
 
         def __init__(self, get_response=None, *args, **kwargs):
             # type: (Optional[Callable[..., Any]], *Any, **Any) -> None

sentry_sdk/integrations/strawberry.py

@@ -37,10 +37,10 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Generator, List, Optional
+    from typing import Any, Callable, Generator, List, Optional, Union
     from graphql import GraphQLError, GraphQLResolveInfo  # type: ignore
     from strawberry.http import GraphQLHTTPResponse
-    from strawberry.types import ExecutionContext, ExecutionResult  # type: ignore
+    from strawberry.types import ExecutionContext, ExecutionResult, SubscriptionExecutionResult  # type: ignore
     from sentry_sdk._types import Event, EventProcessor
 
 
@@ -279,13 +279,13 @@ def _patch_execute():
     old_execute_sync = strawberry_schema.execute_sync
 
     async def _sentry_patched_execute_async(*args, **kwargs):
-        # type: (Any, Any) -> ExecutionResult
+        # type: (Any, Any) -> Union[ExecutionResult, SubscriptionExecutionResult]
         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:
+        if "execution_context" in kwargs:
             scope = sentry_sdk.get_isolation_scope()
             event_processor = _make_request_event_processor(kwargs["execution_context"])
             scope.add_event_processor(event_processor)
@@ -297,7 +297,7 @@ def _sentry_patched_execute_sync(*args, **kwargs):
         # type: (Any, Any) -> ExecutionResult
         result = old_execute_sync(*args, **kwargs)
 
-        if "execution_context" in kwargs and result.errors:
+        if "execution_context" in kwargs:
             scope = sentry_sdk.get_isolation_scope()
             event_processor = _make_request_event_processor(kwargs["execution_context"])
             scope.add_event_processor(event_processor)

tests/integrations/celery/test_celery.py

@@ -10,7 +10,7 @@
 from sentry_sdk import start_transaction, get_current_span
 from sentry_sdk.integrations.celery import (
     CeleryIntegration,
-    _wrap_apply_async,
+    _wrap_task_run,
 )
 from sentry_sdk.integrations.celery.beat import _get_headers
 from tests.conftest import ApproxDict
@@ -568,7 +568,7 @@ def dummy_function(*args, **kwargs):
         assert "sentry-trace" in headers
         assert "baggage" in headers
 
-    wrapped = _wrap_apply_async(dummy_function)
+    wrapped = _wrap_task_run(dummy_function)
     wrapped(mock.MagicMock(), (), headers={})
 
 
@@ -783,3 +783,51 @@ def task(): ...
         assert span["origin"] == "auto.queue.celery"
 
     monkeypatch.setattr(kombu.messaging.Producer, "_publish", old_publish)
+
+
+@pytest.mark.forked
+@mock.patch("celery.Celery.send_task")
+def test_send_task_wrapped(
+    patched_send_task,
+    sentry_init,
+    capture_events,
+    reset_integrations,
+):
+    sentry_init(integrations=[CeleryIntegration()], enable_tracing=True)
+    celery = Celery(__name__, broker="redis://example.com")  # noqa: E231
+
+    events = capture_events()
+
+    with sentry_sdk.start_transaction(name="custom_transaction"):
+        celery.send_task("very_creative_task_name", args=(1, 2), kwargs={"foo": "bar"})
+
+    (call,) = patched_send_task.call_args_list  # We should have exactly one call
+    (args, kwargs) = call
+
+    assert args == (celery, "very_creative_task_name")
+    assert kwargs["args"] == (1, 2)
+    assert kwargs["kwargs"] == {"foo": "bar"}
+    assert set(kwargs["headers"].keys()) == {
+        "sentry-task-enqueued-time",
+        "sentry-trace",
+        "baggage",
+        "headers",
+    }
+    assert set(kwargs["headers"]["headers"].keys()) == {
+        "sentry-trace",
+        "baggage",
+        "sentry-task-enqueued-time",
+    }
+    assert (
+        kwargs["headers"]["sentry-trace"]
+        == kwargs["headers"]["headers"]["sentry-trace"]
+    )
+
+    (event,) = events  # We should have exactly one event (the transaction)
+    assert event["type"] == "transaction"
+    assert event["transaction"] == "custom_transaction"
+
+    (span,) = event["spans"]  # We should have exactly one span
+    assert span["description"] == "very_creative_task_name"
+    assert span["op"] == "queue.submit.celery"
+    assert span["trace_id"] == kwargs["headers"]["sentry-trace"].split("-")[0]

tox.ini

@@ -361,6 +361,8 @@ deps =
 
     # Celery
     celery: redis
+    celery: newrelic
+    celery: pytest<7
     celery-v4: Celery~=4.0
     celery-v5.0: Celery~=5.0.0
     celery-v5.1: Celery~=5.1.0
@@ -370,7 +372,6 @@ deps =
     celery-latest: Celery
 
     {py3.7}-celery: importlib-metadata<5.0
-    {py3.7,py3.8,py3.9,py3.10,py3.11,py3.12}-celery: newrelic
 
     # Chalice
     chalice-v1.16: chalice~=1.16.0