Merge baggage headers (incoming and new created ones)

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:

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
             )