feat: Add trace item stats endpoint (#103217)

Adds an endpoint that is essentially a wrapper around the TraceItemStats
endpoint.
This endpoint is meant to give you information about the shape of your
data.
This will be used to return distributions. As we add more stats types,
we will expose
them in this endpoint.

---------

Co-authored-by: getsantry[bot] <66042841+getsantry[bot]@users.noreply.github.com>

Author: shruthilayaj

src/sentry/api/endpoints/organization_trace_item_stats.py

@@ -0,0 +1,61 @@
+import logging
+
+from rest_framework import serializers
+from rest_framework.request import Request
+from rest_framework.response import Response
+
+from sentry.api.api_owners import ApiOwner
+from sentry.api.api_publish_status import ApiPublishStatus
+from sentry.api.base import region_silo_endpoint
+from sentry.api.bases import NoProjects, OrganizationEventsV2EndpointBase
+from sentry.models.organization import Organization
+from sentry.search.eap.constants import SUPPORTED_STATS_TYPES
+from sentry.search.eap.resolver import SearchResolver
+from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
+from sentry.search.eap.types import SearchResolverConfig
+from sentry.snuba.referrer import Referrer
+from sentry.snuba.spans_rpc import Spans
+
+logger = logging.getLogger(__name__)
+
+
+class OrganizationTraceItemsStatsSerializer(serializers.Serializer):
+    query = serializers.CharField(required=False)
+    statsType = serializers.ListField(
+        child=serializers.ChoiceField(list(SUPPORTED_STATS_TYPES)), required=True
+    )
+
+
+@region_silo_endpoint
+class OrganizationTraceItemsStatsEndpoint(OrganizationEventsV2EndpointBase):
+    publish_status = {
+        "GET": ApiPublishStatus.PRIVATE,
+    }
+    owner = ApiOwner.VISIBILITY
+
+    def get(self, request: Request, organization: Organization) -> Response:
+        try:
+            snuba_params = self.get_snuba_params(request, organization)
+        except NoProjects:
+            return Response({"data": []})
+
+        serializer = OrganizationTraceItemsStatsSerializer(data=request.GET)
+        if not serializer.is_valid():
+            return Response(serializer.errors, status=400)
+        serialized = serializer.validated_data
+
+        resolver_config = SearchResolverConfig()
+        resolver = SearchResolver(
+            params=snuba_params, config=resolver_config, definitions=SPAN_DEFINITIONS
+        )
+
+        stats_results = Spans.run_stats_query(
+            params=snuba_params,
+            stats_types=serialized.get("statsType"),
+            query_string=serialized.get("query", ""),
+            referrer=Referrer.API_SPANS_FREQUENCY_STATS_RPC.value,
+            config=resolver_config,
+            search_resolver=resolver,
+        )
+
+        return Response({"data": stats_results})

src/sentry/api/urls.py

@@ -40,6 +40,7 @@
 from sentry.api.endpoints.organization_trace_item_attributes_ranked import (
     OrganizationTraceItemsAttributesRankedEndpoint,
 )
+from sentry.api.endpoints.organization_trace_item_stats import OrganizationTraceItemsStatsEndpoint
 from sentry.api.endpoints.organization_unsubscribe import (
     OrganizationUnsubscribeIssue,
     OrganizationUnsubscribeProject,
@@ -1724,6 +1725,11 @@ def create_group_urls(name_prefix: str) -> list[URLPattern | URLResolver]:
         OrganizationTraceItemsAttributesRankedEndpoint.as_view(),
         name="sentry-api-0-organization-trace-item-attributes-ranked",
     ),
+    re_path(
+        r"^(?P<organization_id_or_slug>[^/]+)/trace-items/stats/$",
+        OrganizationTraceItemsStatsEndpoint.as_view(),
+        name="sentry-api-0-organization-trace-item-stats",
+    ),
     re_path(
         r"^(?P<organization_id_or_slug>[^/]+)/spans/fields/$",
         OrganizationSpansFieldsEndpoint.as_view(),

src/sentry/search/eap/constants.py

@@ -19,6 +19,8 @@
     SupportedTraceItemType.PROFILE_FUNCTIONS: TraceItemType.TRACE_ITEM_TYPE_PROFILE_FUNCTION,
 }
 
+SUPPORTED_STATS_TYPES = {"attributeDistributions"}
+
 OPERATOR_MAP = {
     "=": ComparisonFilter.OP_EQUALS,
     "!=": ComparisonFilter.OP_NOT_EQUALS,

src/sentry/snuba/rpc_dataset_common.py

@@ -869,6 +869,19 @@ def run_trace_query(
     ) -> list[dict[str, Any]]:
         raise NotImplementedError()
 
+    @classmethod
+    def run_stats_query(
+        cls,
+        *,
+        params: SnubaParams,
+        stats_types: set[str],
+        query_string: str,
+        referrer: str,
+        config: SearchResolverConfig,
+        search_resolver: SearchResolver | None = None,
+    ) -> list[dict[str, Any]]:
+        raise NotImplementedError()
+
 
 def can_force_highest_accuracy(meta: RequestMeta) -> bool:
     # when using MODE_HIGHEST_ACCURACY_FLEXTIME, we cannot force highest accuracy

src/sentry/snuba/spans_rpc.py

@@ -1,17 +1,29 @@
 import logging
+from collections import defaultdict
 from datetime import timedelta
 from typing import Any
 
 import sentry_sdk
 from sentry_protos.snuba.v1.endpoint_get_trace_pb2 import GetTraceRequest
+from sentry_protos.snuba.v1.endpoint_trace_item_stats_pb2 import (
+    AttributeDistributionsRequest,
+    StatsType,
+    TraceItemStatsRequest,
+)
 from sentry_protos.snuba.v1.request_common_pb2 import PageToken, TraceItemType
 
 from sentry.exceptions import InvalidSearchQuery
-from sentry.search.eap.constants import DOUBLE, INT, STRING
+from sentry.search.eap.constants import DOUBLE, INT, STRING, SUPPORTED_STATS_TYPES
 from sentry.search.eap.resolver import SearchResolver
 from sentry.search.eap.sampling import events_meta_from_rpc_request_meta
 from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
-from sentry.search.eap.types import AdditionalQueries, EAPResponse, SearchResolverConfig
+from sentry.search.eap.types import (
+    AdditionalQueries,
+    EAPResponse,
+    SearchResolverConfig,
+    SupportedTraceItemType,
+)
+from sentry.search.eap.utils import can_expose_attribute
 from sentry.search.events.types import SAMPLING_MODES, EventsMeta, SnubaParams
 from sentry.snuba import rpc_dataset_common
 from sentry.snuba.discover import zerofill
@@ -237,3 +249,61 @@ def run_trace_query(
                             )
                 spans.append(span)
         return spans
+
+    @classmethod
+    @sentry_sdk.trace
+    def run_stats_query(
+        cls,
+        *,
+        params: SnubaParams,
+        stats_types: set[str],
+        query_string: str,
+        referrer: str,
+        config: SearchResolverConfig,
+        search_resolver: SearchResolver | None = None,
+    ) -> list[dict[str, Any]]:
+        search_resolver = search_resolver or cls.get_resolver(params, config)
+        stats_filter, _, _ = search_resolver.resolve_query(query_string)
+        meta = search_resolver.resolve_meta(
+            referrer=referrer,
+            sampling_mode=params.sampling_mode,
+        )
+        stats_request = TraceItemStatsRequest(
+            filter=stats_filter,
+            meta=meta,
+            stats_types=[],
+        )
+
+        if not set(stats_types).intersection(SUPPORTED_STATS_TYPES):
+            return []
+
+        if "attributeDistributions" in stats_types:
+            stats_request.stats_types.append(
+                StatsType(
+                    attribute_distributions=AttributeDistributionsRequest(
+                        max_buckets=75,
+                    )
+                )
+            )
+
+        response = snuba_rpc.trace_item_stats_rpc(stats_request)
+        stats = []
+
+        for result in response.results:
+            if "attributeDistributions" in stats_types and result.HasField(
+                "attribute_distributions"
+            ):
+                attributes = defaultdict(list)
+                for attribute in result.attribute_distributions.attributes:
+                    if not can_expose_attribute(
+                        attribute.attribute_name, SupportedTraceItemType.SPANS
+                    ):
+                        continue
+
+                    for bucket in attribute.buckets:
+                        attributes[attribute.attribute_name].append(
+                            {"label": bucket.label, "value": bucket.value}
+                        )
+                stats.append({"attribute_distributions": {"data": attributes}})
+
+        return stats

static/app/utils/api/knownSentryApiUrls.generated.ts

@@ -553,6 +553,7 @@ export type KnownSentryApiUrls =
   | '/organizations/$organizationIdOrSlug/trace-items/attributes/'
   | '/organizations/$organizationIdOrSlug/trace-items/attributes/$key/values/'
   | '/organizations/$organizationIdOrSlug/trace-items/attributes/ranked/'
+  | '/organizations/$organizationIdOrSlug/trace-items/stats/'
   | '/organizations/$organizationIdOrSlug/trace-logs/'
   | '/organizations/$organizationIdOrSlug/trace-meta/$traceId/'
   | '/organizations/$organizationIdOrSlug/trace-summary/'

tests/snuba/api/endpoints/test_organization_trace_item_stats.py

@@ -0,0 +1,78 @@
+from django.urls import reverse
+
+from sentry.testutils.cases import APITransactionTestCase, SnubaTestCase, SpanTestCase
+from sentry.testutils.helpers.datetime import before_now
+
+
+class OrganizationTraceItemsStatsEndpointTest(
+    APITransactionTestCase,
+    SnubaTestCase,
+    SpanTestCase,
+):
+    view = "sentry-api-0-organization-trace-item-stats"
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.login_as(user=self.user)
+        self.ten_mins_ago = before_now(minutes=10)
+        self.ten_mins_ago_iso = self.ten_mins_ago.replace(microsecond=0).isoformat()
+
+    def do_request(self, query=None, features=None, **kwargs):
+        if query:
+            query.setdefault("sampling", "HIGHEST_ACCURACY")
+
+        response = self.client.get(
+            reverse(
+                self.view,
+                kwargs={"organization_id_or_slug": self.organization.slug},
+            ),
+            query,
+            format="json",
+            **kwargs,
+        )
+
+        return response
+
+    def _store_span(self, description=None, tags=None, duration=None):
+        if tags is None:
+            tags = {"foo": "bar"}
+
+        self.store_span(
+            self.create_span(
+                {"description": description or "foo", "sentry_tags": tags},
+                start_ts=self.ten_mins_ago,
+                duration=duration or 1000,
+            ),
+            is_eap=True,
+        )
+
+    def test_no_project(self) -> None:
+        response = self.do_request()
+        assert response.status_code == 200, response.data
+        assert response.data == {"data": []}
+
+    def test_distribution_values(self) -> None:
+        tags = [
+            ({"browser": "chrome", "device": "desktop"}, 500),
+            ({"browser": "chrome", "device": "mobile"}, 100),
+            ({"browser": "chrome", "device": "mobile"}, 100),
+            ({"browser": "chrome", "device": "desktop"}, 100),
+            ({"browser": "safari", "device": "mobile"}, 100),
+            ({"browser": "chrome", "device": "desktop"}, 500),
+            ({"browser": "edge", "device": "desktop"}, 500),
+        ]
+
+        for tag, duration in tags:
+            self._store_span(tags=tag, duration=duration)
+
+        response = self.do_request(
+            query={"query": "span.duration:<=100", "statsType": ["attributeDistributions"]}
+        )
+        assert response.status_code == 200, response.data
+        assert len(response.data["data"]) == 1
+        attribute_distribution = response.data["data"][0]["attribute_distributions"]["data"]
+        device_data = attribute_distribution["sentry.device"]
+        assert {"label": "mobile", "value": 3.0} in device_data
+        assert {"label": "desktop", "value": 1.0} in device_data
+
+        assert response.data