From 15fe602001f22fec25b9a8668653f1f32a911455 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Mon, 10 Nov 2025 11:49:49 -0500 Subject: [PATCH 1/9] feat: Make extrapolation mode an enum --- pyproject.toml | 2 +- .../api/endpoints/organization_events.py | 16 ++++++- .../endpoints/organization_events_stats.py | 10 ++++ .../organization_events_timeseries.py | 10 ++++ src/sentry/search/eap/columns.py | 46 ++++++++++--------- src/sentry/search/eap/constants.py | 9 +++- src/sentry/search/eap/spans/aggregates.py | 9 ++-- src/sentry/search/eap/types.py | 6 ++- 8 files changed, 79 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 134df943ac9941..7f1fa4f1f2177b 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/sentry/api/endpoints/organization_events.py b/src/sentry/api/endpoints/organization_events.py index b2f8582221e00e..946f8190981c84 100644 --- a/src/sentry/api/endpoints/organization_events.py +++ b/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,27 @@ def get_rpc_config(): request.GET.get("disableAggregateExtrapolation", "0") == "1" ) + extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") + + if extrapolation_mode in EXTRAPOLATION_MODE_MAP: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + else: + raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + 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,6 +551,7 @@ 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 @@ -548,17 +559,20 @@ 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 == 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": diff --git a/src/sentry/api/endpoints/organization_events_stats.py b/src/sentry/api/endpoints/organization_events_stats.py index 9e2e9886468e81..699c9d004d7c9a 100644 --- a/src/sentry/api/endpoints/organization_events_stats.py +++ b/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,12 @@ def get_rpc_config(): if scoped_dataset not in RPC_DATASETS: raise NotImplementedError + extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") + if extrapolation_mode in EXTRAPOLATION_MODE_MAP: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + else: + raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + if scoped_dataset == TraceMetrics: # tracemetrics uses aggregate conditions metric_name, metric_type, metric_unit = get_trace_metric_from_request(request) @@ -257,6 +265,7 @@ def get_rpc_config(): "disableAggregateExtrapolation", "0" ) == "1", + extrapolation_mode=extrapolation_mode, ) return SearchResolverConfig( @@ -266,6 +275,7 @@ def get_rpc_config(): "disableAggregateExtrapolation", "0" ) == "1", + extrapolation_mode=extrapolation_mode, ) if top_events > 0: diff --git a/src/sentry/api/endpoints/organization_events_timeseries.py b/src/sentry/api/endpoints/organization_events_timeseries.py index 5c0c9eca9de9c2..2fb62be0045d0f 100644 --- a/src/sentry/api/endpoints/organization_events_timeseries.py +++ b/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 + extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") + if extrapolation_mode in EXTRAPOLATION_MODE_MAP: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + else: + raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + 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: diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index afd6b1ebb405c8..9084cb22e50cda 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -201,7 +201,9 @@ 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 = ( + ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ) 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 +215,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, ) @@ -226,7 +224,9 @@ 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 = ( + ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ) # The condition to filter on filter: TraceItemFilter # The attribute to conditionally aggregate on @@ -242,11 +242,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, ) @@ -280,7 +276,9 @@ class FunctionDefinition: # 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: ExtrapolationMode.ValueType = ( + ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ) # 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` @@ -342,8 +340,10 @@ 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=( + self.search_config.extrapolation_mode + if not search_config.disable_aggregate_extrapolation + else ExtrapolationMode.EXTRAPOLATION_MODE_NONE ), argument=resolved_attribute, ) @@ -407,8 +407,10 @@ 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=( + self.extrapolation_mode + if not search_config.disable_aggregate_extrapolation + else ExtrapolationMode.EXTRAPOLATION_MODE_NONE ), argument=resolved_attribute, ) @@ -446,8 +448,10 @@ def resolve( filter=aggregate_filter, key=key, processor=self.processor, - extrapolation=( - self.extrapolation if not search_config.disable_aggregate_extrapolation else False + extrapolation_mode=( + self.extrapolation_mode + if not search_config.disable_aggregate_extrapolation + else ExtrapolationMode.EXTRAPOLATION_MODE_NONE ), ) @@ -475,8 +479,8 @@ def resolve( ) -> ResolvedFormula: resolver_settings = ResolverSettings( extrapolation_mode=( - ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - if self.extrapolation and not search_config.disable_aggregate_extrapolation + self.extrapolation_mode + if self.extrapolation_mode and not search_config.disable_aggregate_extrapolation else ExtrapolationMode.EXTRAPOLATION_MODE_NONE ), snuba_params=snuba_params, diff --git a/src/sentry/search/eap/constants.py b/src/sentry/search/eap/constants.py index 264977ef408222..c014ad0e8f1599 100644 --- a/src/sentry/search/eap/constants.py +++ b/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 diff --git a/src/sentry/search/eap/spans/aggregates.py b/src/sentry/search/eap/spans/aggregates.py index 614077aa9d2e31..0a15ff9e9487d4 100644 --- a/src/sentry/search/eap/spans/aggregates.py +++ b/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=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=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=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=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, ), "p75": AggregateDefinition( internal_function=Function.FUNCTION_P75, diff --git a/src/sentry/search/eap/types.py b/src/sentry/search/eap/types.py index dbce61609cb7fc..b68339c20f63ab 100644 --- a/src/sentry/search/eap/types.py +++ b/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,10 @@ 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 + # If set to True, do not extrapolate any values regardless of individual aggregate settings + extrapolation_mode: ExtrapolationMode.ValueType = ( + ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ) # Whether to set the timestamp granularities to stable buckets stable_timestamp_quantization: bool = True From 03461a3d0666852068a86df95b7fcb0b9b1bb0f4 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Mon, 10 Nov 2025 12:01:17 -0500 Subject: [PATCH 2/9] fix --- src/sentry/api/endpoints/organization_events.py | 9 ++++----- src/sentry/api/endpoints/organization_events_stats.py | 9 +++++---- .../api/endpoints/organization_events_timeseries.py | 8 ++++---- src/sentry/search/eap/columns.py | 2 +- src/sentry/search/eap/types.py | 4 +--- 5 files changed, 15 insertions(+), 17 deletions(-) diff --git a/src/sentry/api/endpoints/organization_events.py b/src/sentry/api/endpoints/organization_events.py index 946f8190981c84..2b4729aa9059c8 100644 --- a/src/sentry/api/endpoints/organization_events.py +++ b/src/sentry/api/endpoints/organization_events.py @@ -518,12 +518,11 @@ def get_rpc_config(): request.GET.get("disableAggregateExtrapolation", "0") == "1" ) - extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") - - if extrapolation_mode in EXTRAPOLATION_MODE_MAP: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] - else: + extrapolation_mode = request.GET.get("extrapolationMode", None) + if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + elif extrapolation_mode: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] if scoped_dataset == Spans: return SearchResolverConfig( diff --git a/src/sentry/api/endpoints/organization_events_stats.py b/src/sentry/api/endpoints/organization_events_stats.py index 699c9d004d7c9a..5b78142dd08a11 100644 --- a/src/sentry/api/endpoints/organization_events_stats.py +++ b/src/sentry/api/endpoints/organization_events_stats.py @@ -245,11 +245,12 @@ def get_rpc_config(): if scoped_dataset not in RPC_DATASETS: raise NotImplementedError - extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") - if extrapolation_mode in EXTRAPOLATION_MODE_MAP: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] - else: + extrapolation_mode = request.GET.get("extrapolationMode", None) + + if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + elif extrapolation_mode: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] if scoped_dataset == TraceMetrics: # tracemetrics uses aggregate conditions diff --git a/src/sentry/api/endpoints/organization_events_timeseries.py b/src/sentry/api/endpoints/organization_events_timeseries.py index 2fb62be0045d0f..93f958773c6639 100644 --- a/src/sentry/api/endpoints/organization_events_timeseries.py +++ b/src/sentry/api/endpoints/organization_events_timeseries.py @@ -224,11 +224,11 @@ def get_rpc_config(): if dataset not in RPC_DATASETS: raise NotImplementedError - extrapolation_mode = request.GET.get("extrapolationMode", "sampleWeighted") - if extrapolation_mode in EXTRAPOLATION_MODE_MAP: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] - else: + extrapolation_mode = request.GET.get("extrapolationMode", None) + if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") + elif extrapolation_mode: + extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] if dataset == TraceMetrics: # tracemetrics uses aggregate conditions diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index 9084cb22e50cda..132795e5622ad4 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -341,7 +341,7 @@ def resolve( internal_type=self.internal_type, processor=self.processor, extrapolation_mode=( - self.search_config.extrapolation_mode + self.extrapolation_mode if not search_config.disable_aggregate_extrapolation else ExtrapolationMode.EXTRAPOLATION_MODE_NONE ), diff --git a/src/sentry/search/eap/types.py b/src/sentry/search/eap/types.py index b68339c20f63ab..7355fe4a8a89ab 100644 --- a/src/sentry/search/eap/types.py +++ b/src/sentry/search/eap/types.py @@ -32,9 +32,7 @@ class SearchResolverConfig: # If set to True, do not extrapolate any values regardless of individual aggregate settings disable_aggregate_extrapolation: bool = False # If set to True, do not extrapolate any values regardless of individual aggregate settings - extrapolation_mode: ExtrapolationMode.ValueType = ( - ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - ) + extrapolation_mode: ExtrapolationMode.ValueType | None = None # Whether to set the timestamp granularities to stable buckets stable_timestamp_quantization: bool = True From 5564f9ab5bab5f94b1b848f01eec05401864b528 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Mon, 10 Nov 2025 12:14:58 -0500 Subject: [PATCH 3/9] clean up extrapolation mode --- src/sentry/search/eap/columns.py | 25 +++++---------------- src/sentry/search/eap/extrapolation_mode.py | 19 ++++++++++++++++ 2 files changed, 24 insertions(+), 20 deletions(-) create mode 100644 src/sentry/search/eap/extrapolation_mode.py diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index 132795e5622ad4..187bdee300a335 100644 --- a/src/sentry/search/eap/columns.py +++ b/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, SearchResolverConfig from sentry.search.events.types import SnubaParams @@ -340,11 +341,7 @@ def resolve( search_type=search_type, internal_type=self.internal_type, processor=self.processor, - extrapolation_mode=( - self.extrapolation_mode - if not search_config.disable_aggregate_extrapolation - else ExtrapolationMode.EXTRAPOLATION_MODE_NONE - ), + extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), argument=resolved_attribute, ) @@ -407,11 +404,7 @@ def resolve( search_type=search_type, internal_type=self.internal_type, processor=self.processor, - extrapolation_mode=( - self.extrapolation_mode - if not search_config.disable_aggregate_extrapolation - else ExtrapolationMode.EXTRAPOLATION_MODE_NONE - ), + extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), argument=resolved_attribute, ) @@ -448,11 +441,7 @@ def resolve( filter=aggregate_filter, key=key, processor=self.processor, - extrapolation_mode=( - self.extrapolation_mode - if not search_config.disable_aggregate_extrapolation - else ExtrapolationMode.EXTRAPOLATION_MODE_NONE - ), + extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), ) @@ -478,11 +467,7 @@ def resolve( search_config: SearchResolverConfig, ) -> ResolvedFormula: resolver_settings = ResolverSettings( - extrapolation_mode=( - self.extrapolation_mode - if self.extrapolation_mode and not search_config.disable_aggregate_extrapolation - else ExtrapolationMode.EXTRAPOLATION_MODE_NONE - ), + extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), snuba_params=snuba_params, query_result_cache=query_result_cache, search_config=search_config, diff --git a/src/sentry/search/eap/extrapolation_mode.py b/src/sentry/search/eap/extrapolation_mode.py new file mode 100644 index 00000000000000..36c0d136981a0c --- /dev/null +++ b/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: + return search_config.extrapolation_mode + + return ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED From 0a33a34a755ad4df33db62cbfde9462c0b6db277 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Mon, 10 Nov 2025 12:24:40 -0500 Subject: [PATCH 4/9] make it an overrride --- src/sentry/search/eap/columns.py | 28 ++++++++++++----------- src/sentry/search/eap/spans/aggregates.py | 8 +++---- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index 187bdee300a335..b549850e9646d1 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -202,9 +202,7 @@ class ResolvedAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType # Whether to enable extrapolation - extrapolation_mode: ExtrapolationMode.ValueType = ( - ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - ) + extrapolation_mode_override: ExtrapolationMode.ValueType | None = None is_aggregate: bool = field(default=True, init=False) # Only for aggregates, we only support functions with 1 argument right now argument: AttributeKey | None = None @@ -225,9 +223,7 @@ class ResolvedConditionalAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType # Whether to enable extrapolation - extrapolation_mode: ExtrapolationMode.ValueType = ( - ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - ) + extrapolation_mode_override: ExtrapolationMode.ValueType | None = None # The condition to filter on filter: TraceItemFilter # The attribute to conditionally aggregate on @@ -277,9 +273,7 @@ class FunctionDefinition: # 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_mode: ExtrapolationMode.ValueType = ( - ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - ) + 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` @@ -341,7 +335,9 @@ def resolve( search_type=search_type, internal_type=self.internal_type, processor=self.processor, - extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), + extrapolation_mode=resolve_extrapolation_mode( + search_config, self.extrapolation_mode_override + ), argument=resolved_attribute, ) @@ -404,7 +400,9 @@ def resolve( search_type=search_type, internal_type=self.internal_type, processor=self.processor, - extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), + extrapolation_mode=resolve_extrapolation_mode( + search_config, self.extrapolation_mode_override + ), argument=resolved_attribute, ) @@ -441,7 +439,9 @@ def resolve( filter=aggregate_filter, key=key, processor=self.processor, - extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), + extrapolation_mode=resolve_extrapolation_mode( + search_config, self.extrapolation_mode_override + ), ) @@ -467,7 +467,9 @@ def resolve( search_config: SearchResolverConfig, ) -> ResolvedFormula: resolver_settings = ResolverSettings( - extrapolation_mode=resolve_extrapolation_mode(search_config, self.extrapolation_mode), + extrapolation_mode=resolve_extrapolation_mode( + search_config, self.extrapolation_mode_override + ), snuba_params=snuba_params, query_result_cache=query_result_cache, search_config=search_config, diff --git a/src/sentry/search/eap/spans/aggregates.py b/src/sentry/search/eap/spans/aggregates.py index 0a15ff9e9487d4..45fdd738b1dfbd 100644 --- a/src/sentry/search/eap/spans/aggregates.py +++ b/src/sentry/search/eap/spans/aggregates.py @@ -452,7 +452,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace ], aggregate_resolver=resolve_bounded_sample, processor=lambda x: x > 0, - extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, + extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, ), } @@ -509,7 +509,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace default_arg="span.duration", ) ], - extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, + extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, ), "count": AggregateDefinition( internal_function=Function.FUNCTION_COUNT, @@ -550,7 +550,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace default_arg="span.duration", ) ], - extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, + extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, ), "p50": AggregateDefinition( internal_function=Function.FUNCTION_P50, @@ -586,7 +586,7 @@ def resolve_bounded_sample(args: ResolvedArguments) -> tuple[AttributeKey, Trace default_arg="span.duration", ) ], - extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, + extrapolation_mode_override=ExtrapolationMode.EXTRAPOLATION_MODE_NONE, ), "p75": AggregateDefinition( internal_function=Function.FUNCTION_P75, From 670af023938e368080ec42e1421255f23cf493c9 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Mon, 10 Nov 2025 12:45:47 -0500 Subject: [PATCH 5/9] typing --- src/sentry/api/endpoints/organization_events.py | 12 +++++++----- .../api/endpoints/organization_events_stats.py | 11 ++++++----- .../api/endpoints/organization_events_timeseries.py | 10 +++++----- src/sentry/search/eap/columns.py | 4 ++-- 4 files changed, 20 insertions(+), 17 deletions(-) diff --git a/src/sentry/api/endpoints/organization_events.py b/src/sentry/api/endpoints/organization_events.py index 2b4729aa9059c8..383110359b3e79 100644 --- a/src/sentry/api/endpoints/organization_events.py +++ b/src/sentry/api/endpoints/organization_events.py @@ -518,11 +518,13 @@ def get_rpc_config(): request.GET.get("disableAggregateExtrapolation", "0") == "1" ) - extrapolation_mode = request.GET.get("extrapolationMode", None) - if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: - raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") - elif extrapolation_mode: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + 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( diff --git a/src/sentry/api/endpoints/organization_events_stats.py b/src/sentry/api/endpoints/organization_events_stats.py index 5b78142dd08a11..4ee0693a741ced 100644 --- a/src/sentry/api/endpoints/organization_events_stats.py +++ b/src/sentry/api/endpoints/organization_events_stats.py @@ -245,12 +245,13 @@ def get_rpc_config(): if scoped_dataset not in RPC_DATASETS: raise NotImplementedError - extrapolation_mode = request.GET.get("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}") - if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: - raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") - elif extrapolation_mode: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + extrapolation_mode = ( + EXTRAPOLATION_MODE_MAP[requested_mode] if requested_mode else None + ) if scoped_dataset == TraceMetrics: # tracemetrics uses aggregate conditions diff --git a/src/sentry/api/endpoints/organization_events_timeseries.py b/src/sentry/api/endpoints/organization_events_timeseries.py index 93f958773c6639..704e2b5f61c182 100644 --- a/src/sentry/api/endpoints/organization_events_timeseries.py +++ b/src/sentry/api/endpoints/organization_events_timeseries.py @@ -224,11 +224,11 @@ def get_rpc_config(): if dataset not in RPC_DATASETS: raise NotImplementedError - extrapolation_mode = request.GET.get("extrapolationMode", None) - if extrapolation_mode and extrapolation_mode not in EXTRAPOLATION_MODE_MAP: - raise InvalidSearchQuery(f"Unknown extrapolation mode: {extrapolation_mode}") - elif extrapolation_mode: - extrapolation_mode = EXTRAPOLATION_MODE_MAP[extrapolation_mode] + 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 diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index b549850e9646d1..e53ba607530569 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -202,7 +202,7 @@ class ResolvedAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType # Whether to enable extrapolation - extrapolation_mode_override: ExtrapolationMode.ValueType | None = None + 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 @@ -223,7 +223,7 @@ class ResolvedConditionalAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType # Whether to enable extrapolation - extrapolation_mode_override: ExtrapolationMode.ValueType | None = None + extrapolation_mode: ExtrapolationMode.ValueType # The condition to filter on filter: TraceItemFilter # The attribute to conditionally aggregate on From 2aa855ad243b04fa570b532666cb839d0244a8a2 Mon Sep 17 00:00:00 2001 From: "getsantry[bot]" <66042841+getsantry[bot]@users.noreply.github.com> Date: Mon, 10 Nov 2025 17:47:14 +0000 Subject: [PATCH 6/9] :snowflake: re-freeze requirements --- uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/uv.lock b/uv.lock index 518ea1b11abad6..e82377818f5e10 100644 --- a/uv.lock +++ b/uv.lock @@ -2169,7 +2169,7 @@ requires-dist = [ { name = "sentry-forked-email-reply-parser", specifier = ">=0.5.12.post1" }, { name = "sentry-kafka-schemas", specifier = ">=2.1.13" }, { name = "sentry-ophio", specifier = ">=1.1.3" }, - { name = "sentry-protos", specifier = ">=0.4.2" }, + { name = "sentry-protos", specifier = ">=0.4.3" }, { name = "sentry-redis-tools", specifier = ">=0.5.0" }, { name = "sentry-relay", specifier = ">=0.9.19" }, { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.43.0" }, @@ -2366,7 +2366,7 @@ wheels = [ [[package]] name = "sentry-protos" -version = "0.4.2" +version = "0.4.3" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ { name = "grpc-stubs", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -2374,7 +2374,7 @@ dependencies = [ { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] wheels = [ - { url = "https://pypi.devinfra.sentry.io/wheels/sentry_protos-0.4.2-py3-none-any.whl", hash = "sha256:4abbfeb7d5e810878786f341d1ba4b437c9b90d71d0992c3c3343b7e360aec4c" }, + { url = "https://pypi.devinfra.sentry.io/wheels/sentry_protos-0.4.3-py3-none-any.whl", hash = "sha256:ab0124f324c64faf1956cf4ce5219bc92b957c4cd79582f2ca36b00fa9b4d05c" }, ] [[package]] From 0a6028920fff66419de1370d67a7ddf6e87b9374 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Wed, 12 Nov 2025 13:13:00 -0500 Subject: [PATCH 7/9] add timeseries test --- ...st_organization_events_timeseries_spans.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/tests/snuba/api/endpoints/test_organization_events_timeseries_spans.py b/tests/snuba/api/endpoints/test_organization_events_timeseries_spans.py index 3654af08139eb1..879e50ba26d54a 100644 --- a/tests/snuba/api/endpoints/test_organization_events_timeseries_spans.py +++ b/tests/snuba/api/endpoints/test_organization_events_timeseries_spans.py @@ -3,8 +3,10 @@ import pytest from django.urls import reverse +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode from sentry.testutils.helpers.datetime import before_now +from sentry.utils.snuba_rpc import _make_rpc_requests from tests.snuba.api.endpoints.test_organization_events import OrganizationEventsEndpointTestBase from tests.snuba.api.endpoints.test_organization_events_span_indexed import KNOWN_PREFLIGHT_ID @@ -2372,6 +2374,65 @@ def test_disable_extrapolation(self) -> None: "interval": 3_600_000, } + @mock.patch( + "sentry.utils.snuba_rpc._make_rpc_requests", + wraps=_make_rpc_requests, + ) + def test_extrapolation_mode_server_only(self, mock_rpc_request) -> None: + event_counts = [6, 0, 6, 3, 0, 3] + spans = [] + for hour, count in enumerate(event_counts): + spans.extend( + [ + self.create_span( + { + "description": "foo", + "sentry_tags": {"status": "success"}, + "measurements": {"server_sample_rate": {"value": 0.1}}, + }, + start_ts=self.start + timedelta(hours=hour, minutes=minute), + ) + for minute in range(count) + ], + ) + self.store_spans(spans, is_eap=True) + response = self._do_request( + data={ + "start": self.start, + "end": self.end, + "interval": "1h", + "yAxis": "count()", + "project": self.project.id, + "dataset": "spans", + "extrapolationMode": "serverOnly", + }, + ) + + assert ( + mock_rpc_request.call_args.kwargs["timeseries_requests"][0] + .expressions[0] + .aggregation.extrapolation_mode + == ExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY + ) + + assert response.data["meta"] == { + "dataset": "spans", + "start": self.start.timestamp() * 1000, + "end": self.end.timestamp() * 1000, + } + assert len(response.data["timeSeries"]) == 1 + + timeseries = response.data["timeSeries"][0] + assert len(timeseries["values"]) == 6 + assert timeseries["yAxis"] == "count()" + + assert timeseries["meta"] == { + "dataScanned": "full", + "valueType": "integer", + "valueUnit": None, + "interval": 3_600_000, + } + def test_top_events_with_timestamp(self) -> None: """Users shouldn't groupby timestamp for top events""" response = self._do_request( From f2f3ef7ea41e690c70daff3924a3eb2e74f69c5a Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Wed, 12 Nov 2025 14:26:21 -0500 Subject: [PATCH 8/9] none check --- src/sentry/search/eap/extrapolation_mode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sentry/search/eap/extrapolation_mode.py b/src/sentry/search/eap/extrapolation_mode.py index 36c0d136981a0c..9d492f758e8425 100644 --- a/src/sentry/search/eap/extrapolation_mode.py +++ b/src/sentry/search/eap/extrapolation_mode.py @@ -13,7 +13,7 @@ def resolve_extrapolation_mode( if argument_override: return argument_override - if search_config.extrapolation_mode: + if search_config.extrapolation_mode is not None: return search_config.extrapolation_mode return ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED From 3f039d00d94244d1292f5b3a5558a546b8cfdbbc Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Wed, 12 Nov 2025 14:29:50 -0500 Subject: [PATCH 9/9] comments --- src/sentry/search/eap/columns.py | 4 +--- src/sentry/search/eap/types.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index dd26f2ff564991..5dfc743578831b 100644 --- a/src/sentry/search/eap/columns.py +++ b/src/sentry/search/eap/columns.py @@ -201,7 +201,6 @@ class ResolvedAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType - # Whether to enable extrapolation extrapolation_mode: ExtrapolationMode.ValueType is_aggregate: bool = field(default=True, init=False) # Only for aggregates, we only support functions with 1 argument right now @@ -229,7 +228,6 @@ class ResolvedMetricAggregate(ResolvedAggregate): class ResolvedConditionalAggregate(ResolvedFunction): # The internal rpc alias for this column internal_name: Function.ValueType - # Whether to enable extrapolation extrapolation_mode: ExtrapolationMode.ValueType # The condition to filter on filter: TraceItemFilter @@ -279,7 +277,7 @@ 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 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 diff --git a/src/sentry/search/eap/types.py b/src/sentry/search/eap/types.py index 18211c566872b6..b214746e92903d 100644 --- a/src/sentry/search/eap/types.py +++ b/src/sentry/search/eap/types.py @@ -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 - # If set to True, do not extrapolate any values regardless of individual aggregate settings extrapolation_mode: ExtrapolationMode.ValueType | None = None # Whether to set the timestamp granularities to stable buckets stable_timestamp_quantization: bool = True