feat(aci): schema validation for DataCondition.comparison (#83457)

cathteng · andrewshie-sentry · commit 161b48f9ec9b · 2025-01-22T10:52:16.000-08:00
diff --git a/src/sentry/workflow_engine/handlers/condition/age_comparison_handler.py b/src/sentry/workflow_engine/handlers/condition/age_comparison_handler.py
@@ -11,13 +11,27 @@
 
 @condition_handler_registry.register(Condition.AGE_COMPARISON)
 class AgeComparisonConditionHandler(DataConditionHandler[WorkflowJob]):
+    comparison_json_schema = {
+        "type": "object",
+        "properties": {
+            "comparison_type": {
+                "type": "string",
+                "enum": [AgeComparisonType.OLDER, AgeComparisonType.NEWER],
+            },
+            "value": {"type": "integer", "minimum": 0},
+            "time": {"type": "string", "enum": list(timeranges.keys())},
+        },
+        "required": ["comparison_type", "value", "time"],
+        "additionalProperties": False,
+    }
+
     @staticmethod
     def evaluate_value(job: WorkflowJob, comparison: Any) -> bool:
         event = job["event"]
         first_seen = event.group.first_seen
         current_time = timezone.now()
-        comparison_type = comparison.get("comparison_type")
-        time = comparison.get("time")
+        comparison_type = comparison["comparison_type"]
+        time = comparison["time"]
 
         if (
             not comparison_type
@@ -31,7 +45,7 @@ def evaluate_value(job: WorkflowJob, comparison: Any) -> bool:
             return False
 
         try:
-            value = int(comparison.get("value"))
+            value = int(comparison["value"])
         except (TypeError, ValueError):
             return False
 
diff --git a/src/sentry/workflow_engine/migration_helpers/issue_alert_conditions.py b/src/sentry/workflow_engine/migration_helpers/issue_alert_conditions.py
@@ -165,10 +165,9 @@ def create_tagged_event_data_condition(
 def create_age_comparison_data_condition(
     data: dict[str, Any], dcg: DataConditionGroup
 ) -> DataCondition:
-    # TODO: Add comparison validation (error if not enough information)
     comparison = {
         "comparison_type": data["comparison_type"],
-        "value": data["value"],
+        "value": int(data["value"]),
         "time": data["time"],
     }
 
diff --git a/src/sentry/workflow_engine/models/data_condition.py b/src/sentry/workflow_engine/models/data_condition.py
@@ -3,6 +3,9 @@
 from typing import Any, TypeVar, cast
 
 from django.db import models
+from django.db.models.signals import pre_save
+from django.dispatch import receiver
+from jsonschema import ValidationError, validate
 
 from sentry.backup.scopes import RelocationScope
 from sentry.db.models import DefaultFieldsModel, region_silo_model, sane_repr
@@ -40,7 +43,7 @@ class Condition(models.TextChoices):
     ISSUE_PRIORITY_EQUALS = "issue_priority_equals"
 
 
-condition_ops = {
+CONDITION_OPS = {
     Condition.EQUAL: operator.eq,
     Condition.GREATER_OR_EQUAL: operator.ge,
     Condition.GREATER: operator.gt,
@@ -103,9 +106,9 @@ def evaluate_value(self, value: T) -> DataConditionResult:
             )
             return None
 
-        if condition_type in condition_ops:
+        if condition_type in CONDITION_OPS:
             # If the condition is a base type, handle it directly
-            op = condition_ops[Condition(self.type)]
+            op = CONDITION_OPS[Condition(self.type)]
             result = op(cast(Any, value), self.comparison)
             return self.get_condition_result() if result else None
 
@@ -121,3 +124,28 @@ def evaluate_value(self, value: T) -> DataConditionResult:
 
         result = handler.evaluate_value(value, self.comparison)
         return self.get_condition_result() if result else None
+
+
+@receiver(pre_save, sender=DataCondition)
+def enforce_comparison_schema(sender, instance: DataCondition, **kwargs):
+
+    condition_type = Condition(instance.type)
+    if condition_type in CONDITION_OPS:
+        # don't enforce schema for default ops, this can be any type
+        return
+
+    try:
+        handler = condition_handler_registry.get(condition_type)
+    except NoRegistrationExistsError:
+        logger.exception(
+            "No registration exists for condition",
+            extra={"type": instance.type, "id": instance.id},
+        )
+        return None
+
+    schema = handler.comparison_json_schema
+
+    try:
+        validate(instance.comparison, schema)
+    except ValidationError as e:
+        raise ValidationError(f"Invalid config: {e.message}")
diff --git a/src/sentry/workflow_engine/types.py b/src/sentry/workflow_engine/types.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import IntEnum
-from typing import TYPE_CHECKING, Any, Generic, TypedDict, TypeVar
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypedDict, TypeVar
 
 from sentry.types.group import PriorityLevel
 
@@ -55,6 +55,8 @@ def bulk_get_query_object(data_sources) -> dict[int, T | None]:
 
 
 class DataConditionHandler(Generic[T]):
+    comparison_json_schema: ClassVar[dict[str, Any]] = {}
+
     @staticmethod
     def evaluate_value(value: T, comparison: Any) -> DataConditionResult:
         raise NotImplementedError
diff --git a/tests/sentry/workflow_engine/handlers/condition/test_age_comparison_handler.py b/tests/sentry/workflow_engine/handlers/condition/test_age_comparison_handler.py
@@ -1,5 +1,8 @@
 from datetime import datetime, timedelta, timezone
 
+import pytest
+from jsonschema import ValidationError
+
 from sentry.rules.age import AgeComparisonType
 from sentry.rules.filters.age_comparison import AgeComparisonFilter
 from sentry.testutils.helpers.datetime import freeze_time
@@ -36,7 +39,7 @@ def setUp(self):
         )
         self.dc = self.create_data_condition(
             type=self.condition,
-            comparison={"comparison_type": AgeComparisonType.OLDER, "value": "10", "time": "hour"},
+            comparison={"comparison_type": AgeComparisonType.OLDER, "value": 10, "time": "hour"},
             condition_result=True,
         )
 
@@ -47,16 +50,30 @@ def test_dual_write(self):
         assert dc.type == self.condition
         assert dc.comparison == {
             "comparison_type": AgeComparisonType.OLDER,
-            "value": "10",
+            "value": 10,
             "time": "hour",
         }
         assert dc.condition_result is True
         assert dc.condition_group == dcg
 
+    def test_json_schema(self):
+        self.dc.comparison.update({"time": "asdf"})
+        with pytest.raises(ValidationError):
+            self.dc.save()
+
+        self.dc.comparison.update({"value": "bad_value"})
+        with pytest.raises(ValidationError):
+            self.dc.save()
+
+        self.dc.comparison.update({"comparison_type": "bad_value"})
+        with pytest.raises(ValidationError):
+            self.dc.save()
+
     def test_older_applies_correctly(self):
-        self.dc.update(
-            comparison={"comparison_type": AgeComparisonType.OLDER, "value": "10", "time": "hour"}
+        self.dc.comparison.update(
+            {"comparison_type": AgeComparisonType.OLDER, "value": 10, "time": "hour"}
         )
+        self.dc.save()
 
         self.group.update(first_seen=datetime.now(timezone.utc) - timedelta(hours=3))
         self.assert_does_not_pass(self.dc, self.job)
@@ -67,22 +84,13 @@ def test_older_applies_correctly(self):
         self.assert_passes(self.dc, self.job)
 
     def test_newer_applies_correctly(self):
-        self.dc.update(
-            comparison={"comparison_type": AgeComparisonType.NEWER, "value": "10", "time": "hour"}
+        self.dc.comparison.update(
+            {"comparison_type": AgeComparisonType.NEWER, "value": 10, "time": "hour"}
         )
+        self.dc.save()
 
         self.group.update(first_seen=datetime.now(timezone.utc) - timedelta(hours=3))
         self.assert_passes(self.dc, self.job)
 
         self.group.update(first_seen=datetime.now(timezone.utc) - timedelta(hours=10))
         self.assert_does_not_pass(self.dc, self.job)
-
-    def test_fails_on_insufficient_data(self):
-        self.dc.update(comparison={"time": "hour"})
-        self.assert_does_not_pass(self.dc, self.job)
-
-        self.dc.update(comparison={"value": "bad_value"})
-        self.assert_does_not_pass(self.dc, self.job)
-
-        self.dc.update(comparison={"comparison_type": "bad_value"})
-        self.assert_does_not_pass(self.dc, self.job)
diff --git a/tests/sentry/workflow_engine/migration_helpers/test_migrate_rule.py b/tests/sentry/workflow_engine/migration_helpers/test_migrate_rule.py
@@ -84,7 +84,7 @@ def test_create_issue_alert_no_actions(self):
             type=Condition.AGE_COMPARISON,
             comparison={
                 "comparison_type": AgeComparisonType.OLDER,
-                "value": "10",
+                "value": 10,
                 "time": "hour",
             },
             condition_result=True,
diff --git a/tests/sentry/workflow_engine/models/test_data_condition.py b/tests/sentry/workflow_engine/models/test_data_condition.py
@@ -46,14 +46,11 @@ def test(self):
         assert dc.evaluate_value(1) is None
 
     def test_bad_condition(self):
-
-        dc = self.create_data_condition(
-            type="invalid", comparison=1.0, condition_result=DetectorPriorityLevel.HIGH
-        )
-
-        with mock.patch("sentry.workflow_engine.models.data_condition.logger") as mock_logger:
-            assert dc.evaluate_value(2) is None
-            assert mock_logger.exception.call_args[0][0] == "Invalid condition type"
+        with pytest.raises(ValueError):
+            # Raises ValueError because the condition is invalid
+            self.create_data_condition(
+                type="invalid", comparison=1.0, condition_result=DetectorPriorityLevel.HIGH
+            )
 
     def test_bad_comparison(self):
         dc = self.create_data_condition(
