✨ feat(aci): refactor issue digests for ACI (#90268)

Author: iamrajjoshi

src/sentry/digests/notifications.py

@@ -6,14 +6,16 @@
 from typing import Any, NamedTuple, TypeAlias
 
 from sentry import features, tsdb
-from sentry.digests.types import Notification, Record, RecordWithRuleObjects
+from sentry.digests.types import IdentifierKey, Notification, Record, RecordWithRuleObjects
 from sentry.eventstore.models import Event
 from sentry.models.group import Group, GroupStatus
 from sentry.models.project import Project
 from sentry.models.rule import Rule
 from sentry.notifications.types import ActionTargetType, FallthroughChoiceType
 from sentry.notifications.utils.rules import get_key_from_rule_data
 from sentry.tsdb.base import TSDBModel
+from sentry.workflow_engine.models import Workflow
+from sentry.workflow_engine.models.alertrule_workflow import AlertRuleWorkflow
 
 logger = logging.getLogger("sentry.digests")
 
@@ -70,16 +72,20 @@ def event_to_record(
         logger.warning("Creating record for %s that does not contain any rules!", event)
 
     rule_ids = []
-    # TODO(iamrajjoshi): This will only work during the dual write period of the rollout!
-    if features.has("organizations:workflow-engine-trigger-actions", event.organization):
+    identifier_key = IdentifierKey.RULE
+    if features.has("organizations:workflow-engine-ui-links", event.organization):
+        identifier_key = IdentifierKey.WORKFLOW
+        for rule in rules:
+            rule_ids.append(int(get_key_from_rule_data(rule, "workflow_id")))
+    elif features.has("organizations:workflow-engine-trigger-actions", event.organization):
         for rule in rules:
             rule_ids.append(int(get_key_from_rule_data(rule, "legacy_rule_id")))
     else:
         for rule in rules:
             rule_ids.append(rule.id)
     return Record(
         event.event_id,
-        Notification(event, rule_ids, notification_uuid),
+        Notification(event, rule_ids, notification_uuid, identifier_key),
         event.datetime.timestamp(),
     )
 
@@ -159,6 +165,67 @@ def _build_digest_impl(
     return _sort_digest(grouped, event_counts=event_counts, user_counts=user_counts)
 
 
+def get_rules_from_workflows(project: Project, workflow_ids: set[int]) -> dict[int, Rule]:
+    rules: dict[int, Rule] = {}
+    if not workflow_ids:
+        return rules
+
+    # Fetch all workflows in bulk
+    workflows = Workflow.objects.filter(organization_id=project.organization_id).in_bulk(
+        workflow_ids
+    )
+
+    # We are only processing the workflows in the digest if under the new flag
+    # This should be ok since we should only add workflow_ids to redis when under this flag
+    if features.has("organizations:workflow-engine-ui-links", project.organization):
+        for workflow_id, workflow in workflows.items():
+            assert (
+                workflow.organization_id == project.organization_id
+            ), "Workflow must belong to Organization"
+            rules[workflow_id] = Rule(
+                label=workflow.name,
+                id=workflow_id,
+                project_id=project.id,
+                # We need to do this so that the links are built correctly downstream
+                data={"actions": [{"workflow_id": workflow_id}]},
+            )
+    # This is if we had workflows in the digest but the flag is not enabled
+    # This can happen if we rollback the flag, but the records in the digest aren't flushed
+    else:
+        alert_rule_workflows = AlertRuleWorkflow.objects.filter(workflow_id__in=workflow_ids)
+        alert_rule_workflows_map = {awf.workflow_id: awf for awf in alert_rule_workflows}
+
+        rule_ids_to_fetch = {awf.rule_id for awf in alert_rule_workflows}
+
+        bulk_rules = Rule.objects.filter(project_id=project.id).in_bulk(rule_ids_to_fetch)
+
+        for workflow_id in workflow_ids:
+            alert_workflow = alert_rule_workflows_map.get(workflow_id)
+            if not alert_workflow:
+                logger.warning(
+                    "Workflow %s does not have a corresponding AlertRuleWorkflow entry", workflow_id
+                )
+                raise
+
+            rule = bulk_rules.get(alert_workflow.rule_id)
+            if not rule:
+                logger.warning(
+                    "Rule %s linked to Workflow %s not found or does not belong to project %s",
+                    alert_workflow.rule_id,
+                    workflow_id,
+                    project.id,
+                )
+                continue
+
+            assert rule.project_id == project.id, "Rule must belong to Project"
+
+            if features.has("organizations:workflow-engine-trigger-actions", project.organization):
+                rule.data["actions"][0]["legacy_rule_id"] = rule.id
+
+            rules[workflow_id] = rule
+    return rules
+
+
 def build_digest(project: Project, records: Sequence[Record]) -> DigestInfo:
     if not records:
         return DigestInfo({}, {}, {})
@@ -170,15 +237,28 @@ def build_digest(project: Project, records: Sequence[Record]) -> DigestInfo:
     start = records[-1].datetime
     end = records[0].datetime
 
+    rule_ids: set[int] = set()
+    workflow_ids: set[int] = set()
+
+    for record in records:
+        identifier_key = getattr(record.value, "identifier_key", IdentifierKey.RULE)
+        # record.value is Notification, record.value.rules is Sequence[int]
+        ids_to_add = record.value.rules
+        if identifier_key == IdentifierKey.RULE:
+            rule_ids.update(ids_to_add)
+        elif identifier_key == IdentifierKey.WORKFLOW:
+            workflow_ids.update(ids_to_add)
+
     groups = Group.objects.in_bulk(record.value.event.group_id for record in records)
     group_ids = list(groups)
-    rules = Rule.objects.in_bulk(rule_id for record in records for rule_id in record.value.rules)
+    rules = Rule.objects.in_bulk(rule_ids)
 
-    # TODO(iamrajjoshi): This will only work during the dual write period of the rollout!
     if features.has("organizations:workflow-engine-trigger-actions", project.organization):
         for rule in rules.values():
             rule.data["actions"][0]["legacy_rule_id"] = rule.id
 
+    rules.update(get_rules_from_workflows(project, workflow_ids))
+
     for group_id, g in groups.items():
         assert g.project_id == project.id, "Group must belong to Project"
     for rule_id, rule in rules.items():
@@ -199,7 +279,6 @@ def build_digest(project: Project, records: Sequence[Record]) -> DigestInfo:
         end,
         tenant_ids=tenant_ids,
     )
-
     digest = _build_digest_impl(records, groups, rules, event_counts, user_counts)
 
     return DigestInfo(digest, event_counts, user_counts)

src/sentry/digests/types.py

@@ -2,6 +2,7 @@
 
 import datetime as datetime_mod
 from collections.abc import Sequence
+from enum import StrEnum
 from typing import TYPE_CHECKING, NamedTuple
 
 from sentry.utils.dates import to_datetime
@@ -11,16 +12,23 @@
     from sentry.models.rule import Rule
 
 
+class IdentifierKey(StrEnum):
+    RULE = "rule"
+    WORKFLOW = "workflow"
+
+
 class Notification(NamedTuple):
     event: Event
     rules: Sequence[int] = ()
     notification_uuid: str | None = None
+    identifier_key: IdentifierKey = IdentifierKey.RULE
 
     def with_rules(self, rules: list[Rule]) -> NotificationWithRuleObjects:
         return NotificationWithRuleObjects(
             event=self.event,
             rules=rules,
             notification_uuid=self.notification_uuid,
+            # We don't really worry about identifier_key here since this method is not used after we pop record from redis
         )
 
 

src/sentry/notifications/notifications/digest.py

@@ -117,14 +117,20 @@ def get_context(self) -> MutableMapping[str, Any]:
 
         sentry_query_params = self.get_sentry_query_params(ExternalProviders.EMAIL)
 
-        snooze_alert = len(rule_details) > 0
-        snooze_alert_urls = {
-            rule.id: f"{rule.status_url}{sentry_query_params}&{urlencode({'mute': '1'})}"
-            for rule in rule_details
-        }
-
-        context["snooze_alert"] = snooze_alert
-        context["snooze_alert_urls"] = snooze_alert_urls
+        if not features.has("organizations:workflow-engine-ui-links", self.project.organization):
+            # TODO(iamrajjoshi): This actually mutes a rule for a user, something we have not ported over in the new system
+            # By not including this context, the template will not show the mute button
+            snooze_alert = len(rule_details) > 0
+            snooze_alert_urls = {
+                rule.id: f"{rule.status_url}{sentry_query_params}&{urlencode({'mute': '1'})}"
+                for rule in rule_details
+            }
+
+            context["snooze_alert"] = snooze_alert
+            context["snooze_alert_urls"] = snooze_alert_urls
+        else:
+            context["snooze_alert"] = False
+            context["snooze_alert_urls"] = None
 
         return context