diff --git a/asap-common/dependencies/rs/asap_types/src/query_requirements.rs b/asap-common/dependencies/rs/asap_types/src/query_requirements.rs index 56b2b76..fbb9fc8 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_requirements.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_requirements.rs @@ -1,5 +1,13 @@ +use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::Statistic; +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, +}; +use tracing::warn; + +use crate::promql_schema::PromQLSchema; +use crate::utils::normalize_spatial_filter; /// What a query needs in order to be answered by a stored aggregation. #[derive(Debug, Clone)] @@ -28,3 +36,60 @@ pub struct QueryRequirements { /// constrain the sketch weighting); matching ignores it when `None`. pub topk_count_events: Option, } + +/// Build `QueryRequirements` from an already-pattern-matched PromQL query. +/// +/// Shared by the query engine's capability-matching fallback and the planner's +/// AQE extractor, which independently parse and pattern-match a query before +/// calling this. `query` is used only for diagnostics on the unsupported-stats +/// warning. +pub fn build_query_requirements_promql( + query: &str, + match_result: &PromQLMatchResult, + pattern_type: QueryPatternType, + metric_schema: &PromQLSchema, +) -> Option { + let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result); + + let statistics = get_statistics_to_compute(pattern_type, match_result) + .map_err(|err| { + warn!( + query = %query, + error = %err, + "skipping matched query with unsupported statistics" + ); + err + }) + .ok()?; + + let data_range_ms = match pattern_type { + QueryPatternType::OnlySpatial => None, + _ => match_result + .get_range_duration() + .map(|d| d.num_seconds() as u64 * 1000), + }; + + let all_labels = metric_schema + .get_labels(&metric) + .cloned() + .unwrap_or_else(KeyByLabelNames::empty); + + let grouping_labels = match pattern_type { + // OnlyTemporal preserves all labels. + QueryPatternType::OnlyTemporal => all_labels, + // OnlySpatial and OneTemporalOneSpatial encode their output labels in + // the AST's `by (...)` / `without (...)` clause. + QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => { + get_spatial_aggregation_output_labels(match_result, &all_labels) + } + }; + + Some(QueryRequirements { + metric, + statistics, + data_range_ms, + grouping_labels, + spatial_filter_normalized: normalize_spatial_filter(&spatial_filter), + topk_count_events: None, + }) +} diff --git a/asap-planner-rs/src/optimizer/aqe_extractor.rs b/asap-planner-rs/src/optimizer/aqe_extractor.rs index 1048ce0..e4df49a 100644 --- a/asap-planner-rs/src/optimizer/aqe_extractor.rs +++ b/asap-planner-rs/src/optimizer/aqe_extractor.rs @@ -1,13 +1,9 @@ use std::collections::HashMap; -use asap_types::query_requirements::QueryRequirements; -use asap_types::utils::normalize_spatial_filter; +use asap_types::query_requirements::{build_query_requirements_promql, QueryRequirements}; use asap_types::PromQLSchema; 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, -}; +use promql_utilities::query_logics::enums::Statistic; use tracing::warn; use crate::planner::patterns::build_patterns; @@ -160,12 +156,6 @@ fn decompose_to_leaves(query: &str) -> Vec { /// Try to extract `QueryRequirements` from a single leaf PromQL query string. /// Returns `None` if the query cannot be parsed or does not match any pattern. -/// -/// TODO: this duplicates `build_query_requirements_promql` in -/// `asap-query-engine/src/engines/simple_engine/promql.rs:614`. That function -/// is a private `&self` method tied to `SimplePromQLEngine`. The shared logic -/// should be extracted into a free function in `asap_types::query_requirements` -/// and called from both sites. fn extract_requirements(query: &str, metric_schema: &PromQLSchema) -> Option { let ast = promql_parser::parser::parse(query).ok()?; let patterns = build_patterns(); @@ -179,52 +169,7 @@ fn extract_requirements(query: &str, metric_schema: &PromQLSchema) -> Option None, - _ => match_result - .get_range_duration() - .map(|d| d.num_seconds() as u64 * 1000), - }; - - let grouping_labels = match pattern_type { - // OnlyTemporal preserves all labels — look them up in the schema. - // If the metric is unknown, fall back to empty (dedup still works; cost - // model will treat it as a zero-group-count sketch). - QueryPatternType::OnlyTemporal => metric_schema - .get_labels(&metric) - .cloned() - .unwrap_or_else(KeyByLabelNames::empty), - // OnlySpatial and OneTemporalOneSpatial encode their output labels in - // the AST's `by (...)` / `without (...)` clause. - QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => { - let all_labels = metric_schema - .get_labels(&metric) - .cloned() - .unwrap_or_else(KeyByLabelNames::empty); - get_spatial_aggregation_output_labels(&match_result, &all_labels) - } - }; - - Some(QueryRequirements { - metric, - statistics, - data_range_ms, - grouping_labels, - spatial_filter_normalized: normalize_spatial_filter(&spatial_filter), - topk_count_events: None, - }) + build_query_requirements_promql(query, &match_result, pattern_type, metric_schema) } #[cfg(test)] diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 457ac24..3426b1c 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -10,8 +10,8 @@ use super::{ }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; use crate::engines::query_result::{QueryResult, RangeVectorElement}; -use asap_types::query_requirements::QueryRequirements; -use asap_types::utils::normalize_spatial_filter; +use asap_types::query_requirements::build_query_requirements_promql; +use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::get_is_collapsable; @@ -614,56 +614,6 @@ impl SimpleEngine { Some((output_labels, QueryResult::matrix(combined))) } - /// 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, - ) -> Option { - let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result); - - let statistics = get_statistics_to_compute(query_pattern_type, match_result) - .map_err(|err| { - warn!("{}", err); - err - }) - .ok()?; - - let data_range_ms = match query_pattern_type { - QueryPatternType::OnlySpatial => None, - _ => match_result - .get_range_duration() - .map(|d| d.num_seconds() as u64 * 1000), - }; - - let all_labels = match &self.inference_config.read().unwrap().schema { - SchemaConfig::PromQL(schema) => schema - .get_labels(&metric) - .cloned() - .unwrap_or_else(KeyByLabelNames::empty), - _ => 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) - } - }; - - Some(QueryRequirements { - metric, - statistics, - data_range_ms, - grouping_labels, - spatial_filter_normalized: normalize_spatial_filter(&spatial_filter), - // PromQL top-k does not constrain the sketch weighting; leave the - // count/sum discriminator unset so matching does not over-filter. - topk_count_events: None, - }) - } - // /// Try to extract sketch query components from a PromQL query string. // /// // /// Attempts the standard AST parser first. If that fails (e.g. for custom @@ -1085,8 +1035,18 @@ impl SimpleEngine { "No query_config entry for PromQL query '{}'. Attempting capability-based matching.", query ); - let requirements = - self.build_query_requirements_promql(&match_result, query_pattern_type)?; + let inference_config = self.inference_config.read().unwrap(); + let empty_schema = PromQLSchema::new(); + let metric_schema = match &inference_config.schema { + SchemaConfig::PromQL(schema) => schema, + _ => &empty_schema, + }; + let requirements = build_query_requirements_promql( + &query, + &match_result, + query_pattern_type, + metric_schema, + )?; self.streaming_config .read() .unwrap()