fix(eap): Handle contexts in trace-item attributes (#112524)

- Take contexts into account for the trace item attributes so we return
them as possible fields

Author: wmak

src/sentry/api/endpoints/organization_trace_item_attributes.py

@@ -20,7 +20,11 @@
     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, OrFilter, TraceItemFilter
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
+    ExistsFilter,
+    OrFilter,
+    TraceItemFilter,
+)
 
 from sentry import features, options
 from sentry.api.api_owners import ApiOwner
@@ -54,7 +58,11 @@
 from sentry.search.eap.resolver import SearchResolver
 from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
 from sentry.search.eap.trace_metrics.definitions import TRACE_METRICS_DEFINITIONS
-from sentry.search.eap.types import SearchResolverConfig, SupportedTraceItemType
+from sentry.search.eap.types import (
+    AttributeSourceType,
+    SearchResolverConfig,
+    SupportedTraceItemType,
+)
 from sentry.search.eap.utils import (
     can_expose_attribute,
     get_secondary_aliases,
@@ -79,6 +87,10 @@
 POSSIBLE_ATTRIBUTE_TYPES = ["string", "number", "boolean"]
 
 
+class ProxyResolvedAttribute(ResolvedAttribute):
+    pass
+
+
 class TraceItemAttributeKey(TypedDict):
     key: str
     name: str
@@ -218,6 +230,7 @@ def as_attribute_key(
     name: str,
     attr_type: Literal["string", "number", "boolean"],
     item_type: SupportedTraceItemType,
+    is_proxy: bool = False,
 ) -> TraceItemAttributeKey:
     public_key, public_name, attribute_source = translate_internal_to_public_alias(
         name, attr_type, item_type
@@ -226,6 +239,8 @@ def as_attribute_key(
 
     if public_key is not None and public_name is not None:
         pass
+    elif is_proxy:
+        public_key = public_name = name
     elif attr_type == "number":
         public_key = f"tags[{name},number]"
         public_name = name
@@ -237,7 +252,11 @@ def as_attribute_key(
         public_name = name
 
     serialized_source: dict[str, str | bool] = {
-        "source_type": attribute_source["source_type"].value
+        "source_type": (
+            attribute_source["source_type"].value
+            if not is_proxy
+            else AttributeSourceType.SENTRY.value
+        )
     }
     if attribute_source.get("is_transformed_alias"):
         serialized_source["is_transformed_alias"] = True
@@ -378,18 +397,54 @@ def query_trace_attributes(
             all_aliased_attributes = []
             # our aliases don't exist in the db, so filter over our aliases
             # virtually page through defined aliases before we hit the db
-            if substring_match and offset <= len(column_definitions.columns):
-                for index, column in enumerate(column_definitions.columns.values()):
-                    if (
-                        column.proto_type == attr_type
-                        and substring_match in column.public_alias
-                        and not column.secondary_alias
-                        and not column.private
-                    ):
-                        all_aliased_attributes.append(column)
+            if offset <= len(column_definitions.columns) + len(column_definitions.contexts):
+                if substring_match:
+                    for column in column_definitions.columns.values():
+                        if (
+                            column.proto_type == attr_type
+                            and substring_match in column.public_alias
+                            and not column.secondary_alias
+                            and not column.private
+                        ):
+                            all_aliased_attributes.append(column)
+                    for (
+                        public_label,
+                        virtual_context,
+                    ) in column_definitions.contexts.items():
+                        if (
+                            substring_match in public_label
+                            and virtual_context.search_type is not None
+                            and not virtual_context.secondary_alias
+                            and constants.TYPE_MAP[virtual_context.search_type] == attr_type
+                        ):
+                            all_aliased_attributes.append(
+                                ProxyResolvedAttribute(
+                                    public_alias=public_label,
+                                    internal_name=public_label,
+                                    search_type=virtual_context.search_type,
+                                )
+                            )
+                else:
+                    for (
+                        public_label,
+                        virtual_context,
+                    ) in column_definitions.contexts.items():
+                        if (
+                            substring_match in public_label
+                            and virtual_context.search_type is not None
+                            and not virtual_context.secondary_alias
+                            and constants.TYPE_MAP[virtual_context.search_type] == attr_type
+                        ):
+                            all_aliased_attributes.append(
+                                ProxyResolvedAttribute(
+                                    public_alias=public_label,
+                                    internal_name=public_label,
+                                    search_type=virtual_context.search_type,
+                                )
+                            )
             aliased_attributes = all_aliased_attributes[offset : offset + limit]
         with sentry_sdk.start_span(op="query", name="attribute_names") as span:
-            if len(aliased_attributes) < limit - 1:
+            if len(aliased_attributes) < limit:
                 offset -= len(all_aliased_attributes) - len(aliased_attributes)
                 limit -= len(aliased_attributes)
                 rpc_request = TraceItemAttributeNamesRequest(
@@ -419,6 +474,7 @@ def query_trace_attributes(
                 include_internal,
                 substring_match,
                 aliased_attributes,
+                all_aliased_attributes,
             )
 
             sentry_sdk.set_context("api_response", {"attributes": attributes})
@@ -433,7 +489,8 @@ def serialize_trace_attributes_using_sentry_conventions(
         trace_item_type: SupportedTraceItemType,
         include_internal: bool,
         substring_match: str,
-        aliased_attributes: list[ResolvedAttribute],
+        aliased_attributes: list[ResolvedAttribute | ProxyResolvedAttribute],
+        exclude_attributes: list[ResolvedAttribute | ProxyResolvedAttribute],
     ) -> list[TraceItemAttributeKey]:
         attribute_keys = {}
         for attribute in rpc_response.attributes:
@@ -457,11 +514,21 @@ def serialize_trace_attributes_using_sentry_conventions(
                     )
 
                 attribute_keys[attr_key["name"]] = attr_key
+        for aliased_attr in exclude_attributes:
+            attr_key = as_attribute_key(
+                aliased_attr.internal_name,
+                attribute_type,
+                trace_item_type,
+                is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
+            )
+            if attr_key["name"] in attribute_keys:
+                del attribute_keys[attr_key["name"]]
         for aliased_attr in aliased_attributes:
             attr_key = as_attribute_key(
                 aliased_attr.internal_name,
                 attribute_type,
                 trace_item_type,
+                is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
             )
             attribute_keys[attr_key["name"]] = attr_key
 
@@ -476,33 +543,42 @@ def serialize_trace_attributes(
         trace_item_type: SupportedTraceItemType,
         include_internal: bool,
         substring_match: str,
-        aliased_attributes: list[ResolvedAttribute],
+        aliased_attributes: list[ResolvedAttribute | ProxyResolvedAttribute],
+        exclude_attributes: list[ResolvedAttribute | ProxyResolvedAttribute],
     ) -> list[TraceItemAttributeKey]:
-        attributes = list(
-            filter(
-                lambda x: (
-                    not is_sentry_convention_replacement_attribute(x["name"], trace_item_type)
+        attribute_keys = {}
+        for attribute in rpc_response.attributes:
+            if attribute.name and can_expose_attribute(
+                attribute.name,
+                trace_item_type,
+                include_internal=include_internal,
+            ):
+                attr_key = as_attribute_key(
+                    attribute.name,
+                    attribute_type,
+                    trace_item_type,
+                )
+                if (
+                    not is_sentry_convention_replacement_attribute(
+                        attr_key["name"], trace_item_type
+                    )
                     # Remove anything where the public alias doesn't match the substring
                     # This can happen when the public alias is different, but that's handled by
                     # aliased_attributes
-                    and (substring_match in x["name"] if substring_match else True)
-                ),
-                [
-                    as_attribute_key(
-                        attribute.name,
-                        attribute_type,
-                        trace_item_type,
-                    )
-                    for attribute in rpc_response.attributes
-                    if attribute.name
-                    and can_expose_attribute(
-                        attribute.name,
-                        trace_item_type,
-                        include_internal=include_internal,
-                    )
-                ],
+                    and (substring_match in attr_key["name"] if substring_match else True)
+                ):
+                    attribute_keys[attr_key["key"]] = attr_key
+        # We need to exclude any aliased attributes here since because of pagination they might have already been seen
+        # earlier
+        for aliased_attr in exclude_attributes:
+            attr_key = as_attribute_key(
+                aliased_attr.internal_name,
+                attribute_type,
+                trace_item_type,
+                is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
             )
-        )
+            if attr_key["name"] in attribute_keys:
+                del attribute_keys[attr_key["name"]]
         for aliased_attr in aliased_attributes:
             if can_expose_attribute(
                 aliased_attr.public_alias,
@@ -513,8 +589,11 @@ def serialize_trace_attributes(
                     aliased_attr.internal_name,
                     attribute_type,
                     trace_item_type,
+                    is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
                 )
-                attributes.append(attr_key)
+                attribute_keys[attr_key["key"]] = attr_key
+        attributes = list(attribute_keys.values())
+        sentry_sdk.set_context("api_response", {"attributes": attributes})
         return attributes
 
 

src/sentry/search/eap/columns.py

@@ -158,6 +158,9 @@ class AttributeArgumentDefinition(BaseArgumentDefinition):
 @dataclass
 class VirtualColumnDefinition:
     constructor: Callable[[SnubaParams, Any], VirtualColumnContext]
+    # Need a type for the attributes endpoint
+    search_type: constants.SearchType
+    secondary_alias: bool = False
     # Allows additional processing to the term after its been resolved
     term_resolver: (
         Callable[
@@ -255,7 +258,9 @@ class ResolvedTraceMetricAggregate(ResolvedFunction):
     trace_filter: TraceItemFilter | None
 
     @property
-    def proto_definition(self) -> AttributeAggregation | AttributeConditionalAggregation:
+    def proto_definition(
+        self,
+    ) -> AttributeAggregation | AttributeConditionalAggregation:
         if self.trace_metric is None and self.trace_filter is None:
             return AttributeAggregation(
                 aggregate=self.internal_name,
@@ -460,9 +465,11 @@ def resolve(
             ),
             key=cast(AttributeKey, resolved_attribute),
             trace_metric=trace_metric,
-            trace_filter=cast(TraceItemFilter, resolved_arguments[0])
-            if isinstance(self, ConditionalTraceMetricAggregateDefinition)
-            else None,
+            trace_filter=(
+                cast(TraceItemFilter, resolved_arguments[0])
+                if isinstance(self, ConditionalTraceMetricAggregateDefinition)
+                else None
+            ),
         )
 
 
@@ -674,7 +681,9 @@ class ColumnDefinitions:
     column_to_alias: Callable[[str], str | None] | None
 
 
-def attribute_key_to_tuple(attribute_key: AttributeKey) -> tuple[str, AttributeKey.Type.ValueType]:
+def attribute_key_to_tuple(
+    attribute_key: AttributeKey,
+) -> tuple[str, AttributeKey.Type.ValueType]:
     return (attribute_key.name, attribute_key.type)
 
 
@@ -731,9 +740,11 @@ def extract_trace_metric_aggregate_arguments(
         return TraceMetric(
             metric_name=cast(str, resolved_arguments[offset + 1]),
             metric_type=cast(TraceMetricType, resolved_arguments[offset + 2]),
-            metric_unit=None
-            if resolved_arguments[offset + 3] == "-"
-            else cast(str, resolved_arguments[offset + 3]),
+            metric_unit=(
+                None
+                if resolved_arguments[offset + 3] == "-"
+                else cast(str, resolved_arguments[offset + 3])
+            ),
         )
     elif all(resolved_argument == "" for resolved_argument in resolved_arguments[offset + 1 :]):
         # no metrics were specified, assume we query all metrics

src/sentry/search/eap/common_columns.py

@@ -12,6 +12,8 @@
         constructor=project_context_constructor(key),
         term_resolver=project_term_resolver,
         filter_column="project.id",
+        search_type="string",
+        secondary_alias=key != "project",
     )
     for key in constants.PROJECT_FIELDS
 }

src/sentry/search/eap/spans/attributes.py

@@ -682,14 +682,17 @@ def is_starred_segment_context_constructor(
         # TODO: need to change this so the VCC is using it too, but would require rewriting the term_resolver
         default_value="Unknown",
         sort_column="sentry.device.class",
+        search_type="string",
     ),
     "span.module": VirtualColumnDefinition(
         constructor=module_context_constructor,
+        search_type="string",
     ),
     "is_starred_transaction": VirtualColumnDefinition(
         constructor=is_starred_segment_context_constructor,
         default_value="false",
         processor=lambda x: True if x == "true" else False,
+        search_type="boolean",
     ),
     **project_virtual_contexts(),
 }