feat: Make extrapolation mode an enum

pyproject.toml

@@ -83,7 +83,7 @@ dependencies = [
     "sentry-forked-email-reply-parser>=0.5.12.post1",
     "sentry-kafka-schemas>=2.1.15",
     "sentry-ophio>=1.1.3",
-    "sentry-protos>=0.4.2",
+    "sentry-protos>=0.4.3",
     "sentry-redis-tools>=0.5.0",
     "sentry-relay>=0.9.19",
     "sentry-sdk[http2]>=2.43.0",

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.ValueType | 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

@@ -517,18 +517,22 @@ def get_rpc_config():
                     request.GET.get("disableAggregateExtrapolation", "0") == "1"
                 )
 
+                extrapolation_mode = self.get_extrapolation_mode(request)
+
                 if scoped_dataset == Spans:
                     return SearchResolverConfig(
                         auto_fields=True,
                         use_aggregate_conditions=use_aggregate_conditions,
                         fields_acl=FieldsACL(functions={"time_spent_percentage"}),
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
                 elif scoped_dataset == OurLogs:
                     # ourlogs doesn't have use aggregate conditions
                     return SearchResolverConfig(
                         use_aggregate_conditions=False,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
                 elif scoped_dataset == TraceMetrics:
                     # tracemetrics uses aggregate conditions
@@ -541,24 +545,28 @@ def get_rpc_config():
                         use_aggregate_conditions=use_aggregate_conditions,
                         auto_fields=True,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
                 elif scoped_dataset == ProfileFunctions:
                     # profile_functions uses aggregate conditions
                     return SearchResolverConfig(
                         use_aggregate_conditions=use_aggregate_conditions,
                         auto_fields=True,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
                 elif scoped_dataset == uptime_results.UptimeResults:
                     return SearchResolverConfig(
                         use_aggregate_conditions=use_aggregate_conditions,
                         auto_fields=True,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
                 else:
                     return SearchResolverConfig(
                         use_aggregate_conditions=use_aggregate_conditions,
                         disable_aggregate_extrapolation=disable_aggregate_extrapolation,
+                        extrapolation_mode=extrapolation_mode,
                     )
 
             if snuba_params.sampling_mode == "HIGHEST_ACCURACY_FLEX_TIME":

src/sentry/api/endpoints/organization_events_stats.py

@@ -243,6 +243,8 @@ def get_rpc_config():
                 if scoped_dataset not in RPC_DATASETS:
                     raise NotImplementedError
 
+                extrapolation_mode = self.get_extrapolation_mode(request)
+
                 if scoped_dataset == TraceMetrics:
                     # tracemetrics uses aggregate conditions
                     metric_name, metric_type, metric_unit = get_trace_metric_from_request(request)
@@ -257,6 +259,7 @@ def get_rpc_config():
                             "disableAggregateExtrapolation", "0"
                         )
                         == "1",
+                        extrapolation_mode=extrapolation_mode,
                     )
 
                 return SearchResolverConfig(
@@ -266,6 +269,7 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
+                    extrapolation_mode=extrapolation_mode,
                 )
 
             if top_events > 0:

src/sentry/api/endpoints/organization_events_timeseries.py

@@ -222,6 +222,8 @@ def get_rpc_config():
             if dataset not in RPC_DATASETS:
                 raise NotImplementedError
 
+            extrapolation_mode = self.get_extrapolation_mode(request)
+
             if dataset == TraceMetrics:
                 # tracemetrics uses aggregate conditions
                 metric_name, metric_type, metric_unit = get_trace_metric_from_request(request)
@@ -235,6 +237,7 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
+                    extrapolation_mode=extrapolation_mode,
                 )
 
             return SearchResolverConfig(
@@ -244,6 +247,7 @@ def get_rpc_config():
                     "disableAggregateExtrapolation", "0"
                 )
                 == "1",
+                extrapolation_mode=extrapolation_mode,
             )
 
         if top_events > 0:

src/sentry/search/eap/columns.py

@@ -22,6 +22,7 @@
 from sentry.api.event_search import SearchFilter
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap import constants
+from sentry.search.eap.extrapolation_mode import resolve_extrapolation_mode
 from sentry.search.eap.types import EAPResponse, MetricType, SearchResolverConfig
 from sentry.search.events.types import SnubaParams
 
@@ -207,8 +208,7 @@ class ResolvedAggregate(ResolvedFunction):
 
     # The internal rpc alias for this column
     internal_name: Function.ValueType
-    # Whether to enable extrapolation
-    extrapolation: bool = True
+    extrapolation_mode: ExtrapolationMode.ValueType
     is_aggregate: bool = field(default=True, init=False)
     # Only for aggregates, we only support functions with 1 argument right now
     argument: AttributeKey | None = None
@@ -220,11 +220,7 @@ def proto_definition(self) -> AttributeAggregation:
             aggregate=self.internal_name,
             key=self.argument,
             label=self.public_alias,
-            extrapolation_mode=(
-                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-                if self.extrapolation
-                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
-            ),
+            extrapolation_mode=self.extrapolation_mode,
         )
 
 
@@ -239,8 +235,7 @@ class ResolvedTraceMetricAggregate(ResolvedAggregate):
 class ResolvedConditionalAggregate(ResolvedFunction):
     # The internal rpc alias for this column
     internal_name: Function.ValueType
-    # Whether to enable extrapolation
-    extrapolation: bool = True
+    extrapolation_mode: ExtrapolationMode.ValueType
     # The condition to filter on
     filter: TraceItemFilter
     # The attribute to conditionally aggregate on
@@ -256,11 +251,7 @@ def proto_definition(self) -> AttributeConditionalAggregation:
             key=self.key,
             filter=self.filter,
             label=self.public_alias,
-            extrapolation_mode=(
-                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-                if self.extrapolation
-                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
-            ),
+            extrapolation_mode=self.extrapolation_mode,
         )
 
 
@@ -293,8 +284,8 @@ class FunctionDefinition:
     infer_search_type_from_arguments: bool = True
     # The internal rpc type for this function, optional as it can mostly be inferred from search_type
     internal_type: AttributeKey.Type.ValueType | None = None
-    # Whether to request extrapolation or not, should be true for all functions except for _sample functions for debugging
-    extrapolation: bool = True
+    # Extrapolation mode to be used
+    extrapolation_mode_override: ExtrapolationMode.ValueType | None = None
     # Processor is the function run in the post process step to transform a row into the final result
     processor: Callable[[Any], Any] | None = None
     # if a function is private, assume it can't be used unless it's provided in `SearchResolverConfig.functions_acl`
@@ -356,8 +347,8 @@ def resolve(
             search_type=search_type,
             internal_type=self.internal_type,
             processor=self.processor,
-            extrapolation=(
-                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
+            extrapolation_mode=resolve_extrapolation_mode(
+                search_config, self.extrapolation_mode_override
             ),
             argument=resolved_attribute,
         )
@@ -396,8 +387,8 @@ def resolve(
             search_type=search_type,
             internal_type=self.internal_type,
             processor=self.processor,
-            extrapolation=(
-                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
+            extrapolation_mode=resolve_extrapolation_mode(
+                search_config, self.extrapolation_mode_override
             ),
             argument=resolved_attribute,
             metric_name=metric_name,
@@ -438,8 +429,8 @@ def resolve(
             filter=aggregate_filter,
             key=key,
             processor=self.processor,
-            extrapolation=(
-                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
+            extrapolation_mode=resolve_extrapolation_mode(
+                search_config, self.extrapolation_mode_override
             ),
         )
 
@@ -466,10 +457,8 @@ def resolve(
         search_config: SearchResolverConfig,
     ) -> ResolvedFormula:
         resolver_settings = ResolverSettings(
-            extrapolation_mode=(
-                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-                if self.extrapolation and not search_config.disable_aggregate_extrapolation
-                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
+            extrapolation_mode=resolve_extrapolation_mode(
+                search_config, self.extrapolation_mode_override
             ),
             snuba_params=snuba_params,
             query_result_cache=query_result_cache,
@@ -501,10 +490,8 @@ def resolve(
         search_config: SearchResolverConfig,
     ) -> ResolvedFormula:
         resolver_settings = ResolverSettings(
-            extrapolation_mode=(
-                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
-                if self.extrapolation and not search_config.disable_aggregate_extrapolation
-                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
+            extrapolation_mode=resolve_extrapolation_mode(
+                search_config, self.extrapolation_mode_override
             ),
             snuba_params=snuba_params,
             query_result_cache=query_result_cache,

src/sentry/search/eap/constants.py

@@ -3,7 +3,7 @@
 from sentry_protos.snuba.v1.downsampled_storage_pb2 import DownsampledStorageConfig
 from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import AggregationComparisonFilter, Column
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
-from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, ExtrapolationMode
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import ComparisonFilter
 
 from sentry.search.eap.types import SupportedTraceItemType
@@ -50,6 +50,13 @@
     "<=": AggregationComparisonFilter.OP_LESS_THAN_OR_EQUALS,
 }
 
+EXTRAPOLATION_MODE_MAP = {
+    "sampleWeighted": ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED,
+    "serverOnly": ExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY,
+    "unspecified": ExtrapolationMode.EXTRAPOLATION_MODE_UNSPECIFIED,
+    "none": ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
+}
+
 SearchType = (
     SizeUnit
     | DurationUnit

src/sentry/search/eap/extrapolation_mode.py

@@ -0,0 +1,19 @@
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode
+
+from sentry.search.eap.types import SearchResolverConfig
+
+
+def resolve_extrapolation_mode(
+    search_config: SearchResolverConfig,
+    argument_override: ExtrapolationMode.ValueType | None = None,
+) -> ExtrapolationMode.ValueType:
+    if search_config.disable_aggregate_extrapolation:
+        return ExtrapolationMode.EXTRAPOLATION_MODE_NONE
+
+    if argument_override is not None:
+        return argument_override
+
+    if search_config.extrapolation_mode is not None:
+        return search_config.extrapolation_mode
+
+    return ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED

src/sentry/search/eap/spans/aggregates.py

@@ -3,6 +3,7 @@
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import (
     AttributeKey,
     AttributeValue,
+    ExtrapolationMode,
     Function,
     StrArray,
 )
@@ -451,7 +452,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
         ],
         aggregate_resolver=resolve_bounded_sample,
         processor=lambda x: x > 0,
-        extrapolation=False,
+        extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
     ),
 }
 
@@ -508,7 +509,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
                 default_arg="span.duration",
             )
         ],
-        extrapolation=False,
+        extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
     ),
     "count": AggregateDefinition(
         internal_function=Function.FUNCTION_COUNT,
@@ -549,7 +550,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
                 default_arg="span.duration",
             )
         ],
-        extrapolation=False,
+        extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
     ),
     "p50": AggregateDefinition(
         internal_function=Function.FUNCTION_P50,
@@ -585,7 +586,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace
                 default_arg="span.duration",
             )
         ],
-        extrapolation=False,
+        extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE,
     ),
     "p75": AggregateDefinition(
         internal_function=Function.FUNCTION_P75,

src/sentry/search/eap/types.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict
 
 from sentry_protos.snuba.v1.request_common_pb2 import PageToken
-from sentry_protos.snuba.v1.trace_item_attribute_pb2 import Reliability
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode, Reliability
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
 
 from sentry.search.events.types import EventsResponse
@@ -31,6 +31,7 @@ class SearchResolverConfig:
     fields_acl: FieldsACL = field(default_factory=lambda: FieldsACL())
     # If set to True, do not extrapolate any values regardless of individual aggregate settings
     disable_aggregate_extrapolation: bool = False
+    extrapolation_mode: ExtrapolationMode.ValueType | None = None
     # Whether to set the timestamp granularities to stable buckets
     stable_timestamp_quantization: bool = True