use std::path::PathBuf; use rpki::query_db::{ArtifactIndexerConfig, index_artifacts}; #[derive(Debug, PartialEq, Eq)] struct Args { query_db: PathBuf, run_root: Option, run_dir: Option, repo_bytes_db: Option, projection_entry_limit: usize, min_run_seq: Option, retain_indexed_runs: Option, dump_summary: bool, } fn usage() -> &'static str { "Usage: rpki_query_indexer --query-db (--run-root |--run-dir ) [--repo-bytes-db ] [--projection-entry-limit ] [--min-run-seq ] [--retain-indexed-runs ] [--dump-summary]" } fn parse_args(argv: &[String]) -> Result { let mut query_db = None; let mut run_root = None; let mut run_dir = None; let mut repo_bytes_db = None; let mut projection_entry_limit = 50usize; let mut min_run_seq = None; let mut retain_indexed_runs = Some(10usize); let mut dump_summary = false; let mut index = 1usize; while index < argv.len() { match argv[index].as_str() { "--query-db" => { index += 1; query_db = Some(PathBuf::from(value_at(argv, index, "--query-db")?)); } "--run-root" => { index += 1; run_root = Some(PathBuf::from(value_at(argv, index, "--run-root")?)); } "--run-dir" => { index += 1; run_dir = Some(PathBuf::from(value_at(argv, index, "--run-dir")?)); } "--repo-bytes-db" => { index += 1; repo_bytes_db = Some(PathBuf::from(value_at(argv, index, "--repo-bytes-db")?)); } "--projection-entry-limit" => { index += 1; let raw = value_at(argv, index, "--projection-entry-limit")?; projection_entry_limit = raw .parse::() .map_err(|_| format!("invalid --projection-entry-limit: {raw}"))?; } "--min-run-seq" => { index += 1; let raw = value_at(argv, index, "--min-run-seq")?; min_run_seq = Some( raw.parse::() .map_err(|_| format!("invalid --min-run-seq: {raw}"))?, ); } "--retain-indexed-runs" => { index += 1; let raw = value_at(argv, index, "--retain-indexed-runs")?; retain_indexed_runs = Some( raw.parse::() .map_err(|_| format!("invalid --retain-indexed-runs: {raw}"))?, ); } "--dump-summary" => dump_summary = true, "-h" | "--help" => return Err(usage().to_string()), other => return Err(format!("unknown argument: {other}\n{}", usage())), } index += 1; } if run_root.is_some() == run_dir.is_some() { return Err(format!( "exactly one of --run-root or --run-dir is required\n{}", usage() )); } Ok(Args { query_db: query_db.ok_or_else(|| format!("--query-db is required\n{}", usage()))?, run_root, run_dir, repo_bytes_db, projection_entry_limit, min_run_seq, retain_indexed_runs, dump_summary, }) } fn value_at<'a>(argv: &'a [String], index: usize, flag: &str) -> Result<&'a str, String> { argv.get(index) .map(String::as_str) .ok_or_else(|| format!("{flag} requires a value")) } fn main() { let argv = std::env::args().collect::>(); let args = match parse_args(&argv) { Ok(args) => args, Err(err) => { eprintln!("{err}"); std::process::exit(2); } }; let summary = match index_artifacts(&ArtifactIndexerConfig { query_db_path: args.query_db, run_root: args.run_root, run_dir: args.run_dir, repo_bytes_db_path: args.repo_bytes_db, projection_entry_limit: args.projection_entry_limit, min_run_seq: args.min_run_seq, retain_indexed_runs: args.retain_indexed_runs, }) { Ok(summary) => summary, Err(err) => { eprintln!("query index failed: {err}"); std::process::exit(1); } }; if args.dump_summary { println!( "{}", serde_json::to_string_pretty(&summary).expect("serialize summary") ); } else { println!( "indexed_runs={} runs_deleted={} retained_runs={} repos={} publication_points={} objects={} latest_ready_run={}", summary.runs_indexed, summary.runs_deleted, summary.retained_runs, summary.repos_indexed, summary.publication_points_indexed, summary.object_instances_indexed, summary.latest_ready_run.as_deref().unwrap_or("none") ); } if !summary.errors.is_empty() { eprintln!( "index completed with {} per-run errors", summary.errors.len() ); std::process::exit(1); } } #[cfg(test)] mod tests { use super::*; fn argv(args: &[&str]) -> Vec { std::iter::once("rpki_query_indexer") .chain(args.iter().copied()) .map(str::to_string) .collect() } #[test] fn parse_args_accepts_run_root_and_optional_flags() { let args = parse_args(&argv(&[ "--query-db", "query.db", "--run-root", "runs", "--repo-bytes-db", "repo-bytes.db", "--projection-entry-limit", "7", "--min-run-seq", "42", "--retain-indexed-runs", "3", "--dump-summary", ])) .expect("args"); assert_eq!(args.query_db, PathBuf::from("query.db")); assert_eq!(args.run_root.as_deref(), Some(std::path::Path::new("runs"))); assert_eq!(args.run_dir, None); assert_eq!( args.repo_bytes_db.as_deref(), Some(std::path::Path::new("repo-bytes.db")) ); assert_eq!(args.projection_entry_limit, 7); assert_eq!(args.min_run_seq, Some(42)); assert_eq!(args.retain_indexed_runs, Some(3)); assert!(args.dump_summary); } #[test] fn parse_args_accepts_single_run_dir() { let args = parse_args(&argv(&[ "--query-db", "query.db", "--run-dir", "runs/run_0001", ])) .expect("args"); assert_eq!( args.run_dir.as_deref(), Some(std::path::Path::new("runs/run_0001")) ); assert_eq!(args.run_root, None); assert_eq!(args.projection_entry_limit, 50); assert_eq!(args.min_run_seq, None); assert_eq!(args.retain_indexed_runs, Some(10)); } #[test] fn parse_args_rejects_invalid_or_ambiguous_input() { assert!( parse_args(&argv(&["--query-db", "query.db"])) .unwrap_err() .contains("exactly one") ); assert!( parse_args(&argv(&[ "--query-db", "query.db", "--run-root", "runs", "--run-dir", "runs/run_0001", ])) .unwrap_err() .contains("exactly one") ); assert!( parse_args(&argv(&["--run-root", "runs"])) .unwrap_err() .contains("--query-db") ); assert!( parse_args(&argv(&[ "--query-db", "query.db", "--run-root", "runs", "--projection-entry-limit", "bad", ])) .unwrap_err() .contains("invalid --projection-entry-limit") ); assert!( parse_args(&argv(&[ "--query-db", "query.db", "--run-root", "runs", "--unknown", ])) .unwrap_err() .contains("unknown argument") ); assert!( parse_args(&argv(&[ "--query-db", "query.db", "--run-root", "runs", "--retain-indexed-runs", "bad", ])) .unwrap_err() .contains("invalid --retain-indexed-runs") ); assert!( parse_args(&argv(&["--query-db"])) .unwrap_err() .contains("requires a value") ); } }