feat(planner): --planner milp for SQL

asap-planner-rs/src/bin/optimizer_cli.rs

@@ -4,8 +4,8 @@
 use std::path::PathBuf;
 
 use asap_planner::optimizer::{
-    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpInputs,
-    MilpPlan,
+    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpConfig,
+    MilpInputs, MilpPlan,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
@@ -82,6 +82,7 @@ fn main() -> anyhow::Result<()> {
 }
 
 fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
+    let config = MilpConfig::Promql(config);
     // Fail before solving when the plan would be written but can't be.
     if args.output_dir.is_some() {
         reject_unwritable_queries(config)?;

asap-planner-rs/src/main.rs

@@ -1,9 +1,10 @@
 use asap_planner::optimizer::{
-    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpInputs,
+    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpConfig,
+    MilpInputs,
 };
 use asap_planner::{
     Controller, ControllerConfig, ElasticController, ElasticRuntimeOptions, RuntimeOptions,
-    SQLController, SQLRuntimeOptions, StreamingEngine,
+    SQLController, SQLControllerConfig, SQLRuntimeOptions, StreamingEngine,
 };
 use asap_types::enums::QueryLanguage;
 use clap::Parser;
@@ -56,8 +57,9 @@ struct Args {
     #[arg(long = "clickhouse-database", required = false)]
     clickhouse_database: Option<String>,
 
-    /// `milp` plans with sketch-bench's rqe-optimizer: PromQL with
-    /// --input_config only, labels from its `metrics:` hints.
+    /// `milp` plans with sketch-bench's rqe-optimizer: PromQL or SQL with
+    /// --input_config only, labels from its `metrics:` hints or tables'
+    /// `metadata_columns`.
     #[arg(long, value_enum, default_value = "legacy")]
     planner: PlannerArg,
 
@@ -215,16 +217,19 @@ fn main() -> anyhow::Result<()> {
 
 fn run_milp(args: &Args) -> anyhow::Result<()> {
     anyhow::ensure!(
-        matches!(args.query_language, QueryLanguage::promql),
-        "--planner milp supports only --query-language promql"
+        matches!(
+            args.query_language,
+            QueryLanguage::promql | QueryLanguage::sql
+        ),
+        "--planner milp supports only --query-language promql or sql"
     );
     anyhow::ensure!(
         args.query_log.is_none(),
-        "--planner milp needs --input_config: query logs have no `metrics:` hints"
+        "--planner milp needs --input_config: query logs have no label hints"
     );
     anyhow::ensure!(
         args.prometheus_url.is_none(),
-        "--planner milp takes labels from the config's `metrics:` hints, not --prometheus-url"
+        "--planner milp takes labels from the config, not --prometheus-url"
     );
     anyhow::ensure!(
         !args.enable_punting && args.range_duration_ms == 0 && args.step_ms == 0,
@@ -246,11 +251,19 @@ fn run_milp(args: &Args) -> anyhow::Result<()> {
         scrape_interval_ms > 0,
         "--data-ingestion-interval-ms must be positive"
     );
-    let config: ControllerConfig = serde_yaml::from_str(&std::fs::read_to_string(config_path)?)?;
+    let yaml = std::fs::read_to_string(config_path)?;
+    let (promql, sql): (ControllerConfig, SQLControllerConfig);
+    let config = if args.query_language == QueryLanguage::sql {
+        sql = serde_yaml::from_str(&yaml)?;
+        MilpConfig::Sql(&sql)
+    } else {
+        promql = serde_yaml::from_str(&yaml)?;
+        MilpConfig::Promql(&promql)
+    };
     // Fail before solving: the plan is always written.
-    reject_unwritable_queries(&config)?;
+    reject_unwritable_queries(config)?;
     let plan = plan_milp(
-        &config,
+        config,
         &MilpInputs {
             workload_facts: args
                 .workload_facts
@@ -270,7 +283,7 @@ fn run_milp(args: &Args) -> anyhow::Result<()> {
             allow_undeployable_families: false,
         },
     )?;
-    plan_to_planner_output(&config, &plan.workload, &plan.solution)?
+    plan_to_planner_output(config, &plan.workload, &plan.solution)?
         .write_to_dir(&args.output_dir)?;
     Ok(())
 }