fix(planner): handle unsupported PromQL statistics without panicking

asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs

@@ -1,5 +1,3 @@
-use core::panic;
-
 use promql_parser::parser::Expr;
 use tracing::debug;
 
@@ -8,6 +6,30 @@ use crate::ast_matching::PromQLMatchResult;
 use crate::data_model::KeyByLabelNames;
 use crate::query_logics::enums::{AggregationOperator, QueryPatternType, Statistic};
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum StatisticExtractionError {
+    MissingStatistic { pattern_type: QueryPatternType },
+    UnsupportedStatistic { statistic: String },
+}
+
+impl std::fmt::Display for StatisticExtractionError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::MissingStatistic { pattern_type } => {
+                write!(
+                    f,
+                    "No statistic found for query pattern type {pattern_type:?}"
+                )
+            }
+            Self::UnsupportedStatistic { statistic } => {
+                write!(f, "Unsupported statistic: {statistic}")
+            }
+        }
+    }
+}
+
+impl std::error::Error for StatisticExtractionError {}
+
 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();
@@ -50,38 +72,39 @@ pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (Strin
     (metric_name, spatial_filter)
 }
 
-/// Get statistics to compute based on pattern type and tokens
+/// Get statistics to compute based on pattern type and tokens.
+/// Returns a typed error if the matched statistic/function name is not
+/// recognized, so callers can decide whether to skip or fail the query.
 pub fn get_statistics_to_compute(
     pattern_type: QueryPatternType,
     match_result: &PromQLMatchResult,
-) -> Vec<Statistic> {
+) -> Result<Vec<Statistic>, StatisticExtractionError> {
     debug!("Computing statistics for pattern type {:?}", pattern_type);
-    let statistic_to_compute: Option<String> = 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
+    let statistic_to_compute: Option<String> = match pattern_type {
+        QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => {
+            match_result.get_function_name().map(|function_name| {
+                let name = function_name.to_lowercase();
+                name.split('_').next().unwrap_or(&name).to_string()
+            })
+        }
+        QueryPatternType::OnlySpatial => match_result
             .get_aggregation_op()
-            .map(|agg| agg.to_lowercase())
-    } else {
-        panic!("Unsupported query pattern type");
+            .map(|agg| agg.to_lowercase()),
     };
 
-    if let Some(statistic_to_compute) = statistic_to_compute {
-        debug!("Found statistic to compute: {}", statistic_to_compute);
-        if statistic_to_compute.parse::<AggregationOperator>() == Ok(AggregationOperator::Avg) {
-            vec![Statistic::Sum, Statistic::Count]
-        } else if let Ok(stat) = statistic_to_compute.parse::<Statistic>() {
-            vec![stat]
-        } else {
-            panic!("Unsupported statistic: {}", statistic_to_compute);
-        }
+    let Some(statistic_to_compute) = statistic_to_compute else {
+        return Err(StatisticExtractionError::MissingStatistic { pattern_type });
+    };
+
+    debug!("Found statistic to compute: {}", statistic_to_compute);
+    if statistic_to_compute.parse::<AggregationOperator>() == Ok(AggregationOperator::Avg) {
+        Ok(vec![Statistic::Sum, Statistic::Count])
+    } else if let Ok(stat) = statistic_to_compute.parse::<Statistic>() {
+        Ok(vec![stat])
     } else {
-        panic!("No statistic found in the query");
+        Err(StatisticExtractionError::UnsupportedStatistic {
+            statistic: statistic_to_compute,
+        })
     }
 }
 
@@ -133,3 +156,69 @@ pub fn get_spatial_aggregation_output_labels(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use crate::ast_matching::{FunctionToken, TokenData};
+
+    use super::*;
+
+    fn empty_token() -> TokenData {
+        TokenData {
+            metric: None,
+            function: None,
+            aggregation: None,
+            range_vector: None,
+            subquery: None,
+            binary_op: None,
+            number: None,
+        }
+    }
+
+    fn temporal_match(function_name: &str) -> PromQLMatchResult {
+        let mut tokens = HashMap::new();
+        tokens.insert(
+            "function".to_string(),
+            TokenData {
+                function: Some(FunctionToken {
+                    name: function_name.to_string(),
+                    args: vec![],
+                }),
+                ..empty_token()
+            },
+        );
+        PromQLMatchResult::with_tokens(tokens)
+    }
+
+    #[test]
+    fn unsupported_matched_temporal_statistic_returns_typed_error() {
+        let err =
+            get_statistics_to_compute(QueryPatternType::OnlyTemporal, &temporal_match("stddev"))
+                .unwrap_err();
+
+        assert_eq!(
+            err,
+            StatisticExtractionError::UnsupportedStatistic {
+                statistic: "stddev".to_string()
+            }
+        );
+    }
+
+    #[test]
+    fn missing_statistic_returns_typed_error() {
+        let err = get_statistics_to_compute(
+            QueryPatternType::OnlyTemporal,
+            &PromQLMatchResult::with_tokens(HashMap::new()),
+        )
+        .unwrap_err();
+
+        assert_eq!(
+            err,
+            StatisticExtractionError::MissingStatistic {
+                pattern_type: QueryPatternType::OnlyTemporal
+            }
+        );
+    }
+}

asap-planner-rs/src/optimizer/aqe_extractor.rs

@@ -204,7 +204,16 @@ fn extract_requirements(query: &str, metric_schema: &PromQLSchema) -> Option<Que
     })?;
 
     let (metric, spatial_filter) = get_metric_and_spatial_filter(&match_result);
-    let statistics = get_statistics_to_compute(pattern_type, &match_result);
+    let statistics = get_statistics_to_compute(pattern_type, &match_result)
+        .map_err(|err| {
+            warn!(
+                query = %query,
+                error = %err,
+                "aqe_extractor: skipping matched leaf query with unsupported statistics"
+            );
+            err
+        })
+        .ok()?;
 
     let data_range_ms = match pattern_type {
         QueryPatternType::OnlySpatial => None,

asap-planner-rs/src/planner/promql.rs

@@ -226,7 +226,12 @@ impl SingleQueryProcessor {
             .ok_or_else(|| ControllerError::UnknownMetric(metric.clone()))?
             .clone();
 
-        let statistics = get_statistics_to_compute(pattern_type, &match_result);
+        let statistics = get_statistics_to_compute(pattern_type, &match_result).map_err(|err| {
+            ControllerError::PlannerError(format!(
+                "Unsupported statistic for query '{}': {}",
+                self.query, err
+            ))
+        })?;
 
         let mut window_cfg = IntermediateWindowConfig::default();
         set_window_parameters(

asap-query-engine/src/engines/simple_engine/promql.rs

@@ -265,7 +265,12 @@ impl SimpleEngine {
         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);
+        let statistics_to_compute = get_statistics_to_compute(query_pattern_type, match_result)
+            .map_err(|err| {
+                warn!("{}", err);
+                err
+            })
+            .ok()?;
         if statistics_to_compute.len() != 1 {
             warn!(
                 "Expected exactly one statistic to compute, found {}",
@@ -615,10 +620,15 @@ impl SimpleEngine {
         &self,
         match_result: &PromQLMatchResult,
         query_pattern_type: QueryPatternType,
-    ) -> QueryRequirements {
+    ) -> Option<QueryRequirements> {
         let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);
 
-        let statistics = get_statistics_to_compute(query_pattern_type, 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,
@@ -642,7 +652,7 @@ impl SimpleEngine {
             }
         };
 
-        QueryRequirements {
+        Some(QueryRequirements {
             metric,
             statistics,
             data_range_ms,
@@ -651,7 +661,7 @@ impl SimpleEngine {
             // 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.
@@ -1076,7 +1086,7 @@ impl SimpleEngine {
                 query
             );
             let requirements =
-                self.build_query_requirements_promql(&match_result, query_pattern_type);
+                self.build_query_requirements_promql(&match_result, query_pattern_type)?;
             self.streaming_config
                 .read()
                 .unwrap()