279 lines
8.6 KiB
Rust
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")
|
|
);
|
|
}
|
|
}
|