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

- 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
@@ -237,7 +250,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,15 +395,33 @@ 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 substring_match and offset <= len(column_definitions.columns) + len(
+                column_definitions.contexts
+            ):
+                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 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:
@@ -433,7 +468,7 @@ 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],
     ) -> list[TraceItemAttributeKey]:
         attribute_keys = {}
         for attribute in rpc_response.attributes:
@@ -462,6 +497,7 @@ def serialize_trace_attributes_using_sentry_conventions(
                 aliased_attr.internal_name,
                 attribute_type,
                 trace_item_type,
+                is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
             )
             attribute_keys[attr_key["name"]] = attr_key
 
@@ -476,7 +512,7 @@ def serialize_trace_attributes(
         trace_item_type: SupportedTraceItemType,
         include_internal: bool,
         substring_match: str,
-        aliased_attributes: list[ResolvedAttribute],
+        aliased_attributes: list[ResolvedAttribute | ProxyResolvedAttribute],
     ) -> list[TraceItemAttributeKey]:
         attributes = list(
             filter(
@@ -513,6 +549,7 @@ def serialize_trace_attributes(
                     aliased_attr.internal_name,
                     attribute_type,
                     trace_item_type,
+                    is_proxy=isinstance(aliased_attr, ProxyResolvedAttribute),
                 )
                 attributes.append(attr_key)
         return attributes

src/sentry/search/eap/columns.py

@@ -158,6 +158,8 @@ class AttributeArgumentDefinition(BaseArgumentDefinition):
 @dataclass
 class VirtualColumnDefinition:
     constructor: Callable[[SnubaParams, Any], VirtualColumnContext]
+    # Need a type for the attributes endpoint
+    search_type: constants.SearchType
     # Allows additional processing to the term after its been resolved
     term_resolver: (
         Callable[
@@ -255,7 +257,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 +464,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 +680,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 +739,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,7 @@
         constructor=project_context_constructor(key),
         term_resolver=project_term_resolver,
         filter_column="project.id",
+        search_type="string",
     )
     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(),
 }

tests/snuba/api/endpoints/test_organization_trace_item_attributes.py

@@ -980,6 +980,29 @@ def test_aliased_attribute(self) -> None:
         keys = {item["key"] for item in response.data}
         assert len(keys) == 0
 
+    def test_aliased_attribute_project(self) -> None:
+        span1 = self.create_span(
+            {"sentry_tags": {"op": "foo"}}, start_ts=before_now(days=0, minutes=10)
+        )
+        span2 = self.create_span(
+            {"sentry_tags": {"op": "bar"}}, start_ts=before_now(days=0, minutes=10)
+        )
+        self.store_spans([span1, span2])
+
+        response = self.do_request(query={"attributeType": "string"})
+        assert response.status_code == 200, response.content
+
+        keys = {item["key"] for item in response.data}
+        assert len(keys) > 1
+        assert "project" in keys
+
+        response = self.do_request(query={"attributeType": "string", "substringMatch": "pro"})
+        assert response.status_code == 200, response.content
+
+        keys = {item["key"] for item in response.data}
+        assert len(keys) == 5
+        assert "project" in keys
+
     def test_aliased_attribute_with_paging(self) -> None:
         span1 = self.create_span(
             {"tags": {"tag.op": "foo"}}, start_ts=before_now(days=0, minutes=10)