refactor: name the modules after what they hold

This commit is contained in:
2026-09-08 12:10:41 +00:00
parent a3c08f6eaa
commit bb0a2cca44
12 changed files with 1014 additions and 989 deletions

296
src/commands/postgres.rs Normal file
View File

@@ -0,0 +1,296 @@
mod server;
mod shape;
use fs_err::File;
use std::io::{self, Write};
use std::path::Path;
use std::process::Stdio;
use anyhow::{Context, Result, bail};
use self::server::{Database, wait_until_ready, when_ready};
use self::shape::{Dump, HEADER_LEN, Kind};
use crate::cli::Format;
use crate::cmd::{
Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Rm, Start, Stop,
Up,
};
use crate::ctx::Ctx;
use crate::fsops::{remove_file, rename, suffixed};
use crate::output::note;
// unique per run: docker cp will not copy a directory over an existing path, and
// quietly leaves whatever was there for pg_restore to read instead
fn remote_dump() -> String {
format!("/tmp/ahab-dump-{}", std::process::id())
}
fn is_existing_role_error(line: &str) -> bool {
line.starts_with("ERROR:") && line.contains("role \"") && line.ends_with("already exists")
}
fn restore_cluster(ctx: &Ctx, db: &Database, restore: &dyn Cmd, file: &Path) -> Result<()> {
let out = restore
.in_container(&db.container)
.interactive()
.stdin_from_captured(ctx, file)
.context("running psql")?;
io::stdout().write_all(&out.stdout).ok();
let mut existing = 0;
for line in String::from_utf8_lossy(&out.stderr).lines() {
if is_existing_role_error(line) {
existing += 1;
continue;
}
note!("{line}");
}
if existing > 0 {
note!(
"left {existing} existing role{} alone",
if existing == 1 { "" } else { "s" }
);
}
if !out.status.success() {
bail!("psql failed, the database is left empty");
}
Ok(())
}
pub fn import(ctx: &Ctx, file: &Path) -> Result<()> {
let dump = Dump::of(file)?;
let db = Database::resolve(ctx)?;
note!("stopping all containers");
Stop.run(ctx)?;
note!("starting db container");
Start::service(&db.service).run(ctx)?;
let remote = remote_dump();
// a directory cannot be streamed, so it is the one shape that gets copied in
if matches!(dump, Dump::Directory) {
when_ready(
ctx,
&db,
&Cp::into_container(&file.to_string_lossy(), &db.container, &remote),
)?;
}
let kind = match &dump {
Dump::Directory => Kind::Archive,
Dump::Header(header) => Kind::of(header),
Dump::Gzip => {
wait_until_ready(ctx, &db)?;
let out = Gunzip
.pipe(&Head::bytes(HEADER_LEN))
// head closes the pipe once it has its bytes, which kills gunzip
.allow_early_close()
.in_container(&db.container)
.interactive()
.stdin_from_captured(ctx, file)
.context("reading the compressed dump's header")?;
// head exits 0 whatever gunzip did, so an empty header is the only
// sign that decompression failed
if !out.status.success() || out.stdout.is_empty() {
let reason = String::from_utf8_lossy(&out.stderr);
let reason = reason.trim();
bail!(
"could not decompress {}: {}",
file.display(),
if reason.is_empty() {
"no data came out"
} else {
reason
}
);
}
Kind::of(&out.stdout)
}
};
let restore = db.restore_with(kind);
// the name of what actually runs, so the message cannot drift from it
let tool = restore.argv().program().to_string();
note!("restoring database with {tool}");
when_ready(
ctx,
&db,
&DropDb {
username: &db.user,
dbname: &db.name,
}
.in_container(&db.container),
)?;
// a cluster dump creates the database itself, and would trip over one that
// is already there
if kind != Kind::Cluster {
when_ready(
ctx,
&db,
&CreateDb {
username: &db.user,
dbname: &db.name,
}
.in_container(&db.container),
)?;
}
wait_until_ready(ctx, &db)?;
if ctx.dry_run {
note!("would restore with {tool}");
} else {
// a directory dump was copied in whole, so pg_restore reads it from the
// container; every other shape is fed in on stdin, through gunzip when it
// arrives compressed
if matches!(dump, Dump::Directory) {
let status = PgRestore::new(&db.user, &db.name)
.from(&remote)
.in_container(&db.container)
.status(ctx)?;
if !status.success() {
bail!("{tool} failed, the database is left empty");
}
} else {
let restore: Box<dyn Cmd> = match dump {
Dump::Gzip => Box::new(Gunzip.pipe(&*restore)),
_ => restore,
};
if kind == Kind::Cluster {
// psql's output is read rather than streamed here, to keep the
// expected role errors out of the way
restore_cluster(ctx, &db, &*restore, file)?;
} else {
let status = restore
.in_container(&db.container)
.interactive()
.stdin_from(ctx, file)?;
if !status.success() {
bail!("{tool} failed, the database is left empty");
}
}
}
}
if matches!(dump, Dump::Directory) {
let _ = Rm::recursive(&remote).in_container(&db.container).run(ctx);
}
note!("restarting containers");
Stop.run(ctx)?;
Up.run(ctx)?;
Ok(())
}
pub fn dump(ctx: &Ctx, file: &Path, format: Format, gzip: bool) -> Result<()> {
let db = Database::resolve(ctx)?;
if format == Format::Directory {
if gzip {
bail!("a directory dump is a directory of already compressed files, not a stream");
}
return dump_directory(ctx, &db, file);
}
note!("dumping to local file {}", file.to_string_lossy());
// written beside the target and renamed once the dump succeeds, so a failure
// cannot destroy the dump that is already there
let partial = suffixed(file, ".partial");
// a dry run produces no dump, so it must not lay a hand on the target either
let stdout = if ctx.dry_run {
Stdio::null()
} else {
Stdio::from(std::fs::File::from(File::create(&partial)?))
};
let dumping = dump_command(&db, format);
let dumped = if gzip {
dumping
.pipe(&Gzip)
.in_container(&db.container)
.stream_to(ctx, stdout)
} else {
dumping.in_container(&db.container).stream_to(ctx, stdout)
};
if let Err(e) = dumped {
let _ = remove_file(ctx, &partial);
return Err(e);
}
rename(ctx, &partial, file)?;
Ok(())
}
// pg_dump for one database, pg_dumpall for a cluster, which has no format letter
fn dump_command(db: &Database, format: Format) -> Box<dyn Cmd + '_> {
match format.flag() {
Some(flag) => Box::new(PgDump::new(&db.user, &db.name, flag)),
None => Box::new(PgDumpAll { username: &db.user }),
}
}
// pg_dump writes a directory format dump itself rather than to stdout, so it lands
// in the container and comes back with docker cp
fn dump_directory(ctx: &Ctx, db: &Database, target: &Path) -> Result<()> {
if target.exists() {
bail!(
"{} already exists; a directory dump will not be written over it",
target.display()
);
}
note!("dumping to local directory {}", target.display());
let remote = remote_dump();
PgDump::new(&db.user, &db.name, "d")
.to(&remote)
.in_container(&db.container)
.run(ctx)?;
let copied = Cp::out_of_container(&db.container, &remote, &target.to_string_lossy()).run(ctx);
let _ = Rm::recursive(&remote).in_container(&db.container).run(ctx);
copied
}
#[cfg(test)]
mod tests {
use super::is_existing_role_error;
#[test]
fn only_the_existing_role_complaint_is_expected() {
assert!(is_existing_role_error(
r#"ERROR: role "postgres" already exists"#
));
for line in [
r#"ERROR: relation "things" already exists"#,
r#"ERROR: database "db" already exists"#,
"ERROR: syntax error at or near \"slect\"",
"NOTICE: role \"postgres\" already exists",
] {
assert!(!is_existing_role_error(line), "{line}");
}
}
}