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

src/sentry/incidents/endpoints/organization_alert_rule_details.py

@@ -49,6 +49,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
@@ -102,6 +104,25 @@ def fetch_alert_rule(
     return Response(serialized_rule)
 
 
+def _is_invalid_extrapolation_mode(old_extrapolation_mode, new_extrapolation_mode) -> bool:
+    if (
+        new_extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED.name.lower()
+        and old_extrapolation_mode != ExtrapolationMode.SERVER_WEIGHTED.name.lower()
+    ):
+        return True
+    if (
+        new_extrapolation_mode == ExtrapolationMode.NONE.name.lower()
+        and old_extrapolation_mode != ExtrapolationMode.NONE.name.lower()
+    ):
+        return True
+    if (
+        new_extrapolation_mode == ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED.name.lower()
+        or new_extrapolation_mode == ExtrapolationMode.UNKNOWN.name.lower()
+    ):
+        return False
+    return False
+
+
 def update_alert_rule(
     request: Request, organization: Organization, alert_rule: AlertRule
 ) -> Response:
@@ -121,6 +142,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,
@@ -100,6 +101,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
@@ -162,6 +167,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 alerts."
+            )
+
     def get_quota(self) -> DetectorQuota:
         organization = self.context.get("organization")
         request = self.context.get("request")
@@ -201,6 +212,31 @@ def is_editing_transaction_dataset(
             return True
         return False
 
+    def is_invalid_extrapolation_mode(
+        self, snuba_query: SnubaQuery, data_source: SnubaQueryDataSourceType
+    ) -> bool:
+        if data_source.get("dataset") == Dataset.EventsAnalyticsPlatform:
+            new_extrapolation_mode = data_source.get(
+                "extrapolation_mode", snuba_query.extrapolation_mode
+            )
+            old_extrapolation_mode = snuba_query.extrapolation_mode
+            if (
+                new_extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED.value
+                and old_extrapolation_mode != ExtrapolationMode.SERVER_WEIGHTED.value
+            ):
+                return True
+            if (
+                new_extrapolation_mode == ExtrapolationMode.NONE.value
+                and old_extrapolation_mode != ExtrapolationMode.NONE.value
+            ):
+                return True
+            if (
+                new_extrapolation_mode == ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED.value
+                or new_extrapolation_mode == ExtrapolationMode.UNKNOWN.value
+            ):
+                return False
+        return False
+
     def update_data_source(self, instance: Detector, data_source: SnubaQueryDataSourceType):
         try:
             source_instance = DataSource.objects.get(detector=instance)
@@ -224,6 +260,9 @@ 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."
             )
 
+        if self.is_invalid_extrapolation_mode(snuba_query, data_source):
+            raise serializers.ValidationError("Invalid extrapolation mode for this alert type.")
+
         update_snuba_query(
             snuba_query=snuba_query,
             query_type=data_source.get("query_type", snuba_query.type),
@@ -234,6 +273,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]):
@@ -267,6 +309,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:
+            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,
@@ -617,3 +618,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 alerts.",
+        ):
+            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 alert type.",
+        ):
+            update_validator.save()