From 3b792b2dff48af472dd70757df931c6031f7fe54 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Fri, 23 Jun 2023 12:58:27 +0200 Subject: [PATCH 01/12] Do not overwrite existing baggage on outgoing requests --- sentry_sdk/integrations/httpx.py | 17 +++++++++-- tests/integrations/httpx/test_httpx.py | 40 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/httpx.py b/sentry_sdk/integrations/httpx.py index e84a28d165..04db5047b4 100644 --- a/sentry_sdk/integrations/httpx.py +++ b/sentry_sdk/integrations/httpx.py @@ -1,6 +1,7 @@ from sentry_sdk import Hub from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.tracing import BAGGAGE_HEADER_NAME from sentry_sdk.tracing_utils import should_propagate_trace from sentry_sdk.utils import ( SENSITIVE_DATA_SUBSTITUTE, @@ -72,7 +73,13 @@ def send(self, request, **kwargs): key=key, value=value, url=request.url ) ) - request.headers[key] = value + if key == BAGGAGE_HEADER_NAME and request.headers.get( + BAGGAGE_HEADER_NAME + ): + # do not overwrite any existing baggage, just append to it + request.headers[key] += "," + value + else: + request.headers[key] = value rv = real_send(self, request, **kwargs) @@ -119,7 +126,13 @@ async def send(self, request, **kwargs): key=key, value=value, url=request.url ) ) - request.headers[key] = value + if key == BAGGAGE_HEADER_NAME and request.headers.get( + BAGGAGE_HEADER_NAME + ): + # do not overwrite any existing baggage, just append to it + request.headers[key] += "," + value + else: + request.headers[key] = value rv = await real_send(self, request, **kwargs) diff --git a/tests/integrations/httpx/test_httpx.py b/tests/integrations/httpx/test_httpx.py index 72188a23e3..ab1ee0b8a5 100644 --- a/tests/integrations/httpx/test_httpx.py +++ b/tests/integrations/httpx/test_httpx.py @@ -89,6 +89,46 @@ def test_outgoing_trace_headers(sentry_init, httpx_client): ) +@pytest.mark.parametrize( + "httpx_client", + (httpx.Client(), httpx.AsyncClient()), +) +def test_outgoing_trace_headers_append_to_baggage(sentry_init, httpx_client): + sentry_init( + traces_sample_rate=1.0, + integrations=[HttpxIntegration()], + release="d08ebdb9309e1b004c6f52202de58a09c2268e42", + ) + + url = "http://example.com/" + responses.add(responses.GET, url, status=200) + + with start_transaction( + name="/interactions/other-dogs/new-dog", + op="greeting.sniff", + trace_id="01234567890123456789012345678901", + ) as transaction: + if asyncio.iscoroutinefunction(httpx_client.get): + response = asyncio.get_event_loop().run_until_complete( + httpx_client.get(url, headers={"baGGage": "custom-data"}) + ) + else: + response = httpx_client.get(url, headers={"baGGage": "custom-data"}) + + request_span = transaction._span_recorder.spans[-1] + assert response.request.headers[ + "sentry-trace" + ] == "{trace_id}-{parent_span_id}-{sampled}".format( + trace_id=transaction.trace_id, + parent_span_id=request_span.span_id, + sampled=1, + ) + assert ( + response.request.headers["baggage"] + == "custom-data,sentry-trace_id=01234567890123456789012345678901,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0" + ) + + @pytest.mark.parametrize( "httpx_client,trace_propagation_targets,url,trace_propagated", [ From 191b66aed62a461a8e4a5e9f78dd2be45035ae87 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Fri, 23 Jun 2023 13:31:04 +0200 Subject: [PATCH 02/12] . --- tests/integrations/httpx/test_httpx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integrations/httpx/test_httpx.py b/tests/integrations/httpx/test_httpx.py index ab1ee0b8a5..9b7842fbb7 100644 --- a/tests/integrations/httpx/test_httpx.py +++ b/tests/integrations/httpx/test_httpx.py @@ -110,10 +110,10 @@ def test_outgoing_trace_headers_append_to_baggage(sentry_init, httpx_client): ) as transaction: if asyncio.iscoroutinefunction(httpx_client.get): response = asyncio.get_event_loop().run_until_complete( - httpx_client.get(url, headers={"baGGage": "custom-data"}) + httpx_client.get(url, headers={"baGGage": "custom=data"}) ) else: - response = httpx_client.get(url, headers={"baGGage": "custom-data"}) + response = httpx_client.get(url, headers={"baGGage": "custom=data"}) request_span = transaction._span_recorder.spans[-1] assert response.request.headers[ @@ -125,7 +125,7 @@ def test_outgoing_trace_headers_append_to_baggage(sentry_init, httpx_client): ) assert ( response.request.headers["baggage"] - == "custom-data,sentry-trace_id=01234567890123456789012345678901,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0" + == "custom=data,sentry-trace_id=01234567890123456789012345678901,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0" ) From 0d2d15ce749c562db551182ccbd492486d438d2b Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Fri, 23 Jun 2023 17:19:08 +0200 Subject: [PATCH 03/12] celery --- sentry_sdk/integrations/celery.py | 19 +++++++++-- tests/integrations/celery/test_celery.py | 40 +++++++++++++++++------- 2 files changed, 46 insertions(+), 13 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index 741a2c8bb7..764fb8acdc 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -11,7 +11,7 @@ from sentry_sdk.hub import Hub from sentry_sdk.integrations import Integration, DidNotEnable from sentry_sdk.integrations.logging import ignore_logger -from sentry_sdk.tracing import TRANSACTION_SOURCE_TASK +from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TRANSACTION_SOURCE_TASK from sentry_sdk._types import TYPE_CHECKING from sentry_sdk.utils import ( capture_internal_exceptions, @@ -158,14 +158,29 @@ def apply_async(*args, **kwargs): # Note: kwargs can contain headers=None, so no setdefault! # Unsure which backend though. kwarg_headers = kwargs.get("headers") or {} + existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) kwarg_headers.update(headers) + if existing_baggage: + kwarg_headers[BAGGAGE_HEADER_NAME] = "{},{}".format( + existing_baggage, headers[BAGGAGE_HEADER_NAME] + ) # 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) + kwarg_headers.setdefault("headers", {}) + existing_baggage = kwarg_headers["headers"].get( + BAGGAGE_HEADER_NAME + ) + kwarg_headers["headers"].update(headers) + if existing_baggage: + kwarg_headers["headers"][ + BAGGAGE_HEADER_NAME + ] = "{},{}".format( + existing_baggage, headers[BAGGAGE_HEADER_NAME] + ) # Add the Sentry options potentially added in `sentry_apply_entry` # to the headers (done when auto-instrumenting Celery Beat tasks) diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py index d120d34a12..a3aa298f3a 100644 --- a/tests/integrations/celery/test_celery.py +++ b/tests/integrations/celery/test_celery.py @@ -11,7 +11,6 @@ from celery import Celery, VERSION from celery.bin import worker -from celery.signals import task_success try: from unittest import mock # python 3.3 and above @@ -493,17 +492,36 @@ def test_task_headers(celery): "sentry-monitor-check-in-id": "123abc", } - @celery.task(name="dummy_task") - def dummy_task(x, y): - return x + y - - def crons_task_success(sender, **kwargs): - headers = _get_headers(sender) - assert headers == sentry_crons_setup - - task_success.connect(crons_task_success) + @celery.task(name="dummy_task", bind=True) + def dummy_task(self, x, y): + return _get_headers(self) # This is how the Celery Beat auto-instrumentation starts a task # in the monkey patched version of `apply_async` # in `sentry_sdk/integrations/celery.py::_wrap_apply_async()` - dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup) + result = dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup) + assert result.get() == sentry_crons_setup + + +def test_baggage_propagation(init_celery): + celery = init_celery(traces_sample_rate=1.0, release="abcdef") + + @celery.task(name="dummy_task", bind=True) + def dummy_task(self, x, y): + return _get_headers(self) + + with start_transaction() as transaction: + result = dummy_task.apply_async( + args=(1, 0), + headers={"baggage": "custom=value", "headers": {"baggage": "custom=value"}}, + ).get() + + assert sorted(result["baggage"].split(",")) == sorted( + [ + "sentry-release=abcdef", + "sentry-trace_id={}".format(transaction.trace_id), + "sentry-environment=production", + "sentry-sample_rate=1.0", + "custom=value", + ] + ) From b13da1964383be2e50609dad60def1882a7d6f71 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Fri, 23 Jun 2023 17:21:34 +0200 Subject: [PATCH 04/12] cleanup --- tests/integrations/celery/test_celery.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py index a3aa298f3a..304f6c2f04 100644 --- a/tests/integrations/celery/test_celery.py +++ b/tests/integrations/celery/test_celery.py @@ -359,7 +359,7 @@ def dummy_task(self): # TODO: This test is hanging when running test with `tox --parallel auto`. Find out why and fix it! @pytest.mark.skip @pytest.mark.forked -def test_redis_backend_trace_propagation(init_celery, capture_events_forksafe, tmpdir): +def test_redis_backend_trace_propagation(init_celery, capture_events_forksafe): celery = init_celery(traces_sample_rate=1.0, backend="redis", debug=True) events = capture_events_forksafe() @@ -513,7 +513,7 @@ def dummy_task(self, x, y): with start_transaction() as transaction: result = dummy_task.apply_async( args=(1, 0), - headers={"baggage": "custom=value", "headers": {"baggage": "custom=value"}}, + headers={"baggage": "custom=value"}, ).get() assert sorted(result["baggage"].split(",")) == sorted( From 4ac72c8066fa77e5418c31406d7f7293654f4021 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Mon, 26 Jun 2023 10:03:19 +0200 Subject: [PATCH 05/12] prettify --- sentry_sdk/integrations/celery.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index 764fb8acdc..c5e5c7a777 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -158,29 +158,25 @@ def apply_async(*args, **kwargs): # Note: kwargs can contain headers=None, so no setdefault! # Unsure which backend though. kwarg_headers = kwargs.get("headers") or {} + existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) - kwarg_headers.update(headers) + combined_baggage = existing_baggage if existing_baggage: - kwarg_headers[BAGGAGE_HEADER_NAME] = "{},{}".format( - existing_baggage, headers[BAGGAGE_HEADER_NAME] + combined_baggage += ",{}".format( + headers[BAGGAGE_HEADER_NAME] ) + kwarg_headers.update(headers) + 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", {}) - existing_baggage = kwarg_headers["headers"].get( - BAGGAGE_HEADER_NAME - ) kwarg_headers["headers"].update(headers) - if existing_baggage: - kwarg_headers["headers"][ - BAGGAGE_HEADER_NAME - ] = "{},{}".format( - existing_baggage, headers[BAGGAGE_HEADER_NAME] - ) + 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) From d4416726510dbe826ef4463ade18a868229d70cc Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Mon, 26 Jun 2023 12:40:59 +0200 Subject: [PATCH 06/12] get ci unstuck From 124ef7d8e9f04172b94edb584e1749985bd90c5c Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 10:50:33 +0200 Subject: [PATCH 07/12] mypy --- sentry_sdk/integrations/celery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index c5e5c7a777..dfa24435cb 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -161,7 +161,7 @@ def apply_async(*args, **kwargs): existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) combined_baggage = existing_baggage - if existing_baggage: + if existing_baggage and headers[BAGGAGE_HEADER_NAME]: combined_baggage += ",{}".format( headers[BAGGAGE_HEADER_NAME] ) From 2a71e1c6d301485967d01ceae525b40ef28c6181 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 10:55:44 +0200 Subject: [PATCH 08/12] mypy again --- sentry_sdk/integrations/celery.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index dfa24435cb..37e4e8669b 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -160,11 +160,13 @@ def apply_async(*args, **kwargs): kwarg_headers = kwargs.get("headers") or {} existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) - combined_baggage = existing_baggage - if existing_baggage and headers[BAGGAGE_HEADER_NAME]: - combined_baggage += ",{}".format( - headers[BAGGAGE_HEADER_NAME] + if existing_baggage: + combined_baggage = "{},{}".format( + existing_baggage, + headers[BAGGAGE_HEADER_NAME], ) + else: + combined_baggage = headers[BAGGAGE_HEADER_NAME] kwarg_headers.update(headers) kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage From bd0098a5c51f0cb2c72a2b667c258ad1a8f83fa2 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 11:06:11 +0200 Subject: [PATCH 09/12] fixes --- sentry_sdk/integrations/celery.py | 15 ++++++++++----- sentry_sdk/integrations/httpx.py | 4 +++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index 37e4e8669b..bd11e2da8d 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -160,16 +160,18 @@ def apply_async(*args, **kwargs): kwarg_headers = kwargs.get("headers") or {} existing_baggage = kwarg_headers.get(BAGGAGE_HEADER_NAME) - if existing_baggage: + sentry_baggage = headers.get(BAGGAGE_HEADER_NAME) + if sentry_baggage and existing_baggage: combined_baggage = "{},{}".format( existing_baggage, - headers[BAGGAGE_HEADER_NAME], + sentry_baggage, ) else: - combined_baggage = headers[BAGGAGE_HEADER_NAME] + combined_baggage = sentry_baggage or existing_baggage kwarg_headers.update(headers) - kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage + if combined_baggage: + kwarg_headers[BAGGAGE_HEADER_NAME] = combined_baggage # https://github.com/celery/celery/issues/4875 # @@ -178,7 +180,10 @@ def apply_async(*args, **kwargs): # workaround and we don't want to break them. kwarg_headers.setdefault("headers", {}) kwarg_headers["headers"].update(headers) - kwarg_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage + 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) diff --git a/sentry_sdk/integrations/httpx.py b/sentry_sdk/integrations/httpx.py index 04db5047b4..323e7396fa 100644 --- a/sentry_sdk/integrations/httpx.py +++ b/sentry_sdk/integrations/httpx.py @@ -77,7 +77,9 @@ def send(self, request, **kwargs): BAGGAGE_HEADER_NAME ): # do not overwrite any existing baggage, just append to it - request.headers[key] += "," + value + request.headers[key] = "{},{}".format( + request.headers[key], value + ) else: request.headers[key] = value From 2539399745700741cab1d27f04102403eb0bd801 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 11:07:09 +0200 Subject: [PATCH 10/12] consistency --- sentry_sdk/integrations/httpx.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sentry_sdk/integrations/httpx.py b/sentry_sdk/integrations/httpx.py index 323e7396fa..04db5047b4 100644 --- a/sentry_sdk/integrations/httpx.py +++ b/sentry_sdk/integrations/httpx.py @@ -77,9 +77,7 @@ def send(self, request, **kwargs): BAGGAGE_HEADER_NAME ): # do not overwrite any existing baggage, just append to it - request.headers[key] = "{},{}".format( - request.headers[key], value - ) + request.headers[key] += "," + value else: request.headers[key] = value From 87888de95fd6aa556519da5665a9d60cfa03939b Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 11:08:06 +0200 Subject: [PATCH 11/12] simplify --- sentry_sdk/integrations/celery.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index bd11e2da8d..e28f4183b9 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -178,8 +178,7 @@ def apply_async(*args, **kwargs): # 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", {}) - kwarg_headers["headers"].update(headers) + kwarg_headers.setdefault("headers", {}).update(headers) if combined_baggage: kwarg_headers["headers"][ BAGGAGE_HEADER_NAME From 37e9262a358c09aefa3863a76db0244f4fd63549 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Tue, 27 Jun 2023 11:14:51 +0200 Subject: [PATCH 12/12] mypy?? --- sentry_sdk/integrations/celery.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py index e28f4183b9..443fcdad45 100644 --- a/sentry_sdk/integrations/celery.py +++ b/sentry_sdk/integrations/celery.py @@ -161,13 +161,13 @@ def apply_async(*args, **kwargs): 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, ) - else: - combined_baggage = sentry_baggage or existing_baggage kwarg_headers.update(headers) if combined_baggage: