From 5011bde41275883219118c49850836cf7f61ca60 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Thu, 18 Apr 2024 20:36:50 +0200 Subject: [PATCH 01/15] The working start --- sentry_sdk/integrations/celery/__init__.py | 139 +++++++++------------ 1 file changed, 59 insertions(+), 80 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index b3cbfe8acb..0427590d83 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -157,101 +157,80 @@ def event_processor(event, hint): def _wrap_apply_async(f): # type: (F) -> F - """ - Apply_async is always called to put a task in the queue. This is called by the - celery client (for example the Django project or the Celery Beat process) - """ - @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) + propagate_traces = kwarg_headers.pop( + "sentry-propagate-traces", integration.propagate_traces + ) + + if not propagate_traces: + return f(*args, **kwargs) + + try: + task_started_from_beat = args[1][0] == "BEAT" + except (IndexError, TypeError): + task_started_from_beat = False + task = args[0] - # Do not create a span when the task is a Celery Beat task - # (Because we do not have a transaction in that case) span_mgr = ( sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task.name) - if not Scope.get_isolation_scope()._name == "celery-beat" + if not task_started_from_beat else NoOpMgr() ) # type: Union[Span, NoOpMgr] with span_mgr as span: - incoming_headers = kwargs.get("headers") or {} - integration = sentry_sdk.get_client().get_integration(CeleryIntegration) - - # If Sentry Crons monitoring for Celery Beat tasks is enabled - # add start timestamp of task, - if integration is not None and integration.monitor_beat_tasks: - incoming_headers.update( - { - "sentry-monitor-start-timestamp-s": "%.9f" - % _now_seconds_since_epoch(), - } + with capture_internal_exceptions(): + headers = ( + dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) + if span is not None + else {} ) - - # Propagate Sentry trace information into the Celery task if desired - default_propagate_traces = ( - integration.propagate_traces if integration is not None else True - ) - propagate_traces = incoming_headers.pop( - "sentry-propagate-traces", default_propagate_traces - ) - - if propagate_traces: - with capture_internal_exceptions(): - sentry_trace_headers = dict( - Scope.get_current_scope().iter_trace_propagation_headers( - span=span - ) + if integration.monitor_beat_tasks: + headers.update( + { + "sentry-monitor-start-timestamp-s": "%.9f" + % _now_seconds_since_epoch(), + } ) - # Set Sentry trace data in the headers of the Celery task - if sentry_trace_headers: - # Make sure we don't overwrite existing baggage - incoming_baggage = incoming_headers.get(BAGGAGE_HEADER_NAME) - sentry_baggage = sentry_trace_headers.get(BAGGAGE_HEADER_NAME) - - combined_baggage = sentry_baggage or incoming_baggage - if sentry_baggage and incoming_baggage: - combined_baggage = "{},{}".format( - incoming_baggage, - sentry_baggage, - ) - - # Set Sentry trace data to the headers of the Celery task - incoming_headers.update(sentry_trace_headers) - - if combined_baggage: - incoming_headers[BAGGAGE_HEADER_NAME] = combined_baggage - - # Set sentry trace data also to the inner headers of the Celery task - # https://github.com/celery/celery/issues/4875 - # - # Need to setdefault the inner headers too since other - # tracing tools (dd-trace-py) also employ this exact - # workaround and we don't want to break them. - incoming_headers.setdefault("headers", {}).update( - sentry_trace_headers + + if headers: + existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) + sentry_baggage = headers.get(BAGGAGE_HEADER_NAME) + + combined_baggage = sentry_baggage or existing_baggage + if sentry_baggage and existing_baggage: + combined_baggage = "{},{}".format( + existing_baggage, + sentry_baggage, ) - if combined_baggage: - incoming_headers["headers"][ - BAGGAGE_HEADER_NAME - ] = combined_baggage - - # Add the Sentry options potentially added in `sentry_sdk.integrations.beat.sentry_apply_entry` - # to the inner headers (done when auto-instrumenting Celery Beat tasks) - # https://github.com/celery/celery/issues/4875 - # - # Need to setdefault the inner headers too since other - # tracing tools (dd-trace-py) also employ this exact - # workaround and we don't want to break them. - incoming_headers.setdefault("headers", {}) - for key, value in incoming_headers.items(): - if key.startswith("sentry-"): - incoming_headers["headers"][key] = value - - # Run the task (with updated headers in kwargs) - kwargs["headers"] = incoming_headers + + kwarg_headers.update(headers) + if combined_baggage: + kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage + + # https://github.com/celery/celery/issues/4875 + # + # Need to setdefault the inner headers too since other + # tracing tools (dd-trace-py) also employ this exact + # workaround and we don't want to break them. + kwarg_headers.setdefault("headers", {}).update(headers) + if combined_baggage: + kwarg_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage + + # Add the Sentry options potentially added in `sentry_apply_entry` + # to the headers (done when auto-instrumenting Celery Beat tasks) + for key, value in kwarg_headers.items(): + if key.startswith("sentry-"): + kwarg_headers["headers"][key] = value + + kwargs["headers"] = kwarg_headers return f(*args, **kwargs) From 969b62c819268a7bdd62b1ac9e0301895055b7a4 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Thu, 18 Apr 2024 21:00:14 +0200 Subject: [PATCH 02/15] Added debug output --- sentry_sdk/integrations/celery/__init__.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 0427590d83..0457cc1020 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -161,6 +161,12 @@ def _wrap_apply_async(f): @ensure_integration_enabled(CeleryIntegration, f) def apply_async(*args, **kwargs): # type: (*Any, **Any) -> Any + print("\n\n\n\n\n\n\n\n") + print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") + print(f'kwargs incoming:') + from pprint import pprint + pprint(kwargs) + # Note: kwargs can contain headers=None, so no setdefault! # Unsure which backend though. kwarg_headers = kwargs.get("headers") or {} @@ -170,6 +176,10 @@ def apply_async(*args, **kwargs): ) if not propagate_traces: + print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") + print("111 kwargs updated:") + pprint(kwargs) + return f(*args, **kwargs) try: @@ -232,6 +242,10 @@ def apply_async(*args, **kwargs): kwargs["headers"] = kwarg_headers + print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") + print("222 kwargs updated:") + pprint(kwargs) + return f(*args, **kwargs) return apply_async # type: ignore From 8a6d6c6247ede257f415450496f2658bfc552bed Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Thu, 18 Apr 2024 21:25:20 +0200 Subject: [PATCH 03/15] Moved stuff to refactor into function --- sentry_sdk/integrations/celery/__init__.py | 101 ++++++++++++--------- 1 file changed, 56 insertions(+), 45 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 0457cc1020..44dc12dba1 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -155,6 +155,61 @@ def event_processor(event, hint): return event_processor +def do_stuff_with_kwargs(kwargs, span, integration): + kwarg_headers = kwargs.get("headers") or {} + propagate_traces = kwarg_headers.pop( + "sentry-propagate-traces", integration.propagate_traces + ) + + with capture_internal_exceptions(): + headers = ( + dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) + if span is not None + else {} + ) + if integration.monitor_beat_tasks: + headers.update( + { + "sentry-monitor-start-timestamp-s": "%.9f" + % _now_seconds_since_epoch(), + } + ) + + if headers: + existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) + sentry_baggage = headers.get(BAGGAGE_HEADER_NAME) + + combined_baggage = sentry_baggage or existing_baggage + if sentry_baggage and existing_baggage: + combined_baggage = "{},{}".format( + existing_baggage, + sentry_baggage, + ) + + kwarg_headers.update(headers) + if combined_baggage: + kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage + + # https://github.com/celery/celery/issues/4875 + # + # Need to setdefault the inner headers too since other + # tracing tools (dd-trace-py) also employ this exact + # workaround and we don't want to break them. + kwarg_headers.setdefault("headers", {}).update(headers) + if combined_baggage: + kwarg_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage + + # Add the Sentry options potentially added in `sentry_apply_entry` + # to the headers (done when auto-instrumenting Celery Beat tasks) + for key, value in kwarg_headers.items(): + if key.startswith("sentry-"): + kwarg_headers["headers"][key] = value + + kwargs["headers"] = kwarg_headers + + return kwargs + + def _wrap_apply_async(f): # type: (F) -> F @wraps(f) @@ -196,51 +251,7 @@ def apply_async(*args, **kwargs): ) # type: Union[Span, NoOpMgr] with span_mgr as span: - with capture_internal_exceptions(): - headers = ( - dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) - if span is not None - else {} - ) - if integration.monitor_beat_tasks: - headers.update( - { - "sentry-monitor-start-timestamp-s": "%.9f" - % _now_seconds_since_epoch(), - } - ) - - if headers: - existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) - sentry_baggage = headers.get(BAGGAGE_HEADER_NAME) - - combined_baggage = sentry_baggage or existing_baggage - if sentry_baggage and existing_baggage: - combined_baggage = "{},{}".format( - existing_baggage, - sentry_baggage, - ) - - kwarg_headers.update(headers) - if combined_baggage: - kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage - - # https://github.com/celery/celery/issues/4875 - # - # Need to setdefault the inner headers too since other - # tracing tools (dd-trace-py) also employ this exact - # workaround and we don't want to break them. - kwarg_headers.setdefault("headers", {}).update(headers) - if combined_baggage: - kwarg_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage - - # Add the Sentry options potentially added in `sentry_apply_entry` - # to the headers (done when auto-instrumenting Celery Beat tasks) - for key, value in kwarg_headers.items(): - if key.startswith("sentry-"): - kwarg_headers["headers"][key] = value - - kwargs["headers"] = kwarg_headers + kwargs = do_stuff_with_kwargs(kwargs, span, integration) print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") print("222 kwargs updated:") From b4bc09f219a4a63c731b019936035448e582657f Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Thu, 18 Apr 2024 21:43:08 +0200 Subject: [PATCH 04/15] Added some comments to make some sense --- sentry_sdk/integrations/celery/__init__.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 44dc12dba1..a9859e75e0 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -156,10 +156,13 @@ def event_processor(event, hint): def do_stuff_with_kwargs(kwargs, span, integration): + # Note: kwargs can contain headers=None, so no setdefault! + # Unsure which backend though. + kwarg_headers = kwargs.get("headers") or {} propagate_traces = kwarg_headers.pop( "sentry-propagate-traces", integration.propagate_traces - ) + ) # this is popped here, because it is also popped outside of this function with capture_internal_exceptions(): headers = ( @@ -237,6 +240,8 @@ def apply_async(*args, **kwargs): return f(*args, **kwargs) + # This did never work, because the args is never "BEAT"! + # see refactoring in rc4 to see how this is checked now. (in some other code the tracing info is reset, i think) try: task_started_from_beat = args[1][0] == "BEAT" except (IndexError, TypeError): From fd3b304b95bdb1ac85309aed7bf6795e8f0d4b10 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 09:16:50 +0200 Subject: [PATCH 05/15] Only pass the headers into the function, because it only changes the headers --- sentry_sdk/integrations/celery/__init__.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index a9859e75e0..0904dc89e6 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -155,15 +155,12 @@ def event_processor(event, hint): return event_processor -def do_stuff_with_kwargs(kwargs, span, integration): - # Note: kwargs can contain headers=None, so no setdefault! - # Unsure which backend though. - - kwarg_headers = kwargs.get("headers") or {} - propagate_traces = kwarg_headers.pop( - "sentry-propagate-traces", integration.propagate_traces - ) # this is popped here, because it is also popped outside of this function - +def _update_celery_task_headers(kwarg_headers, span, integration): + # type: (dict, Optional[Span], CeleryIntegration) -> dict + """ + Updates the headers of the Celery task with the tracing information + and eventually Sentry Crons monitroing information. + """ with capture_internal_exceptions(): headers = ( dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) @@ -208,9 +205,7 @@ def do_stuff_with_kwargs(kwargs, span, integration): if key.startswith("sentry-"): kwarg_headers["headers"][key] = value - kwargs["headers"] = kwarg_headers - - return kwargs + return kwarg_headers def _wrap_apply_async(f): @@ -256,7 +251,7 @@ def apply_async(*args, **kwargs): ) # type: Union[Span, NoOpMgr] with span_mgr as span: - kwargs = do_stuff_with_kwargs(kwargs, span, integration) + kwargs["headers"] = _update_celery_task_headers(kwarg_headers, span, integration) print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") print("222 kwargs updated:") From 3a9d4702776380831357b3f44741fb525b9cfec1 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 09:21:29 +0200 Subject: [PATCH 06/15] Just pass the flag and not the whole integration into function --- sentry_sdk/integrations/celery/__init__.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 0904dc89e6..7b47f8d087 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -155,11 +155,11 @@ def event_processor(event, hint): return event_processor -def _update_celery_task_headers(kwarg_headers, span, integration): - # type: (dict, Optional[Span], CeleryIntegration) -> dict +def _update_celery_task_headers(kwarg_headers, span, monitor_beat_tasks): + # type: (dict[str, Any], Optional[Span], bool) -> dict[str, Any] """ Updates the headers of the Celery task with the tracing information - and eventually Sentry Crons monitroing information. + and eventually Sentry Crons monitoring information for beat tasks. """ with capture_internal_exceptions(): headers = ( @@ -167,7 +167,7 @@ def _update_celery_task_headers(kwarg_headers, span, integration): if span is not None else {} ) - if integration.monitor_beat_tasks: + if monitor_beat_tasks: headers.update( { "sentry-monitor-start-timestamp-s": "%.9f" @@ -216,8 +216,9 @@ def apply_async(*args, **kwargs): # type: (*Any, **Any) -> Any print("\n\n\n\n\n\n\n\n") print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") - print(f'kwargs incoming:') + print("kwargs incoming:") from pprint import pprint + pprint(kwargs) # Note: kwargs can contain headers=None, so no setdefault! @@ -251,7 +252,9 @@ def apply_async(*args, **kwargs): ) # type: Union[Span, NoOpMgr] with span_mgr as span: - kwargs["headers"] = _update_celery_task_headers(kwarg_headers, span, integration) + kwargs["headers"] = _update_celery_task_headers( + kwarg_headers, span, integration.monitor_beat_tasks + ) print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") print("222 kwargs updated:") From e68343156d57dbec7d19a12ed720e3cfa9f36d50 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 10:11:59 +0200 Subject: [PATCH 07/15] Do not change the passed in dict in the function --- sentry_sdk/integrations/celery/__init__.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 7b47f8d087..1d9e0c8063 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -155,12 +155,13 @@ def event_processor(event, hint): return event_processor -def _update_celery_task_headers(kwarg_headers, span, monitor_beat_tasks): +def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): # type: (dict[str, Any], Optional[Span], bool) -> dict[str, Any] """ Updates the headers of the Celery task with the tracing information and eventually Sentry Crons monitoring information for beat tasks. """ + updated_headers = original_headers.copy() with capture_internal_exceptions(): headers = ( dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) @@ -176,7 +177,7 @@ def _update_celery_task_headers(kwarg_headers, span, monitor_beat_tasks): ) if headers: - existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) + existing_baggage = updated_headers.get(BAGGAGE_HEADER_NAME) sentry_baggage = headers.get(BAGGAGE_HEADER_NAME) combined_baggage = sentry_baggage or existing_baggage @@ -186,26 +187,26 @@ def _update_celery_task_headers(kwarg_headers, span, monitor_beat_tasks): sentry_baggage, ) - kwarg_headers.update(headers) + updated_headers.update(headers) if combined_baggage: - kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage + updated_headers[BAGGAGE_HEADER_NAME] = combined_baggage # https://github.com/celery/celery/issues/4875 # # Need to setdefault the inner headers too since other # tracing tools (dd-trace-py) also employ this exact # workaround and we don't want to break them. - kwarg_headers.setdefault("headers", {}).update(headers) + updated_headers.setdefault("headers", {}).update(headers) if combined_baggage: - kwarg_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage + updated_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage # Add the Sentry options potentially added in `sentry_apply_entry` # to the headers (done when auto-instrumenting Celery Beat tasks) - for key, value in kwarg_headers.items(): + for key, value in updated_headers.items(): if key.startswith("sentry-"): - kwarg_headers["headers"][key] = value + updated_headers["headers"][key] = value - return kwarg_headers + return updated_headers def _wrap_apply_async(f): @@ -255,7 +256,6 @@ def apply_async(*args, **kwargs): kwargs["headers"] = _update_celery_task_headers( kwarg_headers, span, integration.monitor_beat_tasks ) - print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") print("222 kwargs updated:") pprint(kwargs) From 57fa249e167522f14a3970535e84d816d4c714cf Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 14:06:29 +0200 Subject: [PATCH 08/15] Added more tests --- sentry_sdk/integrations/celery/__init__.py | 2 +- .../celery/test_update_celery_task_headers.py | 159 ++++++++++++++++++ 2 files changed, 160 insertions(+), 1 deletion(-) create mode 100644 tests/integrations/celery/test_update_celery_task_headers.py diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 1d9e0c8063..e605bec610 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -164,7 +164,7 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): updated_headers = original_headers.copy() with capture_internal_exceptions(): headers = ( - dict(Scope.get_current_scope().iter_trace_propagation_headers(span)) + dict(Scope.get_current_scope().iter_trace_propagation_headers(span=span)) if span is not None else {} ) diff --git a/tests/integrations/celery/test_update_celery_task_headers.py b/tests/integrations/celery/test_update_celery_task_headers.py new file mode 100644 index 0000000000..0f7e5d98ad --- /dev/null +++ b/tests/integrations/celery/test_update_celery_task_headers.py @@ -0,0 +1,159 @@ +import pytest + +from unittest import mock + +from sentry_sdk.integrations.celery import _update_celery_task_headers +import sentry_sdk + + +BAGGAGE_VALUE = ( + "sentry-trace_id=771a43a4192642f0b136d5159a501700," + "sentry-public_key=49d0f7386ad645858ae85020e393bef3," + "sentry-sample_rate=0.1337," + "custom=value" +) + +SENTRY_TRACE_VALUE = "771a43a4192642f0b136d5159a501700-1234567890abcdef-1" + + +@pytest.mark.parametrize("monitor_beat_tasks", [True, False, None, "", "bla", 1, 0]) +def test_monitor_beat_tasks(monitor_beat_tasks): + headers = {} + span = None + + updated_headers = _update_celery_task_headers(headers, span, monitor_beat_tasks) + + assert headers == {} # left unchanged + + if monitor_beat_tasks: + assert updated_headers == { + "headers": {"sentry-monitor-start-timestamp-s": mock.ANY}, + "sentry-monitor-start-timestamp-s": mock.ANY, + } + else: + assert updated_headers == headers + + +@pytest.mark.parametrize("monitor_beat_tasks", [True, False, None, "", "bla", 1, 0]) +def test_monitor_beat_tasks_with_headers(monitor_beat_tasks): + headers = { + "blub": "foo", + "sentry-something": "bar", + } + span = None + + updated_headers = _update_celery_task_headers(headers, span, monitor_beat_tasks) + + if monitor_beat_tasks: + assert updated_headers == { + "blub": "foo", + "sentry-something": "bar", + "headers": { + "sentry-monitor-start-timestamp-s": mock.ANY, + "sentry-something": "bar", + }, + "sentry-monitor-start-timestamp-s": mock.ANY, + } + else: + assert updated_headers == headers + + +def test_span_with_transaction(sentry_init): + sentry_init(enable_tracing=True) + headers = {} + + with sentry_sdk.start_transaction(name="test_transaction") as transaction: + with sentry_sdk.start_span(op="test_span") as span: + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + assert updated_headers["baggage"] == transaction.get_baggage().serialize() + assert ( + updated_headers["headers"]["baggage"] + == transaction.get_baggage().serialize() + ) + + +def test_span_with_no_transaction(sentry_init): + sentry_init(enable_tracing=True) + headers = {} + + with sentry_sdk.start_span(op="test_span") as span: + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + assert "baggage" not in updated_headers.keys() + assert "baggage" not in updated_headers["headers"].keys() + + +def test_custom_span(sentry_init): + sentry_init(enable_tracing=True) + span = sentry_sdk.tracing.Span() + headers = {} + + with sentry_sdk.start_transaction(name="test_transaction"): + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + assert "baggage" not in updated_headers.keys() + assert "baggage" not in updated_headers["headers"].keys() + + +def test_span_with_transaction_custom_headers(sentry_init): + sentry_init(enable_tracing=True) + headers = { + "baggage": BAGGAGE_VALUE, + "sentry-trace": SENTRY_TRACE_VALUE, + } + + with sentry_sdk.start_transaction(name="test_transaction") as transaction: + with sentry_sdk.start_span(op="test_span") as span: + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + # This is probably the cause for https://github.com/getsentry/sentry-python/issues/2916 + assert ( + updated_headers["baggage"] + == headers["baggage"] + "," + transaction.get_baggage().serialize() + ) + assert ( + updated_headers["headers"]["baggage"] + == headers["baggage"] + "," + transaction.get_baggage().serialize() + ) + + +def test_span_with_no_transaction_custom_headers(sentry_init): + sentry_init(enable_tracing=True) + headers = { + "baggage": BAGGAGE_VALUE, + "sentry-trace": SENTRY_TRACE_VALUE, + } + + with sentry_sdk.start_span(op="test_span") as span: + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + assert updated_headers["baggage"] == headers["baggage"] + assert updated_headers["headers"]["baggage"] == headers["baggage"] + + +def test_custom_span_custom_headers(sentry_init): + sentry_init(enable_tracing=True) + span = sentry_sdk.tracing.Span() + headers = { + "baggage": BAGGAGE_VALUE, + "sentry-trace": SENTRY_TRACE_VALUE, + } + + with sentry_sdk.start_transaction(name="test_transaction"): + updated_headers = _update_celery_task_headers(headers, span, False) + + assert updated_headers["sentry-trace"] == span.to_traceparent() + assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() + assert updated_headers["baggage"] == headers["baggage"] + assert updated_headers["headers"]["baggage"] == headers["baggage"] From b02e389d348eb5be393b2543d24c7005e963edc2 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 14:11:20 +0200 Subject: [PATCH 09/15] Correct check for if the task was started by celery-beat --- sentry_sdk/integrations/celery/__init__.py | 24 +--------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index e605bec610..628641a2b9 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -215,13 +215,6 @@ def _wrap_apply_async(f): @ensure_integration_enabled(CeleryIntegration, f) def apply_async(*args, **kwargs): # type: (*Any, **Any) -> Any - print("\n\n\n\n\n\n\n\n") - print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") - print("kwargs incoming:") - from pprint import pprint - - pprint(kwargs) - # Note: kwargs can contain headers=None, so no setdefault! # Unsure which backend though. kwarg_headers = kwargs.get("headers") or {} @@ -231,24 +224,13 @@ def apply_async(*args, **kwargs): ) if not propagate_traces: - print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") - print("111 kwargs updated:") - pprint(kwargs) - return f(*args, **kwargs) - # This did never work, because the args is never "BEAT"! - # see refactoring in rc4 to see how this is checked now. (in some other code the tracing info is reset, i think) - try: - task_started_from_beat = args[1][0] == "BEAT" - except (IndexError, TypeError): - task_started_from_beat = False - task = args[0] span_mgr = ( sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task.name) - if not task_started_from_beat + if not Scope.get_isolation_scope()._name == "celery-beat" else NoOpMgr() ) # type: Union[Span, NoOpMgr] @@ -256,10 +238,6 @@ def apply_async(*args, **kwargs): kwargs["headers"] = _update_celery_task_headers( kwarg_headers, span, integration.monitor_beat_tasks ) - print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") - print("222 kwargs updated:") - pprint(kwargs) - return f(*args, **kwargs) return apply_async # type: ignore From 55009a765c9829385d59cdefe042af6209dc1087 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 14:17:32 +0200 Subject: [PATCH 10/15] bring the comment back --- sentry_sdk/integrations/celery/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 628641a2b9..7b8b0d2c67 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -228,6 +228,8 @@ def apply_async(*args, **kwargs): task = args[0] + # Do not create a span when the task is a Celery Beat task + # (Because we do not have a transaction in that case) span_mgr = ( sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task.name) if not Scope.get_isolation_scope()._name == "celery-beat" From 99e1cb333b80aff3c7e54fa4e4b6c1991dd4b184 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Fri, 19 Apr 2024 14:20:38 +0200 Subject: [PATCH 11/15] more comments --- tests/integrations/celery/test_update_celery_task_headers.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integrations/celery/test_update_celery_task_headers.py b/tests/integrations/celery/test_update_celery_task_headers.py index 0f7e5d98ad..9312e6a623 100644 --- a/tests/integrations/celery/test_update_celery_task_headers.py +++ b/tests/integrations/celery/test_update_celery_task_headers.py @@ -116,6 +116,9 @@ def test_span_with_transaction_custom_headers(sentry_init): assert updated_headers["sentry-trace"] == span.to_traceparent() assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent() # This is probably the cause for https://github.com/getsentry/sentry-python/issues/2916 + # If incoming baggage includes sentry data, we should not concatenate a new baggage value to it + # but just keep the incoming sentry baggage values and concatenate new third-party items to the baggage + # I have some code somewhere where I have implemented this. assert ( updated_headers["baggage"] == headers["baggage"] + "," + transaction.get_baggage().serialize() From 95a522e369a859093496562a4e22928196721111 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Mon, 22 Apr 2024 15:08:31 +0200 Subject: [PATCH 12/15] trigger ci From b2068d37a1ac186406122e8189925fa3fb8c3913 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Mon, 22 Apr 2024 15:52:34 +0200 Subject: [PATCH 13/15] Updated tests that was green by coincidence --- tests/integrations/celery/test_celery.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py index 7e0b533d4c..000bb76b13 100644 --- a/tests/integrations/celery/test_celery.py +++ b/tests/integrations/celery/test_celery.py @@ -501,7 +501,13 @@ def dummy_task(self, x, y): # in the monkey patched version of `apply_async` # in `sentry_sdk/integrations/celery.py::_wrap_apply_async()` result = dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup) - assert result.get() == sentry_crons_setup + + expected_headers = sentry_crons_setup.copy() + # Newly added headers + expected_headers["sentry-trace"] = mock.ANY + expected_headers["baggage"] = mock.ANY + + assert result.get() == expected_headers def test_baggage_propagation(init_celery): From 1a14efb1a4252b6a664d22ce2857532ebf29be7c Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Tue, 23 Apr 2024 09:12:01 +0200 Subject: [PATCH 14/15] Update sentry_sdk/integrations/celery/__init__.py Co-authored-by: Ivana Kellyerova --- sentry_sdk/integrations/celery/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 7b8b0d2c67..f3ba983e5f 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -163,11 +163,9 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): """ updated_headers = original_headers.copy() with capture_internal_exceptions(): - headers = ( - dict(Scope.get_current_scope().iter_trace_propagation_headers(span=span)) - if span is not None - else {} - ) + headers = {} + if span is not None: + headers = dict(Scope.get_current_scope().iter_trace_propagation_headers(span=span)) if monitor_beat_tasks: headers.update( { From 66b601cd3098e4d972798f7dc6c8c9ab16813884 Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Tue, 23 Apr 2024 09:15:53 +0200 Subject: [PATCH 15/15] Formatting --- sentry_sdk/integrations/celery/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index f3ba983e5f..a9febc2570 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -165,7 +165,9 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): with capture_internal_exceptions(): headers = {} if span is not None: - headers = dict(Scope.get_current_scope().iter_trace_propagation_headers(span=span)) + headers = dict( + Scope.get_current_scope().iter_trace_propagation_headers(span=span) + ) if monitor_beat_tasks: headers.update( {