Repository: getsentry/sentry

Issue:
Implement a rule converter for `count_if`
For on-demand metrics, we need to translate conditions from `count_if` select to Relay rule conditions. 
```[tasklist]
### Tasks
- [ ] https://github.com/getsentry/sentry/pull/52784
```



Discussion:

Resolving patch:
diff --git a/src/sentry/search/events/builder/metrics.py b/src/sentry/search/events/builder/metrics.py
--- a/src/sentry/search/events/builder/metrics.py
+++ b/src/sentry/search/events/builder/metrics.py
@@ -97,15 +97,15 @@ def __init__(
     def _resolve_on_demand_spec(
         self, dataset: Optional[Dataset], selected_cols: List[Optional[str]], query: str
     ) -> Optional[OndemandMetricSpec]:
-        if not is_on_demand_query(dataset, query):
-            return None
-
         field = selected_cols[0] if selected_cols else None
         if not self.is_alerts_query or not field:
             return None
+
+        if not is_on_demand_query(dataset, field, query):
+            return None
+
         try:
             return OndemandMetricSpec(field, query)
-
         except Exception as e:
             sentry_sdk.capture_exception(e)
 
diff --git a/src/sentry/snuba/metrics/extraction.py b/src/sentry/snuba/metrics/extraction.py
--- a/src/sentry/snuba/metrics/extraction.py
+++ b/src/sentry/snuba/metrics/extraction.py
@@ -4,10 +4,10 @@
 from typing import (
     Any,
     Dict,
+    List,
     Literal,
     Optional,
     Sequence,
-    Tuple,
     Type,
     TypedDict,
     TypeVar,
@@ -85,10 +85,21 @@
     ">=": "gte",
 }
 
+# Maps from parsed count_if condition args to Relay rule condition operators.
+_COUNTIF_TO_RELAY_OPERATORS: Dict[str, "CompareOp"] = {
+    "equals": "eq",
+    "notEquals": "eq",
+    "less": "lt",
+    "greater": "gt",
+    "lessOrEquals": "lte",
+    "greaterOrEquals": "gte",
+}
+
 # Maps plain Discover functions to metric aggregation functions. Derived metrics
 # are not part of this mapping.
-_SEARCH_TO_METRIC_AGGREGATES: Dict[str, Optional[MetricOperationType]] = {
+_SEARCH_TO_METRIC_AGGREGATES: Dict[str, MetricOperationType] = {
     "count": "sum",
+    "count_if": "sum",
     "avg": "avg",
     "max": "max",
     "p50": "p50",
@@ -101,7 +112,7 @@
 # Mapping to infer metric type from Discover function.
 _AGGREGATE_TO_METRIC_TYPE = {
     "count": "c",
-    # TODO(ogi): support count_if
+    "count_if": "c",
     "avg": "d",
     "max": "d",
     "p50": "d",
@@ -146,7 +157,7 @@ class LogicalRuleCondition(TypedDict):
     """RuleCondition that applies a logical operator to a sequence of conditions."""
 
     op: Literal["and", "or"]
-    inner: Sequence[RuleCondition]
+    inner: List[RuleCondition]
 
 
 class NotRuleCondition(TypedDict):
@@ -197,27 +208,49 @@ class MetricSpec(TypedDict):
 
 def is_on_demand_snuba_query(snuba_query: SnubaQuery) -> bool:
     """Returns ``True`` if the snuba query can't be supported by standard metrics."""
+    return is_on_demand_query(snuba_query.dataset, snuba_query.aggregate, snuba_query.query)
 
-    return is_on_demand_query(snuba_query.dataset, snuba_query.query)
 
-
-def is_on_demand_query(dataset: Optional[Union[str, Dataset]], query: Optional[str]) -> bool:
+def is_on_demand_query(
+    dataset: Optional[Union[str, Dataset]], aggregate: str, query: Optional[str]
+) -> bool:
     """Returns ``True`` if the dataset is performance metrics and query contains non-standard search fields."""
 
-    if not dataset or not query:
+    if not dataset or Dataset(dataset) != Dataset.PerformanceMetrics:
         return False
 
-    if Dataset(dataset) != Dataset.PerformanceMetrics:
-        return False
+    for field in _get_aggregate_fields(aggregate):
+        if not _is_standard_metrics_field(field):
+            return True
 
     try:
-        return not _is_standard_metrics_compatible(event_search.parse_search_query(query))
+        return not _is_standard_metrics_query(event_search.parse_search_query(query or ""))
     except InvalidSearchQuery:
         logger.error(f"Failed to parse search query: {query}", exc_info=True)
         return False
 
 
-def _is_standard_metrics_compatible(tokens: Sequence[QueryToken]) -> bool:
+def _get_aggregate_fields(aggregate: str) -> Sequence[str]:
+    """
+    Returns any fields referenced by the arguments of supported aggregate
+    functions, otherwise ``None``.
+    """
+
+    # count_if is currently the only supported function, exit early
+    if not aggregate.startswith("count_if("):
+        return []
+
+    try:
+        function, arguments, _ = fields.parse_function(aggregate)
+        if function == "count_if" and arguments:
+            return [arguments[0]]
+    except InvalidSearchQuery:
+        logger.error(f"Failed to parse aggregate: {aggregate}", exc_info=True)
+
+    return []
+
+
+def _is_standard_metrics_query(tokens: Sequence[QueryToken]) -> bool:
     """
     Recursively checks if any of the supplied token contain search filters that can't be handled by standard metrics.
     """
@@ -231,19 +264,35 @@ def _is_standard_metrics_compatible(tokens: Sequence[QueryToken]) -> bool:
 
 def _is_standard_metrics_search_filter(token: QueryToken) -> bool:
     if isinstance(token, SearchFilter):
-        return token.key.name in _STANDARD_METRIC_FIELDS
+        return _is_standard_metrics_field(token.key.name)
 
     if isinstance(token, ParenExpression):
-        return _is_standard_metrics_compatible(token.children)
+        return _is_standard_metrics_query(token.children)
 
     return True
 
 
+def _is_standard_metrics_field(field: str) -> bool:
+    return field in _STANDARD_METRIC_FIELDS
+
+
 class OndemandMetricSpec:
     """
     Contains the information required to query or extract an on-demand metric.
     """
 
+    # The data type of the metric to extract.
+    metric_type: str
+    # The payload field to extract the metric value from. Empty for counters.
+    field: Optional[str]
+    # The aggregation to execute on the metric.
+    op: MetricOperationType
+
+    # The original query used to construct the metric spec.
+    _query: str
+    # Rule condition parsed from the aggregate field expression.
+    _field_condition: Optional[RuleCondition]
+
     def __init__(self, field: str, query: str):
         """
         Parses a selected column and query into an on demand metric spec containing the MRI, field and op.
@@ -258,12 +307,39 @@ def __init__(self, field: str, query: str):
         # On-demand metrics are implicitly transaction metrics. Remove the
         # filter from the query since it can't be translated to a RuleCondition.
         self._query = re.sub(r"event\.type:transaction\s*", "", query)
+        self._init_aggregate(field)
+
+    def _init_aggregate(self, aggregate: str) -> None:
+        """
+        Extracts the field name, metric type and metric operation from a Discover
+        function call.
+
+        This does not support derived metrics such as ``apdex``.
+        """
+
+        # TODO: Add support for derived metrics: failure_rate, apdex, eps, epm, tps, tpm
+        function, arguments, _alias = fields.parse_function(aggregate)
+        assert (
+            function in _AGGREGATE_TO_METRIC_TYPE and function in _SEARCH_TO_METRIC_AGGREGATES
+        ), f"Unsupported aggregate function {function}"
+        self.metric_type = _AGGREGATE_TO_METRIC_TYPE[function]
+        self.op = _SEARCH_TO_METRIC_AGGREGATES[function]
+
+        self.field = None
+        self._field_condition = None
 
-        relay_field, metric_type, op = _extract_field_info(field)
-        self.field = relay_field
-        self.metric_type = metric_type
-        self.mri = f"{metric_type}:{CUSTOM_ALERT_METRIC_NAME}@none"
-        self.op = op
+        if self.metric_type != "c":
+            assert len(arguments) == 1, "Only one parameter is supported"
+            self.field = _map_field_name(arguments[0])
+
+        if function == "count_if":
+            key, op, value = arguments
+            self._field_condition = _convert_countif_filter(key, op, value)
+
+    @property
+    def mri(self) -> str:
+        """The unique identifier of the on-demand metric."""
+        return f"{self.metric_type}:{CUSTOM_ALERT_METRIC_NAME}@none"
 
     def query_hash(self) -> str:
         """Returns a hash of the query and field to be used as a unique identifier for the on-demand metric."""
@@ -276,32 +352,35 @@ def condition(self) -> RuleCondition:
         """Returns a condition that should be fulfilled for the on-demand metric to be extracted."""
 
         tokens = event_search.parse_search_query(self._query)
-        assert tokens, "This query should not use on demand metrics"
-        return SearchQueryConverter(tokens).convert()
+        if not tokens:
+            assert self._field_condition is not None, "This query should not use on demand metrics"
+            return self._field_condition
 
+        condition = SearchQueryConverter(tokens).convert()
+        if not self._field_condition:
+            return condition
 
-def _extract_field_info(aggregate: str) -> Tuple[Optional[str], str, MetricOperationType]:
-    """
-    Extracts the field name, metric type and metric operation from a Discover
-    function call.
+        if condition["op"] != "and":
+            return {"op": "and", "inner": [condition, self._field_condition]}
 
-    This does not support derived metrics such as ``apdex`` and aggregates with
-    filters (``count_if``).
-    """
-    name, arguments, _alias = fields.parse_function(aggregate)
+        condition["inner"].append(self._field_condition)
+        return condition
+
+
+def _convert_countif_filter(key: str, op: str, value: str) -> RuleCondition:
+    """Maps ``count_if`` arguments to a ``RuleCondition``."""
+    assert op in _COUNTIF_TO_RELAY_OPERATORS, f"Unsupported `count_if` operator {op}"
 
-    # TODO: Add support for derived metrics: failure_rate, apdex, eps, epm, tps, tpm
+    condition: RuleCondition = {
+        "op": _COUNTIF_TO_RELAY_OPERATORS[op],
+        "name": _map_field_name(key),
+        "value": fields.normalize_count_if_value({"column": key, "value": value}),
+    }
 
-    metric_type = _AGGREGATE_TO_METRIC_TYPE.get(name)
-    metric_op = _SEARCH_TO_METRIC_AGGREGATES.get(name)
-    assert metric_type and metric_op, f"Unsupported aggregate function {name}"
+    if op == "notEquals":
+        condition = {"op": "not", "inner": condition}
 
-    if metric_type == "c":
-        assert not arguments, "`count()` does not support arguments"
-        return None, metric_type, metric_op
-    else:
-        assert len(arguments) == 1, "Only one parameter is supported"
-        return _map_field_name(arguments[0]), metric_type, metric_op
+    return condition
 
 
 def _map_field_name(search_key: str) -> str:


Test patch:
diff --git a/tests/sentry/relay/config/test_metric_extraction.py b/tests/sentry/relay/config/test_metric_extraction.py
--- a/tests/sentry/relay/config/test_metric_extraction.py
+++ b/tests/sentry/relay/config/test_metric_extraction.py
@@ -6,15 +6,13 @@
 from sentry.snuba.models import SnubaQuery
 
 
-def create_alert(query: str):
+def test_empty_query():
     snuba_query = SnubaQuery(
-        aggregate="p75(measurements.fp)", query=query, dataset=Dataset.PerformanceMetrics.value
+        aggregate="p75(measurements.fp)",
+        query="",
+        dataset=Dataset.PerformanceMetrics.value,
     )
-    return AlertRule(snuba_query=snuba_query)
-
-
-def test_empty_query():
-    alert = create_alert("")
+    alert = AlertRule(snuba_query=snuba_query)
 
     assert convert_query_to_metric(alert.snuba_query) is None
 
@@ -28,119 +26,10 @@ def test_simple_query_count():
     alert = AlertRule(snuba_query=snuba_query)
 
     metric = convert_query_to_metric(alert.snuba_query)
-
-    expected = {
+    assert metric == {
         "category": "transaction",
         "condition": {"name": "event.duration", "op": "gte", "value": 1000.0},
         "field": None,
         "mri": "c:transactions/on_demand@none",
         "tags": [{"key": "query_hash", "value": ANY}],
     }
-
-    assert metric == expected
-
-
-def test_simple_query():
-    alert = create_alert("transaction.duration:>=1000")
-    metric = convert_query_to_metric(alert.snuba_query)
-
-    expected = {
-        "category": "transaction",
-        "condition": {"name": "event.duration", "op": "gte", "value": 1000.0},
-        "field": "event.measurements.fp",
-        "mri": "d:transactions/on_demand@none",
-        "tags": [{"key": "query_hash", "value": ANY}],
-    }
-
-    assert metric == expected
-
-
-def test_or_boolean_condition():
-    alert = create_alert("transaction.duration:>=100 OR transaction.duration:<1000")
-    metric = convert_query_to_metric(alert.snuba_query)
-
-    expected = {
-        "category": "transaction",
-        "condition": {
-            "inner": [
-                {"name": "event.duration", "op": "gte", "value": 100.0},
-                {"name": "event.duration", "op": "lt", "value": 1000.0},
-            ],
-            "op": "or",
-        },
-        "field": "event.measurements.fp",
-        "mri": "d:transactions/on_demand@none",
-        "tags": [{"key": "query_hash", "value": ANY}],
-    }
-
-    assert metric == expected
-
-
-def test_and_boolean_condition():
-    alert = create_alert("release:foo transaction.duration:<10s")
-    metric = convert_query_to_metric(alert.snuba_query)
-
-    expected = {
-        "category": "transaction",
-        "condition": {
-            "inner": [
-                {"name": "event.release", "op": "eq", "value": "foo"},
-                {"name": "event.duration", "op": "lt", "value": 10000.0},
-            ],
-            "op": "and",
-        },
-        "field": "event.measurements.fp",
-        "mri": "d:transactions/on_demand@none",
-        "tags": [{"key": "query_hash", "value": ANY}],
-    }
-
-    assert metric == expected
-
-
-def test_nested_conditions():
-    query = "(release:=a OR transaction.op:=b) transaction.duration:>1s"
-    metric = convert_query_to_metric(create_alert(query).snuba_query)
-
-    expected = {
-        "category": "transaction",
-        "condition": {
-            "op": "and",
-            "inner": [
-                {
-                    "op": "or",
-                    "inner": [
-                        {"name": "event.release", "op": "eq", "value": "a"},
-                        {"name": "event.contexts.trace.op", "op": "eq", "value": "b"},
-                    ],
-                },
-                {"name": "event.duration", "op": "gt", "value": 1000.0},
-            ],
-        },
-        "field": "event.measurements.fp",
-        "mri": "d:transactions/on_demand@none",
-        "tags": [{"key": "query_hash", "value": ANY}],
-    }
-
-    assert metric == expected
-
-
-def test_wildcard_condition():
-    # transaction.duration is required for convert_query_to_metric
-    alert = create_alert("release.version:1.* transaction.duration:1s")
-    metric = convert_query_to_metric(alert.snuba_query)
-
-    expected = {
-        "category": "transaction",
-        "condition": {
-            "op": "and",
-            "inner": [
-                {"name": "event.release.version.short", "op": "glob", "value": ["1.*"]},
-                {"name": "event.duration", "op": "eq", "value": 1000.0},
-            ],
-        },
-        "field": "event.measurements.fp",
-        "mri": "d:transactions/on_demand@none",
-        "tags": [{"key": "query_hash", "value": ANY}],
-    }
-
-    assert metric == expected
diff --git a/tests/sentry/snuba/test_extraction.py b/tests/sentry/snuba/test_extraction.py
--- a/tests/sentry/snuba/test_extraction.py
+++ b/tests/sentry/snuba/test_extraction.py
@@ -1,34 +1,40 @@
 from sentry.snuba.dataset import Dataset
-from sentry.snuba.metrics.extraction import is_on_demand_query
+from sentry.snuba.metrics.extraction import OndemandMetricSpec, is_on_demand_query
 
 
 def test_is_on_demand_query_wrong_dataset():
-    assert is_on_demand_query(Dataset.Transactions, "geo.city:=Vienna") is False
-    assert is_on_demand_query(Dataset.Metrics, "browser.version:=1 os.name:=android") is False
+    assert is_on_demand_query(Dataset.Transactions, "count()", "geo.city:Vienna") is False
+    assert (
+        is_on_demand_query(Dataset.Metrics, "count()", "browser.version:1 os.name:android") is False
+    )
 
 
 def test_is_on_demand_query_no_query():
-    assert is_on_demand_query(Dataset.PerformanceMetrics, "") is False
+    assert is_on_demand_query(Dataset.PerformanceMetrics, "count()", "") is False
 
 
 def test_is_on_demand_query_invalid_query():
-    assert is_on_demand_query(Dataset.PerformanceMetrics, "AND") is False
-    assert is_on_demand_query(Dataset.PerformanceMetrics, "(AND transaction.duration:>=1") is False
-    assert is_on_demand_query(Dataset.PerformanceMetrics, "transaction.duration:>=abc") is False
+    dataset = Dataset.PerformanceMetrics
+
+    assert is_on_demand_query(dataset, "count()", "AND") is False
+    assert is_on_demand_query(dataset, "count()", "(AND transaction.duration:>=1") is False
+    assert is_on_demand_query(dataset, "count()", "transaction.duration:>=abc") is False
 
 
 def test_is_on_demand_query_true():
     dataset = Dataset.PerformanceMetrics
 
     # transaction.duration is a non-standard field
-    assert is_on_demand_query(dataset, "transaction.duration:>=1") is True
+    assert is_on_demand_query(dataset, "count()", "transaction.duration:>=1") is True
     # transaction.duration is a non-standard field
-    assert is_on_demand_query(dataset, "geo.city:=Vienna") is True
+    assert is_on_demand_query(dataset, "count()", "geo.city:Vienna") is True
     # os.name is a standard field, browser.version is not
-    assert is_on_demand_query(dataset, "browser.version:=1 os.name:=android") is True
+    assert is_on_demand_query(dataset, "count()", "browser.version:1 os.name:android") is True
     # os.version is not a standard field
     assert (
-        is_on_demand_query(dataset, "(release:=a OR transaction.op:=b) transaction.duration:>1s")
+        is_on_demand_query(
+            dataset, "count()", "(release:a OR transaction.op:b) transaction.duration:>1s"
+        )
         is True
     )
 
@@ -36,12 +42,136 @@ def test_is_on_demand_query_true():
 def test_is_on_demand_query_false():
     dataset = Dataset.PerformanceMetrics
 
-    assert is_on_demand_query(dataset, "") is False
-    assert is_on_demand_query(dataset, "environment:=dev") is False
-    assert is_on_demand_query(dataset, "release:=initial OR os.name:=android") is False
+    assert is_on_demand_query(dataset, "count()", "") is False
+    assert is_on_demand_query(dataset, "count()", "environment:dev") is False
+    assert is_on_demand_query(dataset, "count()", "release:initial OR os.name:android") is False
     assert (
         is_on_demand_query(
-            dataset, "(http.method:=POST OR http.status_code:=404) browser.name:=chrome"
+            dataset, "count()", "(http.method:POST OR http.status_code:404) browser.name:chrome"
         )
         is False
     )
+
+
+def test_is_on_demand_query_countif():
+    dataset = Dataset.PerformanceMetrics
+
+    assert is_on_demand_query(dataset, "count_if(transaction.duration,equals,300)", "") is True
+    assert is_on_demand_query(dataset, 'count_if(release,equals,"foo")', "") is False
+
+
+def test_spec_simple_query_count():
+    spec = OndemandMetricSpec("count()", "transaction.duration:>1s")
+
+    assert spec.metric_type == "c"
+    assert spec.field is None
+    assert spec.op == "sum"
+    assert spec.condition() == {"name": "event.duration", "op": "gt", "value": 1000.0}
+
+
+def test_spec_simple_query_distribution():
+    spec = OndemandMetricSpec("p75(measurements.fp)", "transaction.duration:>1s")
+
+    assert spec.metric_type == "d"
+    assert spec.field == "event.measurements.fp"
+    assert spec.op == "p75"
+    assert spec.condition() == {"name": "event.duration", "op": "gt", "value": 1000.0}
+
+
+def test_spec_or_condition():
+    spec = OndemandMetricSpec("count()", "transaction.duration:>=100 OR transaction.duration:<1000")
+
+    assert spec.condition() == {
+        "inner": [
+            {"name": "event.duration", "op": "gte", "value": 100.0},
+            {"name": "event.duration", "op": "lt", "value": 1000.0},
+        ],
+        "op": "or",
+    }
+
+
+def test_spec_and_condition():
+    spec = OndemandMetricSpec("count()", "release:foo transaction.duration:<10s")
+    assert spec.condition() == {
+        "inner": [
+            {"name": "event.release", "op": "eq", "value": "foo"},
+            {"name": "event.duration", "op": "lt", "value": 10000.0},
+        ],
+        "op": "and",
+    }
+
+
+def test_spec_nested_condition():
+    spec = OndemandMetricSpec("count()", "(release:a OR transaction.op:b) transaction.duration:>1s")
+    assert spec.condition() == {
+        "op": "and",
+        "inner": [
+            {
+                "op": "or",
+                "inner": [
+                    {"name": "event.release", "op": "eq", "value": "a"},
+                    {"name": "event.contexts.trace.op", "op": "eq", "value": "b"},
+                ],
+            },
+            {"name": "event.duration", "op": "gt", "value": 1000.0},
+        ],
+    }
+
+
+def test_spec_boolean_precedence():
+    spec = OndemandMetricSpec("count()", "release:a OR transaction.op:b transaction.duration:>1s")
+    assert spec.condition() == {
+        "op": "or",
+        "inner": [
+            {"name": "event.release", "op": "eq", "value": "a"},
+            {
+                "op": "and",
+                "inner": [
+                    {"name": "event.contexts.trace.op", "op": "eq", "value": "b"},
+                    {"name": "event.duration", "op": "gt", "value": 1000.0},
+                ],
+            },
+        ],
+    }
+
+
+def test_spec_wildcard():
+    spec = OndemandMetricSpec("count()", "release.version:1.*")
+    assert spec.condition() == {
+        "name": "event.release.version.short",
+        "op": "glob",
+        "value": ["1.*"],
+    }
+
+
+def test_spec_countif():
+    spec = OndemandMetricSpec("count_if(transaction.duration,equals,300)", "")
+
+    assert spec.metric_type == "c"
+    assert spec.field is None
+    assert spec.op == "sum"
+    assert spec.condition() == {
+        "name": "event.duration",
+        "op": "eq",
+        "value": 300.0,
+    }
+
+
+def test_spec_countif_with_query():
+    spec = OndemandMetricSpec(
+        "count_if(transaction.duration,equals,300)", "release:a OR transaction.op:b"
+    )
+
+    assert spec.condition() == {
+        "op": "and",
+        "inner": [
+            {
+                "op": "or",
+                "inner": [
+                    {"name": "event.release", "op": "eq", "value": "a"},
+                    {"name": "event.contexts.trace.op", "op": "eq", "value": "b"},
+                ],
+            },
+            {"name": "event.duration", "op": "eq", "value": 300.0},
+        ],
+    }