From 878bea7455863a22505314d9559c2b2e23eaf52e Mon Sep 17 00:00:00 2001
From: Anton Pirker <anton.pirker@sentry.io>
Date: Tue, 30 Jul 2024 17:10:50 +0200
Subject: [PATCH] Add span data to the transactions trace context (#3374)

Fixes #3372
---
 sentry_sdk/tracing.py      |  9 +++++++++
 tests/tracing/test_misc.py | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
index 1db105fa62..41c998cb99 100644
--- a/sentry_sdk/tracing.py
+++ b/sentry_sdk/tracing.py
@@ -1027,6 +1027,15 @@ def to_json(self):
 
         return rv
 
+    def get_trace_context(self):
+        # type: () -> Any
+        trace_context = super().get_trace_context()
+
+        if self._data:
+            trace_context["data"] = self._data
+
+        return trace_context
+
     def get_baggage(self):
         # type: () -> Baggage
         """Returns the :py:class:`~sentry_sdk.tracing_utils.Baggage`
diff --git a/tests/tracing/test_misc.py b/tests/tracing/test_misc.py
index de25acd7d2..02966642fd 100644
--- a/tests/tracing/test_misc.py
+++ b/tests/tracing/test_misc.py
@@ -60,6 +60,33 @@ def test_transaction_naming(sentry_init, capture_events):
     assert events[2]["transaction"] == "a"
 
 
+def test_transaction_data(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events()
+
+    with start_transaction(name="test-transaction"):
+        span_or_tx = sentry_sdk.get_current_span()
+        span_or_tx.set_data("foo", "bar")
+        with start_span(op="test-span") as span:
+            span.set_data("spanfoo", "spanbar")
+
+    assert len(events) == 1
+
+    transaction = events[0]
+    transaction_data = transaction["contexts"]["trace"]["data"]
+
+    assert "data" not in transaction.keys()
+    assert transaction_data.items() >= {"foo": "bar"}.items()
+
+    assert len(transaction["spans"]) == 1
+
+    span = transaction["spans"][0]
+    span_data = span["data"]
+
+    assert "contexts" not in span.keys()
+    assert span_data.items() >= {"spanfoo": "spanbar"}.items()
+
+
 def test_start_transaction(sentry_init):
     sentry_init(traces_sample_rate=1.0)