feat(tracing): Allow multiple average columns (#67899)

- This updates the averageColumn param to be a list so multiple columns
can be passed, so that the frontend can load both span.self_time and
span.duration

src/sentry/api/endpoints/organization_event_details.py

@@ -16,10 +16,14 @@ from sentry.api.utils import handle_query_errors
 from sentry.constants import ObjectStatus
 from sentry.models.project import Project
 from sentry.search.events.builder import SpansMetricsQueryBuilder
+from sentry.search.events.types import QueryBuilderConfig
 from sentry.snuba.dataset import Dataset
+from sentry.snuba.referrer import Referrer
 
+VALID_AVERAGE_COLUMNS = {"span.self_time", "span.duration"}
 
-def add_comparison_to_event(event, average_column):
+
+def add_comparison_to_event(event, average_columns):
     if "spans" not in event.data:
         return
     group_to_span_map = defaultdict(list)
@@ -46,8 +50,9 @@ def add_comparison_to_event(event, average_column):
             },
             selected_columns=[
                 "span.group",
-                f"avg({average_column}) as avg",
+                *[f"avg({average_column})" for average_column in average_columns],
             ],
+            config=QueryBuilderConfig(transform_alias_to_input_format=True),
             # orderby shouldn't matter, just picking something so results are consistent
             orderby=["span.group"],
         )
@@ -60,13 +65,19 @@ def add_comparison_to_event(event, average_column):
                 )
             ]
         )
-        result = builder.run_query("Get avg for spans")
+        result = builder.process_results(
+            builder.run_query(Referrer.API_PERFORMANCE_ORG_EVENT_AVERAGE_SPAN.value)
+        )
         sentry_sdk.set_measurement("query.groups_found", len(result["data"]))
-        for result in result["data"]:
-            group = result["span.group"]
-            avg = result["avg"]
+        for row in result["data"]:
+            group = row["span.group"]
             for span in group_to_span_map[group]:
-                span["span.average_time"] = avg
+                average_results = {}
+                for col in row:
+                    if col.startswith("avg") and row[col] > 0:
+                        average_results[col] = row[col]
+                if average_results:
+                    span["span.averageResults"] = average_results
 
 
 @region_silo_endpoint
@@ -100,9 +111,12 @@ class OrganizationEventDetailsEndpoint(OrganizationEventsEndpointBase):
         if event is None:
             return Response({"detail": "Event not found"}, status=404)
 
-        average_column = request.GET.get("averageColumn")
-        if average_column in ["span.self_time"]:
-            add_comparison_to_event(event, average_column)
+        average_columns = request.GET.getlist("averageColumn", [])
+        if (
+            all(col in VALID_AVERAGE_COLUMNS for col in average_columns)
+            and len(average_columns) > 0
+        ):
+            add_comparison_to_event(event, average_columns)
 
         # TODO: Remove `for_group` check once performance issues are moved to the issue platform
         if hasattr(event, "for_group") and event.group:

src/sentry/snuba/referrer.py

@@ -349,6 +349,7 @@ class Referrer(Enum):
     )
     API_PERFORMANCE_VITAL_DETAIL = "api.performance.vital-detail"
     API_PERFORMANCE_VITALS_CARDS = "api.performance.vitals-cards"
+    API_PERFORMANCE_ORG_EVENT_AVERAGE_SPAN = "api.performance.org-event-average-span"
     API_PROFILING_LANDING_CHART = "api.profiling.landing-chart"
     API_PROFILING_LANDING_TABLE = "api.profiling.landing-table"
     API_PROFILING_LANDING_FUNCTIONS_CARD = "api.profiling.landing-functions-card"

tests/snuba/api/endpoints/test_organization_event_details.py

@@ -301,6 +301,7 @@ class EventComparisonTest(MetricsEnhancedPerformanceTestCase):
         self.init_snuba()
         self.ten_mins_ago = before_now(minutes=10)
         self.transaction_data = load_data("transaction", timestamp=self.ten_mins_ago)
+        self.RESULT_COLUMN = "span.averageResults"
         event = self.store_event(self.transaction_data, self.project)
         self.url = reverse(
             self.endpoint,
@@ -325,7 +326,53 @@ class EventComparisonTest(MetricsEnhancedPerformanceTestCase):
         for entry in entries:
             if entry["type"] == "spans":
                 for span in entry["data"]:
-                    if span["span_id"] == "26b881987e4bad99":
-                        assert span["span.average_time"] == 1.0
-                    if span["span_id"] == "c048b4fffdc4279d":
-                        assert "span.average_time" not in span
+                    if span["op"] == "db":
+                        assert span[self.RESULT_COLUMN] == {"avg(span.self_time)": 1.0}
+                    if span["op"] == "django.middleware":
+                        assert self.RESULT_COLUMN not in span
+
+    def test_get_multiple_columns(self):
+        self.store_span_metric(
+            2,
+            internal_metric=constants.SPAN_METRICS_MAP["span.duration"],
+            timestamp=self.ten_mins_ago,
+            tags={"span.group": "26b881987e4bad99"},
+        )
+        response = self.client.get(self.url, {"averageColumn": ["span.self_time", "span.duration"]})
+        assert response.status_code == 200, response.content
+        entries = response.data["entries"]  # type: ignore[attr-defined]
+        for entry in entries:
+            if entry["type"] == "spans":
+                for span in entry["data"]:
+                    if span["op"] == "db":
+                        assert span[self.RESULT_COLUMN] == {
+                            "avg(span.self_time)": 1.0,
+                            "avg(span.duration)": 2.0,
+                        }
+                    if span["op"] == "django.middlewares":
+                        assert self.RESULT_COLUMN not in span
+
+    def test_nan_column(self):
+        # If there's nothing stored for a metric, span.duration in this case the query returns nan
+        response = self.client.get(self.url, {"averageColumn": ["span.self_time", "span.duration"]})
+        assert response.status_code == 200, response.content
+        entries = response.data["entries"]  # type: ignore[attr-defined]
+        for entry in entries:
+            if entry["type"] == "spans":
+                for span in entry["data"]:
+                    if span["op"] == "db":
+                        assert span[self.RESULT_COLUMN] == {"avg(span.self_time)": 1.0}
+                    if span["op"] == "django.middlewares":
+                        assert self.RESULT_COLUMN not in span
+
+    def test_invalid_column(self):
+        # If any columns are invalid, ignore average field in results completely
+        response = self.client.get(
+            self.url, {"averageColumn": ["span.self_time", "span.everything"]}
+        )
+        assert response.status_code == 200, response.content
+        entries = response.data["entries"]  # type: ignore[attr-defined]
+        for entry in entries:
+            if entry["type"] == "spans":
+                for span in entry["data"]:
+                    assert self.RESULT_COLUMN not in span