feat(tracemetrics): Lift metric name and type to top level query params (#102512)

It doesn't make a lot of sense to query multiple metrics together and in
fact, the query would be very inefficient. Instead, opt to lift the
metric name and type to a top level query param for easier handling.

Author: Zylphrex

src/sentry/api/endpoints/organization_events.py

@@ -29,6 +29,10 @@
 from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
 from sentry.models.organization import Organization
 from sentry.ratelimits.config import RateLimitConfig
+from sentry.search.eap.trace_metrics.config import (
+    TraceMetricsSearchResolverConfig,
+    get_trace_metric_info,
+)
 from sentry.search.eap.types import AdditionalQueries, FieldsACL, SearchResolverConfig
 from sentry.snuba import (
     discover,
@@ -527,7 +531,11 @@ def get_rpc_config():
                     )
                 elif scoped_dataset == TraceMetrics:
                     # tracemetrics uses aggregate conditions
-                    return SearchResolverConfig(
+                    metric_name, metric_type = get_trace_metric_info(request)
+
+                    return TraceMetricsSearchResolverConfig(
+                        metric_name=metric_name,
+                        metric_type=metric_type,
                         use_aggregate_conditions=use_aggregate_conditions,
                         auto_fields=True,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,

src/sentry/api/endpoints/organization_events_stats.py

@@ -19,6 +19,10 @@
 from sentry.constants import MAX_TOP_EVENTS
 from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
 from sentry.models.organization import Organization
+from sentry.search.eap.trace_metrics.config import (
+    TraceMetricsSearchResolverConfig,
+    get_trace_metric_info,
+)
 from sentry.search.eap.types import SearchResolverConfig
 from sentry.search.events.types import SnubaParams
 from sentry.snuba import (
@@ -233,6 +237,34 @@ def _get_event_stats(
             if should_upsample:
                 final_columns = transform_query_columns_for_error_upsampling(query_columns)
 
+            def get_rpc_config():
+                if scoped_dataset not in RPC_DATASETS:
+                    raise NotImplementedError
+
+                if scoped_dataset == TraceMetrics:
+                    # tracemetrics uses aggregate conditions
+                    metric_name, metric_type = get_trace_metric_info(request)
+
+                    return TraceMetricsSearchResolverConfig(
+                        metric_name=metric_name,
+                        metric_type=metric_type,
+                        auto_fields=False,
+                        use_aggregate_conditions=True,
+                        disable_aggregate_extrapolation=request.GET.get(
+                            "disableAggregateExtrapolation", "0"
+                        )
+                        == "1",
+                    )
+
+                return SearchResolverConfig(
+                    auto_fields=False,
+                    use_aggregate_conditions=True,
+                    disable_aggregate_extrapolation=request.GET.get(
+                        "disableAggregateExtrapolation", "0"
+                    )
+                    == "1",
+                )
+
             if top_events > 0:
                 raw_groupby = self.get_field_list(organization, request)
                 if "timestamp" in raw_groupby:
@@ -247,14 +279,7 @@ def _get_event_stats(
                         limit=top_events,
                         include_other=include_other,
                         referrer=referrer,
-                        config=SearchResolverConfig(
-                            auto_fields=False,
-                            use_aggregate_conditions=True,
-                            disable_aggregate_extrapolation=request.GET.get(
-                                "disableAggregateExtrapolation", "0"
-                            )
-                            == "1",
-                        ),
+                        config=get_rpc_config(),
                         sampling_mode=snuba_params.sampling_mode,
                         equations=self.get_equation_list(organization, request),
                     )
@@ -285,14 +310,7 @@ def _get_event_stats(
                     query_string=query,
                     y_axes=final_columns,
                     referrer=referrer,
-                    config=SearchResolverConfig(
-                        auto_fields=False,
-                        use_aggregate_conditions=True,
-                        disable_aggregate_extrapolation=request.GET.get(
-                            "disableAggregateExtrapolation", "0"
-                        )
-                        == "1",
-                    ),
+                    config=get_rpc_config(),
                     sampling_mode=snuba_params.sampling_mode,
                     comparison_delta=comparison_delta,
                 )

src/sentry/search/eap/columns.py

@@ -22,7 +22,7 @@
 from sentry.api.event_search import SearchFilter
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap import constants
-from sentry.search.eap.types import EAPResponse
+from sentry.search.eap.types import EAPResponse, SearchResolverConfig
 from sentry.search.events.types import SnubaParams
 
 ResolvedArgument: TypeAlias = AttributeKey | str | int | float
@@ -33,6 +33,7 @@ class ResolverSettings(TypedDict):
     extrapolation_mode: ExtrapolationMode.ValueType
     snuba_params: SnubaParams
     query_result_cache: dict[str, EAPResponse]
+    search_config: SearchResolverConfig
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -298,7 +299,7 @@ def resolve(
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
-        extrapolation_override: bool = False,
+        search_config: SearchResolverConfig,
     ) -> ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate:
         raise NotImplementedError()
 
@@ -319,7 +320,7 @@ def resolve(
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
-        extrapolation_override: bool = False,
+        search_config: SearchResolverConfig,
     ) -> ResolvedAggregate:
         if len(resolved_arguments) > 1:
             raise InvalidSearchQuery(
@@ -341,7 +342,9 @@ def resolve(
             search_type=search_type,
             internal_type=self.internal_type,
             processor=self.processor,
-            extrapolation=self.extrapolation if not extrapolation_override else False,
+            extrapolation=(
+                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
+            ),
             argument=resolved_attribute,
         )
 
@@ -367,7 +370,7 @@ def resolve(
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
-        extrapolation_override: bool = False,
+        search_config: SearchResolverConfig,
     ) -> ResolvedConditionalAggregate:
         key, aggregate_filter = self.aggregate_resolver(resolved_arguments)
         return ResolvedConditionalAggregate(
@@ -378,7 +381,9 @@ def resolve(
             filter=aggregate_filter,
             key=key,
             processor=self.processor,
-            extrapolation=self.extrapolation if not extrapolation_override else False,
+            extrapolation=(
+                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
+            ),
         )
 
 
@@ -401,16 +406,17 @@ def resolve(
         resolved_arguments: list[AttributeKey | Any],
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
-        extrapolation_override: bool = False,
+        search_config: SearchResolverConfig,
     ) -> ResolvedFormula:
         resolver_settings = ResolverSettings(
             extrapolation_mode=(
                 ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-                if self.extrapolation and not extrapolation_override
+                if self.extrapolation and not search_config.disable_aggregate_extrapolation
                 else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
             ),
             snuba_params=snuba_params,
             query_result_cache=query_result_cache,
+            search_config=search_config,
         )
 
         return ResolvedFormula(

src/sentry/search/eap/resolver.py

@@ -1064,7 +1064,7 @@ def resolve_function(
             resolved_arguments=resolved_arguments,
             snuba_params=self.params,
             query_result_cache=self._query_result_cache,
-            extrapolation_override=self.config.disable_aggregate_extrapolation,
+            search_config=self.config,
         )
 
         resolved_context = None

src/sentry/search/eap/trace_metrics/config.py

@@ -0,0 +1,87 @@
+from dataclasses import dataclass
+from typing import Literal, cast
+
+from rest_framework.request import Request
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
+    AndFilter,
+    ComparisonFilter,
+    TraceItemFilter,
+)
+
+from sentry.search.eap.resolver import SearchResolver
+from sentry.search.eap.types import SearchResolverConfig
+
+MetricType = Literal["counter", "gauge", "distribution"]
+
+
+@dataclass(frozen=True, kw_only=True)
+class TraceMetricsSearchResolverConfig(SearchResolverConfig):
+    metric_name: str
+    metric_type: MetricType
+
+    def extra_conditions(self, search_resolver: SearchResolver) -> TraceItemFilter | None:
+        if not self.metric_name or not self.metric_type:
+            return None
+
+        metric_name, _ = search_resolver.resolve_column("metric.name")
+        if not isinstance(metric_name.proto_definition, AttributeKey):
+            raise ValueError("Unable to resolve metric.name")
+
+        metric_type, _ = search_resolver.resolve_column("metric.type")
+        if not isinstance(metric_type.proto_definition, AttributeKey):
+            raise ValueError("Unable to resolve metric.type")
+
+        metric_name_filter = TraceItemFilter(
+            comparison_filter=ComparisonFilter(
+                key=metric_name.proto_definition,
+                op=ComparisonFilter.OP_EQUALS,
+                value=AttributeValue(val_str=self.metric_name),
+            )
+        )
+        metric_type_filter = TraceItemFilter(
+            comparison_filter=ComparisonFilter(
+                key=metric_type.proto_definition,
+                op=ComparisonFilter.OP_EQUALS,
+                value=AttributeValue(val_str=self.metric_type),
+            )
+        )
+
+        return TraceItemFilter(
+            and_filter=AndFilter(filters=[metric_name_filter, metric_type_filter])
+        )
+
+
+ALLOWED_METRIC_TYPES: list[MetricType] = ["counter", "gauge", "distribution"]
+
+
+def get_trace_metric_info(
+    request: Request,
+) -> tuple[str, MetricType]:
+    metric_name = request.GET.get("metricName")
+    metric_type = request.GET.get("metricType")
+
+    # TODO: remove this when these args are not optional anymore
+    if not metric_name:
+        metric_name = ""
+    if not metric_type:
+        metric_type = ""
+
+    # errors = {}
+
+    # if not metric_name:
+    #     errors["metricName"] = ErrorDetail("This field is required.", code="required")
+
+    # if not metric_type:
+    #     errors["metricType"] = ErrorDetail("This field is required.", code="required")
+    # elif metric_type not in ALLOWED_METRIC_TYPES:
+    #     errors["metricType"] = ErrorDetail(
+    #         string=f'"{metric_type}" is not a valid choice.', code="invalid_choice"
+    #     )
+
+    # if errors:
+    #     raise ValidationError(errors)
+
+    # assert metric_name
+
+    return metric_name, cast(MetricType, metric_type)

src/sentry/search/eap/trace_metrics/formulas.py

@@ -8,12 +8,14 @@
     Function,
 )
 
+from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.columns import (
     FormulaDefinition,
     ResolvedArguments,
     ResolverSettings,
     ValueArgumentDefinition,
 )
+from sentry.search.eap.trace_metrics.config import TraceMetricsSearchResolverConfig
 
 
 def _rate_internal(
@@ -22,6 +24,10 @@ def _rate_internal(
     """
     Calculate rate per X for trace metrics using the value attribute.
     """
+    search_config = settings["search_config"]
+    if not isinstance(search_config, TraceMetricsSearchResolverConfig):
+        raise InvalidSearchQuery("unexpected search config")
+
     extrapolation_mode = settings["extrapolation_mode"]
     is_timeseries_request = settings["snuba_params"].is_timeseries_request
 
@@ -31,6 +37,8 @@ def _rate_internal(
         else settings["snuba_params"].interval
     )
 
+    metric_type = search_config.metric_type if search_config.metric_type else metric_type
+
     if metric_type == "counter":
         aggregate_func = Function.FUNCTION_SUM
     else:
@@ -55,7 +63,6 @@ def per_second(args: ResolvedArguments, settings: ResolverSettings) -> Column.Bi
     """
     Calculate rate per second for trace metrics using the value attribute.
     """
-
     metric_type = cast(str, args[1]) if len(args) > 1 else "counter"
     return _rate_internal(1, metric_type, settings)
 

src/sentry/search/eap/types.py

@@ -1,20 +1,24 @@
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Literal, NotRequired, TypedDict
+from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict
 
 from sentry_protos.snuba.v1.request_common_pb2 import PageToken
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import Reliability
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
 
 from sentry.search.events.types import EventsResponse
 
+if TYPE_CHECKING:
+    from sentry.search.eap.resolver import SearchResolver
+
 
 @dataclass(frozen=True)
 class FieldsACL:
     functions: set[str] = field(default_factory=set)
     attributes: set[str] = field(default_factory=set)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, kw_only=True)
 class SearchResolverConfig:
     # Automatically add id, etc. if there are no aggregates
     auto_fields: bool = False
@@ -30,6 +34,9 @@ class SearchResolverConfig:
     # Whether to set the timestamp granularities to stable buckets
     stable_timestamp_quantization: bool = True
 
+    def extra_conditions(self, search_resolver: "SearchResolver") -> TraceItemFilter | None:
+        return None
+
 
 CONFIDENCES: dict[Reliability.ValueType, Literal["low", "high"]] = {
     Reliability.RELIABILITY_LOW: "low",