chore(spans-migration): add post and put validation for EAP extrapolation mode (#102970)

We've added a new extrapolation mode field in metrics alerts for EAP. We
want to make sure the `server_weighted` value is not put in by the user.
They also shouldn't be able to set this extrapolation mode if there
already is another mode when updating. I've added tests to confirm but
let me know if there are better ways to set and validate the
extrapolation mode in the new alerts code.

Author: nikkikapadia

src/sentry/incidents/endpoints/organization_alert_rule_details.py

@@ -38,6 +38,7 @@
     delete_alert_rule,
     get_slack_actions_with_async_lookups,
 )
+from sentry.incidents.metric_issue_detector import is_invalid_extrapolation_mode
 from sentry.incidents.models.alert_rule import AlertRule
 from sentry.incidents.serializers import AlertRuleSerializer as DrfAlertRuleSerializer
 from sentry.incidents.utils.sentry_apps import trigger_sentry_app_action_creators_for_incidents
@@ -49,6 +50,8 @@
 from sentry.models.rulesnooze import RuleSnooze
 from sentry.sentry_apps.services.app import app_service
 from sentry.sentry_apps.utils.errors import SentryAppBaseError
+from sentry.snuba.dataset import Dataset
+from sentry.snuba.models import ExtrapolationMode
 from sentry.users.services.user.service import user_service
 from sentry.workflow_engine.migration_helpers.alert_rule import dual_delete_migrated_alert_rule
 from sentry.workflow_engine.models import Detector
@@ -122,6 +125,17 @@ def update_alert_rule(
         partial=True,
     )
     if validator.is_valid():
+        if data.get("dataset") == Dataset.EventsAnalyticsPlatform.value:
+            if data.get("extrapolation_mode"):
+                old_extrapolation_mode = ExtrapolationMode(
+                    alert_rule.snuba_query.extrapolation_mode
+                ).name.lower()
+                new_extrapolation_mode = data.get("extrapolation_mode", old_extrapolation_mode)
+                if is_invalid_extrapolation_mode(old_extrapolation_mode, new_extrapolation_mode):
+                    raise serializers.ValidationError(
+                        "Invalid extrapolation mode for this alert type."
+                    )
+
         try:
             trigger_sentry_app_action_creators_for_incidents(validator.validated_data)
         except SentryAppBaseError as e:

src/sentry/incidents/endpoints/organization_alert_rule_index.py

@@ -73,6 +73,7 @@
 from sentry.sentry_apps.services.app import app_service
 from sentry.sentry_apps.utils.errors import SentryAppBaseError
 from sentry.snuba.dataset import Dataset
+from sentry.snuba.models import ExtrapolationMode
 from sentry.uptime.types import (
     DATA_SOURCE_UPTIME_SUBSCRIPTION,
     GROUP_TYPE_UPTIME_DOMAIN_CHECK_FAILURE,
@@ -101,6 +102,9 @@ def create_metric_alert(
             "Creation of transaction-based alerts is disabled, as we migrate to the span dataset. Create span-based alerts (dataset: events_analytics_platform) with the is_transaction:true filter instead."
         )
 
+    if data.get("extrapolation_mode") == ExtrapolationMode.SERVER_WEIGHTED.name.lower():
+        raise ValidationError("server_weighted extrapolation mode is not supported for new alerts.")
+
     if project:
         data["projects"] = [project.slug]
 

src/sentry/incidents/metric_issue_detector.py

@@ -11,7 +11,12 @@
 from sentry.seer.anomaly_detection.store_data_workflow_engine import send_new_detector_data
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.metrics.extraction import should_use_on_demand_metrics
-from sentry.snuba.models import QuerySubscription, SnubaQuery, SnubaQueryEventType
+from sentry.snuba.models import (
+    ExtrapolationMode,
+    QuerySubscription,
+    SnubaQuery,
+    SnubaQueryEventType,
+)
 from sentry.snuba.snuba_query_validator import SnubaQueryValidator
 from sentry.snuba.subscriptions import update_snuba_query
 from sentry.tasks.relay import schedule_invalidate_project_config
@@ -147,6 +152,19 @@ def validate_conditions(self, value):
         return value
 
 
+def is_invalid_extrapolation_mode(old_extrapolation_mode, new_extrapolation_mode) -> bool:
+    if type(new_extrapolation_mode) is int:
+        new_extrapolation_mode = ExtrapolationMode(new_extrapolation_mode).name.lower()
+    if type(old_extrapolation_mode) is int:
+        old_extrapolation_mode = ExtrapolationMode(old_extrapolation_mode).name.lower()
+    if (
+        new_extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED.name.lower()
+        and old_extrapolation_mode != ExtrapolationMode.SERVER_WEIGHTED.name.lower()
+    ):
+        return True
+    return False
+
+
 class MetricIssueDetectorValidator(BaseDetectorTypeValidator):
     data_sources = serializers.ListField(
         child=SnubaQueryValidator(timeWindowSeconds=True), required=False
@@ -174,6 +192,12 @@ def _validate_transaction_dataset_deprecation(self, dataset: Dataset) -> None:
                     "Creation of transaction-based alerts is disabled, as we migrate to the span dataset. Create span-based alerts (dataset: events_analytics_platform) with the is_transaction:true filter instead."
                 )
 
+    def _validate_extrapolation_mode(self, extrapolation_mode: str) -> None:
+        if extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED.value:
+            raise serializers.ValidationError(
+                "server_weighted extrapolation mode is not supported for new detectors."
+            )
+
     def get_quota(self) -> DetectorQuota:
         organization = self.context.get("organization")
         request = self.context.get("request")
@@ -236,6 +260,16 @@ def update_data_source(self, instance: Detector, data_source: SnubaQueryDataSour
                 "Updates to transaction-based alerts is disabled, as we migrate to the span dataset. Create span-based alerts (dataset: events_analytics_platform) with the is_transaction:true filter instead."
             )
 
+        old_extrapolation_mode = snuba_query.extrapolation_mode
+        new_extrapolation_mode = data_source.get(
+            "extrapolation_mode", snuba_query.extrapolation_mode
+        )
+        if data_source.get("dataset") == Dataset.EventsAnalyticsPlatform:
+            if is_invalid_extrapolation_mode(old_extrapolation_mode, new_extrapolation_mode):
+                raise serializers.ValidationError(
+                    "Invalid extrapolation mode for this detector type."
+                )
+
         update_snuba_query(
             snuba_query=snuba_query,
             query_type=data_source.get("query_type", snuba_query.type),
@@ -246,6 +280,9 @@ def update_data_source(self, instance: Detector, data_source: SnubaQueryDataSour
             resolution=timedelta(seconds=data_source.get("resolution", snuba_query.resolution)),
             environment=data_source.get("environment", snuba_query.environment),
             event_types=data_source.get("event_types", [event_type for event_type in event_types]),
+            extrapolation_mode=data_source.get(
+                "extrapolation_mode", snuba_query.extrapolation_mode
+            ),
         )
 
     def update(self, instance: Detector, validated_data: dict[str, Any]):
@@ -279,6 +316,7 @@ def create(self, validated_data: dict[str, Any]):
         if "data_sources" in validated_data:
             for validated_data_source in validated_data["data_sources"]:
                 self._validate_transaction_dataset_deprecation(validated_data_source.get("dataset"))
+                self._validate_extrapolation_mode(validated_data_source.get("extrapolation_mode"))
 
         detector = super().create(validated_data)
 

src/sentry/snuba/snuba_query_validator.py

@@ -31,6 +31,7 @@
 )
 from sentry.snuba.metrics.naming_layer.mri import is_mri
 from sentry.snuba.models import (
+    ExtrapolationMode,
     QuerySubscription,
     QuerySubscriptionDataSourceHandler,
     SnubaQuery,
@@ -86,6 +87,7 @@ class SnubaQueryValidator(BaseDataSourceValidator[QuerySubscription]):
         required=False,
         allow_empty=False,
     )
+    extrapolation_mode = serializers.IntegerField(required=False, allow_null=True)
 
     class Meta:
         model = QuerySubscription
@@ -98,6 +100,7 @@ class Meta:
             "environment",
             "event_types",
             "group_by",
+            "extrapolation_mode",
         ]
 
     data_source_type_handler = QuerySubscriptionDataSourceHandler
@@ -402,6 +405,9 @@ def _validate_group_by(self, value: Sequence[str] | None) -> Sequence[str] | Non
 
     @override
     def create_source(self, validated_data) -> QuerySubscription:
+        extrapolation_mode = validated_data.get("extrapolation_mode")
+        if extrapolation_mode is not None:
+            extrapolation_mode = ExtrapolationMode(extrapolation_mode)
         snuba_query = create_snuba_query(
             query_type=validated_data["query_type"],
             dataset=validated_data["dataset"],
@@ -412,6 +418,7 @@ def create_source(self, validated_data) -> QuerySubscription:
             environment=validated_data["environment"],
             event_types=validated_data["event_types"],
             group_by=validated_data.get("group_by"),
+            extrapolation_mode=extrapolation_mode,
         )
         return create_snuba_subscription(
             project=self.context["project"],

src/sentry/snuba/subscriptions.py

@@ -7,7 +7,12 @@
 from sentry.models.environment import Environment
 from sentry.models.project import Project
 from sentry.snuba.dataset import Dataset
-from sentry.snuba.models import QuerySubscription, SnubaQuery, SnubaQueryEventType
+from sentry.snuba.models import (
+    ExtrapolationMode,
+    QuerySubscription,
+    SnubaQuery,
+    SnubaQueryEventType,
+)
 from sentry.snuba.tasks import (
     create_subscription_in_snuba,
     delete_subscription_from_snuba,
@@ -27,6 +32,7 @@ def create_snuba_query(
     environment: Environment | None,
     event_types: Collection[SnubaQueryEventType.EventType] = (),
     group_by: Sequence[str] | None = None,
+    extrapolation_mode: ExtrapolationMode | None = None,
 ):
     """
     Constructs a SnubaQuery which is the postgres representation of a query in snuba
@@ -52,6 +58,11 @@ def create_snuba_query(
         resolution=int(resolution.total_seconds()),
         environment=environment,
         group_by=group_by,
+        extrapolation_mode=(
+            extrapolation_mode.value
+            if extrapolation_mode is not None
+            else ExtrapolationMode.UNKNOWN.value
+        ),
     )
     if not event_types:
         if dataset == Dataset.Events:
@@ -78,6 +89,7 @@ def update_snuba_query(
     resolution,
     environment,
     event_types,
+    extrapolation_mode=None,
 ):
     """
     Updates a SnubaQuery. Triggers updates to any related QuerySubscriptions.
@@ -118,6 +130,11 @@ def update_snuba_query(
             time_window=int(time_window.total_seconds()),
             resolution=int(resolution.total_seconds()),
             environment=environment,
+            extrapolation_mode=(
+                extrapolation_mode
+                if extrapolation_mode is not None
+                else ExtrapolationMode.UNKNOWN.value
+            ),
         )
         if new_event_types:
             SnubaQueryEventType.objects.bulk_create(

tests/sentry/incidents/endpoints/test_organization_alert_rule_details.py

@@ -1416,6 +1416,24 @@ def test_change_name_of_existing_alert(self) -> None:
             )
         assert len(audit_log_entry) == 1
 
+    def test_invalid_extrapolation_mode(self) -> None:
+        self.create_member(
+            user=self.user, organization=self.organization, role="owner", teams=[self.team]
+        )
+        self.login_as(self.user)
+        alert_rule = self.alert_rule
+        # We need the IDs to force update instead of create, so we just get the rule using our own API. Like frontend would.
+        alert_rule_dict = deepcopy(self.alert_rule_dict)
+        alert_rule_dict["dataset"] = "events_analytics_platform"
+        alert_rule_dict["alertType"] = "eap_metrics"
+        alert_rule_dict["extrapolation_mode"] = "server_weighted"
+
+        with self.feature("organizations:incidents"):
+            resp = self.get_error_response(
+                self.organization.slug, alert_rule.id, status_code=400, **alert_rule_dict
+            )
+        assert resp.data[0] == "Invalid extrapolation mode for this alert type."
+
 
 class AlertRuleDetailsSlackPutEndpointTest(AlertRuleDetailsBase):
     method = "put"

tests/sentry/incidents/endpoints/test_organization_alert_rule_index.py

@@ -1705,6 +1705,15 @@ def test_generic_metrics_dataset_deprecation_validation(self) -> None:
             == "Creation of transaction-based alerts is disabled, as we migrate to the span dataset. Create span-based alerts (dataset: events_analytics_platform) with the is_transaction:true filter instead."
         )
 
+    def test_invalid_extrapolation_mode(self) -> None:
+        data = deepcopy(self.alert_rule_dict)
+        data["dataset"] = "events_analytics_platform"
+        data["alertType"] = "eap_metrics"
+        data["extrapolation_mode"] = "server_weighted"
+        with self.feature(["organizations:incidents", "organizations:performance-view"]):
+            resp = self.get_error_response(self.organization.slug, status_code=400, **data)
+        assert resp.data[0] == "server_weighted extrapolation mode is not supported for new alerts."
+
 
 @freeze_time()
 class AlertRuleCreateEndpointTestCrashRateAlert(AlertRuleIndexBase):

tests/sentry/incidents/endpoints/validators/test_validators.py

@@ -26,6 +26,7 @@
 )
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.models import (
+    ExtrapolationMode,
     QuerySubscription,
     QuerySubscriptionDataSourceHandler,
     SnubaQuery,
@@ -654,3 +655,74 @@ def test_transaction_dataset_deprecation_generic_metrics_update(self) -> None:
             with_feature("organizations:discover-saved-queries-deprecation"),
         ):
             update_validator.save()
+
+    def test_invalid_extrapolation_mode_create(self) -> None:
+        data = {
+            **self.valid_data,
+            "dataSources": [
+                {
+                    "queryType": SnubaQuery.Type.PERFORMANCE.value,
+                    "dataset": Dataset.EventsAnalyticsPlatform.value,
+                    "query": "test query",
+                    "aggregate": "count()",
+                    "timeWindow": 3600,
+                    "environment": self.environment.name,
+                    "eventTypes": [SnubaQueryEventType.EventType.TRACE_ITEM_SPAN.name.lower()],
+                    "extrapolation_mode": ExtrapolationMode.SERVER_WEIGHTED.value,
+                },
+            ],
+        }
+
+        validator = MetricIssueDetectorValidator(data=data, context=self.context)
+        assert validator.is_valid(), validator.errors
+        with self.assertRaisesMessage(
+            ValidationError,
+            expected_message="server_weighted extrapolation mode is not supported for new detectors.",
+        ):
+            validator.save()
+
+    def test_invalid_extrapolation_mode_update(self) -> None:
+        data = {
+            **self.valid_data,
+            "dataSources": [
+                {
+                    "queryType": SnubaQuery.Type.PERFORMANCE.value,
+                    "dataset": Dataset.EventsAnalyticsPlatform.value,
+                    "query": "test query",
+                    "aggregate": "count()",
+                    "timeWindow": 3600,
+                    "environment": self.environment.name,
+                    "eventTypes": [SnubaQueryEventType.EventType.TRACE_ITEM_SPAN.name.lower()],
+                    "extrapolation_mode": ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED.value,
+                },
+            ],
+        }
+
+        validator = MetricIssueDetectorValidator(data=data, context=self.context)
+        assert validator.is_valid(), validator.errors
+
+        detector = validator.save()
+
+        update_data = {
+            "dataSources": [
+                {
+                    "queryType": SnubaQuery.Type.PERFORMANCE.value,
+                    "dataset": Dataset.EventsAnalyticsPlatform.value,
+                    "query": "updated query",
+                    "aggregate": "count()",
+                    "timeWindow": 3600,
+                    "environment": self.environment.name,
+                    "eventTypes": [SnubaQueryEventType.EventType.TRACE_ITEM_SPAN.name.lower()],
+                    "extrapolation_mode": ExtrapolationMode.SERVER_WEIGHTED.value,
+                }
+            ],
+        }
+        update_validator = MetricIssueDetectorValidator(
+            instance=detector, data=update_data, context=self.context, partial=True
+        )
+        assert update_validator.is_valid(), update_validator.errors
+        with self.assertRaisesMessage(
+            ValidationError,
+            expected_message="Invalid extrapolation mode for this detector type.",
+        ):
+            update_validator.save()