diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index ca552a121f5474..346c195519331a 100644 --- a/src/sentry/search/eap/columns.py +++ b/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,7 +405,7 @@ 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`. @@ -411,8 +413,6 @@ class ConditionalAggregateDefinition(FunctionDefinition): 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] diff --git a/src/sentry/search/eap/ourlogs/definitions.py b/src/sentry/search/eap/ourlogs/definitions.py index 2fe10eb8366534..32ee75573445aa 100644 --- a/src/sentry/search/eap/ourlogs/definitions.py +++ b/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, diff --git a/src/sentry/search/eap/profile_functions/definitions.py b/src/sentry/search/eap/profile_functions/definitions.py index ca60342183f4f3..ac296a546c040e 100644 --- a/src/sentry/search/eap/profile_functions/definitions.py +++ b/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, diff --git a/src/sentry/search/eap/resolver.py b/src/sentry/search/eap/resolver.py index ff354e6f3138e2..bad24f9be249ba 100644 --- a/src/sentry/search/eap/resolver.py +++ b/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, diff --git a/src/sentry/search/eap/spans/aggregates.py b/src/sentry/search/eap/spans/aggregates.py index 59de1ff0ebe6f7..4a9eb1b1cd32dd 100644 --- a/src/sentry/search/eap/spans/aggregates.py +++ b/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", diff --git a/src/sentry/search/eap/spans/definitions.py b/src/sentry/search/eap/spans/definitions.py index 03c6eca4f040e4..6207a00e89f928 100644 --- a/src/sentry/search/eap/spans/definitions.py +++ b/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, diff --git a/src/sentry/search/eap/trace_metrics/config.py b/src/sentry/search/eap/trace_metrics/config.py index ac92ec9b94285b..a9550deb072104 100644 --- a/src/sentry/search/eap/trace_metrics/config.py +++ b/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,7 +69,7 @@ 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, @@ -82,53 +77,7 @@ def _extra_conditions_from_metric( ) -> 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"] diff --git a/src/sentry/search/eap/trace_metrics/definitions.py b/src/sentry/search/eap/trace_metrics/definitions.py index 01b2ecd5007cff..85389cf9bce670 100644 --- a/src/sentry/search/eap/trace_metrics/definitions.py +++ b/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, diff --git a/src/sentry/search/eap/trace_metrics/types.py b/src/sentry/search/eap/trace_metrics/types.py index dcff895495f80e..6a4bb1cf104da6 100644 --- a/src/sentry/search/eap/trace_metrics/types.py +++ b/src/sentry/search/eap/trace_metrics/types.py @@ -1,6 +1,13 @@ from dataclasses import dataclass from typing import Literal +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, +) + TraceMetricType = Literal["counter", "gauge", "distribution"] @@ -9,3 +16,36 @@ class TraceMetric: metric_name: str metric_type: TraceMetricType metric_unit: str | None + + def get_filter(self) -> TraceItemFilter: + filters = [ + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(name="sentry.metric_name", type=AttributeKey.Type.TYPE_STRING), + op=ComparisonFilter.OP_EQUALS, + value=AttributeValue(val_str=self.metric_name), + ) + ), + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(name="sentry.metric_type", type=AttributeKey.Type.TYPE_STRING), + op=ComparisonFilter.OP_EQUALS, + value=AttributeValue(val_str=self.metric_type), + ) + ), + ] + + if self.metric_unit: + filters.append( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey( + name="sentry.metric_unit", type=AttributeKey.Type.TYPE_STRING + ), + op=ComparisonFilter.OP_EQUALS, + value=AttributeValue(val_str=self.metric_unit), + ) + ) + ) + + return TraceItemFilter(and_filter=AndFilter(filters=filters)) diff --git a/src/sentry/search/eap/uptime_results/definitions.py b/src/sentry/search/eap/uptime_results/definitions.py index 2fb6fa159e3e95..666b3235d8b828 100644 --- a/src/sentry/search/eap/uptime_results/definitions.py +++ b/src/sentry/search/eap/uptime_results/definitions.py @@ -5,7 +5,6 @@ UPTIME_RESULT_DEFINITIONS = ColumnDefinitions( aggregates={}, - conditional_aggregates={}, formulas={}, columns=UPTIME_ATTRIBUTE_DEFINITIONS, contexts={},