diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index 8bd27a68ca6f8f..746a8c79e4e862 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -54,6 +54,7 @@ class ResolvedColumn: # Indicates this attribute is a secondary alias for the attribute. # It exists for compatibility or convenience reasons and should NOT be preferred. secondary_alias: bool = False + is_aggregate: bool def process_column(self, value: Any) -> Any: """Given the value from results, return a processed value if a processor is defined otherwise return it""" @@ -81,6 +82,18 @@ def proto_type(self) -> AttributeKey.Type.ValueType: else: return constants.TYPE_MAP[self.search_type] + @property + def proto_definition( + self, + ) -> ( + LiteralValue + | AttributeKey + | AttributeAggregation + | AttributeConditionalAggregation + | Column.BinaryFormula + ): + raise NotImplementedError + @dataclass(frozen=True, kw_only=True) class ResolvedLiteral(ResolvedColumn): @@ -575,17 +588,6 @@ def project_term_resolver( return int(raw_value) -# Any of the resolved attributes, mostly to clean typing up so there's not this giant list all over the code -AnyResolved = ( - ResolvedAttribute - | ResolvedAggregate - | ResolvedConditionalAggregate - | ResolvedFormula - | ResolvedEquation - | ResolvedLiteral -) - - @dataclass(frozen=True) class ColumnDefinitions: aggregates: dict[str, AggregateDefinition] diff --git a/src/sentry/search/eap/resolver.py b/src/sentry/search/eap/resolver.py index 84f580682882b5..ff354e6f3138e2 100644 --- a/src/sentry/search/eap/resolver.py +++ b/src/sentry/search/eap/resolver.py @@ -42,13 +42,13 @@ from sentry.search.eap import constants from sentry.search.eap.columns import ( AggregateDefinition, - AnyResolved, AttributeArgumentDefinition, ColumnDefinitions, ConditionalAggregateDefinition, FormulaDefinition, ResolvedAggregate, ResolvedAttribute, + ResolvedColumn, ResolvedConditionalAggregate, ResolvedEquation, ResolvedFormula, @@ -1081,7 +1081,7 @@ def resolve_function( return self._resolved_function_cache[alias] def resolve_equations(self, equations: list[str]) -> tuple[ - list[AnyResolved], + list[ResolvedColumn], list[VirtualColumnDefinition], ]: formulas = [] @@ -1093,7 +1093,7 @@ def resolve_equations(self, equations: list[str]) -> tuple[ return formulas, contexts def resolve_equation(self, equation: str) -> tuple[ - AnyResolved, + ResolvedColumn, list[VirtualColumnDefinition], ]: """Resolve an equation creating a ResolvedEquation object, we don't just return a Column.BinaryFormula since diff --git a/src/sentry/snuba/rpc_dataset_common.py b/src/sentry/snuba/rpc_dataset_common.py index e84d7dcbf78525..0e60bc0957c749 100644 --- a/src/sentry/snuba/rpc_dataset_common.py +++ b/src/sentry/snuba/rpc_dataset_common.py @@ -7,6 +7,9 @@ import sentry_sdk from google.protobuf.json_format import MessageToJson +from sentry_protos.snuba.v1.attribute_conditional_aggregation_pb2 import ( + AttributeConditionalAggregation, +) from sentry_protos.snuba.v1.downsampled_storage_pb2 import DownsampledStorageConfig from sentry_protos.snuba.v1.endpoint_time_series_pb2 import ( Expression, @@ -19,6 +22,7 @@ TraceItemTableRequest, TraceItemTableResponse, ) +from sentry_protos.snuba.v1.formula_pb2 import Literal from sentry_protos.snuba.v1.request_common_pb2 import ( PageToken, RequestMeta, @@ -26,7 +30,12 @@ TraceItemFilterWithType, TraceItemType, ) -from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue, Function +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( + AttributeAggregation, + AttributeKey, + AttributeValue, + Function, +) from sentry_protos.snuba.v1.trace_item_filter_pb2 import ( AndFilter, ComparisonFilter, @@ -37,16 +46,7 @@ from sentry.api.event_search import SearchFilter, SearchKey, SearchValue from sentry.discover import arithmetic from sentry.exceptions import InvalidSearchQuery -from sentry.search.eap.columns import ( - AnyResolved, - ColumnDefinitions, - ResolvedAggregate, - ResolvedAttribute, - ResolvedConditionalAggregate, - ResolvedEquation, - ResolvedFormula, - ResolvedLiteral, -) +from sentry.search.eap.columns import ColumnDefinitions, ResolvedAttribute, ResolvedColumn from sentry.search.eap.constants import DOUBLE, MAX_ROLLUP_POINTS, VALID_GRANULARITIES from sentry.search.eap.resolver import SearchResolver from sentry.search.eap.rpc_utils import and_trace_item_filters @@ -97,7 +97,7 @@ class TableRequest: """Container for rpc requests""" rpc_request: TraceItemTableRequest - columns: list[AnyResolved] + columns: list[ResolvedColumn] def check_timeseries_has_data(timeseries: SnubaData, y_axes: list[str]): @@ -140,39 +140,48 @@ def get_resolver( @classmethod def categorize_column( cls, - column: AnyResolved, + column: ResolvedColumn, ) -> Column: - # Can't do bare literals, so they're actually formulas with +0 - if isinstance(column, (ResolvedFormula, ResolvedEquation, ResolvedLiteral)): - return Column(formula=column.proto_definition, label=column.public_alias) - elif isinstance(column, ResolvedAggregate): - return Column(aggregation=column.proto_definition, label=column.public_alias) - elif isinstance(column, ResolvedConditionalAggregate): - return Column( - conditional_aggregation=column.proto_definition, label=column.public_alias - ) - else: - return Column(key=column.proto_definition, label=column.public_alias) + proto_definition = column.proto_definition + + if isinstance(proto_definition, AttributeKey): + return Column(key=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, AttributeAggregation): + return Column(aggregation=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, AttributeConditionalAggregation): + return Column(conditional_aggregation=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, Column.BinaryFormula): + return Column(formula=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, Literal): + return Column(literal=proto_definition, label=column.public_alias) + + raise TypeError(f"Unsupported proto definition type: {type(proto_definition)}") @classmethod def categorize_aggregate( cls, - column: AnyResolved, + column: ResolvedColumn, ) -> Expression: - if isinstance(column, (ResolvedFormula, ResolvedEquation)): + proto_definition = column.proto_definition + + if isinstance(proto_definition, AttributeAggregation): + return Expression(aggregation=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, AttributeConditionalAggregation): + return Expression(conditional_aggregation=proto_definition, label=column.public_alias) + + if isinstance(proto_definition, Column.BinaryFormula): # TODO: Remove when https://github.com/getsentry/eap-planning/issues/206 is merged, since we can use formulas in both APIs at that point return Expression( - formula=transform_binary_formula_to_expression(column.proto_definition), + formula=transform_binary_formula_to_expression(proto_definition), label=column.public_alias, ) - elif isinstance(column, ResolvedAggregate): - return Expression(aggregation=column.proto_definition, label=column.public_alias) - elif isinstance(column, ResolvedConditionalAggregate): - return Expression( - conditional_aggregation=column.proto_definition, label=column.public_alias - ) - else: - raise Exception(f"Unknown column type {type(column)}") + + raise TypeError(f"Unsupported proto definition type: {type(proto_definition)}") @classmethod def get_cross_trace_queries(cls, query: TableQuery) -> list[TraceItemFilterWithType]: @@ -251,7 +260,7 @@ def get_table_rpc_request(cls, query: TableQuery) -> TableRequest: # incomplete traces. meta.downsampled_storage_config.mode = DownsampledStorageConfig.MODE_HIGHEST_ACCURACY - all_columns: list[AnyResolved] = [] + all_columns: list[ResolvedColumn] = [] equations, equation_contexts = resolver.resolve_equations( query.equations if query.equations else [] ) @@ -594,7 +603,7 @@ def get_timeseries_query( extra_conditions: TraceItemFilter | None = None, ) -> tuple[ TimeSeriesRequest, - list[AnyResolved], + list[ResolvedColumn], list[ResolvedAttribute], ]: selected_equations, selected_axes = arithmetic.categorize_columns(y_axes)