Skip to content

Commit 625e1b3

Browse files
sentrivanasentrivana
authored
Do not overwrite existing baggage on outgoing requests (getsentry#2191)
1 parent 8b505a1 commit 625e1b3

4 files changed

Lines changed: 103 additions & 15 deletions

File tree

‎sentry_sdk/integrations/celery.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from sentry_sdk.hub import Hub
1212
from sentry_sdk.integrations import Integration, DidNotEnable
1313
from sentry_sdk.integrations.logging import ignore_logger
14-
from sentry_sdk.tracing import TRANSACTION_SOURCE_TASK
14+
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TRANSACTION_SOURCE_TASK
1515
from sentry_sdk._types import TYPE_CHECKING
1616
from sentry_sdk.utils import (
1717
capture_internal_exceptions,
@@ -158,14 +158,31 @@ def apply_async(*args, **kwargs):
158158
# Note: kwargs can contain headers=None, so no setdefault!
159159
# Unsure which backend though.
160160
kwarg_headers = kwargs.get("headers") or {}
161+
162+
existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME)
163+
sentry_baggage = headers.get(BAGGAGE_HEADER_NAME)
164+
165+
combined_baggage = sentry_baggage or existing_baggage
166+
if sentry_baggage and existing_baggage:
167+
combined_baggage = "{},{}".format(
168+
existing_baggage,
169+
sentry_baggage,
170+
)
171+
161172
kwarg_headers.update(headers)
173+
if combined_baggage:
174+
kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage
162175

163176
# https://github.com/celery/celery/issues/4875
164177
#
165178
# Need to setdefault the inner headers too since other
166179
# tracing tools (dd-trace-py) also employ this exact
167180
# workaround and we don't want to break them.
168181
kwarg_headers.setdefault("headers", {}).update(headers)
182+
if combined_baggage:
183+
kwarg_headers["headers"][
184+
BAGGAGE_HEADER_NAME
185+
] = combined_baggage
169186

170187
# Add the Sentry options potentially added in `sentry_apply_entry`
171188
# to the headers (done when auto-instrumenting Celery Beat tasks)

‎sentry_sdk/integrations/httpx.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from sentry_sdk import Hub
22
from sentry_sdk.consts import OP, SPANDATA
33
from sentry_sdk.integrations import Integration, DidNotEnable
4+
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
45
from sentry_sdk.tracing_utils import should_propagate_trace
56
from sentry_sdk.utils import (
67
SENSITIVE_DATA_SUBSTITUTE,
@@ -72,7 +73,13 @@ def send(self, request, **kwargs):
7273
key=key, value=value, url=request.url
7374
)
7475
)
75-
request.headers[key] = value
76+
if key == BAGGAGE_HEADER_NAME and request.headers.get(
77+
BAGGAGE_HEADER_NAME
78+
):
79+
# do not overwrite any existing baggage, just append to it
80+
request.headers[key] += "," + value
81+
else:
82+
request.headers[key] = value
7683

7784
rv = real_send(self, request, **kwargs)
7885

@@ -119,7 +126,13 @@ async def send(self, request, **kwargs):
119126
key=key, value=value, url=request.url
120127
)
121128
)
122-
request.headers[key] = value
129+
if key == BAGGAGE_HEADER_NAME and request.headers.get(
130+
BAGGAGE_HEADER_NAME
131+
):
132+
# do not overwrite any existing baggage, just append to it
133+
request.headers[key] += "," + value
134+
else:
135+
request.headers[key] = value
123136

124137
rv = await real_send(self, request, **kwargs)
125138

‎tests/integrations/celery/test_celery.py‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111

1212
from celery import Celery, VERSION
1313
from celery.bin import worker
14-
from celery.signals import task_success
1514

1615
try:
1716
from unittest import mock # python 3.3 and above
@@ -360,7 +359,7 @@ def dummy_task(self):
360359
# TODO: This test is hanging when running test with `tox --parallel auto`. Find out why and fix it!
361360
@pytest.mark.skip
362361
@pytest.mark.forked
363-
def test_redis_backend_trace_propagation(init_celery, capture_events_forksafe, tmpdir):
362+
def test_redis_backend_trace_propagation(init_celery, capture_events_forksafe):
364363
celery = init_celery(traces_sample_rate=1.0, backend="redis", debug=True)
365364

366365
events = capture_events_forksafe()
@@ -493,17 +492,36 @@ def test_task_headers(celery):
493492
"sentry-monitor-check-in-id": "123abc",
494493
}
495494

496-
@celery.task(name="dummy_task")
497-
def dummy_task(x, y):
498-
return x + y
499-
500-
def crons_task_success(sender, **kwargs):
501-
headers = _get_headers(sender)
502-
assert headers == sentry_crons_setup
503-
504-
task_success.connect(crons_task_success)
495+
@celery.task(name="dummy_task", bind=True)
496+
def dummy_task(self, x, y):
497+
return _get_headers(self)
505498

506499
# This is how the Celery Beat auto-instrumentation starts a task
507500
# in the monkey patched version of `apply_async`
508501
# in `sentry_sdk/integrations/celery.py::_wrap_apply_async()`
509-
dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup)
502+
result = dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup)
503+
assert result.get() == sentry_crons_setup
504+
505+
506+
def test_baggage_propagation(init_celery):
507+
celery = init_celery(traces_sample_rate=1.0, release="abcdef")
508+
509+
@celery.task(name="dummy_task", bind=True)
510+
def dummy_task(self, x, y):
511+
return _get_headers(self)
512+
513+
with start_transaction() as transaction:
514+
result = dummy_task.apply_async(
515+
args=(1, 0),
516+
headers={"baggage": "custom=value"},
517+
).get()
518+
519+
assert sorted(result["baggage"].split(",")) == sorted(
520+
[
521+
"sentry-release=abcdef",
522+
"sentry-trace_id={}".format(transaction.trace_id),
523+
"sentry-environment=production",
524+
"sentry-sample_rate=1.0",
525+
"custom=value",
526+
]
527+
)

‎tests/integrations/httpx/test_httpx.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,46 @@ def test_outgoing_trace_headers(sentry_init, httpx_client):
8989
)
9090

9191

92+
@pytest.mark.parametrize(
93+
"httpx_client",
94+
(httpx.Client(), httpx.AsyncClient()),
95+
)
96+
def test_outgoing_trace_headers_append_to_baggage(sentry_init, httpx_client):
97+
sentry_init(
98+
traces_sample_rate=1.0,
99+
integrations=[HttpxIntegration()],
100+
release="d08ebdb9309e1b004c6f52202de58a09c2268e42",
101+
)
102+
103+
url = "http://example.com/"
104+
responses.add(responses.GET, url, status=200)
105+
106+
with start_transaction(
107+
name="/interactions/other-dogs/new-dog",
108+
op="greeting.sniff",
109+
trace_id="01234567890123456789012345678901",
110+
) as transaction:
111+
if asyncio.iscoroutinefunction(httpx_client.get):
112+
response = asyncio.get_event_loop().run_until_complete(
113+
httpx_client.get(url, headers={"baGGage": "custom=data"})
114+
)
115+
else:
116+
response = httpx_client.get(url, headers={"baGGage": "custom=data"})
117+
118+
request_span = transaction._span_recorder.spans[-1]
119+
assert response.request.headers[
120+
"sentry-trace"
121+
] == "{trace_id}-{parent_span_id}-{sampled}".format(
122+
trace_id=transaction.trace_id,
123+
parent_span_id=request_span.span_id,
124+
sampled=1,
125+
)
126+
assert (
127+
response.request.headers["baggage"]
128+
== "custom=data,sentry-trace_id=01234567890123456789012345678901,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0"
129+
)
130+
131+
92132
@pytest.mark.parametrize(
93133
"httpx_client,trace_propagation_targets,url,trace_propagated",
94134
[

0 commit comments

Comments
 (0)