diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index a9febc2570..7305736d5e 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/sentry_sdk/integrations/celery/__init__.py @@ -16,6 +16,7 @@ from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TRANSACTION_SOURCE_TASK from sentry_sdk._types import TYPE_CHECKING from sentry_sdk.scope import Scope +from sentry_sdk.tracing_utils import Baggage from sentry_sdk.utils import ( capture_internal_exceptions, ensure_integration_enabled, @@ -168,6 +169,7 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): headers = dict( Scope.get_current_scope().iter_trace_propagation_headers(span=span) ) + if monitor_beat_tasks: headers.update( { @@ -182,10 +184,23 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks): combined_baggage = sentry_baggage or existing_baggage if sentry_baggage and existing_baggage: - combined_baggage = "{},{}".format( - existing_baggage, - sentry_baggage, + # Merge incoming and sentry baggage, where the sentry trace information + # in the incoming baggage takes precedence and the third-party items + # are concatenated. + incoming = Baggage.from_incoming_header(existing_baggage) + combined = Baggage.from_incoming_header(sentry_baggage) + combined.sentry_items.update(incoming.sentry_items) + combined.third_party_items = ",".join( + [ + x + for x in [ + combined.third_party_items, + incoming.third_party_items, + ] + if x is not None and x != "" + ] ) + combined_baggage = combined.serialize(include_third_party=True) updated_headers.update(headers) if combined_baggage: diff --git a/tests/integrations/celery/test_update_celery_task_headers.py b/tests/integrations/celery/test_update_celery_task_headers.py index 9312e6a623..b1588e86b8 100644 --- a/tests/integrations/celery/test_update_celery_task_headers.py +++ b/tests/integrations/celery/test_update_celery_task_headers.py @@ -1,9 +1,11 @@ +from copy import copy import pytest from unittest import mock from sentry_sdk.integrations.celery import _update_celery_task_headers import sentry_sdk +from sentry_sdk.tracing_utils import Baggage BAGGAGE_VALUE = ( @@ -115,17 +117,25 @@ 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() + + incoming_baggage = Baggage.from_incoming_header(headers["baggage"]) + combined_baggage = copy(transaction.get_baggage()) + combined_baggage.sentry_items.update(incoming_baggage.sentry_items) + combined_baggage.third_party_items = ",".join( + [ + x + for x in [ + combined_baggage.third_party_items, + incoming_baggage.third_party_items, + ] + if x is not None and x != "" + ] ) - assert ( - updated_headers["headers"]["baggage"] - == headers["baggage"] + "," + transaction.get_baggage().serialize() + assert updated_headers["baggage"] == combined_baggage.serialize( + include_third_party=True + ) + assert updated_headers["headers"]["baggage"] == combined_baggage.serialize( + include_third_party=True )