ref: Reuse get extrapolation mode func

src/sentry/api/bases/organization_events.py

@@ -13,6 +13,7 @@
 from django.utils import timezone
 from rest_framework.exceptions import ParseError, ValidationError
 from rest_framework.request import Request
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode
 from sentry_relay.consts import SPAN_STATUS_CODE_TO_NAME
 
 from sentry import features, quotas
@@ -37,7 +38,11 @@
 from sentry.models.organization import Organization
 from sentry.models.project import Project
 from sentry.models.team import Team
-from sentry.search.eap.constants import SAMPLING_MODE_MAP, VALID_GRANULARITIES
+from sentry.search.eap.constants import (
+    EXTRAPOLATION_MODE_MAP,
+    SAMPLING_MODE_MAP,
+    VALID_GRANULARITIES,
+)
 from sentry.search.events.constants import DURATION_UNITS, SIZE_UNITS
 from sentry.search.events.fields import get_function_alias
 from sentry.search.events.types import SAMPLING_MODES, SnubaParams
@@ -520,6 +525,15 @@ def validate_comparison_delta(
             if retention and comparison_start < timezone.now() - timedelta(days=retention):
                 raise ValidationError("Comparison period is outside your retention window")
 
+    def get_extrapolation_mode(self, request: Request) -> ExtrapolationMode | None:
+        requested_mode = request.GET.get("extrapolationMode", None)
+        if requested_mode is not None and requested_mode not in EXTRAPOLATION_MODE_MAP:
+            raise InvalidSearchQuery(f"Unknown extrapolation mode: {requested_mode}")
+
+        extrapolation_mode = EXTRAPOLATION_MODE_MAP[requested_mode] if requested_mode else None
+
+        return extrapolation_mode
+
     def get_event_stats_data(
         self,
         request: Request,

src/sentry/api/endpoints/organization_events.py

@@ -25,11 +25,10 @@
 from sentry.apidocs.parameters import GlobalParams, OrganizationParams, VisibilityParams
 from sentry.apidocs.utils import inline_sentry_response_serializer
 from sentry.discover.models import DiscoverSavedQuery, DiscoverSavedQueryTypes
-from sentry.exceptions import InvalidParams, InvalidSearchQuery
+from sentry.exceptions import InvalidParams
 from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
 from sentry.models.organization import Organization
 from sentry.ratelimits.config import RateLimitConfig
-from sentry.search.eap.constants import EXTRAPOLATION_MODE_MAP
 from sentry.search.eap.trace_metrics.config import (
     TraceMetricsSearchResolverConfig,
     get_trace_metric_from_request,
@@ -517,14 +516,7 @@ def get_rpc_config():
                 disable_aggregate_extrapolation = (
                     request.GET.get("disableAggregateExtrapolation", "0") == "1"
                 )
-
-                requested_mode = request.GET.get("extrapolationMode", None)
-                if requested_mode is not None and requested_mode not in EXTRAPOLATION_MODE_MAP:
-                    raise InvalidSearchQuery(f"Unknown extrapolation mode: {requested_mode}")
-
-                extrapolation_mode = (
-                    EXTRAPOLATION_MODE_MAP[requested_mode] if requested_mode else None
-                )
+                extrapolation_mode = self.get_extrapolation_mode(request)
 
                 if scoped_dataset == Spans:
                     return SearchResolverConfig(

src/sentry/api/endpoints/organization_events_stats.py

@@ -17,10 +17,8 @@
     transform_query_columns_for_error_upsampling,
 )
 from sentry.constants import MAX_TOP_EVENTS
-from sentry.exceptions import InvalidSearchQuery
 from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
 from sentry.models.organization import Organization
-from sentry.search.eap.constants import EXTRAPOLATION_MODE_MAP
 from sentry.search.eap.trace_metrics.config import (
     TraceMetricsSearchResolverConfig,
     get_trace_metric_from_request,
@@ -245,13 +243,7 @@ def get_rpc_config():
                 if scoped_dataset not in RPC_DATASETS:
                     raise NotImplementedError
 
-                requested_mode = request.GET.get("extrapolationMode", None)
-                if requested_mode is not None and requested_mode not in EXTRAPOLATION_MODE_MAP:
-                    raise InvalidSearchQuery(f"Unknown extrapolation mode: {requested_mode}")
-
-                extrapolation_mode = (
-                    EXTRAPOLATION_MODE_MAP[requested_mode] if requested_mode else None
-                )
+                extrapolation_mode = self.get_extrapolation_mode(request)
 
                 if scoped_dataset == TraceMetrics:
                     # tracemetrics uses aggregate conditions

src/sentry/api/endpoints/organization_events_timeseries.py

@@ -22,9 +22,7 @@
 )
 from sentry.api.utils import handle_query_errors
 from sentry.constants import MAX_TOP_EVENTS
-from sentry.exceptions import InvalidSearchQuery
 from sentry.models.organization import Organization
-from sentry.search.eap.constants import EXTRAPOLATION_MODE_MAP
 from sentry.search.eap.trace_metrics.config import (
     TraceMetricsSearchResolverConfig,
     get_trace_metric_from_request,
@@ -224,11 +222,7 @@ def get_rpc_config():
             if dataset not in RPC_DATASETS:
                 raise NotImplementedError
 
-            requested_mode = request.GET.get("extrapolationMode", None)
-            if requested_mode is not None and requested_mode not in EXTRAPOLATION_MODE_MAP:
-                raise InvalidSearchQuery(f"Unknown extrapolation mode: {requested_mode}")
-
-            extrapolation_mode = EXTRAPOLATION_MODE_MAP[requested_mode] if requested_mode else None
+            extrapolation_mode = self.get_extrapolation_mode(request)
 
             if dataset == TraceMetrics:
                 # tracemetrics uses aggregate conditions