diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 7dd5fef836..46e4c5e41e 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -1,12 +1,12 @@ import sys from collections.abc import Mapping +from contextlib import nullcontext from functools import wraps from typing import TYPE_CHECKING import sentry_sdk from sentry_sdk import isolation_scope -from sentry_sdk.api import continue_trace -from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS +from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version from sentry_sdk.integrations.celery.beat import ( _patch_beat_apply_entry, @@ -16,9 +16,9 @@ from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch from sentry_sdk.integrations.logging import ignore_logger_for_events from sentry_sdk.scope import Scope, should_send_default_pii -from sentry_sdk.traces import SegmentNameSource, StreamedSpan, get_current_span -from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource -from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled +from sentry_sdk.traces import SegmentNameSource, StreamedSpan +from sentry_sdk.tracing import BAGGAGE_HEADER_NAME +from sentry_sdk.tracing_utils import Baggage from sentry_sdk.utils import ( SENSITIVE_DATA_SUBSTITUTE, capture_internal_exceptions, @@ -29,7 +29,7 @@ ) if TYPE_CHECKING: - from typing import Any, Callable, List, Optional, TypeVar, Union + from typing import Any, Callable, List, Optional, TypeVar from sentry_sdk._types import Event, EventProcessor, ExcInfo, Hint @@ -94,16 +94,11 @@ def setup_once() -> None: def _set_status(status: str) -> None: - client = sentry_sdk.get_client() - span_streaming = has_span_streaming_enabled(client.options) - with capture_internal_exceptions(): - scope = sentry_sdk.get_current_scope() + span = sentry_sdk.traces.get_current_span() - if span_streaming and scope.streamed_span is not None: - scope.streamed_span.status = "ok" if status == "ok" else "error" - elif not span_streaming and scope.span is not None: - scope.span.set_status(status) + if span is not None: + span.status = "ok" if status == "ok" else "error" def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None: @@ -175,7 +170,7 @@ def event_processor(event: "Event", hint: "Hint") -> "Optional[Event]": def _update_celery_task_headers( original_headers: "dict[str, Any]", - span: "Optional[Union[StreamedSpan, Span]]", + span: "Optional[StreamedSpan]", monitor_beat_tasks: bool, ) -> "dict[str, Any]": """ @@ -256,14 +251,6 @@ def _update_celery_task_headers( return updated_headers -class NoOpMgr: - def __enter__(self) -> None: - return None - - def __exit__(self, exc_type: "Any", exc_value: "Any", traceback: "Any") -> None: - return None - - def _wrap_task_run(f: "F") -> "F": @wraps(f) def apply_async(*args: "Any", **kwargs: "Any") -> "Any": @@ -289,30 +276,22 @@ def apply_async(*args: "Any", **kwargs: "Any") -> "Any": else: task_name = "" - span_streaming = has_span_streaming_enabled(client.options) - task_started_from_beat = sentry_sdk.get_isolation_scope()._name == "celery-beat" - span_mgr: "Union[StreamedSpan, Span, NoOpMgr]" = NoOpMgr() - if span_streaming: - if not task_started_from_beat and get_current_span() is not None: - span_mgr = sentry_sdk.traces.start_span( - name=task_name, - attributes={ - "sentry.op": OP.QUEUE_SUBMIT_CELERY, - "sentry.origin": CeleryIntegration.origin, - }, - ) - - else: - if not task_started_from_beat: - span_mgr = sentry_sdk.start_span( - op=OP.QUEUE_SUBMIT_CELERY, - name=task_name, - origin=CeleryIntegration.origin, - ) + span = None + if ( + not task_started_from_beat + and sentry_sdk.traces.get_current_span() is not None + ): + span = sentry_sdk.traces.start_span( + name=task_name, + attributes={ + "sentry.op": OP.QUEUE_SUBMIT_CELERY, + "sentry.origin": CeleryIntegration.origin, + }, + ) - with span_mgr as span: + with span if span else nullcontext(): kwargs["headers"] = _update_celery_task_headers( kwarg_headers, span, integration.monitor_beat_tasks ) @@ -334,8 +313,6 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any": if client.get_integration(CeleryIntegration) is None: return f(*args, **kwargs) - span_streaming = has_span_streaming_enabled(client.options) - with isolation_scope() as scope: scope._name = "celery" scope.clear_breadcrumbs() @@ -355,52 +332,33 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any": } } - span: "Union[Span, StreamedSpan]" - span_ctx: "Union[StreamedSpan, Span, NoOpMgr]" = NoOpMgr() + span = 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 {} - if span_streaming: - sentry_sdk.traces.continue_trace(headers) - Scope.set_custom_sampling_context(custom_sampling_context) - span = sentry_sdk.traces.start_span( - name=task_name, - parent_span=None, # make this a segment - attributes={ - "sentry.origin": CeleryIntegration.origin, - "sentry.segment.name.source": SegmentNameSource.TASK.value, - "sentry.op": OP.QUEUE_TASK_CELERY, - }, - ) - - span_ctx = span + sentry_sdk.traces.continue_trace(headers) - else: - span = continue_trace( - headers, - op=OP.QUEUE_TASK_CELERY, - name=task_name, - source=TransactionSource.TASK, - origin=CeleryIntegration.origin, - ) - span.set_status(SPANSTATUS.OK) + Scope.set_custom_sampling_context(custom_sampling_context) - span_ctx = sentry_sdk.start_transaction( - span, - custom_sampling_context=custom_sampling_context, - ) + span = sentry_sdk.traces.start_span( + name=task_name, + parent_span=None, # make this a segment + attributes={ + "sentry.origin": CeleryIntegration.origin, + "sentry.segment.name.source": SegmentNameSource.TASK.value, + "sentry.op": OP.QUEUE_TASK_CELERY, + }, + ) - with span_ctx: + with span if span else nullcontext(): return f(*args, **kwargs) return _inner # type: ignore -def _set_messaging_destination_name( - task: "Any", span: "Union[StreamedSpan, Span]" -) -> None: +def _set_messaging_destination_name(task: "Any", span: "StreamedSpan") -> None: """Set "messaging.destination.name" tag for span""" with capture_internal_exceptions(): delivery_info = task.request.delivery_info @@ -409,10 +367,7 @@ def _set_messaging_destination_name( if delivery_info.get("exchange") == "" and routing_key is not None: # Empty exchange indicates the default exchange, meaning the tasks # are sent to the queue with the same name as the routing key. - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key) - else: - span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key) + span.set_attribute(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key) def _wrap_task_call(task: "Any", f: "F") -> "F": @@ -425,34 +380,17 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any": if client.get_integration(CeleryIntegration) is None: return f(*args, **kwargs) - span_streaming = has_span_streaming_enabled(client.options) - try: - if span_streaming and get_current_span() is None: + if sentry_sdk.traces.get_current_span() is None: return f(*args, **kwargs) - span: "Union[Span, StreamedSpan]" - if span_streaming: - span = sentry_sdk.traces.start_span( - name=task.name, - attributes={ - "sentry.op": OP.QUEUE_PROCESS, - "sentry.origin": CeleryIntegration.origin, - }, - ) - else: - span = sentry_sdk.start_span( - op=OP.QUEUE_PROCESS, - name=task.name, - origin=CeleryIntegration.origin, - ) - - with span: - if isinstance(span, StreamedSpan): - set_on_span = span.set_attribute - else: - set_on_span = span.set_data - + with sentry_sdk.traces.start_span( + name=task.name, + attributes={ + "sentry.op": OP.QUEUE_PROCESS, + "sentry.origin": CeleryIntegration.origin, + }, + ) as span: _set_messaging_destination_name(task, span) latency = None @@ -467,25 +405,28 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any": if latency is not None: latency *= 1000 # milliseconds - set_on_span(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency) + span.set_attribute( + SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency + ) with capture_internal_exceptions(): - set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id) + span.set_attribute(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id) with capture_internal_exceptions(): - set_on_span( + span.set_attribute( SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, task.request.retries, ) with capture_internal_exceptions(): with task.app.connection() as conn: - set_on_span( + span.set_attribute( SPANDATA.MESSAGING_SYSTEM, conn.transport.driver_type, ) return f(*args, **kwargs) + except Exception: exc_info = sys.exc_info() with capture_internal_exceptions(): @@ -560,8 +501,6 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any": if client.get_integration(CeleryIntegration) is None: return original_publish(self, *args, **kwargs) - span_streaming = has_span_streaming_enabled(client.options) - kwargs_headers = kwargs.get("headers", {}) if not isinstance(kwargs_headers, Mapping): # Ensure kwargs_headers is a Mapping, so we can safely call get(). @@ -578,45 +517,29 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any": routing_key = kwargs.get("routing_key") exchange = kwargs.get("exchange") - span: "Union[StreamedSpan, Span, None]" = None - if span_streaming: - if get_current_span() is not None: - span = sentry_sdk.traces.start_span( - name=task_name, - attributes={ - "sentry.op": OP.QUEUE_PUBLISH, - "sentry.origin": CeleryIntegration.origin, - }, - ) - else: - span = sentry_sdk.start_span( - op=OP.QUEUE_PUBLISH, - name=task_name, - origin=CeleryIntegration.origin, - ) - - if span is None: + if sentry_sdk.traces.get_current_span() is None: return original_publish(self, *args, **kwargs) - with span: - if isinstance(span, StreamedSpan): - set_on_span = span.set_attribute - else: - set_on_span = span.set_data - + with sentry_sdk.traces.start_span( + name=task_name, + attributes={ + "sentry.op": OP.QUEUE_PUBLISH, + "sentry.origin": CeleryIntegration.origin, + }, + ) as span: if task_id is not None: - set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task_id) + span.set_attribute(SPANDATA.MESSAGING_MESSAGE_ID, task_id) if exchange == "" and routing_key is not None: # Empty exchange indicates the default exchange, meaning messages are # routed to the queue with the same name as the routing key. - set_on_span(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key) + span.set_attribute(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key) if retries is not None: - set_on_span(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries) + span.set_attribute(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries) with capture_internal_exceptions(): - set_on_span( + span.set_attribute( SPANDATA.MESSAGING_SYSTEM, self.connection.transport.driver_type ) diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py index 0bc22d9509..8a35c94e41 100644 --- a/tests/integrations/celery/test_celery.py +++ b/tests/integrations/celery/test_celery.py @@ -8,15 +8,12 @@ from celery.bin import worker import sentry_sdk -import sentry_sdk.traces -from sentry_sdk import get_current_span, start_transaction from sentry_sdk.integrations.celery import ( CeleryIntegration, _wrap_task_run, ) from sentry_sdk.integrations.celery.beat import _get_headers from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE -from tests.conftest import ApproxDict @pytest.fixture @@ -121,20 +118,17 @@ def celery_invocation(request): return request.param -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("send_default_pii", [True, False]) def test_simple_with_performance( - capture_events, capture_items, init_celery, celery_invocation, - span_streaming, send_default_pii, ): celery = init_celery( traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task(name="dummy_task") @@ -142,30 +136,18 @@ def dummy_task(x, y): foo = 42 # noqa return x / y - if span_streaming: - items = capture_items("event", "span") + items = capture_items("event", "span") - with sentry_sdk.traces.start_span(name="span") as span: - celery_invocation(dummy_task, 1, 2) - _, expected_context = celery_invocation(dummy_task, 1, 0) + with sentry_sdk.traces.start_span(name="span") as span: + celery_invocation(dummy_task, 1, 2) + _, expected_context = celery_invocation(dummy_task, 1, 0) - sentry_sdk.flush() - - error_event = next(item.payload for item in items if item.type == "event") - - assert error_event["contexts"]["trace"]["trace_id"] == span.trace_id - assert error_event["contexts"]["trace"]["span_id"] != span.span_id - else: - events = capture_events() - - with start_transaction(op="unit test transaction") as transaction: - celery_invocation(dummy_task, 1, 2) - _, expected_context = celery_invocation(dummy_task, 1, 0) + sentry_sdk.flush() - (_, error_event, _, _) = events + error_event = next(item.payload for item in items if item.type == "event") - assert error_event["contexts"]["trace"]["trace_id"] == transaction.trace_id - assert error_event["contexts"]["trace"]["span_id"] != transaction.span_id + assert error_event["contexts"]["trace"]["trace_id"] == span.trace_id + assert error_event["contexts"]["trace"]["span_id"] != span.span_id assert error_event["transaction"] == "dummy_task" assert "celery_task_id" in error_event["tags"] @@ -286,12 +268,9 @@ def dummy_task(x, y): pytest.param("keyword", [], {"x": 1, "y": 0}, id="keyword_args"), ], ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_task_args_kwargs_data_collection( - capture_events, capture_items, init_celery, - span_streaming, invocation_style, task_args, task_kwargs, @@ -301,7 +280,7 @@ def test_task_args_kwargs_data_collection( ): init_dict = {"send_default_pii": True, **init_kwargs} celery = init_celery( - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", **init_dict, ) @@ -314,15 +293,10 @@ def dummy_task(x, y): else: invoke = lambda: dummy_task.apply_async(kwargs=dict(x=1, y=0)) - if span_streaming: - items = capture_items("event") - invoke() - sentry_sdk.flush() - (error_event,) = (item.payload for item in items) - else: - events = capture_events() - invoke() - (error_event,) = events + items = capture_items("event") + invoke() + sentry_sdk.flush() + (error_event,) = (item.payload for item in items) celery_job = error_event["extra"]["celery-job"] @@ -337,19 +311,16 @@ def dummy_task(x, y): assert celery_job["kwargs"] == SENSITIVE_DATA_SUBSTITUTE -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("task_fails", [True, False], ids=["error", "success"]) def test_transaction_events( - capture_events, capture_items, init_celery, celery_invocation, task_fails, - span_streaming, ): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task(name="dummy_task") @@ -360,100 +331,44 @@ def dummy_task(x, y): celery_invocation(dummy_task, 1, 1) sentry_sdk.flush() - if span_streaming: - items = capture_items("event", "span") - - with sentry_sdk.traces.start_span(name="submission") as span: - celery_invocation(dummy_task, 1, 0 if task_fails else 1) - - sentry_sdk.flush() - - if task_fails: - error_event = items.pop(0).payload - assert error_event["contexts"]["trace"]["trace_id"] == span.trace_id - assert error_event["exception"]["values"][0]["type"] == "ZeroDivisionError" - - process_span, execution_span, submit_span, submission_span = [ - item.payload for item in items - ] + items = capture_items("event", "span") - assert execution_span["name"] == "dummy_task" - assert execution_span["is_segment"] is True - assert execution_span["attributes"]["sentry.segment.name.source"] == "task" - assert execution_span["trace_id"] == span.trace_id - if task_fails: - assert execution_span["status"] == "error" - else: - assert execution_span["status"] == "ok" + with sentry_sdk.traces.start_span(name="submission") as span: + celery_invocation(dummy_task, 1, 0 if task_fails else 1) - assert process_span["name"] == "dummy_task" - assert process_span["trace_id"] == span.trace_id - assert process_span["attributes"]["sentry.op"] == "queue.process" - assert process_span["parent_span_id"] == execution_span["span_id"] + sentry_sdk.flush() - assert submission_span["name"] == "submission" - assert submission_span["is_segment"] is True + if task_fails: + error_event = items.pop(0).payload + assert error_event["contexts"]["trace"]["trace_id"] == span.trace_id + assert error_event["exception"]["values"][0]["type"] == "ZeroDivisionError" - assert submit_span["name"] == "dummy_task" - assert submit_span["attributes"]["sentry.op"] == "queue.submit.celery" - assert submit_span["attributes"]["sentry.origin"] == "auto.queue.celery" - assert ( - submit_span["parent_span_id"] == submission_span["span_id"] == span.span_id - ) - assert submit_span["trace_id"] == span.trace_id + process_span, execution_span, submit_span, submission_span = [ + item.payload for item in items + ] + assert execution_span["name"] == "dummy_task" + assert execution_span["is_segment"] is True + assert execution_span["attributes"]["sentry.segment.name.source"] == "task" + assert execution_span["trace_id"] == span.trace_id + if task_fails: + assert execution_span["status"] == "error" else: - events = capture_events() - - with start_transaction(name="submission") as transaction: - celery_invocation(dummy_task, 1, 0 if task_fails else 1) + assert execution_span["status"] == "ok" - if task_fails: - error_event = events.pop(0) - assert error_event["contexts"]["trace"]["trace_id"] == transaction.trace_id - assert error_event["exception"]["values"][0]["type"] == "ZeroDivisionError" + assert process_span["name"] == "dummy_task" + assert process_span["trace_id"] == span.trace_id + assert process_span["attributes"]["sentry.op"] == "queue.process" + assert process_span["parent_span_id"] == execution_span["span_id"] - execution_event, submission_event = events - assert execution_event["transaction"] == "dummy_task" - assert execution_event["transaction_info"] == {"source": "task"} + assert submission_span["name"] == "submission" + assert submission_span["is_segment"] is True - assert submission_event["transaction"] == "submission" - assert submission_event["transaction_info"] == {"source": "custom"} - - assert execution_event["type"] == submission_event["type"] == "transaction" - assert execution_event["contexts"]["trace"]["trace_id"] == transaction.trace_id - assert submission_event["contexts"]["trace"]["trace_id"] == transaction.trace_id - - if task_fails: - assert execution_event["contexts"]["trace"]["status"] == "internal_error" - else: - assert execution_event["contexts"]["trace"]["status"] == "ok" - - assert len(execution_event["spans"]) == 1 - assert ( - execution_event["spans"][0].items() - >= { - "trace_id": str(transaction.trace_id), - "same_process_as_parent": True, - "op": "queue.process", - "description": "dummy_task", - "data": ApproxDict(), - }.items() - ) - assert submission_event["spans"] == [ - { - "data": ApproxDict(), - "description": "dummy_task", - "op": "queue.submit.celery", - "origin": "auto.queue.celery", - "parent_span_id": submission_event["contexts"]["trace"]["span_id"], - "same_process_as_parent": True, - "span_id": submission_event["spans"][0]["span_id"], - "start_timestamp": submission_event["spans"][0]["start_timestamp"], - "timestamp": submission_event["spans"][0]["timestamp"], - "trace_id": str(transaction.trace_id), - } - ] + assert submit_span["name"] == "dummy_task" + assert submit_span["attributes"]["sentry.op"] == "queue.submit.celery" + assert submit_span["attributes"]["sentry.origin"] == "auto.queue.celery" + assert submit_span["parent_span_id"] == submission_span["span_id"] == span.span_id + assert submit_span["trace_id"] == span.trace_id def test_no_double_patching(celery): @@ -494,7 +409,7 @@ def test_simple_no_propagation(capture_events, init_celery): def dummy_task(): 1 / 0 - with start_transaction() as transaction: + with sentry_sdk.start_transaction() as transaction: dummy_task.delay() (event,) = events @@ -569,7 +484,7 @@ def dummy_task(self): runs.append(1) 1 / 0 - with start_transaction(name="submit_celery"): + with sentry_sdk.start_transaction(name="submit_celery"): # Curious: Cannot use delay() here or py2.7-celery-4.2 crashes res = dummy_task.apply_async() @@ -648,9 +563,7 @@ def dummy_task(self, x, y): assert celery_invocation(dummy_task, 1, 1)[0].wait() == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_traces_sampler_gets_task_info_in_sampling_context( - span_streaming, init_celery, celery_invocation, DictionaryContaining, # noqa:N803 @@ -658,7 +571,7 @@ def test_traces_sampler_gets_task_info_in_sampling_context( traces_sampler = mock.Mock(return_value=1.0) celery = init_celery( traces_sampler=traces_sampler, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task(name="dog_walk") @@ -694,7 +607,7 @@ def __call__(self, *args, **kwargs): def dummy_task(x, y): return x / y - with start_transaction(): + with sentry_sdk.start_transaction(): celery_invocation(dummy_task, 1, 0) assert not events @@ -737,7 +650,7 @@ def dummy_task(self, x, y): # patch random.randrange to return a predictable sample_rand value with mock.patch("sentry_sdk.tracing_utils.Random.randrange", return_value=500000): - with start_transaction() as transaction: + with sentry_sdk.start_transaction() as transaction: result = dummy_task.apply_async( args=(1, 0), headers={"baggage": "custom=value"}, @@ -756,8 +669,7 @@ def dummy_task(self, x, y): ) -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_sentry_propagate_traces_override(span_streaming, init_celery): +def test_sentry_propagate_traces_override(init_celery): """ Test if the `sentry-propagate-traces` header given to `apply_async` overrides the `propagate_traces` parameter in the integration constructor. @@ -766,50 +678,29 @@ def test_sentry_propagate_traces_override(span_streaming, init_celery): propagate_traces=True, traces_sample_rate=1.0, release="abcdef", - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task(name="dummy_task", bind=True) def dummy_task(self, message): - trace_id = ( - sentry_sdk.traces.get_current_span().trace_id - if span_streaming - else get_current_span().trace_id - ) + trace_id = sentry_sdk.traces.get_current_span().trace_id return trace_id - if span_streaming: - with sentry_sdk.traces.start_span(name="parent") as span: - parent_trace_id = span.trace_id + with sentry_sdk.traces.start_span(name="parent") as span: + parent_trace_id = span.trace_id - # should propagate trace - task_trace_id = dummy_task.apply_async( - args=("some message",), - ).get() - assert parent_trace_id == task_trace_id - - # should NOT propagate trace - task_trace_id = dummy_task.apply_async( - args=("another message",), - headers={"sentry-propagate-traces": False}, - ).get() - assert parent_trace_id != task_trace_id - else: - with start_transaction() as transaction: - transaction_trace_id = transaction.trace_id - - # should propagate trace - task_trace_id = dummy_task.apply_async( - args=("some message",), - ).get() - assert transaction_trace_id == task_trace_id + # should propagate trace + task_trace_id = dummy_task.apply_async( + args=("some message",), + ).get() + assert parent_trace_id == task_trace_id - # should NOT propagate trace - task_trace_id = dummy_task.apply_async( - args=("another message",), - headers={"sentry-propagate-traces": False}, - ).get() - assert transaction_trace_id != task_trace_id + # should NOT propagate trace + task_trace_id = dummy_task.apply_async( + args=("another message",), + headers={"sentry-propagate-traces": False}, + ).get() + assert parent_trace_id != task_trace_id def test_apply_async_manually_span(sentry_init): @@ -841,47 +732,35 @@ def example_task(): assert result.get() == "success" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("routing_key", ("celery", "custom")) @mock.patch("celery.app.task.Task.request") def test_messaging_destination_name_default_exchange( mock_request, routing_key, - span_streaming, init_celery, - capture_events, capture_items, ): celery_app = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) mock_request.delivery_info = {"routing_key": routing_key, "exchange": ""} @celery_app.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert ( - process_span.payload["attributes"]["messaging.destination.name"] - == routing_key - ) - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert span["data"]["messaging.destination.name"] == routing_key + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert ( + process_span.payload["attributes"]["messaging.destination.name"] == routing_key + ) -@pytest.mark.parametrize("span_streaming", [True, False]) @mock.patch("celery.app.task.Task.request") def test_messaging_destination_name_nondefault_exchange( - mock_request, span_streaming, init_celery, capture_events, capture_items + mock_request, init_celery, capture_items ): """ Currently, we only capture the routing key as the messaging.destination.name when @@ -891,112 +770,76 @@ def test_messaging_destination_name_nondefault_exchange( """ celery_app = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) mock_request.delivery_info = {"routing_key": "celery", "exchange": "custom"} @celery_app.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert "messaging.destination.name" not in process_span.payload["attributes"] - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert "messaging.destination.name" not in span["data"] + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert "messaging.destination.name" not in process_span.payload["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_messaging_id(span_streaming, init_celery, capture_events, capture_items): +def test_messaging_id(init_celery, capture_items): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task def example_task(): ... - if span_streaming: - items = capture_items("span") - example_task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert "messaging.message.id" in process_span.payload["attributes"] - else: - events = capture_events() - example_task.apply_async() - (event,) = events - (span,) = event["spans"] - assert "messaging.message.id" in span["data"] + items = capture_items("span") + example_task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert "messaging.message.id" in process_span.payload["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_retry_count_zero(span_streaming, init_celery, capture_events, capture_items): +def test_retry_count_zero(init_celery, capture_items): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert process_span.payload["attributes"]["messaging.message.retry.count"] == 0 - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert span["data"]["messaging.message.retry.count"] == 0 + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert process_span.payload["attributes"]["messaging.message.retry.count"] == 0 -@pytest.mark.parametrize("span_streaming", [True, False]) @mock.patch("celery.app.task.Task.request") -def test_retry_count_nonzero( - mock_request, span_streaming, init_celery, capture_events, capture_items -): +def test_retry_count_nonzero(mock_request, init_celery, capture_items): mock_request.retries = 3 celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert process_span.payload["attributes"]["messaging.message.retry.count"] == 3 - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert span["data"]["messaging.message.retry.count"] == 3 + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert process_span.payload["attributes"]["messaging.message.retry.count"] == 3 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("system", ("redis", "amqp")) -def test_messaging_system( - system, span_streaming, init_celery, capture_events, capture_items -): +def test_messaging_system(system, init_celery, capture_items): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) # Does not need to be a real URL, since we use always eager @@ -1005,25 +848,15 @@ def test_messaging_system( @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert process_span.payload["attributes"]["messaging.system"] == system - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert span["data"]["messaging.system"] == system + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert process_span.payload["attributes"]["messaging.system"] == system -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("system", ("amqp", "redis")) -def test_producer_span_data( - system, span_streaming, monkeypatch, sentry_init, capture_events, capture_items -): +def test_producer_span_data(system, monkeypatch, sentry_init, capture_items): old_publish = kombu.messaging.Producer._publish def publish(*args, **kwargs): @@ -1034,111 +867,69 @@ def publish(*args, **kwargs): sentry_init( integrations=[CeleryIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) celery = Celery(__name__, broker=f"{system}://example.com") # noqa: E231 @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span(name="producer test"): - task.apply_async() + items = capture_items("span") - sentry_sdk.flush() - - span_items = [item.payload for item in items] - publish_span = next( - s for s in span_items if s["attributes"].get("sentry.op") == "queue.publish" - ) - - assert publish_span["attributes"]["messaging.system"] == system - assert publish_span["attributes"]["messaging.destination.name"] == "celery" - assert "messaging.message.id" in publish_span["attributes"] - assert publish_span["attributes"]["messaging.message.retry.count"] == 0 - else: - events = capture_events() + with sentry_sdk.traces.start_span(name="producer test"): + task.apply_async() - with start_transaction(): - task.apply_async() + sentry_sdk.flush() - (event,) = events - span = next(span for span in event["spans"] if span["op"] == "queue.publish") + span_items = [item.payload for item in items] + publish_span = next( + s for s in span_items if s["attributes"].get("sentry.op") == "queue.publish" + ) - assert span["data"]["messaging.system"] == system - assert span["data"]["messaging.destination.name"] == "celery" - assert "messaging.message.id" in span["data"] - assert span["data"]["messaging.message.retry.count"] == 0 + assert publish_span["attributes"]["messaging.system"] == system + assert publish_span["attributes"]["messaging.destination.name"] == "celery" + assert "messaging.message.id" in publish_span["attributes"] + assert publish_span["attributes"]["messaging.message.retry.count"] == 0 monkeypatch.setattr(kombu.messaging.Producer, "_publish", old_publish) -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_receive_latency(span_streaming, init_celery, capture_events, capture_items): +def test_receive_latency(init_celery, capture_items): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, _execution_span = items - assert "messaging.message.receive.latency" in process_span.payload["attributes"] - assert ( - process_span.payload["attributes"]["messaging.message.receive.latency"] > 0 - ) - else: - events = capture_events() - task.apply_async() - (event,) = events - (span,) = event["spans"] - assert "messaging.message.receive.latency" in span["data"] - assert span["data"]["messaging.message.receive.latency"] > 0 + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, _execution_span = items + assert "messaging.message.receive.latency" in process_span.payload["attributes"] + assert process_span.payload["attributes"]["messaging.message.receive.latency"] > 0 -@pytest.mark.parametrize("span_streaming", [True, False]) -def tests_span_origin_consumer( - span_streaming, init_celery, capture_events, capture_items -): +def tests_span_origin_consumer(init_celery, capture_items): celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) celery.conf.broker_url = "redis://example.com" # noqa: E231 @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - task.apply_async() - sentry_sdk.flush() - process_span, execution_span = items - assert ( - execution_span.payload["attributes"]["sentry.origin"] == "auto.queue.celery" - ) - assert ( - process_span.payload["attributes"]["sentry.origin"] == "auto.queue.celery" - ) - else: - events = capture_events() - task.apply_async() - (event,) = events - assert event["contexts"]["trace"]["origin"] == "auto.queue.celery" - assert event["spans"][0]["origin"] == "auto.queue.celery" + items = capture_items("span") + task.apply_async() + sentry_sdk.flush() + process_span, execution_span = items + assert execution_span.payload["attributes"]["sentry.origin"] == "auto.queue.celery" + assert process_span.payload["attributes"]["sentry.origin"] == "auto.queue.celery" -@pytest.mark.parametrize("span_streaming", [True, False]) -def tests_span_origin_producer( - span_streaming, monkeypatch, sentry_init, capture_events, capture_items -): +def tests_span_origin_producer(monkeypatch, sentry_init, capture_items): old_publish = kombu.messaging.Producer._publish def publish(*args, **kwargs): @@ -1149,74 +940,49 @@ def publish(*args, **kwargs): sentry_init( integrations=[CeleryIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) celery = Celery(__name__, broker="redis://example.com") # noqa: E231 @celery.task() def task(): ... - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span(name="custom parent"): - task.apply_async() - - sentry_sdk.flush() - - parent = items.pop(-1).payload - assert parent["name"] == "custom parent" - assert parent["attributes"]["sentry.origin"] == "manual" - - for item in items: - assert item.payload["attributes"]["sentry.origin"] == "auto.queue.celery" - else: - events = capture_events() + items = capture_items("span") - with start_transaction(name="custom_transaction"): - task.apply_async() + with sentry_sdk.traces.start_span(name="custom parent"): + task.apply_async() - (event,) = events + sentry_sdk.flush() - assert event["contexts"]["trace"]["origin"] == "manual" + parent = items.pop(-1).payload + assert parent["name"] == "custom parent" + assert parent["attributes"]["sentry.origin"] == "manual" - for span in event["spans"]: - assert span["origin"] == "auto.queue.celery" + for item in items: + assert item.payload["attributes"]["sentry.origin"] == "auto.queue.celery" monkeypatch.setattr(kombu.messaging.Producer, "_publish", old_publish) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.forked @mock.patch("celery.Celery.send_task") def test_send_task_wrapped( patched_send_task, - span_streaming, sentry_init, - capture_events, capture_items, reset_integrations, ): sentry_init( integrations=[CeleryIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) celery = Celery(__name__, broker="redis://example.com") # noqa: E231 - if span_streaming: - items = capture_items("span") - with sentry_sdk.traces.start_span(name="custom parent") as outer_span: - celery.send_task( - "very_creative_task_name", args=(1, 2), kwargs={"foo": "bar"} - ) - sentry_sdk.flush() - else: - events = capture_events() - with sentry_sdk.start_transaction(name="custom_transaction"): - celery.send_task( - "very_creative_task_name", args=(1, 2), kwargs={"foo": "bar"} - ) + items = capture_items("span") + with sentry_sdk.traces.start_span(name="custom parent") as outer_span: + celery.send_task("very_creative_task_name", args=(1, 2), kwargs={"foo": "bar"}) + sentry_sdk.flush() (call,) = patched_send_task.call_args_list # We should have exactly one call (args, kwargs) = call @@ -1240,32 +1006,18 @@ def test_send_task_wrapped( == kwargs["headers"]["headers"]["sentry-trace"] ) - if span_streaming: - submit_span, outer = [item.payload for item in items] - - assert outer["name"] == "custom parent" - assert outer["is_segment"] is True - - assert submit_span["name"] == "very_creative_task_name" - assert submit_span["attributes"]["sentry.op"] == "queue.submit.celery" - assert submit_span["trace_id"] == outer_span.trace_id - assert ( - submit_span["trace_id"] == kwargs["headers"]["sentry-trace"].split("-")[0] - ) + submit_span, outer = [item.payload for item in items] - else: - (event,) = events - assert event["type"] == "transaction" - assert event["transaction"] == "custom_transaction" + assert outer["name"] == "custom parent" + assert outer["is_segment"] is True - (span,) = event["spans"] - assert span["description"] == "very_creative_task_name" - assert span["op"] == "queue.submit.celery" - assert span["trace_id"] == kwargs["headers"]["sentry-trace"].split("-")[0] + assert submit_span["name"] == "very_creative_task_name" + assert submit_span["attributes"]["sentry.op"] == "queue.submit.celery" + assert submit_span["trace_id"] == outer_span.trace_id + assert submit_span["trace_id"] == kwargs["headers"]["sentry-trace"].split("-")[0] -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_user_custom_headers_accessible_in_task(span_streaming, init_celery): +def test_user_custom_headers_accessible_in_task(init_celery): """ Regression test for https://github.com/getsentry/sentry-python/issues/5566 @@ -1274,7 +1026,7 @@ def test_user_custom_headers_accessible_in_task(span_streaming, init_celery): """ celery = init_celery( traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @celery.task(name="custom_headers_task", bind=True) @@ -1287,12 +1039,8 @@ def custom_headers_task(self): "tenant_id": "tenant-42", } - if span_streaming: - with sentry_sdk.traces.start_span(name="test"): - result = custom_headers_task.apply_async(headers=custom_headers) - else: - with start_transaction(name="test"): - result = custom_headers_task.apply_async(headers=custom_headers) + with sentry_sdk.traces.start_span(name="test"): + result = custom_headers_task.apply_async(headers=custom_headers) received_headers = result.get() for key, value in custom_headers.items():