Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@ pub fn walk_ast_and_extract_info(ast: &dsl::Search) -> Option<ElasticDSLQueryInf
// Extract information from the bool query
walk_bool_query_and_extract_info(&bool_query)
}
Some(dsl::Query::Json(json_query)) => {
// Serde's untagged enum deserialization falls through to Json when typed
// deserialization fails. Extract predicates from the raw JSON structure.
walk_json_query_and_extract_info(&json_query.0)
}
Some(other) => {
// Predicates may just be specified directly without enclosing bool context.
if let Some(predicate) = extract_predicates_from_query(&other) {
Expand All @@ -44,6 +49,57 @@ pub fn walk_ast_and_extract_info(ast: &dsl::Search) -> Option<ElasticDSLQueryInf
))
}

fn walk_json_query_and_extract_info(json: &serde_json::Value) -> Vec<Predicate> {
let Some(filter) = json
.get("bool")
.and_then(|b| b.get("filter"))
.and_then(|f| f.as_array())
else {
return Vec::new();
};
filter
.iter()
.filter_map(extract_predicate_from_json)
.collect()
}

fn extract_predicate_from_json(json: &serde_json::Value) -> Option<Predicate> {
if let Some(range_obj) = json.get("range").and_then(|v| v.as_object()) {
let (field, bounds) = range_obj.iter().next()?;
let field = strip_keyword_suffix(field).to_owned();
let gte = bounds
.get("gte")
.and_then(|v| v.as_str())
.map(|s| TermValue::String(s.to_owned()));
let lte = bounds
.get("lte")
.and_then(|v| v.as_str())
.map(|s| TermValue::String(s.to_owned()));
return Some(Predicate::Range { field, gte, lte });
}
if let Some(term_obj) = json.get("term").and_then(|v| v.as_object()) {
let (field, value_obj) = term_obj.iter().next()?;
let field = strip_keyword_suffix(field).to_owned();
let raw = value_obj.get("value").unwrap_or(value_obj);
let value = match raw {
serde_json::Value::String(s) => TermValue::String(s.clone()),
serde_json::Value::Bool(b) => TermValue::Boolean(*b),
serde_json::Value::Number(n) => {
if let Some(u) = n.as_u64() {
TermValue::UnsignedInt(u)
} else if let Some(i) = n.as_i64() {
TermValue::Int(i)
} else {
TermValue::Float(n.as_f64()?)
}
}
_ => return None,
};
return Some(Predicate::Term { field, value });
}
None
}

fn walk_bool_query_and_extract_info(bool_query: &dsl::BoolQuery) -> Vec<Predicate> {
// Placeholder for walking the filter context of the AST and extracting relevant information
// This would involve traversing the filter nodes and applying logic to determine label filters, time ranges, etc.
Expand Down
53 changes: 36 additions & 17 deletions asap-planner-rs/src/planner/elastic_dsl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ use indexmap::IndexSet;
pub struct ElasticSingleQueryProcessor {
query_string: String,
t_repeat_ms: u64,
#[allow(dead_code)]
data_ingestion_interval_ms: u64,
index_schema: ElasticIndexSchemaBuilder,
#[allow(dead_code)]
Expand Down Expand Up @@ -66,9 +65,43 @@ impl ElasticSingleQueryProcessor {
let (treatment_type, statistics) = get_elastic_statistics(&query_info.aggregation)?;

let t_repeat_ms = self.t_repeat_ms;

// Validate and resolve the time-range predicate before building configs.
let time_field = self.index_schema.time_field.clone();
let time_range = query_info
.predicates
.iter()
.find(|p| match p {
Predicate::Range { field, .. } => field == &time_field,
_ => false,
})
.and_then(|p| range_query_to_time_range(p, 0));
let (t_lookback_ms, window_size_ms) = match time_range {
None => {
return Err(ControllerError::UnsupportedElasticDSLQuery(
"query must include a time-range predicate on the time field".to_string(),
))
}
Some(tr) => {
let duration = tr.duration_ms().unwrap_or(t_repeat_ms);
if duration < self.data_ingestion_interval_ms {
return Err(ControllerError::UnsupportedElasticDSLQuery(format!(
"time-range duration {}ms is shorter than the data ingestion interval {}ms",
duration, self.data_ingestion_interval_ms
)));
}
// Spatial query: duration spans exactly one ingestion interval → use interval as window
let window = if duration == self.data_ingestion_interval_ms {
self.data_ingestion_interval_ms
} else {
t_repeat_ms
};
(duration, window)
}
};
let window_cfg = IntermediateWindowConfig {
window_size_ms: t_repeat_ms,
slide_interval_ms: t_repeat_ms,
window_size_ms,
slide_interval_ms: window_size_ms,
window_type: WindowType::Tumbling,
};

Expand Down Expand Up @@ -121,20 +154,6 @@ impl ElasticSingleQueryProcessor {
)
.map_err(ControllerError::ElasticDSLParse)?;

let time_field = self.index_schema.time_field.clone();
let time_range = query_info
.predicates
.iter()
.find(|p| match p {
Predicate::Range { field, .. } => field == &time_field,
_ => false,
})
.and_then(|p| range_query_to_time_range(p, 0));
let t_lookback_ms = match time_range {
Some(tr) => tr.duration_ms().unwrap_or(t_repeat_ms),
None => t_repeat_ms,
};

// Calculate cleanup param based on query's time window
let cleanup_param = if self.cleanup_policy == CleanupPolicy::NoCleanup {
None
Expand Down
96 changes: 95 additions & 1 deletion asap-planner-rs/tests/elastic_dsl_integration.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use asap_planner::{ElasticController, ElasticRuntimeOptions, PlannerOutput, StreamingEngine};
use asap_planner::{
ControllerError, ElasticController, ElasticRuntimeOptions, PlannerOutput, StreamingEngine,
};
use asap_types::{QueryLanguage, SchemaConfig};
use std::io::Write;
use std::path::Path;
Expand Down Expand Up @@ -41,6 +43,23 @@ fn elastic_output(index: &str, time_field: &str, query: &str, t_repeat_ms: u64)
elastic_output_with_interval(index, time_field, query, t_repeat_ms, 15_000)
}

fn try_elastic_with_interval(
index: &str,
time_field: &str,
query: &str,
t_repeat_ms: u64,
data_ingestion_interval_ms: u64,
) -> Result<PlannerOutput, ControllerError> {
let yaml = elastic_yaml(index, time_field, query, t_repeat_ms);
let mut file = NamedTempFile::new().unwrap();
file.write_all(yaml.as_bytes()).unwrap();
let opts = ElasticRuntimeOptions {
streaming_engine: StreamingEngine::Arroyo,
data_ingestion_interval_ms,
};
ElasticController::from_file(Path::new(file.path()), opts)?.generate()
}

fn elastic_output_with_interval(
index: &str,
time_field: &str,
Expand Down Expand Up @@ -493,3 +512,78 @@ fn sub_second_repetition_delay_ms() {
assert_eq!(out.streaming_aggregation_count(), 2);
assert!(out.all_tumbling_window_sizes_eq(500));
}

// ── time-range validation ─────────────────────────────────────────────────────

fn time_range_query(duration: &str) -> String {
format!(
r#"{{
"aggs": {{ "sum_cpu": {{ "sum": {{ "field": "cpu_usage" }} }} }},
"query": {{ "bool": {{ "filter": [{{ "range": {{ "@timestamp": {{ "gte": "now-{duration}", "lte": "now" }} }} }}] }} }}
}}"#
)
}

const QUERY_NO_TIME_RANGE: &str = r#"{
"aggs": { "sum_cpu": { "sum": { "field": "cpu_usage" } } }
}"#;

#[test]
fn no_time_range_predicate_is_rejected() {
let result = try_elastic_with_interval(
"metrics",
"\"@timestamp\"",
QUERY_NO_TIME_RANGE,
300_000,
15_000,
);
assert!(matches!(
result,
Err(ControllerError::UnsupportedElasticDSLQuery(_))
));
}

#[test]
fn time_range_shorter_than_ingestion_interval_is_rejected() {
// data_ingestion_interval_ms = 15_000ms (15s), query range = 5s → duration < interval
let result = try_elastic_with_interval(
"metrics",
"\"@timestamp\"",
&time_range_query("5s"),
300_000,
15_000,
);
assert!(matches!(
result,
Err(ControllerError::UnsupportedElasticDSLQuery(_))
));
}

#[test]
fn time_range_equal_to_ingestion_interval_uses_interval_as_window_size() {
// data_ingestion_interval_ms = 15_000ms (15s), query range = 15s → Spatial: window = interval
let out = try_elastic_with_interval(
"metrics",
"\"@timestamp\"",
&time_range_query("15s"),
300_000,
15_000,
)
.unwrap();
assert!(out.all_tumbling_window_sizes_eq(15_000));
}

#[test]
fn time_range_longer_than_ingestion_interval_uses_t_repeat_as_window_size() {
// data_ingestion_interval_ms = 15_000ms (15s), query range = 5m > 15s → Temporal: window = t_repeat_ms
let t_repeat_ms = 300_000;
let out = try_elastic_with_interval(
"metrics",
"\"@timestamp\"",
&time_range_query("5m"),
t_repeat_ms,
15_000,
)
.unwrap();
assert!(out.all_tumbling_window_sizes_eq(t_repeat_ms));
}
Loading