feat(metrics): Add equation support to tracemetrics (#112609)

- This adds equation support to the tracemetrics dataset
- I'm really not sure what the extra_conditions are doing here, they add
a top level condition on the metric details, but we're already adding
them per-aggregate anyways? I'm just going to leave this alone (but add
equation support) for now and move on

Author: wmak

src/sentry/search/eap/trace_metrics/config.py

@@ -4,6 +4,7 @@
 from rest_framework.request import Request
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
 
+from sentry.discover.arithmetic import parse_arithmetic
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.columns import ResolvedTraceMetricAggregate, ResolvedTraceMetricFormula
 from sentry.search.eap.resolver import SearchResolver
@@ -41,32 +42,39 @@ def _extra_conditions_from_columns(
         selected_columns: list[str] | None,
         equations: list[str] | None,
     ) -> TraceItemFilter | None:
+        """While we also add the metric conditions inside -if combinators for each aggregate, adding it to the top level
+        where should help clickhouse prune more rows"""
         aggregate_all_metrics = False
         selected_metrics: set[TraceMetric] = set()
+        columns: set[str] = set()
 
         if selected_columns:
             stripped_columns = [column.strip() for column in selected_columns]
             for column in stripped_columns:
                 match = fields.is_function(column)
-                if not match:
-                    continue
+                if match:
+                    columns.add(column)
 
-                resolved_function, _ = search_resolver.resolve_function(column)
+        if equations:
+            for equation in equations:
+                _, _, terms = parse_arithmetic(equation)
+                for term in terms:
+                    columns.add(term)
 
-                if not isinstance(
-                    resolved_function, ResolvedTraceMetricAggregate
-                ) and not isinstance(resolved_function, ResolvedTraceMetricFormula):
-                    continue
+        for column in columns:
+            resolved_function, _ = search_resolver.resolve_function(column)
 
-                if resolved_function.trace_metric is None:
-                    # found an aggregation across all metrics, not just 1
-                    aggregate_all_metrics = True
-                    continue
+            if not isinstance(resolved_function, ResolvedTraceMetricAggregate) and not isinstance(
+                resolved_function, ResolvedTraceMetricFormula
+            ):
+                continue
 
-                selected_metrics.add(resolved_function.trace_metric)
+            if resolved_function.trace_metric is None:
+                # found an aggregation across all metrics, not just 1
+                aggregate_all_metrics = True
+                continue
 
-        if equations:
-            raise InvalidSearchQuery("Cannot support equations on trace metrics yet")
+            selected_metrics.add(resolved_function.trace_metric)
 
         # no selected metrics, no filter needed
         if not selected_metrics:

src/sentry/snuba/trace_metrics.py

@@ -42,7 +42,8 @@ def run_table_query(
         additional_queries: AdditionalQueries | None = None,
     ) -> EAPResponse:
         """timestamp_precise is always displayed in the UI in lieu of timestamp but since the TraceItem table isn't a DateTime64
-        so we need to always order by it regardless of what is actually passed to the orderby."""
+        so we need to always order by it regardless of what is actually passed to the orderby.
+        """
         if (
             orderby is not None
             and len(orderby) == 1
@@ -59,6 +60,7 @@ def run_table_query(
             rpc_dataset_common.TableQuery(
                 query_string=query_string,
                 selected_columns=selected_columns,
+                equations=equations,
                 orderby=orderby,
                 offset=offset,
                 limit=limit,

tests/snuba/api/endpoints/test_organization_events_trace_metrics.py

@@ -6,7 +6,9 @@
 
 from sentry.conf.types.sentry_config import SentryMode
 from sentry.utils.snuba_rpc import table_rpc
-from tests.snuba.api.endpoints.test_organization_events import OrganizationEventsEndpointTestBase
+from tests.snuba.api.endpoints.test_organization_events import (
+    OrganizationEventsEndpointTestBase,
+)
 
 
 class OrganizationEventsTraceMetricsEndpointTest(OrganizationEventsEndpointTestBase):
@@ -283,7 +285,9 @@ def test_per_second_formula_with_gauge_metric_type(self) -> None:
             "per_second(value, cpu_usage, gauge, -)": pytest.approx(2 / 600, abs=0.001)
         }
 
-    def test_per_second_formula_with_gauge_metric_type_without_top_level_metric_type(self) -> None:
+    def test_per_second_formula_with_gauge_metric_type_without_top_level_metric_type(
+        self,
+    ) -> None:
         gauge_metrics = [
             self.create_trace_metric("cpu_usage", 75.0, "gauge"),
             self.create_trace_metric("cpu_usage", 80.0, "gauge"),
@@ -777,6 +781,40 @@ def test_count_if_with_arbitrary_search_term(self):
         assert len(data) == 1
         assert data[0]["count_if(`release:abcdef`,value,request_duration,distribution,none)"] == 1
 
+    def test_count_if_in_equation(self):
+        trace_metrics = [
+            self.create_trace_metric("request_duration", 200.0, "distribution", metric_unit="none"),
+            self.create_trace_metric(
+                "request_duration",
+                200.0,
+                "distribution",
+                metric_unit="none",
+                attributes={"release": "abcdef"},
+            ),
+        ]
+        self.store_eap_items(trace_metrics)
+
+        response = self.do_request(
+            {
+                "field": [
+                    "equation|count_if(`release:abcdef`,value,request_duration,distribution,none) + count_if(`!release:abcdef`,value,request_duration,distribution,none)",
+                ],
+                "project": self.project.id,
+                "dataset": self.dataset,
+                "statsPeriod": "10m",
+            }
+        )
+        assert response.status_code == 200, response.content
+        data = response.data["data"]
+
+        assert len(data) == 1
+        assert (
+            data[0][
+                "equation|count_if(`release:abcdef`,value,request_duration,distribution,none) + count_if(`!release:abcdef`,value,request_duration,distribution,none)"
+            ]
+            == 2
+        )
+
     def test_p50_if_with_arbitrary_search_term(self):
         trace_metrics = [
             self.create_trace_metric(