Skip to main content

scouter_types/custom/
alert.rs

1use crate::error::TypeError;
2use crate::{
3    dispatch::AlertDispatchType, AlertDispatchConfig, AlertThreshold, CommonCrons,
4    DispatchAlertDescription, OpsGenieDispatchConfig, PyHelperFuncs, SlackDispatchConfig,
5    ValidateAlertConfig,
6};
7use crate::{AlertCondition, AlertMap};
8use core::fmt::Debug;
9use pyo3::prelude::*;
10use pyo3::types::PyString;
11
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14
15#[pyclass]
16#[derive(Debug, Serialize, Deserialize, Clone)]
17pub struct CustomMetric {
18    #[pyo3(get, set)]
19    pub name: String,
20
21    #[pyo3(get, set)]
22    pub baseline_value: f64,
23
24    #[pyo3(get, set)]
25    pub alert_condition: AlertCondition,
26}
27
28#[pymethods]
29impl CustomMetric {
30    #[new]
31    #[pyo3(signature = (name, baseline_value, alert_threshold, delta=None))]
32    pub fn new(
33        name: &str,
34        baseline_value: f64,
35        alert_threshold: AlertThreshold,
36        delta: Option<f64>,
37    ) -> Result<Self, TypeError> {
38        let custom_condition = AlertCondition::new(baseline_value, alert_threshold, delta);
39
40        Ok(Self {
41            name: name.to_lowercase(),
42            baseline_value,
43            alert_condition: custom_condition,
44        })
45    }
46
47    pub fn __str__(&self) -> String {
48        // serialize the struct to a string
49        PyHelperFuncs::__str__(self)
50    }
51
52    #[getter]
53    pub fn alert_threshold(&self) -> AlertThreshold {
54        self.alert_condition.alert_threshold.clone()
55    }
56
57    #[getter]
58    pub fn delta(&self) -> Option<f64> {
59        self.alert_condition.delta
60    }
61}
62
63#[pyclass]
64#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
65pub struct CustomMetricAlertConfig {
66    pub dispatch_config: AlertDispatchConfig,
67
68    #[pyo3(get, set)]
69    pub schedule: String,
70
71    #[pyo3(get, set)]
72    pub alert_conditions: Option<HashMap<String, AlertCondition>>,
73}
74
75impl CustomMetricAlertConfig {
76    pub fn set_alert_conditions(&mut self, metrics: &[CustomMetric]) {
77        self.alert_conditions = Some(
78            metrics
79                .iter()
80                .map(|m| (m.name.clone(), m.alert_condition.clone()))
81                .collect(),
82        );
83    }
84}
85
86impl ValidateAlertConfig for CustomMetricAlertConfig {}
87
88#[pymethods]
89impl CustomMetricAlertConfig {
90    #[new]
91    #[pyo3(signature = (schedule=None, dispatch_config=None))]
92    pub fn new(
93        schedule: Option<&Bound<'_, PyAny>>,
94        dispatch_config: Option<&Bound<'_, PyAny>>,
95    ) -> Result<Self, TypeError> {
96        let alert_dispatch_config = match dispatch_config {
97            None => AlertDispatchConfig::default(),
98            Some(config) => {
99                if config.is_instance_of::<SlackDispatchConfig>() {
100                    AlertDispatchConfig::Slack(config.extract::<SlackDispatchConfig>()?)
101                } else if config.is_instance_of::<OpsGenieDispatchConfig>() {
102                    AlertDispatchConfig::OpsGenie(config.extract::<OpsGenieDispatchConfig>()?)
103                } else {
104                    AlertDispatchConfig::default()
105                }
106            }
107        };
108
109        let schedule = match schedule {
110            Some(schedule) => {
111                if schedule.is_instance_of::<PyString>() {
112                    schedule.to_string()
113                } else if schedule.is_instance_of::<CommonCrons>() {
114                    schedule.extract::<CommonCrons>().unwrap().cron()
115                } else {
116                    return Err(TypeError::InvalidScheduleError)?;
117                }
118            }
119            None => CommonCrons::EveryDay.cron(),
120        };
121
122        let schedule = Self::resolve_schedule(&schedule);
123
124        Ok(Self {
125            schedule,
126            dispatch_config: alert_dispatch_config,
127            alert_conditions: None,
128        })
129    }
130
131    #[getter]
132    pub fn dispatch_type(&self) -> AlertDispatchType {
133        self.dispatch_config.dispatch_type()
134    }
135
136    #[getter]
137    pub fn dispatch_config<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
138        self.dispatch_config.config(py)
139    }
140}
141
142impl Default for CustomMetricAlertConfig {
143    fn default() -> CustomMetricAlertConfig {
144        Self {
145            dispatch_config: AlertDispatchConfig::default(),
146            schedule: CommonCrons::EveryDay.cron(),
147            alert_conditions: None,
148        }
149    }
150}
151
152#[derive(Serialize, Deserialize, Debug, Default, Clone)]
153pub struct ComparisonMetricAlert {
154    pub metric_name: String,
155    pub baseline_value: f64,
156    pub observed_value: f64,
157    pub delta: Option<f64>,
158    pub alert_threshold: AlertThreshold,
159}
160
161impl From<ComparisonMetricAlert> for AlertMap {
162    fn from(val: ComparisonMetricAlert) -> Self {
163        AlertMap::Custom(val)
164    }
165}
166
167impl ComparisonMetricAlert {
168    fn alert_description_header(&self) -> String {
169        let below_threshold = |delta: Option<f64>| match delta {
170            Some(b) => format!(
171                "The observed {} metric value has dropped below the threshold (initial value - {})",
172                self.metric_name, b
173            ),
174            None => format!(
175                "The {} metric value has dropped below the initial value",
176                self.metric_name
177            ),
178        };
179
180        let above_threshold = |delta: Option<f64>| match delta {
181            Some(b) => format!(
182                "The {} metric value has increased beyond the threshold (initial value + {})",
183                self.metric_name, b
184            ),
185            None => format!(
186                "The {} metric value has increased beyond the initial value",
187                self.metric_name
188            ),
189        };
190
191        let outside_threshold = |delta: Option<f64>| match delta {
192            Some(b) => format!(
193                "The {} metric value has fallen outside the threshold (initial value ± {})",
194                self.metric_name, b,
195            ),
196            None => format!(
197                "The metric value has fallen outside the initial value for {}",
198                self.metric_name
199            ),
200        };
201
202        match self.alert_threshold {
203            AlertThreshold::Below => below_threshold(self.delta),
204            AlertThreshold::Above => above_threshold(self.delta),
205            AlertThreshold::Outside => outside_threshold(self.delta),
206        }
207    }
208}
209
210impl DispatchAlertDescription for ComparisonMetricAlert {
211    // TODO make pretty per dispatch type
212    fn create_alert_description(&self, _dispatch_type: AlertDispatchType) -> String {
213        let mut alert_description = String::new();
214        let header = format!("{}\n", self.alert_description_header());
215        alert_description.push_str(&header);
216
217        let current_metric = format!("Current Metric Value: {}\n", self.observed_value);
218        let historical_metric = format!("Initial Metric Value: {}\n", self.baseline_value);
219
220        alert_description.push_str(&historical_metric);
221        alert_description.push_str(&current_metric);
222
223        alert_description
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn test_alert_config() {
233        //test console alert config
234        let dispatch_config = AlertDispatchConfig::OpsGenie(OpsGenieDispatchConfig {
235            team: "test-team".to_string(),
236            priority: "P5".to_string(),
237        });
238        let schedule = "0 0 * * * *".to_string();
239        let mut alert_config = CustomMetricAlertConfig {
240            dispatch_config,
241            schedule,
242            ..Default::default()
243        };
244        assert_eq!(alert_config.dispatch_type(), AlertDispatchType::OpsGenie);
245
246        let custom_metrics = vec![
247            CustomMetric::new("mae", 12.4, AlertThreshold::Above, Some(2.3)).unwrap(),
248            CustomMetric::new("accuracy", 0.85, AlertThreshold::Below, None).unwrap(),
249        ];
250
251        alert_config.set_alert_conditions(&custom_metrics);
252
253        if let Some(alert_conditions) = alert_config.alert_conditions.as_ref() {
254            assert_eq!(
255                alert_conditions["mae"].alert_threshold,
256                AlertThreshold::Above
257            );
258            assert_eq!(alert_conditions["mae"].delta, Some(2.3));
259            assert_eq!(
260                alert_conditions["accuracy"].alert_threshold,
261                AlertThreshold::Below
262            );
263            assert_eq!(alert_conditions["accuracy"].delta, None);
264        } else {
265            panic!("alert_conditions should not be None");
266        }
267    }
268
269    #[test]
270    fn test_create_alert_description() {
271        let alert_above_threshold = ComparisonMetricAlert {
272            metric_name: "mse".to_string(),
273            baseline_value: 12.5,
274            observed_value: 14.0,
275            delta: Some(1.0),
276            alert_threshold: AlertThreshold::Above,
277        };
278
279        let description =
280            alert_above_threshold.create_alert_description(AlertDispatchType::Console);
281        assert!(description.contains(
282            "The mse metric value has increased beyond the threshold (initial value + 1)"
283        ));
284        assert!(description.contains("Initial Metric Value: 12.5"));
285        assert!(description.contains("Current Metric Value: 14"));
286
287        let alert_below_threshold = ComparisonMetricAlert {
288            metric_name: "accuracy".to_string(),
289            baseline_value: 0.9,
290            observed_value: 0.7,
291            delta: None,
292            alert_threshold: AlertThreshold::Below,
293        };
294
295        let description =
296            alert_below_threshold.create_alert_description(AlertDispatchType::Console);
297        assert!(
298            description.contains("The accuracy metric value has dropped below the initial value")
299        );
300        assert!(description.contains("Initial Metric Value: 0.9"));
301        assert!(description.contains("Current Metric Value: 0.7"));
302
303        let alert_outside_threshold = ComparisonMetricAlert {
304            metric_name: "mae".to_string(),
305            baseline_value: 12.5,
306            observed_value: 22.0,
307            delta: Some(2.0),
308            alert_threshold: AlertThreshold::Outside,
309        };
310
311        let description =
312            alert_outside_threshold.create_alert_description(AlertDispatchType::Console);
313        assert!(description
314            .contains("The mae metric value has fallen outside the threshold (initial value ± 2)"));
315        assert!(description.contains("Initial Metric Value: 12.5"));
316        assert!(description.contains("Current Metric Value: 22"));
317    }
318}