diff --git a/src/sentry/incidents/utils/process_update_helpers.py b/src/sentry/incidents/utils/process_update_helpers.py index 1e2bfe3e24f009..8bd697de046971 100644 --- a/src/sentry/incidents/utils/process_update_helpers.py +++ b/src/sentry/incidents/utils/process_update_helpers.py @@ -1,5 +1,5 @@ import logging -from datetime import timedelta +from datetime import datetime, timedelta from snuba_sdk import Column, Condition, Limit, Op @@ -8,6 +8,7 @@ from sentry.snuba.dataset import Dataset from sentry.snuba.entity_subscription import ( ENTITY_TIME_COLUMNS, + EntitySubscription, get_entity_key_from_query_builder, get_entity_subscription_from_snuba_query, ) @@ -74,6 +75,96 @@ def get_aggregation_value_helper(subscription_update: QuerySubscriptionUpdate) - return aggregation_value +def get_eap_aggregation_value( + entity_subscription: EntitySubscription, + subscription_update: QuerySubscriptionUpdate, + snuba_query: SnubaQuery, + project_ids: list[int], + organization_id: int, + start: datetime, + end: datetime, + alert_rule_id: int | None = None, +) -> float | None: + comparison_aggregate: None | float = None + try: + rpc_time_series_request = entity_subscription.build_rpc_request( + query=snuba_query.query, + project_ids=project_ids, + environment=snuba_query.environment, + params={ + "organization_id": organization_id, + "project_id": project_ids, + }, + referrer="subscription_processor.comparison_query", + ) + + rpc_time_series_request = add_start_end_conditions(rpc_time_series_request, start, end) + + rpc_response = snuba_rpc.timeseries_rpc([rpc_time_series_request])[0] + if len(rpc_response.result_timeseries): + comparison_aggregate = rpc_response.result_timeseries[0].data_points[0].data + + except Exception: + logger.exception( + "Failed to run RPC comparison query", + extra={ + "alert_rule_id": alert_rule_id, + "subscription_id": subscription_update.get("subscription_id"), + "organization_id": organization_id, + }, + ) + return None + return comparison_aggregate + + +def get_aggregation_value( + entity_subscription: EntitySubscription, + subscription_update: QuerySubscriptionUpdate, + snuba_query: SnubaQuery, + project_ids: list[int], + organization_id: int, + start: datetime, + end: datetime, + alert_rule_id: int | None = None, +) -> float | None: + comparison_aggregate: None | float = None + try: + # TODO: determine whether we need to include the subscription query_extra here + query_builder = entity_subscription.build_query_builder( + query=snuba_query.query, + project_ids=project_ids, + environment=snuba_query.environment, + params={ + "organization_id": organization_id, + "project_id": project_ids, + "start": start, + "end": end, + }, + ) + time_col = ENTITY_TIME_COLUMNS[get_entity_key_from_query_builder(query_builder)] + query_builder.add_conditions( + [ + Condition(Column(time_col), Op.GTE, start), + Condition(Column(time_col), Op.LT, end), + ] + ) + query_builder.limit = Limit(1) + results = query_builder.run_query(referrer="subscription_processor.comparison_query") + comparison_aggregate = list(results["data"][0].values())[0] + + except Exception: + logger.exception( + "Failed to run comparison query", + extra={ + "alert_rule_id": alert_rule_id, + "subscription_id": subscription_update.get("subscription_id"), + "organization_id": organization_id, + }, + ) + return None + return comparison_aggregate + + def get_comparison_aggregation_value( subscription_update: QuerySubscriptionUpdate, snuba_query: SnubaQuery, @@ -97,72 +188,29 @@ def get_comparison_aggregation_value( dataset = Dataset(snuba_query.dataset) query_type = SnubaQuery.Type(snuba_query.type) - comparison_aggregate: None | float = None if query_type == SnubaQuery.Type.PERFORMANCE and dataset == Dataset.EventsAnalyticsPlatform: - try: - rpc_time_series_request = entity_subscription.build_rpc_request( - query=snuba_query.query, - project_ids=project_ids, - environment=snuba_query.environment, - params={ - "organization_id": organization_id, - "project_id": project_ids, - }, - referrer="subscription_processor.comparison_query", - ) - - rpc_time_series_request = add_start_end_conditions(rpc_time_series_request, start, end) - - rpc_response = snuba_rpc.timeseries_rpc([rpc_time_series_request])[0] - if len(rpc_response.result_timeseries): - comparison_aggregate = rpc_response.result_timeseries[0].data_points[0].data - - except Exception: - logger.exception( - "Failed to run RPC comparison query", - extra={ - "alert_rule_id": alert_rule_id, - "subscription_id": subscription_update.get("subscription_id"), - "organization_id": organization_id, - }, - ) - return None + comparison_aggregate = get_eap_aggregation_value( + entity_subscription, + subscription_update, + snuba_query, + project_ids, + organization_id, + start, + end, + alert_rule_id, + ) else: - try: - # TODO: determine whether we need to include the subscription query_extra here - query_builder = entity_subscription.build_query_builder( - query=snuba_query.query, - project_ids=project_ids, - environment=snuba_query.environment, - params={ - "organization_id": organization_id, - "project_id": project_ids, - "start": start, - "end": end, - }, - ) - time_col = ENTITY_TIME_COLUMNS[get_entity_key_from_query_builder(query_builder)] - query_builder.add_conditions( - [ - Condition(Column(time_col), Op.GTE, start), - Condition(Column(time_col), Op.LT, end), - ] - ) - query_builder.limit = Limit(1) - results = query_builder.run_query(referrer="subscription_processor.comparison_query") - comparison_aggregate = list(results["data"][0].values())[0] - - except Exception: - logger.exception( - "Failed to run comparison query", - extra={ - "alert_rule_id": alert_rule_id, - "subscription_id": subscription_update.get("subscription_id"), - "organization_id": organization_id, - }, - ) - return None + comparison_aggregate = get_aggregation_value( + entity_subscription, + subscription_update, + snuba_query, + project_ids, + organization_id, + start, + end, + alert_rule_id, + ) if not comparison_aggregate: metrics.incr("incidents.alert_rules.skipping_update_comparison_value_invalid")