#![forbid(unsafe_code)]
use haematite::db::migrate::{Options as MigrateOptions, chunking as migrate_chunking};
use haematite::db::vacuum::{
Options as VacuumOptions, SweepProgress, run_with_progress as vacuum_sweep_with_progress,
stats as vacuum_stats,
};
use std::io::Write;
use std::path::PathBuf;
use std::process::ExitCode;
mod cli_usage;
mod render;
mod render_migration;
use cli_usage::USAGE;
fn main() -> ExitCode {
let mut args = std::env::args().skip(1);
let command = args.next();
#[cfg(unix)]
match command.as_deref() {
Some("call") => return haematite_cli::client::run(args),
Some("door") => return haematite_cli::door::cli::dispatch(args),
Some("assistant") => return haematite_cli::assistant::cli::run(args),
Some("mcp") => return haematite_cli::mcp::remote::dispatch(args),
Some("mcp-http") => return haematite_cli::mcp::http::cli::run(args),
Some("mcp-pass") => return haematite_cli::mcp::pass_header::run(args),
Some("write") => return haematite_cli::client::write::run(args),
Some("query") => return haematite_cli::client::query::run(args),
_ => {}
}
if command.as_deref() == Some("history-check") {
return haematite_cli::history_check::run(args);
}
if command.as_deref() == Some("observe-get") {
return haematite_cli::observe::run(args);
}
match run(command, args) {
Ok(()) => ExitCode::SUCCESS,
Err(message) => {
eprintln!("{message}");
ExitCode::FAILURE
}
}
}
fn run(command: Option<String>, args: impl Iterator<Item = String>) -> Result<(), String> {
let command = command.ok_or_else(|| USAGE.to_owned())?;
match command.as_str() {
"--help" | "-h" => {
let mut stdout = std::io::stdout().lock();
writeln!(stdout, "{USAGE}")
.and_then(|()| stdout.flush())
.map_err(|error| format!("failed to write or flush help: {error}"))
}
#[cfg(unix)]
"serve" => haematite_cli::serve::cli::run(args),
"lys-schema" => haematite_cli::serve::lys_schema::run(args),
#[cfg(unix)]
"install" => haematite_cli::install::run(args),
#[cfg(unix)]
"release" => haematite_cli::serve::cli::release(args),
#[cfg(unix)]
"branch-rename" => haematite_cli::serve::cli::branch_rename(args),
"stats" => {
let (options, json) = parse_stats_args(args)?;
let report = vacuum_stats(&options).map_err(|error| format!("refused: {error}"))?;
emit(&report, json)
}
"vacuum" => {
let (options, execute, json) = parse_vacuum_args(args)?;
run_vacuum(&options, execute, json)
}
"migrate-chunking" => {
let (options, json) = parse_migrate_args(args)?;
let report = migrate_chunking(&options).map_err(|error| format!("refused: {error}"))?;
emit_migration(&report, json)
}
other => Err(format!("unknown command {other:?}\n\n{USAGE}")),
}
}
fn run_vacuum(options: &VacuumOptions, execute: bool, json: bool) -> Result<(), String> {
let mut progress_error = None;
let outcome = vacuum_sweep_with_progress(options, execute, &mut |progress| {
if progress_error.is_none() {
progress_error = emit_sweep_progress(progress).err();
}
});
let report = match outcome {
Ok(report) => report,
Err(error) => {
let refusal = format!("refused: {error}");
return Err(match progress_error {
Some(progress) => format!("{refusal}; {progress}"),
None => refusal,
});
}
};
match (emit(&report, json), progress_error) {
(Ok(()), None) => Ok(()),
(Ok(()), Some(error)) | (Err(error), None) => Err(error),
(Err(report), Some(progress)) => Err(format!("{report}; {progress}")),
}
}
fn emit_sweep_progress(progress: SweepProgress) -> Result<(), String> {
if progress.removed_nodes != 0
&& !progress.removed_nodes.is_multiple_of(10_000)
&& progress.removed_nodes != progress.planned_nodes
{
return Ok(());
}
let mut stderr = std::io::stderr().lock();
let written = if progress.removed_nodes == 0 {
writeln!(
stderr,
"sweep planned {} nodes, {} bytes",
progress.planned_nodes, progress.planned_bytes
)
} else {
writeln!(
stderr,
"swept {} of {} nodes",
progress.removed_nodes, progress.planned_nodes
)
};
let flushed = written.and_then(|()| stderr.flush());
drop(stderr);
flushed.map_err(|error| format!("failed to write or flush sweep progress: {error}"))
}
fn emit_migration(
report: &haematite::db::migrate::MigrationReport,
json: bool,
) -> Result<(), String> {
if json {
let rendered = serde_json::to_string_pretty(report)
.map_err(|error| format!("failed to serialise report: {error}"))?;
println!("{rendered}");
} else {
print!(
"{}",
render_migration::render(report)
.map_err(|error| format!("failed to render report: {error}"))?
);
}
Ok(())
}
fn parse_migrate_args(
args: impl Iterator<Item = String>,
) -> Result<(MigrateOptions, bool), String> {
let mut data_dir: Option<PathBuf> = None;
let mut json = false;
let mut abort: Option<bool> = None;
for arg in args {
match arg.as_str() {
"--json" => json = true,
"--migrate" => set_migrate_mode(&mut abort, false)?,
"--abort" => set_migrate_mode(&mut abort, true)?,
flag if flag.starts_with("--") => {
return Err(format!("unknown flag {flag:?}\n\n{USAGE}"));
}
positional => {
if data_dir.is_some() {
return Err(format!("unexpected argument {positional:?}\n\n{USAGE}"));
}
data_dir = Some(PathBuf::from(positional));
}
}
}
let data_dir = data_dir.ok_or_else(|| format!("missing <data_dir>\n\n{USAGE}"))?;
let abort = abort.ok_or_else(|| {
format!(
"haematite migrate-chunking requires exactly one of --migrate or --abort\n\n{USAGE}"
)
})?;
let options = if abort {
MigrateOptions::abort(data_dir)
} else {
MigrateOptions::migrate(data_dir)
};
Ok((options, json))
}
fn set_migrate_mode(abort: &mut Option<bool>, value: bool) -> Result<(), String> {
if abort.is_some() {
return Err(format!(
"--migrate and --abort are mutually exclusive\n\n{USAGE}"
));
}
*abort = Some(value);
Ok(())
}
fn emit(report: &haematite::db::vacuum::Report, json: bool) -> Result<(), String> {
if json {
let rendered = serde_json::to_string_pretty(report)
.map_err(|error| format!("failed to serialise report: {error}"))?;
println!("{rendered}");
} else {
print!(
"{}",
render::render(report).map_err(|error| format!("failed to render report: {error}"))?
);
}
Ok(())
}
fn parse_vacuum_args(
args: impl Iterator<Item = String>,
) -> Result<(VacuumOptions, bool, bool), String> {
let mut data_dir: Option<PathBuf> = None;
let mut refs_dirs = Vec::new();
let mut snapshot_files = Vec::new();
let mut attest = false;
let mut json = false;
let mut execute: Option<bool> = None;
let mut args = args.peekable();
while let Some(arg) = args.next() {
match arg.as_str() {
"--json" => json = true,
"--attest-metadata-complete" => attest = true,
"--dry-run" => set_mode(&mut execute, false)?,
"--sweep" => set_mode(&mut execute, true)?,
"--refs-dir" => refs_dirs.push(required_value(&mut args, "--refs-dir")?),
"--snapshot-file" => {
snapshot_files.push(required_value(&mut args, "--snapshot-file")?);
}
flag if flag.starts_with("--") => {
return Err(format!("unknown flag {flag:?}\n\n{USAGE}"));
}
positional => {
if data_dir.is_some() {
return Err(format!("unexpected argument {positional:?}\n\n{USAGE}"));
}
data_dir = Some(PathBuf::from(positional));
}
}
}
let data_dir = data_dir.ok_or_else(|| format!("missing <data_dir>\n\n{USAGE}"))?;
let execute = execute.ok_or_else(|| {
format!("haematite vacuum requires exactly one of --dry-run or --sweep\n\n{USAGE}")
})?;
Ok((
VacuumOptions {
data_dir,
refs_dirs,
snapshot_files,
attest_metadata_complete: attest,
},
execute,
json,
))
}
fn set_mode(execute: &mut Option<bool>, value: bool) -> Result<(), String> {
if execute.is_some() {
return Err(format!(
"--dry-run and --sweep are mutually exclusive\n\n{USAGE}"
));
}
*execute = Some(value);
Ok(())
}
fn parse_stats_args(args: impl Iterator<Item = String>) -> Result<(VacuumOptions, bool), String> {
let mut data_dir: Option<PathBuf> = None;
let mut refs_dirs = Vec::new();
let mut snapshot_files = Vec::new();
let mut json = false;
let mut args = args.peekable();
while let Some(arg) = args.next() {
match arg.as_str() {
"--json" => json = true,
"--refs-dir" => refs_dirs.push(required_value(&mut args, "--refs-dir")?),
"--snapshot-file" => {
snapshot_files.push(required_value(&mut args, "--snapshot-file")?);
}
flag if flag.starts_with("--") => {
return Err(format!("unknown flag {flag:?}\n\n{USAGE}"));
}
positional => {
if data_dir.is_some() {
return Err(format!("unexpected argument {positional:?}\n\n{USAGE}"));
}
data_dir = Some(PathBuf::from(positional));
}
}
}
let data_dir = data_dir.ok_or_else(|| format!("missing <data_dir>\n\n{USAGE}"))?;
Ok((
VacuumOptions {
data_dir,
refs_dirs,
snapshot_files,
attest_metadata_complete: false,
},
json,
))
}
fn required_value(
args: &mut std::iter::Peekable<impl Iterator<Item = String>>,
flag: &str,
) -> Result<PathBuf, String> {
match args.next() {
Some(value) if value.starts_with("--") => Err(format!(
"{flag} requires a path, found {value:?}\n\n{USAGE}"
)),
Some(value) => Ok(PathBuf::from(value)),
None => Err(format!("{flag} requires a path\n\n{USAGE}")),
}
}
#[cfg(test)]
mod tests {
use super::{parse_migrate_args, parse_stats_args};
use haematite::db::migrate::Mode as MigrateMode;
fn args(list: &[&str]) -> impl Iterator<Item = String> + use<> {
list.iter()
.map(|arg| (*arg).to_owned())
.collect::<Vec<_>>()
.into_iter()
}
#[test]
fn parses_paths_and_json() -> Result<(), String> {
let (options, json) = parse_stats_args(args(&[
"/data",
"--refs-dir",
"/refs",
"--snapshot-file",
"/snaps.hsr",
"--json",
]))?;
assert_eq!(options.data_dir, std::path::Path::new("/data"));
assert_eq!(options.refs_dirs, vec![std::path::PathBuf::from("/refs")]);
assert_eq!(
options.snapshot_files,
vec![std::path::PathBuf::from("/snaps.hsr")]
);
assert!(!options.attest_metadata_complete, "stats never attests");
assert!(json);
Ok(())
}
#[test]
fn a_following_flag_is_not_a_path() {
let result = parse_stats_args(args(&["/data", "--refs-dir", "--json"]));
assert!(result.is_err(), "--json must not become a refs dir");
}
#[test]
fn missing_value_missing_dir_and_extras_are_rejected() {
assert!(parse_stats_args(args(&["/data", "--refs-dir"])).is_err());
assert!(parse_stats_args(args(&["--json"])).is_err(), "no data_dir");
assert!(parse_stats_args(args(&["/a", "/b"])).is_err(), "two dirs");
assert!(parse_stats_args(args(&["/a", "--bogus"])).is_err());
}
#[test]
fn flags_may_precede_the_positional() -> Result<(), String> {
let (options, json) = parse_stats_args(args(&["--json", "/data"]))?;
assert_eq!(options.data_dir, std::path::Path::new("/data"));
assert!(json);
Ok(())
}
#[test]
fn migrate_parses_direction_and_json() -> Result<(), String> {
let (options, json) = parse_migrate_args(args(&["/data", "--migrate", "--json"]))?;
assert_eq!(options.data_dir, std::path::Path::new("/data"));
assert_eq!(options.mode, MigrateMode::Migrate);
assert!(json);
let (abort, _) = parse_migrate_args(args(&["/data", "--abort"]))?;
assert_eq!(abort.mode, MigrateMode::Abort);
Ok(())
}
#[test]
fn migrate_requires_exactly_one_direction() {
assert!(
parse_migrate_args(args(&["/data"])).is_err(),
"no direction"
);
assert!(
parse_migrate_args(args(&["/data", "--migrate", "--abort"])).is_err(),
"both directions"
);
assert!(
parse_migrate_args(args(&["--migrate"])).is_err(),
"no data_dir"
);
assert!(parse_migrate_args(args(&["/data", "--migrate", "--bogus"])).is_err());
}
}