feat(trace-spans): Allow querying additional attributes (#95101)

Allow querying additional attributes when fetching the spans of a trace.

Author: ArthurKnaus

src/sentry/api/endpoints/organization_trace.py

@@ -71,6 +71,7 @@ class SerializedSpan(SerializedEvent):
     start_timestamp: datetime
     is_transaction: bool
     transaction_id: str
+    additional_attributes: NotRequired[dict[str, Any]]
 
 
 @region_silo_endpoint
@@ -169,12 +170,21 @@ def _qualify_short_id(project: str, short_id: int | None) -> str | None:
             raise Exception(f"Unknown event encountered in trace: {event.get('event_type')}")
 
     def serialize_rpc_event(
-        self, event: dict[str, Any], group_cache: dict[int, Group]
+        self,
+        event: dict[str, Any],
+        group_cache: dict[int, Group],
+        additional_attributes: list[str] | None = None,
     ) -> SerializedEvent | SerializedIssue:
         if event.get("event_type") == "span":
+            attribute_dict = {
+                attribute: event[attribute]
+                for attribute in additional_attributes or []
+                if attribute in event
+            }
             return SerializedSpan(
                 children=[
-                    self.serialize_rpc_event(child, group_cache) for child in event["children"]
+                    self.serialize_rpc_event(child, group_cache, additional_attributes)
+                    for child in event["children"]
                 ],
                 errors=[self.serialize_rpc_issue(error, group_cache) for error in event["errors"]],
                 occurrences=[
@@ -204,6 +214,7 @@ def serialize_rpc_event(
                 op=event["span.op"],
                 name=event["span.name"],
                 event_type="span",
+                additional_attributes=attribute_dict,
             )
         else:
             return self.serialize_rpc_issue(event, group_cache)
@@ -306,7 +317,11 @@ def run_perf_issues_query(self, occurrence_query: DiscoverQueryBuilder):
 
     @sentry_sdk.tracing.trace
     def query_trace_data(
-        self, snuba_params: SnubaParams, trace_id: str, error_id: str | None = None
+        self,
+        snuba_params: SnubaParams,
+        trace_id: str,
+        error_id: str | None = None,
+        additional_attributes: list[str] | None = None,
     ) -> list[SerializedEvent]:
         """Queries span/error data for a given trace"""
         # This is a hack, long term EAP will store both errors and performance_issues eventually but is not ready
@@ -326,6 +341,7 @@ def query_trace_data(
             snuba_params,
             Referrer.API_TRACE_VIEW_GET_EVENTS.value,
             SearchResolverConfig(),
+            additional_attributes,
         )
         errors_future = _query_thread_pool.submit(
             self.run_errors_query,
@@ -376,7 +392,9 @@ def query_trace_data(
         for errors in id_to_error.values():
             result.extend(errors)
         group_cache: dict[int, Group] = {}
-        return [self.serialize_rpc_event(root, group_cache) for root in result]
+        return [
+            self.serialize_rpc_event(root, group_cache, additional_attributes) for root in result
+        ]
 
     def has_feature(self, organization: Organization, request: Request) -> bool:
         return bool(
@@ -393,6 +411,8 @@ def get(self, request: Request, organization: Organization, trace_id: str) -> Ht
         except NoProjects:
             return Response(status=404)
 
+        additional_attributes = request.GET.getlist("additional_attributes", [])
+
         update_snuba_params_with_timestamp(request, snuba_params)
 
         error_id = request.GET.get("errorId")
@@ -402,7 +422,9 @@ def get(self, request: Request, organization: Organization, trace_id: str) -> Ht
         def data_fn(offset: int, limit: int) -> list[SerializedEvent]:
             """offset and limit don't mean anything on this endpoint currently"""
             with handle_query_errors():
-                spans = self.query_trace_data(snuba_params, trace_id, error_id)
+                spans = self.query_trace_data(
+                    snuba_params, trace_id, error_id, additional_attributes
+                )
             return spans
 
         return self.paginate(

src/sentry/snuba/spans_rpc.py

@@ -182,7 +182,11 @@ def run_trace_query(
     params: SnubaParams,
     referrer: str,
     config: SearchResolverConfig,
+    additional_attributes: list[str] | None = None,
 ) -> list[dict[str, Any]]:
+    if additional_attributes is None:
+        additional_attributes = []
+
     trace_attributes = [
         "parent_span",
         "description",
@@ -201,6 +205,7 @@ def run_trace_query(
         "sdk.name",
         "measurements.time_to_initial_display",
         "measurements.time_to_full_display",
+        *additional_attributes,
     ]
     for key in {
         "lcp",

tests/snuba/api/endpoints/test_organization_events_trace.py

@@ -56,6 +56,10 @@ def load_trace(self, is_eap=False):
                     "description": f"GET gen1-{i}",
                     "span_id": root_span_id,
                     "trace_id": self.trace_id,
+                    "data": {
+                        "gen_ai.request.model": "gpt-4o",
+                        "gen_ai.usage.total_tokens": 100,
+                    },
                 }
                 for i, root_span_id in enumerate(self.root_span_ids)
             ],

tests/snuba/api/endpoints/test_organization_trace.py

@@ -234,6 +234,29 @@ def test_with_only_errors(self):
         assert len(data) == 1
         assert data[0]["event_id"] == error.event_id
 
+    def test_with_additional_attributes(self):
+        self.load_trace(is_eap=True)
+        with self.feature(self.FEATURES):
+            response = self.client_get(
+                data={
+                    "timestamp": self.day_ago,
+                    "additional_attributes": [
+                        "gen_ai.request.model",
+                        "gen_ai.usage.total_tokens",
+                    ],
+                },
+            )
+        assert response.status_code == 200, response.content
+        data = response.data
+        assert len(data) == 1
+
+        # The root span doesn't have any of the additional attributes and returns defaults
+        assert data[0]["additional_attributes"]["gen_ai.request.model"] == ""
+        assert data[0]["additional_attributes"]["gen_ai.usage.total_tokens"] == 0
+
+        assert data[0]["children"][0]["additional_attributes"]["gen_ai.request.model"] == "gpt-4o"
+        assert data[0]["children"][0]["additional_attributes"]["gen_ai.usage.total_tokens"] == 100
+
     def test_with_target_error(self):
         start, _ = self.get_start_end_from_day_ago(1000)
         error_data = load_data(