ref(eap): Consolidate conditional aggregates

src/sentry/search/eap/columns.py

@@ -29,6 +29,13 @@
 
 ResolvedArgument: TypeAlias = AttributeKey | str | int | float
 ResolvedArguments: TypeAlias = list[ResolvedArgument]
+ProtoDefinition: TypeAlias = (
+    LiteralValue
+    | AttributeKey
+    | AttributeAggregation
+    | AttributeConditionalAggregation
+    | Column.BinaryFormula
+)
 
 
 class ResolverSettings(TypedDict):
@@ -86,13 +93,7 @@ def proto_type(self) -> AttributeKey.Type.ValueType:
     @property
     def proto_definition(
         self,
-    ) -> (
-        LiteralValue
-        | AttributeKey
-        | AttributeAggregation
-        | AttributeConditionalAggregation
-        | Column.BinaryFormula
-    ):
+    ) -> ProtoDefinition:
         raise NotImplementedError
 
 
@@ -315,12 +316,13 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate:
+    ) -> ResolvedFunction:
         raise NotImplementedError()
 
 
 @dataclass(kw_only=True)
 class AggregateDefinition(FunctionDefinition):
+    # The type of aggregation (ex. sum, avg)
     internal_function: Function.ValueType
     # An optional function that takes in the resolved argument and returns the
     # attribute key to aggregate on. If not provided, assumes the aggregate is
@@ -335,7 +337,7 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedAggregate:
+    ) -> ResolvedFunction:
         if len(resolved_arguments) > 1:
             raise InvalidSearchQuery(
                 f"Aggregates expects exactly 1 argument, got {len(resolved_arguments)}"
@@ -376,7 +378,7 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedAggregate:
+    ) -> ResolvedFunction:
         if not isinstance(resolved_arguments[0], AttributeKey):
             raise InvalidSearchQuery(
                 "Trace metric aggregates expect argument 0 to be of type AttributeArgumentDefinition"
@@ -403,16 +405,14 @@ def resolve(
 
 
 @dataclass(kw_only=True)
-class ConditionalAggregateDefinition(FunctionDefinition):
+class ConditionalAggregateDefinition(AggregateDefinition):
     """
     The definition of a conditional aggregation,
     Conditionally aggregates the `key`, if it passes the `filter`.
     The type of aggregation is defined by the `internal_name`.
     The `filter` is returned by the `filter_resolver` function which takes in the args from the user and returns a `TraceItemFilter`.
     """
 
-    # The type of aggregation (ex. sum, avg)
-    internal_function: Function.ValueType
     # A function that takes in the resolved argument and returns the condition to filter on and the key to aggregate on
     aggregate_resolver: Callable[[ResolvedArguments], tuple[AttributeKey, TraceItemFilter]]
 
@@ -424,7 +424,7 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedConditionalAggregate:
+    ) -> ResolvedFunction:
         key, aggregate_filter = self.aggregate_resolver(resolved_arguments)
         return ResolvedConditionalAggregate(
             public_alias=alias,
@@ -460,7 +460,7 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedFormula:
+    ) -> ResolvedFunction:
         resolver_settings = ResolverSettings(
             extrapolation_mode=resolve_extrapolation_mode(
                 search_config, self.extrapolation_mode_override
@@ -493,7 +493,7 @@ def resolve(
         snuba_params: SnubaParams,
         query_result_cache: dict[str, EAPResponse],
         search_config: SearchResolverConfig,
-    ) -> ResolvedFormula:
+    ) -> ResolvedFunction:
         resolver_settings = ResolverSettings(
             extrapolation_mode=resolve_extrapolation_mode(
                 search_config, self.extrapolation_mode_override
@@ -579,7 +579,6 @@ def project_term_resolver(
 @dataclass(frozen=True)
 class ColumnDefinitions:
     aggregates: dict[str, AggregateDefinition]
-    conditional_aggregates: dict[str, ConditionalAggregateDefinition]
     formulas: dict[str, FormulaDefinition]
     columns: dict[str, ResolvedAttribute]
     contexts: dict[str, VirtualColumnDefinition]

src/sentry/search/eap/ourlogs/definitions.py

@@ -11,7 +11,6 @@
 
 OURLOG_DEFINITIONS = ColumnDefinitions(
     aggregates=LOG_AGGREGATE_DEFINITIONS,
-    conditional_aggregates={},
     formulas={},
     columns=OURLOG_ATTRIBUTE_DEFINITIONS,
     contexts=OURLOG_VIRTUAL_CONTEXTS,

src/sentry/search/eap/profile_functions/definitions.py

@@ -9,7 +9,6 @@
 
 PROFILE_FUNCTIONS_DEFINITIONS = ColumnDefinitions(
     aggregates=PROFILE_FUNCTIONS_AGGREGATE_DEFINITIONS,
-    conditional_aggregates={},
     formulas={},
     columns=PROFILE_FUNCTIONS_ATTRIBUTE_DEFINITIONS,
     contexts=PROFILE_FUNCTIONS_VIRTUAL_CONTEXTS,

src/sentry/search/eap/resolver.py

@@ -44,14 +44,11 @@
     AggregateDefinition,
     AttributeArgumentDefinition,
     ColumnDefinitions,
-    ConditionalAggregateDefinition,
     FormulaDefinition,
-    ResolvedAggregate,
     ResolvedAttribute,
     ResolvedColumn,
-    ResolvedConditionalAggregate,
     ResolvedEquation,
-    ResolvedFormula,
+    ResolvedFunction,
     ResolvedLiteral,
     ValueArgumentDefinition,
     VirtualColumnDefinition,
@@ -84,20 +81,18 @@ class SearchResolver:
     _resolved_function_cache: dict[
         str,
         tuple[
-            ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate,
+            ResolvedFunction,
             VirtualColumnDefinition | None,
         ],
     ] = field(default_factory=dict)
 
     def get_function_definition(
         self, function_name: str
-    ) -> ConditionalAggregateDefinition | FormulaDefinition | AggregateDefinition:
+    ) -> FormulaDefinition | AggregateDefinition:
         if function_name in self.definitions.aggregates:
             return self.definitions.aggregates[function_name]
         elif function_name in self.definitions.formulas:
             return self.definitions.formulas[function_name]
-        elif function_name in self.definitions.conditional_aggregates:
-            return self.definitions.conditional_aggregates[function_name]
         else:
             raise InvalidSearchQuery(f"Unknown function {function_name}")
 
@@ -790,9 +785,7 @@ def resolve_contexts(
 
     @sentry_sdk.trace
     def resolve_columns(self, selected_columns: list[str], has_aggregates: bool = False) -> tuple[
-        list[
-            ResolvedAttribute | ResolvedAggregate | ResolvedConditionalAggregate | ResolvedFormula
-        ],
+        list[ResolvedAttribute | ResolvedFunction],
         list[VirtualColumnDefinition | None],
     ]:
         """Given a list of columns resolve them and get their context if applicable
@@ -832,7 +825,7 @@ def resolve_column(
         match: Match[str] | None = None,
         public_alias_override: str | None = None,
     ) -> tuple[
-        ResolvedAttribute | ResolvedAggregate | ResolvedConditionalAggregate | ResolvedFormula,
+        ResolvedAttribute | ResolvedFunction,
         VirtualColumnDefinition | None,
     ]:
         """Column is either an attribute or an aggregate, this function will determine which it is and call the relevant
@@ -938,7 +931,7 @@ def resolve_attribute(
 
     @sentry_sdk.trace
     def resolve_functions(self, columns: list[str]) -> tuple[
-        list[ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate],
+        list[ResolvedFunction],
         list[VirtualColumnDefinition | None],
     ]:
         """Helper function to resolve a list of functions instead of 1 attribute at a time"""
@@ -954,10 +947,7 @@ def resolve_function(
         column: str,
         match: Match[str] | None = None,
         public_alias_override: str | None = None,
-    ) -> tuple[
-        ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate,
-        VirtualColumnDefinition | None,
-    ]:
+    ) -> tuple[ResolvedFunction, VirtualColumnDefinition | None]:
         if match is None:
             match = fields.is_function(column)
             if match is None:
@@ -1171,18 +1161,25 @@ def _resolve_operation(self, operation: arithmetic.OperandType) -> tuple[
             )
         elif isinstance(operation, float):
             return Column(literal=LiteralValue(val_double=operation)), []
-        else:
-            # Resolve the column, and turn it into a RPC Column so it can be used in a BinaryFormula
-            col, context = self.resolve_column(operation)
-            contexts = [context] if context is not None else []
-            if isinstance(col, ResolvedAttribute):
-                return Column(key=col.proto_definition), contexts
-            elif isinstance(col, ResolvedAggregate):
-                return Column(aggregation=col.proto_definition), contexts
-            elif isinstance(col, ResolvedConditionalAggregate):
-                return Column(conditional_aggregation=col.proto_definition), contexts
-            elif isinstance(col, ResolvedFormula):
-                return Column(formula=col.proto_definition), contexts
+
+        # Resolve the column, and turn it into a RPC Column so it can be used in a BinaryFormula
+        col, context = self.resolve_column(operation)
+        contexts = [context] if context is not None else []
+        proto_definition = col.proto_definition
+
+        if isinstance(proto_definition, AttributeKey):
+            return Column(key=proto_definition), contexts
+
+        if isinstance(proto_definition, AttributeAggregation):
+            return Column(aggregation=proto_definition), contexts
+
+        if isinstance(proto_definition, AttributeConditionalAggregation):
+            return Column(conditional_aggregation=proto_definition), contexts
+
+        if isinstance(proto_definition, Column.BinaryFormula):
+            return Column(formula=proto_definition), contexts
+
+        raise TypeError(f"Unsupported proto definition type: {type(proto_definition)}")
 
     def resolve_dataset_conditions(
         self,

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

@@ -180,7 +180,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
     return (attribute, filter)
 
 
-SPAN_CONDITIONAL_AGGREGATE_DEFINITIONS = {
+SPAN_AGGREGATE_DEFINITIONS = {
     "count_op": ConditionalAggregateDefinition(
         internal_function=Function.FUNCTION_COUNT,
         default_search_type="integer",
@@ -455,9 +455,6 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
         processor=lambda x: x > 0,
         extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
     ),
-}
-
-SPAN_AGGREGATE_DEFINITIONS = {
     "sum": AggregateDefinition(
         internal_function=Function.FUNCTION_SUM,
         default_search_type="duration",

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

@@ -1,17 +1,13 @@
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
 
 from sentry.search.eap.columns import ColumnDefinitions
-from sentry.search.eap.spans.aggregates import (
-    SPAN_AGGREGATE_DEFINITIONS,
-    SPAN_CONDITIONAL_AGGREGATE_DEFINITIONS,
-)
+from sentry.search.eap.spans.aggregates import SPAN_AGGREGATE_DEFINITIONS
 from sentry.search.eap.spans.attributes import SPAN_ATTRIBUTE_DEFINITIONS, SPAN_VIRTUAL_CONTEXTS
 from sentry.search.eap.spans.filter_aliases import SPAN_FILTER_ALIAS_DEFINITIONS
 from sentry.search.eap.spans.formulas import SPAN_FORMULA_DEFINITIONS
 
 SPAN_DEFINITIONS = ColumnDefinitions(
     aggregates=SPAN_AGGREGATE_DEFINITIONS,
-    conditional_aggregates=SPAN_CONDITIONAL_AGGREGATE_DEFINITIONS,
     formulas=SPAN_FORMULA_DEFINITIONS,
     columns=SPAN_ATTRIBUTE_DEFINITIONS,
     contexts=SPAN_VIRTUAL_CONTEXTS,

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

@@ -2,12 +2,7 @@
 from typing import 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_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
 
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.columns import ResolvedTraceMetricAggregate, ResolvedTraceMetricFormula
@@ -74,61 +69,15 @@ def _extra_conditions_from_columns(
 
         selected_metric = selected_metrics.pop()
 
-        return get_metric_filter(search_resolver, selected_metric)
+        return selected_metric.get_filter()
 
     def _extra_conditions_from_metric(
         self,
         search_resolver: SearchResolver,
     ) -> TraceItemFilter | None:
         if self.metric is None:
             return None
-        return get_metric_filter(search_resolver, self.metric)
-
-
-def get_metric_filter(
-    search_resolver: SearchResolver,
-    metric: TraceMetric,
-) -> TraceItemFilter:
-    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")
-
-    filters = [
-        TraceItemFilter(
-            comparison_filter=ComparisonFilter(
-                key=metric_name.proto_definition,
-                op=ComparisonFilter.OP_EQUALS,
-                value=AttributeValue(val_str=metric.metric_name),
-            )
-        ),
-        TraceItemFilter(
-            comparison_filter=ComparisonFilter(
-                key=metric_type.proto_definition,
-                op=ComparisonFilter.OP_EQUALS,
-                value=AttributeValue(val_str=metric.metric_type),
-            )
-        ),
-    ]
-
-    if metric.metric_unit:
-        metric_unit, _ = search_resolver.resolve_column("metric.unit")
-        if not isinstance(metric_unit.proto_definition, AttributeKey):
-            raise ValueError("Unable to resolve metric.unit")
-        filters.append(
-            TraceItemFilter(
-                comparison_filter=ComparisonFilter(
-                    key=metric_unit.proto_definition,
-                    op=ComparisonFilter.OP_EQUALS,
-                    value=AttributeValue(val_str=metric.metric_unit),
-                )
-            )
-        )
-
-    return TraceItemFilter(and_filter=AndFilter(filters=filters))
+        return self.metric.get_filter()
 
 
 ALLOWED_METRIC_TYPES: list[TraceMetricType] = ["counter", "gauge", "distribution"]

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

@@ -10,7 +10,6 @@
 
 TRACE_METRICS_DEFINITIONS = ColumnDefinitions(
     aggregates=TRACE_METRICS_AGGREGATE_DEFINITIONS,
-    conditional_aggregates={},
     formulas=TRACE_METRICS_FORMULA_DEFINITIONS,
     columns=TRACE_METRICS_ATTRIBUTE_DEFINITIONS,
     contexts=TRACE_METRICS_VIRTUAL_CONTEXTS,