Revert "feat: Make extrapolation mode an enum (#103071)"

This reverts commit a093a68.

Co-authored-by: shruthilayaj <[email protected]>

Author: getsentry-bot

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.3",
+    "sentry-protos>=0.4.2",
     "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,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,
@@ -518,28 +517,18 @@ 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
@@ -552,28 +541,24 @@ 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,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,14 +243,6 @@ 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)
@@ -267,7 +257,6 @@ def get_rpc_config():
                             "disableAggregateExtrapolation", "0"
                         )
                         == "1",
-                        extrapolation_mode=extrapolation_mode,
                     )
 
                 return SearchResolverConfig(
@@ -277,7 +266,6 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
-                    extrapolation_mode=extrapolation_mode,
                 )
 
             if top_events > 0:

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,12 +222,6 @@ 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)
@@ -243,7 +235,6 @@ def get_rpc_config():
                         "disableAggregateExtrapolation", "0"
                     )
                     == "1",
-                    extrapolation_mode=extrapolation_mode,
                 )
 
             return SearchResolverConfig(
@@ -253,7 +244,6 @@ def get_rpc_config():
                     "disableAggregateExtrapolation", "0"
                 )
                 == "1",
-                extrapolation_mode=extrapolation_mode,
             )
 
         if top_events > 0:

src/sentry/search/eap/columns.py

@@ -22,7 +22,6 @@
 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
 
@@ -208,7 +207,8 @@ class ResolvedAggregate(ResolvedFunction):
 
     # The internal rpc alias for this column
     internal_name: Function.ValueType
-    extrapolation_mode: ExtrapolationMode.ValueType
+    # Whether to enable extrapolation
+    extrapolation: bool = True
     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,7 +220,11 @@ def proto_definition(self) -> AttributeAggregation:
             aggregate=self.internal_name,
             key=self.argument,
             label=self.public_alias,
-            extrapolation_mode=self.extrapolation_mode,
+            extrapolation_mode=(
+                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
+                if self.extrapolation
+                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
+            ),
         )
 
 
@@ -235,7 +239,8 @@ class ResolvedTraceMetricAggregate(ResolvedAggregate):
 class ResolvedConditionalAggregate(ResolvedFunction):
     # The internal rpc alias for this column
     internal_name: Function.ValueType
-    extrapolation_mode: ExtrapolationMode.ValueType
+    # Whether to enable extrapolation
+    extrapolation: bool = True
     # The condition to filter on
     filter: TraceItemFilter
     # The attribute to conditionally aggregate on
@@ -251,7 +256,11 @@ def proto_definition(self) -> AttributeConditionalAggregation:
             key=self.key,
             filter=self.filter,
             label=self.public_alias,
-            extrapolation_mode=self.extrapolation_mode,
+            extrapolation_mode=(
+                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
+                if self.extrapolation
+                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
+            ),
         )
 
 
@@ -284,8 +293,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
-    # Extrapolation mode to be used
-    extrapolation_mode_override: ExtrapolationMode.ValueType | None = None
+    # Whether to request extrapolation or not, should be true for all functions except for _sample functions for debugging
+    extrapolation: bool = True
     # 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`
@@ -347,8 +356,8 @@ def resolve(
             search_type=search_type,
             internal_type=self.internal_type,
             processor=self.processor,
-            extrapolation_mode=resolve_extrapolation_mode(
-                search_config, self.extrapolation_mode_override
+            extrapolation=(
+                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
             ),
             argument=resolved_attribute,
         )
@@ -387,8 +396,8 @@ def resolve(
             search_type=search_type,
             internal_type=self.internal_type,
             processor=self.processor,
-            extrapolation_mode=resolve_extrapolation_mode(
-                search_config, self.extrapolation_mode_override
+            extrapolation=(
+                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
             ),
             argument=resolved_attribute,
             metric_name=metric_name,
@@ -429,8 +438,8 @@ def resolve(
             filter=aggregate_filter,
             key=key,
             processor=self.processor,
-            extrapolation_mode=resolve_extrapolation_mode(
-                search_config, self.extrapolation_mode_override
+            extrapolation=(
+                self.extrapolation if not search_config.disable_aggregate_extrapolation else False
             ),
         )
 
@@ -457,8 +466,10 @@ def resolve(
         search_config: SearchResolverConfig,
     ) -> ResolvedFormula:
         resolver_settings = ResolverSettings(
-            extrapolation_mode=resolve_extrapolation_mode(
-                search_config, self.extrapolation_mode_override
+            extrapolation_mode=(
+                ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED
+                if self.extrapolation and not search_config.disable_aggregate_extrapolation
+                else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
             ),
             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, ExtrapolationMode
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import ComparisonFilter
 
 from sentry.search.eap.types import SupportedTraceItemType
@@ -50,13 +50,6 @@
     "<=": 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

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

src/sentry/search/eap/spans/aggregates.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 ExtrapolationMode, Reliability
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import Reliability
 from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
 
 from sentry.search.events.types import EventsResponse
@@ -31,7 +31,6 @@ 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