Added support for asap-planner-rs to infer query repetitions and config from Prometheus query log (#254)

Author: milindsrivastava1997

Cargo.lock

@@ -25,6 +25,7 @@ tracing.workspace = true
 tracing-subscriber.workspace = true
 clap.workspace = true
 indexmap.workspace = true
+chrono.workspace = true
 promql-parser = "0.5.0"
 
 [dev-dependencies]

asap-planner-rs/Cargo.toml

@@ -13,10 +13,17 @@ pub struct QueryGroup {
     pub id: Option<u32>,
     pub queries: Vec<String>,
     pub repetition_delay: u64,
+    #[serde(default)]
     pub controller_options: ControllerOptions,
+    /// Per-group step override (seconds). Falls back to `RuntimeOptions::step` when None.
+    #[serde(default)]
+    pub step: Option<u64>,
+    /// Per-group range_duration override (seconds). Falls back to `RuntimeOptions::range_duration` when None.
+    #[serde(default)]
+    pub range_duration: Option<u64>,
 }
 
-#[derive(Debug, Clone, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Default)]
 pub struct ControllerOptions {
     pub accuracy_sla: f64,
     pub latency_sla: f64,

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

@@ -2,6 +2,7 @@ pub mod config;
 pub mod error;
 pub mod output;
 pub mod planner;
+pub mod query_log;
 
 use serde_yaml::Value as YamlValue;
 use std::path::Path;
@@ -291,6 +292,30 @@ impl Controller {
         })
     }
 
+    /// Build a `Controller` from a Prometheus query log file and a metrics config YAML.
+    ///
+    /// - `log_path`: newline-delimited JSON query log (Prometheus `--query.log-file` output)
+    /// - `metrics_path`: YAML file with a `metrics:` section listing metric names and labels
+    pub fn from_query_log(
+        log_path: &Path,
+        metrics_path: &Path,
+        opts: RuntimeOptions,
+    ) -> Result<Self, ControllerError> {
+        let entries = query_log::parse_log_file(log_path)?;
+        let (instants, ranges) =
+            query_log::infer_queries(&entries, opts.prometheus_scrape_interval);
+
+        let metrics_yaml = std::fs::read_to_string(metrics_path)?;
+        let metrics_config: query_log::MetricsConfig = serde_yaml::from_str(&metrics_yaml)?;
+
+        let config = query_log::to_controller_config(instants, ranges, metrics_config.metrics);
+
+        Ok(Self {
+            config,
+            options: opts,
+        })
+    }
+
     pub fn generate(&self) -> Result<PlannerOutput, ControllerError> {
         let output = output::generator::generate_plan(&self.config, &self.options)?;
         Ok(PlannerOutput {

asap-planner-rs/src/lib.rs

@@ -6,8 +6,17 @@ use std::path::PathBuf;
 #[derive(Parser, Debug)]
 #[command(name = "asap-planner", about = "ASAP Query Planner")]
 struct Args {
-    #[arg(long = "input_config")]
-    input_config: PathBuf,
+    /// Path to a hand-authored YAML workload config. Mutually exclusive with --query-log.
+    #[arg(long = "input_config", conflicts_with = "query_log")]
+    input_config: Option<PathBuf>,
+
+    /// Path to a Prometheus query log file (newline-delimited JSON). Mutually exclusive with --input_config.
+    #[arg(long = "query-log", conflicts_with = "input_config")]
+    query_log: Option<PathBuf>,
+
+    /// Path to a metrics config YAML (required when using --query-log).
+    #[arg(long = "metrics-config", requires = "query_log")]
+    metrics_config: Option<PathBuf>,
 
     #[arg(long = "output_dir")]
     output_dir: PathBuf,
@@ -71,20 +80,33 @@ fn main() -> anyhow::Result<()> {
                 range_duration: args.range_duration,
                 step: args.step,
             };
-            let controller = Controller::from_file(&args.input_config, opts)?;
+            let controller = match (args.input_config, args.query_log) {
+                (Some(config_path), None) => Controller::from_file(&config_path, opts)?,
+                (None, Some(log_path)) => {
+                    let metrics_path = args
+                        .metrics_config
+                        .expect("--metrics-config is required when using --query-log");
+                    Controller::from_query_log(&log_path, &metrics_path, opts)?
+                }
+                _ => anyhow::bail!(
+                    "exactly one of --input_config or --query-log must be provided for PromQL mode"
+                ),
+            };
             controller.generate_to_dir(&args.output_dir)?;
         }
         QueryLanguage::sql | QueryLanguage::elastic_sql => {
             let interval = args.data_ingestion_interval.ok_or_else(|| {
                 anyhow::anyhow!("--data-ingestion-interval is required for SQL mode")
             })?;
+            let config_path = args
+                .input_config
+                .ok_or_else(|| anyhow::anyhow!("--input_config is required for SQL mode"))?;
             let opts = SQLRuntimeOptions {
                 streaming_engine: engine,
                 query_evaluation_time: None,
                 data_ingestion_interval: interval,
             };
-            SQLController::from_file(&args.input_config, opts)?
-                .generate_to_dir(&args.output_dir)?;
+            SQLController::from_file(&config_path, opts)?.generate_to_dir(&args.output_dir)?;
         }
         QueryLanguage::elastic_querydsl => {
             anyhow::bail!("ElasticQueryDSL is not yet supported");

asap-planner-rs/src/main.rs

@@ -57,8 +57,8 @@ pub fn generate_plan(
                 metric_schema.clone(),
                 opts.streaming_engine,
                 controller_config.sketch_parameters.clone(),
-                opts.range_duration,
-                opts.step,
+                qg.range_duration.unwrap_or(opts.range_duration),
+                qg.step.unwrap_or(opts.step),
                 cleanup_policy,
             );
 
@@ -73,14 +73,25 @@ pub fn generate_plan(
             }
 
             if should_process {
-                let (configs, cleanup_param) = processor.get_streaming_aggregation_configs()?;
-                let mut keys_for_query = Vec::new();
-                for config in configs {
-                    let key = config.identifying_key();
-                    keys_for_query.push((key.clone(), cleanup_param));
-                    dedup_map.entry(key).or_insert(config);
+                match processor.get_streaming_aggregation_configs() {
+                    Ok((configs, cleanup_param)) => {
+                        let mut keys_for_query = Vec::new();
+                        for config in configs {
+                            let key = config.identifying_key();
+                            keys_for_query.push((key.clone(), cleanup_param));
+                            dedup_map.entry(key).or_insert(config);
+                        }
+                        query_keys_map.insert(query_string.clone(), keys_for_query);
+                    }
+                    Err(ControllerError::UnknownMetric(ref metric)) => {
+                        tracing::warn!(
+                            query = %query_string,
+                            metric = %metric,
+                            "skipping query referencing unknown metric"
+                        );
+                    }
+                    Err(e) => return Err(e),
                 }
-                query_keys_map.insert(query_string.clone(), keys_for_query);
             }
         }
     }

asap-planner-rs/src/output/generator.rs

@@ -0,0 +1,55 @@
+use serde::Deserialize;
+
+use crate::config::input::{
+    AggregateCleanupConfig, ControllerConfig, MetricDefinition, QueryGroup,
+};
+
+use super::frequency::{InstantQueryInfo, RangeQueryInfo};
+
+/// Subset of ControllerConfig used when loading from a metrics-only YAML file.
+#[derive(Deserialize)]
+pub struct MetricsConfig {
+    pub metrics: Vec<MetricDefinition>,
+}
+
+/// Build a `ControllerConfig` from extracted instant and range queries plus a metrics definition.
+///
+/// Each query becomes its own `QueryGroup` (one query per group, no SLA fields needed).
+pub fn to_controller_config(
+    instants: Vec<InstantQueryInfo>,
+    ranges: Vec<RangeQueryInfo>,
+    metrics: Vec<MetricDefinition>,
+) -> ControllerConfig {
+    let mut query_groups: Vec<QueryGroup> = Vec::new();
+
+    for info in instants {
+        query_groups.push(QueryGroup {
+            id: None,
+            queries: vec![info.query],
+            repetition_delay: info.repetition_delay,
+            controller_options: Default::default(),
+            step: None,
+            range_duration: None,
+        });
+    }
+
+    for info in ranges {
+        query_groups.push(QueryGroup {
+            id: None,
+            queries: vec![info.query],
+            repetition_delay: info.repetition_delay,
+            controller_options: Default::default(),
+            step: Some(info.step),
+            range_duration: Some(info.range_duration),
+        });
+    }
+
+    ControllerConfig {
+        query_groups,
+        metrics,
+        sketch_parameters: None,
+        aggregate_cleanup: Some(AggregateCleanupConfig {
+            policy: Some("read_based".to_string()),
+        }),
+    }
+}