fix(aci): fix Group N+1 query in delayed workflow query handlers (#95161)

Author: cathteng

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

@@ -6,7 +6,6 @@
 from typing import Any, ClassVar, Literal, Protocol, TypedDict
 
 from django.core.cache import cache
-from django.db.models import QuerySet
 from snuba_sdk import Op
 
 from sentry import release_health, tsdb
@@ -16,7 +15,6 @@
     get_issue_tsdb_user_group_model,
 )
 from sentry.issues.grouptype import GroupCategory, get_group_type_by_type_id
-from sentry.models.group import Group
 from sentry.rules.conditions.event_attribute import ATTR_CHOICES
 from sentry.rules.conditions.event_frequency import (
     MIN_SESSIONS_TO_FIRE,
@@ -35,6 +33,13 @@
 QueryResult = dict[int, int | float]
 
 
+class GroupValues(TypedDict):
+    id: int
+    type: int
+    project_id: int
+    project__organization_id: int
+
+
 class TSDBFunction(Protocol):
     def __call__(
         self,
@@ -61,13 +66,6 @@ class InvalidFilter(Exception):
     pass
 
 
-class _QSTypedDict(TypedDict):
-    id: int
-    type: int
-    project_id: int
-    project__organization_id: int
-
-
 class BaseEventFrequencyQueryHandler(ABC):
     intervals: ClassVar[dict[str, tuple[str, timedelta]]] = STANDARD_INTERVALS
 
@@ -154,7 +152,7 @@ def get_chunked_result(
 
     def get_group_ids_by_category(
         self,
-        groups: QuerySet[Group, _QSTypedDict],
+        groups: list[GroupValues],
     ) -> dict[GroupCategory, list[int]]:
         """
         Separate group ids into error group ids and generic group ids
@@ -170,7 +168,7 @@ def get_group_ids_by_category(
 
     def get_value_from_groups(
         self,
-        groups: QuerySet[Group, _QSTypedDict] | None,
+        groups: list[GroupValues],
         value: Literal["id", "project_id", "project__organization_id"],
     ) -> int | None:
         result = None
@@ -270,7 +268,7 @@ def convert_filter_to_snuba_condition(
     @abstractmethod
     def batch_query(
         self,
-        group_ids: set[int],
+        groups: list[GroupValues],
         start: datetime,
         end: datetime,
         environment_id: int | None,
@@ -285,7 +283,7 @@ def batch_query(
     def get_rate_bulk(
         self,
         duration: timedelta,
-        group_ids: set[int],
+        groups: list[GroupValues],
         environment_id: int | None,
         current_time: datetime,
         comparison_interval: timedelta | None,
@@ -306,7 +304,7 @@ def get_rate_bulk(
 
         with self.disable_consistent_snuba_mode(duration):
             result = self.batch_query(
-                group_ids=group_ids,
+                groups=groups,
                 start=start,
                 end=end,
                 environment_id=environment_id,
@@ -325,16 +323,13 @@ def get_rate_bulk(
 class EventFrequencyQueryHandler(BaseEventFrequencyQueryHandler):
     def batch_query(
         self,
-        group_ids: set[int],
+        groups: list[GroupValues],
         start: datetime,
         end: datetime,
         environment_id: int | None,
         filters: list[QueryFilter] | None = None,
     ) -> QueryResult:
         batch_sums: QueryResult = defaultdict(int)
-        groups = Group.objects.filter(id__in=group_ids).values(
-            "id", "type", "project_id", "project__organization_id"
-        )
         category_group_ids = self.get_group_ids_by_category(groups)
         organization_id = self.get_value_from_groups(groups, "project__organization_id")
 
@@ -371,16 +366,13 @@ def batch_query(
 class EventUniqueUserFrequencyQueryHandler(BaseEventFrequencyQueryHandler):
     def batch_query(
         self,
-        group_ids: set[int],
+        groups: list[GroupValues],
         start: datetime,
         end: datetime,
         environment_id: int | None,
         filters: list[QueryFilter] | None = None,
     ) -> QueryResult:
         batch_sums: QueryResult = defaultdict(int)
-        groups = Group.objects.filter(id__in=group_ids).values(
-            "id", "type", "project_id", "project__organization_id"
-        )
         category_group_ids = self.get_group_ids_by_category(groups)
         organization_id = self.get_value_from_groups(groups, "project__organization_id")
 
@@ -442,16 +434,13 @@ def get_session_interval(self, session_count: int, duration: timedelta) -> int |
 
     def batch_query(
         self,
-        group_ids: set[int],
+        groups: list[GroupValues],
         start: datetime,
         end: datetime,
         environment_id: int | None,
         filters: list[QueryFilter] | None = None,
     ) -> QueryResult:
         batch_percents: QueryResult = {}
-        groups = Group.objects.filter(id__in=group_ids).values(
-            "id", "type", "project_id", "project__organization_id"
-        )
         category_group_ids = self.get_group_ids_by_category(groups)
         project_id = self.get_value_from_groups(groups, "project_id")
 

src/sentry/workflow_engine/processors/delayed_workflow.py

@@ -39,6 +39,7 @@
 from sentry.utils.retries import ConditionalRetryPolicy, exponential_delay
 from sentry.workflow_engine.handlers.condition.event_frequency_query_handlers import (
     BaseEventFrequencyQueryHandler,
+    GroupValues,
     QueryFilter,
     QueryResult,
     slow_condition_query_handler_registry,
@@ -423,9 +424,21 @@ def get_condition_group_results(
     condition_group_results = {}
     current_time = timezone.now()
 
+    all_group_ids: set[GroupId] = set()
+    # bulk gather groups and fetch them
+    for time_and_groups in queries_to_groups.values():
+        all_group_ids.update(time_and_groups.group_ids)
+
+    all_groups: list[GroupValues] = list(
+        Group.objects.filter(id__in=all_group_ids).values(
+            "id", "type", "project_id", "project__organization_id"
+        )
+    )
+
     for unique_condition, time_and_groups in queries_to_groups.items():
         handler = unique_condition.handler()
         group_ids = time_and_groups.group_ids
+        groups_to_query = [group for group in all_groups if group["id"] in group_ids]
         time = time_and_groups.timestamp or current_time
 
         _, duration = handler.intervals[unique_condition.interval]
@@ -438,7 +451,7 @@ def get_condition_group_results(
 
         result = handler.get_rate_bulk(
             duration=duration,
-            group_ids=group_ids,
+            groups=groups_to_query,
             environment_id=unique_condition.environment_id,
             current_time=time,
             comparison_interval=comparison_interval,

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

@@ -3,6 +3,7 @@
 from uuid import uuid4
 
 from sentry.issues.grouptype import PerformanceNPlusOneGroupType
+from sentry.models.group import Group
 from sentry.testutils.cases import PerformanceIssueTestCase, RuleTestCase, SnubaTestCase
 from sentry.testutils.helpers.datetime import before_now
 from sentry.utils.samples import load_data
@@ -129,3 +130,14 @@ def setUp(self):
             fingerprint=fingerprint,
         )
         self.data = {"interval": "5m", "value": 30}
+
+        self.groups = list(
+            Group.objects.filter(
+                id__in={self.event.group_id, self.event2.group_id, self.perf_event.group_id}
+            ).values("id", "type", "project_id", "project__organization_id")
+        )
+        self.group_3 = list(
+            Group.objects.filter(id=self.event3.group_id).values(
+                "id", "type", "project_id", "project__organization_id"
+            )
+        )