feat(suspect-spans): Add endpoint for span stats (#30919)

This introduces a new endpoint for getting time series data for spans.

src/sentry/api/endpoints/organization_events_spans_performance.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import dataclasses
+from datetime import datetime
 from itertools import chain
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
+import sentry_sdk
 from rest_framework import serializers
 from rest_framework.exceptions import ParseError
 from rest_framework.request import Request
@@ -14,15 +16,16 @@ from snuba_sdk.function import Function
 from snuba_sdk.orderby import LimitBy
 
 from sentry import eventstore, features
-from sentry.api.bases import NoProjects, OrganizationEventsEndpointBase
+from sentry.api.bases import NoProjects, OrganizationEventsV2EndpointBase
 from sentry.api.paginator import GenericOffsetPaginator
 from sentry.api.serializers.rest_framework import ListField
 from sentry.discover.arithmetic import is_equation, strip_equation
 from sentry.models import Organization
-from sentry.search.events.builder import QueryBuilder
+from sentry.search.events.builder import QueryBuilder, TimeseriesQueryBuilder
 from sentry.search.events.types import ParamsType
+from sentry.snuba import discover
 from sentry.utils.cursors import Cursor, CursorResult
-from sentry.utils.snuba import Dataset, raw_snql_query
+from sentry.utils.snuba import Dataset, SnubaTSResult, raw_snql_query
 from sentry.utils.time_window import TimeWindow, remove_time_windows, union_time_windows
 from sentry.utils.validators import INVALID_SPAN_ID, is_span_id
 
@@ -78,7 +81,7 @@ SPAN_PERFORMANCE_COLUMNS: Dict[str, SpanPerformanceColumn] = {
 }
 
 
-class OrganizationEventsSpansEndpointBase(OrganizationEventsEndpointBase):  # type: ignore
+class OrganizationEventsSpansEndpointBase(OrganizationEventsV2EndpointBase):  # type: ignore
     def has_feature(self, request: Request, organization: Organization) -> bool:
         return bool(
             features.has(
@@ -194,7 +197,7 @@ class OrganizationEventsSpansPerformanceEndpoint(OrganizationEventsSpansEndpoint
             )
 
 
-class SpansExamplesSerializer(serializers.Serializer):  # type: ignore
+class SpansSerializer(serializers.Serializer):  # type: ignore
     query = serializers.CharField(required=False, allow_null=True)
     span = ListField(child=serializers.CharField(), required=True, allow_null=False, max_length=10)
 
@@ -205,7 +208,7 @@ class SpansExamplesSerializer(serializers.Serializer):  # type: ignore
             raise serializers.ValidationError(str(e))
 
 
-class OrganizationEventsSpansEndpoint(OrganizationEventsSpansEndpointBase):
+class OrganizationEventsSpansExamplesEndpoint(OrganizationEventsSpansEndpointBase):
     def get(self, request: Request, organization: Organization) -> Response:
         if not self.has_feature(request, organization):
             return Response(status=404)
@@ -215,7 +218,7 @@ class OrganizationEventsSpansEndpoint(OrganizationEventsSpansEndpointBase):
         except NoProjects:
             return Response(status=404)
 
-        serializer = SpansExamplesSerializer(data=request.GET)
+        serializer = SpansSerializer(data=request.GET)
         if not serializer.is_valid():
             return Response(serializer.errors, status=400)
         serialized = serializer.validated_data
@@ -279,6 +282,91 @@ class SpanExamplesPaginator:
         )
 
 
+class OrganizationEventsSpansStatsEndpoint(OrganizationEventsSpansEndpointBase):
+    def get(self, request: Request, organization: Organization) -> Response:
+        if not self.has_feature(request, organization):
+            return Response(status=404)
+
+        serializer = SpansSerializer(data=request.GET)
+        if not serializer.is_valid():
+            return Response(serializer.errors, status=400)
+        serialized = serializer.validated_data
+
+        # TODO: The serializer currently supports >1 span.
+        # This will no longer be the case once the examples query is split up.
+        spans = serialized["span"]
+        if len(spans) != 1:
+            raise ParseError(detail="Can only specify 1 span.")
+        span = spans[0]
+
+        def get_event_stats(
+            query_columns: Sequence[str],
+            query: str,
+            params: Dict[str, str],
+            rollup: int,
+            zerofill_results: bool,
+            comparison_delta: Optional[datetime] = None,
+        ) -> SnubaTSResult:
+            with sentry_sdk.start_span(
+                op="discover.discover", description="timeseries.filter_transform"
+            ):
+                builder = TimeseriesQueryBuilder(
+                    Dataset.Discover,
+                    params,
+                    rollup,
+                    query=query,
+                    selected_columns=query_columns,
+                    functions_acl=["array_join", "percentileArray", "sumArray"],
+                )
+
+                span_op_column = builder.resolve_function("array_join(spans_op)")
+                span_group_column = builder.resolve_function("array_join(spans_group)")
+
+                # Adding spans.op and spans.group to the group by because
+                # We need them in the query to help the array join optimizer
+                # in snuba take effect but the TimeseriesQueryBuilder
+                # removes all non aggregates from the select clause.
+                builder.groupby.extend([span_op_column, span_group_column])
+
+                builder.add_conditions(
+                    [
+                        Condition(
+                            Function("tuple", [span_op_column, span_group_column]),
+                            Op.IN,
+                            Function("tuple", [Function("tuple", [span.op, span.group])]),
+                        ),
+                    ]
+                )
+
+                snql_query = builder.get_snql_query()
+                results = raw_snql_query(
+                    snql_query, "api.organization-events-spans-performance-stats"
+                )
+
+            with sentry_sdk.start_span(
+                op="discover.discover", description="timeseries.transform_results"
+            ):
+                result = discover.zerofill(
+                    results["data"],
+                    params["start"],
+                    params["end"],
+                    rollup,
+                    "time",
+                )
+
+            return SnubaTSResult({"data": result}, params["start"], params["end"], rollup)
+
+        return Response(
+            self.get_event_stats_data(
+                request,
+                organization,
+                get_event_stats,
+                query_column="sumArray(spans_exclusive_time)",
+            ),
+            status=200,
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class ExampleSpan:
     id: str

src/sentry/api/urls.py

@@ -186,8 +186,9 @@ from .endpoints.organization_events_meta import (
 )
 from .endpoints.organization_events_span_ops import OrganizationEventsSpanOpsEndpoint
 from .endpoints.organization_events_spans_performance import (
-    OrganizationEventsSpansEndpoint,
+    OrganizationEventsSpansExamplesEndpoint,
     OrganizationEventsSpansPerformanceEndpoint,
+    OrganizationEventsSpansStatsEndpoint,
 )
 from .endpoints.organization_events_stats import OrganizationEventsStatsEndpoint
 from .endpoints.organization_events_trace import (
@@ -1044,7 +1045,7 @@ urlpatterns = [
                 ),
                 url(
                     r"^(?P<organization_slug>[^\/]+)/events-spans/$",
-                    OrganizationEventsSpansEndpoint.as_view(),
+                    OrganizationEventsSpansExamplesEndpoint.as_view(),
                     name="sentry-api-0-organization-events-spans",
                 ),
                 url(
@@ -1052,6 +1053,11 @@ urlpatterns = [
                     OrganizationEventsSpansPerformanceEndpoint.as_view(),
                     name="sentry-api-0-organization-events-spans-performance",
                 ),
+                url(
+                    r"^(?P<organization_slug>[^\/]+)/events-spans-stats/$",
+                    OrganizationEventsSpansStatsEndpoint.as_view(),
+                    name="sentry-api-0-organization-events-spans-stats",
+                ),
                 url(
                     r"^(?P<organization_slug>[^\/]+)/events-meta/$",
                     OrganizationEventsMetaEndpoint.as_view(),

src/sentry/search/events/builder.py

@@ -127,9 +127,6 @@ class QueryBuilder(QueryFilter):  # type: ignore
                     )
                 )
 
-    def add_conditions(self, conditions: List[Condition]) -> None:
-        self.where += conditions
-
     def get_snql_query(self) -> Query:
         self.validate_having_clause()
 
@@ -160,13 +157,14 @@ class TimeseriesQueryBuilder(QueryFilter):  # type: ignore
         query: Optional[str] = None,
         selected_columns: Optional[List[str]] = None,
         equations: Optional[List[str]] = None,
+        functions_acl: Optional[List[str]] = None,
         limit: Optional[int] = 10000,
     ):
         super().__init__(
             dataset,
             params,
             auto_fields=False,
-            functions_acl=[],
+            functions_acl=functions_acl,
             equation_config={"auto_add": True, "aggregates_only": True},
         )
         self.where, self.having = self.resolve_conditions(query, use_aggregate_conditions=False)
@@ -253,6 +251,7 @@ class TopEventsQueryBuilder(TimeseriesQueryBuilder):
         selected_columns: Optional[List[str]] = None,
         timeseries_columns: Optional[List[str]] = None,
         equations: Optional[List[str]] = None,
+        functions_acl: Optional[List[str]] = None,
         limit: Optional[int] = 10000,
     ):
         selected_columns = [] if selected_columns is None else selected_columns
@@ -266,6 +265,7 @@ class TopEventsQueryBuilder(TimeseriesQueryBuilder):
             query=query,
             selected_columns=list(set(selected_columns + timeseries_functions)),
             equations=list(set(equations + timeseries_equations)),
+            functions_acl=functions_acl,
             limit=limit,
         )
 
@@ -333,6 +333,7 @@ class TopEventsQueryBuilder(TimeseriesQueryBuilder):
                 else:
                     values.add(event.get(alias))
             values_list = list(values)
+
             if values_list:
                 if field == "timestamp" or field.startswith("timestamp.to_"):
                     if not other:

src/sentry/search/events/filter.py

@@ -1091,6 +1091,9 @@ class QueryFilter(QueryFields):
             SEMVER_BUILD_ALIAS: self._semver_build_filter_converter,
         }
 
+    def add_conditions(self, conditions: List[Condition]) -> None:
+        self.where += conditions
+
     def parse_query(self, query: Optional[str]) -> ParsedTerms:
         """Given a user's query, string construct a list of filters that can be
         then used to construct the conditions of the Query"""

src/sentry/snuba/discover.py

@@ -485,6 +485,7 @@ def timeseries_query(
     referrer: Optional[str] = None,
     zerofill_results: bool = True,
     comparison_delta: Optional[timedelta] = None,
+    functions_acl: Optional[Sequence[str]] = None,
     use_snql: Optional[bool] = False,
 ):
     """
@@ -526,6 +527,7 @@ def timeseries_query(
                 query=query,
                 selected_columns=columns,
                 equations=equations,
+                functions_acl=functions_acl,
             )
             query_list = [base_builder]
             if comparison_delta:
@@ -683,6 +685,7 @@ def top_events_timeseries(
     allow_empty=True,
     zerofill_results=True,
     include_other=False,
+    functions_acl=None,
     use_snql=False,
 ):
     """
@@ -737,6 +740,7 @@ def top_events_timeseries(
             selected_columns=selected_columns,
             timeseries_columns=timeseries_columns,
             equations=equations,
+            functions_acl=functions_acl,
         )
         if len(top_events["data"]) == limit and include_other:
             other_events_builder = TopEventsQueryBuilder(

tests/snuba/api/endpoints/test_organization_events_spans_performance.py

@@ -30,6 +30,7 @@ class OrganizationEventsSpansEndpointTestBase(APITestCase, SnubaTestCase):
         )
 
         self.min_ago = before_now(minutes=1).replace(microsecond=0)
+        self.day_ago = before_now(days=1).replace(hour=10, minute=0, second=0, microsecond=0)
 
     def update_snuba_config_ensure(self, config, poll=60, wait=1):
         self.snuba_update_config(config)
@@ -1156,7 +1157,7 @@ class OrganizationEventsSpansPerformanceEndpointTest(OrganizationEventsSpansEndp
         self.assert_suspect_span(response.data, [results])
 
 
-class OrganizationEventsSpansEndpointTest(OrganizationEventsSpansEndpointTestBase):
+class OrganizationEventsSpansExamplesEndpointTest(OrganizationEventsSpansEndpointTestBase):
     URL = "sentry-api-0-organization-events-spans"
 
     def test_no_feature(self):
@@ -1325,3 +1326,141 @@ class OrganizationEventsSpansEndpointTest(OrganizationEventsSpansEndpointTestBas
             response.data,
             [self.span_example_results("http.server", event)],
         )
+
+
+class OrganizationEventsSpansStatsEndpointTest(OrganizationEventsSpansEndpointTestBase):
+    URL = "sentry-api-0-organization-events-spans-stats"
+
+    def test_no_feature(self):
+        response = self.client.get(self.url, format="json")
+        assert response.status_code == 404, response.content
+
+    def test_require_span_param(self):
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": self.project.id},
+                format="json",
+            )
+
+        assert response.status_code == 400, response.content
+        assert response.data == {"span": [ErrorDetail("This field is required.", code="required")]}
+
+    def test_bad_span_param(self):
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": self.project.id, "span": ["http.server"]},
+                format="json",
+            )
+
+        assert response.status_code == 400, response.content
+        assert response.data == {
+            "span": [
+                ErrorDetail(
+                    "span must consist of of a span op and a valid 16 character hex delimited by a colon (:)",
+                    code="invalid",
+                )
+            ]
+        }
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": self.project.id, "span": ["http.server:ab"]},
+                format="json",
+            )
+
+        assert response.status_code == 400, response.content
+        assert response.data == {
+            "span": [
+                ErrorDetail(
+                    "spanGroup must be a valid 16 character hex (containing only digits, or a-f characters)",
+                    code="invalid",
+                )
+            ]
+        }
+
+    def test_multiple_span_param(self):
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={
+                    "project": self.project.id,
+                    "span": [f"http.server:{'ab' * 8}", f"django.middleware:{'cd' * 8}"],
+                },
+                format="json",
+            )
+
+        assert response.status_code == 400, response.content
+        assert response.data == {"detail": "Can only specify 1 span."}
+
+    @patch("sentry.api.endpoints.organization_events_spans_performance.raw_snql_query")
+    def test_one_span(self, mock_raw_snql_query):
+        mock_raw_snql_query.side_effect = [{"data": []}]
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={
+                    "project": self.project.id,
+                    "span": f"http.server:{'ab' * 8}",
+                    "yAxis": [
+                        "percentileArray(spans_exclusive_time, 0.75)",
+                        "percentileArray(spans_exclusive_time, 0.95)",
+                        "percentileArray(spans_exclusive_time, 0.99)",
+                    ],
+                    "start": iso_format(self.day_ago),
+                    "end": iso_format(self.day_ago + timedelta(hours=2)),
+                    "interval": "1h",
+                },
+                format="json",
+            )
+
+        assert response.status_code == 200, response.content
+
+        # ensure that the result is a proper time series
+        data = response.data
+        series_names = [
+            f"percentileArray(spans_exclusive_time, 0.{percentile})"
+            for percentile in ["75", "95", "99"]
+        ]
+        assert set(data.keys()) == set(series_names)
+        for i, series_name in enumerate(series_names):
+            series = data[series_name]
+            assert series["order"] == i
+            assert [attrs for _, attrs in series["data"]] == [
+                [{"count": 0}],
+                [{"count": 0}],
+            ]
+
+        assert mock_raw_snql_query.call_count == 1
+        query = mock_raw_snql_query.call_args_list[0][0][0]
+
+        # ensure the specified y axes are in the select
+        for percentile in ["75", "95", "99"]:
+            assert (
+                Function(
+                    f"quantile(0.{percentile.rstrip('0')})",
+                    [Function("arrayJoin", [Column("spans.exclusive_time")])],
+                    f"percentileArray_spans_exclusive_time_0_{percentile}",
+                )
+                in query.select
+            )
+
+        spans_op = Function("arrayJoin", [Column("spans.op")], "array_join_spans_op")
+        spans_group = Function("arrayJoin", [Column("spans.group")], "array_join_spans_group")
+
+        # ensure the two span columns are in the group by
+        for column in [spans_op, spans_group]:
+            assert column in query.groupby
+
+        # ensure there is a condition on the span
+        assert (
+            Condition(
+                Function("tuple", [spans_op, spans_group]),
+                Op.IN,
+                Function("tuple", [Function("tuple", ["http.server", "ab" * 8])]),
+            )
+            in query.where
+        )