feat(upsampling) - Upsample Error counts for ACI Event Frequency Alerts (#97517)

- Thread project_ids through ACI Event Frequency handlers to enable
Snuba upsampling.
- Add tests: upsampling enabled (10) vs disabled (2).
- Ended up adding project_ids to all tsdb_functions to avoid the ambiguity in code, at the expense of adding it where it isnt relevant. We have other parameters that follow this pattern there, so it seemed the lesser evil.

Author: yuvmen

src/sentry/rules/conditions/event_frequency.py

@@ -351,34 +351,19 @@ def get_snuba_query_result(
         group_on_time: bool = False,
         project_ids: list[int] | None = None,
     ) -> Mapping[int, int]:
-        kwargs = {
-            "model": model,
-            "keys": keys,
-            "start": start,
-            "end": end,
-            "environment_id": environment_id,
-            "use_cache": True,
-            "jitter_value": group_id,
-            "tenant_ids": {"organization_id": organization_id},
-            "referrer_suffix": referrer_suffix,
-            "group_on_time": group_on_time,
-        }
-
-        # Try to pass project_ids if provided, but fall back gracefully if not supported
-        result: Mapping[int, int]
-        if project_ids is not None:
-            try:
-                kwargs["project_ids"] = project_ids
-                result = tsdb_function(**kwargs)
-            except TypeError as e:
-                if "project_ids" in str(e):
-                    # Function doesn't support project_ids, try without it
-                    kwargs.pop("project_ids", None)
-                    result = tsdb_function(**kwargs)
-                else:
-                    raise
-        else:
-            result = tsdb_function(**kwargs)
+        result: Mapping[int, int] = tsdb_function(
+            model=model,
+            keys=keys,
+            start=start,
+            end=end,
+            environment_id=environment_id,
+            use_cache=True,
+            jitter_value=group_id,
+            tenant_ids={"organization_id": organization_id},
+            referrer_suffix=referrer_suffix,
+            group_on_time=group_on_time,
+            project_ids=project_ids,
+        )
         return result
 
     def get_chunked_result(

src/sentry/tsdb/base.py

@@ -588,6 +588,7 @@ def get_distinct_counts_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[int, list[tuple[int, Any]]]:
         """
         Fetch counts of distinct items for each rollup interval within the range.
@@ -608,6 +609,7 @@ def get_distinct_counts_totals(
         referrer_suffix: str | None = None,
         conditions: list[SnubaCondition] | None = None,
         group_on_time: bool = False,
+        project_ids: Sequence[int] | None = None,
     ) -> Mapping[TSDBKey, int]:
         """
         Count distinct items during a time range with optional conditions
@@ -665,6 +667,7 @@ def get_frequency_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]:
         """
         Retrieve the frequency of known items in a table over time.

src/sentry/tsdb/dummy.py

@@ -58,6 +58,7 @@ def get_distinct_counts_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[int, list[tuple[int, Any]]]:
         self.validate_arguments([model], [environment_id])
         _, series = self.get_optimal_rollup_series(start, end, rollup)
@@ -77,6 +78,7 @@ def get_distinct_counts_totals(
         referrer_suffix=None,
         conditions=None,
         group_on_time: bool = False,
+        project_ids: Sequence[int] | None = None,
     ):
         self.validate_arguments([model], [environment_id])
         return {k: 0 for k in keys}
@@ -108,6 +110,7 @@ def get_frequency_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]:
         self.validate_arguments([model], [environment_id])
         rollup, series = self.get_optimal_rollup_series(start, end, rollup)

src/sentry/tsdb/redis.py

@@ -501,6 +501,7 @@ def get_distinct_counts_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[int, list[tuple[int, Any]]]:
         """
         Fetch counts of distinct items for each rollup interval within the range.
@@ -542,6 +543,7 @@ def get_distinct_counts_totals(
         referrer_suffix: str | None = None,
         conditions: list[SnubaCondition] | None = None,
         group_on_time: bool = False,
+        project_ids: Sequence[int] | None = None,
     ) -> Mapping[TSDBKey, int]:
         """
         Count distinct items during a time range.
@@ -753,6 +755,7 @@ def get_frequency_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]:
         self.validate_arguments([model], [environment_id])
 

src/sentry/tsdb/snuba.py

@@ -839,6 +839,7 @@ def get_distinct_counts_series(
         rollup=None,
         environment_id=None,
         tenant_ids=None,
+        project_ids: Sequence[int] | None = None,
     ):
         result = self.get_data(
             model,
@@ -871,6 +872,7 @@ def get_distinct_counts_totals(
         referrer_suffix=None,
         conditions=None,
         group_on_time: bool = False,
+        project_ids: Sequence[int] | None = None,
     ) -> Mapping[TSDBKey, int]:
         return self.get_data(
             model,
@@ -897,6 +899,7 @@ def get_frequency_series(
         rollup: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
+        project_ids: Sequence[int] | None = None,
     ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]:
         result = self.get_data(
             model,

src/sentry/workflow_engine/handlers/condition/event_frequency_query_handlers.py

@@ -55,6 +55,7 @@ def __call__(
         referrer_suffix: str | None = None,
         conditions: list[SnubaCondition] | None = None,
         group_on_time: bool = False,
+        project_ids: list[int] | None = None,
     ) -> Mapping[TSDBKey, int]: ...
 
 
@@ -101,6 +102,7 @@ def get_snuba_query_result(
         referrer_suffix: str,
         conditions: list[SnubaCondition] | None = None,
         group_on_time: bool = False,
+        project_ids: list[int] | None = None,
     ) -> Mapping[int, int]:
         result: Mapping[int, int] = tsdb_function(
             model=model,
@@ -114,6 +116,7 @@ def get_snuba_query_result(
             referrer_suffix=referrer_suffix,
             conditions=conditions,
             group_on_time=group_on_time,
+            project_ids=project_ids,
         )
         return result
 
@@ -129,6 +132,7 @@ def get_chunked_result(
         referrer_suffix: str,
         filters: list[QueryFilter] | None = None,
         group_on_time: bool = False,
+        project_ids: list[int] | None = None,
     ) -> dict[int, int]:
         batch_totals: dict[int, int] = defaultdict(int)
         group_id = group_ids[0]
@@ -146,6 +150,7 @@ def get_chunked_result(
                 referrer_suffix=referrer_suffix,
                 conditions=conditions,
                 group_on_time=group_on_time,
+                project_ids=project_ids,
             )
             batch_totals.update(result)
         return batch_totals
@@ -332,6 +337,8 @@ def batch_query(
         batch_sums: QueryResult = defaultdict(int)
         category_group_ids = self.get_group_ids_by_category(groups)
         organization_id = self.get_value_from_groups(groups, "project__organization_id")
+        # Build project_ids list from incoming groups
+        project_ids = list({g["project_id"] for g in groups}) if groups else []
 
         if not organization_id:
             return batch_sums
@@ -350,6 +357,7 @@ def batch_query(
                     referrer_suffix="wf_batch_alert_event_frequency",
                     filters=filters,
                     group_on_time=False,
+                    project_ids=project_ids,
                 )
             except InvalidFilter:
                 # Filter is not supported for this issue type

tests/sentry/workflow_engine/handlers/condition/test_event_frequency_query_handlers.py

@@ -42,6 +42,86 @@ def test_batch_query(self) -> None:
         )
         assert batch_query == {self.event3.group_id: 1}
 
+    def test_batch_query_with_upsampling_enabled_counts_upsampled(self) -> None:
+        # Create two sampled error events in a dedicated group
+        event_a = self.store_event(
+            data={
+                "event_id": "d" * 32,
+                "environment": self.environment.name,
+                "timestamp": before_now(seconds=20).isoformat(),
+                "fingerprint": ["upsampled-group"],
+                "contexts": {"error_sampling": {"client_sample_rate": 0.2}},
+                "exception": {"values": [{"type": "ValueError", "value": "a"}]},
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": "e" * 32,
+                "environment": self.environment.name,
+                "timestamp": before_now(seconds=10).isoformat(),
+                "fingerprint": ["upsampled-group"],
+                "contexts": {"error_sampling": {"client_sample_rate": 0.2}},
+                "exception": {"values": [{"type": "ValueError", "value": "b"}]},
+            },
+            project_id=self.project.id,
+        )
+
+        groups = list(
+            Group.objects.filter(id=event_a.group_id).values(
+                "id", "type", "project_id", "project__organization_id"
+            )
+        )
+
+        with self.options({"issues.client_error_sampling.project_allowlist": [self.project.id]}):
+            batch_query = self.handler().batch_query(
+                groups=groups,
+                start=self.start,
+                end=self.end,
+                environment_id=self.environment.id,
+            )
+        # Expect 2 events upsampled by 5x => 10
+        assert batch_query[event_a.group_id] == 10
+
+    def test_batch_query_without_upsampling_counts_raw(self) -> None:
+        # Same setup as above but without allowlist; expect raw count of 2
+        event_a = self.store_event(
+            data={
+                "event_id": "f" * 32,
+                "environment": self.environment.name,
+                "timestamp": before_now(seconds=20).isoformat(),
+                "fingerprint": ["upsampled-group-raw"],
+                "contexts": {"error_sampling": {"client_sample_rate": 0.2}},
+                "exception": {"values": [{"type": "ValueError", "value": "a"}]},
+            },
+            project_id=self.project.id,
+        )
+        self.store_event(
+            data={
+                "event_id": "1" * 32,
+                "environment": self.environment.name,
+                "timestamp": before_now(seconds=10).isoformat(),
+                "fingerprint": ["upsampled-group-raw"],
+                "contexts": {"error_sampling": {"client_sample_rate": 0.2}},
+                "exception": {"values": [{"type": "ValueError", "value": "b"}]},
+            },
+            project_id=self.project.id,
+        )
+
+        groups = list(
+            Group.objects.filter(id=event_a.group_id).values(
+                "id", "type", "project_id", "project__organization_id"
+            )
+        )
+
+        batch_query = self.handler().batch_query(
+            groups=groups,
+            start=self.start,
+            end=self.end,
+            environment_id=self.environment.id,
+        )
+        assert batch_query[event_a.group_id] == 2
+
     def test_batch_query__tag_conditions__equal(self) -> None:
         batch_query = self.handler().batch_query(
             groups=self.groups,