feat(ddm): Add support for count_unique in metrics layer (#61138)

src/sentry/snuba/metrics_layer/query.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import replace
 from datetime import datetime
-from typing import Any, Mapping, Union
+from typing import Any, List, Mapping, Union, cast
 
 from snuba_sdk import (
     AliasedExpression,
@@ -15,6 +15,7 @@ from snuba_sdk import (
     Request,
     Timeseries,
 )
+from snuba_sdk.formula import FormulaParameterGroup
 
 from sentry.api.utils import InvalidParams
 from sentry.sentry_metrics.use_case_id_registry import UseCaseID
@@ -40,6 +41,7 @@ AGGREGATE_ALIASES = {
     "p90": ("quantiles", [0.9]),
     "p95": ("quantiles", [0.95]),
     "p99": ("quantiles", [0.99]),
+    "count_unique": ("uniq", None),
 }
 
 
@@ -412,15 +414,17 @@ def _resolve_aggregate_aliases(metrics_query: MetricsQuery) -> MetricsQuery:
             aggregate, parameters = AGGREGATE_ALIASES[series.aggregate]
             series = series.set_aggregate(aggregate, parameters)
             return metrics_query.set_query(series)
-
     elif isinstance(metrics_query.query, Formula):
-        parameters = metrics_query.query.parameters
-        for i, p in enumerate(parameters):
+        if not metrics_query.query.parameters:
+            return metrics_query
+
+        aliased_parameters = cast(List[FormulaParameterGroup], metrics_query.query.parameters)
+        for i, p in enumerate(aliased_parameters):
             if isinstance(p, Timeseries):
                 if p.aggregate in AGGREGATE_ALIASES:
-                    aggregate, parameters = AGGREGATE_ALIASES[p.aggregate]
-                    parameters[i] = p.set_aggregate(aggregate, parameters)
+                    new_aggregate, new_parameters = AGGREGATE_ALIASES[p.aggregate]
+                    aliased_parameters[i] = p.set_aggregate(new_aggregate, new_parameters)
 
-        return metrics_query.set_query(metrics_query.query.set_parameters(parameters))
+        return metrics_query.set_query(metrics_query.query.set_parameters(aliased_parameters))
 
     return metrics_query

tests/sentry/sentry_metrics/querying/test_api.py

@@ -352,6 +352,39 @@ class MetricsAPITestCase(TestCase, BaseMetricsTestCase):
                 referrer="metrics.data.api",
             )
 
+    def test_query_with_custom_set(self):
+        mri = "s:custom/user_click@none"
+        for user in ("marco", "marco", "john"):
+            self.store_metric(
+                self.project.organization.id,
+                self.project.id,
+                "set",
+                mri,
+                {},
+                self.ts(self.now()),
+                user,
+                UseCaseID.CUSTOM,
+            )
+
+        field = f"count_unique({mri})"
+        results = run_metrics_query(
+            fields=[field],
+            query=None,
+            group_bys=None,
+            start=self.now() - timedelta(minutes=30),
+            end=self.now() + timedelta(hours=1, minutes=30),
+            interval=3600,
+            organization=self.project.organization,
+            projects=[self.project],
+            environments=[],
+            referrer="metrics.data.api",
+        )
+        groups = results["groups"]
+        assert len(groups) == 1
+        assert groups[0]["by"] == {}
+        assert groups[0]["series"] == {field: [0, 2, 0]}
+        assert groups[0]["totals"] == {field: 2}
+
     @pytest.mark.skip(reason="sessions are not supported in the new metrics layer")
     def test_with_sessions(self) -> None:
         self.store_session(