From 9c68dc498304fa0552ae18b5f30e2980ef651f1e Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Fri, 21 Nov 2025 12:59:42 -0500 Subject: [PATCH 1/2] feat: Add support for extrapolation modes in entity subscription --- src/sentry/snuba/entity_subscription.py | 30 ++++++++++++++- src/sentry/snuba/snuba_query_validator.py | 1 + src/sentry/snuba/tasks.py | 5 ++- .../sentry/snuba/test_entity_subscriptions.py | 37 ++++++++++++++++++- 4 files changed, 69 insertions(+), 4 deletions(-) diff --git a/src/sentry/snuba/entity_subscription.py b/src/sentry/snuba/entity_subscription.py index 76ca1bf41e39b9..990b838b449005 100644 --- a/src/sentry/snuba/entity_subscription.py +++ b/src/sentry/snuba/entity_subscription.py @@ -8,6 +8,9 @@ from typing import Any, TypedDict, Union from sentry_protos.snuba.v1.endpoint_time_series_pb2 import TimeSeriesRequest +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( + ExtrapolationMode as ProtoExtrapolationMode, +) from snuba_sdk import Column, Condition, Entity, Join, Op, Request from sentry import features @@ -27,7 +30,7 @@ from sentry.snuba.dataset import Dataset, EntityKey from sentry.snuba.metrics.extraction import MetricSpecType from sentry.snuba.metrics.naming_layer.mri import SessionMRI -from sentry.snuba.models import SnubaQuery, SnubaQueryEventType +from sentry.snuba.models import ExtrapolationMode, SnubaQuery, SnubaQueryEventType from sentry.snuba.ourlogs import OurLogs from sentry.snuba.referrer import Referrer from sentry.snuba.rpc_dataset_common import RPCBase @@ -67,6 +70,14 @@ "timestamp.to_day", } +# Mapping from model ExtrapolationMode to proto ExtrapolationMode +MODEL_TO_PROTO_EXTRAPOLATION_MODE = { + ExtrapolationMode.UNKNOWN: ProtoExtrapolationMode.EXTRAPOLATION_MODE_UNSPECIFIED, + ExtrapolationMode.NONE: ProtoExtrapolationMode.EXTRAPOLATION_MODE_NONE, + ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, + ExtrapolationMode.SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY, +} + def apply_dataset_query_conditions( query_type: SnubaQuery.Type, @@ -106,6 +117,7 @@ def apply_dataset_query_conditions( class _EntitySpecificParams(TypedDict, total=False): org_id: int + extrapolation_mode: ExtrapolationMode | None event_types: list[SnubaQueryEventType.EventType] | None @@ -265,9 +277,11 @@ def __init__( self.aggregate = aggregate self.event_types = None self.time_window = time_window + self.extrapolation_mode = None if extra_fields: self.org_id = extra_fields.get("org_id") self.event_types = extra_fields.get("event_types") + self.extrapolation_mode = extra_fields.get("extrapolation_mode") def build_rpc_request( self, @@ -301,8 +315,19 @@ def build_rpc_request( end=now, granularity_secs=self.time_window, ) + + # Convert model ExtrapolationMode to proto ExtrapolationMode + proto_extrapolation_mode = None + if self.extrapolation_mode is not None: + model_mode = ExtrapolationMode(self.extrapolation_mode) + proto_extrapolation_mode = MODEL_TO_PROTO_EXTRAPOLATION_MODE.get(model_mode) + search_resolver = dataset_module.get_resolver( - snuba_params, SearchResolverConfig(stable_timestamp_quantization=False) + snuba_params, + SearchResolverConfig( + stable_timestamp_quantization=False, + extrapolation_mode=proto_extrapolation_mode, + ), ) rpc_request, _, _ = dataset_module.get_timeseries_query( @@ -641,6 +666,7 @@ def get_entity_subscription_from_snuba_query( extra_fields={ "org_id": organization_id, "event_types": snuba_query.event_types, + "extrapolation_mode": ExtrapolationMode(snuba_query.extrapolation_mode), }, ) diff --git a/src/sentry/snuba/snuba_query_validator.py b/src/sentry/snuba/snuba_query_validator.py index c5bc55d8d036e9..a78295d4d2610b 100644 --- a/src/sentry/snuba/snuba_query_validator.py +++ b/src/sentry/snuba/snuba_query_validator.py @@ -284,6 +284,7 @@ def _validate_query(self, data): extra_fields={ "org_id": projects[0].organization_id, "event_types": data.get("event_types"), + "extrapolation_mode": data.get("extrapolation_mode"), }, ) except UnsupportedQuerySubscription as e: diff --git a/src/sentry/snuba/tasks.py b/src/sentry/snuba/tasks.py index 7fa6420fec5bfb..0957a0e990f277 100644 --- a/src/sentry/snuba/tasks.py +++ b/src/sentry/snuba/tasks.py @@ -19,7 +19,7 @@ get_entity_subscription, get_entity_subscription_from_snuba_query, ) -from sentry.snuba.models import QuerySubscription, SnubaQuery +from sentry.snuba.models import ExtrapolationMode, QuerySubscription, SnubaQuery from sentry.snuba.utils import build_query_strings from sentry.tasks.base import instrumented_task from sentry.taskworker.namespaces import alerts_tasks @@ -133,6 +133,9 @@ def update_subscription_in_snuba( extra_fields={ "org_id": subscription.project.organization_id, "event_types": subscription.snuba_query.event_types, + "extrapolation_mode": ExtrapolationMode( + subscription.snuba_query.extrapolation_mode + ), }, ) old_entity_key = ( diff --git a/tests/sentry/snuba/test_entity_subscriptions.py b/tests/sentry/snuba/test_entity_subscriptions.py index 6729173ecf07eb..3913d821993ed4 100644 --- a/tests/sentry/snuba/test_entity_subscriptions.py +++ b/tests/sentry/snuba/test_entity_subscriptions.py @@ -1,4 +1,7 @@ import pytest +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( + ExtrapolationMode as ProtoExtrapolationMode, +) from snuba_sdk import And, Column, Condition, Entity, Function, Join, Op, Relationship from sentry.exceptions import InvalidQuerySubscription, UnsupportedQuerySubscription @@ -20,7 +23,7 @@ get_entity_subscription_from_snuba_query, ) from sentry.snuba.metrics.naming_layer.mri import SessionMRI -from sentry.snuba.models import SnubaQuery +from sentry.snuba.models import ExtrapolationMode, SnubaQuery from sentry.testutils.cases import TestCase from sentry.testutils.skips import requires_snuba @@ -456,6 +459,38 @@ def test_get_entity_subscription_for_eap_rpc_query(self) -> None: assert rpc_timeseries_request.granularity_secs == 3600 assert rpc_timeseries_request.filter.comparison_filter.value.val_str == "http.client" assert rpc_timeseries_request.expressions[0].aggregation.label == "count(span.duration)" + assert ( + rpc_timeseries_request.expressions[0].aggregation.extrapolation_mode + == ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ) + + def test_get_entity_subscription_for_eap_with_extrapolation_mode(self) -> None: + aggregate = "count(span.duration)" + query = "span.op:http.client" + + # Test with SERVER_WEIGHTED extrapolation mode + entity_subscription = get_entity_subscription( + query_type=SnubaQuery.Type.PERFORMANCE, + dataset=Dataset.EventsAnalyticsPlatform, + aggregate=aggregate, + time_window=3600, + extra_fields={ + "org_id": self.organization.id, + "extrapolation_mode": ExtrapolationMode.SERVER_WEIGHTED, + }, + ) + assert isinstance(entity_subscription, PerformanceSpansEAPRpcEntitySubscription) + assert entity_subscription.extrapolation_mode == ExtrapolationMode.SERVER_WEIGHTED + + rpc_timeseries_request = entity_subscription.build_rpc_request( + query, [self.project.id], None + ) + + # Verify the extrapolation mode is passed to the RPC request + assert ( + rpc_timeseries_request.expressions[0].aggregation.extrapolation_mode + == ProtoExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY + ) class GetEntitySubscriptionFromSnubaQueryTest(TestCase): From fdb348742297efaf1ad484b57344140eaf02e919 Mon Sep 17 00:00:00 2001 From: Shruthilaya Jaganathan Date: Fri, 21 Nov 2025 15:12:21 -0500 Subject: [PATCH 2/2] unknown extrapolation mode is sample weighted --- src/sentry/snuba/entity_subscription.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sentry/snuba/entity_subscription.py b/src/sentry/snuba/entity_subscription.py index 990b838b449005..8cbd27c0164c7f 100644 --- a/src/sentry/snuba/entity_subscription.py +++ b/src/sentry/snuba/entity_subscription.py @@ -72,7 +72,7 @@ # Mapping from model ExtrapolationMode to proto ExtrapolationMode MODEL_TO_PROTO_EXTRAPOLATION_MODE = { - ExtrapolationMode.UNKNOWN: ProtoExtrapolationMode.EXTRAPOLATION_MODE_UNSPECIFIED, + ExtrapolationMode.UNKNOWN: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, ExtrapolationMode.NONE: ProtoExtrapolationMode.EXTRAPOLATION_MODE_NONE, ExtrapolationMode.CLIENT_AND_SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, ExtrapolationMode.SERVER_WEIGHTED: ProtoExtrapolationMode.EXTRAPOLATION_MODE_SERVER_ONLY,