feat: eap support formulas in timeseries endpoint (#6854)

kylemumma · web-flow · commit 6d8de6dd990d · 2025-02-11T10:01:10.000-08:00
this PR implements support for formulas in the timeseries endpoint. it closes this ticket getsentry/eap-planning#27 major changes: * auto-convert `TimeSeriesRequest.aggregations` to `TimeSeriesRequest.expressions` * implement support for formula tests: * I have a test for a non-extrapolated formula as well as an extrapolated one design decisions * reliability doesnt work with formulas * formulas dont work w uptime checks or logs
diff --git a/requirements.txt b/requirements.txt
@@ -29,7 +29,7 @@ python-rapidjson==1.8
 redis==4.5.4
 sentry-arroyo==2.19.12
 sentry-kafka-schemas==1.0.4
-sentry-protos==0.1.58
+sentry-protos==0.1.59
 sentry-redis-tools==0.3.0
 sentry-relay==0.9.5
 sentry-sdk==2.18.0
diff --git a/snuba/web/rpc/v1/endpoint_time_series.py b/snuba/web/rpc/v1/endpoint_time_series.py
@@ -3,6 +3,7 @@
 from typing import Type
 
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
+    Expression,
     TimeSeriesRequest,
     TimeSeriesResponse,
 )
@@ -75,6 +76,19 @@ def _validate_time_buckets(request: TimeSeriesRequest) -> None:
         )
 
 
+def _convert_aggregations_to_expressions(
+    request: TimeSeriesRequest,
+) -> TimeSeriesRequest:
+    if len(request.aggregations) > 0:
+        new_req = TimeSeriesRequest()
+        new_req.CopyFrom(request)
+        new_req.ClearField("aggregations")
+        for agg in request.aggregations:
+            new_req.expressions.append(Expression(aggregation=agg, label=agg.label))
+        return new_req
+    return request
+
+
 class EndpointTimeSeries(RPCEndpoint[TimeSeriesRequest, TimeSeriesResponse]):
     @classmethod
     def version(cls) -> str:
@@ -107,5 +121,6 @@ def _execute(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse:
             raise BadSnubaRPCRequestException(
                 "This endpoint requires meta.trace_item_type to be set (are you requesting spans? logs?)"
             )
+        in_msg = _convert_aggregations_to_expressions(in_msg)
         resolver = self.get_resolver(in_msg.meta.trace_item_type)
         return resolver.resolve(in_msg)
diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py
@@ -1,12 +1,16 @@
 import uuid
 from collections import defaultdict
+from dataclasses import replace
 from datetime import datetime
 from typing import Any, Dict, Iterable
 
 from google.protobuf.json_format import MessageToDict
 from google.protobuf.timestamp_pb2 import Timestamp
+from sentry_protos.snuba.v1.endpoint_time_series_pb2 import DataPoint
+from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
+    Expression as ProtoExpression,
+)
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
-    DataPoint,
     TimeSeries,
     TimeSeriesRequest,
     TimeSeriesResponse,
@@ -23,6 +27,7 @@
 from snuba.query.data_source.simple import Entity
 from snuba.query.dsl import Functions as f
 from snuba.query.dsl import column
+from snuba.query.expressions import Expression
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
@@ -48,6 +53,13 @@
     attribute_key_to_expression,
 )
 
+OP_TO_EXPR = {
+    ProtoExpression.BinaryFormula.OP_ADD: f.plus,
+    ProtoExpression.BinaryFormula.OP_SUBTRACT: f.minus,
+    ProtoExpression.BinaryFormula.OP_MULTIPLY: f.multiply,
+    ProtoExpression.BinaryFormula.OP_DIVIDE: f.divide,
+}
+
 
 def _convert_result_timeseries(
     request: TimeSeriesRequest, data: list[Dict[str, Any]]
@@ -94,7 +106,8 @@ def _convert_result_timeseries(
 
     # to convert the results, need to know which were the groupby columns and which ones
     # were aggregations
-    aggregation_labels = set([agg.label for agg in request.aggregations])
+    aggregation_labels = set([expr.label for expr in request.expressions])
+
     group_by_labels = set([attr.name for attr in request.group_by])
 
     # create a mapping with (all the group by attribute key,val pairs as strs, label name)
@@ -154,7 +167,7 @@ def _convert_result_timeseries(
                 extrapolation_context = ExtrapolationContext.from_row(
                     timeseries.label, row_data
                 )
-                if extrapolation_context.is_data_present:
+                if row_data.get(timeseries.label, None) is not None:
                     timeseries.data_points.append(
                         DataPoint(
                             data=row_data[timeseries.label],
@@ -169,23 +182,19 @@ def _convert_result_timeseries(
     return result_timeseries.values()
 
 
-def _build_query(request: TimeSeriesRequest) -> Query:
-    # TODO: This is hardcoded still
-    entity = Entity(
-        key=EntityKey("eap_spans"),
-        schema=get_entity(EntityKey("eap_spans")).get_data_model(),
-        sample=None,
-    )
+def _get_reliability_context_columns(
+    expressions: Iterable[ProtoExpression],
+) -> list[SelectedExpression]:
+    # this reliability logic ignores formulas, meaning formulas may not properly support reliability
+    additional_context_columns = []
 
-    aggregation_columns = [
-        SelectedExpression(
-            name=aggregation.label, expression=aggregation_to_expression(aggregation)
-        )
-        for aggregation in request.aggregations
-    ]
+    aggregates = []
+    for e in expressions:
+        if e.WhichOneof("expression") == "aggregation":
+            # ignore formulas
+            aggregates.append(e.aggregation)
 
-    additional_context_columns = []
-    for aggregation in request.aggregations:
+    for aggregation in aggregates:
         if (
             aggregation.extrapolation_mode
             == ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
@@ -211,6 +220,41 @@ def _build_query(request: TimeSeriesRequest) -> Query:
         additional_context_columns.append(
             SelectedExpression(name=count_column.alias, expression=count_column)
         )
+    return additional_context_columns
+
+
+def _proto_expression_to_ast_expression(expr: ProtoExpression) -> Expression:
+    match expr.WhichOneof("expression"):
+        case "aggregation":
+            return aggregation_to_expression(expr.aggregation)
+        case "formula":
+            formula_expr = OP_TO_EXPR[expr.formula.op](
+                _proto_expression_to_ast_expression(expr.formula.left),
+                _proto_expression_to_ast_expression(expr.formula.right),
+            )
+            formula_expr = replace(formula_expr, alias=expr.label)
+            return formula_expr
+        case default:
+            raise ValueError(f"Unknown expression type: {default}")
+
+
+def _build_query(request: TimeSeriesRequest) -> Query:
+    # TODO: This is hardcoded still
+    entity = Entity(
+        key=EntityKey("eap_spans"),
+        schema=get_entity(EntityKey("eap_spans")).get_data_model(),
+        sample=None,
+    )
+
+    aggregation_columns = [
+        SelectedExpression(
+            name=expr.label,
+            expression=_proto_expression_to_ast_expression(expr),
+        )
+        for expr in request.expressions
+    ]
+
+    additional_context_columns = _get_reliability_context_columns(request.expressions)
 
     groupby_columns = [
         SelectedExpression(
@@ -302,6 +346,10 @@ def trace_item_type(cls) -> TraceItemType.ValueType:
         return TraceItemType.TRACE_ITEM_TYPE_SPAN
 
     def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse:
+        # aggregations field is deprecated, it gets converted to request.expressions
+        # if the user passes it in
+        assert len(in_msg.aggregations) == 0
+
         snuba_request = _build_snuba_request(in_msg)
         res = run_query(
             dataset=PluggableDataset(name="eap", all_entities=[]),
diff --git a/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_time_series.py b/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_time_series.py
@@ -7,6 +7,7 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
     DataPoint,
+    Expression,
     TimeSeries,
     TimeSeriesRequest,
     TimeSeriesResponse,
@@ -31,6 +32,7 @@
     extract_response_meta,
     setup_trace_query_settings,
 )
+from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.resolvers import ResolverTimeSeries
 from snuba.web.rpc.v1.resolvers.common.aggregation import aggregation_to_expression
 from snuba.web.rpc.v1.resolvers.R_uptime_checks.common.common import (
@@ -40,6 +42,18 @@
 )
 
 
+def _get_aggregation_label(expr: Expression) -> str:
+    match expr.WhichOneof("expression"):
+        case "aggregation":
+            return expr.aggregation.label
+        case "formula":
+            raise BadSnubaRPCRequestException(
+                "formulas are not supported for uptime checks"
+            )
+        case default:
+            raise BadSnubaRPCRequestException(f"Unknown expression type: {default}")
+
+
 def _convert_result_timeseries(
     request: TimeSeriesRequest, data: list[Dict[str, Any]]
 ) -> Iterable[TimeSeries]:
@@ -85,7 +99,10 @@ def _convert_result_timeseries(
 
     # to convert the results, need to know which were the groupby columns and which ones
     # were aggregations
-    aggregation_labels = set([agg.label for agg in request.aggregations])
+    aggregation_labels = set(
+        [_get_aggregation_label(expr) for expr in request.expressions]
+    )
+
     group_by_labels = set([attr.name for attr in request.group_by])
 
     # create a mapping with (all the group by attribute key,val pairs as strs, label name)
@@ -158,16 +175,25 @@ def _build_query(request: TimeSeriesRequest) -> Query:
         sample=None,
     )
 
-    aggregation_columns = [
-        SelectedExpression(
-            name=aggregation.label,
-            expression=aggregation_to_expression(
-                aggregation,
-                attribute_key_to_expression(aggregation.key),
-            ),
-        )
-        for aggregation in request.aggregations
-    ]
+    aggregation_columns = []
+    for expr in request.expressions:
+        match expr.WhichOneof("expression"):
+            case "aggregation":
+                aggregation_columns.append(
+                    SelectedExpression(
+                        name=expr.aggregation.label,
+                        expression=aggregation_to_expression(
+                            expr.aggregation,
+                            attribute_key_to_expression(expr.aggregation.key),
+                        ),
+                    )
+                )
+            case "formula":
+                raise BadSnubaRPCRequestException(
+                    "formulas are not supported for uptime checks"
+                )
+            case default:
+                raise BadSnubaRPCRequestException(f"Unknown expression type: {default}")
 
     groupby_columns = [
         SelectedExpression(
diff --git a/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py b/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py
@@ -9,6 +9,7 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
     DataPoint,
+    Expression,
     TimeSeries,
     TimeSeriesRequest,
 )
@@ -903,6 +904,98 @@ def test_OOM(self, monkeypatch: Any) -> None:
             sentry_sdk_mock.assert_called_once()
             assert metrics_mock.increment.call_args_list.count(call("OOM_query")) == 1
 
+    def test_formula(self) -> None:
+        # store a a test metric with a value of 1, every second of one hour
+        granularity_secs = 300
+        query_duration = 60 * 30
+        store_spans_timeseries(
+            BASE_TIME,
+            1,
+            3600,
+            metrics=[DummyMetric("test_metric", get_value=lambda x: 1)],
+        )
+        message = TimeSeriesRequest(
+            meta=RequestMeta(
+                project_ids=[1, 2, 3],
+                organization_id=1,
+                cogs_category="something",
+                referrer="something",
+                start_timestamp=Timestamp(seconds=int(BASE_TIME.timestamp())),
+                end_timestamp=Timestamp(
+                    seconds=int(BASE_TIME.timestamp() + query_duration)
+                ),
+                trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN,
+            ),
+            expressions=[
+                Expression(
+                    formula=Expression.BinaryFormula(
+                        op=Expression.BinaryFormula.OP_ADD,
+                        left=Expression(
+                            aggregation=AttributeAggregation(
+                                aggregate=Function.FUNCTION_SUM,
+                                key=AttributeKey(
+                                    type=AttributeKey.TYPE_FLOAT, name="test_metric"
+                                ),
+                                label="sum",
+                                extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
+                            )
+                        ),
+                        right=Expression(
+                            aggregation=AttributeAggregation(
+                                aggregate=Function.FUNCTION_AVG,
+                                key=AttributeKey(
+                                    type=AttributeKey.TYPE_FLOAT, name="test_metric"
+                                ),
+                                label="avg",
+                                extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
+                            )
+                        ),
+                    ),
+                    label="sum + avg",
+                ),
+            ],
+            granularity_secs=granularity_secs,
+        )
+        response = EndpointTimeSeries().execute(message)
+        expected_buckets = [
+            Timestamp(seconds=int(BASE_TIME.timestamp()) + secs)
+            for secs in range(0, query_duration, granularity_secs)
+        ]
+        expected_avg_timeseries = TimeSeries(
+            label="avg",
+            buckets=expected_buckets,
+            data_points=[
+                DataPoint(data=1, data_present=True, sample_count=300)
+                for _ in range(len(expected_buckets))
+            ],
+        )
+        expected_sum_timeseries = TimeSeries(
+            label="sum",
+            buckets=expected_buckets,
+            data_points=[
+                DataPoint(data=300, data_present=True)
+                for _ in range(len(expected_buckets))
+            ],
+        )
+        expected_formula_timeseries = TimeSeries(
+            label="sum + avg",
+            buckets=expected_buckets,
+            data_points=[
+                DataPoint(
+                    data=sum_datapoint.data + avg_datapoint.data,
+                    data_present=True,
+                    sample_count=sum_datapoint.sample_count,
+                )
+                for sum_datapoint, avg_datapoint in zip(
+                    expected_sum_timeseries.data_points,
+                    expected_avg_timeseries.data_points,
+                )
+            ],
+        )
+        assert sorted(response.result_timeseries, key=lambda x: x.label) == [
+            expected_formula_timeseries
+        ]
+
 
 class TestUtils:
     def test_no_duplicate_labels(self) -> None:
diff --git a/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series_extrapolation.py b/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series_extrapolation.py
