Skip to content

Commit bc5ab71

Browse files
milindsrivastava1997milindsrivastava1997
authored
refactor: planner yaml logic (#283)
1 parent 4a0cba8 commit bc5ab71

7 files changed

Lines changed: 189 additions & 180 deletions

File tree

‎asap-planner-rs/src/config/input.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use asap_types::enums::CleanupPolicy;
12
use asap_types::PromQLSchema;
23
use promql_utilities::data_model::KeyByLabelNames;
34
use serde::Deserialize;
@@ -58,7 +59,7 @@ pub struct MetricDefinition {
5859

5960
#[derive(Debug, Clone, Deserialize)]
6061
pub struct AggregateCleanupConfig {
61-
pub policy: Option<String>,
62+
pub policy: Option<CleanupPolicy>,
6263
}
6364

6465
#[derive(Debug, Clone, Deserialize, Default)]

‎asap-planner-rs/src/lib.rs‎

Lines changed: 91 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ pub use config::input::ControllerConfig;
1818
pub use config::input::SQLControllerConfig;
1919
pub use error::ControllerError;
2020
pub use output::generator::{GeneratorOutput, PuntedQuery};
21+
use output::generator::{
22+
KEY_AGGREGATIONS, KEY_AGG_SUB_TYPE, KEY_AGG_TYPE, KEY_LABELS, KEY_NUM_AGG_TO_RETAIN,
23+
KEY_QUERIES, KEY_QUERY, KEY_READ_COUNT_THRESHOLD, KEY_TABLE_NAME, KEY_VALUE_COLUMN,
24+
KEY_WINDOW_SIZE,
25+
};
2126
pub use output::sql_generator::SQLRuntimeOptions;
2227
pub use prometheus_client::build_schema_from_prometheus;
2328

@@ -61,20 +66,28 @@ impl PlannerOutput {
6166
self.query_count
6267
}
6368

64-
pub fn has_aggregation_type(&self, t: &str) -> bool {
69+
fn streaming_aggs_slice(&self) -> Option<&[YamlValue]> {
6570
if let YamlValue::Mapping(root) = &self.streaming_yaml {
66-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
67-
return aggs.iter().any(|agg| {
68-
if let YamlValue::Mapping(m) = agg {
69-
if let Some(YamlValue::String(agg_type)) = m.get("aggregationType") {
70-
return agg_type == t;
71-
}
72-
}
73-
false
74-
});
71+
if let Some(YamlValue::Sequence(aggs)) = root.get(KEY_AGGREGATIONS) {
72+
return Some(aggs.as_slice());
7573
}
7674
}
77-
false
75+
None
76+
}
77+
78+
fn find_aggregation_by_type(&self, agg_type: &str) -> Option<&serde_yaml::Mapping> {
79+
self.streaming_aggs_slice()?.iter().find_map(|agg| {
80+
if let YamlValue::Mapping(m) = agg {
81+
if matches!(m.get(KEY_AGG_TYPE), Some(YamlValue::String(s)) if s == agg_type) {
82+
return Some(m);
83+
}
84+
}
85+
None
86+
})
87+
}
88+
89+
pub fn has_aggregation_type(&self, t: &str) -> bool {
90+
self.find_aggregation_by_type(t).is_some()
7891
}
7992

8093
pub fn all_tumbling_window_sizes_eq(&self, s: u64) -> bool {
@@ -86,11 +99,11 @@ impl PlannerOutput {
8699
}
87100

88101
fn check_tumbling_window_sizes(&self, predicate: impl Fn(u64) -> bool) -> bool {
89-
if let YamlValue::Mapping(root) = &self.streaming_yaml {
90-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
91-
return aggs.iter().all(|agg| {
102+
self.streaming_aggs_slice()
103+
.map(|aggs| {
104+
aggs.iter().all(|agg| {
92105
if let YamlValue::Mapping(m) = agg {
93-
if let Some(val) = m.get("windowSize") {
106+
if let Some(val) = m.get(KEY_WINDOW_SIZE) {
94107
let size = match val {
95108
YamlValue::Number(n) => n.as_u64().unwrap_or(0),
96109
_ => 0,
@@ -99,59 +112,61 @@ impl PlannerOutput {
99112
}
100113
}
101114
false
102-
});
103-
}
104-
}
105-
false
115+
})
116+
})
117+
.unwrap_or(false)
106118
}
107119

108120
/// Returns the sorted labels for the first aggregation matching `agg_type`,
109121
/// for the given `label_kind` ("rollup", "grouping", or "aggregated").
110122
pub fn aggregation_labels(&self, agg_type: &str, label_kind: &str) -> Vec<String> {
111-
if let YamlValue::Mapping(root) = &self.streaming_yaml {
112-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
113-
for agg in aggs {
114-
if let YamlValue::Mapping(m) = agg {
115-
if let Some(YamlValue::String(t)) = m.get("aggregationType") {
116-
if t == agg_type {
117-
if let Some(YamlValue::Mapping(labels)) = m.get("labels") {
118-
if let Some(YamlValue::Sequence(seq)) = labels.get(label_kind) {
119-
let mut result: Vec<String> = seq
120-
.iter()
121-
.filter_map(|v| {
122-
if let YamlValue::String(s) = v {
123-
Some(s.clone())
124-
} else {
125-
None
126-
}
127-
})
128-
.collect();
129-
result.sort();
130-
return result;
131-
}
132-
}
133-
}
134-
}
135-
}
123+
let Some(seq) = self
124+
.find_aggregation_by_type(agg_type)
125+
.and_then(|m| m.get(KEY_LABELS))
126+
.and_then(|v| {
127+
if let YamlValue::Mapping(lm) = v {
128+
Some(lm)
129+
} else {
130+
None
136131
}
137-
}
138-
}
139-
vec![]
132+
})
133+
.and_then(|lm| lm.get(label_kind))
134+
.and_then(|v| {
135+
if let YamlValue::Sequence(seq) = v {
136+
Some(seq)
137+
} else {
138+
None
139+
}
140+
})
141+
else {
142+
return vec![];
143+
};
144+
let mut result: Vec<String> = seq
145+
.iter()
146+
.filter_map(|v| {
147+
if let YamlValue::String(s) = v {
148+
Some(s.clone())
149+
} else {
150+
None
151+
}
152+
})
153+
.collect();
154+
result.sort();
155+
result
140156
}
141157

142158
/// Returns the cleanup param (read_count_threshold or num_aggregates_to_retain)
143159
/// for the first aggregation entry of the given query string.
144160
pub fn inference_cleanup_param(&self, query: &str) -> Option<u64> {
145161
if let YamlValue::Mapping(root) = &self.inference_yaml {
146-
if let Some(YamlValue::Sequence(queries)) = root.get("queries") {
162+
if let Some(YamlValue::Sequence(queries)) = root.get(KEY_QUERIES) {
147163
for q in queries {
148164
if let YamlValue::Mapping(qm) = q {
149-
if let Some(YamlValue::String(qs)) = qm.get("query") {
165+
if let Some(YamlValue::String(qs)) = qm.get(KEY_QUERY) {
150166
if qs == query {
151-
if let Some(YamlValue::Sequence(aggs)) = qm.get("aggregations") {
167+
if let Some(YamlValue::Sequence(aggs)) = qm.get(KEY_AGGREGATIONS) {
152168
if let Some(YamlValue::Mapping(agg)) = aggs.first() {
153-
for key in
154-
["read_count_threshold", "num_aggregates_to_retain"]
169+
for key in [KEY_READ_COUNT_THRESHOLD, KEY_NUM_AGG_TO_RETAIN]
155170
{
156171
if let Some(YamlValue::Number(n)) = agg.get(key) {
157172
return n.as_u64();
@@ -193,72 +208,44 @@ impl PlannerOutput {
193208

194209
/// Returns the table_name field of the first aggregation matching agg_type.
195210
pub fn aggregation_table_name(&self, agg_type: &str) -> Option<String> {
196-
if let YamlValue::Mapping(root) = &self.streaming_yaml {
197-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
198-
for agg in aggs {
199-
if let YamlValue::Mapping(m) = agg {
200-
if let Some(YamlValue::String(t)) = m.get("aggregationType") {
201-
if t == agg_type {
202-
if let Some(YamlValue::String(name)) = m.get("table_name") {
203-
return Some(name.clone());
204-
}
205-
}
206-
}
207-
}
211+
self.find_aggregation_by_type(agg_type)
212+
.and_then(|m| m.get(KEY_TABLE_NAME))
213+
.and_then(|v| {
214+
if let YamlValue::String(s) = v {
215+
Some(s.clone())
216+
} else {
217+
None
208218
}
209-
}
210-
}
211-
None
219+
})
212220
}
213221

214222
/// Returns the value_column field of the first aggregation matching agg_type.
215223
pub fn aggregation_value_column(&self, agg_type: &str) -> Option<String> {
216-
if let YamlValue::Mapping(root) = &self.streaming_yaml {
217-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
218-
for agg in aggs {
219-
if let YamlValue::Mapping(m) = agg {
220-
if let Some(YamlValue::String(t)) = m.get("aggregationType") {
221-
if t == agg_type {
222-
if let Some(YamlValue::String(col)) = m.get("value_column") {
223-
return Some(col.clone());
224-
}
225-
}
226-
}
227-
}
224+
self.find_aggregation_by_type(agg_type)
225+
.and_then(|m| m.get(KEY_VALUE_COLUMN))
226+
.and_then(|v| {
227+
if let YamlValue::String(s) = v {
228+
Some(s.clone())
229+
} else {
230+
None
228231
}
229-
}
230-
}
231-
None
232+
})
232233
}
233234

234235
/// Returns true if any aggregation has the matching type AND sub_type.
235236
pub fn has_aggregation_type_and_sub_type(&self, agg_type: &str, sub_type: &str) -> bool {
236-
if let YamlValue::Mapping(root) = &self.streaming_yaml {
237-
if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
238-
return aggs.iter().any(|agg| {
237+
self.streaming_aggs_slice()
238+
.map(|aggs| {
239+
aggs.iter().any(|agg| {
239240
if let YamlValue::Mapping(m) = agg {
240-
let type_matches = m.get("aggregationType").and_then(|v| {
241-
if let YamlValue::String(s) = v {
242-
Some(s.as_str())
243-
} else {
244-
None
245-
}
246-
}) == Some(agg_type);
247-
let sub_matches = m.get("aggregationSubType").and_then(|v| {
248-
if let YamlValue::String(s) = v {
249-
Some(s.as_str())
250-
} else {
251-
None
252-
}
253-
}) == Some(sub_type);
254-
type_matches && sub_matches
241+
matches!(m.get(KEY_AGG_TYPE), Some(YamlValue::String(s)) if s == agg_type)
242+
&& matches!(m.get(KEY_AGG_SUB_TYPE), Some(YamlValue::String(s)) if s == sub_type)
255243
} else {
256244
false
257245
}
258-
});
259-
}
260-
}
261-
false
246+
})
247+
})
248+
.unwrap_or(false)
262249
}
263250
}
264251

0 commit comments

Comments
 (0)