diff --git a/pyproject.toml b/pyproject.toml index 6062fd6bb27795..d4aebc0d9b54d4 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..383110359b3e79 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,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,6 +552,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 +560,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..4ee0693a741ced 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,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: diff --git a/src/sentry/api/endpoints/organization_events_timeseries.py b/src/sentry/api/endpoints/organization_events_timeseries.py index 5c0c9eca9de9c2..704e2b5f61c182 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 + 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: diff --git a/src/sentry/search/eap/columns.py b/src/sentry/search/eap/columns.py index a3a07fd5bca1ec..5dfc743578831b 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, 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, 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/extrapolation_mode.py b/src/sentry/search/eap/extrapolation_mode.py new file mode 100644 index 00000000000000..9d492f758e8425 --- /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 is not None: + return search_config.extrapolation_mode + + return ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED diff --git a/src/sentry/search/eap/spans/aggregates.py b/src/sentry/search/eap/spans/aggregates.py index 614077aa9d2e31..45fdd738b1dfbd 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_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, diff --git a/src/sentry/search/eap/types.py b/src/sentry/search/eap/types.py index bc97b73be9fa16..b214746e92903d 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,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 diff --git a/tests/snuba/api/endpoints/test_organization_events_span_indexed.py b/tests/snuba/api/endpoints/test_organization_events_span_indexed.py index 4d3ebe4c0fbb8b..ebcaff7940deeb 100644 --- a/tests/snuba/api/endpoints/test_organization_events_span_indexed.py +++ b/tests/snuba/api/endpoints/test_organization_events_span_indexed.py @@ -6,10 +6,12 @@ import pytest import urllib3 from django.utils.timezone import now +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode from sentry.insights.models import InsightsStarredSegment from sentry.testutils.helpers import parse_link_header 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 # Downsampling is deterministic, so unless the algorithm changes we can find a known id that will appear in the @@ -2038,6 +2040,56 @@ def test_extrapolation(self) -> None: assert data[1]["count()"] == 1 assert confidence[1]["count()"] in ("high", "low") + @mock.patch( + "sentry.utils.snuba_rpc._make_rpc_requests", + wraps=_make_rpc_requests, + ) + def test_extrapolation_mode_server_only(self, mock_rpc_request: mock.MagicMock) -> None: + spans = [] + spans.append( + self.create_span( + { + "description": "foo", + "sentry_tags": {"status": "success"}, + "measurements": {"server_sample_rate": {"value": 0.1}}, + }, + start_ts=self.ten_mins_ago, + ) + ) + spans.append( + self.create_span( + { + "description": "bar", + "sentry_tags": {"status": "success"}, + }, + start_ts=self.ten_mins_ago, + ) + ) + self.store_spans(spans, is_eap=True) + response = self.do_request( + { + "field": ["description", "count()"], + "orderby": "-count()", + "query": "", + "project": self.project.id, + "dataset": "spans", + "extrapolationMode": "serverOnly", + } + ) + + assert response.status_code == 200, response.content + + assert ( + mock_rpc_request.call_args.kwargs["table_requests"][0] + .columns[1] + .aggregation.extrapolation_mode + == ExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY + ) + + # TODO: Ensure server only extrapolation actually gets applied + data = response.data["data"] + assert len(data) == 2 + def test_span_duration(self) -> None: spans = [ self.create_span( 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..b9f2feb563461c 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,66 @@ 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()" + + # TODO: Ensure server only extrapolation actually gets applied + 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( diff --git a/uv.lock b/uv.lock index fdda926fc5e520..9e3961ca412f5a 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]]