Skip to content

Commit e1c3676

Browse files
authored
ref(ACI): Pull get_comparison_aggregation_value code into smaller functions (#91909)
Follow up to #91783 to pull code into 2 smaller functions for EAP and non EAP comparison aggregates.
1 parent 1b12414 commit e1c3676

File tree

1 file changed

+112
-64
lines changed

1 file changed

+112
-64
lines changed

src/sentry/incidents/utils/process_update_helpers.py

Lines changed: 112 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import logging
2-
from datetime import timedelta
2+
from datetime import datetime, timedelta
33

44
from snuba_sdk import Column, Condition, Limit, Op
55

@@ -8,6 +8,7 @@
88
from sentry.snuba.dataset import Dataset
99
from sentry.snuba.entity_subscription import (
1010
ENTITY_TIME_COLUMNS,
11+
EntitySubscription,
1112
get_entity_key_from_query_builder,
1213
get_entity_subscription_from_snuba_query,
1314
)
@@ -74,6 +75,96 @@ def get_aggregation_value_helper(subscription_update: QuerySubscriptionUpdate) -
7475
return aggregation_value
7576

7677

78+
def get_eap_aggregation_value(
79+
entity_subscription: EntitySubscription,
80+
subscription_update: QuerySubscriptionUpdate,
81+
snuba_query: SnubaQuery,
82+
project_ids: list[int],
83+
organization_id: int,
84+
start: datetime,
85+
end: datetime,
86+
alert_rule_id: int | None = None,
87+
) -> float | None:
88+
comparison_aggregate: None | float = None
89+
try:
90+
rpc_time_series_request = entity_subscription.build_rpc_request(
91+
query=snuba_query.query,
92+
project_ids=project_ids,
93+
environment=snuba_query.environment,
94+
params={
95+
"organization_id": organization_id,
96+
"project_id": project_ids,
97+
},
98+
referrer="subscription_processor.comparison_query",
99+
)
100+
101+
rpc_time_series_request = add_start_end_conditions(rpc_time_series_request, start, end)
102+
103+
rpc_response = snuba_rpc.timeseries_rpc([rpc_time_series_request])[0]
104+
if len(rpc_response.result_timeseries):
105+
comparison_aggregate = rpc_response.result_timeseries[0].data_points[0].data
106+
107+
except Exception:
108+
logger.exception(
109+
"Failed to run RPC comparison query",
110+
extra={
111+
"alert_rule_id": alert_rule_id,
112+
"subscription_id": subscription_update.get("subscription_id"),
113+
"organization_id": organization_id,
114+
},
115+
)
116+
return None
117+
return comparison_aggregate
118+
119+
120+
def get_aggregation_value(
121+
entity_subscription: EntitySubscription,
122+
subscription_update: QuerySubscriptionUpdate,
123+
snuba_query: SnubaQuery,
124+
project_ids: list[int],
125+
organization_id: int,
126+
start: datetime,
127+
end: datetime,
128+
alert_rule_id: int | None = None,
129+
) -> float | None:
130+
comparison_aggregate: None | float = None
131+
try:
132+
# TODO: determine whether we need to include the subscription query_extra here
133+
query_builder = entity_subscription.build_query_builder(
134+
query=snuba_query.query,
135+
project_ids=project_ids,
136+
environment=snuba_query.environment,
137+
params={
138+
"organization_id": organization_id,
139+
"project_id": project_ids,
140+
"start": start,
141+
"end": end,
142+
},
143+
)
144+
time_col = ENTITY_TIME_COLUMNS[get_entity_key_from_query_builder(query_builder)]
145+
query_builder.add_conditions(
146+
[
147+
Condition(Column(time_col), Op.GTE, start),
148+
Condition(Column(time_col), Op.LT, end),
149+
]
150+
)
151+
query_builder.limit = Limit(1)
152+
results = query_builder.run_query(referrer="subscription_processor.comparison_query")
153+
comparison_aggregate = list(results["data"][0].values())[0]
154+
155+
except Exception:
156+
logger.exception(
157+
"Failed to run comparison query",
158+
extra={
159+
"alert_rule_id": alert_rule_id,
160+
"subscription_id": subscription_update.get("subscription_id"),
161+
"organization_id": organization_id,
162+
},
163+
)
164+
return None
165+
return comparison_aggregate
166+
167+
77168
def get_comparison_aggregation_value(
78169
subscription_update: QuerySubscriptionUpdate,
79170
snuba_query: SnubaQuery,
@@ -97,72 +188,29 @@ def get_comparison_aggregation_value(
97188
dataset = Dataset(snuba_query.dataset)
98189
query_type = SnubaQuery.Type(snuba_query.type)
99190

100-
comparison_aggregate: None | float = None
101191
if query_type == SnubaQuery.Type.PERFORMANCE and dataset == Dataset.EventsAnalyticsPlatform:
102-
try:
103-
rpc_time_series_request = entity_subscription.build_rpc_request(
104-
query=snuba_query.query,
105-
project_ids=project_ids,
106-
environment=snuba_query.environment,
107-
params={
108-
"organization_id": organization_id,
109-
"project_id": project_ids,
110-
},
111-
referrer="subscription_processor.comparison_query",
112-
)
113-
114-
rpc_time_series_request = add_start_end_conditions(rpc_time_series_request, start, end)
115-
116-
rpc_response = snuba_rpc.timeseries_rpc([rpc_time_series_request])[0]
117-
if len(rpc_response.result_timeseries):
118-
comparison_aggregate = rpc_response.result_timeseries[0].data_points[0].data
119-
120-
except Exception:
121-
logger.exception(
122-
"Failed to run RPC comparison query",
123-
extra={
124-
"alert_rule_id": alert_rule_id,
125-
"subscription_id": subscription_update.get("subscription_id"),
126-
"organization_id": organization_id,
127-
},
128-
)
129-
return None
192+
comparison_aggregate = get_eap_aggregation_value(
193+
entity_subscription,
194+
subscription_update,
195+
snuba_query,
196+
project_ids,
197+
organization_id,
198+
start,
199+
end,
200+
alert_rule_id,
201+
)
130202

131203
else:
132-
try:
133-
# TODO: determine whether we need to include the subscription query_extra here
134-
query_builder = entity_subscription.build_query_builder(
135-
query=snuba_query.query,
136-
project_ids=project_ids,
137-
environment=snuba_query.environment,
138-
params={
139-
"organization_id": organization_id,
140-
"project_id": project_ids,
141-
"start": start,
142-
"end": end,
143-
},
144-
)
145-
time_col = ENTITY_TIME_COLUMNS[get_entity_key_from_query_builder(query_builder)]
146-
query_builder.add_conditions(
147-
[
148-
Condition(Column(time_col), Op.GTE, start),
149-
Condition(Column(time_col), Op.LT, end),
150-
]
151-
)
152-
query_builder.limit = Limit(1)
153-
results = query_builder.run_query(referrer="subscription_processor.comparison_query")
154-
comparison_aggregate = list(results["data"][0].values())[0]
155-
156-
except Exception:
157-
logger.exception(
158-
"Failed to run comparison query",
159-
extra={
160-
"alert_rule_id": alert_rule_id,
161-
"subscription_id": subscription_update.get("subscription_id"),
162-
"organization_id": organization_id,
163-
},
164-
)
165-
return None
204+
comparison_aggregate = get_aggregation_value(
205+
entity_subscription,
206+
subscription_update,
207+
snuba_query,
208+
project_ids,
209+
organization_id,
210+
start,
211+
end,
212+
alert_rule_id,
213+
)
166214

167215
if not comparison_aggregate:
168216
metrics.incr("incidents.alert_rules.skipping_update_comparison_value_invalid")

0 commit comments

Comments
 (0)