use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;
use anyhow::{bail, Context, Result};
pub const RUNNABLE: &[&str] = &[
"purge",
"census",
"bulk-purge",
"orphan-sweep",
"credentials",
"proofs",
"oob",
"basic-messages",
"question-answer",
"tenants",
];
pub struct Args<'a> {
pub envs: &'a [PathBuf],
pub dir: Option<&'a Path>,
pub run: &'a str,
pub stop_on_error: bool,
pub server_side_delete: bool,
pub no_audit: bool,
}
fn env_files_in(dir: &Path) -> Result<Vec<PathBuf>> {
let mut files: Vec<PathBuf> = std::fs::read_dir(dir)
.with_context(|| format!("could not read {}", dir.display()))?
.filter_map(|e| e.ok().map(|e| e.path()))
.filter(|p| p.is_file() && p.extension().is_some_and(|x| x == "env"))
.collect();
files.sort();
Ok(files)
}
fn identity(path: &Path) -> Result<(String, String)> {
let vars: HashMap<String, String> = dotenvy::from_path_iter(path)
.with_context(|| format!("could not read env file {}", path.display()))?
.collect::<std::result::Result<_, _>>()
.with_context(|| format!("could not parse env file {}", path.display()))?;
let get = |k: &str| vars.get(k).cloned().filter(|v| !v.is_empty());
let Some(wallet) = get("WALLET_ID") else {
bail!("{} has no WALLET_ID", path.display());
};
Ok((wallet, get("DB_HOST").unwrap_or_default()))
}
struct Outcome {
file: PathBuf,
wallet: String,
status: String,
secs: u64,
ok: bool,
}
pub fn run(args: Args<'_>) -> Result<std::process::ExitCode> {
if !RUNNABLE.contains(&args.run) {
bail!("`batch --run {}` is not supported; choose one of: {}", args.run, RUNNABLE.join(", "));
}
let mut files: Vec<PathBuf> = args.envs.to_vec();
if let Some(dir) = args.dir {
files.extend(env_files_in(dir)?);
}
if files.is_empty() {
bail!("no env files: pass --env <FILE>... and/or --dir <DIR> (a directory of *.env files)");
}
let mut seen: HashMap<(String, String), PathBuf> = HashMap::new();
let mut agents = Vec::new();
for f in &files {
let (wallet, host) = identity(f)?;
if let Some(first) = seen.insert((wallet.clone(), host), f.clone()) {
bail!(
"{} and {} point at the same wallet ({wallet}) on the same database host; \
refusing to drain one agent twice",
first.display(),
f.display()
);
}
agents.push((f.clone(), wallet));
}
let exe = std::env::current_exe().context("could not locate the owpurge binary")?;
println!("\nbatch: `{}` across {} agent(s)", args.run, agents.len());
let mut outcomes: Vec<Outcome> = Vec::new();
for (i, (file, wallet)) in agents.iter().enumerate() {
println!("\n{}", "#".repeat(60));
println!("# agent {}/{}: {} (wallet {wallet})", i + 1, agents.len(), file.display());
println!("{}", "#".repeat(60));
let mut cmd = Command::new(&exe);
cmd.env_clear();
if args.server_side_delete {
cmd.arg("--server-side-delete");
}
if args.no_audit {
cmd.arg("--no-audit");
}
cmd.arg("--env-file").arg(file).arg(args.run);
let started = Instant::now();
let (ok, status) = match cmd.status() {
Ok(s) if s.success() => (true, "ok".to_string()),
Ok(s) => (false, format!("failed ({s})")),
Err(e) => (false, format!("could not start: {e}")),
};
outcomes.push(Outcome {
file: file.clone(),
wallet: wallet.clone(),
status,
secs: started.elapsed().as_secs(),
ok,
});
if !ok && args.stop_on_error {
eprintln!("\nstopping: --stop-on-error and this agent failed");
break;
}
}
println!("\n{}", "=".repeat(60));
println!("batch summary ({})", args.run);
println!("{}", "=".repeat(60));
for o in &outcomes {
println!(
" {:<8} {:>5}s {} [{}]",
if o.ok { "ok" } else { "FAILED" },
o.secs,
o.file.display(),
o.wallet
);
if !o.ok {
println!(" {}", o.status);
}
}
let not_run = agents.len() - outcomes.len();
if not_run > 0 {
println!(" {not_run} agent(s) not run (stopped early)");
}
let failed = outcomes.iter().filter(|o| !o.ok).count();
if failed > 0 || not_run > 0 {
eprintln!("\n{failed} agent(s) failed.");
return Ok(std::process::ExitCode::FAILURE);
}
println!("\nAll {} agent(s) completed.", outcomes.len());
Ok(std::process::ExitCode::SUCCESS)
}