feat: Add support for extrapolation modes in entity subscription

src/sentry/snuba/entity_subscription.py

@@ -8,6 +8,9 @@
 from typing import Any, TypedDict, Union
 
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import TimeSeriesRequest
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import (
+    ExtrapolationMode as ProtoExtrapolationMode,
+)
 from snuba_sdk import Column, Condition, Entity, Join, Op, Request
 
 from sentry import features
@@ -27,7 +30,7 @@
 from sentry.snuba.dataset import Dataset, EntityKey
 from sentry.snuba.metrics.extraction import MetricSpecType
 from sentry.snuba.metrics.naming_layer.mri import SessionMRI
-from sentry.snuba.models import SnubaQuery, SnubaQueryEventType
+from sentry.snuba.models import ExtrapolationMode, SnubaQuery, SnubaQueryEventType
 from sentry.snuba.ourlogs import OurLogs
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.rpc_dataset_common import RPCBase
@@ -67,6 +70,14 @@
     "timestamp.to_day",
 }
 
+# Mapping from model ExtrapolationMode to proto ExtrapolationMode
+MODEL_TO_PROTO_EXTRAPOLATION_MODE = {
+    ExtrapolationMode.UNKNOWN: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED,
+    ExtrapolationMode.NONE: ProtoExtrapolationMode.EXTRAPOLATION_MODE_NONE,
+    ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED,
+    ExtrapolationMode.SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY,
+}
+
 
 def apply_dataset_query_conditions(
     query_type: SnubaQuery.Type,
@@ -106,6 +117,7 @@ def apply_dataset_query_conditions(
 
 class _EntitySpecificParams(TypedDict, total=False):
     org_id: int
+    extrapolation_mode: ExtrapolationMode | None
     event_types: list[SnubaQueryEventType.EventType] | None
 
 
@@ -265,9 +277,11 @@ def __init__(
         self.aggregate = aggregate
         self.event_types = None
         self.time_window = time_window
+        self.extrapolation_mode = None
         if extra_fields:
             self.org_id = extra_fields.get("org_id")
             self.event_types = extra_fields.get("event_types")
+            self.extrapolation_mode = extra_fields.get("extrapolation_mode")
 
     def build_rpc_request(
         self,
@@ -301,8 +315,19 @@ def build_rpc_request(
             end=now,
             granularity_secs=self.time_window,
         )
+
+        # Convert model ExtrapolationMode to proto ExtrapolationMode
+        proto_extrapolation_mode = None
+        if self.extrapolation_mode is not None:
+            model_mode = ExtrapolationMode(self.extrapolation_mode)
+            proto_extrapolation_mode = MODEL_TO_PROTO_EXTRAPOLATION_MODE.get(model_mode)
+
         search_resolver = dataset_module.get_resolver(
-            snuba_params, SearchResolverConfig(stable_timestamp_quantization=False)
+            snuba_params,
+            SearchResolverConfig(
+                stable_timestamp_quantization=False,
+                extrapolation_mode=proto_extrapolation_mode,
+            ),
         )
 
         rpc_request, _, _ = dataset_module.get_timeseries_query(
@@ -641,6 +666,7 @@ def get_entity_subscription_from_snuba_query(
         extra_fields={
             "org_id": organization_id,
             "event_types": snuba_query.event_types,
+            "extrapolation_mode": ExtrapolationMode(snuba_query.extrapolation_mode),
         },
     )
 

src/sentry/snuba/snuba_query_validator.py

@@ -284,6 +284,7 @@ def _validate_query(self, data):
                 extra_fields={
                     "org_id": projects[0].organization_id,
                     "event_types": data.get("event_types"),
+                    "extrapolation_mode": data.get("extrapolation_mode"),
                 },
             )
         except UnsupportedQuerySubscription as e:

src/sentry/snuba/tasks.py

@@ -19,7 +19,7 @@
     get_entity_subscription,
     get_entity_subscription_from_snuba_query,
 )
-from sentry.snuba.models import QuerySubscription, SnubaQuery
+from sentry.snuba.models import ExtrapolationMode, QuerySubscription, SnubaQuery
 from sentry.snuba.utils import build_query_strings
 from sentry.tasks.base import instrumented_task
 from sentry.taskworker.namespaces import alerts_tasks
@@ -133,6 +133,9 @@ def update_subscription_in_snuba(
             extra_fields={
                 "org_id": subscription.project.organization_id,
                 "event_types": subscription.snuba_query.event_types,
+                "extrapolation_mode": ExtrapolationMode(
+                    subscription.snuba_query.extrapolation_mode
+                ),
             },
         )
         old_entity_key = (

tests/sentry/snuba/test_entity_subscriptions.py

@@ -1,4 +1,7 @@
 import pytest
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import (
+    ExtrapolationMode as ProtoExtrapolationMode,
+)
 from snuba_sdk import And, Column, Condition, Entity, Function, Join, Op, Relationship
 
 from sentry.exceptions import InvalidQuerySubscription, UnsupportedQuerySubscription
@@ -20,7 +23,7 @@
     get_entity_subscription_from_snuba_query,
 )
 from sentry.snuba.metrics.naming_layer.mri import SessionMRI
-from sentry.snuba.models import SnubaQuery
+from sentry.snuba.models import ExtrapolationMode, SnubaQuery
 from sentry.testutils.cases import TestCase
 from sentry.testutils.skips import requires_snuba
 
@@ -456,6 +459,38 @@ def test_get_entity_subscription_for_eap_rpc_query(self) -> None:
         assert rpc_timeseries_request.granularity_secs == 3600
         assert rpc_timeseries_request.filter.comparison_filter.value.val_str == "http.client"
         assert rpc_timeseries_request.expressions[0].aggregation.label == "count(span.duration)"
+        assert (
+            rpc_timeseries_request.expressions[0].aggregation.extrapolation_mode
+            == ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
+        )
+
+    def test_get_entity_subscription_for_eap_with_extrapolation_mode(self) -> None:
+        aggregate = "count(span.duration)"
+        query = "span.op:http.client"
+
+        # Test with SERVER_WEIGHTED extrapolation mode
+        entity_subscription = get_entity_subscription(
+            query_type=SnubaQuery.Type.PERFORMANCE,
+            dataset=Dataset.EventsAnalyticsPlatform,
+            aggregate=aggregate,
+            time_window=3600,
+            extra_fields={
+                "org_id": self.organization.id,
+                "extrapolation_mode": ExtrapolationMode.SERVER_WEIGHTED,
+            },
+        )
+        assert isinstance(entity_subscription, PerformanceSpansEAPRpcEntitySubscription)
+        assert entity_subscription.extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED
+
+        rpc_timeseries_request = entity_subscription.build_rpc_request(
+            query, [self.project.id], None
+        )
+
+        # Verify the extrapolation mode is passed to the RPC request
+        assert (
+            rpc_timeseries_request.expressions[0].aggregation.extrapolation_mode
+            == ProtoExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY
+        )
 
 
 class GetEntitySubscriptionFromSnubaQueryTest(TestCase):