diff --git a/src/sentry/seer/explorer/tools.py b/src/sentry/seer/explorer/tools.py index 7764c7cb40b1c3..c5d5dc1df390df 100644 --- a/src/sentry/seer/explorer/tools.py +++ b/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 {} diff --git a/tests/sentry/seer/explorer/test_tools.py b/tests/sentry/seer/explorer/test_tools.py index 0b87b5ae009c43..0639121baff067 100644 --- a/tests/sentry/seer/explorer/test_tools.py +++ b/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: