ref(rules): Add batch_query_hook to EventFrequencyCondition (#68559)

Add a `batch_query_hook` to the `EventFrequencyCondition` so that when
we process "slow" rule conditions we can batch the queries by group ids.

Partially closes
getsentry/team-core-product-foundations#239

Author: ceorourke

src/sentry/rules/conditions/event_frequency.py

@@ -4,6 +4,7 @@
 import contextlib
 import logging
 import re
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timedelta
 from typing import Any
@@ -15,15 +16,19 @@
 from sentry import release_health, tsdb
 from sentry.eventstore.models import GroupEvent
 from sentry.issues.constants import get_issue_tsdb_group_model, get_issue_tsdb_user_group_model
+from sentry.issues.grouptype import GroupCategory
+from sentry.models.group import Group
 from sentry.receivers.rules import DEFAULT_RULE_LABEL, DEFAULT_RULE_LABEL_NEW
 from sentry.rules import EventState
 from sentry.rules.conditions.base import EventCondition
+from sentry.tsdb.base import TSDBModel
 from sentry.types.condition_activity import (
     FREQUENCY_CONDITION_BUCKET_SIZE,
     ConditionActivity,
     round_to_five_minute,
 )
 from sentry.utils import metrics
+from sentry.utils.iterators import chunked
 from sentry.utils.snuba import options_override
 
 standard_intervals = {
@@ -49,6 +54,7 @@
     COMPARISON_TYPE_COUNT: COMPARISON_TYPE_COUNT,
     COMPARISON_TYPE_PERCENT: COMPARISON_TYPE_PERCENT,
 }
+SNUBA_LIMIT = 10000
 
 
 class EventFrequencyForm(forms.Form):
@@ -241,6 +247,29 @@ def get_rate(self, event: GroupEvent, interval: str, environment_id: str) -> int
 
         return result
 
+    def get_sums(
+        self,
+        keys: list[int],
+        group: Group,
+        model: TSDBModel,
+        start: datetime,
+        end: datetime,
+        environment_id: str,
+        referrer_suffix: str,
+    ) -> Mapping[int, int]:
+        sums: Mapping[int, int] = self.tsdb.get_sums(
+            model=model,
+            keys=keys,
+            start=start,
+            end=end,
+            environment_id=environment_id,
+            use_cache=True,
+            jitter_value=group.id,
+            tenant_ids={"organization_id": group.project.organization_id},
+            referrer_suffix=referrer_suffix,
+        )
+        return sums
+
     @property
     def is_guessed_to_be_created_on_project_creation(self) -> bool:
         """
@@ -265,19 +294,70 @@ class EventFrequencyCondition(BaseEventFrequencyCondition):
     def query_hook(
         self, event: GroupEvent, start: datetime, end: datetime, environment_id: str
     ) -> int:
-        sums: Mapping[int, int] = self.tsdb.get_sums(
-            model=get_issue_tsdb_group_model(event.group.issue_category),
+        sums: Mapping[int, int] = self.get_sums(
             keys=[event.group_id],
+            group=event.group,
+            model=get_issue_tsdb_group_model(event.group.issue_category),
             start=start,
             end=end,
             environment_id=environment_id,
-            use_cache=True,
-            jitter_value=event.group_id,
-            tenant_ids={"organization_id": event.group.project.organization_id},
             referrer_suffix="alert_event_frequency",
         )
         return sums[event.group_id]
 
+    def get_chunked_sums(
+        self,
+        groups: list[Group],
+        start: datetime,
+        end: datetime,
+        environment_id: str,
+        referrer_suffix: str,
+    ) -> dict[int, int]:
+        batch_sums: dict[int, int] = defaultdict(int)
+        for group_chunk in chunked(groups, SNUBA_LIMIT):
+            group = groups[0]
+            sums = self.get_sums(
+                keys=[group.id for group in groups],
+                group=group,
+                model=get_issue_tsdb_group_model(group.issue_category),
+                start=start,
+                end=end,
+                environment_id=environment_id,
+                referrer_suffix=referrer_suffix,
+            )
+            batch_sums.update(sums)
+        return batch_sums
+
+    def batch_query_hook(
+        self, group_ids: Sequence[int], start: datetime, end: datetime, environment_id: str
+    ) -> dict[int, int]:
+        batch_sums: dict[int, int] = defaultdict(int)
+        groups = Group.objects.filter(id__in=group_ids)
+        error_issues = [group for group in groups if group.issue_category == GroupCategory.ERROR]
+        generic_issues = [group for group in groups if group.issue_category != GroupCategory.ERROR]
+
+        if error_issues:
+            error_sums = self.get_chunked_sums(
+                groups=error_issues,
+                start=start,
+                end=end,
+                environment_id=environment_id,
+                referrer_suffix="alert_event_frequency",
+            )
+            batch_sums.update(error_sums)
+
+        if generic_issues:
+            generic_sums = self.get_chunked_sums(
+                groups=generic_issues,
+                start=start,
+                end=end,
+                environment_id=environment_id,
+                referrer_suffix="alert_event_frequency",
+            )
+            batch_sums.update(generic_sums)
+
+        return batch_sums
+
     def get_preview_aggregate(self) -> tuple[str, str]:
         return "count", "roundedTime"
 
@@ -397,16 +477,13 @@ def query_hook(
                 percent_intervals[self.get_option("interval")][1].total_seconds() // 60
             )
             avg_sessions_in_interval = session_count_last_hour / (60 / interval_in_minutes)
-
-            issue_count = self.tsdb.get_sums(
-                model=get_issue_tsdb_group_model(event.group.issue_category),
+            issue_count = self.get_sums(
                 keys=[event.group_id],
+                group=event.group,
+                model=get_issue_tsdb_group_model(event.group.issue_category),
                 start=start,
                 end=end,
                 environment_id=environment_id,
-                use_cache=True,
-                jitter_value=event.group_id,
-                tenant_ids={"organization_id": event.group.project.organization_id},
                 referrer_suffix="alert_event_frequency_percent",
             )[event.group_id]
             if issue_count > avg_sessions_in_interval:

tests/snuba/rules/conditions/test_event_frequency.py

@@ -28,6 +28,76 @@
 pytestmark = [pytest.mark.sentry_metrics, requires_snuba]
 
 
+class EventFrequencyQueryTest(SnubaTestCase, RuleTestCase, PerformanceIssueTestCase):
+    rule_cls = EventFrequencyCondition
+
+    def test_batch_query(self):
+        event = self.store_event(
+            data={
+                "event_id": "a" * 32,
+                "environment": self.environment.name,
+                "timestamp": iso_format(before_now(seconds=30)),
+                "fingerprint": ["group-1"],
+            },
+            project_id=self.project.id,
+        )
+        event2 = self.store_event(
+            data={
+                "event_id": "b" * 32,
+                "environment": self.environment.name,
+                "timestamp": iso_format(before_now(seconds=12)),
+                "fingerprint": ["group-2"],
+            },
+            project_id=self.project.id,
+        )
+        environment2 = self.create_environment(name="staging")
+        event3 = self.store_event(
+            data={
+                "event_id": "c" * 32,
+                "environment": environment2.name,
+                "timestamp": iso_format(before_now(seconds=12)),
+                "fingerprint": ["group-3"],
+            },
+            project_id=self.project.id,
+        )
+
+        fingerprint = f"{PerformanceNPlusOneGroupType.type_id}-something_random"
+        event_data = load_data(
+            "transaction-n-plus-one",
+            timestamp=before_now(seconds=12),
+            start_timestamp=before_now(seconds=12),
+            fingerprint=[fingerprint],
+        )
+        event_data["user"] = {"id": uuid4().hex}
+        event_data["environment"] = self.environment.name
+
+        # Store a performance event
+        perf_event = self.create_performance_issue(
+            event_data=event_data,
+            project_id=self.project.id,
+            fingerprint=fingerprint,
+        )
+        start = before_now(minutes=1)
+        end = timezone.now()
+
+        condition_inst = self.rule_cls(event.group.project)
+        batch_query = condition_inst.batch_query_hook(
+            group_ids=[event.group_id, event2.group_id, perf_event.group_id],
+            start=start,
+            end=end,
+            environment_id=self.environment.id,
+        )
+        assert batch_query == {event.group_id: 1, event2.group_id: 1, perf_event.group_id: 1}
+
+        batch_query = condition_inst.batch_query_hook(
+            group_ids=[event3.group_id],
+            start=start,
+            end=end,
+            environment_id=environment2.id,
+        )
+        assert batch_query == {event3.group_id: 1}
+
+
 class ErrorEventMixin(SnubaTestCase):
     def add_event(self, data, project_id, timestamp):
         data["timestamp"] = iso_format(timestamp)
@@ -72,7 +142,7 @@ def add_event(self, data, project_id, timestamp):
 
 
 @pytest.mark.snuba_ci
-class StandardIntervalTestBase(SnubaTestCase, RuleTestCase):
+class StandardIntervalTestBase(SnubaTestCase, RuleTestCase, PerformanceIssueTestCase):
     __test__ = Abstract(__module__, __qualname__)
 
     def add_event(self, data, project_id, timestamp):