feat(alerts): ACI dual write alert rule helpers (#82400)

ceorourke · web-flow · commit 526a0594d1d2 · 2024-12-20T11:42:45.000-08:00
A smaller chunk of #81953 that creates the helper methods to migrate the `AlertRule` and not the `AlertRuleTrigger` or `AlertRuleTriggerAction` just yet.
diff --git a/src/sentry/workflow_engine/migration_helpers/alert_rule.py b/src/sentry/workflow_engine/migration_helpers/alert_rule.py
@@ -0,0 +1,156 @@
+from sentry.incidents.grouptype import MetricAlertFire
+from sentry.incidents.models.alert_rule import AlertRule
+from sentry.snuba.models import QuerySubscription, SnubaQuery
+from sentry.users.services.user import RpcUser
+from sentry.workflow_engine.models import (
+    AlertRuleDetector,
+    AlertRuleWorkflow,
+    DataConditionGroup,
+    DataSource,
+    Detector,
+    DetectorState,
+    DetectorWorkflow,
+    Workflow,
+    WorkflowDataConditionGroup,
+)
+from sentry.workflow_engine.types import DetectorPriorityLevel
+
+
+def create_metric_alert_lookup_tables(
+    alert_rule: AlertRule,
+    detector: Detector,
+    workflow: Workflow,
+    data_source: DataSource,
+    data_condition_group: DataConditionGroup,
+) -> tuple[AlertRuleDetector, AlertRuleWorkflow, DetectorWorkflow, WorkflowDataConditionGroup]:
+    alert_rule_detector = AlertRuleDetector.objects.create(alert_rule=alert_rule, detector=detector)
+    alert_rule_workflow = AlertRuleWorkflow.objects.create(alert_rule=alert_rule, workflow=workflow)
+    detector_workflow = DetectorWorkflow.objects.create(detector=detector, workflow=workflow)
+    workflow_data_condition_group = WorkflowDataConditionGroup.objects.create(
+        condition_group=data_condition_group, workflow=workflow
+    )
+    return (
+        alert_rule_detector,
+        alert_rule_workflow,
+        detector_workflow,
+        workflow_data_condition_group,
+    )
+
+
+def create_data_source(
+    organization_id: int, snuba_query: SnubaQuery | None = None
+) -> DataSource | None:
+    if not snuba_query:
+        return None
+
+    try:
+        query_subscription = QuerySubscription.objects.get(snuba_query=snuba_query.id)
+    except QuerySubscription.DoesNotExist:
+        return None
+
+    return DataSource.objects.create(
+        organization_id=organization_id,
+        query_id=query_subscription.id,
+        type="snuba_query_subscription",
+    )
+
+
+def create_data_condition_group(organization_id: int) -> DataConditionGroup:
+    return DataConditionGroup.objects.create(
+        logic_type=DataConditionGroup.Type.ANY,
+        organization_id=organization_id,
+    )
+
+
+def create_workflow(
+    name: str,
+    organization_id: int,
+    data_condition_group: DataConditionGroup,
+    user: RpcUser | None = None,
+) -> Workflow:
+    return Workflow.objects.create(
+        name=name,
+        organization_id=organization_id,
+        when_condition_group=data_condition_group,
+        enabled=True,
+        created_by_id=user.id if user else None,
+    )
+
+
+def create_detector(
+    alert_rule: AlertRule,
+    project_id: int,
+    data_condition_group: DataConditionGroup,
+    user: RpcUser | None = None,
+) -> Detector:
+    return Detector.objects.create(
+        project_id=project_id,
+        enabled=True,
+        created_by_id=user.id if user else None,
+        name=alert_rule.name,
+        workflow_condition_group=data_condition_group,
+        type=MetricAlertFire.slug,
+        description=alert_rule.description,
+        owner_user_id=alert_rule.user_id,
+        owner_team=alert_rule.team,
+        config={  # TODO create a schema
+            "threshold_period": alert_rule.threshold_period,
+            "sensitivity": alert_rule.sensitivity,
+            "seasonality": alert_rule.seasonality,
+            "comparison_delta": alert_rule.comparison_delta,
+        },
+    )
+
+
+def migrate_alert_rule(
+    alert_rule: AlertRule,
+    user: RpcUser | None = None,
+) -> (
+    tuple[
+        DataSource,
+        DataConditionGroup,
+        Workflow,
+        Detector,
+        DetectorState,
+        AlertRuleDetector,
+        AlertRuleWorkflow,
+        DetectorWorkflow,
+        WorkflowDataConditionGroup,
+    ]
+    | None
+):
+    organization_id = alert_rule.organization_id
+    project = alert_rule.projects.first()
+    if not project:
+        return None
+
+    data_source = create_data_source(organization_id, alert_rule.snuba_query)
+    if not data_source:
+        return None
+
+    data_condition_group = create_data_condition_group(organization_id)
+    workflow = create_workflow(alert_rule.name, organization_id, data_condition_group, user)
+    detector = create_detector(alert_rule, project.id, data_condition_group, user)
+
+    data_source.detectors.set([detector])
+    detector_state = DetectorState.objects.create(
+        detector=detector,
+        active=False,
+        state=DetectorPriorityLevel.OK,
+    )
+    alert_rule_detector, alert_rule_workflow, detector_workflow, workflow_data_condition_group = (
+        create_metric_alert_lookup_tables(
+            alert_rule, detector, workflow, data_source, data_condition_group
+        )
+    )
+    return (
+        data_source,
+        data_condition_group,
+        workflow,
+        detector,
+        detector_state,
+        alert_rule_detector,
+        alert_rule_workflow,
+        detector_workflow,
+        workflow_data_condition_group,
+    )
diff --git a/tests/sentry/workflow_engine/migration_helpers/test_migrate_alert_rule.py b/tests/sentry/workflow_engine/migration_helpers/test_migrate_alert_rule.py
@@ -0,0 +1,83 @@
+from sentry.incidents.grouptype import MetricAlertFire
+from sentry.snuba.models import QuerySubscription
+from sentry.testutils.cases import APITestCase
+from sentry.users.services.user.service import user_service
+from sentry.workflow_engine.migration_helpers.alert_rule import migrate_alert_rule
+from sentry.workflow_engine.models import (
+    AlertRuleDetector,
+    AlertRuleWorkflow,
+    DataSource,
+    DataSourceDetector,
+    Detector,
+    DetectorState,
+    DetectorWorkflow,
+    Workflow,
+    WorkflowDataConditionGroup,
+)
+from sentry.workflow_engine.types import DetectorPriorityLevel
+
+
+class AlertRuleMigrationHelpersTest(APITestCase):
+    def setUp(self):
+        self.metric_alert = self.create_alert_rule()
+        self.rpc_user = user_service.get_user(user_id=self.user.id)
+
+    def test_create_metric_alert(self):
+        """
+        Test that when we call the helper methods we create all the ACI models correctly for an alert rule
+        """
+        migrate_alert_rule(self.metric_alert, self.rpc_user)
+
+        alert_rule_workflow = AlertRuleWorkflow.objects.get(alert_rule=self.metric_alert)
+        alert_rule_detector = AlertRuleDetector.objects.get(alert_rule=self.metric_alert)
+
+        workflow = Workflow.objects.get(id=alert_rule_workflow.workflow.id)
+        assert workflow.name == self.metric_alert.name
+        assert self.metric_alert.organization
+        assert workflow.organization_id == self.metric_alert.organization.id
+        detector = Detector.objects.get(id=alert_rule_detector.detector.id)
+        assert detector.name == self.metric_alert.name
+        assert detector.project_id == self.project.id
+        assert detector.enabled is True
+        assert detector.description == self.metric_alert.description
+        assert detector.owner_user_id == self.metric_alert.user_id
+        assert detector.owner_team == self.metric_alert.team
+        assert detector.type == MetricAlertFire.slug
+        assert detector.config == {
+            "threshold_period": self.metric_alert.threshold_period,
+            "sensitivity": None,
+            "seasonality": None,
+            "comparison_delta": None,
+        }
+
+        detector_workflow = DetectorWorkflow.objects.get(detector=detector)
+        assert detector_workflow.workflow == workflow
+
+        workflow_data_condition_group = WorkflowDataConditionGroup.objects.get(workflow=workflow)
+        assert workflow_data_condition_group.condition_group == workflow.when_condition_group
+
+        assert self.metric_alert.snuba_query
+        query_subscription = QuerySubscription.objects.get(
+            snuba_query=self.metric_alert.snuba_query.id
+        )
+        data_source = DataSource.objects.get(
+            organization_id=self.metric_alert.organization_id, query_id=query_subscription.id
+        )
+        assert data_source.type == "snuba_query_subscription"
+        detector_state = DetectorState.objects.get(detector=detector)
+        assert detector_state.active is False
+        assert detector_state.state == str(DetectorPriorityLevel.OK.value)
+
+        data_source_detector = DataSourceDetector.objects.get(data_source=data_source)
+        assert data_source_detector.detector == detector
+
+    def test_create_metric_alert_no_data_source(self):
+        """
+        Test that when we return None and don't create any ACI models if the data source can't be created
+        """
+        self.metric_alert.update(snuba_query=None)
+        migrated = migrate_alert_rule(self.metric_alert, self.rpc_user)
+        assert migrated is None
+        assert len(DataSource.objects.all()) == 0
+        assert not AlertRuleWorkflow.objects.filter(alert_rule=self.metric_alert).exists()
+        assert not AlertRuleDetector.objects.filter(alert_rule=self.metric_alert).exists()
