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.13",
     "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/endpoints/organization_events.py

@@ -25,10 +25,11 @@
 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
+from sentry.exceptions import InvalidParams, InvalidSearchQuery
 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,18 +518,28 @@ def get_rpc_config():
                     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
+                )
+
                 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 +552,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

@@ -17,8 +17,10 @@
     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,
@@ -243,6 +245,14 @@ 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
+                )
+
                 if scoped_dataset == TraceMetrics:
                     # tracemetrics uses aggregate conditions
                     metric_name, metric_type, metric_unit = get_trace_metric_from_request(request)
@@ -257,6 +267,7 @@ def get_rpc_config():
                             "disableAggregateExtrapolation", "0"
                         )
                         == "1",
+                        extrapolation_mode=extrapolation_mode,
                     )
 
                 return SearchResolverConfig(
@@ -266,6 +277,7 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
+                    extrapolation_mode=extrapolation_mode,
                 )
 
             if top_events > 0:

src/sentry/api/endpoints/organization_events_timeseries.py

@@ -22,7 +22,9 @@
 )
 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,
@@ -222,6 +224,12 @@ 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
+
             if dataset == TraceMetrics:
                 # tracemetrics uses aggregate conditions
                 metric_name, metric_type, metric_unit = get_trace_metric_from_request(request)
@@ -235,6 +243,7 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
+                    extrapolation_mode=extrapolation_mode,
                 )
 
             return SearchResolverConfig(
@@ -244,6 +253,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
 
@@ -200,8 +201,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
@@ -213,11 +213,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,
         )
 
 
@@ -232,8 +228,7 @@ class ResolvedMetricAggregate(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
@@ -249,11 +244,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,
         )
 
 
@@ -286,8 +277,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`
@@ -349,8 +340,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,
         )
@@ -422,8 +413,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,
@@ -464,8 +455,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
             ),
         )
 
@@ -492,10 +483,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
@@ -48,6 +48,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:
+        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