feat(eap): Use RPC to fetch tag keys (#77367)

This uses RPC calls to fetch tag keys from eap spans.

Depends on getsentry/snuba#6301

src/sentry/api/endpoints/organization_spans_fields.py

@@ -1,6 +1,17 @@
+from datetime import timedelta
+
 import sentry_sdk
+from google.protobuf.timestamp_pb2 import Timestamp
+from rest_framework import serializers
+from rest_framework.exceptions import ParseError
 from rest_framework.request import Request
 from rest_framework.response import Response
+from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
+    TraceItemAttributesRequest,
+    TraceItemAttributesResponse,
+)
+from sentry_protos.snuba.v1alpha.request_common_pb2 import RequestMeta, TraceItemName
+from sentry_protos.snuba.v1alpha.trace_item_attribute_pb2 import AttributeKey
 from sentry_relay.consts import SPAN_STATUS_CODE_TO_NAME
 from snuba_sdk import Condition, Op
 
@@ -18,6 +29,7 @@
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.referrer import Referrer
 from sentry.tagstore.types import TagKey, TagValue
+from sentry.utils import snuba
 
 
 class OrganizationSpansFieldsEndpointBase(OrganizationEventsV2EndpointBase):
@@ -27,6 +39,18 @@ class OrganizationSpansFieldsEndpointBase(OrganizationEventsV2EndpointBase):
     owner = ApiOwner.PERFORMANCE
 
 
+class OrganizationSpansFieldsEndpointSerializer(serializers.Serializer):
+    dataset = serializers.ChoiceField(
+        ["spans", "spansIndexed"], required=False, default="spansIndexed"
+    )
+    type = serializers.ChoiceField(["string", "number"], required=False)
+
+    def validate(self, attrs):
+        if attrs["dataset"] == "spans" and attrs.get("type") is None:
+            raise ParseError(detail='type is required when using dataset="spans"')
+        return attrs
+
+
 @region_silo_endpoint
 class OrganizationSpansFieldsEndpoint(OrganizationSpansFieldsEndpointBase):
     snuba_methods = ["GET"]
@@ -45,8 +69,58 @@ def get(self, request: Request, organization) -> Response:
                 paginator=ChainPaginator([]),
             )
 
+        serializer = OrganizationSpansFieldsEndpointSerializer(data=request.GET)
+        if not serializer.is_valid():
+            return Response(serializer.errors, status=400)
+        serialized = serializer.validated_data
+
         max_span_tags = options.get("performance.spans-tags-key.max")
 
+        if serialized["dataset"] == "spans" and features.has(
+            "organizations:visibility-explore-dataset", organization, actor=request.user
+        ):
+            start_timestamp = Timestamp()
+            start_timestamp.FromDatetime(
+                snuba_params.start_date.replace(hour=0, minute=0, second=0, microsecond=0)
+            )
+
+            end_timestamp = Timestamp()
+            end_timestamp.FromDatetime(
+                snuba_params.end_date.replace(hour=0, minute=0, second=0, microsecond=0)
+                + timedelta(days=1)
+            )
+
+            rpc_request = TraceItemAttributesRequest(
+                meta=RequestMeta(
+                    organization_id=organization.id,
+                    cogs_category="performance",
+                    referrer=Referrer.API_SPANS_TAG_KEYS_RPC.value,
+                    project_ids=snuba_params.project_ids,
+                    start_timestamp=start_timestamp,
+                    end_timestamp=end_timestamp,
+                    trace_item_name=TraceItemName.TRACE_ITEM_NAME_EAP_SPANS,
+                ),
+                limit=max_span_tags,
+                offset=0,
+                type=AttributeKey.Type.TYPE_STRING,
+            )
+            rpc_response = snuba.rpc(rpc_request, TraceItemAttributesResponse)
+
+            paginator = ChainPaginator(
+                [
+                    [TagKey(tag.name) for tag in rpc_response.tags if tag.name],
+                ],
+                max_limit=max_span_tags,
+            )
+
+            return self.paginate(
+                request=request,
+                paginator=paginator,
+                on_results=lambda results: serialize(results, request.user),
+                default_per_page=max_span_tags,
+                max_per_page=max_span_tags,
+            )
+
         with handle_query_errors():
             # This has the limitations that we cannot paginate and
             # we do not provide any guarantees around which tag keys

src/sentry/snuba/referrer.py

@@ -472,6 +472,7 @@ class Referrer(Enum):
     API_TRACE_EXPLORER_TRACES_OCCURRENCES = "api.trace-explorer.traces-occurrences"
     API_TRACE_EXPLORER_TRACE_SPANS_LIST = "api.trace-explorer.trace-spans-list"
     API_SPANS_TAG_KEYS = "api.spans.tags-keys"
+    API_SPANS_TAG_KEYS_RPC = "api.spans.tags-keys.rpc"
     API_SPANS_TRACE_VIEW = "api.spans.trace-view"
 
     # Performance Mobile UI Module

src/sentry/testutils/cases.py

@@ -1525,6 +1525,8 @@ def store_segment(
         sdk_name: str | None = None,
         op: str | None = None,
         status: str | None = None,
+        organization_id: int = 1,
+        is_eap: bool = False,
     ):
         if span_id is None:
             span_id = self._random_span_id()
@@ -1533,7 +1535,7 @@ def store_segment(
 
         payload = {
             "project_id": project_id,
-            "organization_id": 1,
+            "organization_id": organization_id,
             "span_id": span_id,
             "trace_id": trace_id,
             "duration_ms": int(duration),
@@ -1568,7 +1570,7 @@ def store_segment(
         if status is not None:
             payload["sentry_tags"]["status"] = status
 
-        self.store_span(payload)
+        self.store_span(payload, is_eap=is_eap)
 
         if "_metrics_summary" in payload:
             self.store_metrics_summary(payload)

tests/sentry/api/endpoints/test_organization_spans_fields.py

@@ -7,6 +7,7 @@
 
 
 class OrganizationSpansTagsEndpointTest(BaseSpansTestCase, APITestCase):
+    is_eap = False
     view = "sentry-api-0-organization-spans-fields"
 
     def setUp(self):
@@ -33,20 +34,23 @@ def test_no_project(self):
         assert response.status_code == 200, response.data
         assert response.data == []
 
-    def test_tags(self):
+    def test_tags_list(self):
         for tag in ["foo", "bar", "baz"]:
             self.store_segment(
                 self.project.id,
                 uuid4().hex,
                 uuid4().hex,
                 span_id=uuid4().hex[:15],
+                organization_id=self.organization.id,
                 parent_span_id=None,
                 timestamp=before_now(days=0, minutes=10).replace(microsecond=0),
                 transaction="foo",
                 duration=100,
                 exclusive_time=100,
                 tags={tag: tag},
+                is_eap=self.is_eap,
             )
+
         for features in [
             None,  # use the default features
             ["organizations:performance-trace-explorer"],
@@ -60,6 +64,29 @@ def test_tags(self):
             ]
 
 
+class OrganizationEAPSpansTagsEndpointTest(OrganizationSpansTagsEndpointTest):
+    is_eap = True
+
+    def do_request(self, query=None, features=None, **kwargs):
+        if features is None:
+            features = ["organizations:performance-trace-explorer"]
+
+        features.append("organizations:visibility-explore-dataset")
+
+        if query is None:
+            query = {}
+        query["dataset"] = "spans"
+        query["type"] = "string"
+
+        with self.feature(features):
+            return self.client.get(
+                reverse(self.view, kwargs={"organization_id_or_slug": self.organization.slug}),
+                query,
+                format="json",
+                **kwargs,
+            )
+
+
 class OrganizationSpansTagKeyValuesEndpointTest(BaseSpansTestCase, APITestCase):
     view = "sentry-api-0-organization-spans-fields-values"