refactor: name the modules after what they hold
This commit is contained in:
296
src/commands/postgres.rs
Normal file
296
src/commands/postgres.rs
Normal 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user