feat(eap-spans): Support formulas with time series endpoint (#87031)

Closes getsentry/team-visibility#36

1. Adds support for formulas and conditional_aggregations in time series
requests
2. Converts `Column.BinaryFormula` -> `Expression.BinaryFormula` (until
getsentry/eap-planning#206) so that custom
functions support timeseries.
3. I'm also now passing all aggregations via `expressions` instead of
`aggregations`. This better matches what we already do in the table
query `rpc_dataset_common.py`

---------

Co-authored-by: getsantry[bot] <66042841+getsantry[bot]@users.noreply.github.com>

Author: DominikB2014

src/sentry/search/eap/utils.py

@@ -3,10 +3,20 @@
 from typing import Any
 
 from google.protobuf.timestamp_pb2 import Timestamp
-from sentry_protos.snuba.v1.endpoint_time_series_pb2 import TimeSeriesRequest
+from sentry_protos.snuba.v1.endpoint_time_series_pb2 import Expression, TimeSeriesRequest
+from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import Column
 
 from sentry.exceptions import InvalidSearchQuery
 
+# TODO: Remove when https://github.com/getsentry/eap-planning/issues/206 is merged, since we can use formulas in both APIs at that point
+BINARY_FORMULA_OPERATOR_MAP = {
+    Column.BinaryFormula.OP_ADD: Expression.BinaryFormula.OP_ADD,
+    Column.BinaryFormula.OP_SUBTRACT: Expression.BinaryFormula.OP_SUBTRACT,
+    Column.BinaryFormula.OP_MULTIPLY: Expression.BinaryFormula.OP_MULTIPLY,
+    Column.BinaryFormula.OP_DIVIDE: Expression.BinaryFormula.OP_DIVIDE,
+    Column.BinaryFormula.OP_UNSPECIFIED: Expression.BinaryFormula.OP_UNSPECIFIED,
+}
+
 
 def literal_validator(values: list[Any]) -> Callable[[str], bool]:
     def _validator(input: str) -> bool:
@@ -28,3 +38,29 @@ def add_start_end_conditions(
     in_msg.meta.end_timestamp.CopyFrom(end_time_proto)
 
     return in_msg
+
+
+def transform_binary_formula_to_expression(
+    column: Column.BinaryFormula,
+) -> Expression.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.BinaryFormula(
+        left=transform_column_to_expression(column.left),
+        right=transform_column_to_expression(column.right),
+        op=BINARY_FORMULA_OPERATOR_MAP[column.op],
+    )
+
+
+def transform_column_to_expression(column: Column) -> Expression:
+    """TODO: Remove when https://github.com/getsentry/eap-planning/issues/206 is merged, since we can use formulas in both APIs at that point"""
+    if column.formula.op != Column.BinaryFormula.OP_UNSPECIFIED:
+        return Expression(
+            formula=transform_binary_formula_to_expression(column.formula),
+            label=column.label,
+        )
+
+    return Expression(
+        aggregation=column.aggregation,
+        conditional_aggregation=column.conditional_aggregation,
+        label=column.label,
+    )

src/sentry/snuba/spans_rpc.py

@@ -6,9 +6,13 @@
 
 import sentry_sdk
 from sentry_protos.snuba.v1.endpoint_get_trace_pb2 import GetTraceRequest
-from sentry_protos.snuba.v1.endpoint_time_series_pb2 import TimeSeries, TimeSeriesRequest
+from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
+    Expression,
+    TimeSeries,
+    TimeSeriesRequest,
+)
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
-from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeAggregation, AttributeKey
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import AndFilter, OrFilter, TraceItemFilter
 
 from sentry.api.event_search import SearchFilter, SearchKey, SearchValue
@@ -23,6 +27,7 @@
 from sentry.search.eap.resolver import SearchResolver
 from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
 from sentry.search.eap.types import CONFIDENCES, EAPResponse, SearchResolverConfig
+from sentry.search.eap.utils import transform_binary_formula_to_expression
 from sentry.search.events.fields import is_function
 from sentry.search.events.types import EventsMeta, SnubaData, SnubaParams
 from sentry.snuba import rpc_dataset_common
@@ -33,6 +38,23 @@
 logger = logging.getLogger("sentry.snuba.spans_rpc")
 
 
+def categorize_aggregate(
+    column: ResolvedAggregate | ResolvedConditionalAggregate | ResolvedFormula,
+) -> Expression:
+    if isinstance(column, ResolvedFormula):
+        # 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),
+            label=column.public_alias,
+        )
+    if isinstance(column, ResolvedAggregate):
+        return Expression(aggregation=column.proto_definition, label=column.public_alias)
+    if isinstance(column, ResolvedConditionalAggregate):
+        return Expression(
+            conditional_aggregation=column.proto_definition, label=column.public_alias
+        )
+
+
 @dataclass
 class ProcessedTimeseries:
     timeseries: SnubaData = field(default_factory=list)
@@ -89,7 +111,7 @@ def get_timeseries_query(
     resolver = get_resolver(params=params, config=config)
     meta = resolver.resolve_meta(referrer=referrer)
     query, _, query_contexts = resolver.resolve_query(query_string)
-    (aggregations, _) = resolver.resolve_functions(y_axes)
+    (functions, _) = resolver.resolve_functions(y_axes)
     (groupbys, _) = resolver.resolve_attributes(groupby)
     if extra_conditions is not None:
         if query is not None:
@@ -101,19 +123,15 @@ def get_timeseries_query(
         TimeSeriesRequest(
             meta=meta,
             filter=query,
-            aggregations=[
-                agg.proto_definition
-                for agg in aggregations
-                if isinstance(agg.proto_definition, AttributeAggregation)
-            ],
+            expressions=[categorize_aggregate(fn) for fn in functions if fn.is_aggregate],
             group_by=[
                 groupby.proto_definition
                 for groupby in groupbys
                 if isinstance(groupby.proto_definition, AttributeKey)
             ],
             granularity_secs=granularity_secs,
         ),
-        aggregations,
+        functions,
         groupbys,
     )
 
@@ -178,7 +196,7 @@ def run_timeseries_query(
         )
 
     if comparison_delta is not None:
-        if len(rpc_request.aggregations) != 1:
+        if len(rpc_request.expressions) != 1:
             raise InvalidSearchQuery("Only one column can be selected for comparison queries")
 
         comp_query_params = params.copy()
@@ -305,7 +323,7 @@ def run_top_events_timeseries_query(
         params,
         query_string,
         y_axes,
-        [],  # in the other series, we want eveything in a single group, so remove the group by
+        [],  # in the other series, we want eveything in a single group, so the group by
         referrer,
         config,
         granularity_secs,

tests/sentry/snuba/test_entity_subscriptions.py

@@ -412,7 +412,7 @@ def test_get_entity_subscription_for_eap_rpc_query(self) -> None:
 
         assert rpc_timeseries_request.granularity_secs == 3600
         assert rpc_timeseries_request.filter.comparison_filter.value.val_str == "http.client"
-        assert rpc_timeseries_request.aggregations[0].label == "count(span.duration)"
+        assert rpc_timeseries_request.expressions[0].aggregation.label == "count(span.duration)"
 
 
 class GetEntitySubscriptionFromSnubaQueryTest(TestCase):

tests/sentry/snuba/test_tasks.py

@@ -306,11 +306,11 @@ def test_eap_rpc_query_count(self):
                 == "http.client"
             )
             assert (
-                createSubscriptionRequest.time_series_request.aggregations[0].aggregate
+                createSubscriptionRequest.time_series_request.expressions[0].aggregation.aggregate
                 == FUNCTION_COUNT
             )
             assert (
-                createSubscriptionRequest.time_series_request.aggregations[0].key.name
+                createSubscriptionRequest.time_series_request.expressions[0].aggregation.key.name
                 == "sentry.duration_ms"
             )
             # Validate that the spm function uses the correct time window

tests/snuba/api/endpoints/test_organization_events_stats_span_indexed.py

@@ -1287,3 +1287,97 @@ def test_interval_larger_than_period_uses_default_period(self):
         data = response.data["data"]
         assert len(data) == 73
         assert response.data["meta"]["dataset"] == self.dataset
+
+    def test_cache_miss_rate(self):
+        self.store_spans(
+            [
+                self.create_span(
+                    {
+                        "data": {"cache.hit": False},
+                    },
+                    start_ts=self.day_ago + timedelta(minutes=1),
+                ),
+                self.create_span(
+                    {
+                        "data": {"cache.hit": True},
+                    },
+                    start_ts=self.day_ago + timedelta(minutes=2),
+                ),
+                self.create_span(
+                    {
+                        "data": {"cache.hit": False},
+                    },
+                    start_ts=self.day_ago + timedelta(minutes=2),
+                ),
+                self.create_span(
+                    {
+                        "data": {"cache.hit": True},
+                    },
+                    start_ts=self.day_ago + timedelta(minutes=2),
+                ),
+                self.create_span(
+                    {
+                        "data": {"cache.hit": True},
+                    },
+                    start_ts=self.day_ago + timedelta(minutes=2),
+                ),
+            ],
+            is_eap=self.is_eap,
+        )
+
+        response = self._do_request(
+            data={
+                "start": self.day_ago,
+                "end": self.day_ago + timedelta(minutes=3),
+                "interval": "1m",
+                "yAxis": "cache_miss_rate()",
+                "project": self.project.id,
+                "dataset": self.dataset,
+            },
+        )
+        assert response.status_code == 200, response.content
+        data = response.data["data"]
+        assert len(data) == 3
+
+        assert data[0][1][0]["count"] == 0.0
+        assert data[1][1][0]["count"] == 1.0
+        assert data[2][1][0]["count"] == 0.25
+        assert response.data["meta"]["dataset"] == self.dataset
+
+    def test_count_op(self):
+        self.store_spans(
+            [
+                self.create_span(
+                    {"op": "queue.process", "sentry_tags": {"op": "queue.publish"}},
+                    start_ts=self.day_ago + timedelta(minutes=1),
+                ),
+                self.create_span(
+                    {"op": "queue.process", "sentry_tags": {"op": "queue.publish"}},
+                    start_ts=self.day_ago + timedelta(minutes=1),
+                ),
+                self.create_span(
+                    {"op": "queue.publish", "sentry_tags": {"op": "queue.publish"}},
+                    start_ts=self.day_ago + timedelta(minutes=2),
+                ),
+            ],
+            is_eap=self.is_eap,
+        )
+
+        response = self._do_request(
+            data={
+                "start": self.day_ago,
+                "end": self.day_ago + timedelta(minutes=3),
+                "interval": "1m",
+                "yAxis": "count_op(queue.publish)",
+                "project": self.project.id,
+                "dataset": self.dataset,
+            },
+        )
+        assert response.status_code == 200, response.content
+        data = response.data["data"]
+        assert len(data) == 3
+
+        assert data[0][1][0]["count"] == 0.0
+        assert data[1][1][0]["count"] == 2.0
+        assert data[2][1][0]["count"] == 1.0
+        assert response.data["meta"]["dataset"] == self.dataset