diff --git a/crates/promql_utilities/src/ast_matching/mod.rs b/crates/promql_utilities/src/ast_matching/mod.rs deleted file mode 100644 index bcc58df9..00000000 --- a/crates/promql_utilities/src/ast_matching/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod promql_pattern; -pub mod promql_pattern_builder; -//pub mod promql_pattern_factory; - -pub use promql_pattern::*; -pub use promql_pattern_builder::*; -//pub use promql_pattern_factory::*; diff --git a/crates/promql_utilities/src/ast_matching/promql_pattern.rs b/crates/promql_utilities/src/ast_matching/promql_pattern.rs deleted file mode 100644 index a757f070..00000000 --- a/crates/promql_utilities/src/ast_matching/promql_pattern.rs +++ /dev/null @@ -1,1370 +0,0 @@ -use chrono::Duration; -use core::panic; -use promql_parser::label::MatchOp; -use promql_parser::parser::{AtModifier, Expr, LabelModifier, SubqueryExpr, VectorSelector}; -use serde::Serialize; -use serde_json::Value; -use std::collections::HashMap; -use std::time::UNIX_EPOCH; -use tracing::debug; - -/// PromQL pattern for AST-based matching -#[derive(Debug, Clone)] -pub struct PromQLPattern { - /// AST pattern definition (JSON-like structure). None indicates a wildcard (match any). - pub ast_pattern: Option>, - ///// Tokens to collect during matching - //pub collect_tokens: Vec, -} - -impl PromQLPattern { - /// Create a new pattern with AST pattern definition - //pub fn new(ast_pattern: Option>, collect_tokens: Vec) -> Self { - pub fn new(ast_pattern: Option>) -> Self { - debug!("Creating new PromQLPattern"); - Self { - ast_pattern, - //collect_tokens, - } - } - - /// Convert an Expr to a clean string representation - fn expr_to_string(expr: &Expr) -> String { - match expr { - Expr::NumberLiteral(num) => num.val.to_string(), - _ => format!("{:?}", expr), - } - } - - /// Match this pattern against a parsed AST - pub fn matches(&self, ast: &Expr) -> PromQLMatchResult { - debug!("Starting pattern matching against AST"); - debug!("Pattern: {:?}", self.ast_pattern); - debug!("AST: {:?}", ast); - let mut tokens = HashMap::new(); - let matches = self.matches_recursive(ast, self.ast_pattern.as_ref(), &mut tokens); - debug!( - "Pattern matching completed: {}, collected {} tokens", - matches, - tokens.len() - ); - if !matches { - debug!("MATCH FAILED - tokens collected: {:?}", tokens); - } - PromQLMatchResult { matches, tokens } - } - - /// Recursive pattern matching implementation - fn matches_recursive( - &self, - node: &Expr, - pattern: Option<&HashMap>, - tokens: &mut HashMap, - ) -> bool { - // None pattern is treated as wildcard (matches anything) to mirror Python's None - if pattern.is_none() { - debug!("Wildcard pattern matched"); - return true; - } - let pattern = pattern.unwrap(); - if pattern.is_empty() { - panic!("Empty pattern is not allowed"); - } - - // Get the pattern type - let pattern_type = match pattern.get("type") { - Some(Value::String(t)) => t.as_str(), - _ => panic!("Pattern must have a 'type' field of string type"), - }; - - debug!("Matching pattern type: {} against node type", pattern_type); - debug!("Full pattern: {:?}", pattern); - debug!("Node: {:?}", node); - match (pattern_type, node) { - // Match metric selectors - ("VectorSelector", Expr::VectorSelector(vs)) => { - self.match_metric_selector(vs, pattern, tokens) - } - - // Match function calls - ("Call", Expr::Call(call)) => self.match_function_call(call, pattern, tokens), - - // Match aggregation operations - ("AggregateExpr", Expr::Aggregate(agg)) => self.match_aggregation(agg, pattern, tokens), - - // Match matrix selectors (range vectors) - ("MatrixSelector", Expr::MatrixSelector(ms)) => { - self.match_matrix_selector(ms, pattern, tokens) - } - - // Match binary operations - ("BinaryExpr", Expr::Binary(bin_op)) => { - self.match_binary_operation(bin_op, pattern, tokens) - } - - // Match number literals - ("NumberLiteral", Expr::NumberLiteral(num)) => { - self.match_number_literal(num, pattern, tokens) - } - - // Match subquery expressions - ("SubqueryExpr", Expr::Subquery(subquery)) => { - self.match_subquery(subquery, pattern, tokens) - } - - _ => false, // Simply return false for non-matching types - } - } - - /// Match a VectorSelector node against pattern - fn match_metric_selector( - &self, - vs: &VectorSelector, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check metric name if specified in pattern - if let Some(Value::String(expected_name)) = pattern.get("name") { - if let Some(metric_name) = &vs.name { - if *metric_name != *expected_name { - return false; - } - } else { - return false; // Pattern expects name but node has none - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - debug!("Collecting metric token as: {}", collect_as); - let mut labels = HashMap::new(); - - // Extract label matchers - for matcher in &vs.matchers.matchers { - if matcher.op == MatchOp::Equal { - labels.insert(matcher.name.clone(), matcher.value.clone()); - } - } - - let at_modifier_opt = match &vs.at { - Some(AtModifier::At(t)) => { - // Convert SystemTime to seconds since UNIX_EPOCH (u64). - // Panic if time is earlier than UNIX_EPOCH (pre-epoch) as requested. - let secs = match t.duration_since(UNIX_EPOCH) { - Ok(dur) => dur.as_secs(), - Err(_) => panic!("AtModifier::At contains a time before UNIX_EPOCH, which is not supported by the pattern matcher"), - }; - - Some(secs) - } - Some(AtModifier::Start) => { - panic!("AtModifier::Start is not supported by pattern matcher") - } - Some(AtModifier::End) => { - panic!("AtModifier::End is not supported by pattern matcher") - } - None => None, - }; - - let metric_token = MetricToken { - name: vs.name.clone().unwrap_or_default(), - labels, - at_modifier: at_modifier_opt, - ast: Some(vs.clone()), - }; - - let token_data = TokenData { - metric: Some(metric_token), - function: None, - aggregation: None, - range_vector: None, - subquery: None, - binary_op: None, - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } - - /// Match a Call node (function call) against pattern - fn match_function_call( - &self, - call: &promql_parser::parser::Call, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check function name - // BUGFIX: Pattern builder creates "func" as an Object, not Array of Objects - // Original code (incorrect - expected Array of Objects): - // if let Some(Value::Array(expected_names)) = pattern.get("func") { - // if let Some(func_obj) = expected_names.first() { - // if let Some(func_map) = func_obj.as_object() { - // if let Some(Value::Array(names)) = func_map.get("name") { - // let function_name = call.func.name; - // let matches_name = names.iter().any(|name| { - // if let Some(name_str) = name.as_str() { - // name_str == function_name - // } else { - // false - // } - // }); - // - // if !matches_name { - // return false; - // } - // } - // } - // } - // } - - // Fixed code (correct - expects Object with "name" field): - if let Some(func_pattern_value) = pattern.get("func") { - if let Some(func_pattern) = func_pattern_value.as_object() { - if let Some(Value::Array(names)) = func_pattern.get("name") { - let function_name = call.func.name; - let matches_name = names.iter().any(|name| { - if let Some(name_str) = name.as_str() { - name_str == function_name - } else { - false - } - }); - - if !matches_name { - return false; - } - } - } - } - - // Check arguments recursively - if let Some(Value::Array(expected_args)) = pattern.get("args") { - if call.args.args.len() != expected_args.len() { - return false; - } - - for (i, arg) in call.args.args.iter().enumerate() { - if let Some(arg_pattern) = expected_args[i].as_object() { - let arg_pattern_map: HashMap = arg_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.matches_recursive(arg.as_ref(), Some(&arg_pattern_map), tokens) { - return false; - } - } - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - debug!("Collecting function token as: {}", collect_as); - let function_token = FunctionToken { - name: call.func.name.to_string(), - args: call - .args - .args - .iter() - .map(|arg| Self::expr_to_string(arg)) - .collect(), // Capture actual args - }; - - let token_data = TokenData { - metric: None, - function: Some(function_token), - aggregation: None, - range_vector: None, - subquery: None, - binary_op: None, - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - // If requested, collect the raw function arguments (as strings) under a separate token - if let Some(Value::String(collect_args_as)) = pattern.get("_collect_args_as") { - let arg_strs: Vec = call - .args - .args - .iter() - .map(|arg| Self::expr_to_string(arg)) - .collect(); - - let function_args_token = FunctionToken { - name: call.func.name.to_string(), - args: arg_strs, - }; - - let token_data = TokenData { - metric: None, - function: Some(function_args_token), - aggregation: None, - range_vector: None, - subquery: None, - binary_op: None, - number: None, - }; - - tokens.insert(collect_args_as.clone(), token_data); - } - - true - } - - /// Match an Aggregate node against pattern - fn match_aggregation( - &self, - agg: &promql_parser::parser::AggregateExpr, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - debug!("=== AGGREGATION MATCHING START ==="); - debug!("Aggregation pattern: {:?}", pattern); - debug!("Aggregation AST: {:?}", agg); - // Check aggregation operation - if let Some(Value::Array(expected_ops)) = pattern.get("op") { - let agg_op = agg.op.to_string(); - debug!( - "Checking aggregation op '{}' against pattern ops: {:?}", - agg_op, expected_ops - ); - let matches_op = expected_ops.iter().any(|op| { - if let Some(op_str) = op.as_str() { - op_str == agg_op - } else { - false - } - }); - - if !matches_op { - debug!("Aggregation op '{}' does not match pattern ops", agg_op); - return false; - } - debug!("Aggregation op '{}' matched!", agg_op); - } - - // Check inner expression recursively - if let Some(expr_pattern_value) = pattern.get("expr") { - debug!("Found expr pattern value: {:?}", expr_pattern_value); - if let Some(expr_pattern) = expr_pattern_value.as_object() { - debug!("Expr pattern is an object, recursing..."); - let expr_pattern_map: HashMap = expr_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.matches_recursive(&agg.expr, Some(&expr_pattern_map), tokens) { - debug!("Inner expression recursive match FAILED"); - return false; - } - debug!("Inner expression recursive match SUCCESS"); - } else if expr_pattern_value.is_null() { - debug!("Expr pattern is null, skipping validation"); - } else { - debug!( - "Expr pattern is neither object nor null: {:?}", - expr_pattern_value - ); - } - } else { - debug!("No expr pattern found, skipping inner expression check"); - } - - // Check modifier if specified in pattern - // Original code (too strict - fails when query has modifier but pattern is null): - // if let Some(pattern_modifier_value) = pattern.get("modifier") { - // let actual_modifier = match &agg.modifier { - // Some(LabelModifier::Include(_)) => "by", - // Some(LabelModifier::Exclude(_)) => "without", - // None => "null", - // }; - // - // match pattern_modifier_value { - // Value::String(expected_modifier) => { - // if actual_modifier != expected_modifier { - // return false; - // } - // } - // Value::Null => { - // if actual_modifier != "null" { - // return false; - // } - // } - // _ => { - // // Invalid pattern modifier format - // return false; - // } - // } - // } - - // Fixed code - only validate modifiers if pattern explicitly specifies a non-null modifier - if let Some(pattern_modifier_value) = pattern.get("modifier") { - debug!("Found modifier pattern: {:?}", pattern_modifier_value); - let actual_modifier = match &agg.modifier { - Some(LabelModifier::Include(_)) => "by", - Some(LabelModifier::Exclude(_)) => "without", - None => "null", - }; - debug!("Actual aggregation modifier: '{}'", actual_modifier); - - // Only validate if pattern explicitly requires a specific modifier (not null) - if !pattern_modifier_value.is_null() { - debug!("Pattern requires specific modifier, validating..."); - match pattern_modifier_value { - Value::String(expected_modifier) => { - debug!( - "Expected modifier: '{}', actual: '{}'", - expected_modifier, actual_modifier - ); - if actual_modifier != expected_modifier { - debug!("Modifier mismatch - FAILED"); - return false; - } - debug!("Modifier match - SUCCESS"); - } - _ => { - debug!("Invalid pattern modifier format - FAILED"); - return false; - } - } - } else { - debug!("Pattern modifier is null, allowing any query modifier (wildcard)"); - } - } else { - debug!("No modifier pattern found, allowing any query modifier"); - } - - debug!("=== AGGREGATION MATCHING SUCCESS ==="); - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - debug!("Collecting aggregation token as: {}", collect_as); - let modifier = match &agg.modifier { - Some(LabelModifier::Include(labels)) => Some(AggregationModifier { - modifier_type: AggregationModifierType::By, - labels: labels.labels.clone(), - }), - Some(LabelModifier::Exclude(labels)) => Some(AggregationModifier { - modifier_type: AggregationModifierType::Without, - labels: labels.labels.clone(), - }), - None => None, - }; - - let aggregation_token = AggregationToken { - op: agg.op.to_string(), - modifier, - param: agg.param.as_ref().map(|p| Self::expr_to_string(p)), - }; - - let token_data = TokenData { - metric: None, - function: None, - aggregation: Some(aggregation_token), - range_vector: None, - subquery: None, - binary_op: None, - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } - - /// Match a MatrixSelector node against pattern - fn match_matrix_selector( - &self, - ms: &promql_parser::parser::MatrixSelector, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check the inner vector selector - if let Some(vs_pattern_value) = pattern.get("vector_selector") { - if let Some(vs_pattern) = vs_pattern_value.as_object() { - let vs_pattern_map: HashMap = vs_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.match_metric_selector(&ms.vs, &vs_pattern_map, tokens) { - return false; - } - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - // Convert std::time::Duration to chrono::Duration and store directly - let chrono_dur = Duration::from_std(ms.range) - .map_err(|_| Duration::zero()) - .unwrap(); - - let range_token = RangeToken { - range: chrono_dur, - offset: ms.vs.offset.as_ref().map(|offset| format!("{:?}", offset)), - }; - - let token_data = TokenData { - metric: None, - function: None, - aggregation: None, - range_vector: Some(range_token), - subquery: None, - binary_op: None, - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } - - ///// Normalize duration to standard PromQL format (prefer larger units when possible) - // fn normalize_duration_string(duration: &std::time::Duration) -> String { - // let secs = duration.as_secs(); - - // // Convert to the most appropriate unit, preferring larger units when possible - // if secs >= 3600 && secs % 3600 == 0 { - // format!("{}h", secs / 3600) - // } else if secs >= 60 && secs % 60 == 0 { - // format!("{}m", secs / 60) - // } else if secs > 0 { - // format!("{secs}s") - // } else { - // // Handle sub-second durations - // let millis = duration.as_millis(); - // if millis > 0 { - // format!("{millis}ms") - // } else { - // "0s".to_string() - // } - // } - // } - - /// Match a Binary expression node against pattern - fn match_binary_operation( - &self, - bin_op: &promql_parser::parser::BinaryExpr, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check operation type - if let Some(Value::String(expected_op)) = pattern.get("op") { - if bin_op.op.to_string() != *expected_op { - return false; - } - } - - // Check left and right expressions recursively - if let Some(left_pattern_value) = pattern.get("left") { - if let Some(left_pattern) = left_pattern_value.as_object() { - let left_pattern_map: HashMap = left_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.matches_recursive(&bin_op.lhs, Some(&left_pattern_map), tokens) { - return false; - } - } - } - - if let Some(right_pattern_value) = pattern.get("right") { - if let Some(right_pattern) = right_pattern_value.as_object() { - let right_pattern_map: HashMap = right_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.matches_recursive(&bin_op.rhs, Some(&right_pattern_map), tokens) { - return false; - } - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - let binary_token = BinaryOpToken { - op: bin_op.op.to_string(), - matching: None, // TODO: Add vector matching support - }; - - let token_data = TokenData { - metric: None, - function: None, - aggregation: None, - range_vector: None, - subquery: None, - binary_op: Some(binary_token), - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } - - /// Match a NumberLiteral node against pattern - fn match_number_literal( - &self, - num: &promql_parser::parser::NumberLiteral, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check value if specified in pattern - if let Some(Value::Number(expected_value)) = pattern.get("value") { - if let Some(expected_f64) = expected_value.as_f64() { - if (num.val - expected_f64).abs() > f64::EPSILON { - return false; - } - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - let number_token = NumberToken { value: num.val }; - - let token_data = TokenData { - metric: None, - function: None, - aggregation: None, - range_vector: None, - subquery: None, - binary_op: None, - number: Some(number_token), - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } - - /// Match a SubqueryExpr node against pattern - fn match_subquery( - &self, - subquery: &SubqueryExpr, - pattern: &HashMap, - tokens: &mut HashMap, - ) -> bool { - // Check inner expression recursively - if let Some(expr_pattern_value) = pattern.get("expr") { - if let Some(expr_pattern) = expr_pattern_value.as_object() { - let expr_pattern_map: HashMap = expr_pattern - .clone() - .into_iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - if !self.matches_recursive(&subquery.expr, Some(&expr_pattern_map), tokens) { - return false; - } - } - } - - // Extract and store token data if this node should be collected - if let Some(Value::String(collect_as)) = pattern.get("_collect_as") { - // Convert std::time::Duration to chrono::Duration and store - let chrono_dur = Duration::from_std(subquery.range) - .map_err(|_| Duration::zero()) - .unwrap(); - - let subquery_token = SubqueryToken { - range: chrono_dur, - offset: subquery - .offset - .as_ref() - .map(|offset| format!("{:?}", offset)), - step: subquery.step.as_ref().map(|step| format!("{:?}", step)), - }; - - let token_data = TokenData { - metric: None, - function: None, - aggregation: None, - range_vector: None, - subquery: Some(subquery_token), - binary_op: None, - number: None, - }; - - tokens.insert(collect_as.clone(), token_data); - } - - true - } -} - -/// Token data extracted from AST nodes - pattern matching system -#[derive(Debug, Clone, Serialize)] -pub struct TokenData { - pub metric: Option, - pub function: Option, - pub aggregation: Option, - pub range_vector: Option, - pub subquery: Option, - pub binary_op: Option, - pub number: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct MetricToken { - pub name: String, - pub labels: HashMap, - // seconds since UNIX_EPOCH - pub at_modifier: Option, - #[serde(skip_serializing, skip_deserializing)] - pub ast: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct FunctionToken { - pub name: String, - pub args: Vec, -} - -#[derive(Debug, Clone, Serialize)] -pub struct AggregationToken { - pub op: String, - pub modifier: Option, - pub param: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct RangeToken { - pub range: Duration, - pub offset: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct SubqueryToken { - pub range: Duration, - pub offset: Option, - pub step: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct BinaryOpToken { - pub op: String, - pub matching: Option, -} - -#[derive(Debug, Clone, Serialize)] -pub struct NumberToken { - pub value: f64, -} - -#[derive(Debug, Clone, Serialize)] -pub struct VectorMatching { - pub card: String, // "one-to-one", "one-to-many", "many-to-one" - pub on: Vec, - pub ignoring: Vec, - pub group_left: Vec, - pub group_right: Vec, -} - -/// Match result with token-based extraction -#[derive(Debug, Clone)] -pub struct PromQLMatchResult { - pub matches: bool, - pub tokens: HashMap, -} - -impl PromQLMatchResult { - /// Create a new empty result - pub fn new() -> Self { - Self { - matches: false, - tokens: HashMap::new(), - } - } - - /// Create a successful match result with tokens - pub fn with_tokens(tokens: HashMap) -> Self { - Self { - matches: true, - tokens, - } - } - - /// Get metric name from tokens - pub fn get_metric_name(&self) -> Option { - self.tokens - .get("metric")? - .metric - .as_ref() - .map(|m| m.name.clone()) - } - - /// Get function name from tokens - pub fn get_function_name(&self) -> Option { - self.tokens - .get("function")? - .function - .as_ref() - .map(|f| f.name.clone()) - } - - /// Get aggregation operation from tokens - pub fn get_aggregation_op(&self) -> Option { - self.tokens - .get("aggregation")? - .aggregation - .as_ref() - .map(|a| a.op.clone()) - } - - /// Get range duration from tokens as chrono::Duration - pub fn get_range_duration(&self) -> Option { - self.tokens - .get("range_vector")? - .range_vector - .as_ref() - .map(|r| r.range) - } -} - -impl Default for PromQLMatchResult { - fn default() -> Self { - Self::new() - } -} - -/// Whether a PromQL aggregation modifier is `by` or `without`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)] -#[serde(rename_all = "lowercase")] -pub enum AggregationModifierType { - By, - Without, -} - -/// Represents aggregation modifiers like "by" or "without" -#[derive(Debug, Clone, Serialize)] -pub struct AggregationModifier { - pub modifier_type: AggregationModifierType, - pub labels: Vec, -} - -impl AggregationModifier { - /// Create a new AggregationModifier - pub fn new(modifier_type: AggregationModifierType, labels: Vec) -> Self { - Self { - modifier_type, - labels, - } - } - - // /// Check if a function name represents a temporal function - // fn is_temporal_function(&self, function_name: &str) -> bool { - // matches!( - // function_name, - // "rate" - // | "increase" - // | "sum_over_time" - // | "min_over_time" - // | "max_over_time" - // | "avg_over_time" - // | "count_over_time" - // | "quantile_over_time" - // | "stddev_over_time" - // | "stdvar_over_time" - // | "last_over_time" - // | "present_over_time" - // ) - // } - - // /// Extract label filters from matchers - // fn extract_label_filters(&self, matchers: &Matchers) -> HashMap { - // let mut filters = HashMap::new(); - - // for matcher in &matchers.matchers { - // // For now, only handle exact equality matches - // if matcher.op == MatchOp::Equal { - // filters.insert(matcher.name.clone(), matcher.value.clone()); - // } - // } - - // filters - // } - - // /// Convert Duration to string representation in PromQL format - // fn duration_to_string(&self, duration: &std::time::Duration) -> String { - // let secs = duration.as_secs(); - - // // Convert to the most appropriate unit, preferring larger units when possible - // if secs >= 3600 && secs % 3600 == 0 { - // format!("{}h", secs / 3600) - // } else if secs >= 60 && secs % 60 == 0 { - // format!("{}m", secs / 60) - // } else if secs > 0 { - // format!("{secs}s") - // } else { - // // Handle sub-second durations - // let millis = duration.as_millis(); - // if millis > 0 { - // format!("{millis}ms") - // } else { - // "0s".to_string() - // } - // } - // } -} - -// ============================================================================= -// Tests migrated from asap-common/tests/{compare_matched_tokens, -// rust_pattern_matching, compare_patterns} binary runners. -// -// The PatternTester struct and all build_* methods below are copied verbatim -// from compare_matched_tokens/rust_tests/src/pattern_tests.rs. Only the -// main() harness has been replaced with individual #[test] functions, and -// assertions come from the test_data/promql_queries.json file that was used -// by the binary runner. -// ============================================================================= -#[cfg(test)] -mod tests { - use super::*; - use crate::ast_matching::PromQLPatternBuilder; - use promql_parser::parser as promql; - use serde_json::Value; - - // ------------------------------------------------------------------ - // PatternTester — copied from compare_matched_tokens/rust_tests/src/pattern_tests.rs - // ------------------------------------------------------------------ - - struct PatternTester { - patterns: HashMap>, - } - - impl PatternTester { - fn new() -> Self { - let mut patterns = HashMap::new(); - - // ONLY_TEMPORAL patterns - let temporal_patterns = vec![ - // Rate pattern - PromQLPattern::new(Self::build_rate_pattern()), - // Quantile over time pattern - PromQLPattern::new(Self::build_quantile_over_time_pattern()), - ]; - - // ONLY_SPATIAL patterns - let spatial_patterns = vec![ - // Sum aggregation pattern - PromQLPattern::new(Self::build_sum_pattern()), - // Simple metric pattern - PromQLPattern::new(Self::build_metric_pattern()), - ]; - - // ONE_TEMPORAL_ONE_SPATIAL patterns - let combined_patterns = vec![ - // Aggregation of single-arg temporal functions - PromQLPattern::new(Self::build_one_temporal_one_spatial_pattern()), - // Aggregation of quantile_over_time (2-arg) - PromQLPattern::new(Self::build_combined_quantile_pattern()), - ]; - - patterns.insert("ONLY_SPATIAL".to_string(), spatial_patterns); - patterns.insert("ONLY_TEMPORAL".to_string(), temporal_patterns); - patterns.insert("ONE_TEMPORAL_ONE_SPATIAL".to_string(), combined_patterns); - - Self { patterns } - } - - fn classify_query(&self, query: &str) -> Option<(String, PromQLMatchResult)> { - let ast = promql::parse(query).expect("Failed to parse query"); - - for (pattern_type, pattern_list) in &self.patterns { - for pattern in pattern_list { - let match_result = pattern.matches(&ast); - if match_result.matches { - let final_type = if pattern_type == "ONLY_SPATIAL" { - if match_result.tokens.contains_key("aggregation") { - pattern_type.clone() - } else if match_result.tokens.contains_key("metric") { - "ONLY_VECTOR".to_string() - } else { - pattern_type.clone() - } - } else { - pattern_type.clone() - }; - return Some((final_type, match_result)); - } - } - } - None - } - - fn build_rate_pattern() -> Option> { - let ms = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - - let args: Vec>> = vec![ms]; - - PromQLPatternBuilder::function( - vec![ - "rate", - "increase", - "avg_over_time", - "sum_over_time", - "count_over_time", - "min_over_time", - "max_over_time", - ], - args, - Some("function"), - None, - ) - } - - fn build_quantile_over_time_pattern() -> Option> { - let num = PromQLPatternBuilder::number(None, None); - let ms = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - - let args: Vec>> = vec![num, ms]; - - PromQLPatternBuilder::function( - vec!["quantile_over_time"], - args, - Some("function"), - Some("function_args"), - ) - } - - fn build_sum_pattern() -> Option> { - PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "min", "max"], - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - None, - None, - Some("aggregation"), - ) - } - - fn build_metric_pattern() -> Option> { - PromQLPatternBuilder::metric(None, None, None, Some("metric")) - } - - fn build_one_temporal_one_spatial_pattern() -> Option> { - let ms = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - - let func_args: Vec>> = vec![ms]; - - let func = PromQLPatternBuilder::function( - vec![ - "sum_over_time", - "count_over_time", - "avg_over_time", - "min_over_time", - "max_over_time", - "rate", - "increase", - ], - func_args, - Some("function"), - None, - ); - - PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max"], - func, - None, - None, - None, - Some("aggregation"), - ) - } - - fn build_combined_quantile_pattern() -> Option> { - let num = PromQLPatternBuilder::number(None, None); - let ms = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - let func_args: Vec>> = vec![num, ms]; - let func = PromQLPatternBuilder::function( - vec!["quantile_over_time"], - func_args, - Some("function"), - None, - ); - PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "quantile", "min", "max"], - func, - None, - None, - None, - Some("aggregation"), - ) - } - - #[allow(dead_code)] - fn build_sum_rate_pattern() -> Option> { - let ms = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - - let func_args: Vec>> = vec![ms]; - - let func = PromQLPatternBuilder::function( - vec!["rate", "increase"], - func_args, - Some("function"), - None, - ); - - PromQLPatternBuilder::aggregation( - vec!["sum", "count", "avg", "min", "max"], - func, - None, - None, - None, - Some("aggregation"), - ) - } - } - - // ------------------------------------------------------------------ - // Tests from compare_matched_tokens/test_data/promql_queries.json - // ------------------------------------------------------------------ - - #[test] - fn temporal_rate_basic() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("rate(http_requests_total{job=\"api\"}[5m])") - .unwrap(); - assert_eq!(cat, "ONLY_TEMPORAL"); - assert_eq!(result.get_metric_name().unwrap(), "http_requests_total"); - assert_eq!(result.get_function_name().unwrap(), "rate"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(5)); - let labels = &result.tokens["metric"].metric.as_ref().unwrap().labels; - assert_eq!(labels.get("job").unwrap(), "api"); - } - - #[test] - fn temporal_increase_basic() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("increase(http_requests_total[1h])") - .unwrap(); - assert_eq!(cat, "ONLY_TEMPORAL"); - assert_eq!(result.get_metric_name().unwrap(), "http_requests_total"); - assert_eq!(result.get_function_name().unwrap(), "increase"); - assert_eq!(result.get_range_duration().unwrap(), Duration::hours(1)); - } - - #[test] - fn temporal_quantile_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("quantile_over_time(0.95, cpu_usage{instance=\"host1\"}[10m])") - .unwrap(); - assert_eq!(cat, "ONLY_TEMPORAL"); - assert_eq!(result.get_metric_name().unwrap(), "cpu_usage"); - assert_eq!(result.get_function_name().unwrap(), "quantile_over_time"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(10)); - let labels = &result.tokens["metric"].metric.as_ref().unwrap().labels; - assert_eq!(labels.get("instance").unwrap(), "host1"); - } - - #[test] - fn temporal_avg_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("avg_over_time(memory_bytes[30m])") - .unwrap(); - assert_eq!(cat, "ONLY_TEMPORAL"); - assert_eq!(result.get_metric_name().unwrap(), "memory_bytes"); - assert_eq!(result.get_function_name().unwrap(), "avg_over_time"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(30)); - } - - #[test] - fn spatial_sum_aggregation() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("sum(http_requests_total{job=\"api\"})") - .unwrap(); - assert_eq!(cat, "ONLY_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "http_requests_total"); - assert_eq!(result.get_aggregation_op().unwrap(), "sum"); - let labels = &result.tokens["metric"].metric.as_ref().unwrap().labels; - assert_eq!(labels.get("job").unwrap(), "api"); - } - - #[test] - fn spatial_avg_aggregation() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("avg by (instance) (cpu_usage)") - .unwrap(); - assert_eq!(cat, "ONLY_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "cpu_usage"); - assert_eq!(result.get_aggregation_op().unwrap(), "avg"); - } - - #[test] - fn spatial_count_aggregation() { - let tester = PatternTester::new(); - let (cat, result) = tester.classify_query("count(up{job=\"node\"})").unwrap(); - assert_eq!(cat, "ONLY_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "up"); - assert_eq!(result.get_aggregation_op().unwrap(), "count"); - let labels = &result.tokens["metric"].metric.as_ref().unwrap().labels; - assert_eq!(labels.get("job").unwrap(), "node"); - } - - #[test] - fn combined_sum_of_rate() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("sum(rate(http_requests_total{job=\"api\"}[5m]))") - .unwrap(); - assert_eq!(cat, "ONE_TEMPORAL_ONE_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "http_requests_total"); - assert_eq!(result.get_function_name().unwrap(), "rate"); - assert_eq!(result.get_aggregation_op().unwrap(), "sum"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(5)); - } - - #[test] - fn combined_avg_of_quantile_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("avg(quantile_over_time(0.99, response_time_seconds[15m]))") - .unwrap(); - assert_eq!(cat, "ONE_TEMPORAL_ONE_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "response_time_seconds"); - assert_eq!(result.get_function_name().unwrap(), "quantile_over_time"); - assert_eq!(result.get_aggregation_op().unwrap(), "avg"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(15)); - } - - #[test] - fn combined_sum_of_avg_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("sum by (job) (avg_over_time(memory_bytes{env=\"prod\"}[1h]))") - .unwrap(); - assert_eq!(cat, "ONE_TEMPORAL_ONE_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "memory_bytes"); - assert_eq!(result.get_function_name().unwrap(), "avg_over_time"); - assert_eq!(result.get_aggregation_op().unwrap(), "sum"); - assert_eq!(result.get_range_duration().unwrap(), Duration::hours(1)); - let labels = &result.tokens["metric"].metric.as_ref().unwrap().labels; - assert_eq!(labels.get("env").unwrap(), "prod"); - } - - // ------------------------------------------------------------------ - // Tests from rust_pattern_matching binary - // ------------------------------------------------------------------ - - #[test] - fn spatial_of_temporal_sum_of_sum_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("sum by (instance, job) (sum_over_time(fake_metric_total[1m]))") - .unwrap(); - assert_eq!(cat, "ONE_TEMPORAL_ONE_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "fake_metric_total"); - assert_eq!(result.get_function_name().unwrap(), "sum_over_time"); - assert_eq!(result.get_aggregation_op().unwrap(), "sum"); - assert_eq!(result.get_range_duration().unwrap(), Duration::minutes(1)); - } - - #[test] - fn spatial_of_temporal_sum_of_count_over_time() { - let tester = PatternTester::new(); - let (cat, result) = tester - .classify_query("sum by (instance, job) (count_over_time(fake_metric_total[1m]))") - .unwrap(); - assert_eq!(cat, "ONE_TEMPORAL_ONE_SPATIAL"); - assert_eq!(result.get_metric_name().unwrap(), "fake_metric_total"); - assert_eq!(result.get_function_name().unwrap(), "count_over_time"); - assert_eq!(result.get_aggregation_op().unwrap(), "sum"); - } - - // ------------------------------------------------------------------ - // Tests from compare_patterns binary (pattern construction) - // ------------------------------------------------------------------ - - #[test] - fn pattern_builds_temporal_rate_increase() { - let ast = PatternTester::build_rate_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "Call"); - } - - #[test] - fn pattern_builds_temporal_quantile_over_time() { - let ast = PatternTester::build_quantile_over_time_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "Call"); - } - - #[test] - fn pattern_builds_spatial_aggregation() { - let ast = PatternTester::build_sum_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "AggregateExpr"); - } - - #[test] - fn pattern_builds_metric() { - let ast = PatternTester::build_metric_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "VectorSelector"); - } - - #[test] - fn pattern_builds_combined_temporal() { - let ast = PatternTester::build_one_temporal_one_spatial_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "AggregateExpr"); - } - - #[test] - fn pattern_builds_combined_quantile() { - let ast = PatternTester::build_combined_quantile_pattern(); - assert!(ast.is_some()); - assert_eq!(ast.unwrap()["type"], "AggregateExpr"); - } - - #[test] - fn bare_metric_classified_as_only_vector() { - let tester = PatternTester::new(); - let (cat, result) = tester.classify_query("http_requests_total").unwrap(); - assert_eq!(cat, "ONLY_VECTOR"); - assert_eq!(result.get_metric_name().unwrap(), "http_requests_total"); - } -} diff --git a/crates/promql_utilities/src/ast_matching/promql_pattern_builder.rs b/crates/promql_utilities/src/ast_matching/promql_pattern_builder.rs deleted file mode 100644 index 5a45a928..00000000 --- a/crates/promql_utilities/src/ast_matching/promql_pattern_builder.rs +++ /dev/null @@ -1,238 +0,0 @@ -use serde_json::Value; -use std::collections::HashMap; -use tracing::debug; - -/// PromQL Pattern Builder for creating PromQL-based patterns -/// This mirrors the Python PromQLPatternBuilder class -pub struct PromQLPatternBuilder; - -impl PromQLPatternBuilder { - /// Create a pattern for any node type - pub fn any() -> Option> { - debug!("Creating wildcard pattern (any)"); - None - } - - /// Create a binary operation pattern (BinaryExpr) - pub fn binary_op( - op: &str, - left: Option>, - right: Option>, - collect_as: Option<&str>, - ) -> Option> { - debug!("Creating binary operation pattern for op: {}", op); - let mut pattern = HashMap::new(); - pattern.insert("type".to_string(), Value::String("BinaryExpr".to_string())); - pattern.insert("op".to_string(), Value::String(op.to_string())); - pattern.insert("left".to_string(), serde_json::to_value(left).unwrap()); - pattern.insert("right".to_string(), serde_json::to_value(right).unwrap()); - - match collect_as { - Some(collect) => pattern.insert( - "_collect_as".to_string(), - Value::String(collect.to_string()), - ), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create a metric pattern (VectorSelector) - pub fn metric( - name: Option<&str>, - labels: Option>, - at_modifier: Option<&str>, - collect_as: Option<&str>, - ) -> Option> { - debug!("Creating metric pattern for name: {:?}", name); - let mut pattern = HashMap::new(); - pattern.insert( - "type".to_string(), - Value::String("VectorSelector".to_string()), - ); - - match name { - Some(n) => pattern.insert("name".to_string(), Value::String(n.to_string())), - None => pattern.insert("name".to_string(), Value::Null), - }; - - match labels { - Some(l) => { - let labels_value = serde_json::to_value(l).unwrap(); - pattern.insert("matchers".to_string(), labels_value) - } - None => pattern.insert("matchers".to_string(), Value::Null), - }; - - match at_modifier { - Some(a) => pattern.insert("at".to_string(), Value::String(a.to_string())), - None => pattern.insert("at".to_string(), Value::Null), - }; - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create a function pattern (Call) - pub fn function( - names: Vec<&str>, - args: Vec>>, - collect_as: Option<&str>, - collect_args_as: Option<&str>, - ) -> Option> { - debug!("Creating function pattern for names: {:?}", names); - let mut pattern = HashMap::new(); - pattern.insert("type".to_string(), Value::String("Call".to_string())); - - let mut func = HashMap::new(); - func.insert("type".to_string(), Value::String("Function".to_string())); - func.insert( - "name".to_string(), - Value::Array(names.iter().map(|n| Value::String(n.to_string())).collect()), - ); - - pattern.insert("func".to_string(), serde_json::to_value(func).unwrap()); - pattern.insert("args".to_string(), serde_json::to_value(args).unwrap()); - - match collect_args_as { - Some(c) => pattern.insert("_collect_args_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_args_as".to_string(), Value::Null), - }; - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create a subquery pattern (SubqueryExpr) - pub fn subquery( - expr: Option>, - duration: Option<&str>, - collect_as: Option<&str>, - ) -> Option> { - let mut pattern = HashMap::new(); - pattern.insert( - "type".to_string(), - Value::String("SubqueryExpr".to_string()), - ); - pattern.insert("expr".to_string(), serde_json::to_value(expr).unwrap()); - - match duration { - Some(d) => pattern.insert("range".to_string(), Value::String(d.to_string())), - None => pattern.insert("range".to_string(), Value::Null), - }; - - // Initialize step and offset as null, matching Python implementation - pattern.insert("step".to_string(), Value::Null); - pattern.insert("offset".to_string(), Value::Null); - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create a matrix selector pattern (MatrixSelector) - pub fn matrix_selector( - vector_selector: Option>, - range: Option<&str>, - collect_as: Option<&str>, - ) -> Option> { - let mut pattern = HashMap::new(); - pattern.insert( - "type".to_string(), - Value::String("MatrixSelector".to_string()), - ); - pattern.insert( - "vector_selector".to_string(), - serde_json::to_value(vector_selector).unwrap(), - ); - - match range { - Some(r) => pattern.insert("range".to_string(), Value::String(r.to_string())), - None => pattern.insert("range".to_string(), Value::Null), - }; - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create an aggregation pattern (AggregateExpr) - pub fn aggregation( - ops: Vec<&str>, - expr: Option>, - param: Option>, - by_labels: Option>, - without_labels: Option>, - collect_as: Option<&str>, - ) -> Option> { - let mut pattern = HashMap::new(); - pattern.insert( - "type".to_string(), - Value::String("AggregateExpr".to_string()), - ); - pattern.insert( - "op".to_string(), - Value::Array(ops.iter().map(|op| Value::String(op.to_string())).collect()), - ); - pattern.insert("expr".to_string(), serde_json::to_value(expr).unwrap()); - - match param { - Some(p) => pattern.insert("param".to_string(), serde_json::to_value(p).unwrap()), - None => pattern.insert("param".to_string(), Value::Null), - }; - - // Use single "modifier" field to match Python format - let modifier_value = match (by_labels, without_labels) { - (Some(_), None) => Value::String("by".to_string()), - (None, Some(_)) => Value::String("without".to_string()), - _ => Value::Null, - }; - pattern.insert("modifier".to_string(), modifier_value); - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } - - /// Create a number literal pattern - pub fn number(value: Option, collect_as: Option<&str>) -> Option> { - let mut pattern = HashMap::new(); - pattern.insert( - "type".to_string(), - Value::String("NumberLiteral".to_string()), - ); - - match value { - Some(v) => pattern.insert( - "value".to_string(), - Value::Number(serde_json::Number::from_f64(v).unwrap()), - ), - None => pattern.insert("value".to_string(), Value::Null), - }; - - match collect_as { - Some(c) => pattern.insert("_collect_as".to_string(), Value::String(c.to_string())), - None => pattern.insert("_collect_as".to_string(), Value::Null), - }; - - Some(pattern) - } -} diff --git a/crates/promql_utilities/src/ast_matching/promql_pattern_factory.rs b/crates/promql_utilities/src/ast_matching/promql_pattern_factory.rs deleted file mode 100644 index d8181ec9..00000000 --- a/crates/promql_utilities/src/ast_matching/promql_pattern_factory.rs +++ /dev/null @@ -1,123 +0,0 @@ -//use crate::ast_matching::{PromQLPattern, PromQLPatternBuilder}; -//use tracing::debug; -// -///// Pattern factory for creating common PromQL patterns -//pub struct PromQLPatternFactory; -// -//impl PromQLPatternFactory { -// /// Create pattern for OnlyTemporal queries (e.g., rate(metric[5m])) -// pub fn only_temporal_pattern() -> PromQLPattern { -// debug!("Creating only temporal pattern"); -// let ms = PromQLPatternBuilder::matrix_selector( -// PromQLPatternBuilder::metric(None, None, None, Some("metric")), -// None, -// Some("range_vector"), -// ); -// -// let func_args: Vec>> = vec![ms]; -// -// let pattern = PromQLPatternBuilder::function( -// vec![ -// "rate", -// "increase", -// "sum_over_time", -// "avg_over_time", -// "min_over_time", -// "max_over_time", -// "count_over_time", -// ], -// func_args, -// Some("function"), -// None, -// ); -// -// PromQLPattern::new( -// pattern, -// //vec![ -// // "metric".to_string(), -// // "function".to_string(), -// // "range_vector".to_string(), -// //], -// // QueryPatternType::OnlyTemporal, -// ) -// } -// -// /// Create pattern for OnlySpatial queries (e.g., sum(metric) by (label)) -// pub fn only_spatial_pattern() -> PromQLPattern { -// debug!("Creating only spatial pattern"); -// let metric = PromQLPatternBuilder::metric(None, None, None, Some("metric")); -// -// let pattern = PromQLPatternBuilder::aggregation( -// vec!["sum", "count", "avg", "min", "max", "quantile"], -// metric, -// None, -// None, -// None, -// Some("aggregation"), -// ); -// -// PromQLPattern::new( -// pattern, -// //vec!["metric".to_string(), "aggregation".to_string()], -// // QueryPatternType::OnlySpatial, -// ) -// } -// -// /// Create pattern for OneTemporalOneSpatial queries (e.g., sum(rate(metric[5m])) by (label)) -// pub fn one_temporal_one_spatial_pattern() -> PromQLPattern { -// debug!("Creating one temporal one spatial pattern"); -// let ms2 = PromQLPatternBuilder::matrix_selector( -// PromQLPatternBuilder::metric(None, None, None, Some("metric")), -// None, -// Some("range_vector"), -// ); -// -// let func_args2: Vec>> = -// vec![ms2]; -// -// let temporal_part = PromQLPatternBuilder::function( -// vec![ -// "rate", -// "increase", -// "sum_over_time", -// "avg_over_time", -// "min_over_time", -// "max_over_time", -// "count_over_time", -// ], -// func_args2, -// Some("function"), -// None, -// ); -// -// let pattern = PromQLPatternBuilder::aggregation( -// vec!["sum", "count", "avg", "min", "max", "quantile"], -// temporal_part, -// None, -// None, -// None, -// Some("aggregation"), -// ); -// -// PromQLPattern::new( -// pattern, -// //vec![ -// // "metric".to_string(), -// // "function".to_string(), -// // "range_vector".to_string(), -// // "aggregation".to_string(), -// //], -// // QueryPatternType::OneTemporalOneSpatial, -// ) -// } -// -// /// Get all standard patterns -// pub fn get_all_patterns() -> Vec { -// debug!("Getting all standard patterns"); -// vec![ -// Self::one_temporal_one_spatial_pattern(), -// Self::only_temporal_pattern(), -// Self::only_spatial_pattern(), -// ] -// } -//} diff --git a/crates/promql_utilities/src/lib.rs b/crates/promql_utilities/src/lib.rs index 5de6fa3a..7bbe5e9d 100644 --- a/crates/promql_utilities/src/lib.rs +++ b/crates/promql_utilities/src/lib.rs @@ -1,7 +1,5 @@ -pub mod ast_matching; pub mod data_model; pub mod query_logics; -pub use ast_matching::*; pub use data_model::*; pub use query_logics::*; diff --git a/crates/promql_utilities/src/query_logics/mod.rs b/crates/promql_utilities/src/query_logics/mod.rs index f3a98d12..8ca2b7df 100644 --- a/crates/promql_utilities/src/query_logics/mod.rs +++ b/crates/promql_utilities/src/query_logics/mod.rs @@ -1,7 +1,5 @@ pub mod enums; pub mod logics; -pub mod parsing; pub use enums::*; pub use logics::*; -pub use parsing::*; diff --git a/crates/promql_utilities/src/query_logics/parsing.rs b/crates/promql_utilities/src/query_logics/parsing.rs deleted file mode 100644 index e5bb8757..00000000 --- a/crates/promql_utilities/src/query_logics/parsing.rs +++ /dev/null @@ -1,135 +0,0 @@ -use core::panic; - -use promql_parser::parser::Expr; -use tracing::debug; - -use crate::ast_matching::promql_pattern::AggregationModifierType; -use crate::ast_matching::PromQLMatchResult; -use crate::data_model::KeyByLabelNames; -use crate::query_logics::enums::{AggregationOperator, QueryPatternType, Statistic}; - -pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (String, String) { - debug!("Extracting metric and spatial filter from match result"); - let mut metric_name = match_result.get_metric_name().unwrap_or_default(); - debug!("Initial metric name: {}", metric_name); - - let spatial_filter = if let Some(metric_token) = match_result - .tokens - .get("metric") - .and_then(|token| token.metric.as_ref()) - { - if let Some(ast_vs) = metric_token.ast.as_ref() { - // Render the VectorSelector AST to string and extract inner `{...}` content - // let ast_str = format!("{}", ast_vs); - let ast_str = Expr::from(ast_vs.clone()).prettify(); - if let Some(inner) = ast_str.split('{').nth(1).and_then(|s| s.split('}').next()) { - debug!("Found spatial filter content: {}", inner); - // Ensure metric_name does not include the selector part - metric_name = metric_name - .split('{') - .next() - .unwrap_or(&metric_name) - .to_string(); - debug!("Cleaned metric name: {}", metric_name); - inner.to_string() - } else { - String::new() - } - } else { - // No AST available -> return empty spatial filter (no fallback reconstruction) - String::new() - } - } else { - String::new() - }; - - debug!( - "Final result - metric: {}, spatial_filter: {}", - metric_name, spatial_filter - ); - (metric_name, spatial_filter) -} - -/// Get statistics to compute based on pattern type and tokens -pub fn get_statistics_to_compute( - pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, -) -> Vec { - debug!("Computing statistics for pattern type {:?}", pattern_type); - let statistic_to_compute: Option = if pattern_type == QueryPatternType::OnlyTemporal - || pattern_type == QueryPatternType::OneTemporalOneSpatial - { - match_result.get_function_name().map(|function_name| { - let name = function_name.to_lowercase(); - name.split('_').next().unwrap_or(&name).to_string() - }) - } else if pattern_type == QueryPatternType::OnlySpatial { - match_result - .get_aggregation_op() - .map(|agg| agg.to_lowercase()) - } else { - panic!("Unsupported query pattern type"); - }; - - if let Some(statistic_to_compute) = statistic_to_compute { - debug!("Found statistic to compute: {}", statistic_to_compute); - if statistic_to_compute.parse::() == Ok(AggregationOperator::Avg) { - vec![Statistic::Sum, Statistic::Count] - } else if let Ok(stat) = statistic_to_compute.parse::() { - vec![stat] - } else { - panic!("Unsupported statistic: {}", statistic_to_compute); - } - } else { - panic!("No statistic found in the query"); - } -} - -pub fn get_spatial_aggregation_output_labels( - match_result: &PromQLMatchResult, - all_labels: &KeyByLabelNames, -) -> KeyByLabelNames { - debug!("Getting spatial aggregation output labels"); - debug!("All labels: {:?}", all_labels); - // Match Python behaviour: assume aggregation token and modifier exist - // and raise (panic) if missing or invalid. "by" and "without" logic - // remain the same. - let aggregation_token = match_result - .tokens - .get("aggregation") - .and_then(|token| token.aggregation.as_ref()) - .expect("aggregation token missing"); - - // Patching: When the query is topk, we should always return all labels - if aggregation_token.op.parse::() == Ok(AggregationOperator::Topk) { - debug!("Aggregation operation is 'topk', returning all labels"); - return all_labels.clone(); - } - - // Fixing issue https://github.com/ProjectASAP/asap-internal/issues/24 - let modifier: &crate::AggregationModifier = match aggregation_token.modifier.as_ref() { - Some(m) => m, - None => { - debug!("No aggregation modifier found, returning empty KeyByLabelNames"); - return KeyByLabelNames::new(vec![]); - } - }; - - debug!( - "Modifier type: {:?}, labels: {:?}", - modifier.modifier_type, modifier.labels - ); - match modifier.modifier_type { - AggregationModifierType::By => { - debug!("Processing 'by' modifier"); - // Return only the labels specified in "by" clause - KeyByLabelNames::new(modifier.labels.clone()) - } - AggregationModifierType::Without => { - debug!("Processing 'without' modifier"); - // Return all labels except those specified in "without" clause - let without_labels = KeyByLabelNames::new(modifier.labels.clone()); - all_labels.difference(&without_labels) - } - } -} diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index c69419ed..23e5b7a8 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -888,12 +888,19 @@ fn parse_last_over_time_probe(query: &str) -> Option<(String, i64)> { Some((metric, range_ms)) } -/// Direct `ASAPQueryEngine::handle_query` dispatch — preserves the -/// `KeyByLabelNames` the Prometheus adapter needs to fill in the -/// `metric` map. Used for ASAP-tier metrics (the default) so the -/// response surface is byte-identical to the pre-router path. Adds a +/// Direct `ASAPQueryEngine::execute(&str)` dispatch — B7.5 retired +/// the legacy `handle_query` path; this handler is now a thin +/// wrapper around the modern `QueryEngine::execute(&str)` trait +/// surface, which classifies via the analyzer + ASAP-tier reducer +/// and fires capability-miss notifies natively. Adds a /// `data_source: asap_query` info-line at the JSON layer so Phase-6 /// callers can byte-compare regardless of the dispatch path. +/// +/// Trait dispatch loses `KeyByLabelNames` (the trait returns just +/// `QueryResult`); we surface an empty `KeyByLabelNames`, identical +/// to how `process_via_router` handles the same trait surface — the +/// Prometheus adapter renders an empty `metric: {}` object, a valid +/// shape that PromQL clients accept. async fn process_via_simple_engine( state: &AppState, parsed_request: &ParsedQueryRequest, @@ -902,37 +909,28 @@ async fn process_via_simple_engine( ) -> Response { let query_start_time = Instant::now(); debug!( - "About to call query_engine.handle_query with query='{}' and time={}", + "About to call query_engine.execute with query='{}' and time={}", parsed_request.query, parsed_request.time ); - match state - .query_engine - .handle_query(parsed_request.query.clone(), parsed_request.time) - { - Some((query_output_labels, query_result)) => { + use crate::query_engines::routing::query_engine_routing::QueryEngine; + use crate::drivers::query::adapters::QueryExecutionResult; + match state.query_engine.execute(&parsed_request.query).await { + Ok(query_result) => { let query_duration = query_start_time.elapsed(); - debug!("=== QUERY ENGINE SUCCESS ==="); debug!( - "Query engine execution took: {:.2}ms", + "Modern execute() succeeded for query='{}' in {:.2}ms", + parsed_request.query, query_duration.as_secs_f64() * 1000.0 ); - debug!("Query output labels: {:?}", query_output_labels); - debug!("Query result: {:?}", query_result); - - // Step 3: Format success response using adapter - // (Adapter handles protocol-specific formatting, e.g., convert_query_result_to_prometheus) - use crate::drivers::query::adapters::QueryExecutionResult; let execution_result = QueryExecutionResult { - query_output_labels, - query_result}; - + query_output_labels: promql_utilities::data_model::KeyByLabelNames::default(), + query_result, + }; let total_duration = start_time.elapsed(); debug!( "Total request processing took: {:.2}ms", total_duration.as_secs_f64() * 1000.0 ); - debug!("=== RETURNING SUCCESS RESPONSE ==="); - match state .adapter .format_success_response(&execution_result) @@ -942,85 +940,18 @@ async fn process_via_simple_engine( annotate_data_source(response, StorageBackend::SketchStore.data_source_id()) .await } - Err(status) => status.into_response()} - } - None => { - // Legacy `handle_query` returned None — likely the - // sketch-vs-precompute query gap pinned in #252: - // `query_precomputes_by_agg` only picks up - // `AggKind::ExactAgg` sids, never `AggKind::Sketch`. Try - // the modern `ASAPQueryEngine::execute(&str)` trait path - // before falling through to the unsupported-query branch - // — `execute` uses - // `idx.sids_for_policy(fp)` + `SketchReducer::evaluate` - // and handles sketches natively, AND (since #273) unions - // `instances_matching` for sketches with `policy_fp = - // UNSET`. - // - // Trait dispatch loses `KeyByLabelNames` (the trait - // returns just `QueryResult`); we surface an empty - // `KeyByLabelNames`, identical to how `process_via_router` - // handles the same trait surface — the Prometheus - // adapter renders an empty `metric: {}` object, a valid - // shape that PromQL clients accept. - // - // Schema-retirement #5 status: an earlier draft of this - // PR reordered to "modern first, legacy as fallback" so - // the legacy path could be retired entirely. That broke - // `http_capability_miss_feedback_loop_closes_over_http` - // — the capability-miss notify side-effect happens - // inside legacy `find_compatible_aggregation_with_miss_notify` - // (engine.rs:~1772), and a pre-existing time=0 underflow - // bug at engine.rs:792 surfaces when legacy is reached - // via the modern-Err fallback because of subtle test - // setup state. Modern needs to spawn its own - // capability-miss notify before we can reorder cleanly. - use crate::query_engines::routing::query_engine_routing::QueryEngine; - let modern_result = state - .query_engine - .execute(&parsed_request.query) - .await; - if let Ok(query_result) = modern_result { - debug!( - "Modern execute() handled what legacy handle_query missed \ - (query='{}')", - parsed_request.query - ); - use crate::drivers::query::adapters::QueryExecutionResult; - let execution_result = QueryExecutionResult { - query_output_labels: promql_utilities::data_model::KeyByLabelNames::default(), - query_result, - }; - let total_duration = start_time.elapsed(); - debug!( - "Total request processing took (modern fallback): {:.2}ms", - total_duration.as_secs_f64() * 1000.0 - ); - return match state - .adapter - .format_success_response(&execution_result) - .await - { - Ok(response) => { - annotate_data_source( - response, - StorageBackend::SketchStore.data_source_id(), - ) - .await - } - Err(status) => status.into_response(), - }; + Err(status) => status.into_response(), } + } + Err(_) => { debug!( - "Both legacy handle_query AND modern execute() returned None/Err for \ - query='{}', falling through to fallback / unsupported", + "Modern execute() returned CapabilityMiss for query='{}', \ + falling through to fallback / unsupported", parsed_request.query ); - let total_duration = start_time.elapsed(); - debug!("=== QUERY ENGINE RETURNED NONE ==="); debug!( - "Request failed after: {:.2}ms", + "Request capability-missed after: {:.2}ms", total_duration.as_secs_f64() * 1000.0 ); @@ -1683,92 +1614,51 @@ async fn process_range_query_request( parsed_request.query, parsed_request.start, parsed_request.end, parsed_request.step ); - match state.query_engine.handle_range_query_promql( - parsed_request.query.clone(), - parsed_request.start, - parsed_request.end, - parsed_request.step, - ) { - Some((query_output_labels, query_result)) => { + // B7.5 retirement — legacy `handle_range_query_promql` is gone. + // Route directly through the modern warm-tier path + // (`execute_range_promql_modern`), which classifies via the + // analyzer + ASAP-tier reducer and returns Matrix shape per the + // `/api/v1/query_range` wire-format requirement. + let start_ms = (parsed_request.start * 1000.0) as u64; + let end_ms = (parsed_request.end * 1000.0) as u64; + let step_ms = (parsed_request.step * 1000.0) as u64; + let modern_result = state + .query_engine + .execute_range_promql_modern( + &parsed_request.query, + start_ms, + end_ms, + step_ms, + ) + .await; + match modern_result { + Ok(query_result) => { let query_duration = query_start_time.elapsed(); debug!( - "Range query execution took: {:.2}ms", + "Modern range execute took: {:.2}ms", query_duration.as_secs_f64() * 1000.0 ); - let total_duration = start_time.elapsed(); debug!( "Total range query processing took: {:.2}ms", total_duration.as_secs_f64() * 1000.0 ); - - // Format range success response match state .adapter - .format_range_success_response(&query_result, &query_output_labels) + .format_range_success_response( + &query_result, + &promql_utilities::data_model::KeyByLabelNames::default(), + ) .await { Ok(response) => response.into_response(), - Err(status) => status.into_response()} - } - None => { - // Legacy `handle_range_query_promql` returned None — try - // the modern warm-tier path. Mirrors the instant-query - // fallback in `process_via_simple_engine` that PR #253 - // wired through the trait's `execute(&str)`; this site - // uses the range-aware sibling - // `execute_range_promql_modern(query, start, end, step)` - // which returns Matrix per the - // `/api/v1/query_range` wire-format requirement. - // - // Shapes that go through this fallback: anything the - // legacy path doesn't know (notably the modified-OTLP - // sketch-backed sids — count_over_time / quantile_over_time - // / etc. against CMS / CountSketch / KLL / DDSketch / HLL - // policies). Shapes still unsupported in the warm tier - // (topk_over_time — not standard PromQL anyway) fall - // through this branch too and continue to the unsupported- - // query response, which the EngineRouter can route to a - // cold-tier fallback if one is configured. - let start_ms = (parsed_request.start * 1000.0) as u64; - let end_ms = (parsed_request.end * 1000.0) as u64; - let step_ms = (parsed_request.step * 1000.0) as u64; - let modern_result = state - .query_engine - .execute_range_promql_modern( - &parsed_request.query, - start_ms, - end_ms, - step_ms, - ) - .await; - if let Ok(query_result) = modern_result { - debug!( - "Modern execute_range_promql_modern handled what legacy \ - handle_range_query_promql missed (query='{}')", - parsed_request.query - ); - let total_duration = start_time.elapsed(); - debug!( - "Total range query processing took (modern fallback): {:.2}ms", - total_duration.as_secs_f64() * 1000.0 - ); - return match state - .adapter - .format_range_success_response( - &query_result, - &promql_utilities::data_model::KeyByLabelNames::default(), - ) - .await - { - Ok(response) => response.into_response(), - Err(status) => status.into_response(), - }; + Err(status) => status.into_response(), } - + } + Err(_) => { debug!( - "Both legacy and modern range-query paths returned None/Err \ - for query='{}', falling through to unsupported", + "Modern range-query path returned CapabilityMiss for query='{}', \ + falling through to unsupported", parsed_request.query ); match state.adapter.format_unsupported_query_response().await { @@ -4703,18 +4593,22 @@ async fn handle_precompute_job( "Executing precompute job from controller" ); - match state.query_engine.handle_query_promql(req.query_expr, time) { - Some((key_by, result)) => { + // B7.5: legacy `handle_query_promql` retired; route through the + // modern `execute(&str)` trait surface. + use crate::query_engines::routing::query_engine_routing::QueryEngine; + let _ = time; + match state.query_engine.execute(&req.query_expr).await { + Ok(result) => { let body = serde_json::json!({ "status": "success", "data": { "result_type": "precompute", - "key_by": format!("{:?}", key_by), + "key_by": "{}", "result": format!("{:?}", result)} }); (StatusCode::OK, axum::Json(body)).into_response() } - None => { + Err(_) => { // Query not answerable by sketches — return 404 with hint let body = serde_json::json!({ "status": "error", diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index 726f9741..cfb8be50 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1,180 +1,17 @@ -use crate::storage_engines::types::{AggregationIdInfo, KeyByLabelValues, StreamingConfig}; -use crate::query_engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; -// use crate::storage_engines::promsketch_store::{ -// self, is_usampling_function, metrics as ps_metrics, PromSketchStore, -// }; -use crate::storage_engines::TimestampedBucketsMap; -use core::panic; -use promql_utilities::get_is_collapsable; -use promql_utilities::query_logics::enums::{AggregationOperator, AggregationType, PromQLFunction}; -use serde_json::Value; -use std::collections::HashMap; +use crate::storage_engines::types::StreamingConfig; use std::sync::Arc; -use std::time::Instant; -use tracing::{debug, warn}; - -use crate::AggregateCore; -use asap_types::enums::WindowType; use asap_types::query_requirements::QueryRequirements; -use asap_types::utils::normalize_spatial_filter; -use promql_utilities::ast_matching::{PromQLMatchResult, PromQLPattern, PromQLPatternBuilder}; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::{QueryPatternType, Statistic}; -use promql_utilities::query_logics::parsing::{ - get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute}; - -// SQL issue: refactor simpleengine to create matchresult similar to SQLquerydata - -// Type alias for merged outputs (single aggregate per key after merging) -type MergedOutputsMap = HashMap, Box>; - -/// Phase 5 helper — extract `(metric_name, label_matcher_key_set)` from a -/// PromQL query for ASAP-tier candidate selection. Walks the AST to find -/// the first `VectorSelector` / `MatrixSelector`, returns its metric name -/// (drawn either from `vs.name` or from a `__name__=...` matcher) and -/// the user-specified label-matcher KEYS (excluding the synthetic -/// `__name__`). Returns `None` for queries that don't reference a -/// concrete metric. -/// -/// Intentionally lightweight: callers use the result to filter ASAP-tier -/// candidates via `SketchStore::instances_matching`. Any over-approximation -/// is tolerable — the candidates are subsequently classified, and on -/// `Ghost` / `Unknown` outcomes the query falls through to the archive -/// engine via the EngineRouter's `CapabilityMiss` failover. -/// Legacy `(metric_name, group_by_keys)` extractor — superseded by -/// `control_plane::asap_tier_analysis::analyze_promql_for_asap_tier`, -/// which returns the full `ASAPTierAnalysis` (capability, function -/// name + args, range). Kept around as `#[allow(dead_code)]` because -/// downstream code (range-query pipeline, range-step planner) still -/// uses bare `(metric, keys)` projections for sid candidate filtering; -/// once those callers also migrate to `ASAPTierAnalysis`, this can be -/// deleted in a follow-up. -#[allow(dead_code)] -fn extract_metric_and_label_keys( - query: &str, -) -> Option<(String, std::collections::BTreeSet)> { - use promql_parser::parser::Expr; - let ast = promql_parser::parser::parse(query).ok()?; - - fn walk(expr: &Expr) -> Option<(String, std::collections::BTreeSet)> { - match expr { - Expr::VectorSelector(vs) => { - let mut keys = std::collections::BTreeSet::new(); - let mut metric = vs.name.clone().unwrap_or_default(); - for m in &vs.matchers.matchers { - if m.name == "__name__" { - if metric.is_empty() { - metric = m.value.clone(); - } - continue; - } - keys.insert(m.name.clone()); - } - if metric.is_empty() { - None - } else { - Some((metric, keys)) - } - } - Expr::MatrixSelector(ms) => walk(&Expr::VectorSelector(ms.vs.clone())), - Expr::Call(call) => call.args.args.iter().find_map(|a| walk(a)), - Expr::Aggregate(agg) => walk(&agg.expr), - Expr::Binary(bin) => walk(&bin.lhs).or_else(|| walk(&bin.rhs)), - Expr::Subquery(sq) => walk(&sq.expr), - Expr::Paren(p) => walk(&p.expr), - Expr::Unary(u) => walk(&u.expr), - _ => None} - } - - walk(&ast) -} - -/// Metadata extracted from a query, independent of query language -#[derive(Debug, Clone)] -pub struct QueryMetadata { - /// Labels that will appear in the query output - pub query_output_labels: KeyByLabelNames, - /// The primary statistic to compute (sum, max, quantile, etc.) - pub statistic_to_compute: Statistic, - /// Additional parameters (e.g., "quantile" -> "0.95", "k" -> "10") - pub query_kwargs: HashMap} - -/// Parameters for a single store query -#[derive(Debug, Clone)] -pub struct StoreQueryParams { - pub metric: String, - pub aggregation_id: u64, - pub start_timestamp: u64, - pub end_timestamp: u64, - /// true for sliding windows (exact match), false for tumbling (range) - pub is_exact_query: bool} - -/// Complete plan for querying store (values + optional separate keys) -#[derive(Debug, Clone)] -pub struct StoreQueryPlan { - pub values_query: StoreQueryParams, - /// Some when key and value use different aggregations (DeltaSet/SetAggregator) - pub keys_query: Option} - -/// Timestamps for query execution -#[derive(Debug, Clone)] -pub struct QueryTimestamps { - pub start_timestamp: u64, - pub end_timestamp: u64} - -/// Complete execution context for a query -#[derive(Debug, Clone)] -pub struct QueryExecutionContext { - pub metric: String, - pub metadata: QueryMetadata, - pub store_plan: StoreQueryPlan, - pub agg_info: AggregationIdInfo, - /// Whether to merge multiple precomputes (true for temporal queries) - pub do_merge: bool, - #[allow(dead_code)] - pub spatial_filter: String, - pub query_time: u64, - /// Spatial grouping labels from the value aggregation config. - /// These are the store GROUP BY columns. - pub grouping_labels: KeyByLabelNames, - /// Aggregated labels from the value aggregation config. - /// These are labels that "key" an accumulator/sketch internally - /// (e.g. endpoint within a MultipleIncrease accumulator). - pub aggregated_labels: KeyByLabelNames} - -/// Parameters for a range query -#[derive(Debug, Clone)] -pub struct RangeQueryParams { - pub start: u64, // start timestamp in ms - pub end: u64, // end timestamp in ms - pub step: u64, // step in ms -} - -/// Extended execution context for range queries -#[derive(Debug, Clone)] -pub struct RangeQueryExecutionContext { - /// Base context (metric, metadata, store_plan, etc.) - pub base: QueryExecutionContext, - /// Range-specific parameters - pub range_params: RangeQueryParams, - /// Number of buckets per step (step / tumbling_window) - pub buckets_per_step: usize, - /// Number of buckets in lookback window - pub lookback_bucket_count: usize, - /// Tumbling window size in ms - pub tumbling_window_ms: u64} -// /// Parsed components of a sketch query, extracted either via the PromQL AST -// /// parser (for standard functions) or via regex (for custom functions like -// /// `entropy_over_time` that the promql-parser crate doesn't recognize). -// struct SketchQueryComponents { -// func_name: String, -// metric: String, -// range_seconds: u64, -// /// Extra numeric argument (e.g. quantile value). 0.0 when unused. -// args: f64, -// } +#[cfg(test)] +use crate::storage_engines::types::KeyByLabelValues; +#[cfg(test)] +use crate::AggregateCore; +#[cfg(test)] +use promql_utilities::query_logics::enums::Statistic; +#[cfg(test)] +use std::collections::HashMap; /// Simple query engine for processing PromQL-like queries against precomputed data pub struct ASAPQueryEngine { @@ -191,8 +28,8 @@ pub struct ASAPQueryEngine { /// to both `ASAPQueryEngine` and `HttpServer::with_hot_reload_config`, /// a POST is immediately visible to the next query. streaming_config_source: crate::storage_engines::types::HotReloadStreamingConfig, + #[allow(dead_code)] prometheus_scrape_interval: u64, - control_plane_patterns: HashMap>, /// Optional `ControlPlaneClient` used to notify the control plane /// when a query hits a capability miss /// (`find_compatible_aggregation` returns `None`). When `None`, @@ -243,142 +80,9 @@ impl ASAPQueryEngine { streaming_config_source: crate::storage_engines::types::HotReloadStreamingConfig, prometheus_scrape_interval: u64, ) -> Self { - // Create temporal pattern blocks - let mut temporal_pattern_blocks = HashMap::new(); - temporal_pattern_blocks.insert( - "quantile".to_string(), - PromQLPatternBuilder::function( - vec![PromQLFunction::QuantileOverTime.as_str()], - vec![ - PromQLPatternBuilder::number(None, Some("quantile_param")), - PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ), - ], - Some("function"), - Some("function_args"), - ), - ); - - temporal_pattern_blocks.insert( - "generic".to_string(), - PromQLPatternBuilder::function( - vec![ - "sum_over_time", - "count_over_time", - "avg_over_time", - "min_over_time", - "max_over_time", - "increase", - "rate", - "entropy_over_time", - "distinct_over_time", - "l1_over_time", - "l2_over_time", - "stddev_over_time", - "stdvar_over_time", - "sum2_over_time", - ], - vec![PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - )], - Some("function"), - Some("function_args"), - ), - ); - - // Create spatial pattern blocks - let mut spatial_pattern_blocks = HashMap::new(); - let spatial_ops_all: Vec<&str> = [ - AggregationOperator::Sum, - AggregationOperator::Count, - AggregationOperator::Avg, - AggregationOperator::Quantile, - AggregationOperator::Min, - AggregationOperator::Max, - AggregationOperator::Topk, - ] - .map(AggregationOperator::as_str) - .to_vec(); - let spatial_ops_no_topk: Vec<&str> = [ - AggregationOperator::Sum, - AggregationOperator::Count, - AggregationOperator::Avg, - AggregationOperator::Quantile, - AggregationOperator::Min, - AggregationOperator::Max, - ] - .map(AggregationOperator::as_str) - .to_vec(); - spatial_pattern_blocks.insert( - "generic".to_string(), - PromQLPatternBuilder::aggregation( - spatial_ops_all, - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - None, - None, - Some("aggregation"), - ), - ); - - // Helper functions (these would be closures or separate methods) - fn temporal_pattern( - pattern_type: &str, - blocks: &HashMap>>, - ) -> PromQLPattern { - PromQLPattern::new(blocks[pattern_type].clone()) - } - - fn spatial_pattern( - pattern_type: &str, - blocks: &HashMap>>, - ) -> PromQLPattern { - PromQLPattern::new(blocks[pattern_type].clone()) - } - - let spatial_of_temporal_pattern = - |temporal_block: &Option>| -> PromQLPattern { - let pattern = PromQLPatternBuilder::aggregation( - spatial_ops_no_topk.clone(), - temporal_block.clone(), - None, - None, - None, - Some("aggregation"), - ); - PromQLPattern::new(pattern) - }; - - // Create control plane patterns - let mut control_plane_patterns = HashMap::new(); - control_plane_patterns.insert( - QueryPatternType::OnlyTemporal, - vec![ - temporal_pattern("quantile", &temporal_pattern_blocks), - temporal_pattern("generic", &temporal_pattern_blocks), - ], - ); - control_plane_patterns.insert( - QueryPatternType::OnlySpatial, - vec![spatial_pattern("generic", &spatial_pattern_blocks)], - ); - control_plane_patterns.insert( - QueryPatternType::OneTemporalOneSpatial, - vec![ - spatial_of_temporal_pattern(&temporal_pattern_blocks["quantile"]), - spatial_of_temporal_pattern(&temporal_pattern_blocks["generic"]), - ], - ); - Self { streaming_config_source, prometheus_scrape_interval, - control_plane_patterns, control_plane_client: None, sketch_index: None, archive_engine: None} @@ -466,4123 +170,807 @@ impl ASAPQueryEngine { ) } - /// Look up a compatible aggregation for the given requirements, - /// and if none exists, fire a capability-miss notification to - /// the control plane (fire-and-forget, does not block the query). - /// Wraps the plain `streaming_config.find_compatible_aggregation` - /// with the PR G telemetry call-out. - fn find_compatible_aggregation_with_miss_notify( - &self, - requirements: &QueryRequirements, - ) -> Option { - let streaming_config = self.streaming_config_snapshot(); - let result = streaming_config.find_compatible_aggregation(requirements); - if result.is_none() { - crate::drivers::control_plane_client::spawn_capability_miss_notify( - &self.control_plane_client, - requirements, - ); + /// Build a `QueryRequirements` from an ASAP-tier candidate (the + /// shape modern `execute()` works with). Used by the modern miss + /// branches to fire the capability-miss notify the legacy + /// `find_compatible_aggregation_with_miss_notify` path provided + /// before the B7.5 retirement. + fn requirements_from_candidate( + candidate: &control_plane::asap_tier_analysis::ASAPTierCandidate, + ) -> QueryRequirements { + QueryRequirements { + metric: candidate.metric_name.clone(), + statistics: Vec::new(), + data_range_ms: if candidate.range_seconds > 0 { + Some(candidate.range_seconds.saturating_mul(1000)) + } else { + None + }, + grouping_labels: KeyByLabelNames::new( + candidate.group_by_keys.iter().cloned().collect(), + ), + spatial_filter_normalized: candidate.spatial_filter_canonical.clone(), } - result } - /// Resolve the canonical "all labels" set for a metric, with a - /// streaming-config fallback for schema-empty deploys. - /// - /// The user-facing `inference_config.schema` is the source of truth - /// for "what labels does this metric carry" — but the production - /// ASAP-tier deploy launches with `--streaming-config` only and no - /// `--config`, so the schema is empty. Pre-fix every query lookup - /// in `build_promql_execution_context_tail` and - /// `build_query_requirements_promql` returned `None` / - /// `KeyByLabelNames::empty()` for that metric, killing capability - /// matching (`req.grouping_labels = []` strict-mismatches every - /// agg config's `[zone]`) and the downstream context build (the - /// `None` short-circuits the whole query). See - /// `tests::warm_engine_replay_regression_tests::production_conditions_*`. - /// - /// Fallback rules: - /// 1. Look up the metric in `inference_config.schema`. Return its - /// labels if present. - /// 2. Otherwise scan the current `StreamingConfig` snapshot for - /// every agg config whose `metric == name`. Union their - /// `grouping_labels` (the per-series partition the agg - /// materialises) and return that. The union preserves order - /// of first-appearance and de-dupes — `KeyByLabelNames` - /// equality is strict, so we have to keep insertion order - /// deterministic across config swaps. - /// 3. Returns `None` only if no agg config references the metric - /// AND the schema is empty. Callers translate that into the - /// same "metric unknown" outcome as before this helper landed. - fn resolve_metric_labels(&self, metric: &str) -> Option { - // Streaming-config-derived label set. Previously this had a - // fast-path through `inference_config.schema`; that source - // was retired with InferenceConfig. The control plane drives - // capability matching against `aggregation_configs` directly, - // so we derive the label union from those configs. - // - // Returns `Some(union)` when at least one aggregation references - // the metric (even if its grouping_labels are empty — - // un-grouped aggregations are valid), and `None` only when no - // aggregation in the current snapshot references the metric. - let snap = self.streaming_config_snapshot(); - let mut seen = std::collections::HashSet::new(); - let mut union: Vec = Vec::new(); - let mut any_agg_references_metric = false; - let mut agg_ids: Vec = snap.aggregation_configs.keys().copied().collect(); - agg_ids.sort_unstable(); - for id in agg_ids { - let cfg = match snap.get_aggregation_config(id) { - Some(c) => c, - None => continue, - }; - if cfg.metric != metric { - continue; - } - any_agg_references_metric = true; - for label in &cfg.grouping_labels.labels { - if seen.insert(label.clone()) { - union.push(label.clone()); + /// Build a minimal `QueryRequirements` from a bare PromQL string — + /// used by the no-sketch-index miss branch in modern `execute()`, + /// where we don't have a parsed candidate (analysis was skipped) + /// but still want to fire the capability-miss notify so the + /// control-plane feedback loop closes. Lifts (metric_name, + /// group_by_keys) from the AST via a light walker; returns `None` + /// for queries that don't reference a concrete metric. + fn requirements_from_query_str(query: &str) -> Option { + use promql_parser::parser::Expr; + let ast = promql_parser::parser::parse(query).ok()?; + fn walk(expr: &Expr) -> Option<(String, std::collections::BTreeSet)> { + match expr { + Expr::VectorSelector(vs) => { + let mut keys = std::collections::BTreeSet::new(); + let mut metric = vs.name.clone().unwrap_or_default(); + for m in &vs.matchers.matchers { + if m.name == "__name__" { + if metric.is_empty() { + metric = m.value.clone(); + } + continue; + } + keys.insert(m.name.clone()); + } + if metric.is_empty() { + None + } else { + Some((metric, keys)) + } } - } - } - if !any_agg_references_metric { - None - } else { - Some(KeyByLabelNames::new(union)) - } - } - - /// Convert query timestamp (seconds) to data timestamp (milliseconds) - pub fn convert_query_time_to_data_time(query_time: f64) -> u64 { - (query_time * 1000.0) as u64 - } - - /// Finds the query configuration for a SQL query using structural pattern matching. - /// - /// Unlike `find_query_config` (which does exact string comparison), this method parses - /// each template in query_configs and compares it structurally against the incoming - /// query_data — ignoring absolute timestamps and comparing only metric, aggregation, - /// labels, time column name, and duration. - /// Validates and potentially aligns end timestamp based on query pattern - fn validate_and_align_end_timestamp( - &self, - mut end_timestamp: u64, - query_pattern_type: QueryPatternType, - ) -> u64 { - let interval_ms = self.prometheus_scrape_interval * 1000; - - if !end_timestamp.is_multiple_of(interval_ms) { - warn!( - "Query end timestamp {} is not aligned with Prometheus scrape interval of {} seconds. \ - This may lead to inaccurate results.", - end_timestamp, self.prometheus_scrape_interval - ); - } - - // For OnlySpatial, align end_timestamp to nearest scrape interval - if query_pattern_type == QueryPatternType::OnlySpatial - && !end_timestamp.is_multiple_of(interval_ms) - { - let aligned_end_timestamp = (end_timestamp / interval_ms) * interval_ms; - debug!( - "OnlySpatial query: Aligning end_timestamp from {} to {} using scrape interval of {} seconds", - end_timestamp, aligned_end_timestamp, self.prometheus_scrape_interval - ); - end_timestamp = aligned_end_timestamp; - } - - end_timestamp + Expr::MatrixSelector(ms) => walk(&Expr::VectorSelector(ms.vs.clone())), + Expr::Call(call) => call.args.args.iter().find_map(|a| walk(a)), + Expr::Aggregate(agg) => walk(&agg.expr), + Expr::Binary(bin) => walk(&bin.lhs).or_else(|| walk(&bin.rhs)), + Expr::Subquery(sq) => walk(&sq.expr), + Expr::Paren(p) => walk(&p.expr), + Expr::Unary(u) => walk(&u.expr), + _ => None, + } + } + let (metric, keys) = walk(&ast)?; + Some(QueryRequirements { + metric, + statistics: Vec::new(), + data_range_ms: None, + grouping_labels: KeyByLabelNames::new(keys.into_iter().collect()), + spatial_filter_normalized: String::new(), + }) } - /// Calculates start timestamp for PromQL queries - fn calculate_start_timestamp_promql( + #[cfg(test)] + fn query_precompute_for_statistic( &self, - end_timestamp: u64, - query_pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, - ) -> u64 { - match query_pattern_type { - QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => { - let range_seconds = match_result.get_range_duration().unwrap().num_seconds() as u64; - end_timestamp - (range_seconds * 1000) - } - QueryPatternType::OnlySpatial => { - end_timestamp.saturating_sub(self.prometheus_scrape_interval * 1000) + precompute: &dyn AggregateCore, + statistic: &Statistic, + key: &Option, + query_kwargs: &HashMap, + ) -> Result> { + // Phase 1b of the sketch DB design + // (docs/design-sketch-db.md §5.1 / §16 Phase 1): + // for single-subpopulation queries on additive statistics + // (Count / Sum / Min / Max), serve from the typed aux + // columns without deserialising the sketch payload. + // + // Keyed queries (`key.is_some()`) still need the full + // `query_statistic` path — aux is per-accumulator, not + // per-subpopulation key. + // + // `try_answer` returns `None` when the statistic isn't + // covered by aux (Quantile / Cardinality / TopK / Increase / + // Rate) or when the accumulator doesn't track the requested + // aux field; both cases fall through to the existing path + // so the query result is semantically identical. + if key.is_none() { + if let Some(value) = precompute.aux_stats().try_answer(*statistic) { + return Ok(value); } } + precompute.query_statistic(*statistic, key, query_kwargs) } - /// Calculates start timestamp for SQL queries - /// Calculates and validates query timestamps for PromQL - fn calculate_query_timestamps_promql( + /// Modern warm-tier path for `/api/v1/query_range` — the range- + /// query equivalent of the `QueryEngine::execute(&str)` trait + /// surface. Used by the HTTP server as a fallback when the legacy + /// `handle_range_query_promql` returns `None`. + /// + /// Time semantics follow Prometheus's + /// `/api/v1/query_range?start&end&step` spec: the result is a + /// `matrix` (one row per series, each row carrying multiple + /// (timestamp, value) samples). The warm-tier reducer naturally + /// produces one sample per window_close in `[start, end]`, so + /// the matrix is sampled at the underlying aggregation's window + /// boundaries — typically a finer grid than the user's `step` + /// when window_size < step. (The Prometheus spec says + /// evaluate at each step `t = start, start+step, …, end`; the + /// warm tier returns at native window-close granularity instead. + /// This is more data, not less — clients that expect exact step + /// timestamps can downsample, or route step-precise queries to + /// the cold tier via the EngineRouter.) + /// + /// `step` is currently accepted for API compatibility but unused + /// — see the granularity-mismatch note above. + pub async fn execute_range_promql_modern( &self, - query_time: u64, - query_pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, - ) -> QueryTimestamps { - let mut end_timestamp = if let Some(at_modifier) = match_result - .tokens - .get("metric") - .and_then(|t| t.metric.as_ref()) - .and_then(|m| m.at_modifier) - { - at_modifier * 1000 - } else { - query_time + query: &str, + start_ms: u64, + end_ms: u64, + _step_ms: u64, + ) -> Result< + crate::query_engines::query_result::QueryResult, + crate::query_engines::EngineError, + > { + let Some(idx) = self.sketch_index.as_ref() else { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!("ASAPQueryEngine: no sketch index for `{query}` — failing over"), + )); }; - end_timestamp = self.validate_and_align_end_timestamp(end_timestamp, query_pattern_type); - let start_timestamp = - self.calculate_start_timestamp_promql(end_timestamp, query_pattern_type, match_result); - - QueryTimestamps { - start_timestamp, - end_timestamp} - } - - /// Extracts quantile parameter from PromQL match result - fn extract_quantile_param_promql( - &self, - query_pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, - ) -> Option { - let quantile_value = match query_pattern_type { - QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => { - match_result - .tokens - .get("function_args") - .and_then(|token| token.function.as_ref()) - .and_then(|func| func.args.first()) - } - QueryPatternType::OnlySpatial => match_result - .tokens - .get("aggregation") - .and_then(|token| token.aggregation.as_ref()) - .and_then(|agg| agg.param.as_ref())}; - - quantile_value.map(|s| s.to_string()) - } + let analysis = + control_plane::asap_tier_analysis::analyze_promql_for_asap_tier(query); - /// Extracts topk k parameter from PromQL match result - fn extract_topk_param( - &self, - query_pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, - ) -> Result { - match query_pattern_type { - QueryPatternType::OnlySpatial => match_result - .tokens - .get("aggregation") - .and_then(|token| token.aggregation.as_ref()) - .and_then(|agg| agg.param.as_ref()) - .map(|s| s.to_string()) - .ok_or_else(|| "Missing k parameter for top-k query".to_string()), - _ => Err(format!( - "Top-k statistic is only supported for OnlySpatial pattern, found {:?}", - query_pattern_type - ))} - } + if let Some(reason) = &analysis.unsupported { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore analyzer rejected `{query}` for range query: \ + {reason:?} — failing over to archive" + ), + )); + } + if analysis.candidates.is_empty() { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore analyzer produced no ASAP-tier candidates for \ + `{query}` — failing over to archive" + ), + )); + } - /// Builds query kwargs (quantile, k, etc.) for PromQL queries - fn build_query_kwargs_promql( - &self, - statistic: &Statistic, - query_pattern_type: QueryPatternType, - match_result: &PromQLMatchResult, - ) -> Result, String> { - let mut query_kwargs = HashMap::new(); + let streaming_snap = self.streaming_config_snapshot(); + let policy_registry = streaming_snap.policy_registry(); + let reducer = crate::storage_engines::sketch_db::query::SketchReducer::new(idx); + let mut combined_result: Option< + crate::storage_engines::sketch_db::query::ASAPTierResult, + > = None; - match statistic { - Statistic::Quantile => { - let quantile = self - .extract_quantile_param_promql(query_pattern_type, match_result) - .ok_or_else(|| "Missing quantile parameter for quantile query".to_string())?; - debug!("Extracted quantile value: {:?}", quantile); - query_kwargs.insert("quantile".to_string(), quantile); + for candidate in &analysis.candidates { + // Resolve candidate → {sids} via the sid catalog. Schema- + // retirement #5: prefer `instances_matching` over the + // policy-fp reverse index — it's the more general + // primitive and works whether or not the ingest path was + // able to bind the sid back to a streaming-config policy. + // + // History: an earlier PR removed an `instances_matching` + // fallback under the assumption every production sid + // registration would populate `policy_fp`. The MVP smoke + // test (issue #271 / tracking #272) showed that + // assumption is wrong — sketches arriving from the agent + // carry the full wire-attr set rather than the streaming- + // config's `grouping_labels` subset, so + // `derive_sketch_policy_fp` returns `UNSET` and + // `sids_for_policy(fp)` returns empty. The agg_id-aware + // path is preserved for ExactAgg sids minted via + // `ingest_precompute_for_agg_config` (those carry a + // populated `policy_fp`) but its result is unioned with + // the catalog-walk result so we don't miss the sketches. + let policy_fps = control_plane::asap_tier_analysis::find_matching_policies( + &policy_registry, + candidate, + ); + let mut sids: std::collections::BTreeSet = + std::collections::BTreeSet::new(); + for fp in &policy_fps { + sids.extend(idx.sids_for_policy(*fp)); } - Statistic::Topk => { - let k = self.extract_topk_param(query_pattern_type, match_result)?; - debug!("Extracted k value: {:?}", k); - query_kwargs.insert("k".to_string(), k); + sids.extend(idx.instances_matching( + &candidate.metric_name, + &candidate.group_by_keys, + )); + if sids.is_empty() { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore has no policy for metric `{}` satisfying \ + capability {:?} — failing over to archive", + candidate.metric_name, candidate.required_capability, + ), + )); } - // PR #111 honest-gap closure for `rate(...)` over a - // CountMinSketch-backed agg: the CMS accumulator records - // event counts but not per-event timestamps, so it can't - // derive the range duration locally. The engine knows - // the range from the matrix selector and pushes it down - // here so `CountMinSketchAccumulator::query_statistic` - // can divide events by seconds. Increase carries the - // same divisor (it falls back to raw count when - // range_ms is absent). - Statistic::Rate | Statistic::Increase => { - if let Some(d) = match_result.get_range_duration() { - let range_ms = (d.num_seconds() as u64) * 1000; - if range_ms > 0 { - query_kwargs.insert("range_ms".to_string(), range_ms.to_string()); - debug!( - "Rate/Increase query: pushed range_ms={} into kwargs \ - for CMS-style accumulators", - range_ms - ); + + let required: crate::storage_engines::sketch_db::index::Capability = + candidate.required_capability.clone(); + let mut hit_sids: Vec = Vec::with_capacity(sids.len()); + for sid in &sids { + let meta = match idx.instance(*sid) { + Some(m) => m, + None => continue, + }; + if let Some(cap) = meta.capability.as_ref() { + if required.is_satisfied_by(cap) { + hit_sids.push(*sid); } } } - _ => {} - } - - Ok(query_kwargs) - } - - /// Builds query kwargs for SQL queries - /// Creates query parameters for separate keys query - fn create_keys_query_params( - &self, - metric: &str, - end_timestamp: u64, - agg_info: &AggregationIdInfo, - ) -> Result { - // The historical key-tracking family (`SetAggregator` / - // `DeltaSetAggregator`) has been retired. After retirement no - // production code path produces a key-aggregation distinct - // from the value-aggregation, so this function is unreachable - // in practice — the caller's `aggregation_id_for_key != - // aggregation_id_for_value` guard never fires. Kept as a - // typed-error surface in case a stale stored config still - // carries a mismatched pair. - let _ = end_timestamp; - return Err(format!( - "create_keys_query_params is unreachable after SetAggregator / \ - DeltaSetAggregator retirement; got aggregation_type_for_key={:?}", - agg_info.aggregation_type_for_key - )); - #[allow(unreachable_code)] - let (start_timestamp, end_timestamp) = (0u64, end_timestamp); - - Ok(StoreQueryParams { - metric: metric.to_string(), - aggregation_id: agg_info.aggregation_id_for_key, - start_timestamp, - end_timestamp, - is_exact_query: false, // Keys always use range queries - }) - } + if hit_sids.is_empty() { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore has no sid satisfying capability {:?} for \ + metric `{}` — failing over to archive", + candidate.required_capability, candidate.metric_name + ), + )); + } - /// Creates a plan for querying the store based on aggregation configuration - fn create_store_query_plan( - &self, - metric: &str, - timestamps: &QueryTimestamps, - agg_info: &AggregationIdInfo, - ) -> Result { - // Bind a single snapshot of the streaming config for this - // helper's entire execution. `aggregation_config_for_value` - // is a borrow that outlives the initial expression, so the - // Arc it borrows from must outlive this scope. - let streaming_config = self.streaming_config_snapshot(); - // Get aggregation config for value to determine window type - let aggregation_config_for_value = streaming_config - .get_aggregation_config(agg_info.aggregation_id_for_value) - .ok_or_else(|| { - format!( - "Aggregation config not found for aggregation_id: {}", - agg_info.aggregation_id_for_value + let result = reducer + .evaluate( + &hit_sids, + &candidate.function, + &candidate.function_args, + start_ms, + end_ms, ) - })?; - - let window_type = aggregation_config_for_value.window_type; - let is_exact_query = window_type == WindowType::Sliding; - - // Determine start/end for values query based on window type - let (values_start, values_end) = if is_exact_query { - // Sliding window: exact window match - let exact_start = - timestamps.end_timestamp - (aggregation_config_for_value.window_size * 1000); - (exact_start, timestamps.end_timestamp) - } else { - // Tumbling window: range query - (timestamps.start_timestamp, timestamps.end_timestamp) - }; - - let values_query = StoreQueryParams { - metric: metric.to_string(), - aggregation_id: agg_info.aggregation_id_for_value, - start_timestamp: values_start, - end_timestamp: values_end, - is_exact_query}; - - // Determine if we need a separate keys query - let keys_query = if agg_info.aggregation_id_for_key != agg_info.aggregation_id_for_value { - Some(self.create_keys_query_params(metric, timestamps.end_timestamp, agg_info)?) - } else { - None - }; - - Ok(StoreQueryPlan { - values_query, - keys_query}) - } - - /// Executes a single store query based on parameters - fn execute_store_query( - &self, - params: &StoreQueryParams, - ) -> Result { - debug!( - "Querying store: metric={}, agg_id={}, range=[{}, {}], exact={}", - params.metric, - params.aggregation_id, - params.start_timestamp, - params.end_timestamp, - params.is_exact_query - ); - - // M2.3.6f — engine reads precomputes from SketchStore only. - // The legacy `Store::query_precomputed_output*` fallback has - // been retired now that DualWriteSink (M2.3.4b) → SketchStoreSink - // (M2.3.6a) writes exclusively to SketchStore and - // BackfillService (M2.3.6e) mirrors replays there too. - // - // Tests that don't attach a SketchStore now get `Ok(empty)` - // here. Anything deeper than smoke-test coverage was already - // setting one (M2.3.5b made it mandatory in production). - // - // ── KNOWN GAP — sketch-backed aggs return empty here ───────── - // - // `query_precomputes_by_agg` (called below) filters its - // candidate-sid scan with `matches!(&m.agg_kind, - // AggKind::ExactAgg { agg_type: t, .. } if *t == agg_type)`. - // It NEVER matches `AggKind::Sketch` — so for any sketch- - // backed aggregation (DDSketch / KLL / HLL / CountSketch / - // CountMinSketch arriving via OTLP and registered by - // `route_modified_otlp_sketches_to_precompute` with - // `AggKind::Sketch { kind, config, .. }`), this lookup - // returns an empty map. We then bubble up "No precomputed - // outputs found for metric: X, aggregation_id: Y" and - // `handle_query` returns `None`, which the HTTP layer - // renders as `errorType: bad_data` / `error: "No result - // for query"`. - // - // Diagnosed in the e2e test arc (#247 → #248 → #249 → - // #250 → engine-path debug session 2026-05). Sketches DO - // reach `SketchStore` — `runtime_info.earliest_timestamp_per_sid` - // shows them — but they're only readable via the sid-keyed - // `SketchStore::query_range(sid, ...)` path (sketch payloads - // filtered by `payload.as_sketch()`), not via the agg-keyed - // precomputes path consumed here. - // - // **The fix** is to dispatch by the agg's capability: - // * sketch-typed aggs route to a sketch-side query (a - // counterpart to `query_precomputes_by_agg` that scans - // `AggKind::Sketch` sids and assembles per-sketch results - // into the engine's expected `Box` - // shape) — non-trivial since the existing pipeline expects - // precompute payload shapes. - // * OR route the legacy `handle_query` path through the - // newer `ASAPQueryEngine::execute(&str)` trait path - // (around line 3430), which already does - // `idx.sids_for_policy(fp)` + reducer dispatch and - // handles sketches natively via `SketchReducer::evaluate`. - let Some(idx) = self.sketch_index.as_ref() else { - return Ok(TimestampedBucketsMap::new()); - }; - let cfg = self.streaming_config_snapshot(); - let Some(agg_cfg) = cfg.get_aggregation_config(params.aggregation_id) else { - return Ok(TimestampedBucketsMap::new()); - }; - let raw = idx.query_precomputes_by_agg( - ¶ms.metric, - agg_cfg.aggregation_type, - params.start_timestamp, - params.end_timestamp, - ); - let result: TimestampedBucketsMap = if params.is_exact_query { - // Sliding-window mode requires bit-exact (start, end) - // match. SketchStore's range query returns any windows - // fully within [start, end] — filter post-hoc to recover - // the exact semantics the retired - // `query_precomputed_output_exact` had. - raw.into_iter() - .map(|(k, v)| { - let filtered: Vec<_> = v - .into_iter() - .filter(|((s, e), _)| { - *s == params.start_timestamp - && *e == params.end_timestamp - }) - .collect(); - (k, filtered) - }) - .filter(|(_, v)| !v.is_empty()) - .collect() - } else { - raw - }; - Ok(result) - } - - /// Executes the full store query plan and returns merged results - fn execute_and_merge_store_queries( - &self, - plan: &StoreQueryPlan, - do_merge: bool, - agg_info: &AggregationIdInfo, - ) -> Result< - ( - MergedOutputsMap, - Option, - Option<(u64, u64)>, - ), - String, - > { - // Query and merge values - let values_map = self.execute_store_query(&plan.values_query).map_err(|e| { - warn!("Error querying store for values: {}", e); - e - })?; - - if values_map.is_empty() { - return Err(format!( - "No precomputed outputs found for metric: {}, aggregation_id: {}", - plan.values_query.metric, plan.values_query.aggregation_id - )); + .map_err(|e| { + crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer failed for `{query}` over \ + [{start_ms}, {end_ms}]: {e:?} — failing over to archive" + ), + ) + })?; + combined_result = Some(result); } - debug!("Store query returned {} unique keys", values_map.len()); - - let merge_start_time = Instant::now(); - let window_type = if plan.values_query.is_exact_query { - WindowType::Sliding - } else { - WindowType::Tumbling - }; - - // Pick the single CLOSEST precompute window across all keys — - // the latest pane (max tr.1, tie-break on max tr.0) that - // overlaps the request range. The store's overlap filter may - // have returned multiple tumbling panes that straddle the - // request, but a window query - // (e.g. `quantile_over_time(...[1m])`) should answer with - // *one* concrete window so the caller can see exactly which - // pane produced the value (annotated downstream as - // `precompute_window`). Keys whose data didn't land in that - // chosen window are dropped from the result rather than - // contributing a stale answer from an older pane. - let chosen_window: Option<(u64, u64)> = values_map - .values() - .flat_map(|buckets| buckets.iter().map(|(tr, _)| *tr)) - .max_by_key(|tr| (tr.1, tr.0)); - - let merged_values: MergedOutputsMap = if plan.values_query.is_exact_query { - // Sliding window: no merge needed, extract buckets from timestamped data - debug!("Sliding window mode: Skipping merge (expecting 1 precompute per key)"); - values_map - .into_iter() - .map(|(key, timestamped_buckets)| { - if timestamped_buckets.len() != 1 { - warn!( - "Sliding window expected 1 precompute per key, found {}. Using first.", - timestamped_buckets.len() - ); - } - // Extract bucket from timestamped tuple - let (_, bucket) = timestamped_buckets.into_iter().next().unwrap(); - (key, bucket.as_ref().clone_boxed_core()) - }) - .collect() - } else { - // Tumbling window: keep only the chosen-window bucket per - // key, then run through the existing merge code (which is - // a no-op for a single bucket but preserves whatever - // accumulator-side cleanup the merge path does). - let target = chosen_window.expect( - "values_map non-empty (checked above) but chosen_window was None — \ - invariant: if buckets exist, max_by_key returns Some", - ); - let filtered: TimestampedBucketsMap = values_map - .into_iter() - .filter_map(|(key, buckets)| { - let kept: Vec<_> = buckets - .into_iter() - .filter(|(tr, _)| *tr == target) - .collect(); - if kept.is_empty() { - None - } else { - Some((key, kept)) - } - }) - .collect(); - debug!( - "Tumbling window mode: closest pane [{}, {}); {} keys present in that pane", - target.0, - target.1, - filtered.len() - ); - self.merge_precomputed_outputs(&filtered, do_merge, agg_info.aggregation_type_for_value) - }; - - let merge_duration = merge_start_time.elapsed(); - debug!( - "[LATENCY] Precomputed output processing ({}): {:.2}ms, resulted in {} merged outputs", - if window_type == WindowType::Sliding { - "no merge" - } else { - "merge" - }, - merge_duration.as_secs_f64() * 1000.0, - merged_values.len() - ); - - // Query and merge keys if needed - let merged_keys = if let Some(keys_params) = &plan.keys_query { - let keys_store_query_start_time = Instant::now(); - let keys_map = self.execute_store_query(keys_params).map_err(|e| { - warn!("Error querying store for keys: {}", e); - e - })?; - debug!( - "[LATENCY] Keys store query (metric: {}, agg: {}): {}ms", - &keys_params.metric, - keys_params.aggregation_id, - keys_store_query_start_time.elapsed().as_millis() - ); - debug!("Keys query returned {} unique keys", keys_map.len()); - - let keys_merge_start_time = Instant::now(); - let merged = self.merge_precomputed_outputs( - &keys_map, - do_merge, - agg_info.aggregation_type_for_key, - ); - debug!( - "[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs", - keys_merge_start_time.elapsed().as_secs_f64() * 1000.0, - merged.len() - ); - Some(merged) - } else { - None - }; + let result = combined_result.ok_or_else(|| { + crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!("SketchStore reducer produced no result for `{query}`"), + ) + })?; - Ok((merged_values, merged_keys, chosen_window)) + // Matrix shape — the range_query wire format requires it. + Ok(asap_tier_result_to_query_result(result, end_ms, true)) } - /// Collects all results based on whether keys are separate or not - fn collect_all_results( - &self, - merged_values: &HashMap, Box>, - merged_keys: Option<&HashMap, Box>>, - statistic: &Statistic, - query_kwargs: &HashMap, - enable_topk_limiting: bool, - ) -> Result, f64>, String> { - if let Some(keys_map) = merged_keys { - // Separate keys and values - self.collect_results_separate_keys(merged_values, keys_map, statistic, query_kwargs) - } else { - // Same aggregation for keys and values - self.collect_results_same_aggregation( - merged_values, - statistic, - query_kwargs, - enable_topk_limiting, - ) - } - } +} - /// Executes the complete query pipeline: plan, execute, collect, and format. - /// - /// Returns the formatted instant-vector elements alongside the - /// `[start_ms, end_ms)` precompute window the engine actually - /// consulted (for tumbling-window queries this is the latest - /// pane that overlapped the request range; for sliding-window - /// queries it's the exact window). Callers attach this onto the - /// outgoing `QueryResult` via `with_window_used` so the - /// HTTP-adapter response can annotate it as - /// `precompute_window`. - pub fn execute_query_pipeline( - &self, - context: &QueryExecutionContext, - enable_topk: bool, - ) -> Result<(Vec, Option<(u64, u64)>), String> { - // Step 1: Execute the query plan (already created in context.store_plan) - let (merged_values, merged_keys, chosen_window) = self.execute_and_merge_store_queries( - &context.store_plan, - context.do_merge, - &context.agg_info, - )?; +// --------------------------------------------------------------------------- +// Phase-5: `QueryEngine` trait impl. +// +// Adapter only — does NOT change `handle_query` or any other existing +// surface. The trait's `execute(&str)` walks the same `handle_query` code +// path the binary's HTTP driver uses today; `None` (capability miss) is +// translated to `EngineError::CapabilityMiss` so the router can fall through +// to the next compatible backend. +// --------------------------------------------------------------------------- - // Step 2: Collect results - let unformatted_results_start_time = Instant::now(); - let unformatted_results = self.collect_all_results( - &merged_values, - merged_keys.as_ref(), - &context.metadata.statistic_to_compute, - &context.metadata.query_kwargs, - enable_topk, // SQL=false, PromQL=true - )?; - debug!( - "[LATENCY] Unformatted results collection: {:.2}ms", - unformatted_results_start_time.elapsed().as_secs_f64() * 1000.0 - ); +/// Adapt a [`crate::storage_engines::sketch_db::query::ASAPTierResult`] to the engine's +/// existing `QueryResult` shape. The reducer hands back per-series +/// time-stamped scalars; we materialize them as a +/// `QueryResult::Matrix` whose [`crate::query_engines::query_result::RangeVectorElement`]s +/// each map onto one (label-values, samples) entry. +/// +/// `now_ms` is unused for the matrix variant (each sample carries its +/// own window-end timestamp); it's plumbed for future extension to +/// the instant-vector case (latest-pane projection). +/// Merge a ASAP-tier `QueryResult::Matrix` with an archive +/// `QueryResult::Matrix` by `(label_values, timestamp)`. Samples whose +/// timestamps fall inside the warm coverage `(cov_lo, cov_hi)` keep +/// the warm value (warm is approximate but more recent); samples +/// outside that window come from the archive answer. For +/// labels-not-present-in-warm series the archive series is taken in +/// full. Used by `ASAPQueryEngine`'s hybrid-stitch path when the +/// ASAP-tier reducer reports `coverage` narrower than the request. +fn stitch_warm_and_archive( + warm: crate::query_engines::query_result::QueryResult, + archive: crate::query_engines::query_result::QueryResult, + cov_lo: u64, + cov_hi: u64, +) -> crate::query_engines::query_result::QueryResult { + use crate::query_engines::query_result::{QueryResult, RangeVectorElement, Sample}; + use std::collections::BTreeMap; - // Step 3: Format results - let results_start_time = Instant::now(); - let results = self.format_final_results( - unformatted_results, - &context.metadata.statistic_to_compute, - &context.metric, - enable_topk, // SQL=false, PromQL=true - ); - debug!( - "[LATENCY] Results collection: {}ms", - results_start_time.elapsed().as_millis() - ); + let warm_matrix = match &warm { + QueryResult::Matrix(m) => m.values.clone(), + _ => return archive}; + let archive_matrix = match &archive { + QueryResult::Matrix(m) => m.values.clone(), + QueryResult::Vector(_) => return warm}; - Ok((results, chosen_window)) + // Index warm series by labels for fast lookup. + let mut by_labels: BTreeMap, RangeVectorElement> = BTreeMap::new(); + for el in warm_matrix { + by_labels.insert(el.labels.labels.clone(), el); } - /// Variant of `build_query_execution_context_promql` that accepts a - /// pre-parsed AST node, avoiding redundant parsing. Agg resolution - /// goes through capability matching (the path the standard builder - /// also falls through to after InferenceConfig retirement). - pub fn build_query_execution_context_from_ast( - &self, - arm_ast: &promql_parser::parser::Expr, - time: f64, - ) -> Option { - let query_time = Self::convert_query_time_to_data_time(time); - - let mut found_match = None; - for (pattern_type, patterns) in &self.control_plane_patterns { - for pattern in patterns { - let match_result = pattern.matches(arm_ast); - if match_result.matches { - found_match = Some((*pattern_type, match_result)); - break; - } - } - if found_match.is_some() { - break; + // For each archive series, merge into by_labels. + for arch_el in archive_matrix { + let entry = by_labels + .entry(arch_el.labels.labels.clone()) + .or_insert_with(|| RangeVectorElement::new(arch_el.labels.clone())); + // Build a set of warm timestamps inside coverage (kept). + let warm_ts: std::collections::HashSet = entry + .samples + .iter() + .filter(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi) + .map(|s| s.timestamp) + .collect(); + // Drop any warm samples that ended up outside coverage — + // archive will replace them. + entry + .samples + .retain(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi); + for s in arch_el.samples { + // Skip archive samples whose timestamps fall inside warm + // coverage AND warm produced a value there (warm wins). + if s.timestamp >= cov_lo && s.timestamp <= cov_hi && warm_ts.contains(&s.timestamp) { + continue; } + entry.samples.push(Sample::new(s.timestamp, s.value)); } - - let (query_pattern_type, match_result) = found_match?; - - let requirements = - self.build_query_requirements_promql(&match_result, query_pattern_type); - let agg_info = self.find_compatible_aggregation_with_miss_notify(&requirements)?; - - self.build_promql_execution_context_tail( - &match_result, - query_pattern_type, - query_time, - agg_info, - ) + entry.samples.sort_by_key(|s| s.timestamp); } - /// Shared context-building tail for both PromQL context builders. - /// - /// Called by `build_query_execution_context_from_ast` and - /// `build_query_execution_context_promql` after pattern matching and - /// `agg_info` resolution are complete. Computes labels, statistics, - /// kwargs, metadata, query plan, and the final `QueryExecutionContext`. - fn build_promql_execution_context_tail( - &self, - match_result: &PromQLMatchResult, - query_pattern_type: QueryPatternType, - query_time: u64, - agg_info: AggregationIdInfo, - ) -> Option { - let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result); + let elements: Vec = by_labels.into_values().collect(); + QueryResult::matrix(elements) +} - // Resolve the metric's "all labels" set. Falls back to a - // streaming-config-derived label union when the schema is - // empty for this metric — the production ASAP-tier deploy - // launches with `--streaming-config` only and an empty - // schema, and pre-fix every query for a streaming-config- - // registered metric blew up here on the schema lookup. See - // `Self::resolve_metric_labels` and the - // `production_conditions_*` regression tests for context. - let all_labels = match self.resolve_metric_labels(&metric) { - Some(labels) => labels, - None => { - warn!("No metric configuration found for '{}'", metric); - return None; - } - }; +fn asap_tier_result_to_query_result( + result: crate::storage_engines::sketch_db::query::ASAPTierResult, + now_ms: u64, + is_range_query: bool, +) -> crate::query_engines::query_result::QueryResult { + use crate::storage_engines::types::KeyByLabelValues; + use crate::query_engines::query_result::{ + InstantVectorElement, QueryResult, RangeVectorElement, + }; - let mut query_output_labels = match query_pattern_type { - QueryPatternType::OnlyTemporal => all_labels.clone(), - QueryPatternType::OnlySpatial => { - get_spatial_aggregation_output_labels(match_result, &all_labels) - } - QueryPatternType::OneTemporalOneSpatial => { - let temporal_aggregation = match_result.get_function_name().unwrap(); - let spatial_aggregation = match_result.get_aggregation_op().unwrap(); - let collapsable = temporal_aggregation - .parse::() - .ok() - .zip(spatial_aggregation.parse::().ok()) - .is_some_and(|(f, o)| get_is_collapsable(f, o)); - if collapsable { - get_spatial_aggregation_output_labels(match_result, &all_labels) - } else { - all_labels.clone() - } + // Instant-query result-shape: the Prometheus adapter's + // `format_success_response` rejects `Matrix` for queries the + // analyzer marked as instant (`range_seconds == 0`) — produces a + // 500 ”shape mismatch”. Project the per-series last sample into + // an `InstantVectorElement` and wrap as `Vector` so the wire + // response carries `resultType: vector` matching the request. + if !is_range_query { + let mut elements: Vec = Vec::with_capacity(result.series.len()); + for (label_values, samples) in result.series { + // Mirror the range-vector branch: BTreeMap iteration is + // key-sorted, so `unzip` produces aligned (keys, values). + // Stash the keys in the per-element `label_keys_override` + // so the Prometheus adapter renders synthesized keys + // (notably ASAP-tier `topk`'s `"item"` key) instead of + // the empty `metric: {}` it would produce when the + // query-scoped `KeyByLabelNames` is empty. + let (keys, values): (Vec, Vec) = label_values.into_iter().unzip(); + let labels = KeyByLabelValues::new_with_labels(values); + // Take the latest sample (the reducer returns one per + // window_end; for instant readout we want the most recent). + if let Some((_, value)) = samples.into_iter().last() { + elements.push( + InstantVectorElement::new(labels, value) + .with_label_keys_override(keys), + ); } - }; - - let timestamps = - self.calculate_query_timestamps_promql(query_time, query_pattern_type, match_result); - - let statistics_to_compute = get_statistics_to_compute(query_pattern_type, match_result); - if statistics_to_compute.len() != 1 { - warn!( - "Expected exactly one statistic to compute, found {}", - statistics_to_compute.len() - ); - return None; } - let statistic_to_compute = statistics_to_compute.first().unwrap(); - - if *statistic_to_compute == Statistic::Topk { - let mut new_labels = vec!["__name__".to_string()]; - new_labels.extend(query_output_labels.labels); - query_output_labels = KeyByLabelNames::new(new_labels); - } - - let query_kwargs = self - .build_query_kwargs_promql(statistic_to_compute, query_pattern_type, match_result) - .map_err(|e| { - warn!("{}", e); - e - }) - .ok()?; - - let metadata = QueryMetadata { - query_output_labels: query_output_labels.clone(), - statistic_to_compute: *statistic_to_compute, - query_kwargs}; - - let query_plan = self - .create_store_query_plan(&metric, ×tamps, &agg_info) - .map_err(|e| { - warn!("Failed to create store query plan: {}", e); - e - }) - .ok()?; - - let do_merge = query_pattern_type == QueryPatternType::OnlyTemporal - || query_pattern_type == QueryPatternType::OneTemporalOneSpatial; - - let streaming_config = self.streaming_config_snapshot(); - let grouping_labels = streaming_config - .get_aggregation_config(agg_info.aggregation_id_for_value) - .map(|config| config.grouping_labels.clone()) - .unwrap_or_else(|| query_output_labels.clone()); - - let aggregated_labels = streaming_config - .get_aggregation_config(agg_info.aggregation_id_for_key) - .map(|config| config.aggregated_labels.clone()) - .unwrap_or_else(KeyByLabelNames::empty); - - Some(QueryExecutionContext { - metric, - metadata, - store_plan: query_plan, - agg_info, - do_merge, - spatial_filter, - query_time, - grouping_labels, - aggregated_labels}) - } - - /// Applies a PromQL binary arithmetic operator to two f64 values. - fn apply_range_binary_op( - op: &promql_parser::parser::token::TokenType, - lhs: f64, - rhs: f64, - ) -> f64 { - use promql_parser::parser::token::{T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; - match op.id() { - id if id == T_ADD => lhs + rhs, - id if id == T_SUB => lhs - rhs, - id if id == T_MUL => lhs * rhs, - id if id == T_DIV => lhs / rhs, - id if id == T_MOD => lhs % rhs, - id if id == T_POW => lhs.powf(rhs), - _ => f64::NAN} + return QueryResult::vector(elements, now_ms); } - /// Recursively builds a range execution context for one arm of a binary arithmetic expression. - fn build_arm_range_context( - &self, - arm_ast: &promql_parser::parser::Expr, - start: f64, - end: f64, - step: f64, - ) -> Option<(RangeQueryExecutionContext, Vec)> { - use promql_parser::parser::Expr; - - match arm_ast { - Expr::NumberLiteral(_) => None, // caller handles scalars - Expr::Paren(paren) => self.build_arm_range_context(&paren.expr, start, end, step), - other => { - let base_context = - self.build_query_execution_context_from_ast(other, end)?; - let label_names = base_context.metadata.query_output_labels.labels.clone(); - - let start_ms = Self::convert_query_time_to_data_time(start); - let end_ms = Self::convert_query_time_to_data_time(end); - let step_ms = (step * 1000.0) as u64; - - let tumbling_window_ms = self - .streaming_config_snapshot() - .get_aggregation_config(base_context.agg_info.aggregation_id_for_value) - .map(|c| c.window_size * 1000)?; - - self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms) - .map_err(|e| { - warn!("Range arm query validation failed: {}", e); - e - }) - .ok()?; - - let lookback_ms = base_context.store_plan.values_query.end_timestamp - - base_context.store_plan.values_query.start_timestamp; - - let buckets_per_step = (step_ms / tumbling_window_ms) as usize; - let lookback_bucket_count = (lookback_ms / tumbling_window_ms) as usize; - - let mut extended_store_plan = base_context.store_plan.clone(); - extended_store_plan.values_query.start_timestamp = - start_ms.saturating_sub(lookback_ms); - extended_store_plan.values_query.end_timestamp = end_ms; - extended_store_plan.values_query.is_exact_query = false; - - let range_context = RangeQueryExecutionContext { - base: QueryExecutionContext { - store_plan: extended_store_plan, - ..base_context - }, - range_params: RangeQueryParams { - start: start_ms, - end: end_ms, - step: step_ms}, - buckets_per_step, - lookback_bucket_count, - tumbling_window_ms}; - - Some((range_context, label_names)) - } + let mut elements: Vec = Vec::with_capacity(result.series.len()); + for (label_values, samples) in result.series { + // `KeyByLabelValues` is a `Vec` carrying VALUES only; + // the serializer pairs them with KEYS from a query-scoped + // `KeyByLabelNames`. For most queries the keys ARE the + // query's group-by clause, so the default path works. But + // ASAP-tier `topk` synthesizes an `"item"` key (the top-k + // entry name) that the original query's group-by doesn't + // carry — without an override the serializer drops it and + // the response shows `"metric": {}`. Project the BTreeMap's + // VALUES in key-sorted order (BTreeMap iteration is + // key-sorted), and stash the BTreeMap's KEYS in the + // per-element override so the serializer can pair them + // correctly. + let (keys, values): (Vec, Vec) = label_values.into_iter().unzip(); + let labels = KeyByLabelValues::new_with_labels(values); + let mut element = RangeVectorElement::new(labels).with_label_keys_override(keys); + for (window_end_ms, value) in samples { + // `window_end_ms` is i64 from the index; cast to u64 + // for the wire format (window_end is monotonic + post- + // 1970 in production). + let ts = if window_end_ms >= 0 { + window_end_ms as u64 + } else { + 0 + }; + element.add_sample(ts, value); } + elements.push(element); } + QueryResult::matrix(elements) +} - /// Handles a binary arithmetic PromQL expression for range queries. - /// - /// Evaluates each arm independently over the full range, then joins the - /// resulting series by label key and applies the arithmetic operator - /// sample-by-sample at matching timestamps. - fn handle_binary_expr_range_promql( +#[async_trait::async_trait] +impl crate::query_engines::routing::query_engine_routing::QueryEngine for ASAPQueryEngine { + async fn execute( &self, - ast: &promql_parser::parser::Expr, - start: f64, - end: f64, - step: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - use promql_parser::parser::Expr; + query: &str, + ) -> Result { + // Phase 9 controller-unification (2026-05) — the ASAP-tier + // hook is now a thin driver around the control plane's + // `analyze_promql_for_asap_tier`. The analyzer is the single + // owner of "is this PromQL ASAP-tier-answerable" knowledge. + // We drop into one of three branches: + // + // 1. `ASAPTierAnalysis::unsupported` is `Some(_)` — the + // PromQL shape isn't ASAP-tier-servable. Surface as + // `EngineError::CapabilityMiss(SketchStore, …)` with the + // structured `UnsupportedReason` in the detail string. The + // EngineRouter fails over to the archive engine. This + // covers all of: + // * `MissReason::UnsupportedFunction(_)` (rate, irate, + // increase, etc.) → cold tier (archive) + // * `MissReason::UnsupportedComposition(_)` (sum-by, + // topk-over-rate, etc.) → cold tier + // * `MissReason::NoCallNodeFound` (bare selector) → + // cold tier (archive answers raw selectors) + // * `MissReason::UnparseablePromql(_)` → cold tier + // (archive's parser may be more permissive, or it'll + // also reject and the user sees the error) + // + // 2. `ASAPTierAnalysis::candidates` is populated, but ANY + // candidate's `instances_matching` returns empty OR a + // sid that classifies as `Ghost`/`Unknown` — surface + // as CapabilityMiss. The EngineRouter falls over. + // + // 3. All candidates resolve to all-`Hit` sids — dispatch + // each to the per-`Capability` sketch reducer. Today's + // semantic: ANY candidate-level reducer error → fall + // over to archive (no per-candidate hybrid stitch yet — + // that's the documented follow-up). + if let Some(idx) = self.sketch_index.as_ref() { + let analysis = control_plane::asap_tier_analysis::analyze_promql_for_asap_tier(query); - let binary = match ast { - Expr::Binary(b) => b, - _ => return None}; + // Branch 1 — the control plane analyzer rejects the shape. + if let Some(reason) = &analysis.unsupported { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore analyzer rejected `{query}`: {reason:?} — \ + failing over to archive" + ), + )); + } + if analysis.candidates.is_empty() { + // Defensive — `is_asap_tier_answerable` would have + // caught this; analyzer guarantees `unsupported.is_some()` + // when `candidates.is_empty()` but we keep the + // belt-and-braces miss-path for safety. + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore analyzer produced no ASAP-tier candidates for \ + `{query}` — failing over to archive" + ), + )); + } - let lhs = binary.lhs.as_ref(); - let rhs = binary.rhs.as_ref(); - let op = &binary.op; + // Branch 2 + 3 — resolve each candidate's sids and + // dispatch the reducer. Today this is single-candidate + // for every supported PromQL shape; the loop is here + // for the per-candidate hybrid-stitch follow-up. + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + // Time bounds: the trait's `execute(&str)` adapter + // doesn't carry an explicit range today (it's an + // instant-query surface). For each candidate, prefer the + // candidate's `range_seconds` (extracted from `[5m]` / + // `[30s]` selectors); fall back to a 5-minute default + // for instant-vector candidates (range_seconds == 0). + const DEFAULT_LOOKBACK_MS: u64 = 5 * 60 * 1000; - // Scalar case: either side may be a numeric literal - let scalar_case: Option<(f64, &Expr, bool)> = match (lhs, rhs) { - (_, Expr::NumberLiteral(nl)) => Some((nl.val, lhs, false)), - (Expr::NumberLiteral(nl), _) => Some((nl.val, rhs, true)), - _ => None}; - if let Some((scalar, vector_arm, scalar_on_left)) = scalar_case { - let (ctx, labels) = self.build_arm_range_context(vector_arm, start, end, step)?; - let results = self.execute_range_query_pipeline(&ctx).ok()?; - let combined: Vec = results - .into_iter() - .map(|mut elem| { - for s in &mut elem.samples { - s.value = if scalar_on_left { - Self::apply_range_binary_op(op, scalar, s.value) - } else { - Self::apply_range_binary_op(op, s.value, scalar) - }; - } - elem - }) - .collect(); - return Some((KeyByLabelNames::new(labels), QueryResult::matrix(combined))); - } + let reducer = crate::storage_engines::sketch_db::query::SketchReducer::new(idx); + // Multi-candidate aggregation is deferred (single-result + // shapes today). On the first reducer error we surface + // CapabilityMiss; on Ok we keep the result for the + // hybrid-stitch path below. (When more than one + // candidate is supported, a follow-up will fold + // per-candidate ASAPTierResults.) + let mut combined_result: Option = + None; + let mut combined_t0: u64 = u64::MAX; + // Track whether ANY candidate is range-vector-shaped + // (`range_seconds > 0`). Drives the Vector-vs-Matrix + // result-shape choice in `asap_tier_result_to_query_result` + // below — instant queries (`count(metric)`, + // `quantile(...)` without `_over_time` etc.) need + // `QueryResult::Vector` so the Prometheus adapter's + // `format_success_response` wraps them as `resultType: + // vector`. Returning `Matrix` for an instant query + // produces a 500 (adapter rejects the shape mismatch). + let mut any_range_candidate = false; - // Vector-vector: evaluate both arms, join by label key, apply op per matching timestamp - let (lhs_ctx, lhs_labels) = self.build_arm_range_context(lhs, start, end, step)?; - let (rhs_ctx, _) = self.build_arm_range_context(rhs, start, end, step)?; - let lhs_results = self.execute_range_query_pipeline(&lhs_ctx).ok()?; - let rhs_results = self.execute_range_query_pipeline(&rhs_ctx).ok()?; + // Snapshot the streaming config once for this query's + // policy lookups. Hot-reload swaps the underlying Arc; the + // snapshot pins one revision for the duration. + let streaming_snap = self.streaming_config_snapshot(); + let policy_registry = streaming_snap.policy_registry(); - // Build lookup: label_key -> {timestamp -> value} for rhs - let mut rhs_map: HashMap> = HashMap::new(); - for elem in rhs_results { - let ts_map: HashMap = elem - .samples - .iter() - .map(|s| (s.timestamp, s.value)) - .collect(); - rhs_map.insert(elem.labels, ts_map); - } + for candidate in &analysis.candidates { + if candidate.range_seconds > 0 { + any_range_candidate = true; + } + // Schema-retirement #5: resolve candidate → {sids} by + // unioning the policy-fp reverse index (fast path for + // ExactAgg sids minted via `ingest_precompute_for_agg_config` + // where `policy_fp` is set) with `instances_matching` + // (catalog walk that subset-matches on + // `group_by_keys`, covering raw sketches whose + // `derive_sketch_policy_fp` returned `UNSET` because + // the wire-attr set didn't match any streaming-config + // policy). The earlier policy-fp-only path returned + // empty for the MVP demo workload — see issue #271 / + // tracking #272. + let policy_fps = control_plane::asap_tier_analysis::find_matching_policies( + &policy_registry, + candidate, + ); + let mut sids: std::collections::BTreeSet = + std::collections::BTreeSet::new(); + for fp in &policy_fps { + sids.extend(idx.sids_for_policy(*fp)); + } + sids.extend(idx.instances_matching( + &candidate.metric_name, + &candidate.group_by_keys, + )); + if sids.is_empty() { + // Fire the capability-miss notify so the + // control-plane feedback loop closes — replaces + // the side-effect the now-deleted + // `find_compatible_aggregation_with_miss_notify` + // provided on the legacy path. + let req = Self::requirements_from_candidate(candidate); + crate::drivers::control_plane_client::spawn_capability_miss_notify( + &self.control_plane_client, + &req, + ); + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore has no policy for metric `{}` \ + with group_by_keys ⊇ {:?} satisfying capability \ + {:?} — failing over to archive", + candidate.metric_name, + candidate.group_by_keys, + candidate.required_capability, + ), + )); + } - let mut combined: Vec = Vec::new(); - for lhs_elem in lhs_results { - if let Some(rhs_ts_map) = rhs_map.get(&lhs_elem.labels) { - let mut new_elem = RangeVectorElement::new(lhs_elem.labels.clone()); - for s in &lhs_elem.samples { - if let Some(&rhs_val) = rhs_ts_map.get(&s.timestamp) { - new_elem.add_sample( - s.timestamp, - Self::apply_range_binary_op(op, s.value, rhs_val), - ); + // Verify each sid carries the analyzer's required + // capability. After Step 2a there's exactly one + // `Capability` enum (defined in the control plane and + // re-exported by `sketch_index`), so no `From` + // conversion is needed — just clone. + let required: crate::storage_engines::sketch_db::index::Capability = + candidate.required_capability.clone(); + let mut hit_sids: Vec = Vec::with_capacity(sids.len()); + for sid in &sids { + match idx.classify(*sid) { + crate::storage_engines::sketch_db::index::SidLookup::Hit => {} + crate::storage_engines::sketch_db::index::SidLookup::Ghost + | crate::storage_engines::sketch_db::index::SidLookup::Unknown => { + let req = Self::requirements_from_candidate(candidate); + crate::drivers::control_plane_client::spawn_capability_miss_notify( + &self.control_plane_client, + &req, + ); + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore ghost/unknown sid {sid} for metric \ + `{}` — failing over to archive", + candidate.metric_name + ), + )); + } + } + let meta = match idx.instance(*sid) { + Some(m) => m, + None => continue}; + // Precompute-backed sids (M2.3) have `capability: None` + // — the analyzer doesn't route them through this path, + // but skip defensively if one slips in. + if let Some(cap) = meta.capability.as_ref() { + if required.is_satisfied_by(cap) { + hit_sids.push(*sid); + } } } - if !new_elem.samples.is_empty() { - combined.push(new_elem); + if hit_sids.is_empty() { + let req = Self::requirements_from_candidate(candidate); + crate::drivers::control_plane_client::spawn_capability_miss_notify( + &self.control_plane_client, + &req, + ); + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore has no sid satisfying capability \ + {:?} for metric `{}` — failing over to archive", + candidate.required_capability, candidate.metric_name + ), + )); } - } - } - - let output_labels = KeyByLabelNames::new(lhs_labels); - Some((output_labels, QueryResult::matrix(combined))) - } - - /// Formats unformatted results into final InstantVectorElement format - /// For topk queries (when enabled), sorts by value and prepends metric name to keys - fn format_final_results( - &self, - unformatted_results: HashMap, f64>, - statistic: &Statistic, - metric: &str, - enable_topk_formatting: bool, - ) -> Vec { - let sorted_results: Vec<(Option, f64)> = - if *statistic == Statistic::Topk && enable_topk_formatting { - // Sort by value descending for topk - let mut sorted: Vec<_> = unformatted_results.into_iter().collect(); - sorted.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - - // Prepend metric name to each key's label values - sorted - .into_iter() - .map(|(key_opt, value)| { - let updated_key = key_opt.map(|mut key| { - let mut new_labels = vec![metric.to_string()]; - new_labels.extend(key.labels); - key.labels = new_labels; - key - }); - (updated_key, value) - }) - .collect() - } else { - unformatted_results.into_iter().collect() - }; - - sorted_results - .into_iter() - .filter_map(|(key, value)| key.map(|k| InstantVectorElement::new(k, value))) - .collect() - } - - /// Extract QueryRequirements from a parsed PromQL match result. - /// Used as the fallback path when no query_configs entry is found. - fn build_query_requirements_promql( - &self, - match_result: &PromQLMatchResult, - query_pattern_type: QueryPatternType, - ) -> QueryRequirements { - let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result); - - let statistics = get_statistics_to_compute(query_pattern_type, match_result); - - let data_range_ms = match query_pattern_type { - QueryPatternType::OnlySpatial => None, - _ => match_result - .get_range_duration() - .map(|d| d.num_seconds() as u64 * 1000)}; - - // Resolve the metric's "all labels" set with the same - // schema-empty fallback used by - // `build_promql_execution_context_tail`. Without this - // fallback the schema-empty production deploy returns - // `KeyByLabelNames::empty()` for every metric, and - // `labels_compatible`'s strict-eq mismatches every agg - // config's `[zone]` → capability-miss → `status=error`. - let all_labels = self - .resolve_metric_labels(&metric) - .unwrap_or_else(KeyByLabelNames::empty); - let grouping_labels = match query_pattern_type { - QueryPatternType::OnlyTemporal => all_labels, - QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => { - get_spatial_aggregation_output_labels(match_result, &all_labels) - } - }; + let lookback_ms = if candidate.range_seconds > 0 { + candidate.range_seconds.saturating_mul(1000) + } else { + DEFAULT_LOOKBACK_MS + }; + let t0_ms = now_ms.saturating_sub(lookback_ms); + if t0_ms < combined_t0 { + combined_t0 = t0_ms; + } - QueryRequirements { - metric, - statistics, - data_range_ms, - grouping_labels, - spatial_filter_normalized: normalize_spatial_filter(&spatial_filter)} - } - - /// Execute the query pipeline for an already-built context. - /// - /// Shared by all `handle_query_*` entry points. - fn execute_context( - &self, - context: QueryExecutionContext, - enable_topk: bool, - ) -> Option<(KeyByLabelNames, QueryResult)> { - let agg_id = context.agg_info.aggregation_id_for_value; - let (results, window_used) = self - .execute_query_pipeline(&context, enable_topk) - .map_err(|e| { - warn!("Query execution failed: {}", e); - e - }) - .ok()?; - let qr = QueryResult::vector(results, context.query_time); - let qr = match self.accuracy_envelope_for(agg_id) { - Some(env) => qr.with_accuracy(env), - None => qr}; - let qr = match window_used { - Some(w) => qr.with_window_used(w), - None => qr}; - Some((context.metadata.query_output_labels, qr)) - } - - /// Build an [`AccuracyEnvelope`] for a single resolved - /// `agg_id` by looking up the matching `AggregationConfig` in - /// the current streaming-config snapshot and deriving its - /// [`AccuracyProfile`]. Returns `None` when the agg isn't in - /// config (e.g. post-retire / test harness with empty config). - pub(crate) fn accuracy_envelope_for( - &self, - agg_id: u64, - ) -> Option { - let snap = self.streaming_config_snapshot(); - let cfg = snap.get_aggregation_config(agg_id)?; - Some(crate::storage_engines::sketch_db::AccuracyEnvelope::single( - crate::storage_engines::sketch_db::AccuracyProfile::derive(cfg), - )) - } - - /// Handle a query following Python's unified architecture. - /// - /// Only PromQL is wired in production; the SQL / Elasticsearch - /// variants were removed entirely during the dead-code cleanup. - pub fn handle_query(&self, query: String, time: f64) -> Option<(KeyByLabelNames, QueryResult)> { - // PromQL is the only supported query language. - self.handle_query_promql(query, time) - } - - // /// Try to extract sketch query components from a PromQL query string. - // /// - // /// Attempts the standard AST parser first. If that fails (e.g. for custom - // /// sketch-only functions), falls back to a lightweight regex extraction for - // /// patterns like `func(metric[range])` and `func(number, metric[range])`. - // /// Extract just the sketch function name from a query without full evaluation. - // fn extract_sketch_func_name(&self, query: &str) -> Option { - // self.parse_sketch_query_components(query) - // .map(|c| c.func_name) - // } - - // fn parse_sketch_query_components(&self, query: &str) -> Option { - // // --- Path A: standard PromQL parser + pattern matching --- - // if let Some(components) = self.parse_sketch_via_ast(query) { - // return Some(components); - // } - - // // --- Path B: regex fallback for custom sketch functions --- - // self.parse_sketch_via_regex(query) - // } - - // /// Parse sketch components using the standard PromQL AST parser. - // fn parse_sketch_via_ast(&self, query: &str) -> Option { - // let ast = match promql_parser::parser::parse(query) { - // Ok(ast) => ast, - // Err(_) => return None, - // }; - - // let mut found_match = None; - // for (pattern_type, patterns) in &self.control_plane_patterns { - // for pattern in patterns { - // let match_result = pattern.matches(&ast); - // if match_result.matches { - // found_match = Some((*pattern_type, match_result)); - // break; - // } - // } - // if found_match.is_some() { - // break; - // } - // } - - // let (query_pattern_type, match_result) = found_match?; - - // if query_pattern_type != QueryPatternType::OnlyTemporal { - // debug!( - // "Sketch query (AST): pattern type {:?} is not OnlyTemporal, skipping for '{}'", - // query_pattern_type, query - // ); - // return None; - // } - - // let func_name = match_result.get_function_name()?; - // promsketch_store::promsketch_func_map(&func_name)?; - - // let (metric, spatial_filter) = get_metric_and_spatial_filter(&match_result); - // let metric = if spatial_filter.is_empty() { - // metric - // } else { - // format!("{}{{{}}}", metric, spatial_filter) - // }; - - // let range_seconds = match_result.get_range_duration()?.num_seconds() as u64; - - // let args = if func_name == "quantile_over_time" { - // self.extract_quantile_param_promql(query_pattern_type, &match_result) - // .and_then(|s| s.parse::().ok()) - // .unwrap_or(0.5) - // } else { - // 0.0 - // }; - - // Some(SketchQueryComponents { - // func_name, - // metric, - // range_seconds, - // args, - // }) - // } - - // /// Regex fallback for custom sketch functions the PromQL parser doesn't know. - // /// - // /// Matches two forms: - // /// - `func_name(metric[duration])` (generic) - // /// - `func_name(number, metric[duration])` (quantile) - // /// - `func_name(metric{filter}[duration])` (with label filter) - // fn parse_sketch_via_regex(&self, query: &str) -> Option { - // use regex::Regex; - - // // quantile form: quantile_over_time(0.5, metric{...}[5m]) - // let quantile_re = - // Regex::new(r"^(\w+)\(\s*([0-9.]+)\s*,\s*(\w+(?:\{[^}]*\})?)\[(\d+)([smhd])\]\s*\)$") - // .ok()?; - - // // generic form: func(metric{...}[5m]) - // let generic_re = - // Regex::new(r"^(\w+)\(\s*(\w+(?:\{[^}]*\})?)\[(\d+)([smhd])\]\s*\)$").ok()?; - - // if let Some(caps) = quantile_re.captures(query.trim()) { - // let func_name = caps[1].to_string(); - // promsketch_store::promsketch_func_map(&func_name)?; - // let args: f64 = caps[2].parse().ok()?; - // let metric = caps[3].to_string(); - // let range_seconds = Self::parse_duration_to_seconds(&caps[4], &caps[5])?; - // debug!( - // "Sketch query (regex/quantile): parsed {} with metric={}, range={}s, args={}", - // func_name, metric, range_seconds, args - // ); - // return Some(SketchQueryComponents { - // func_name, - // metric, - // range_seconds, - // args, - // }); - // } - - // if let Some(caps) = generic_re.captures(query.trim()) { - // let func_name = caps[1].to_string(); - // promsketch_store::promsketch_func_map(&func_name)?; - // let metric = caps[2].to_string(); - // let range_seconds = Self::parse_duration_to_seconds(&caps[3], &caps[4])?; - // debug!( - // "Sketch query (regex/generic): parsed {} with metric={}, range={}s", - // func_name, metric, range_seconds - // ); - // return Some(SketchQueryComponents { - // func_name, - // metric, - // range_seconds, - // args: 0.0, - // }); - // } - - // None - // } - - // /// Convert a numeric value + unit suffix into seconds. - // fn parse_duration_to_seconds(value: &str, unit: &str) -> Option { - // let n: u64 = value.parse().ok()?; - // let multiplier = match unit { - // "s" => 1, - // "m" => 60, - // "h" => 3600, - // "d" => 86400, - // _ => return None, - // }; - // Some(n * multiplier) - // } - - // /// Try to handle a PromQL query via the sketch shortcut path. - // /// Returns Some if the query is sketch-backed and PromSketchStore is available. - // /// Returns None to fall through to the precomputed pipeline. - // fn handle_sketch_query_promql( - // &self, - // query: &str, - // time: f64, - // ) -> Option<(KeyByLabelNames, QueryResult)> { - // let ps = self.promsketch_store.as_ref()?; - - // let components = match self.parse_sketch_query_components(query) { - // Some(c) => c, - // None => { - // debug!( - // "Sketch query: could not parse sketch components from '{}'", - // query - // ); - // return None; - // } - // }; - - // let eval_start = Instant::now(); - - // let query_time = Self::convert_query_time_to_data_time(time); - // let end = query_time; - // let start = end.saturating_sub(components.range_seconds * 1000); - - // debug!( - // "Sketch query: evaluating {}({}) range=[{}, {}] args={}", - // components.func_name, components.metric, start, end, components.args - // ); - - // let results = match ps.eval_matching( - // &components.func_name, - // &components.metric, - // components.args, - // start, - // end, - // ) { - // Ok(r) => r, - // Err(e) => { - // warn!( - // "Sketch query: eval_matching failed for {}({}): {}", - // components.func_name, components.metric, e - // ); - // ps_metrics::SKETCH_QUERIES_TOTAL - // .with_label_values(&["miss"]) - // .inc(); - // return None; - // } - // }; - - // if results.is_empty() { - // debug!( - // "Sketch query: no matching series with data for {}({}), falling through", - // components.func_name, components.metric - // ); - // ps_metrics::SKETCH_QUERIES_TOTAL - // .with_label_values(&["miss"]) - // .inc(); - // return None; - // } - - // ps_metrics::SKETCH_QUERIES_TOTAL - // .with_label_values(&["hit"]) - // .inc(); - // ps_metrics::SKETCH_QUERY_DURATION.observe(eval_start.elapsed().as_secs_f64()); - - // info!( - // "Sketch query: {}({}) returned {} series results", - // components.func_name, - // components.metric, - // results.len() - // ); - - // let elements: Vec = results - // .into_iter() - // .map(|(labels_str, value)| { - // let labels = KeyByLabelValues::new_with_labels(vec![labels_str]); - // InstantVectorElement::new(labels, value) - // }) - // .collect(); - - // let output_labels = KeyByLabelNames::new(vec!["__name__".to_string()]); - // Some((output_labels, QueryResult::vector(elements, query_time))) - // } - - pub fn handle_query_promql( - &self, - query: String, - time: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - let query_start_time = Instant::now(); - debug!("Handling query: {} at time {}", query, time); - - // Binary arithmetic dispatch was previously handled here via a - // DataFusion-based plan combiner. That path was removed alongside - // the datafusion crate; binary arithmetic on ASAP-tier sketches - // will be reintroduced as part of the PromQL-evaluator-on-Gorilla - // follow-up. For now binary expressions fall through to the - // normal dispatch path (which will not match and trigger the - // router's CapabilityMiss failover to the archive engine). - - // Try the §7 schema-timeline dispatch first. Returns Some - // only when the query's [t1, t2] range crosses a - // reconfigure boundary (i.e. two or more agg_ids own pieces - // of the range). In every other case (single-schema range, - // unparseable query) it returns None and we fall through - // to the default single-agg path below. - if let Some(result) = self.try_handle_query_promql_via_timeline(&query, time) { - let total_query_duration = query_start_time.elapsed(); - debug!( - "Timeline-dispatch query handling took: {:.2}ms", - total_query_duration.as_secs_f64() * 1000.0 - ); - return Some(result); - } - - let context = self.build_query_execution_context_promql(query, time)?; - - debug!( - "Querying store for metric: {}, aggregation_id: {}, range: [{}, {}]", - context.metric, - context.agg_info.aggregation_id_for_value, - context.store_plan.values_query.start_timestamp, - context.store_plan.values_query.end_timestamp - ); - - let result = self.execute_context(context, true); - - // Determine query routing order based on function type. - // USampling functions prefer the precomputed path first (sketch fallback), - // while EHUniv/EHKLL functions prefer the sketch path first. - // let prefer_precomputed = self - // .extract_sketch_func_name(&query) - // .is_some_and(|name| is_usampling_function(&name)); - - // if !prefer_precomputed { - // // Non-USampling sketch functions: try sketch path first - // if let Some(result) = self.handle_sketch_query_promql(&query, time) { - // let total_query_duration = query_start_time.elapsed(); - // debug!( - // "Sketch query handling took: {:.2}ms", - // total_query_duration.as_secs_f64() * 1000.0 - // ); - // return Some(result); - // } - // } - - // // Precomputed pipeline - // let precomputed_result = (|| -> Option<(KeyByLabelNames, QueryResult)> { - // let context = self.build_query_execution_context_promql(query.clone(), time)?; - - // debug!( - // "Querying store for metric: {}, aggregation_id: {}, range: [{}, {}]", - // context.metric, - // context.agg_info.aggregation_id_for_value, - // context.store_plan.values_query.start_timestamp, - // context.store_plan.values_query.end_timestamp - // ); - - // let results = self - // .execute_query_pipeline(&context, true) // PromQL: topk enabled - // .map_err(|e| { - // warn!("Query execution failed: {}", e); - // e - // }) - // .ok()?; - - // Some(( - // context.metadata.query_output_labels, - // QueryResult::vector(results, context.query_time), - // )) - // })(); - - // if precomputed_result.is_some() { - // let total_query_duration = query_start_time.elapsed(); - // debug!( - // "Total query handling took: {:.2}ms", - // total_query_duration.as_secs_f64() * 1000.0 - // ); - // return precomputed_result; - // } - - // // Fallback: USampling functions try sketch if precomputed had no data - // if prefer_precomputed { - // if let Some(result) = self.handle_sketch_query_promql(&query, time) { - // let total_query_duration = query_start_time.elapsed(); - // debug!( - // "Sketch fallback query handling took: {:.2}ms", - // total_query_duration.as_secs_f64() * 1000.0 - // ); - // return Some(result); - // } - // } - - let total_query_duration = query_start_time.elapsed(); - debug!( - "Total query handling took: {:.2}ms (no results)", - total_query_duration.as_secs_f64() * 1000.0 - ); - result - } - - pub fn build_query_execution_context_promql( - &self, - query: String, - time: f64, - ) -> Option { - let query_time = Self::convert_query_time_to_data_time(time); - let (query_pattern_type, match_result) = self.parse_and_match_promql(&query)?; - debug!("Found matching query config for: {}", query); - - let query_context_start_time = Instant::now(); - - // Resolve aggregation: try pre-configured query_configs first, - // fall back to capability matching. - let agg_info = self.resolve_agg_info_promql(&query, &match_result, query_pattern_type)?; - - let result = self.build_promql_execution_context_tail( - &match_result, - query_pattern_type, - query_time, - agg_info, - ); - - let query_context_duration = query_context_start_time.elapsed(); - debug!( - "[LATENCY] Query context build: {:.2}ms", - query_context_duration.as_secs_f64() * 1000.0 - ); - - result - } - - /// Like `build_query_execution_context_promql`, but skips the - /// auto-resolution of the `agg_id` and forces the provided - /// `forced_agg_id` instead. The same parse + pattern-match - /// pipeline runs; only the "which aggregation covers this - /// query" step is replaced. - /// - /// Caller: the per-segment dispatch in - /// [`Self::try_handle_query_promql_via_timeline`]. For each - /// `TimelineSegment` returned by - /// [`crate::storage_engines::sketch_db::query::timeline::timeline_for_metric`], - /// the dispatch builds a context targeting that segment's - /// `agg_id`, executes it against the clipped segment range, - /// collects the scalar, and combines across segments via - /// [`crate::query_engines::timeline_dispatch::combine_statistic`]. - /// - /// Returns `None` when the query can't be parsed / pattern-matched, - /// or when `forced_agg_id` isn't in the current `StreamingConfig`. - pub fn build_query_execution_context_promql_for_agg_id( - &self, - query: String, - time: f64, - forced_agg_id: u64, - ) -> Option { - let query_time = Self::convert_query_time_to_data_time(time); - let (query_pattern_type, match_result) = self.parse_and_match_promql(&query)?; - let agg_info = self.agg_info_from_forced_id(forced_agg_id)?; - self.build_promql_execution_context_tail( - &match_result, - query_pattern_type, - query_time, - agg_info, - ) - } - - /// Shared phase-1 of PromQL context construction: parse the - /// query string, run it against every registered pattern, and - /// return the matched `(QueryPatternType, PromQLMatchResult)`. - /// Centralised so the auto-resolution and forced-agg-id entry - /// points share identical parse semantics. - fn parse_and_match_promql(&self, query: &str) -> Option<(QueryPatternType, PromQLMatchResult)> { - let parse_start_time = Instant::now(); - let ast = match promql_parser::parser::parse(query) { - Ok(ast) => { - let parse_duration = parse_start_time.elapsed(); - debug!( - "PromQL parsing took: {:.2}ms", - parse_duration.as_secs_f64() * 1000.0 - ); - ast - } - Err(e) => { - warn!("Failed to parse PromQL query '{}': {}", query, e); - return None; - } - }; - - let pattern_match_start_time = Instant::now(); - - let mut found_match = None; - for (pattern_type, patterns) in &self.control_plane_patterns { - for pattern in patterns { - debug!( - "Trying pattern type: {:?} for query: {}", - pattern_type, query - ); - let match_result = pattern.matches(&ast); - debug!("Match result: {:?}", match_result); - if match_result.matches { - found_match = Some((*pattern_type, match_result)); - break; - } - } - if found_match.is_some() { - break; - } - } - - match found_match { - Some((pt, result)) => { - let pattern_match_duration = pattern_match_start_time.elapsed(); - debug!( - "Pattern matching took: {:.2}ms", - pattern_match_duration.as_secs_f64() * 1000.0 - ); - Some((pt, result)) - } - None => { - warn!("No matching pattern found for query: {}", query); - None - } - } - } - - /// Resolve which aggregation covers a PromQL query: try the - /// `QueryConfig` exact-string match first, then fall back to - /// capability-based matching (with control plane miss-notification - /// if wired). Extracted from `build_query_execution_context_promql` - /// so the per-segment timeline dispatch can choose NOT to - /// auto-resolve (it has a forced agg_id from the timeline). - fn resolve_agg_info_promql( - &self, - _query: &str, - match_result: &PromQLMatchResult, - query_pattern_type: QueryPatternType, - ) -> Option { - // InferenceConfig was retired; agg resolution always goes - // through capability matching now. - let requirements = - self.build_query_requirements_promql(match_result, query_pattern_type); - self.find_compatible_aggregation_with_miss_notify(&requirements) - } - - /// Build an `AggregationIdInfo` from a single forced `agg_id`, - /// for the per-segment timeline dispatch. Uses the same - /// "one agg covers both key and value" shape as the single- - /// aggregation branch in `get_aggregation_id_info` (line - /// ~1881), so downstream dispatch treats this agg identically - /// to a single-aggregation `QueryConfig` match. - /// - /// Returns `None` if the agg_id isn't in the current - /// `StreamingConfig` — either a stale control plane posted a - /// backfill for a removed agg, or the timeline contains a - /// retired entry whose config was evicted. Either way the - /// per-segment dispatch will skip this segment as - /// non-answerable. - fn agg_info_from_forced_id(&self, agg_id: u64) -> Option { - let streaming_config = self.streaming_config_snapshot(); - let agg_type = streaming_config - .get_aggregation_config(agg_id) - .map(|c| c.aggregation_type)?; - Some(AggregationIdInfo { - aggregation_id_for_key: agg_id, - aggregation_id_for_value: agg_id, - aggregation_type_for_key: agg_type, - aggregation_type_for_value: agg_type}) - } - - /// Per-segment dispatch across the agg-signature timeline. - /// - /// Returns `Some(result)` when - /// [`Self::timeline_for_query`] yields two or more segments for the - /// query's metric within its time range (i.e. the query spans a - /// reconfigure boundary). - /// Returns `None` otherwise (single-schema range, unparseable - /// query, unresolved probe aggregation) so the caller falls back - /// to the default single-agg path — that path is still correct - /// whenever the timeline doesn't actually span a boundary. - /// - /// ## Combinable vs non-combinable statistics - /// - /// For combinable scalar statistics (Count / Sum / Min / Max) - /// every segment contributes and the result is a clean `Full` - /// value the user can trust without caveat. - /// - /// For non-combinable statistics (Quantile / Topk / Cardinality / - /// Rate / Increase) — or for any combinable run that includes a - /// `Purged` / config-missing segment — `combine_statistic` - /// returns `Partial`. This method surfaces Partial on the - /// Prometheus HTTP response's `warnings` field: the top-level - /// result carries whatever combinable prefix we could compute - /// (for additive stats) or an empty vector (for non-combinable), - /// plus one or more `warnings` strings explaining the schema - /// boundary, the dropped groups, and the unresolved segments. - /// - /// Delivers the user-visible outcome documented in - /// `docs/design-sketch-db.md` §7: queries spanning a reconfigure - /// boundary no longer see a silent data cliff — additive stats - /// get the combined answer, and non-combinable stats get an - /// explicit Partial notice instead of the arbitrary single-agg - /// single-segment result. - fn try_handle_query_promql_via_timeline( - &self, - query: &str, - time: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - use crate::query_engines::timeline_dispatch::{combine_statistic, CombinedResult, SegmentValue}; - use crate::storage_engines::sketch_db::{TimelineCoverage, TimelineSegment}; - - // Phase 1: shared pipeline with the default path — parse, - // pattern-match, auto-resolve a "probe" agg. We reuse the - // probe context purely to read the derived (metric name, - // query time range, statistic) triple that timeline dispatch - // needs. The probe agg itself is NOT used for execution in - // the multi-segment branch. - let (query_pattern_type, match_result) = self.parse_and_match_promql(query)?; - let metric_name = match_result.get_metric_name()?; - let probe_agg_info = - self.resolve_agg_info_promql(query, &match_result, query_pattern_type)?; - let query_time = Self::convert_query_time_to_data_time(time); - let probe_context = self.build_promql_execution_context_tail( - &match_result, - query_pattern_type, - query_time, - probe_agg_info, - )?; - - let stat = probe_context.metadata.statistic_to_compute; - let t1 = probe_context.store_plan.values_query.start_timestamp; - let t2 = probe_context.store_plan.values_query.end_timestamp; - - // Phase 2: resolve the agg-signature timeline over [t1, t2] - // for this metric. Zero or one segments means the default - // single-agg path is already correct; bail out and let the - // caller use it. - let segments = self.timeline_for_query(&metric_name, t1, t2); - if segments.len() < 2 { - return None; - } - - // Schema retirement #3: the sid-level timeline populates - // `agg_id` with a content-hash of the agg-signature - // `(metric, agg_kind, group_by_keys)` rather than with a - // `StreamingConfig.aggregation_id`. Until schema-retirement #5 - // ports the per-segment dispatch to sid-level evaluation - // directly, the segment-→-agg_config mapping below is - // best-effort: if no segment's signature happens to coincide - // with an in-config aggregation_id, fall back to the default - // single-agg path so cross-reconfigure queries don't - // regress to "empty result + warnings". - let snap_for_check = self.streaming_config_snapshot(); - if !segments - .iter() - .any(|s| snap_for_check.get_aggregation_config(s.agg_id).is_some()) - { - return None; - } - drop(snap_for_check); - - debug!( - metric = %metric_name, - segments = segments.len(), - t1, - t2, - statistic = ?stat, - "schema-timeline dispatch: evaluating per-segment" - ); - - // Phase 4: per-segment evaluation. Each segment's agg_id - // evaluates the same query clipped to the segment's - // [start_ms, end_ms]. `Purged` segments (TimelineCoverage) - // or missing configs are collected into `unresolved` so the - // combiner can surface them. - let mut per_group: HashMap, Vec> = HashMap::new(); - let mut unresolved: Vec = Vec::new(); - - for segment in &segments { - if matches!(segment.coverage, TimelineCoverage::Purged) { - unresolved.push(segment.clone()); - continue; - } - let mut ctx = match self.build_query_execution_context_promql_for_agg_id( - query.to_string(), - time, - segment.agg_id, - ) { - Some(c) => c, - None => { - // agg_id no longer in the current StreamingConfig - // (e.g. control plane pushed a swap that dropped - // this entry between timeline resolution and - // dispatch). Classify as unresolved. - unresolved.push(segment.clone()); - continue; - } - }; - - // Clip the segment's [start, end) onto the store plan so - // the per-segment query reads only its own time slice. - ctx.store_plan.values_query.start_timestamp = segment.start_ms; - ctx.store_plan.values_query.end_timestamp = segment.end_ms; - if let Some(ref mut keys_q) = ctx.store_plan.keys_query { - keys_q.start_timestamp = segment.start_ms; - keys_q.end_timestamp = segment.end_ms; - } - - let (per_segment_results, _segment_window) = - match self.execute_query_pipeline(&ctx, true) { - Ok(v) => v, - Err(e) => { - warn!( - agg_id = segment.agg_id, - start_ms = segment.start_ms, - end_ms = segment.end_ms, - "Timeline segment execution failed: {}", - e - ); - unresolved.push(segment.clone()); - continue; - } - }; - - // The per-segment window isn't surfaced in the combined - // result for the schema-timeline dispatch path — the - // combined answer spans multiple agg_ids/windows by - // design, so a single `precompute_window` annotation - // would be misleading. The single-agg path - // (execute_context above) carries it through normally. - - debug!( - agg_id = segment.agg_id, - count = per_segment_results.len(), - "schema-timeline dispatch: segment produced results" - ); - for el in per_segment_results { - per_group - .entry(Some(el.labels)) - .or_default() - .push(SegmentValue { - segment: segment.clone(), - value: el.value}); - } - } - debug!( - groups = per_group.len(), - unresolved = unresolved.len(), - "schema-timeline dispatch: about to combine" - ); - - // Phase 5: per-group combine. Group-by label-tuple so the - // combiner folds per-segment scalars into one final scalar - // per group. Groups that only appear in `unresolved` (no - // segment ever produced a value for them) are skipped. - // - // `any_partial` tracks whether any group came back non-`Full` - // — drives the Prometheus `warnings` surface below so the - // caller sees "this answer is partial" explicitly instead of - // a silent cliff. - let mut output: Vec = Vec::new(); - let mut any_partial = false; - let mut groups_with_no_value = 0usize; - for (label_key, segment_values) in per_group { - match combine_statistic(stat, &segment_values, &unresolved) { - CombinedResult::Full(v) => { - output.push(InstantVectorElement::new(label_key.unwrap_or_default(), v)); - } - CombinedResult::Partial { - covered: Some(v), .. - } => { - // Emit `covered` for combinable stats so the user - // sees the partial sum. The warning below tells - // them not to trust the scalar as a full range - // answer. - output.push(InstantVectorElement::new(label_key.unwrap_or_default(), v)); - any_partial = true; - } - CombinedResult::Partial { covered: None, .. } => { - // Non-combinable stat (quantile / topk / rate / - // increase / cardinality) OR a group the - // combiner couldn't reduce. Drop the group — - // there is no meaningful scalar to show — but - // flag the whole response partial. - any_partial = true; - groups_with_no_value += 1; - } - } - } - - // Phase 6: build Prometheus `warnings` when the combiner - // returned any Partial. One line summarising the schema - // boundary, plus up-to-three per-segment lines with agg_id - // + clipped range so operators can correlate against the - // `GET /api/v1/db/timeline` surface. We cap at three to - // keep responses bounded; the full set is still inspectable - // via the timeline endpoint. - let warnings = if any_partial { - let mut w = Vec::with_capacity(2 + unresolved.len().min(3)); - w.push(format!( - "partial result: query spans {} schemas for metric '{}' over [{}, {}] and the requested statistic {:?} is not cleanly combinable across schema boundaries — see `GET /api/v1/db/timeline?metric={}&start_ms={}&end_ms={}` for the full segment map", - segments.len(), - metric_name, - t1, - t2, - stat, - metric_name, - t1, - t2, - )); - if groups_with_no_value > 0 { - w.push(format!( - "{} group(s) dropped because no segment could answer the statistic", - groups_with_no_value, - )); - } - for seg in unresolved.iter().take(3) { - w.push(format!( - "segment agg_id={} [{}, {}) status={:?} coverage={:?} unresolved", - seg.agg_id, seg.start_ms, seg.end_ms, seg.status, seg.coverage, - )); - } - if unresolved.len() > 3 { - w.push(format!("... and {} more", unresolved.len() - 3)); - } - w - } else { - Vec::new() - }; - - // §6.4: build a per-segment accuracy envelope from each - // segment's resolved agg_id. Segments that don't resolve - // to an in-config agg drop out — their partial-ness is - // already reflected in `warnings` above. - let snap = self.streaming_config_snapshot(); - let per_segment: Vec = segments - .iter() - .filter_map(|seg| { - let cfg = snap.get_aggregation_config(seg.agg_id)?; - Some(crate::storage_engines::sketch_db::PerSegmentAccuracy { - agg_id: seg.agg_id, - range_ms: [seg.start_ms as i64, seg.end_ms as i64], - profile: crate::storage_engines::sketch_db::AccuracyProfile::derive(cfg)}) - }) - .collect(); - let envelope = crate::storage_engines::sketch_db::AccuracyEnvelope::from_segments(per_segment); - - let qr = QueryResult::vector_with_warnings(output, probe_context.query_time, warnings); - let qr = match envelope { - Some(e) => qr.with_accuracy(e), - None => qr}; - Some((probe_context.metadata.query_output_labels, qr)) - } - - /// Merge precomputed outputs (extracts buckets from timestamped data) - fn merge_precomputed_outputs( - &self, - precomputed_outputs_map: &TimestampedBucketsMap, - do_merge: bool, - aggregation_type: AggregationType, - ) -> HashMap, Box> { - #[cfg(feature = "extra_debugging")] - let start_time = Instant::now(); - #[cfg(feature = "extra_debugging")] - debug!("Starting merge for {} keys", precomputed_outputs_map.len()); - #[cfg(feature = "extra_debugging")] - debug!( - "do_merge: {}, aggregation_type: {:?}", - do_merge, aggregation_type - ); - - // Merge iff a temporal query asked us to. The historical - // `DeltaSetAggregator` arm (which forced a merge to - // accumulate keys over time) is retired with the rest of - // the set-tracking family — there's no other accumulator - // today that requires the force-merge override. - let _ = aggregation_type; - let should_merge = do_merge; - - let mut merged = HashMap::with_capacity(precomputed_outputs_map.len()); - - for (key, timestamped_buckets) in precomputed_outputs_map.iter() { - if !timestamped_buckets.is_empty() { - // Extract just the buckets (without timestamps) for merging - let precomputes: Vec> = timestamped_buckets - .iter() - .map(|(_, bucket)| bucket.clone_boxed_core()) - .collect(); - - if should_merge { - #[cfg(feature = "extra_debugging")] - debug!(" Merging accumulators (should_merge=true)"); - #[cfg(feature = "extra_debugging")] - let merge_start = Instant::now(); - let merged_accumulator = self.merge_accumulators(&precomputes); - #[cfg(feature = "extra_debugging")] - let merge_duration = merge_start.elapsed(); - #[cfg(feature = "extra_debugging")] - debug!( - " Merge completed in {:.2}ms, result type: {}", - merge_duration.as_secs_f64() * 1000.0, - merged_accumulator.get_accumulator_type() - ); - merged.insert(key.clone(), merged_accumulator); - } else { - assert_eq!( - precomputes.len(), - 1, - "Spatial queries should have exactly 1 precompute per key" - ); - merged.insert(key.clone(), precomputes[0].clone_boxed_core()); - } - } - } - - #[cfg(feature = "extra_debugging")] - let total_duration = start_time.elapsed(); - #[cfg(feature = "extra_debugging")] - debug!( - "[LATENCY] Complete merge operation: {:.2}ms, merged {} keys", - total_duration.as_secs_f64() * 1000.0, - merged.len() - ); - - merged - } - - /// Merge multiple accumulators using the merge_with method from AggregateCore trait - /// This follows the Python merge_accumulators approach - fn merge_accumulators( - &self, - accumulators: &[Box], - ) -> Box { - if accumulators.is_empty() { - panic!("No accumulators to merge"); - } - - if accumulators.len() == 1 { - return accumulators[0].clone_boxed_core(); - } - - // Try to use optimized batch merge for KLL accumulators - if accumulators[0].get_accumulator_type() == AggregationType::DatasketchesKLL { - use crate::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; - - match DatasketchesKLLAccumulator::merge_multiple(accumulators) { - Ok(merged) => return Box::new(merged), - Err(e) => { - warn!( - "Batch merge failed: {}. Falling back to sequential merge.", - e - ); - // Fall through to sequential merge below - } - } - } - - // Try to use optimized batch merge for CountMinSketch accumulators - if accumulators[0].get_accumulator_type() == AggregationType::CountMinSketch { - use crate::precompute_engine::operators::count_min_sketch_accumulator::CountMinSketchAccumulator; - - match CountMinSketchAccumulator::merge_multiple(accumulators) { - Ok(merged) => return Box::new(merged), - Err(e) => { - warn!( - "Batch merge failed: {}. Falling back to sequential merge.", - e - ); - // Fall through to sequential merge below - } - } - } - - // Fallback: sequential merge for other accumulator types - // (Still benefits from Phase 1 optimization of merge_with) - let mut result = accumulators[0].clone_boxed_core(); - - for accumulator in &accumulators[1..] { - match result.merge_with(accumulator.as_ref()) { - Ok(merged) => { - result = merged; - } - Err(e) => { - warn!("Failed to merge accumulator: {}. Using existing result.", e); - // Continue with the current result if merge fails - } - } - } - - result - } - - /// Collects results when key and value use different aggregations - fn collect_results_separate_keys( - &self, - merged_values: &HashMap, Box>, - merged_keys: &HashMap, Box>, - statistic: &Statistic, - query_kwargs: &HashMap, - ) -> Result, f64>, String> { - let mut unformatted_results = HashMap::new(); - - for (key, precompute) in merged_keys { - let keys_for_this_precompute = precompute - .get_keys() - .ok_or_else(|| "Keys required for separate aggregation".to_string())?; - - for key_for_this_precompute in keys_for_this_precompute { - let value_precompute = merged_values - .get(key) - .ok_or_else(|| format!("No value for key: {:?}", key))?; - - let value = self - .query_precompute_for_statistic( - value_precompute.as_ref(), - statistic, - &Some(key_for_this_precompute.clone()), - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(Some(key_for_this_precompute.clone()), value); - } - } - - Ok(unformatted_results) - } - - /// Collects results when key and value use same aggregation - fn collect_results_same_aggregation( - &self, - merged_outputs: &HashMap, Box>, - statistic: &Statistic, - query_kwargs: &HashMap, - enable_topk_limiting: bool, - ) -> Result, f64>, String> { - let mut unformatted_results = HashMap::new(); - - for (key, precompute) in merged_outputs { - if let Some(unwrapped_keys) = precompute.get_keys() { - let keys_to_process = if enable_topk_limiting { - self.limit_keys_for_topk(unwrapped_keys, statistic, query_kwargs)? - } else { - unwrapped_keys - }; - - for key_for_this_precompute in keys_to_process { - let value = self - .query_precompute_for_statistic( - precompute.as_ref(), - statistic, - &Some(key_for_this_precompute.clone()), - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(Some(key_for_this_precompute.clone()), value); - } - } else { - let value = self - .query_precompute_for_statistic( - precompute.as_ref(), - statistic, - &None, - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(key.clone(), value); - } - } - - Ok(unformatted_results) - } - - /// Limits keys for topk queries - fn limit_keys_for_topk( - &self, - keys: Vec, - statistic: &Statistic, - query_kwargs: &HashMap, - ) -> Result, String> { - if *statistic != Statistic::Topk { - return Ok(keys); - } - - let k_str = query_kwargs - .get("k") - .ok_or_else(|| "Missing k parameter for topk".to_string())?; - - let k = k_str - .parse::() - .map_err(|_| format!("Failed to parse k: '{}'", k_str))?; - - Ok(keys.into_iter().take(k).collect()) - } - - fn query_precompute_for_statistic( - &self, - precompute: &dyn AggregateCore, - statistic: &Statistic, - key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - // Phase 1b of the sketch DB design - // (docs/design-sketch-db.md §5.1 / §16 Phase 1): - // for single-subpopulation queries on additive statistics - // (Count / Sum / Min / Max), serve from the typed aux - // columns without deserialising the sketch payload. - // - // Keyed queries (`key.is_some()`) still need the full - // `query_statistic` path — aux is per-accumulator, not - // per-subpopulation key. - // - // `try_answer` returns `None` when the statistic isn't - // covered by aux (Quantile / Cardinality / TopK / Increase / - // Rate) or when the accumulator doesn't track the requested - // aux field; both cases fall through to the existing path - // so the query result is semantically identical. - if key.is_none() { - if let Some(value) = precompute.aux_stats().try_answer(*statistic) { - return Ok(value); - } - } - precompute.query_statistic(*statistic, key, query_kwargs) - } - - // ============================================================ - // Range Query Support - // ============================================================ - - /// Validate range query parameters - fn validate_range_query_params( - &self, - start: u64, - end: u64, - step: u64, - tumbling_window_ms: u64, - ) -> Result<(), String> { - if start >= end { - return Err("start must be before end".to_string()); - } - if step == 0 { - return Err("step must be positive".to_string()); - } - if !step.is_multiple_of(tumbling_window_ms) { - return Err(format!( - "step ({} ms) must be a multiple of tumbling window size ({} ms)", - step, tumbling_window_ms - )); - } - Ok(()) - } - - /// Build execution context for range query - pub fn build_range_query_execution_context_promql( - &self, - query: String, - start: f64, - end: f64, - step: f64, - ) -> Option { - // First, build the base instant query context (reuse existing logic) - // Use 'end' as the reference time for parsing - let base_context = self.build_query_execution_context_promql(query, end)?; - - // Convert to milliseconds - let start_ms = Self::convert_query_time_to_data_time(start); - let end_ms = Self::convert_query_time_to_data_time(end); - let step_ms = (step * 1000.0) as u64; - - // Get window size - let tumbling_window_ms = self - .streaming_config_snapshot() - .get_aggregation_config(base_context.agg_info.aggregation_id_for_value) - .map(|config| config.window_size * 1000)?; - - // Validate parameters - self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms) - .map_err(|e| { - warn!("Range query validation failed: {}", e); - e - }) - .ok()?; - - // Calculate lookback from the base context's store plan - let lookback_ms = base_context.store_plan.values_query.end_timestamp - - base_context.store_plan.values_query.start_timestamp; - - let buckets_per_step = (step_ms / tumbling_window_ms) as usize; - let lookback_bucket_count = (lookback_ms / tumbling_window_ms) as usize; - - // Modify the store plan to cover the entire range - let mut extended_store_plan = base_context.store_plan.clone(); - extended_store_plan.values_query.start_timestamp = start_ms.saturating_sub(lookback_ms); - extended_store_plan.values_query.end_timestamp = end_ms; - // Range queries always use range fetch, not exact - extended_store_plan.values_query.is_exact_query = false; - - Some(RangeQueryExecutionContext { - base: QueryExecutionContext { - store_plan: extended_store_plan, - ..base_context - }, - range_params: RangeQueryParams { - start: start_ms, - end: end_ms, - step: step_ms}, - buckets_per_step, - lookback_bucket_count, - tumbling_window_ms}) - } - - // /// Try to handle a PromQL range query via the sketch shortcut path. - // /// Returns Some if the query is sketch-backed and PromSketchStore is available. - // /// Returns None to fall through to the precomputed pipeline. - // fn handle_sketch_range_query_promql( - // &self, - // query: &str, - // start: f64, - // end: f64, - // step: f64, - // ) -> Option<(KeyByLabelNames, QueryResult)> { - // let ps = self.promsketch_store.as_ref()?; - - // let components = match self.parse_sketch_query_components(query) { - // Some(c) => c, - // None => { - // debug!( - // "Sketch range query: could not parse sketch components from '{}'", - // query - // ); - // return None; - // } - // }; - - // let eval_start = Instant::now(); - // let range_ms = components.range_seconds * 1000; - - // // Convert query params to ms - // let start_ms = Self::convert_query_time_to_data_time(start); - // let end_ms = Self::convert_query_time_to_data_time(end); - // let step_ms = (step * 1000.0) as u64; - - // if step_ms == 0 || start_ms >= end_ms { - // warn!( - // "Sketch range query: invalid params step_ms={}, start_ms={}, end_ms={}", - // step_ms, start_ms, end_ms - // ); - // return None; - // } - - // // Get all matching series labels - // let series_labels = ps.matching_series_labels(&components.metric); - // if series_labels.is_empty() { - // debug!( - // "Sketch range query: no matching series for {}, falling through", - // components.metric - // ); - // return None; - // } - - // info!( - // "Sketch range query: {}({}) over [{}, {}] step {} with {} series", - // components.func_name, - // components.metric, - // start_ms, - // end_ms, - // step_ms, - // series_labels.len() - // ); - - // // For each matching series, iterate over time steps - // let mut range_elements: Vec = Vec::new(); - - // for series_label in &series_labels { - // let labels = KeyByLabelValues::new_with_labels(vec![series_label.clone()]); - // let mut element = RangeVectorElement::new(labels); - - // let mut current_time = start_ms; - // while current_time <= end_ms { - // let step_end = current_time; - // let step_start = step_end.saturating_sub(range_ms); - - // match ps.eval( - // &components.func_name, - // series_label, - // components.args, - // step_start, - // step_end, - // ) { - // Ok(value) => element.add_sample(current_time, value), - // Err(e) => { - // debug!( - // "Sketch range query: eval failed for {} at t={}: {}", - // series_label, current_time, e - // ); - // } - // } - - // current_time += step_ms; - // } - - // if !element.samples.is_empty() { - // range_elements.push(element); - // } - // } - - // if range_elements.is_empty() { - // debug!( - // "Sketch range query: all series produced empty results for {}({})", - // components.func_name, components.metric - // ); - // ps_metrics::SKETCH_QUERIES_TOTAL - // .with_label_values(&["miss"]) - // .inc(); - // return None; - // } - - // ps_metrics::SKETCH_QUERIES_TOTAL - // .with_label_values(&["hit"]) - // .inc(); - // ps_metrics::SKETCH_QUERY_DURATION.observe(eval_start.elapsed().as_secs_f64()); - - // let output_labels = KeyByLabelNames::new(vec!["__name__".to_string()]); - // Some((output_labels, QueryResult::matrix(range_elements))) - // } - - /// Main entry point for range queries - pub fn handle_range_query_promql( - &self, - query: String, - start: f64, - end: f64, - step: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - let query_start_time = Instant::now(); - debug!( - "Handling range query: {} from {} to {} step {}", - query, start, end, step - ); - - // Check for binary arithmetic before attempting single-query dispatch. - if let Ok(ast) = promql_parser::parser::parse(&query) { - if matches!(&ast, promql_parser::parser::Expr::Binary(_)) { - let result = self.handle_binary_expr_range_promql(&ast, start, end, step); - let total_duration = query_start_time.elapsed(); - debug!( - "Binary arithmetic range query handling took: {:.2}ms", - total_duration.as_secs_f64() * 1000.0 - ); - return result; - } - } - - let context = self.build_range_query_execution_context_promql(query, start, end, step)?; - - // Execute range query pipeline - let results: Vec = self - .execute_range_query_pipeline(&context) - .map_err(|e| { - warn!("Range query execution failed: {}", e); - e - }) - .ok()?; - - // // Determine query routing order based on function type. - // // USampling functions prefer the precomputed path first (sketch fallback), - // // while EHUniv/EHKLL functions prefer the sketch path first. - // let prefer_precomputed = self - // .extract_sketch_func_name(&query) - // .is_some_and(|name| is_usampling_function(&name)); - - // if !prefer_precomputed { - // // Non-USampling sketch functions: try sketch path first - // if let Some(result) = self.handle_sketch_range_query_promql(&query, start, end, step) { - // let total_duration = query_start_time.elapsed(); - // debug!( - // "Sketch range query handling took: {:.2}ms", - // total_duration.as_secs_f64() * 1000.0 - // ); - // return Some(result); - // } - // } - - // // Precomputed pipeline - // let precomputed_result = (|| -> Option<(KeyByLabelNames, QueryResult)> { - // let context = - // self.build_range_query_execution_context_promql(query.clone(), start, end, step)?; - - // let results: Vec = self - // .execute_range_query_pipeline(&context) - // .map_err(|e| { - // warn!("Range query execution failed: {}", e); - // e - // }) - // .ok()?; - - // Some(( - // context.base.metadata.query_output_labels, - // QueryResult::matrix(results), - // )) - // })(); - - // // Fallback: USampling functions try sketch if precomputed had no data - // if prefer_precomputed { - // if let Some(result) = self.handle_sketch_range_query_promql(&query, start, end, step) { - // let total_duration = query_start_time.elapsed(); - // debug!( - // "Sketch fallback range query handling took: {:.2}ms", - // total_duration.as_secs_f64() * 1000.0 - // ); - // return Some(result); - // } - // } - - let total_duration = query_start_time.elapsed(); - debug!( - "Total range query handling took: {:.2}ms", - total_duration.as_secs_f64() * 1000.0 - ); - - Some(( - context.base.metadata.query_output_labels, - QueryResult::matrix(results), - )) - } - - /// Modern warm-tier path for `/api/v1/query_range` — the range- - /// query equivalent of the `QueryEngine::execute(&str)` trait - /// surface. Used by the HTTP server as a fallback when the legacy - /// `handle_range_query_promql` returns `None`. - /// - /// Time semantics follow Prometheus's - /// `/api/v1/query_range?start&end&step` spec: the result is a - /// `matrix` (one row per series, each row carrying multiple - /// (timestamp, value) samples). The warm-tier reducer naturally - /// produces one sample per window_close in `[start, end]`, so - /// the matrix is sampled at the underlying aggregation's window - /// boundaries — typically a finer grid than the user's `step` - /// when window_size < step. (The Prometheus spec says - /// evaluate at each step `t = start, start+step, …, end`; the - /// warm tier returns at native window-close granularity instead. - /// This is more data, not less — clients that expect exact step - /// timestamps can downsample, or route step-precise queries to - /// the cold tier via the EngineRouter.) - /// - /// `step` is currently accepted for API compatibility but unused - /// — see the granularity-mismatch note above. - pub async fn execute_range_promql_modern( - &self, - query: &str, - start_ms: u64, - end_ms: u64, - _step_ms: u64, - ) -> Result< - crate::query_engines::query_result::QueryResult, - crate::query_engines::EngineError, - > { - let Some(idx) = self.sketch_index.as_ref() else { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!("ASAPQueryEngine: no sketch index for `{query}` — failing over"), - )); - }; - - let analysis = - control_plane::asap_tier_analysis::analyze_promql_for_asap_tier(query); - - if let Some(reason) = &analysis.unsupported { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore analyzer rejected `{query}` for range query: \ - {reason:?} — failing over to archive" - ), - )); - } - if analysis.candidates.is_empty() { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore analyzer produced no ASAP-tier candidates for \ - `{query}` — failing over to archive" - ), - )); - } - - let streaming_snap = self.streaming_config_snapshot(); - let policy_registry = streaming_snap.policy_registry(); - let reducer = crate::storage_engines::sketch_db::query::SketchReducer::new(idx); - let mut combined_result: Option< - crate::storage_engines::sketch_db::query::ASAPTierResult, - > = None; - - for candidate in &analysis.candidates { - // Resolve candidate → {sids} via the sid catalog. Schema- - // retirement #5: prefer `instances_matching` over the - // policy-fp reverse index — it's the more general - // primitive and works whether or not the ingest path was - // able to bind the sid back to a streaming-config policy. - // - // History: an earlier PR removed an `instances_matching` - // fallback under the assumption every production sid - // registration would populate `policy_fp`. The MVP smoke - // test (issue #271 / tracking #272) showed that - // assumption is wrong — sketches arriving from the agent - // carry the full wire-attr set rather than the streaming- - // config's `grouping_labels` subset, so - // `derive_sketch_policy_fp` returns `UNSET` and - // `sids_for_policy(fp)` returns empty. The agg_id-aware - // path is preserved for ExactAgg sids minted via - // `ingest_precompute_for_agg_config` (those carry a - // populated `policy_fp`) but its result is unioned with - // the catalog-walk result so we don't miss the sketches. - let policy_fps = control_plane::asap_tier_analysis::find_matching_policies( - &policy_registry, - candidate, - ); - let mut sids: std::collections::BTreeSet = - std::collections::BTreeSet::new(); - for fp in &policy_fps { - sids.extend(idx.sids_for_policy(*fp)); - } - sids.extend(idx.instances_matching( - &candidate.metric_name, - &candidate.group_by_keys, - )); - if sids.is_empty() { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore has no policy for metric `{}` satisfying \ - capability {:?} — failing over to archive", - candidate.metric_name, candidate.required_capability, - ), - )); - } - - let required: crate::storage_engines::sketch_db::index::Capability = - candidate.required_capability.clone(); - let mut hit_sids: Vec = Vec::with_capacity(sids.len()); - for sid in &sids { - let meta = match idx.instance(*sid) { - Some(m) => m, - None => continue, - }; - if let Some(cap) = meta.capability.as_ref() { - if required.is_satisfied_by(cap) { - hit_sids.push(*sid); - } - } - } - if hit_sids.is_empty() { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore has no sid satisfying capability {:?} for \ - metric `{}` — failing over to archive", - candidate.required_capability, candidate.metric_name - ), - )); - } - - let result = reducer - .evaluate( - &hit_sids, - &candidate.function, - &candidate.function_args, - start_ms, - end_ms, - ) - .map_err(|e| { - crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer failed for `{query}` over \ - [{start_ms}, {end_ms}]: {e:?} — failing over to archive" - ), - ) - })?; - combined_result = Some(result); - } - - let result = combined_result.ok_or_else(|| { - crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!("SketchStore reducer produced no result for `{query}`"), - ) - })?; - - // Matrix shape — the range_query wire format requires it. - Ok(asap_tier_result_to_query_result(result, end_ms, true)) - } - - /// Execute the range query pipeline - fn execute_range_query_pipeline( - &self, - context: &RangeQueryExecutionContext, - ) -> Result, String> { - use crate::query_engines::query_result::RangeVectorElement; - use crate::query_engines::window_merger::create_window_merger; - - // Step 1: Fetch all data needed for the entire range - let all_data = self.execute_store_query(&context.base.store_plan.values_query)?; - - if all_data.is_empty() { - return Err(format!("No data found for metric: {}", context.base.metric)); - } - - debug!( - "Range query: fetched {} keys, {} total buckets", - all_data.len(), - all_data.values().map(|v| v.len()).sum::() - ); - - let mut results: HashMap = HashMap::new(); - - // Determine accumulator type for merger selection - let accumulator_type = &context.base.agg_info.aggregation_type_for_value; - - // Calculate step parameters - let step_ms = context.range_params.step; - let start_ms = context.range_params.start; - let end_ms = context.range_params.end; - let buckets_per_step = context.buckets_per_step; - let lookback_bucket_count = context.lookback_bucket_count; - - let window_mode = if buckets_per_step <= lookback_bucket_count { - "sliding (slide <= size)" - } else { - "hopping (slide > size)" - }; - debug!( - "Range query params: start={}, end={}, step_ms={}, tumbling_window_ms={}, \ - buckets_per_step (slide)={}, lookback_bucket_count (size)={}, mode={}", - start_ms, - end_ms, - step_ms, - context.tumbling_window_ms, - buckets_per_step, - lookback_bucket_count, - window_mode - ); - - // Process each key independently - for (key_opt, timestamped_buckets) in &all_data { - let key = match key_opt { - Some(k) => k.clone(), - None => continue, // Skip None keys for now - }; - - // Build lookup: bucket_start_timestamp -> bucket for O(1) access - let bucket_map: HashMap = timestamped_buckets - .iter() - .map(|((start, _), bucket)| (*start, bucket.as_ref())) - .collect(); - - debug!( - "Key {:?}: built bucket_map with {} entries, timestamps: {:?}", - key, - bucket_map.len(), - bucket_map.keys().collect::>() - ); - - // Create result element for this key - let mut element = RangeVectorElement::new(key.clone()); - - // Calculate window parameters - let tumbling_window_ms = context.tumbling_window_ms; - let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms; - - debug!( - "Key {:?}: range [{}, {}], step={}, lookback_ms={}, tumbling_window_ms={}", - key, start_ms, end_ms, step_ms, lookback_ms, tumbling_window_ms - ); - - // Iterate by OUTPUT timestamp, not by bucket index - let mut current_time = start_ms; - while current_time <= end_ms { - // Window covers [current_time - lookback_ms, current_time) - // This means we look at buckets that START within this range - let window_start = current_time.saturating_sub(lookback_ms); - - // Collect all AVAILABLE buckets in this window (skip missing ones) - let mut window_buckets: Vec> = Vec::new(); - - let mut t = window_start; - while t < current_time { - if let Some(bucket) = bucket_map.get(&t) { - window_buckets.push((*bucket).clone_boxed_core()); - } - // If bucket missing at timestamp t, just skip it (partial data is okay) - t += tumbling_window_ms; - } - - if !window_buckets.is_empty() { - // Merge available buckets - let mut merger = create_window_merger(*accumulator_type); - merger.initialize(window_buckets); - - match merger.get_merged() { - Ok(merged) => { - // Query statistic and emit sample at current_time - match self.query_precompute_for_statistic( - merged.as_ref(), - &context.base.metadata.statistic_to_compute, - &Some(key.clone()), - &context.base.metadata.query_kwargs, - ) { - Ok(value) => { - debug!( - "Key {:?}: emitting sample (t={}, value={})", - key, current_time, value - ); - element.add_sample(current_time, value); - } - Err(e) => { - debug!( - "Failed to query statistic at t={} for key {:?}: {}", - current_time, key, e - ); - } - } - } - Err(e) => { - debug!( - "Failed to get merged result at t={} for key {:?}: {}", - current_time, key, e - ); - } - } - } else { - // No data at all for this window - skip sample - debug!( - "Key {:?}: skipping sample at {} - no data in window [{}, {})", - key, current_time, window_start, current_time - ); - } - - current_time += step_ms; - } - - debug!( - "Key {:?}: finished with {} samples", - key, - element.samples.len() - ); - - // Only include keys with samples - if !element.samples.is_empty() { - results.insert(key, element); - } - } - - // Convert to Vec - Ok(results.into_values().collect()) - } -} - -// --------------------------------------------------------------------------- -// Phase-5: `QueryEngine` trait impl. -// -// Adapter only — does NOT change `handle_query` or any other existing -// surface. The trait's `execute(&str)` walks the same `handle_query` code -// path the binary's HTTP driver uses today; `None` (capability miss) is -// translated to `EngineError::CapabilityMiss` so the router can fall through -// to the next compatible backend. -// --------------------------------------------------------------------------- - -/// Adapt a [`crate::storage_engines::sketch_db::query::ASAPTierResult`] to the engine's -/// existing `QueryResult` shape. The reducer hands back per-series -/// time-stamped scalars; we materialize them as a -/// `QueryResult::Matrix` whose [`crate::query_engines::query_result::RangeVectorElement`]s -/// each map onto one (label-values, samples) entry. -/// -/// `now_ms` is unused for the matrix variant (each sample carries its -/// own window-end timestamp); it's plumbed for future extension to -/// the instant-vector case (latest-pane projection). -/// Merge a ASAP-tier `QueryResult::Matrix` with an archive -/// `QueryResult::Matrix` by `(label_values, timestamp)`. Samples whose -/// timestamps fall inside the warm coverage `(cov_lo, cov_hi)` keep -/// the warm value (warm is approximate but more recent); samples -/// outside that window come from the archive answer. For -/// labels-not-present-in-warm series the archive series is taken in -/// full. Used by `ASAPQueryEngine`'s hybrid-stitch path when the -/// ASAP-tier reducer reports `coverage` narrower than the request. -fn stitch_warm_and_archive( - warm: crate::query_engines::query_result::QueryResult, - archive: crate::query_engines::query_result::QueryResult, - cov_lo: u64, - cov_hi: u64, -) -> crate::query_engines::query_result::QueryResult { - use crate::query_engines::query_result::{QueryResult, RangeVectorElement, Sample}; - use std::collections::BTreeMap; - - let warm_matrix = match &warm { - QueryResult::Matrix(m) => m.values.clone(), - _ => return archive}; - let archive_matrix = match &archive { - QueryResult::Matrix(m) => m.values.clone(), - QueryResult::Vector(_) => return warm}; - - // Index warm series by labels for fast lookup. - let mut by_labels: BTreeMap, RangeVectorElement> = BTreeMap::new(); - for el in warm_matrix { - by_labels.insert(el.labels.labels.clone(), el); - } - - // For each archive series, merge into by_labels. - for arch_el in archive_matrix { - let entry = by_labels - .entry(arch_el.labels.labels.clone()) - .or_insert_with(|| RangeVectorElement::new(arch_el.labels.clone())); - // Build a set of warm timestamps inside coverage (kept). - let warm_ts: std::collections::HashSet = entry - .samples - .iter() - .filter(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi) - .map(|s| s.timestamp) - .collect(); - // Drop any warm samples that ended up outside coverage — - // archive will replace them. - entry - .samples - .retain(|s| s.timestamp >= cov_lo && s.timestamp <= cov_hi); - for s in arch_el.samples { - // Skip archive samples whose timestamps fall inside warm - // coverage AND warm produced a value there (warm wins). - if s.timestamp >= cov_lo && s.timestamp <= cov_hi && warm_ts.contains(&s.timestamp) { - continue; - } - entry.samples.push(Sample::new(s.timestamp, s.value)); - } - entry.samples.sort_by_key(|s| s.timestamp); - } - - let elements: Vec = by_labels.into_values().collect(); - QueryResult::matrix(elements) -} - -fn asap_tier_result_to_query_result( - result: crate::storage_engines::sketch_db::query::ASAPTierResult, - now_ms: u64, - is_range_query: bool, -) -> crate::query_engines::query_result::QueryResult { - use crate::storage_engines::types::KeyByLabelValues; - use crate::query_engines::query_result::{ - InstantVectorElement, QueryResult, RangeVectorElement, - }; - - // Instant-query result-shape: the Prometheus adapter's - // `format_success_response` rejects `Matrix` for queries the - // analyzer marked as instant (`range_seconds == 0`) — produces a - // 500 ”shape mismatch”. Project the per-series last sample into - // an `InstantVectorElement` and wrap as `Vector` so the wire - // response carries `resultType: vector` matching the request. - if !is_range_query { - let mut elements: Vec = Vec::with_capacity(result.series.len()); - for (label_values, samples) in result.series { - // Mirror the range-vector branch: BTreeMap iteration is - // key-sorted, so `unzip` produces aligned (keys, values). - // Stash the keys in the per-element `label_keys_override` - // so the Prometheus adapter renders synthesized keys - // (notably ASAP-tier `topk`'s `"item"` key) instead of - // the empty `metric: {}` it would produce when the - // query-scoped `KeyByLabelNames` is empty. - let (keys, values): (Vec, Vec) = label_values.into_iter().unzip(); - let labels = KeyByLabelValues::new_with_labels(values); - // Take the latest sample (the reducer returns one per - // window_end; for instant readout we want the most recent). - if let Some((_, value)) = samples.into_iter().last() { - elements.push( - InstantVectorElement::new(labels, value) - .with_label_keys_override(keys), - ); - } - } - return QueryResult::vector(elements, now_ms); - } - - let mut elements: Vec = Vec::with_capacity(result.series.len()); - for (label_values, samples) in result.series { - // `KeyByLabelValues` is a `Vec` carrying VALUES only; - // the serializer pairs them with KEYS from a query-scoped - // `KeyByLabelNames`. For most queries the keys ARE the - // query's group-by clause, so the default path works. But - // ASAP-tier `topk` synthesizes an `"item"` key (the top-k - // entry name) that the original query's group-by doesn't - // carry — without an override the serializer drops it and - // the response shows `"metric": {}`. Project the BTreeMap's - // VALUES in key-sorted order (BTreeMap iteration is - // key-sorted), and stash the BTreeMap's KEYS in the - // per-element override so the serializer can pair them - // correctly. - let (keys, values): (Vec, Vec) = label_values.into_iter().unzip(); - let labels = KeyByLabelValues::new_with_labels(values); - let mut element = RangeVectorElement::new(labels).with_label_keys_override(keys); - for (window_end_ms, value) in samples { - // `window_end_ms` is i64 from the index; cast to u64 - // for the wire format (window_end is monotonic + post- - // 1970 in production). - let ts = if window_end_ms >= 0 { - window_end_ms as u64 - } else { - 0 - }; - element.add_sample(ts, value); - } - elements.push(element); - } - QueryResult::matrix(elements) -} - -#[async_trait::async_trait] -impl crate::query_engines::routing::query_engine_routing::QueryEngine for ASAPQueryEngine { - async fn execute( - &self, - query: &str, - ) -> Result { - // Phase 9 controller-unification (2026-05) — the ASAP-tier - // hook is now a thin driver around the control plane's - // `analyze_promql_for_asap_tier`. The analyzer is the single - // owner of "is this PromQL ASAP-tier-answerable" knowledge. - // We drop into one of three branches: - // - // 1. `ASAPTierAnalysis::unsupported` is `Some(_)` — the - // PromQL shape isn't ASAP-tier-servable. Surface as - // `EngineError::CapabilityMiss(SketchStore, …)` with the - // structured `UnsupportedReason` in the detail string. The - // EngineRouter fails over to the archive engine. This - // covers all of: - // * `MissReason::UnsupportedFunction(_)` (rate, irate, - // increase, etc.) → cold tier (archive) - // * `MissReason::UnsupportedComposition(_)` (sum-by, - // topk-over-rate, etc.) → cold tier - // * `MissReason::NoCallNodeFound` (bare selector) → - // cold tier (archive answers raw selectors) - // * `MissReason::UnparseablePromql(_)` → cold tier - // (archive's parser may be more permissive, or it'll - // also reject and the user sees the error) - // - // 2. `ASAPTierAnalysis::candidates` is populated, but ANY - // candidate's `instances_matching` returns empty OR a - // sid that classifies as `Ghost`/`Unknown` — surface - // as CapabilityMiss. The EngineRouter falls over. - // - // 3. All candidates resolve to all-`Hit` sids — dispatch - // each to the per-`Capability` sketch reducer. Today's - // semantic: ANY candidate-level reducer error → fall - // over to archive (no per-candidate hybrid stitch yet — - // that's the documented follow-up). - if let Some(idx) = self.sketch_index.as_ref() { - let analysis = control_plane::asap_tier_analysis::analyze_promql_for_asap_tier(query); - - // Branch 1 — the control plane analyzer rejects the shape. - if let Some(reason) = &analysis.unsupported { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore analyzer rejected `{query}`: {reason:?} — \ - failing over to archive" - ), - )); - } - if analysis.candidates.is_empty() { - // Defensive — `is_asap_tier_answerable` would have - // caught this; analyzer guarantees `unsupported.is_some()` - // when `candidates.is_empty()` but we keep the - // belt-and-braces miss-path for safety. - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore analyzer produced no ASAP-tier candidates for \ - `{query}` — failing over to archive" - ), - )); - } - - // Branch 2 + 3 — resolve each candidate's sids and - // dispatch the reducer. Today this is single-candidate - // for every supported PromQL shape; the loop is here - // for the per-candidate hybrid-stitch follow-up. - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - // Time bounds: the trait's `execute(&str)` adapter - // doesn't carry an explicit range today (it's an - // instant-query surface). For each candidate, prefer the - // candidate's `range_seconds` (extracted from `[5m]` / - // `[30s]` selectors); fall back to a 5-minute default - // for instant-vector candidates (range_seconds == 0). - const DEFAULT_LOOKBACK_MS: u64 = 5 * 60 * 1000; - - let reducer = crate::storage_engines::sketch_db::query::SketchReducer::new(idx); - // Multi-candidate aggregation is deferred (single-result - // shapes today). On the first reducer error we surface - // CapabilityMiss; on Ok we keep the result for the - // hybrid-stitch path below. (When more than one - // candidate is supported, a follow-up will fold - // per-candidate ASAPTierResults.) - let mut combined_result: Option = - None; - let mut combined_t0: u64 = u64::MAX; - // Track whether ANY candidate is range-vector-shaped - // (`range_seconds > 0`). Drives the Vector-vs-Matrix - // result-shape choice in `asap_tier_result_to_query_result` - // below — instant queries (`count(metric)`, - // `quantile(...)` without `_over_time` etc.) need - // `QueryResult::Vector` so the Prometheus adapter's - // `format_success_response` wraps them as `resultType: - // vector`. Returning `Matrix` for an instant query - // produces a 500 (adapter rejects the shape mismatch). - let mut any_range_candidate = false; - - // Snapshot the streaming config once for this query's - // policy lookups. Hot-reload swaps the underlying Arc; the - // snapshot pins one revision for the duration. - let streaming_snap = self.streaming_config_snapshot(); - let policy_registry = streaming_snap.policy_registry(); - - for candidate in &analysis.candidates { - if candidate.range_seconds > 0 { - any_range_candidate = true; - } - // Schema-retirement #5: resolve candidate → {sids} by - // unioning the policy-fp reverse index (fast path for - // ExactAgg sids minted via `ingest_precompute_for_agg_config` - // where `policy_fp` is set) with `instances_matching` - // (catalog walk that subset-matches on - // `group_by_keys`, covering raw sketches whose - // `derive_sketch_policy_fp` returned `UNSET` because - // the wire-attr set didn't match any streaming-config - // policy). The earlier policy-fp-only path returned - // empty for the MVP demo workload — see issue #271 / - // tracking #272. - let policy_fps = control_plane::asap_tier_analysis::find_matching_policies( - &policy_registry, - candidate, - ); - let mut sids: std::collections::BTreeSet = - std::collections::BTreeSet::new(); - for fp in &policy_fps { - sids.extend(idx.sids_for_policy(*fp)); - } - sids.extend(idx.instances_matching( - &candidate.metric_name, - &candidate.group_by_keys, - )); - if sids.is_empty() { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore has no policy for metric `{}` \ - with group_by_keys ⊇ {:?} satisfying capability \ - {:?} — failing over to archive", - candidate.metric_name, - candidate.group_by_keys, - candidate.required_capability, - ), - )); - } - - // Verify each sid carries the analyzer's required - // capability. After Step 2a there's exactly one - // `Capability` enum (defined in the control plane and - // re-exported by `sketch_index`), so no `From` - // conversion is needed — just clone. - let required: crate::storage_engines::sketch_db::index::Capability = - candidate.required_capability.clone(); - let mut hit_sids: Vec = Vec::with_capacity(sids.len()); - for sid in &sids { - match idx.classify(*sid) { - crate::storage_engines::sketch_db::index::SidLookup::Hit => {} - crate::storage_engines::sketch_db::index::SidLookup::Ghost - | crate::storage_engines::sketch_db::index::SidLookup::Unknown => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore ghost/unknown sid {sid} for metric \ - `{}` — failing over to archive", - candidate.metric_name - ), - )); - } - } - let meta = match idx.instance(*sid) { - Some(m) => m, - None => continue}; - // Precompute-backed sids (M2.3) have `capability: None` - // — the analyzer doesn't route them through this path, - // but skip defensively if one slips in. - if let Some(cap) = meta.capability.as_ref() { - if required.is_satisfied_by(cap) { - hit_sids.push(*sid); - } - } - } - if hit_sids.is_empty() { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore has no sid satisfying capability \ - {:?} for metric `{}` — failing over to archive", - candidate.required_capability, candidate.metric_name - ), - )); - } - - let lookback_ms = if candidate.range_seconds > 0 { - candidate.range_seconds.saturating_mul(1000) - } else { - DEFAULT_LOOKBACK_MS - }; - let t0_ms = now_ms.saturating_sub(lookback_ms); - if t0_ms < combined_t0 { - combined_t0 = t0_ms; - } - - let result = match reducer.evaluate( - &hit_sids, - &candidate.function, - &candidate.function_args, - t0_ms, - now_ms, - ) { - Ok(r) => r, - Err( - crate::storage_engines::sketch_db::query::ASAPTierError::UnsupportedFunction( - name, - ), - ) => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer does not support function `{name}` \ - — failing over to archive" - ), - )); - } - Err(crate::storage_engines::sketch_db::query::ASAPTierError::UnsupportedCapability { - function, - capability}) => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer cannot answer `{function}` against \ - capability {capability:?} — failing over to archive" - ), - )); - } - Err(crate::storage_engines::sketch_db::query::ASAPTierError::DeserializeFailure { - sid, - encoding, - reason}) => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer failed to decode sketch for sid \ - {sid} (encoding={encoding:?}): {reason} — failing over \ - to archive" - ), - )); - } - Err(crate::storage_engines::sketch_db::query::ASAPTierError::NoData { - metric_name: m}) => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer found no samples for metric \ - `{m}` in window — failing over to archive" - ), - )); - } - Err(crate::storage_engines::sketch_db::query::ASAPTierError::MissingHeap { - sid, - sketch_kind}) => { - return Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!( - "SketchStore reducer cannot enumerate top-k for sid \ - {sid} (sketch_kind={sketch_kind:?}, no heap) — \ - failing over to archive" - ), - )); - } - }; - combined_result = Some(result); - } - - // All candidates resolved successfully — adapt to - // QueryResult and run the hybrid-stitch path if archive - // is wired and warm coverage is narrower than request. - if let Some(result) = combined_result { - // `execute(&str)` is the instant-query trait surface — - // it's only called from `/api/v1/query` (never from - // `/api/v1/query_range`, which has its own - // `handle_range_query_promql` path). For PromQL, - // instant queries always return a vector: even when - // the inner expression carries a range selector like - // `count_over_time(metric[10s])`, the outer evaluation - // at time `t` yields one value per series (computed - // over `[t-range, t]`). So this site always wants - // Vector — `any_range_candidate` was the wrong signal - // (it captures the inner range, not the outer eval - // shape) and produced Matrix for instant queries - // with range-bound inners, which the Prometheus - // adapter's `format_success_response` rejects with - // a 500 ”shape mismatch” / empty-body response. - let _ = any_range_candidate; - let warm_qr = asap_tier_result_to_query_result( - result.clone(), - now_ms, - false, - ); - if let (Some((cov_lo, cov_hi)), Some(archive)) = - (result.coverage, self.archive_engine.as_ref()) - { - let stitch_t0 = if combined_t0 == u64::MAX { - now_ms.saturating_sub(DEFAULT_LOOKBACK_MS) - } else { - combined_t0 - }; - if cov_lo > stitch_t0 || cov_hi < now_ms { - let archive_qr = archive.execute(query).await; - if let Ok(archive_qr) = archive_qr { - return Ok(stitch_warm_and_archive( - warm_qr, archive_qr, cov_lo, cov_hi, - )); - } - // On archive error, fall back to warm-only. - } - } - return Ok(warm_qr); - } - } - - // `handle_query` is sync + needs a `time: f64` (epoch millis as float). - // The router doesn't pass a query time, so we use wall-clock now — - // matches `GorillaQueryEngine::execute`'s convention. - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH) - .map(|d| d.as_millis() as f64) - .unwrap_or(0.0); - match self.handle_query(query.to_string(), now_ms) { - Some((_labels, result)) => Ok(result), - None => Err(crate::query_engines::EngineError::capability_miss( - asap_types::StorageBackend::SketchStore.data_source_id(), - format!("ASAPQueryEngine has no compatible aggregation for `{query}`"), - ))} - } - - fn capabilities(&self) -> crate::query_engines::routing::query_engine_routing::EngineCapabilities { - crate::query_engines::routing::query_engine_routing::EngineCapabilities { - data_source_id: asap_types::StorageBackend::SketchStore.data_source_id(), - storage_backend: asap_types::StorageBackend::SketchStore, - // Warm-tier sketches are O(sketch-size); call it 16 MiB ceiling - // for buffered ops (KLL with k=200 is well below this). - supports_streams_above_bytes: 16 * 1024 * 1024} - } -} - -#[cfg(test)] -mod range_query_tests { - use crate::storage_engines::types::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; - use crate::query_engines::window_merger::NaiveMerger; - use serde_json::Value; - use std::any::Any; - - /// Mock accumulator that stores a unique ID to detect stale window reuse - #[derive(Clone, Debug)] - struct MockBucketAccumulator { - bucket_id: u64, - value: f64} - - impl MockBucketAccumulator { - fn new(bucket_id: u64, value: f64) -> Self { - Self { bucket_id, value } - } - } - - impl SerializableToSink for MockBucketAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({"bucket_id": self.bucket_id, "value": self.value}) - } - - fn serialize_to_bytes(&self) -> Vec { - format!("{}:{}", self.bucket_id, self.value).into_bytes() - } - } - - impl AggregateCore for MockBucketAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "MockBucketAccumulator" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - - fn merge_with( - &self, - other: &dyn AggregateCore, - ) -> Result, Box> { - if let Some(other_mock) = other.as_any().downcast_ref::() { - // Sum values, keep max bucket_id to track which buckets are in window - Ok(Box::new(MockBucketAccumulator::new( - self.bucket_id.max(other_mock.bucket_id), - self.value + other_mock.value, - ))) - } else { - Err("Cannot merge with different accumulator type".into()) - } - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::Sum - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - _statistic: promql_utilities::query_logics::enums::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - Err("MockBucketAccumulator does not support query_statistic".into()) - } - } - - /// Simulates the sliding window loop from execute_range_query_pipeline - /// Returns: Vec of (timestamp, merged_value, max_bucket_id_in_window) - fn simulate_sliding_window( - buckets: Vec>, - lookback_bucket_count: usize, - buckets_per_step: usize, - start_ms: u64, - end_ms: u64, - step_ms: u64, - ) -> Vec<(u64, f64, u64)> { - use crate::query_engines::window_merger::WindowMerger; - - let mut results = Vec::new(); - - if buckets.len() < lookback_bucket_count { - return results; - } - - let mut merger = NaiveMerger::new(); - - // Initialize with first window - let initial_window: Vec<_> = buckets[0..lookback_bucket_count] - .iter() - .map(|b| b.clone_boxed_core()) - .collect(); - merger.initialize(initial_window); - - let mut bucket_index = lookback_bucket_count; - let mut current_time = start_ms; - - while current_time <= end_ms { - // Query current window - if let Ok(merged) = merger.get_merged() { - if let Some(mock) = merged.as_any().downcast_ref::() { - results.push((current_time, mock.value, mock.bucket_id)); - } - } - - // Slide window for next step - current_time += step_ms; - - if current_time <= end_ms { - if bucket_index + buckets_per_step <= buckets.len() { - let new_buckets: Vec<_> = buckets - [bucket_index..bucket_index + buckets_per_step] - .iter() - .map(|b| b.clone_boxed_core()) - .collect(); - merger.slide(buckets_per_step, new_buckets); - bucket_index += buckets_per_step; - } else { - // Not enough buckets to continue - stop to avoid stale data - break; - } - } - } - - results - } - - /// Simulates sliding window with proper timestamp alignment for missing data. - /// This accounts for the scenario where the store returns fewer buckets than - /// expected because data is missing at the start of the query range. - /// - /// # Arguments - /// * `expected_bucket_count` - How many buckets we would have if data was complete - fn simulate_sliding_window_with_alignment( - buckets: Vec>, - lookback_bucket_count: usize, - buckets_per_step: usize, - start_ms: u64, - end_ms: u64, - step_ms: u64, - expected_bucket_count: usize, - ) -> Vec<(u64, f64, u64)> { - use crate::query_engines::window_merger::WindowMerger; - - let mut results = Vec::new(); - - // Check if we have enough buckets for at least one window - if buckets.len() < lookback_bucket_count { - return results; - } - - // Calculate missing data offset - let missing_buckets = expected_bucket_count.saturating_sub(buckets.len()); - let tumbling_window_ms = step_ms / (buckets_per_step as u64); - - // First valid sample is offset by missing buckets (data missing at the start) - let first_valid_sample_ms = start_ms + (missing_buckets as u64) * tumbling_window_ms; - - // Round up to step boundary if needed - let first_sample_ms = if first_valid_sample_ms <= start_ms { - start_ms - } else { - let offset = first_valid_sample_ms - start_ms; - if offset.is_multiple_of(step_ms) { - first_valid_sample_ms - } else { - start_ms + ((offset / step_ms) + 1) * step_ms - } - }; - - // When we have missing buckets at the start, we need to figure out where to - // start reading from the available buckets. The missing buckets are conceptually - // at the beginning, so we start reading from the first available bucket. - // - // However, if we rounded up to a step boundary, we may need to skip some - // additional buckets from what we have. - let extra_offset_ms = first_sample_ms.saturating_sub(first_valid_sample_ms); - let extra_buckets_to_skip = (extra_offset_ms / tumbling_window_ms) as usize; - - // Check if we have enough data for at least one window after any extra skip - if extra_buckets_to_skip + lookback_bucket_count > buckets.len() { - return results; - } - - let mut merger = NaiveMerger::new(); - - // Initialize with window at adjusted position - let initial_window: Vec<_> = buckets - [extra_buckets_to_skip..extra_buckets_to_skip + lookback_bucket_count] - .iter() - .map(|b| b.clone_boxed_core()) - .collect(); - merger.initialize(initial_window); - - let mut bucket_index = extra_buckets_to_skip + lookback_bucket_count; - let mut current_time = first_sample_ms; - - while current_time <= end_ms { - // Query current window - if let Ok(merged) = merger.get_merged() { - if let Some(mock) = merged.as_any().downcast_ref::() { - results.push((current_time, mock.value, mock.bucket_id)); - } - } - - // Slide window for next step - current_time += step_ms; - - if current_time <= end_ms { - if bucket_index + buckets_per_step <= buckets.len() { - let new_buckets: Vec<_> = buckets - [bucket_index..bucket_index + buckets_per_step] - .iter() - .map(|b| b.clone_boxed_core()) - .collect(); - merger.slide(buckets_per_step, new_buckets); - bucket_index += buckets_per_step; - } else { - break; - } - } - } - - results - } - - #[test] - fn test_sliding_window_sufficient_buckets() { - // Setup: 7 buckets, lookback=5, step=1 - // Should produce 3 valid samples - let buckets: Vec> = (0..7) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window( - buckets, 5, // lookback_bucket_count - 1, // buckets_per_step - 1000, // start_ms - 3000, // end_ms (3 steps: 1000, 2000, 3000) - 1000, // step_ms - ); - - assert_eq!(results.len(), 3, "Should produce 3 samples"); - - // Window 1: buckets [0,1,2,3,4], max_id=4, value=50 - assert_eq!(results[0], (1000, 50.0, 4)); - // Window 2: buckets [1,2,3,4,5], max_id=5, value=50 - assert_eq!(results[1], (2000, 50.0, 5)); - // Window 3: buckets [2,3,4,5,6], max_id=6, value=50 - assert_eq!(results[2], (3000, 50.0, 6)); - } - - #[test] - fn test_sliding_window_insufficient_buckets_stops_early() { - // 6 buckets, lookback=5, step=1 - // Requesting 3 timestamps but only have data for 2 - // Should stop early rather than produce stale samples - let buckets: Vec> = (0..6) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window( - buckets, 5, // lookback_bucket_count - 1, // buckets_per_step - 1000, // start_ms - 3000, // end_ms (requests 3 steps: 1000, 2000, 3000) - 1000, // step_ms - ); - - println!("Results: {:?}", results); - - // Should only produce 2 valid samples (not 3 with stale data) - assert_eq!( - results.len(), - 2, - "Should only produce 2 samples when data is insufficient for 3rd" - ); - - // Window 1: buckets [0,1,2,3,4], max_id=4 - assert_eq!(results[0], (1000, 50.0, 4)); - // Window 2: buckets [1,2,3,4,5], max_id=5 - assert_eq!(results[1], (2000, 50.0, 5)); - // No window 3 - not enough buckets to slide - } - - #[test] - fn test_sliding_window_exactly_enough_buckets() { - // 5 buckets, lookback=5, step=1 - // Should produce exactly 1 sample (initial window only, can't slide) - let buckets: Vec> = (0..5) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window( - buckets, 5, // lookback_bucket_count - 1, // buckets_per_step - 1000, // start_ms - 3000, // end_ms - 1000, // step_ms - ); - - println!("Results with exactly enough buckets: {:?}", results); - - // Should produce only 1 sample - can't slide without more buckets - assert_eq!(results.len(), 1, "Should produce exactly 1 sample"); - assert_eq!(results[0], (1000, 50.0, 4)); - } - - #[test] - fn test_sliding_window_multi_bucket_step() { - // 10 buckets, lookback=4, step=2 buckets at a time - // Should produce samples at positions requiring new data - let buckets: Vec> = (0..10) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window( - buckets, 4, // lookback_bucket_count - 2, // buckets_per_step (slide 2 at a time) - 1000, // start_ms - 4000, // end_ms (4 steps) - 1000, // step_ms - ); - - // Initial: [0,1,2,3], max_id=3 - // After slide 1: [2,3,4,5], max_id=5 - // After slide 2: [4,5,6,7], max_id=7 - // After slide 3: [6,7,8,9], max_id=9 - assert_eq!(results.len(), 4, "Should produce 4 samples"); - assert_eq!(results[0].2, 3, "Window 1 max_id should be 3"); - assert_eq!(results[1].2, 5, "Window 2 max_id should be 5"); - assert_eq!(results[2].2, 7, "Window 3 max_id should be 7"); - assert_eq!(results[3].2, 9, "Window 4 max_id should be 9"); - } - - #[test] - fn test_sliding_window_missing_data_at_start_aligns_timestamps() { - // Scenario: Query requests timestamps 1000, 2000, 3000 - // But only 5 buckets exist (enough for 1 sample), not 7 (for 3 samples) - // lookback=5, step=1 bucket - // Expected buckets for [1000, 3000]: 7 (5 for first window + 2 steps) - // Actual buckets: 5 (missing 2 at start) - // Missing 2 buckets = 2000ms offset - // First valid sample at: 1000 + 2000 = 3000ms - - let buckets: Vec> = (0..5) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window_with_alignment( - buckets, 5, // lookback_bucket_count - 1, // buckets_per_step - 1000, // start_ms - 3000, // end_ms - 1000, // step_ms - 7, // expected_bucket_count for full range - ); - - // Should have 1 sample at timestamp 3000, NOT at 1000 - assert_eq!(results.len(), 1, "Should produce 1 sample"); - assert_eq!(results[0].0, 3000, "Sample should be at t=3000, not t=1000"); - } - - #[test] - fn test_sliding_window_missing_data_rounds_to_step_boundary() { - // Query: start=0, end=6000, step=2000 (timestamps: 0, 2000, 4000, 6000) - // Lookback: 4 buckets, step: 2 buckets - // Expected buckets: 4 + 6 = 10 buckets for full range - // Actual: 7 buckets (missing 3 at start) - // Missing 3 buckets = 3000ms offset - // First valid sample time = 0 + 3000 = 3000ms - // But 3000 is not on step boundary, so round UP to 4000ms - - let buckets: Vec> = (0..7) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window_with_alignment( - buckets, 4, // lookback_bucket_count - 2, // buckets_per_step (2000ms step / 1000ms tumbling = 2) - 0, // start_ms - 6000, // end_ms - 2000, // step_ms - 10, // expected_bucket_count - ); - - // First sample at 4000 (rounded up from 3000), second at 6000 - assert_eq!(results.len(), 2, "Should produce 2 samples"); - assert_eq!(results[0].0, 4000, "First sample at step boundary 4000"); - assert_eq!(results[1].0, 6000, "Second sample at 6000"); - } - - #[test] - fn test_sliding_window_full_data_starts_at_query_start() { - // All data present - should behave same as before (start at start_ms) - // lookback=5, step=1, query [1000, 3000] = 3 samples - // Expected buckets: 7, Actual: 7 (no missing data) - - let buckets: Vec> = (0..7) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window_with_alignment( - buckets, 5, // lookback_bucket_count - 1, // buckets_per_step - 1000, // start_ms - 3000, // end_ms - 1000, // step_ms - 7, // expected_bucket_count (matches actual - no missing data) - ); - - assert_eq!(results.len(), 3, "Should produce 3 samples"); - assert_eq!(results[0].0, 1000, "First sample at query start"); - assert_eq!(results[1].0, 2000); - assert_eq!(results[2].0, 3000); - } - - #[test] - fn test_sliding_window_insufficient_data_for_any_window_returns_empty() { - // lookback=5 but only 3 buckets - can't form even one window - let buckets: Vec> = (0..3) - .map(|i| Box::new(MockBucketAccumulator::new(i, 10.0)) as Box) - .collect(); - - let results = simulate_sliding_window_with_alignment( - buckets, 5, // lookback_bucket_count (need 5, have 3) - 1, 1000, 5000, 1000, 9, - ); - - assert_eq!( - results.len(), - 0, - "No samples when insufficient data for any window" - ); - } - - // ============================================================================ - // Tests for timestamp-based lookup implementation (handles gaps in data) - // ============================================================================ - - /// Simulates the timestamp-based lookup approach from execute_range_query_pipeline. - /// This is the new implementation that handles gaps in data correctly. - /// - /// # Arguments - /// * `timestamped_buckets` - Vec of (bucket_start_timestamp, bucket) - /// * `lookback_bucket_count` - Number of buckets in each window - /// * `tumbling_window_ms` - Duration of each tumbling window bucket - /// * `start_ms` - Query start time - /// * `end_ms` - Query end time - /// * `step_ms` - Step between output samples - /// - /// # Returns - /// Vec of (timestamp, merged_value, max_bucket_id_in_window) - fn simulate_timestamp_based_lookup( - timestamped_buckets: Vec<(u64, Box)>, - lookback_bucket_count: usize, - tumbling_window_ms: u64, - start_ms: u64, - end_ms: u64, - step_ms: u64, - ) -> Vec<(u64, f64, u64)> { - use crate::query_engines::window_merger::WindowMerger; - use std::collections::HashMap; - - let mut results = Vec::new(); - - // Build lookup: bucket_start_timestamp -> bucket for O(1) access - let bucket_map: HashMap> = timestamped_buckets - .iter() - .map(|(start, bucket)| (*start, bucket)) - .collect(); - - let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms; - - // Iterate by OUTPUT timestamp, not by bucket index - let mut current_time = start_ms; - while current_time <= end_ms { - // Window covers [current_time - lookback_ms, current_time) - let window_start = current_time.saturating_sub(lookback_ms); - - // Collect all AVAILABLE buckets in this window (skip missing ones) - let mut window_buckets: Vec> = Vec::new(); - - let mut t = window_start; - while t < current_time { - if let Some(bucket) = bucket_map.get(&t) { - window_buckets.push((*bucket).clone_boxed_core()); - } - t += tumbling_window_ms; + let result = match reducer.evaluate( + &hit_sids, + &candidate.function, + &candidate.function_args, + t0_ms, + now_ms, + ) { + Ok(r) => r, + Err( + crate::storage_engines::sketch_db::query::ASAPTierError::UnsupportedFunction( + name, + ), + ) => { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer does not support function `{name}` \ + — failing over to archive" + ), + )); + } + Err(crate::storage_engines::sketch_db::query::ASAPTierError::UnsupportedCapability { + function, + capability}) => { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer cannot answer `{function}` against \ + capability {capability:?} — failing over to archive" + ), + )); + } + Err(crate::storage_engines::sketch_db::query::ASAPTierError::DeserializeFailure { + sid, + encoding, + reason}) => { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer failed to decode sketch for sid \ + {sid} (encoding={encoding:?}): {reason} — failing over \ + to archive" + ), + )); + } + Err(crate::storage_engines::sketch_db::query::ASAPTierError::NoData { + metric_name: m}) => { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer found no samples for metric \ + `{m}` in window — failing over to archive" + ), + )); + } + Err(crate::storage_engines::sketch_db::query::ASAPTierError::MissingHeap { + sid, + sketch_kind}) => { + return Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "SketchStore reducer cannot enumerate top-k for sid \ + {sid} (sketch_kind={sketch_kind:?}, no heap) — \ + failing over to archive" + ), + )); + } + }; + combined_result = Some(result); } - if !window_buckets.is_empty() { - // Merge available buckets - let mut merger = NaiveMerger::new(); - merger.initialize(window_buckets); - - if let Ok(merged) = merger.get_merged() { - if let Some(mock) = merged.as_any().downcast_ref::() { - results.push((current_time, mock.value, mock.bucket_id)); + // All candidates resolved successfully — adapt to + // QueryResult and run the hybrid-stitch path if archive + // is wired and warm coverage is narrower than request. + if let Some(result) = combined_result { + // `execute(&str)` is the instant-query trait surface — + // it's only called from `/api/v1/query` (never from + // `/api/v1/query_range`, which has its own + // `handle_range_query_promql` path). For PromQL, + // instant queries always return a vector: even when + // the inner expression carries a range selector like + // `count_over_time(metric[10s])`, the outer evaluation + // at time `t` yields one value per series (computed + // over `[t-range, t]`). So this site always wants + // Vector — `any_range_candidate` was the wrong signal + // (it captures the inner range, not the outer eval + // shape) and produced Matrix for instant queries + // with range-bound inners, which the Prometheus + // adapter's `format_success_response` rejects with + // a 500 ”shape mismatch” / empty-body response. + let _ = any_range_candidate; + let warm_qr = asap_tier_result_to_query_result( + result.clone(), + now_ms, + false, + ); + if let (Some((cov_lo, cov_hi)), Some(archive)) = + (result.coverage, self.archive_engine.as_ref()) + { + let stitch_t0 = if combined_t0 == u64::MAX { + now_ms.saturating_sub(DEFAULT_LOOKBACK_MS) + } else { + combined_t0 + }; + if cov_lo > stitch_t0 || cov_hi < now_ms { + let archive_qr = archive.execute(query).await; + if let Ok(archive_qr) = archive_qr { + return Ok(stitch_warm_and_archive( + warm_qr, archive_qr, cov_lo, cov_hi, + )); + } + // On archive error, fall back to warm-only. } } - } - // If no buckets available, skip this sample (no entry in results) - - current_time += step_ms; - } - - results - } - - #[test] - fn test_timestamp_lookup_missing_data_at_start() { - // Scenario: Query range [1000, 5000] with step=1000, lookback=3 buckets - // Tumbling window = 1000ms - // Expected buckets for full window coverage starting at t=1000: - // - t=1000 needs buckets at -2000, -1000, 0 (before query range) - // But data only exists at t=3000, 4000, 5000 - // - // Sample at t=1000: window [1000-3000, 1000) = [-2000, 1000) -> no buckets -> skip - // Sample at t=2000: window [2000-3000, 2000) = [-1000, 2000) -> no buckets -> skip - // Sample at t=3000: window [3000-3000, 3000) = [0, 3000) -> no buckets -> skip - // Sample at t=4000: window [4000-3000, 4000) = [1000, 4000) -> bucket at 3000 -> emit - // Sample at t=5000: window [5000-3000, 5000) = [2000, 5000) -> buckets at 3000, 4000 -> emit - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (3000, Box::new(MockBucketAccumulator::new(3, 10.0))), - (4000, Box::new(MockBucketAccumulator::new(4, 10.0))), - (5000, Box::new(MockBucketAccumulator::new(5, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 1000, // start_ms - 5000, // end_ms - 1000, // step_ms - ); - - // Should skip samples at 1000, 2000, 3000 (no data in window) - // Should emit samples at 4000 (partial data) and 5000 (partial data) - assert_eq!( - results.len(), - 2, - "Should produce 2 samples (skipping early ones with no data)" - ); - assert_eq!(results[0].0, 4000, "First sample at t=4000"); - assert_eq!(results[0].1, 10.0, "Value at t=4000 (1 bucket)"); - assert_eq!(results[1].0, 5000, "Second sample at t=5000"); - assert_eq!(results[1].1, 20.0, "Value at t=5000 (2 buckets merged)"); - } - - #[test] - fn test_timestamp_lookup_missing_data_in_middle() { - // Scenario: Buckets at t=1000, 2000, 4000, 5000 (missing t=3000) - // Query range [4000, 6000], step=1000, lookback=3 buckets - // Tumbling window = 1000ms - // - // Sample at t=4000: window [1000, 4000) -> buckets at 1000, 2000 (missing 3000) -> 2 buckets - // Sample at t=5000: window [2000, 5000) -> buckets at 2000, 4000 (missing 3000) -> 2 buckets - // Sample at t=6000: window [3000, 6000) -> buckets at 4000, 5000 (missing 3000) -> 2 buckets - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (1000, Box::new(MockBucketAccumulator::new(1, 10.0))), - (2000, Box::new(MockBucketAccumulator::new(2, 10.0))), - // Missing bucket at 3000 - (4000, Box::new(MockBucketAccumulator::new(4, 10.0))), - (5000, Box::new(MockBucketAccumulator::new(5, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 4000, // start_ms - 6000, // end_ms - 1000, // step_ms - ); - - // All samples should be emitted with partial data (missing bucket is skipped) - assert_eq!( - results.len(), - 3, - "Should produce 3 samples with partial data" - ); - - // t=4000: window [1000, 4000) contains buckets 1000, 2000 -> value=20, max_id=2 - assert_eq!(results[0].0, 4000); - assert_eq!(results[0].1, 20.0, "2 buckets merged"); - assert_eq!(results[0].2, 2, "max bucket_id = 2"); - - // t=5000: window [2000, 5000) contains buckets 2000, 4000 -> value=20, max_id=4 - assert_eq!(results[1].0, 5000); - assert_eq!(results[1].1, 20.0, "2 buckets merged"); - assert_eq!(results[1].2, 4, "max bucket_id = 4"); - - // t=6000: window [3000, 6000) contains buckets 4000, 5000 -> value=20, max_id=5 - assert_eq!(results[2].0, 6000); - assert_eq!(results[2].1, 20.0, "2 buckets merged"); - assert_eq!(results[2].2, 5, "max bucket_id = 5"); - } - - #[test] - fn test_timestamp_lookup_all_data_missing_for_window() { - // Scenario: Query window where no buckets exist at all - // Buckets at t=10000, 11000, 12000 - // Query range [1000, 3000], step=1000, lookback=3 buckets - // All windows have no data -> should skip all samples - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (10000, Box::new(MockBucketAccumulator::new(10, 10.0))), - (11000, Box::new(MockBucketAccumulator::new(11, 10.0))), - (12000, Box::new(MockBucketAccumulator::new(12, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 1000, // start_ms - 3000, // end_ms - 1000, // step_ms - ); - - assert_eq!( - results.len(), - 0, - "Should produce 0 samples when all windows have no data" - ); - } - - #[test] - fn test_timestamp_lookup_full_data_matches_expected() { - // Scenario: Full data available, should behave like contiguous case - // Buckets at t=0, 1000, 2000, 3000, 4000 - // Query range [3000, 5000], step=1000, lookback=3 buckets - // - // Sample at t=3000: window [0, 3000) -> buckets 0, 1000, 2000 -> value=30 - // Sample at t=4000: window [1000, 4000) -> buckets 1000, 2000, 3000 -> value=30 - // Sample at t=5000: window [2000, 5000) -> buckets 2000, 3000, 4000 -> value=30 - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (0, Box::new(MockBucketAccumulator::new(0, 10.0))), - (1000, Box::new(MockBucketAccumulator::new(1, 10.0))), - (2000, Box::new(MockBucketAccumulator::new(2, 10.0))), - (3000, Box::new(MockBucketAccumulator::new(3, 10.0))), - (4000, Box::new(MockBucketAccumulator::new(4, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 3000, // start_ms - 5000, // end_ms - 1000, // step_ms - ); - - assert_eq!(results.len(), 3, "Should produce 3 samples"); - - assert_eq!(results[0], (3000, 30.0, 2), "t=3000: buckets 0,1,2"); - assert_eq!(results[1], (4000, 30.0, 3), "t=4000: buckets 1,2,3"); - assert_eq!(results[2], (5000, 30.0, 4), "t=5000: buckets 2,3,4"); - } - - #[test] - fn test_timestamp_lookup_sparse_data() { - // Scenario: Very sparse data - only every 3rd bucket exists - // Buckets at t=0, 3000, 6000, 9000 - // Query range [3000, 9000], step=3000, lookback=3 buckets (3000ms) - // - // Sample at t=3000: window [0, 3000) -> bucket 0 -> value=10 - // Sample at t=6000: window [3000, 6000) -> bucket 3000 -> value=10 - // Sample at t=9000: window [6000, 9000) -> bucket 6000 -> value=10 - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (0, Box::new(MockBucketAccumulator::new(0, 10.0))), - (3000, Box::new(MockBucketAccumulator::new(3, 10.0))), - (6000, Box::new(MockBucketAccumulator::new(6, 10.0))), - (9000, Box::new(MockBucketAccumulator::new(9, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 3000, // start_ms - 9000, // end_ms - 3000, // step_ms - ); - - assert_eq!( - results.len(), - 3, - "Should produce 3 samples with sparse data" - ); + return Ok(warm_qr); + } + } - // Each window only has 1 bucket because data is sparse - assert_eq!( - results[0], - (3000, 10.0, 0), - "t=3000: only bucket 0 in window" - ); - assert_eq!( - results[1], - (6000, 10.0, 3), - "t=6000: only bucket 3 in window" - ); - assert_eq!( - results[2], - (9000, 10.0, 6), - "t=9000: only bucket 6 in window" - ); + // CRITICAL #4: no-sketch-index fallback. The legacy + // `handle_query` path used to live here and provide the + // capability-miss notify side-effect through + // `find_compatible_aggregation_with_miss_notify`. With legacy + // retired (B7.5), this branch fires the notify directly so the + // control-plane feedback loop still closes — required by + // `capability_miss_http_e2e_tests::http_capability_miss_feedback_loop_closes_over_http`, + // which wires the engine WITHOUT `.with_sketch_index(...)` and + // therefore lands here on every miss. + if let Some(req) = Self::requirements_from_query_str(query) { + crate::drivers::control_plane_client::spawn_capability_miss_notify( + &self.control_plane_client, + &req, + ); + } + Err(crate::query_engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchStore.data_source_id(), + format!( + "ASAPQueryEngine: no sketch index for `{query}` — failing over to archive" + ), + )) } - #[test] - fn test_timestamp_lookup_missing_data_at_end() { - // Scenario: Data missing at end of query range - // Buckets at t=0, 1000, 2000 - // Query range [3000, 6000], step=1000, lookback=3 buckets - // - // Sample at t=3000: window [0, 3000) -> buckets 0, 1000, 2000 -> full data - // Sample at t=4000: window [1000, 4000) -> buckets 1000, 2000 -> partial (missing 3000) - // Sample at t=5000: window [2000, 5000) -> bucket 2000 -> partial - // Sample at t=6000: window [3000, 6000) -> no buckets -> skip - - let timestamped_buckets: Vec<(u64, Box)> = vec![ - (0, Box::new(MockBucketAccumulator::new(0, 10.0))), - (1000, Box::new(MockBucketAccumulator::new(1, 10.0))), - (2000, Box::new(MockBucketAccumulator::new(2, 10.0))), - ]; - - let results = simulate_timestamp_based_lookup( - timestamped_buckets, - 3, // lookback_bucket_count - 1000, // tumbling_window_ms - 3000, // start_ms - 6000, // end_ms - 1000, // step_ms - ); - - assert_eq!( - results.len(), - 3, - "Should produce 3 samples (last one skipped)" - ); - - assert_eq!(results[0], (3000, 30.0, 2), "t=3000: full window"); - assert_eq!( - results[1], - (4000, 20.0, 2), - "t=4000: partial window (2 buckets)" - ); - assert_eq!( - results[2], - (5000, 10.0, 2), - "t=5000: partial window (1 bucket)" - ); - // t=6000 is skipped because no data + fn capabilities(&self) -> crate::query_engines::routing::query_engine_routing::EngineCapabilities { + crate::query_engines::routing::query_engine_routing::EngineCapabilities { + data_source_id: asap_types::StorageBackend::SketchStore.data_source_id(), + storage_backend: asap_types::StorageBackend::SketchStore, + // Warm-tier sketches are O(sketch-size); call it 16 MiB ceiling + // for buffered ops (KLL with k=200 is well below this). + supports_streams_above_bytes: 16 * 1024 * 1024} } } @@ -4950,320 +1338,6 @@ mod hot_reload_phase2_tests { } } -// ─── End-to-end feedback loop test ───────────────────────────────────── -// -// The minimum-viable integration test for the full miss → notify → -// plan-push → next-query-hit loop. Covers every seam landed in PR #10 -// (HotReloadStreamingConfig endpoint), PR #11 (fire-and-forget -// capability-miss notification), PR #12 (ASAPQueryEngine per-query -// re-snapshot), and mirrors the DataCollector controller side from -// DataCollector PR #156 via an in-process mock client. -// -// What this test does NOT exercise: real HTTP traffic between real -// binaries. The mock controller is an in-process closure that directly -// swaps the `HotReloadStreamingConfig` handle. This is deliberate — -// each component is tested on its own in other suites, and the seams -// between them (`ASAPQueryEngine` field types, the shared `ArcSwap`, -// the `spawn_capability_miss_notify` helper) are what this test -// validates. -// -// The cross-process e2e (real collector, real backend, real query) -// is tracked as a separate operational follow-up and is bounded by -// the pre-existing DataCollector go.mod module-resolution issues. -#[cfg(test)] -mod e2e_feedback_loop_tests { - use super::*; - use crate::storage_engines::types::{ - AggregationType, CleanupPolicy, HotReloadStreamingConfig, - StreamingConfig, WindowType}; - use crate::drivers::control_plane_client::ControlPlaneClient; - use async_trait::async_trait; - use promql_utilities::data_model::key_by_label_names::KeyByLabelNames; - use promql_utilities::query_logics::enums::Statistic; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Mutex; - use std::time::{Duration, Instant}; - - fn agg_for_metric(_id: u64, metric: &str) -> crate::storage_engines::types::AggregationConfig { - // `_id` is unused after PR 5 — identity is content-addressed. - crate::storage_engines::types::AggregationConfig::new( - AggregationType::Sum, - String::new(), - std::collections::HashMap::new(), - KeyByLabelNames::empty(), - KeyByLabelNames::empty(), - KeyByLabelNames::empty(), - String::new(), - 60, - 60, - WindowType::Tumbling, - String::new(), - metric.to_string(), - None, - None, - None, - ) - } - - fn streaming_config_with(metric: &str, id: u64) -> StreamingConfig { - let mut map = std::collections::HashMap::new(); - let cfg = agg_for_metric(id, metric); - map.insert(cfg.aggregation_id(), cfg); - StreamingConfig::new(map) - } - - /// Mock controller that stands in for DataCollector's - /// `controller/src/main.rs`: - /// - /// On each `notify_capability_miss` call it: - /// 1. Records the requirement (for test assertions). - /// 2. Invokes a user-supplied "plan generator" closure that - /// produces a fresh `StreamingConfig` from the requirements. - /// 3. Swaps the backend's `HotReloadStreamingConfig` handle — - /// mirroring what happens when DataCollector PR #156's - /// `BackendClient::push_streaming_config` POSTs to the - /// backend's `/api/v1/streaming-config` endpoint on a real - /// cross-binary deployment. - struct InProcessMockControlPlane { - calls: Mutex>, - call_count: AtomicUsize, - hot_reload: HotReloadStreamingConfig, - planner: Box< - dyn Fn(&asap_types::query_requirements::QueryRequirements) -> StreamingConfig - + Send - + Sync, - >} - - impl InProcessMockControlPlane { - fn new( - hot_reload: HotReloadStreamingConfig, - planner: impl Fn(&asap_types::query_requirements::QueryRequirements) -> StreamingConfig - + Send - + Sync - + 'static, - ) -> Self { - Self { - calls: Mutex::new(Vec::new()), - call_count: AtomicUsize::new(0), - hot_reload, - planner: Box::new(planner)} - } - } - - #[async_trait] - impl ControlPlaneClient for InProcessMockControlPlane { - async fn notify_capability_miss( - &self, - requirements: &asap_types::query_requirements::QueryRequirements, - ) -> Result<(), String> { - self.call_count.fetch_add(1, Ordering::Relaxed); - self.calls.lock().unwrap().push(requirements.clone()); - let new_config = (self.planner)(requirements); - self.hot_reload.swap(new_config); - Ok(()) - } - } - - /// Exercises the full PR #10/#11/#12 + DC #156 feedback loop end - /// to end: - /// - /// 1. Start with an empty `HotReloadStreamingConfig`. - /// 2. Build a `ASAPQueryEngine` wired to the handle (PR #12) and - /// to an in-process mock controller client (PR #11 + - /// DC #156 mirror). - /// 3. Observe initial snapshot: empty. - /// 4. Call `find_compatible_aggregation_with_miss_notify` with - /// a requirement that will not match anything. This fires - /// the fire-and-forget notification which runs the mock - /// controller's planner closure and swaps the handle. - /// 5. Poll `streaming_config_snapshot` until the swap lands. - /// 6. Assert final state has the new aggregation_id the - /// planner returned. - /// - /// This test simulates, inside a single process, exactly what a - /// real backend ↔ controller deployment does across HTTP. The - /// observable contract is: once the controller acts on a miss, - /// the next `ASAPQueryEngine` query snapshot reflects the new - /// plan. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn capability_miss_feedback_loop_closes() { - // 1. Empty initial config. - let hot_reload = HotReloadStreamingConfig::new(StreamingConfig::default()); - - // 2. Mock controller: when a miss comes in, generate a config - // that covers the requested metric. This mirrors DC's - // replanner running and POSTing via its BackendClient. - let mock = Arc::new(InProcessMockControlPlane::new(hot_reload.clone(), |req| { - // Mock controller mints an explicit id here just to keep - // the test self-contained. In production, the - // `control_plane::emit::asapquery_backend` emitter no longer - // writes `aggregationId` (M2.2) and the backend derives - // one via `compute_agg_config_id`; explicit ids in the - // YAML are still honored for backwards compatibility. - let id: u64 = { - use std::collections::hash_map::DefaultHasher; - use std::hash::{Hash, Hasher}; - let mut h = DefaultHasher::new(); - req.metric.hash(&mut h); - h.finish().saturating_add(1) - }; - streaming_config_with(&req.metric, id) - })); - - // 3. Build ASAPQueryEngine with the handle and mock controller. - let engine = ASAPQueryEngine::new_with_hot_reload(hot_reload.clone(), 15000) - .with_control_plane_client(mock.clone() as Arc); - - // 4. Initial snapshot: empty. - let snap_before = engine.streaming_config_snapshot(); - assert_eq!( - snap_before.aggregation_configs.len(), - 0, - "precondition: initial config should be empty" - ); - - // 5. Trigger a capability miss via the private helper. This - // is the same entry point the live query paths in - // simple_engine.rs call. - let requirements = asap_types::query_requirements::QueryRequirements { - metric: "http_requests_total".to_string(), - statistics: vec![Statistic::Sum], - data_range_ms: Some(60_000), - grouping_labels: KeyByLabelNames::new(vec!["service".to_string()]), - spatial_filter_normalized: String::new()}; - let miss_result = engine.find_compatible_aggregation_with_miss_notify(&requirements); - assert!( - miss_result.is_none(), - "miss handler should return None when no agg matches" - ); - - // 6. The notification is fire-and-forget via `tokio::spawn`, - // so yield and poll until the swap lands (or timeout). - let deadline = Instant::now() + Duration::from_secs(2); - loop { - tokio::task::yield_now().await; - if mock.call_count.load(Ordering::Relaxed) > 0 { - // Give the spawned task a moment to complete its - // async body — the call_count is bumped at the - // start of notify_capability_miss, but the swap() - // happens synchronously inside the same call, so - // once call_count > 0 the swap is already visible. - break; - } - if Instant::now() >= deadline { - ::std::panic!( - "feedback loop did not fire within 2s; call_count={}", - mock.call_count.load(Ordering::Relaxed) - ); - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - - // 7. Next query snapshot should reflect the new plan. - // This is the PR #12 per-query re-snapshot contract. - let snap_after = engine.streaming_config_snapshot(); - assert_eq!( - snap_after.aggregation_configs.len(), - 1, - "feedback loop should have populated the config — \ - call_count={}, recorded_calls={:?}", - mock.call_count.load(Ordering::Relaxed), - mock.calls.lock().unwrap().len() - ); - - // 8. Validate the controller received the exact requirements. - let recorded = mock.calls.lock().unwrap(); - assert_eq!(recorded.len(), 1); - assert_eq!(recorded[0].metric, "http_requests_total"); - assert_eq!(recorded[0].statistics, vec![Statistic::Sum]); - assert_eq!(recorded[0].data_range_ms, Some(60_000)); - - // 9. PR 5: the aggregation_id on the wire is now a content- - // addressed `PolicyFingerprint` derived from the agg's - // metric / type / parameters. Pinning the exact id would - // couple the test to the fingerprint algorithm; instead - // assert it's deterministic-non-zero. - let new_ids: Vec = snap_after.aggregation_configs.keys().copied().collect(); - assert_eq!(new_ids.len(), 1); - assert_ne!(new_ids[0], 0, "fingerprint is never the 0 sentinel"); - } - - /// Second-order check: after the loop closes, a repeat miss on - /// the **same** requirements must not spawn a second plan - /// generation — the existing config already covers it. This - /// verifies the loop is idempotent under the common replay - /// pattern where a query client retries. - /// - /// Note: this doesn't test the "query now hits" path directly - /// because calling into the matching engine from here requires - /// AggregationIdInfo plumbing that isn't easy to stub. The - /// observable proxy is: `find_compatible_aggregation_with_miss_notify` - /// returns `Some` on the second call, meaning the config has - /// the aggregation AND the miss-notify does NOT fire again. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn capability_miss_idempotent_on_repeat() { - let hot_reload = HotReloadStreamingConfig::new(StreamingConfig::default()); - let mock = Arc::new(InProcessMockControlPlane::new(hot_reload.clone(), |req| { - streaming_config_with(&req.metric, 42) - })); - - let engine = ASAPQueryEngine::new_with_hot_reload(hot_reload.clone(), 15000) - .with_control_plane_client(mock.clone() as Arc); - - let requirements = asap_types::query_requirements::QueryRequirements { - metric: "latency_ms".to_string(), - statistics: vec![Statistic::Sum], - data_range_ms: Some(60_000), - grouping_labels: KeyByLabelNames::new(vec!["host".to_string()]), - spatial_filter_normalized: String::new()}; - - // First call — miss, loop closes. - let first = engine.find_compatible_aggregation_with_miss_notify(&requirements); - assert!(first.is_none()); - - // Wait for the swap to land. - let deadline = Instant::now() + Duration::from_secs(2); - while mock.call_count.load(Ordering::Relaxed) == 0 && Instant::now() < deadline { - tokio::time::sleep(Duration::from_millis(10)).await; - } - assert_eq!(mock.call_count.load(Ordering::Relaxed), 1); - - // Second call on the same requirements — the swap should be - // visible via per-query re-snapshot. The agg_id is 42 (the - // planner closure pinned it above), and the metric matches, - // so `find_compatible_aggregation` should return Some. - // - // Whether the StreamingConfig's capability matcher actually - // accepts these requirements depends on its internal logic; - // if it rejects them for a reason unrelated to the metric - // being present, the miss-notify fires a second time. The - // test accepts either outcome but pins that the second - // attempt produces behavior consistent with PR #12's - // re-snapshot semantics. - let second = engine.find_compatible_aggregation_with_miss_notify(&requirements); - let after_count = mock.call_count.load(Ordering::Relaxed); - - // At minimum: the snapshot is populated. PR 5: keys are - // content-addressed fingerprints, so we just assert the entry - // count rather than a specific u64. - let snap = engine.streaming_config_snapshot(); - assert_eq!(snap.aggregation_configs.len(), 1); - - // Second call fires at most once more — the point is that - // the runtime doesn't spin into a loop retrying the same - // miss. Either the capability matcher found the new agg - // (after_count == 1), or it rejected the agg and re-notified - // (after_count == 2). Both are acceptable; unbounded retry - // would be a regression. - assert!( - after_count <= 2, - "idempotency: unexpected notification count {} (expected ≤ 2)", - after_count - ); - let _ = second; - } -} - // ============================================================ // Phase 1b tests: AuxStats pushdown on `query_precompute_for_statistic` // ============================================================ @@ -5272,11 +1346,21 @@ mod e2e_feedback_loop_tests { // query path returns the aux value without ever calling the // accumulator's `query_statistic` method. When the statistic is NOT // covered, the code falls through to `query_statistic`. +// +// B7.5 retirement note: the in-process `capability_miss_feedback_loop_closes` +// + `capability_miss_idempotent_on_repeat` tests that previously lived +// adjacent to this module were retired alongside their probe surface +// (`find_compatible_aggregation_with_miss_notify`). The +// capability-miss feedback loop is now exercised end-to-end via +// `crate::tests::capability_miss_http_e2e_tests::http_capability_miss_feedback_loop_closes_over_http`, +// which round-trips a real HTTP capability-miss through the modern +// `execute()` path's `spawn_capability_miss_notify` calls. #[cfg(test)] mod aux_pushdown_tests { use super::*; use crate::precompute_engine::operators::{ min_max_accumulator::MinMaxAccumulator, sum_accumulator::SumAccumulator}; + use crate::storage_engines::types::AggregationType; use promql_utilities::query_logics::enums::Statistic; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -5445,142 +1529,6 @@ mod aux_pushdown_tests { // ── build_query_execution_context_promql_for_agg_id (forced-agg) tests ── -#[cfg(test)] -mod forced_agg_id_tests { - use super::*; - use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; - use crate::tests::test_utilities::engine_factories::create_engine_single_pop; - - /// Sanity: the refactored auto-resolve path still produces a - /// valid context for a simple sum_over_time query — ensures - /// the `parse_and_match_promql` / `resolve_agg_info_promql` - /// extraction hasn't broken existing behaviour. - #[test] - fn auto_resolve_still_works_after_refactor() { - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )]; - let query = "sum_over_time(http_requests[60s])"; - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - ); - let ctx = engine.build_query_execution_context_promql(query.to_string(), 1000.0); - assert!(ctx.is_some(), "auto-resolve should yield a context"); - } - - /// The new forced-agg-id entry point produces a context when - /// the forced `agg_id` matches the one the auto-resolver - /// would have picked. Basic smoke test; §7 timeline dispatch - /// per-segment timeline dispatch tests exercise it against multiple - /// agg_ids. - #[test] - fn forced_agg_id_produces_context_for_known_id() { - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )]; - let query = "sum_over_time(http_requests[60s])"; - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - ); - // PR 5: `create_engine_single_pop` keys the streaming config - // on the policy fingerprint; pull the only id out of the live - // snapshot rather than hardcoding `1`. - let agg_id = *engine - .streaming_config_snapshot() - .aggregation_configs - .keys() - .next() - .expect("one agg config registered"); - let ctx = engine.build_query_execution_context_promql_for_agg_id( - query.to_string(), - 1000.0, - agg_id, - ); - assert!(ctx.is_some(), "forced valid agg_id should yield a context"); - let ctx = ctx.unwrap(); - assert_eq!(ctx.agg_info.aggregation_id_for_value, agg_id); - assert_eq!(ctx.agg_info.aggregation_id_for_key, agg_id); - } - - /// Unknown `agg_id` returns `None` without panicking or - /// polluting the auto-resolver state. Per-segment timeline dispatch relies on - /// this to gracefully skip timeline segments whose agg_id - /// disappeared from the StreamingConfig mid-query. - #[test] - fn forced_agg_id_returns_none_for_unknown_id() { - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )]; - let query = "sum_over_time(http_requests[60s])"; - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - ); - let ctx = engine.build_query_execution_context_promql_for_agg_id( - query.to_string(), - 1000.0, - 9999, // not in StreamingConfig - ); - assert!(ctx.is_none(), "unknown agg_id → None"); - } - - /// Forced and auto-resolved contexts should be observably - /// equivalent for the common one-agg case (where the - /// auto-resolver would have picked the same id). The - /// invariant that matters for per-segment timeline dispatch: dispatching - /// through the forced path against the single covering - /// segment yields the same answer as the existing path. - #[test] - fn forced_and_auto_resolve_produce_same_agg_info_for_single_agg() { - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )]; - let query = "sum_over_time(http_requests[60s])"; - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - ); - let agg_id = *engine - .streaming_config_snapshot() - .aggregation_configs - .keys() - .next() - .expect("one agg config registered"); - let auto = engine - .build_query_execution_context_promql(query.to_string(), 1000.0) - .unwrap(); - let forced = engine - .build_query_execution_context_promql_for_agg_id(query.to_string(), 1000.0, agg_id) - .unwrap(); - assert_eq!( - auto.agg_info.aggregation_id_for_value, - forced.agg_info.aggregation_id_for_value - ); - assert_eq!( - auto.agg_info.aggregation_type_for_value, - forced.agg_info.aggregation_type_for_value - ); - } -} - // =========================================================================== // HLL count() — capability matching + accumulator query round-trip. @@ -5591,225 +1539,17 @@ mod forced_agg_id_tests { // `query_statistic` returns the cardinality estimate. This is the // runtime contract the wire-side _hll alias resolver above relies on. // =========================================================================== -#[cfg(test)] -mod hll_count_query_tests { - use super::*; - use crate::precompute_engine::operators::HllSketchAccumulator; - use crate::tests::test_utilities::engine_factories::create_engine_single_pop; - use asap_sketchlib::sketches::hll::HllVariant; - - fn hll_with_observations(observations: &[u64]) -> HllSketchAccumulator { - // Build an HLL with precision 8 (256 registers) and populate - // its register array directly. Backend's `HllSketch` is a - // pure data carrier (no `insert_with_hash` surface) — the - // wire decoder unpacks raw registers from the modified-OTLP - // proto, and queries read those registers via the canonical - // `α_m × m² / Σ 2^(-r)` HLL estimator. To exercise the - // estimator we mimic what the agent's hashing pipeline would - // produce: for each observation, derive a (bucket, leading- - // zeros) pair from a SplitMix64-style spread of the input - // and write `max(register[bucket], leading_zeros)`. This is - // exactly the math `HyperLogLogImpl::insert_with_hash` uses, - // performed inline. - let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 8); - let m = 1u64 << 8; // 256 registers - for &v in observations { - let h = v.wrapping_mul(0x9E37_79B9_7F4A_7C15); - let bucket = (h >> (64 - 8)) as usize; // top 8 bits - // Remaining 56 bits — count leading zeros + 1 (capped at 64). - let rem = h << 8; - let lz = if rem == 0 { - 64 - 8 - } else { - rem.leading_zeros() - } as u8 - + 1; - if (bucket as u64) < m { - let r = &mut acc.inner.registers[bucket]; - if lz > *r { - *r = lz; - } - } - } - acc - } - - #[test] - fn count_over_hll_returns_cardinality() { - // Insert 100 distinct observations and verify HLL's - // `query_statistic(Count)` returns a cardinality estimate - // close to the truth. ε ≈ 1.04/√m for HLL precision 8 → m=256 - // → ≈ 6.5 % standard error, generous bound below. - let acc = hll_with_observations(&(1..=100).collect::>()); - let trait_obj: &dyn AggregateCore = &acc; - let v = trait_obj - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("HLL answers Statistic::Count"); - assert!( - (v - 100.0).abs() < 30.0, - "HLL cardinality estimate diverged: got {v} for n=100" - ); - } - - #[test] - fn count_over_empty_hll_returns_zero() { - let acc = HllSketchAccumulator::new(HllVariant::Regular, 8); - let trait_obj: &dyn AggregateCore = &acc; - let v = trait_obj - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .expect("empty HLL still answers Count"); - // Linear-counting branch returns 0 when all registers are 0. - assert!(v.abs() < 1e-9, "empty HLL cardinality should be 0, got {v}"); - } - - #[test] - fn cardinality_is_an_alias_of_count() { - let acc = hll_with_observations(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); - let trait_obj: &dyn AggregateCore = &acc; - let by_count = trait_obj - .query_statistic(Statistic::Count, &None, &HashMap::new()) - .unwrap(); - let by_card = trait_obj - .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) - .unwrap(); - assert!( - (by_count - by_card).abs() < 1e-9, - "Cardinality and Count must produce the same HLL estimate" - ); - } - - #[test] - fn capability_matching_resolves_count_to_hll() { - // End-to-end through the ASAPQueryEngine: register an HLL agg for - // `unique_users_per_min_hll`, run the `count(...)` query - // through `build_query_execution_context_promql`, and assert - // the resolved agg is the HLL one. Regression guard for the - // PR #111 honest-gap closure on HLL-Count capability. - let acc = hll_with_observations(&(1..=50).collect::>()); - let data = vec![(None, Box::new(acc) as Box)]; - // No `by (...)` modifier on the query → empty grouping. The - // engine factory's HLL agg is registered with empty grouping - // labels; this matches the ASAP-tier production shape. - let engine = create_engine_single_pop( - "unique_users_per_min_hll", - AggregationType::HLL, - vec![], - data, - "count(unique_users_per_min_hll)", - ); - let ctx = engine - .build_query_execution_context_promql( - "count(unique_users_per_min_hll)".to_string(), - 1.0, - ) - .expect("count(HLL_metric) should produce a context"); - assert_eq!( - ctx.agg_info.aggregation_type_for_value, - AggregationType::HLL - ); - assert_eq!(ctx.metadata.statistic_to_compute, Statistic::Count); - } -} - // =========================================================================== // KLL quantile — pin that DatasketchesKLL is in the Quantile capability // list and the accumulator answers `Statistic::Quantile`. Mirrors the // HLL-Count contract; closes the wire-side ingest gap diagnosis. // =========================================================================== -#[cfg(test)] -mod kll_quantile_query_tests { - use super::*; - use crate::precompute_engine::operators::DatasketchesKLLAccumulator; - use crate::tests::test_utilities::engine_factories::create_engine_single_pop; - - #[test] - fn capability_matching_resolves_quantile_to_kll() { - // KLL is one of the canonical quantile approximators (along - // with HydraKLL and DDSketch). Register a KLL agg for - // `request_size_bytes_quantile` and verify - // `quantile_over_time(0.99, ...)` resolves to it. - let acc = DatasketchesKLLAccumulator::new(200); - let data = vec![(None, Box::new(acc) as Box)]; - let engine = create_engine_single_pop( - "request_size_bytes_quantile", - AggregationType::DatasketchesKLL, - vec![], - data, - "quantile_over_time(0.99, request_size_bytes_quantile[30s])", - ); - let ctx = engine - .build_query_execution_context_promql( - "quantile_over_time(0.99, request_size_bytes_quantile[30s])".to_string(), - 30.0, - ) - .expect("quantile_over_time(KLL_metric) should produce a context"); - assert_eq!( - ctx.agg_info.aggregation_type_for_value, - AggregationType::DatasketchesKLL - ); - assert_eq!(ctx.metadata.statistic_to_compute, Statistic::Quantile); - assert_eq!( - ctx.metadata - .query_kwargs - .get("quantile") - .map(String::as_str), - Some("0.99") - ); - } -} - // =========================================================================== // Capability matching — Rate over CountMinSketch (PR #111 honest-gap // closure). With the new `Statistic::Rate` arm in // `compatible_agg_types`, `rate([])` against a CMS-only // agg config now matches. // =========================================================================== -#[cfg(test)] -mod cms_rate_capability_tests { - use super::*; - use crate::precompute_engine::operators::CountMinSketchAccumulator; - use crate::tests::test_utilities::engine_factories::create_engine_single_pop; - - // TODO: after InferenceConfig retirement this test regressed — - // capability-matching path returns None where the old find_query_config - // path returned the same agg. Functionality unchanged in production - // (capability matching is the only path now), but the test expectation - // needs the test factory updated. Mark ignored pending investigation. - #[test] - #[ignore = "regression after InferenceConfig retirement; see TODO"] - fn capability_matching_resolves_rate_to_count_min_sketch() { - let acc = CountMinSketchAccumulator::new(4, 64); - let data = vec![(None, Box::new(acc) as Box)]; - let engine = create_engine_single_pop( - "endpoint_request_freq", - AggregationType::CountMinSketch, - vec![], - data, - "rate(endpoint_request_freq[60s])", - ); - let ctx = engine - .build_query_execution_context_promql( - "rate(endpoint_request_freq[60s])".to_string(), - 60.0, - ) - .expect("rate over CMS should produce a context"); - assert_eq!( - ctx.agg_info.aggregation_type_for_value, - AggregationType::CountMinSketch - ); - assert_eq!(ctx.metadata.statistic_to_compute, Statistic::Rate); - // The engine pushes range_ms into kwargs so the CMS - // accumulator can divide events by seconds at query time. - assert_eq!( - ctx.metadata - .query_kwargs - .get("range_ms") - .map(String::as_str), - Some("60000") - ); - } -} - /// Phase 5 — `QueryEngine::execute` ASAP-tier classification tests. /// Pre-Phase-5 the trait adapter unconditionally delegated to /// `handle_query`. After Phase 5 wire-in, when a `SketchStore` is @@ -6101,292 +1841,3 @@ mod hybrid_stitch_tests { } } -#[cfg(test)] -mod analyzer_parity_tests { - //! PR-α parity tests — capture both PromQL → asap-tier analyzers - //! side by side and pin the output. - //! - //! Two analyzers exist today and the analyzer-unification chain - //! (α→β→γ→δ→ε) is going to collapse them. To make that collapse - //! verifiable, α (this test) freezes how each analyzer answers - //! every shape in an 18-query corpus. γ rewrites the engine path - //! as a shim over the control plane path; δ deletes the engine - //! analyzer. Both stages must preserve the *engine column* of - //! this table — that's the parity contract. - //! - //! The two analyzers: - //! - //! 1. **Control plane** — - //! `control_plane::asap_tier_analysis::analyze_promql_for_asap_tier`. - //! Pipeline: `query_parser::parse_query` → - //! `intent_algebra::lower::lower_parsed_query` → - //! `capability_for(&AggIntent)`. Output: - //! `ASAPTierAnalysis { candidates, unsupported }` — speaks - //! `Capability` + `AggIntent` (L3). - //! - //! 2. **Engine** — `ASAPQueryEngine::parse_and_match_promql` + - //! `build_query_requirements_promql`. Pipeline: - //! `promql_parser::parser::parse` → match against - //! `controller_patterns: HashMap>` - //! built at `new_with_hot_reload`. Output: - //! `(QueryPatternType, PromQLMatchResult)` + `QueryRequirements` - //! — speaks `Statistic` + `QueryPatternType` (physical - //! sketch-storage table). - //! - //! Known divergences pinned by this corpus (see - //! `control_plane/docs/analyzer-parity-matrix.md` for the - //! per-query explanation): - //! - //! - `count_over_time(m[r])` without an outer `count by`: - //! control plane → `MISS(UnsupportedAggIntent("count"))`, - //! engine → `OK pattern=only_temporal stats=[count]`. - //! - `histogram_quantile(phi, m[…])`: control plane substitutes to - //! `Quantile` at the parser site (γ5 / PR #144); engine has no - //! `histogram_quantile` pattern and falls through to `MISS`. - //! - `irate(m[r])`: engine pattern list omits `irate` so it - //! misses; control plane rejects it as `UnsupportedAggIntent("rate")`. - //! - `topk(k, sum_by(…))` (topk wrapping a spatial agg, no - //! metric leaf at the call site): engine's `topk` pattern only - //! accepts a bare metric; control plane accepts via the topk - //! bridge. - //! - Bare selectors (`m`, `m{l=v}`): control plane → `NoCallNodeFound`; - //! engine → `MISS(NoPattern)` (no aggregation / function node). - - use super::*; - use crate::storage_engines::types::{HotReloadStreamingConfig, StreamingConfig}; - - /// Build a parity-test `ASAPQueryEngine`. The engine analyzer's - /// `parse_and_match_promql` depends only on `self.controller_patterns` - /// (built inside `new_with_hot_reload` from a static table), so an - /// empty `StreamingConfig` is sufficient. `build_query_requirements_promql` - /// calls `resolve_metric_labels(&metric)` which returns `None` against - /// an empty config and falls back to `KeyByLabelNames::empty()` — we - /// want exactly that fallback so the parity output is deterministic - /// and independent of any schema registry state. - fn make_engine() -> ASAPQueryEngine { - let sc = Arc::new(StreamingConfig::new(HashMap::new())); - let hr = HotReloadStreamingConfig::from_arc(sc); - ASAPQueryEngine::new_with_hot_reload(hr, 60) - } - - /// The 18-query parity corpus. Each row is `(id, promql)`. The id - /// is the row anchor in `control_plane/docs/analyzer-parity-matrix.md`; - /// keep them aligned when adding queries. - const CORPUS: &[(&str, &str)] = &[ - ("q01", "quantile_over_time(0.99, http_latency_ms[5m])"), - ("q02", "quantile_over_time(0.5, m[30s])"), - ("q03", "quantile_over_time(0.99, m[2h])"), - ("q04", "sum by (zone) (http_requests_total)"), - ("q05", "sum by (zone, region) (http_requests_total)"), - ("q06", "topk(5, http_requests_total)"), - ("q07", "topk(10, sum by (svc) (m))"), - ("q08", "count_over_time(http_requests_total[5m])"), - ("q09", "count by (zone) (count_over_time(http_requests_total[5m]))"), - ("q10", "histogram_quantile(0.99, sum by (le) (rate(http_latency_bucket[5m])))"), - ("q11", "histogram_quantile(0.99, http_latency_bucket)"), - ("q12", "http_requests_total"), - ("q13", "http_requests_total{zone=\"z0\"}"), - ("q14", "rate(http_requests_total[5m])"), - ("q15", "irate(http_requests_total[5m])"), - ("q16", "increase(http_requests_total[5m])"), - ("q17", "sum(rate(http_requests_total[5m]))"), - ("q18", "@@@ not promql @@@"), - ]; - - /// One-line stable summary of `ASAPTierAnalysis`. `MISS(reason)` on - /// the unsupported path; `OK [cand, ...]` on the supported path - /// with the full candidate shape so γ can be checked against this - /// without ambiguity. - fn summarize_controller(q: &str) -> String { - let a = control_plane::asap_tier_analysis::analyze_promql_for_asap_tier(q); - if let Some(reason) = &a.unsupported { - return format!("MISS({:?})", reason); - } - if a.candidates.is_empty() { - return "MISS(NoCandidates)".to_string(); - } - let cands: Vec = a - .candidates - .iter() - .map(|c| { - let gbk: Vec<&str> = c.group_by_keys.iter().map(|s| s.as_str()).collect(); - format!( - "metric={} gbk={:?} cap={:?} fn={} args={:?} range_s={}", - c.metric_name, - gbk, - c.required_capability, - c.function, - c.function_args, - c.range_seconds, - ) - }) - .collect(); - format!("OK [{}]", cands.join(" | ")) - } - - /// One-line stable summary of the engine analyzer's - /// `(QueryPatternType, PromQLMatchResult)` + `QueryRequirements`. - /// `MISS(NoPattern)` when no pattern in `controller_patterns` - /// matches the AST; `OK pattern=… stats=[…] …` otherwise. - fn summarize_engine(eng: &ASAPQueryEngine, q: &str) -> String { - match eng.parse_and_match_promql(q) { - None => "MISS(NoPattern)".to_string(), - Some((pt, mr)) => { - let req = eng.build_query_requirements_promql(&mr, pt); - let stats: Vec = - req.statistics.iter().map(|s| s.to_string()).collect(); - let fn_name = mr.get_function_name().unwrap_or_default(); - let agg_op = mr.get_aggregation_op().unwrap_or_default(); - let range_s = mr - .get_range_duration() - .map(|d| d.num_seconds().to_string()) - .unwrap_or_else(|| "-".to_string()); - format!( - "OK pattern={pattern} stats=[{stats}] metric={metric} fn={fn_name} \ - agg_op={agg_op} range_s={range_s} range_ms={range_ms:?} \ - spatial={spatial:?} grouping={grouping:?}", - pattern = pt, - stats = stats.join(","), - metric = req.metric, - fn_name = fn_name, - agg_op = agg_op, - range_s = range_s, - range_ms = req.data_range_ms, - spatial = req.spatial_filter_normalized, - grouping = req.grouping_labels.labels, - ) - } - } - } - - fn build_parity_table() -> String { - let eng = make_engine(); - let mut out = String::new(); - for (id, q) in CORPUS { - out.push_str(&format!("─── {id}: {q}\n")); - out.push_str(&format!(" ctrl {}\n", summarize_controller(q))); - out.push_str(&format!(" engine {}\n", summarize_engine(&eng, q))); - } - out - } - - /// Embedded golden master — captured against `origin/main` at - /// commit `6557fb8` (post-PR #187), re-verified byte-for-byte on - /// the rebase onto `origin/main` post-PR #211. Replace whenever an - /// analyzer output changes intentionally: re-run the test, copy the - /// printed `=== ACTUAL ===` block, and update - /// `control_plane/docs/analyzer-parity-matrix.md` in the same PR. - /// - /// Each row records what the analyzer **today** answers. β/γ MUST - /// preserve every `engine ...` row (the parity contract); δ MAY - /// change them only if the matching `ctrl ...` row already matches - /// the new behavior. The two paths must converge, not drift apart. - const GOLDEN: &str = "\ -─── q01: quantile_over_time(0.99, http_latency_ms[5m]) - ctrl OK [metric=http_latency_ms gbk=[] cap=QuantileApprox(Any) fn=quantile_over_time args=[0.99] range_s=300] - engine OK pattern=only_temporal stats=[quantile] metric=http_latency_ms fn=quantile_over_time agg_op= range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[] -─── q02: quantile_over_time(0.5, m[30s]) - ctrl OK [metric=m gbk=[] cap=QuantileApprox(Any) fn=quantile_over_time args=[0.5] range_s=30] - engine OK pattern=only_temporal stats=[quantile] metric=m fn=quantile_over_time agg_op= range_s=30 range_ms=Some(30000) spatial=\"\" grouping=[] -─── q03: quantile_over_time(0.99, m[2h]) - ctrl OK [metric=m gbk=[] cap=QuantileApprox(Any) fn=quantile_over_time args=[0.99] range_s=7200] - engine OK pattern=only_temporal stats=[quantile] metric=m fn=quantile_over_time agg_op= range_s=7200 range_ms=Some(7200000) spatial=\"\" grouping=[] -─── q04: sum by (zone) (http_requests_total) - ctrl OK [metric=http_requests_total gbk=[\"zone\"] cap=ExactAgg(Sum) fn=sum args=[] range_s=0] - engine OK pattern=only_spatial stats=[sum] metric=http_requests_total fn= agg_op=sum range_s=- range_ms=None spatial=\"\" grouping=[\"zone\"] -─── q05: sum by (zone, region) (http_requests_total) - ctrl OK [metric=http_requests_total gbk=[\"region\", \"zone\"] cap=ExactAgg(Sum) fn=sum args=[] range_s=0] - engine OK pattern=only_spatial stats=[sum] metric=http_requests_total fn= agg_op=sum range_s=- range_ms=None spatial=\"\" grouping=[\"region\", \"zone\"] -─── q06: topk(5, http_requests_total) - ctrl OK [metric=http_requests_total gbk=[] cap=FrequencyTopk(Any) fn=topk args=[5.0] range_s=0] - engine OK pattern=only_spatial stats=[topk] metric=http_requests_total fn= agg_op=topk range_s=- range_ms=None spatial=\"\" grouping=[] -─── q07: topk(10, sum by (svc) (m)) - ctrl OK [metric=m gbk=[\"svc\"] cap=FrequencyTopk(Any) fn=topk args=[10.0] range_s=0] - engine MISS(NoPattern) -─── q08: count_over_time(http_requests_total[5m]) - ctrl OK [metric=http_requests_total gbk=[] cap=FrequencyEstimate(Any) fn=count_over_time args=[] range_s=300] - engine OK pattern=only_temporal stats=[count] metric=http_requests_total fn=count_over_time agg_op= range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[] -─── q09: count by (zone) (count_over_time(http_requests_total[5m])) - ctrl OK [metric=http_requests_total gbk=[\"zone\"] cap=CardinalityApprox fn=count args=[] range_s=300 | metric=http_requests_total gbk=[\"zone\"] cap=CardinalityApprox fn=count args=[] range_s=300] - engine OK pattern=one_temporal_one_spatial stats=[count] metric=http_requests_total fn=count_over_time agg_op=count range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[\"zone\"] -─── q10: histogram_quantile(0.99, sum by (le) (rate(http_latency_bucket[5m]))) - ctrl MISS(UnparseableMetricsql(\"expected MatrixSelector, got Discriminant(0)\")) - engine MISS(NoPattern) -─── q11: histogram_quantile(0.99, http_latency_bucket) - ctrl MISS(UnparseableMetricsql(\"expected MatrixSelector, got Discriminant(7)\")) - engine MISS(NoPattern) -─── q12: http_requests_total - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn= args=[] range_s=0] - engine MISS(NoPattern) -─── q13: http_requests_total{zone=\"z0\"} - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn= args=[] range_s=0] - engine MISS(NoPattern) -─── q14: rate(http_requests_total[5m]) - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn=rate args=[] range_s=300] - engine OK pattern=only_temporal stats=[rate] metric=http_requests_total fn=rate agg_op= range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[] -─── q15: irate(http_requests_total[5m]) - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn=irate args=[] range_s=300] - engine MISS(NoPattern) -─── q16: increase(http_requests_total[5m]) - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn=increase args=[] range_s=300] - engine OK pattern=only_temporal stats=[increase] metric=http_requests_total fn=increase agg_op= range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[] -─── q17: sum(rate(http_requests_total[5m])) - ctrl OK [metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn=sum args=[] range_s=300] - engine OK pattern=one_temporal_one_spatial stats=[rate] metric=http_requests_total fn=rate agg_op=sum range_s=300 range_ms=Some(300000) spatial=\"\" grouping=[] -─── q18: @@@ not promql @@@ - ctrl MISS(UnparseableMetricsql(\"PromQL parse error: invalid promql query\")) - engine MISS(NoPattern) -"; - - /// Run all 18 queries through both analyzers, format as a parity - /// table, and pin against the embedded golden. A mismatch here is - /// the parity-violation signal β/γ must not trip — and is what - /// PR #144 (`histogram_quantile` parser substitution) regression- - /// guards against. - #[test] - fn analyzer_parity_18_query_corpus() { - let actual = build_parity_table(); - if actual != GOLDEN { - eprintln!("=== ACTUAL ===\n{actual}=== END ACTUAL ==="); - } - assert_eq!( - actual, GOLDEN, - "analyzer parity drifted — update control_plane/docs/analyzer-parity-matrix.md \ - and replace GOLDEN with the new ACTUAL block above" - ); - } -} - -#[cfg(test)] -mod calculate_start_timestamp_promql_tests { - use super::*; - use crate::storage_engines::types::{HotReloadStreamingConfig, StreamingConfig}; - - #[test] - fn calculate_start_timestamp_promql_handles_time_zero_without_underflow() { - // Capability-miss probe queries fire with time=0 (Unix epoch). - // Pre-fix: u64 subtraction underflows and panics with - // "attempt to subtract with overflow". Post-fix: saturating_sub - // clamps to 0, which the downstream store query treats as a - // [0, 0]-width range — degenerates to an empty result, the - // right answer when the probe is looking for capability-miss - // signal not data. - - let hot_reload = HotReloadStreamingConfig::from_arc(Arc::new(StreamingConfig::default())); - let engine = ASAPQueryEngine::new_with_hot_reload(hot_reload, 15000); - - // Synthesize a minimal OnlySpatial match_result. The body of - // calculate_start_timestamp_promql for OnlySpatial only reads - // `self.prometheus_scrape_interval`; the match_result arg is - // unused in that branch. A default-constructed - // PromQLMatchResult is fine. - let mr = PromQLMatchResult::new(); - - let start = engine.calculate_start_timestamp_promql( - 0, // end_timestamp = 0 (the bug trigger) - QueryPatternType::OnlySpatial, - &mr, - ); - assert_eq!(start, 0, "saturating_sub should clamp to 0, not panic"); - } -} diff --git a/data_plane/src/query_engines/asap_query_engine/mod.rs b/data_plane/src/query_engines/asap_query_engine/mod.rs index ada153e8..06891a44 100644 --- a/data_plane/src/query_engines/asap_query_engine/mod.rs +++ b/data_plane/src/query_engines/asap_query_engine/mod.rs @@ -18,7 +18,4 @@ pub use crate::storage_engines::sketch_db::query as asap_tier; #[cfg(test)] pub mod tests; -pub use engine::{ - ASAPQueryEngine, QueryExecutionContext, QueryMetadata, QueryTimestamps, - StoreQueryParams, StoreQueryPlan, -}; +pub use engine::ASAPQueryEngine; diff --git a/data_plane/src/tests/capability_matching_tests.rs b/data_plane/src/tests/capability_matching_tests.rs deleted file mode 100644 index 1dd22217..00000000 --- a/data_plane/src/tests/capability_matching_tests.rs +++ /dev/null @@ -1,285 +0,0 @@ -//! Integration tests for capability-based aggregation matching. -//! -//! These tests verify that when no pre-configured query_config entry exists, -//! the engine falls back to searching StreamingConfig by capability, and that -//! the existing query_config path still takes priority when an entry is present. - -use crate::drivers::ingest::series_resolver::SeriesIdResolver; -use crate::storage_engines::types::{ - AggregationConfig, AggregationType, PrecomputedOutput, StreamingConfig, WindowType}; -use crate::query_engines::asap_query_engine::engine::ASAPQueryEngine; -use crate::precompute_engine::operators::count_min_sketch_accumulator::CountMinSketchAccumulator; -use crate::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; -use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; -use crate::storage_engines::sketch_db::index::SketchStore; -use promql_utilities::data_model::KeyByLabelNames; -use std::collections::HashMap; -use std::sync::Arc; - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Build a minimal `AggregationConfig`. -fn make_agg_config( - _id: u64, - metric: &str, - agg_type: AggregationType, - window_size_s: u64, - window_type: WindowType, - grouping: &[&str], -) -> AggregationConfig { - // `_id` is unused after PR 5 — identity is content-addressed via - // `PolicyFingerprint::from_config`. - AggregationConfig { - aggregation_type: agg_type, - aggregation_sub_type: String::new(), - parameters: HashMap::new(), - grouping_labels: KeyByLabelNames::new(grouping.iter().map(|s| s.to_string()).collect()), - aggregated_labels: KeyByLabelNames::empty(), - rollup_labels: KeyByLabelNames::empty(), - original_yaml: String::new(), - window_size: window_size_s, - slide_interval: window_size_s, - window_type, - spatial_filter: String::new(), - spatial_filter_normalized: String::new(), - metric: metric.to_string(), - num_aggregates_to_retain: None, - table_name: None, - value_column: None} -} - -/// Build a `ASAPQueryEngine` with an explicit list of `AggregationConfig`s and no query_configs. -/// Data is inserted at timestamp 1_000_000 with a window covering [1_000_000 - window_ms, 1_000_000]. -fn engine_no_query_configs( - metric: &str, - schema_labels: &[&str], - agg_configs: Vec, -) -> ASAPQueryEngine { - let mut agg_map = HashMap::new(); - for c in &agg_configs { - agg_map.insert(c.aggregation_id(), c.clone()); - } - let streaming_config = Arc::new(StreamingConfig { - aggregation_configs: agg_map, - storage_backend: Default::default()}); - let sketch_index = Arc::new(SketchStore::new()); - - // Insert a data point for each aggregation so queries can actually execute. - let ts = 1_000_000_u64; - for c in &agg_configs { - let window_ms = c.window_size * 1000; - let output = PrecomputedOutput::new(ts - window_ms, ts, None, asap_types::PolicyFingerprint(c.aggregation_id())); - let acc: Box = match c.aggregation_type.as_str() { - "DatasketchesKLL" => { - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(1.0); - Box::new(kll) - } - "CountMinSketch" => { - let cms = CountMinSketchAccumulator::new(4, 1000); - Box::new(cms) - } - _ => Box::new(SumAccumulator::with_sum(42.0))}; - let resolver = Arc::new(SeriesIdResolver::new()); - sketch_index.ingest_precompute_for_agg_config( - |m, fp, ak| resolver.resolve(m, fp, ak), - c, - &output, - acc.as_ref(), - ); - } - - let schema_label_names = - KeyByLabelNames::new(schema_labels.iter().map(|s| s.to_string()).collect()); - - ASAPQueryEngine::new(streaming_config, 1).with_sketch_index(sketch_index) -} - -/// Build a `ASAPQueryEngine` with both a query_config entry AND a streaming aggregation. -fn engine_with_query_config( - metric: &str, - schema_labels: &[&str], - agg_config: AggregationConfig, - promql_query: &str, -) -> ASAPQueryEngine { - let agg_id = agg_config.aggregation_id(); - let mut agg_map = HashMap::new(); - agg_map.insert(agg_id, agg_config.clone()); - let streaming_config = Arc::new(StreamingConfig { - aggregation_configs: agg_map, - storage_backend: Default::default()}); - let sketch_index = Arc::new(SketchStore::new()); - - let ts = 1_000_000_u64; - let window_ms = agg_config.window_size * 1000; - let output = PrecomputedOutput::new(ts - window_ms, ts, None, asap_types::PolicyFingerprint(agg_id)); - let acc = SumAccumulator::with_sum(99.0); - let resolver = Arc::new(SeriesIdResolver::new()); - sketch_index.ingest_precompute_for_agg_config( - |m, fp, ak| resolver.resolve(m, fp, ak), - &agg_config, - &output, - &acc, - ); - - let schema_label_names = - KeyByLabelNames::new(schema_labels.iter().map(|s| s.to_string()).collect()); - - - - ASAPQueryEngine::new(streaming_config, 1).with_sketch_index(sketch_index) -} - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -/// When no query_config entry exists but a compatible Sum aggregation does, -/// capability matching should route to it and return a valid context. -#[test] -fn capability_fallback_fires_when_no_config() { - let agg = make_agg_config( - 1, - "cpu", - AggregationType::Sum, - 300, - WindowType::Tumbling, - &[], - ); - let expected = agg.aggregation_id(); - let engine = engine_no_query_configs("cpu", &[], vec![agg]); - - // sum_over_time(cpu[5m]) — 5 min = 300 s matches the 300 s tumbling config - let ctx = - engine.build_query_execution_context_promql("sum_over_time(cpu[5m])".to_string(), 1000.0); - assert!( - ctx.is_some(), - "Expected capability matching to find a compatible aggregation" - ); - assert_eq!(ctx.unwrap().agg_info.aggregation_id_for_value, expected); -} - -/// When a query_config entry exists, the engine must use it (not capability matching). -/// We verify by giving the config a different agg_id than any compatible-by-type config. -#[test] -fn config_path_takes_priority_over_capability_matching() { - let agg = make_agg_config( - 42, - "cpu", - AggregationType::Sum, - 300, - WindowType::Tumbling, - &[], - ); - let expected = agg.aggregation_id(); - let engine = engine_with_query_config("cpu", &[], agg, "sum_over_time(cpu[5m])"); - - let ctx = engine - .build_query_execution_context_promql("sum_over_time(cpu[5m])".to_string(), 1000.0) - .expect("should succeed via config path"); - - // The config path routes via the config's policy fingerprint. - assert_eq!(ctx.agg_info.aggregation_id_for_value, expected); -} - -/// A query for quantile(0.5) and quantile(0.9) should both resolve to the same -/// KLL aggregation when no query_configs are present. -#[test] -fn quantile_different_values_resolve_to_same_aggregation() { - let kll = make_agg_config( - 7, - "latency", - AggregationType::DatasketchesKLL, - 300, - WindowType::Tumbling, - &[], - ); - let engine = engine_no_query_configs("latency", &[], vec![kll]); - - let q50 = engine.build_query_execution_context_promql( - "quantile_over_time(0.5, latency[5m])".to_string(), - 1000.0, - ); - let q90 = engine.build_query_execution_context_promql( - "quantile_over_time(0.9, latency[5m])".to_string(), - 1000.0, - ); - - assert!( - q50.is_some(), - "quantile(0.5) should resolve via capability matching" - ); - assert!( - q90.is_some(), - "quantile(0.9) should resolve via capability matching" - ); - assert_eq!( - q50.unwrap().agg_info.aggregation_id_for_value, - q90.unwrap().agg_info.aggregation_id_for_value, - "Both quantile queries should route to the same KLL aggregation" - ); -} - -/// When no config entry exists and no compatible aggregation exists, return None. -#[test] -fn no_match_returns_none() { - // KLL config present, but query asks for Sum — incompatible - let kll = make_agg_config( - 1, - "cpu", - AggregationType::DatasketchesKLL, - 300, - WindowType::Tumbling, - &[], - ); - let engine = engine_no_query_configs("cpu", &[], vec![kll]); - - let ctx = - engine.build_query_execution_context_promql("sum_over_time(cpu[5m])".to_string(), 1000.0); - assert!( - ctx.is_none(), - "Should return None when no compatible aggregation exists" - ); -} - -/// When multiple compatible aggregations exist, the largest window should be preferred. -#[test] -fn priority_largest_window_wins() { - let small = make_agg_config( - 1, - "cpu", - AggregationType::Sum, - 300, - WindowType::Tumbling, - &[], - ); - let large = make_agg_config( - 2, - "cpu", - AggregationType::Sum, - 900, - WindowType::Tumbling, - &[], - ); - let expected_large = large.aggregation_id(); - let engine = engine_no_query_configs("cpu", &[], vec![small, large]); - - // sum_over_time(cpu[15m]) = 900 s — both 300 s and 900 s configs match (900 = 3×300), - // but the largest window should be preferred. - let ctx = engine - .build_query_execution_context_promql("sum_over_time(cpu[15m])".to_string(), 1000.0) - .expect("should find a compatible aggregation"); - - assert_eq!( - ctx.agg_info.aggregation_id_for_value, expected_large, - "The 900 s aggregation should be preferred over the 300 s one" - ); -} - -// `cms_only_backend_resolves_sum_over_time_via_capability_matching` -// retired alongside the `SetAggregator` / `DeltaSetAggregator` family. -// The test asserted that a CMS aggregation paired with a -// `DeltaSetAggregator` key aggregation would resolve `sum_over_time` -// via capability matching; the pairing is no longer expressible. diff --git a/data_plane/src/tests/mod.rs b/data_plane/src/tests/mod.rs index f33c2fe7..bc38c61e 100644 --- a/data_plane/src/tests/mod.rs +++ b/data_plane/src/tests/mod.rs @@ -1,12 +1,14 @@ pub mod accuracy_empirical_validation_tests; pub mod accuracy_in_promql_response_tests; -pub mod capability_matching_tests; pub mod capability_miss_http_e2e_tests; // M2.3.6g — legacy SketchStore-specific suites retired: // persist_format_versioning_tests, persistence_integration_tests, // persistence_perf_tests, store_correctness_tests. +// B7.5 retirement — `capability_matching_tests` and +// `schema_timeline_dispatch_tests` exercised the now-deleted +// legacy `build_query_execution_context_promql` / +// `handle_query_promql` paths; removed alongside the engine fns. pub mod prometheus_forwarding_tests; -pub mod schema_timeline_dispatch_tests; pub mod trait_design_tests; #[cfg(test)] diff --git a/data_plane/src/tests/schema_timeline_dispatch_tests.rs b/data_plane/src/tests/schema_timeline_dispatch_tests.rs deleted file mode 100644 index 7483c123..00000000 --- a/data_plane/src/tests/schema_timeline_dispatch_tests.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! End-to-end tests for the schema-timeline query dispatcher. -//! -//! Exercises the full path from a PromQL query → sid catalog -//! timeline lookup → per-segment store query → `combine_statistic` -//! → Prometheus `warnings`, on a real `ASAPQueryEngine` + -//! `SketchStore`. -//! -//! Contract validated: queries that span a reconfigure boundary -//! do not see a silent data cliff. Combinable statistics (Count / -//! Sum / Min / Max) get the stitched answer; non-combinable or -//! Purged segments surface as explicit `warnings` on the response -//! so the caller knows the answer is partial. -//! -//! Lives inside the crate (not `tests/`) so we can reach the -//! crate-private helpers (`seed_sum_at`) directly without leaking -//! a test-only surface. - -use std::collections::HashMap; -use std::sync::Arc; - -use asap_types::aggregation_config::AggregationConfig; -use asap_types::enums::{AggregationType, WindowType}; -use promql_utilities::data_model::key_by_label_names::KeyByLabelNames; - -use crate::storage_engines::types::{ - HotReloadStreamingConfig, KeyByLabelValues, PrecomputedOutput, StreamingConfig}; -use crate::query_engines::{QueryResult, ASAPQueryEngine}; -use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; - -const METRIC: &str = "sensor_reading"; - -// Timeline layout used by the tests. Picked so that an instant -// query at `QUERY_TIME_SEC` produces a range that straddles the -// reconfigure boundary between `agg_1` and `agg_2`. -const QUERY_TIME_SEC: f64 = 501.0; -const QUERY_TIME_MS: u64 = 501_000; - -fn make_agg_config(_id: u64) -> AggregationConfig { - // `_id` is unused after PR 5 — identity is content-addressed. - AggregationConfig::new( - AggregationType::Sum, - String::new(), - HashMap::new(), - KeyByLabelNames::new(vec!["host".to_string()]), - KeyByLabelNames::empty(), - KeyByLabelNames::empty(), - String::new(), - 1, - 1, - WindowType::Tumbling, - String::new(), - METRIC.to_string(), - None, - None, - None, - ) -} - -/// Instant PromQL query used by the tests. Runs through the -/// OnlySpatial aggregation pattern (op=sum) with a `by (host)` -/// modifier — the engine's `format_final_results` path only -/// emits keyed output elements, so the query must be grouped for -/// the result vector to be non-empty. -const TEST_QUERY: &str = "sum by (host) (sensor_reading)"; - -fn build_engine( - streaming_config: Arc, - sketch_index: Arc, -) -> ASAPQueryEngine { - let hot_reload = HotReloadStreamingConfig::from_arc(streaming_config); - ASAPQueryEngine::new_with_hot_reload(hot_reload, 1) - .with_sketch_index(sketch_index) -} - -/// Insert a single `SumAccumulator` window at `ts` into `agg_id`. -/// M2.3.6g — SketchStore-only after the legacy SketchStore retirement. -/// Registers the sid in the catalog as a side-effect via -/// `ingest_precompute_for_agg_config`, so callers do not need to -/// pre-populate any schema/registry — the sid timeline is built -/// directly from these ingests. -fn seed_sum_at( - sketch_index: &crate::storage_engines::sketch_db::index::SketchStore, - streaming_config: &StreamingConfig, - agg_id: u64, - ts: u64, - host: &str, - sum: f64, -) { - let key = Some(KeyByLabelValues { - labels: vec![host.to_string()]}); - use crate::drivers::ingest::series_resolver::SeriesIdResolver; - use std::sync::Arc as StdArc; - let output = PrecomputedOutput::new(ts, ts, key, asap_types::PolicyFingerprint(agg_id)); - let acc = SumAccumulator::with_sum(sum); - if let Some(agg_cfg) = streaming_config.get_aggregation_config(agg_id) { - thread_local! { - static RESOLVER: StdArc = StdArc::new(SeriesIdResolver::new()); - } - let resolver = RESOLVER.with(|r| r.clone()); - sketch_index.ingest_precompute_for_agg_config( - |m, fp, ak| resolver.resolve(m, fp, ak), - agg_cfg, - &output, - &acc, - ); - } - let _ = (ts, host); -} - -/// Two schemas for the same metric, both answerable. Sum is -/// combinable, so the dispatcher folds 10.0 + 20.0 into -/// `Full(30.0)` — no warnings, no data cliff. -#[ignore] -#[test] -fn sum_query_across_reconfigure_boundary_returns_combined_full_result() { - panic!("ignored: schema retirement #5 follow-up — re-enable when sid-level cross-reconfigure dispatch lands"); -} - -/// agg_1 Expired (coverage=Purged, unresolved); agg_2 Active with -/// data. The dispatcher must surface the partial through -/// `QueryResult::warnings()`. -#[ignore] -#[test] -fn sum_query_with_purged_segment_returns_partial_with_warnings() { - panic!("ignored: schema retirement #5 follow-up — re-enable when sid-level cross-reconfigure dispatch lands"); -} - -/// Single-schema regression guard: when the timeline has only one -/// segment, the dispatcher returns `None`, the default single-agg -/// path handles the query, no warnings attach. -#[test] -fn single_schema_query_falls_through_to_default_path() { - let cfg = make_agg_config(7); - let agg_id = cfg.aggregation_id(); - let mut agg_map = HashMap::new(); - agg_map.insert(agg_id, cfg); - let streaming_config = Arc::new(StreamingConfig::new(agg_map)); - - let sketch_index = Arc::new(crate::storage_engines::sketch_db::index::SketchStore::new()); - // Single ingest registers exactly one sid in the catalog → the - // sid-level `timeline_for_metric` returns one segment → the - // dispatcher bails to the default single-agg path. - seed_sum_at(&sketch_index, &streaming_config, agg_id, QUERY_TIME_MS, "A", 42.0); - - let engine = build_engine(streaming_config, sketch_index); - - let (_labels, qr) = engine - .handle_query_promql(TEST_QUERY.to_string(), QUERY_TIME_SEC) - .expect("query must produce a result"); - - assert!( - qr.warnings().is_empty(), - "single-schema path must not attach warnings: {:?}", - qr.warnings() - ); - match qr { - QueryResult::Vector(iv) => { - assert_eq!(iv.values.len(), 1); - assert!( - (iv.values[0].value - 42.0).abs() < 1e-9, - "single-agg answer should be 42.0, got {}", - iv.values[0].value - ); - } - other => panic!("expected instant vector, got {other:?}")} -} diff --git a/data_plane/src/tests/test_utilities/comparison.rs b/data_plane/src/tests/test_utilities/comparison.rs deleted file mode 100644 index 4e6a7d3b..00000000 --- a/data_plane/src/tests/test_utilities/comparison.rs +++ /dev/null @@ -1,252 +0,0 @@ -//! Comparison utilities for query equivalence tests -//! -//! Provides assertion helpers for deep equality checking of query execution contexts. - -use crate::storage_engines::types::{AggregationIdInfo, AggregationType}; -use crate::query_engines::asap_query_engine::engine::{ - QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, -}; -use promql_utilities::data_model::KeyByLabelNames; - -/// Assert that two QueryExecutionContext objects are equivalent -/// -/// Compares all fields and provides detailed error messages on mismatch -pub fn assert_execution_context_equivalent( - context1: &QueryExecutionContext, - context2: &QueryExecutionContext, - test_name: &str, -) { - // Compare metric - assert_eq!( - context1.metric, context2.metric, - "{}: Metric mismatch", - test_name - ); - - // Compare do_merge - assert_eq!( - context1.do_merge, context2.do_merge, - "{}: do_merge mismatch", - test_name - ); - - // Compare metadata - assert_metadata_equivalent(&context1.metadata, &context2.metadata, test_name); - - // Compare store plans - assert_store_plan_equivalent(&context1.store_plan, &context2.store_plan, test_name); - - // Compare aggregation info - assert_agg_info_equivalent(&context1.agg_info, &context2.agg_info, test_name); - - // Note: We don't compare spatial_filter as it may have different representations - // that are semantically equivalent (e.g., different string formats) -} - -/// Assert that two QueryMetadata objects are equivalent -pub fn assert_metadata_equivalent(meta1: &QueryMetadata, meta2: &QueryMetadata, test_name: &str) { - // Compare output labels (KeyByLabelNames maintains sorted order, so direct comparison works) - assert_label_names_equivalent( - &meta1.query_output_labels, - &meta2.query_output_labels, - test_name, - ); - - // Compare statistic - assert_eq!( - meta1.statistic_to_compute, meta2.statistic_to_compute, - "{}: Statistic mismatch - PromQL={:?}, SQL={:?}", - test_name, meta1.statistic_to_compute, meta2.statistic_to_compute - ); - - // Compare kwargs - assert_eq!( - meta1.query_kwargs, meta2.query_kwargs, - "{}: Query kwargs mismatch - PromQL={:?}, SQL={:?}", - test_name, meta1.query_kwargs, meta2.query_kwargs - ); -} - -/// Assert that two StoreQueryPlan objects are equivalent -pub fn assert_store_plan_equivalent( - plan1: &StoreQueryPlan, - plan2: &StoreQueryPlan, - test_name: &str, -) { - // Compare values query - assert_store_params_equivalent(&plan1.values_query, &plan2.values_query, test_name); - - // Compare keys query (both Some or both None) - match (&plan1.keys_query, &plan2.keys_query) { - (Some(k1), Some(k2)) => assert_store_params_equivalent(k1, k2, test_name), - (None, None) => {} - (Some(_), None) => panic!( - "{}: Keys query presence mismatch - PromQL has keys query, SQL doesn't", - test_name - ), - (None, Some(_)) => panic!( - "{}: Keys query presence mismatch - SQL has keys query, PromQL doesn't", - test_name - ), - } -} - -/// Assert that two StoreQueryParams objects are equivalent -pub fn assert_store_params_equivalent( - params1: &StoreQueryParams, - params2: &StoreQueryParams, - test_name: &str, -) { - assert_eq!( - params1.metric, params2.metric, - "{}: Metric mismatch - PromQL='{}', SQL='{}'", - test_name, params1.metric, params2.metric - ); - - assert_eq!( - params1.aggregation_id, params2.aggregation_id, - "{}: Aggregation ID mismatch - PromQL={}, SQL={}", - test_name, params1.aggregation_id, params2.aggregation_id - ); - - assert_eq!( - params1.start_timestamp, params2.start_timestamp, - "{}: Start timestamp mismatch - PromQL={}, SQL={}", - test_name, params1.start_timestamp, params2.start_timestamp - ); - - assert_eq!( - params1.end_timestamp, params2.end_timestamp, - "{}: End timestamp mismatch - PromQL={}, SQL={}", - test_name, params1.end_timestamp, params2.end_timestamp - ); - - assert_eq!( - params1.is_exact_query, params2.is_exact_query, - "{}: Query type mismatch - PromQL={}, SQL={}", - test_name, params1.is_exact_query, params2.is_exact_query - ); -} - -/// Assert that two KeyByLabelNames objects are equivalent -pub fn assert_label_names_equivalent( - labels1: &KeyByLabelNames, - labels2: &KeyByLabelNames, - test_name: &str, -) { - // KeyByLabelNames maintains sorted order, so direct comparison works - assert_eq!( - labels1, labels2, - "{}: Label names mismatch - PromQL={:?}, SQL={:?}", - test_name, labels1.labels, labels2.labels - ); -} - -/// Assert that two AggregationIdInfo objects are equivalent -pub fn assert_agg_info_equivalent( - agg1: &AggregationIdInfo, - agg2: &AggregationIdInfo, - test_name: &str, -) { - assert_eq!( - agg1.aggregation_id_for_key, agg2.aggregation_id_for_key, - "{}: Aggregation ID for key mismatch - PromQL={}, SQL={}", - test_name, agg1.aggregation_id_for_key, agg2.aggregation_id_for_key - ); - - assert_eq!( - agg1.aggregation_id_for_value, agg2.aggregation_id_for_value, - "{}: Aggregation ID for value mismatch - PromQL={}, SQL={}", - test_name, agg1.aggregation_id_for_value, agg2.aggregation_id_for_value - ); - - assert_eq!( - agg1.aggregation_type_for_key, agg2.aggregation_type_for_key, - "{}: Aggregation type for key mismatch - PromQL='{}', SQL='{}'", - test_name, agg1.aggregation_type_for_key, agg2.aggregation_type_for_key - ); - - assert_eq!( - agg1.aggregation_type_for_value, agg2.aggregation_type_for_value, - "{}: Aggregation type for value mismatch - PromQL='{}', SQL='{}'", - test_name, agg1.aggregation_type_for_value, agg2.aggregation_type_for_value - ); -} - -#[cfg(test)] -mod tests { - use super::*; - use promql_utilities::query_logics::enums::Statistic; - use std::collections::HashMap; - - fn create_test_context() -> QueryExecutionContext { - QueryExecutionContext { - metric: "test_metric".to_string(), - metadata: QueryMetadata { - query_output_labels: KeyByLabelNames::new(vec!["L1".to_string(), "L2".to_string()]), - statistic_to_compute: Statistic::Sum, - query_kwargs: HashMap::new(), - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: "test_metric".to_string(), - aggregation_id: 1, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: false, - }, - keys_query: None, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: 1, - aggregation_id_for_value: 1, - aggregation_type_for_key: AggregationType::Sum, - aggregation_type_for_value: AggregationType::Sum, - }, - do_merge: true, // OnlyTemporal queries merge - spatial_filter: String::new(), - query_time: 2_000_000, // query timestamp in milliseconds - grouping_labels: KeyByLabelNames::new(vec!["L1".to_string(), "L2".to_string()]), - aggregated_labels: KeyByLabelNames::empty(), - } - } - - #[test] - fn test_identical_contexts_are_equivalent() { - let ctx1 = create_test_context(); - let ctx2 = create_test_context(); - - // Should not panic - assert_execution_context_equivalent(&ctx1, &ctx2, "test_identical"); - } - - #[test] - #[should_panic(expected = "Metric mismatch")] - fn test_different_metrics_fail() { - let ctx1 = create_test_context(); - let mut ctx2 = create_test_context(); - ctx2.metric = "different_metric".to_string(); - - assert_execution_context_equivalent(&ctx1, &ctx2, "test_different_metrics"); - } - - #[test] - #[should_panic(expected = "Statistic mismatch")] - fn test_different_statistics_fail() { - let ctx1 = create_test_context(); - let mut ctx2 = create_test_context(); - ctx2.metadata.statistic_to_compute = Statistic::Max; - - assert_execution_context_equivalent(&ctx1, &ctx2, "test_different_stats"); - } - - #[test] - #[should_panic(expected = "Start timestamp mismatch")] - fn test_different_timestamps_fail() { - let ctx1 = create_test_context(); - let mut ctx2 = create_test_context(); - ctx2.store_plan.values_query.start_timestamp = 5000; - - assert_execution_context_equivalent(&ctx1, &ctx2, "test_different_timestamps"); - } -} diff --git a/data_plane/src/tests/test_utilities/mod.rs b/data_plane/src/tests/test_utilities/mod.rs index feaf31b7..8515aa5f 100644 --- a/data_plane/src/tests/test_utilities/mod.rs +++ b/data_plane/src/tests/test_utilities/mod.rs @@ -1,11 +1,11 @@ //! Test utilities for query equivalence testing //! -//! This module provides utilities for testing that semantically equivalent -//! PromQL and SQL queries produce equivalent internal logic in the QueryEngine. +//! Provides engine-construction helpers shared by the surviving +//! capability-miss + asap-tier integration tests. The +//! `comparison.rs` module was retired with B7.5 because it +//! reached into the now-deleted `QueryExecutionContext` / +//! `StoreQueryPlan` legacy types. -pub mod comparison; pub mod engine_factories; -// Re-export commonly used items -pub use comparison::*; pub use engine_factories::*;