rpki/src/bin/rpki_query_indexer.rs

279 lines
8.6 KiB
Rust

use std::path::PathBuf;
use rpki::query_db::{ArtifactIndexerConfig, index_artifacts};
#[derive(Debug, PartialEq, Eq)]
struct Args {
query_db: PathBuf,
run_root: Option<PathBuf>,
run_dir: Option<PathBuf>,
repo_bytes_db: Option<PathBuf>,
projection_entry_limit: usize,
min_run_seq: Option<u64>,
retain_indexed_runs: Option<usize>,
dump_summary: bool,
}
fn usage() -> &'static str {
"Usage: rpki_query_indexer --query-db <path> (--run-root <path>|--run-dir <path>) [--repo-bytes-db <path>] [--projection-entry-limit <n>] [--min-run-seq <n>] [--retain-indexed-runs <n>] [--dump-summary]"
}
fn parse_args(argv: &[String]) -> Result<Args, String> {
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::<usize>()
.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::<u64>()
.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::<usize>()
.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::<Vec<_>>();
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<String> {
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")
);
}
}