ref(eap): Remove AnyResolved type

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]

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

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,14 +22,20 @@
     TraceItemTableRequest,
     TraceItemTableResponse,
 )
+from sentry_protos.snuba.v1.formula_pb2 import Literal
 from sentry_protos.snuba.v1.request_common_pb2 import (
     PageToken,
     RequestMeta,
     ResponseMeta,
     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)