feat(ourlogs): Add trace items attributes endpoints (#87343)

k-fish · web-flow · commit caa841c7fd8b · 2025-03-19T09:27:19.000-04:00
### Summary
Since logs and spans (etc.) are moving towards the generic trace items,
this adds an endpoint for use with search autocomplete on the frontend,
which we can eventually use for any search relating to eap items.

This mostly ports the EAP side of the `/spans/fields/` endpoints
(organization_spans_fields.py) and makes both endpoints accept an
item_type in lieu of dataset. Returning all attr keys and fetching attr
values are not yet supported for EAP trace items (see skipped tests).
diff --git a/requirements-base.txt b/requirements-base.txt
@@ -68,7 +68,7 @@ rfc3986-validator>=0.1.1
 sentry-arroyo>=2.20.0
 sentry-kafka-schemas>=1.1.2
 sentry-ophio==1.0.0
-sentry-protos>=0.1.62
+sentry-protos>=0.1.63
 sentry-redis-tools>=0.1.7
 sentry-relay>=0.9.6
 sentry-sdk[http2]>=2.23.1
diff --git a/requirements-dev-frozen.txt b/requirements-dev-frozen.txt
@@ -189,7 +189,7 @@ sentry-forked-djangorestframework-stubs==3.15.3.post1
 sentry-forked-email-reply-parser==0.5.12.post1
 sentry-kafka-schemas==1.1.2
 sentry-ophio==1.0.0
-sentry-protos==0.1.62
+sentry-protos==0.1.63
 sentry-redis-tools==0.1.7
 sentry-relay==0.9.6
 sentry-sdk==2.23.1
diff --git a/requirements-frozen.txt b/requirements-frozen.txt
@@ -127,7 +127,7 @@ sentry-arroyo==2.20.0
 sentry-forked-email-reply-parser==0.5.12.post1
 sentry-kafka-schemas==1.1.2
 sentry-ophio==1.0.0
-sentry-protos==0.1.62
+sentry-protos==0.1.63
 sentry-redis-tools==0.1.7
 sentry-relay==0.9.6
 sentry-sdk==2.23.1
diff --git a/src/sentry/api/endpoints/organization_trace_item_attributes.py b/src/sentry/api/endpoints/organization_trace_item_attributes.py
@@ -0,0 +1,341 @@
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Literal
+
+import sentry_sdk
+from google.protobuf.timestamp_pb2 import Timestamp
+from rest_framework import serializers
+from rest_framework.request import Request
+from rest_framework.response import Response
+from sentry_protos.snuba.v1.endpoint_trace_item_attributes_pb2 import (
+    TraceItemAttributeNamesRequest,
+    TraceItemAttributeValuesRequest,
+)
+from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType as ProtoTraceItemType
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import ExistsFilter, TraceItemFilter
+
+from sentry import features, options
+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.api.endpoints.organization_spans_fields import BaseSpanFieldValuesAutocompletionExecutor
+from sentry.api.event_search import translate_escape_sequences
+from sentry.api.paginator import ChainPaginator
+from sentry.api.serializers import serialize
+from sentry.models.organization import Organization
+from sentry.search.eap import constants
+from sentry.search.eap.columns import ColumnDefinitions
+from sentry.search.eap.ourlogs.definitions import OURLOG_DEFINITIONS
+from sentry.search.eap.resolver import SearchResolver
+from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
+from sentry.search.eap.spans.utils import translate_internal_to_public_alias
+from sentry.search.eap.types import SearchResolverConfig
+from sentry.search.events.types import SnubaParams
+from sentry.snuba.referrer import Referrer
+from sentry.tagstore.types import TagValue
+from sentry.utils import snuba_rpc
+
+
+class TraceItemType(str, Enum):
+    LOGS = "logs"
+    SPANS = "spans"
+
+
+# Mapping from our enum types to the protobuf enum types
+TRACE_ITEM_TYPE_MAP = {
+    TraceItemType.LOGS: ProtoTraceItemType.TRACE_ITEM_TYPE_LOG,
+    TraceItemType.SPANS: ProtoTraceItemType.TRACE_ITEM_TYPE_SPAN,
+}
+
+
+class OrganizationTraceItemAttributesEndpointBase(OrganizationEventsV2EndpointBase):
+    publish_status = {
+        "GET": ApiPublishStatus.PRIVATE,
+    }
+    owner = ApiOwner.PERFORMANCE
+    feature_flag = "organizations:ourlogs-enabled"  # Can be changed to performance-trace-explorer once spans work.
+
+
+class OrganizationTraceItemAttributesEndpointSerializer(serializers.Serializer):
+    item_type = serializers.ChoiceField([e.value for e in TraceItemType], required=True)
+    attribute_type = serializers.ChoiceField(["string", "number"], required=True)
+    substring_match = serializers.CharField(required=False)
+    query = serializers.CharField(required=False)
+
+
+def is_valid_item_type(item_type: str) -> bool:
+    return item_type in [e.value for e in TraceItemType]
+
+
+def resolve_attribute_referrer(item_type: str, attribute_type: str) -> Referrer:
+    return (
+        Referrer.API_SPANS_TAG_KEYS_RPC
+        if item_type == TraceItemType.SPANS.value
+        else Referrer.API_LOGS_TAG_KEYS_RPC
+    )
+
+
+def resolve_attribute_values_referrer(item_type: str) -> Referrer:
+    return (
+        Referrer.API_SPANS_TAG_VALUES_RPC
+        if item_type == TraceItemType.SPANS.value
+        else Referrer.API_LOGS_TAG_VALUES_RPC
+    )
+
+
+def as_attribute_key(name: str, type: Literal["string", "number"]):
+    key = translate_internal_to_public_alias(name, type)
+
+    if key is not None:
+        name = key
+    elif type == "number":
+        key = f"tags[{name},number]"
+    else:
+        key = name
+
+    return {
+        # key is what will be used to query the API
+        "key": key,
+        # name is what will be used to display the tag nicely in the UI
+        "name": name,
+    }
+
+
+def empty_filter(trace_item_type: TraceItemType):
+    column_name = "sentry.body" if trace_item_type == TraceItemType.LOGS else "sentry.description"
+    return TraceItemFilter(
+        exists_filter=ExistsFilter(
+            key=AttributeKey(name=column_name),
+        )
+    )
+
+
+@region_silo_endpoint
+class OrganizationTraceItemAttributesEndpoint(OrganizationTraceItemAttributesEndpointBase):
+    def get(self, request: Request, organization: Organization) -> Response:
+        if not features.has(self.feature_flag, organization, actor=request.user):
+            return Response(status=404)
+
+        serializer = OrganizationTraceItemAttributesEndpointSerializer(data=request.GET)
+        if not serializer.is_valid():
+            return Response(serializer.errors, status=400)
+
+        try:
+            snuba_params = self.get_snuba_params(request, organization)
+        except NoProjects:
+            return self.paginate(
+                request=request,
+                paginator=ChainPaginator([]),
+            )
+
+        serialized = serializer.validated_data
+        substring_match = serialized.get("substring_match", "")
+        query_string = serialized.get("query")
+        attribute_type = serialized.get("attribute_type")
+        item_type = serialized.get("item_type")
+
+        max_attributes = options.get("performance.spans-tags-key.max")
+        value_substring_match = translate_escape_sequences(substring_match)
+        item_type_type = TraceItemType(item_type)
+        referrer = resolve_attribute_referrer(item_type_type, attribute_type)
+        resolver = SearchResolver(
+            params=snuba_params, config=SearchResolverConfig(), definitions=SPAN_DEFINITIONS
+        )
+        filter, _, _ = resolver.resolve_query(query_string)
+        meta = resolver.resolve_meta(referrer=referrer.value)
+        meta.trace_item_type = TRACE_ITEM_TYPE_MAP.get(
+            item_type_type, ProtoTraceItemType.TRACE_ITEM_TYPE_SPAN
+        )
+
+        adjusted_start_date, adjusted_end_date = adjust_start_end_window(
+            snuba_params.start_date, snuba_params.end_date
+        )
+        snuba_params.start = adjusted_start_date
+        snuba_params.end = adjusted_end_date
+
+        filter = filter or empty_filter(item_type_type)
+        attr_type = (
+            AttributeKey.Type.TYPE_DOUBLE
+            if attribute_type == "number"
+            else AttributeKey.Type.TYPE_STRING
+        )
+
+        rpc_request = TraceItemAttributeNamesRequest(
+            meta=meta,
+            limit=max_attributes,
+            offset=0,
+            type=attr_type,
+            value_substring_match=value_substring_match,
+            intersecting_attributes_filter=filter,
+        )
+
+        rpc_response = snuba_rpc.attribute_names_rpc(rpc_request)
+
+        paginator = ChainPaginator(
+            [
+                [
+                    as_attribute_key(attribute.name, serialized["attribute_type"])
+                    for attribute in rpc_response.attributes
+                    if attribute.name
+                ],
+            ],
+            max_limit=max_attributes,
+        )
+
+        return self.paginate(
+            request=request,
+            paginator=paginator,
+            on_results=lambda results: serialize(results, request.user),
+            default_per_page=max_attributes,
+            max_per_page=max_attributes,
+        )
+
+
+@region_silo_endpoint
+class OrganizationTraceItemAttributeValuesEndpoint(OrganizationTraceItemAttributesEndpointBase):
+    def get(self, request: Request, organization: Organization, key: str) -> Response:
+        if not features.has(self.feature_flag, organization, actor=request.user):
+            return Response(status=404)
+
+        serializer = OrganizationTraceItemAttributesEndpointSerializer(data=request.GET)
+        if not serializer.is_valid():
+            return Response(serializer.errors, status=400)
+
+        try:
+            snuba_params = self.get_snuba_params(request, organization)
+        except NoProjects:
+            return self.paginate(
+                request=request,
+                paginator=ChainPaginator([]),
+            )
+
+        sentry_sdk.set_tag("query.attribute_key", key)
+
+        serialized = serializer.validated_data
+        item_type = serialized.get("item_type")
+        substring_match = serialized.get("substring_match", "")
+
+        max_attribute_values = options.get("performance.spans-tags-values.max")
+
+        definitions = (
+            SPAN_DEFINITIONS if item_type == TraceItemType.SPANS.value else OURLOG_DEFINITIONS
+        )
+
+        executor = TraceItemAttributeValuesAutocompletionExecutor(
+            organization=organization,
+            snuba_params=snuba_params,
+            key=key,
+            query=substring_match,
+            max_span_tag_values=max_attribute_values,
+            definitions=definitions,
+        )
+
+        tag_values = executor.execute()
+        tag_values.sort(key=lambda tag: tag.value)
+
+        paginator = ChainPaginator([tag_values], max_limit=max_attribute_values)
+
+        return self.paginate(
+            request=request,
+            paginator=paginator,
+            on_results=lambda results: serialize(results, request.user),
+            default_per_page=max_attribute_values,
+            max_per_page=max_attribute_values,
+        )
+
+
+class TraceItemAttributeValuesAutocompletionExecutor(BaseSpanFieldValuesAutocompletionExecutor):
+    def __init__(
+        self,
+        organization: Organization,
+        snuba_params: SnubaParams,
+        key: str,
+        query: str | None,
+        max_span_tag_values: int,
+        definitions: ColumnDefinitions,
+    ):
+        super().__init__(organization, snuba_params, key, query, max_span_tag_values)
+        self.resolver = SearchResolver(
+            params=snuba_params, config=SearchResolverConfig(), definitions=definitions
+        )
+        self.search_type, self.attribute_key = self.resolve_attribute_key(key, snuba_params)
+
+    def resolve_attribute_key(
+        self, key: str, snuba_params: SnubaParams
+    ) -> tuple[constants.SearchType, AttributeKey]:
+        resolved_attr, _ = self.resolver.resolve_attribute(key)
+        return resolved_attr.search_type, resolved_attr.proto_definition
+
+    def execute(self) -> list[TagValue]:
+        if self.key in self.PROJECT_ID_KEYS:
+            return self.project_id_autocomplete_function()
+
+        if self.key in self.PROJECT_SLUG_KEYS:
+            return self.project_slug_autocomplete_function()
+
+        if self.search_type == "boolean":
+            return self.boolean_autocomplete_function()
+
+        if self.search_type == "string":
+            return self.string_autocomplete_function()
+
+        return []
+
+    def boolean_autocomplete_function(self) -> list[TagValue]:
+        return [
+            TagValue(
+                key=self.key,
+                value="false",
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            ),
+            TagValue(
+                key=self.key,
+                value="true",
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            ),
+        ]
+
+    def string_autocomplete_function(self) -> list[TagValue]:
+        adjusted_start_date, adjusted_end_date = adjust_start_end_window(
+            self.snuba_params.start_date, self.snuba_params.end_date
+        )
+        start_timestamp = Timestamp()
+        start_timestamp.FromDatetime(adjusted_start_date)
+
+        end_timestamp = Timestamp()
+        end_timestamp.FromDatetime(adjusted_end_date)
+
+        query = translate_escape_sequences(self.query)
+
+        meta = self.resolver.resolve_meta(referrer=Referrer.API_SPANS_TAG_VALUES_RPC.value)
+        rpc_request = TraceItemAttributeValuesRequest(
+            meta=meta,
+            key=self.attribute_key,
+            value_substring_match=query,
+            limit=self.max_span_tag_values,
+        )
+        rpc_response = snuba_rpc.attribute_values_rpc(rpc_request)
+
+        return [
+            TagValue(
+                key=self.key,
+                value=value,
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            )
+            for value in rpc_response.values
+            if value
+        ]
+
+
+def adjust_start_end_window(start_date: datetime, end_date: datetime) -> tuple[datetime, datetime]:
+    start_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
+    end_date = end_date.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
+    return start_date, end_date
diff --git a/src/sentry/api/urls.py b/src/sentry/api/urls.py
@@ -27,6 +27,10 @@
 )
 from sentry.api.endpoints.organization_spans_aggregation import OrganizationSpansAggregationEndpoint
 from sentry.api.endpoints.organization_stats_summary import OrganizationStatsSummaryEndpoint
+from sentry.api.endpoints.organization_trace_item_attributes import (
+    OrganizationTraceItemAttributesEndpoint,
+    OrganizationTraceItemAttributeValuesEndpoint,
+)
 from sentry.api.endpoints.organization_unsubscribe import (
     OrganizationUnsubscribeIssue,
     OrganizationUnsubscribeProject,
@@ -1462,6 +1466,16 @@ def create_group_urls(name_prefix: str) -> list[URLPattern | URLResolver]:
         OrganizationTracesStatsEndpoint.as_view(),
         name="sentry-api-0-organization-traces-stats",
     ),
+    re_path(
+        r"^(?P<organization_id_or_slug>[^\/]+)/trace-items/attributes/$",
+        OrganizationTraceItemAttributesEndpoint.as_view(),
+        name="sentry-api-0-organization-trace-item-attributes",
+    ),
+    re_path(
+        r"^(?P<organization_id_or_slug>[^\/]+)/trace-items/attributes/(?P<key>[^/]+)/values/$",
+        OrganizationTraceItemAttributeValuesEndpoint.as_view(),
+        name="sentry-api-0-organization-trace-item-attribute-values",
+    ),
     re_path(
         r"^(?P<organization_id_or_slug>[^\/]+)/spans/fields/$",
         OrganizationSpansFieldsEndpoint.as_view(),
diff --git a/src/sentry/snuba/referrer.py b/src/sentry/snuba/referrer.py
diff --git a/tests/snuba/api/endpoints/test_organization_trace_item_attributes.py b/tests/snuba/api/endpoints/test_organization_trace_item_attributes.py
