feat(llm-detection): Accept additional_attributes param in get_trace_waterfall

src/sentry/seer/explorer/tools.py

@@ -371,7 +371,9 @@ def execute_trace_table_query(
         raise
 
 
-def get_trace_waterfall(trace_id: str, organization_id: int) -> EAPTrace | None:
+def get_trace_waterfall(
+    trace_id: str, organization_id: int, additional_attributes: list[str] | None = None
+) -> EAPTrace | None:
     """
     Get the full span waterfall and connected errors for a trace.
 
@@ -382,6 +384,8 @@ def get_trace_waterfall(trace_id: str, organization_id: int) -> EAPTrace | None:
     Returns:
         The spans and errors in the trace, along with the full 32-character trace ID.
     """
+    if additional_attributes is None:
+        additional_attributes = ["span.status_code"]
 
     try:
         organization = Organization.objects.get(id=organization_id)
@@ -416,7 +420,7 @@ def get_trace_waterfall(trace_id: str, organization_id: int) -> EAPTrace | None:
     events = query_trace_data(
         snuba_params,
         full_trace_id,
-        additional_attributes=["span.status_code"],
+        additional_attributes=additional_attributes,
         referrer=Referrer.SEER_EXPLORER_TOOLS,
         organization=organization,
     )
@@ -428,8 +432,10 @@ def get_trace_waterfall(trace_id: str, organization_id: int) -> EAPTrace | None:
     )
 
 
-def rpc_get_trace_waterfall(trace_id: str, organization_id: int) -> dict[str, Any]:
-    trace = get_trace_waterfall(trace_id, organization_id)
+def rpc_get_trace_waterfall(
+    trace_id: str, organization_id: int, additional_attributes: list[str] | None = None
+) -> dict[str, Any]:
+    trace = get_trace_waterfall(trace_id, organization_id, additional_attributes)
     return trace.dict() if trace else {}
 
 

tests/sentry/seer/explorer/test_tools.py

@@ -760,11 +760,10 @@ def test_get_trace_waterfall_sliding_window_beyond_limit(self) -> None:
         assert result is None
 
     def test_get_trace_waterfall_includes_status_code(self) -> None:
-        """Test that span.status_code is included in additional_attributes."""
+        """Test that span.status_code is included if additional_attributes is not provided"""
         transaction_name = "api/test/status"
         trace_id = uuid.uuid4().hex
 
-        # Create a span with status_code
         span = self.create_span(
             {
                 "description": "http-request",
@@ -782,11 +781,38 @@ def test_get_trace_waterfall_includes_status_code(self) -> None:
         result = get_trace_waterfall(trace_id, self.organization.id)
         assert isinstance(result, EAPTrace)
 
-        # Find the span and verify additional_attributes contains status_code
         root_span = result.trace[0]
         assert "additional_attributes" in root_span
         assert root_span["additional_attributes"].get("span.status_code") == "500"
 
+    def test_get_trace_waterfall_includes_additional_attributes(self) -> None:
+        """Test that additional_attributes passed into the function are included on returned traces"""
+        transaction_name = "api/test/status"
+        trace_id = uuid.uuid4().hex
+
+        span = self.create_span(
+            {
+                "description": "http-request",
+                "sentry_tags": {
+                    "transaction": transaction_name,
+                    "status_code": "500",
+                    "request.url": "best-url-ev3r.biz",
+                },
+                "trace_id": trace_id,
+                "is_segment": True,
+            },
+            start_ts=self.ten_mins_ago,
+        )
+        self.store_spans([span])
+
+        result = get_trace_waterfall(trace_id, self.organization.id, ["request.url"])
+        assert isinstance(result, EAPTrace)
+
+        root_span = result.trace[0]
+        assert "additional_attributes" in root_span
+        assert root_span["additional_attributes"].get("span.status_code") is None
+        assert root_span["additional_attributes"].get("request.url") == "best-url-ev3r.biz"
+
 
 class TestTraceTableQuery(APITransactionTestCase, SnubaTestCase, SpanTestCase):
     def setUp(self) -> None: