diff --git a/README.md b/README.md index e86031f..dc094f9 100644 --- a/README.md +++ b/README.md @@ -83,10 +83,23 @@ ahab django make-command ## postgres ```bash -ahab postgres dump # pg_dump, custom format -ahab postgres import # drop, create, pg_restore +ahab postgres dump # pg_dump, custom format +ahab postgres dump -F plain # pg_dump, plain sql +ahab postgres import # drop, create, then restore ``` +The format of a dump being imported is read from the file rather than its name. +A custom format dump starts with `PGDMP` and a tar one with `toc.dat`, both of +which go to `pg_restore`, as does a directory produced by `pg_dump -Fd`. +Anything else is treated as sql and fed to `psql` with `ON_ERROR_STOP` and +`--single-transaction`, so a bad file rolls back instead of half applying. +A whole cluster dump from `pg_dumpall` is recognised by its header and handled +differently again: it creates its own databases and carries role statements, so +the database is dropped but not recreated, the dump goes to `psql` connected to +`postgres`, and it runs without `ON_ERROR_STOP` because roles that already exist +report errors that are expected. Gzipped dumps are decompressed on the way in, +whichever of the three they hold. + ## link `ahab link` moves untracked paths out of the repository into an out-of-repo diff --git a/src/cli/mod.rs b/src/cli/mod.rs index a4fd174..65361e1 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -6,4 +6,4 @@ mod postgres; pub use ahab::{Ahab, Commands}; pub use django::Django; pub use link::Link; -pub use postgres::Postgres; +pub use postgres::{Format, Postgres}; diff --git a/src/cli/postgres.rs b/src/cli/postgres.rs index 7924e80..5047b20 100644 --- a/src/cli/postgres.rs +++ b/src/cli/postgres.rs @@ -1,12 +1,39 @@ use std::path::PathBuf; -use clap::Subcommand; +use clap::{Subcommand, ValueEnum}; #[derive(Subcommand, Debug)] pub enum Postgres { - /// Import dump via pg_restore + /// Import a dump, custom format or plain sql Import { path: PathBuf }, - /// Dump via pg_dump with format=c - Dump { path: PathBuf }, + /// Dump via pg_dump + Dump { + path: PathBuf, + + /// Dump format + #[arg(short = 'F', long, value_enum, default_value_t = Format::Custom)] + format: Format, + }, +} + +#[derive(ValueEnum, Clone, Copy, Debug, Default)] +pub enum Format { + /// pg_dump --format=c, restored with pg_restore + #[default] + #[value(alias = "c")] + Custom, + + /// pg_dump --format=p, restored with psql + #[value(alias = "p")] + Plain, +} + +impl Format { + pub fn flag(&self) -> &'static str { + match self { + Self::Custom => "c", + Self::Plain => "p", + } + } } diff --git a/src/command_builder.rs b/src/command_builder.rs index e57697d..b7a9d78 100644 --- a/src/command_builder.rs +++ b/src/command_builder.rs @@ -46,6 +46,11 @@ impl CommandBuilder { Self::default().args("docker compose") } + pub fn arg(mut self, arg: impl Into) -> Self { + self.args.push(arg.into()); + self + } + pub fn args(mut self, args: T) -> Self where Args: From, diff --git a/src/main.rs b/src/main.rs index b6a10ed..38d7def 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,7 @@ fn main() -> Result<()> { }, cli::Commands::Postgres { command } => match command { cli::Postgres::Import { path } => scripts::postgres::import(&path), - cli::Postgres::Dump { path } => scripts::postgres::dump(&path), + cli::Postgres::Dump { path, format } => scripts::postgres::dump(&path, format), }, cli::Commands::Link { command } => match command { cli::Link::Add { paths, force } => scripts::link::add(&paths, force), diff --git a/src/scripts/postgres.rs b/src/scripts/postgres.rs index b817525..7408cfc 100644 --- a/src/scripts/postgres.rs +++ b/src/scripts/postgres.rs @@ -1,6 +1,7 @@ -use anyhow::{Result, anyhow}; +use anyhow::{Context, Result, anyhow, bail}; use std::{ fs::File, + io::Read, path::{Path, PathBuf}, process::Stdio, thread, @@ -8,8 +9,22 @@ use std::{ }; use super::docker_compose; +use crate::cli::Format; use crate::{command_builder::CommandBuilder, compose::Compose, debug_eprintln}; +const CUSTOM_MAGIC: &[u8] = b"PGDMP"; +const TAR_MAGIC: &[u8] = b"toc.dat"; +const GZIP_MAGIC: &[u8] = b"\x1f\x8b"; +const CLUSTER_MARKER: &str = "PostgreSQL database cluster dump"; +// wide enough for the cluster marker, which sits a few bytes into the file +const HEADER_LEN: usize = 512; + +// 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()) +} + struct Database { service: String, container: String, @@ -41,11 +56,124 @@ impl Database { name, }) } + + fn restore_with(&self, kind: Kind) -> String { + match kind { + Kind::Archive => format!("pg_restore -U {} --dbname={}", self.user, self.name), + Kind::Sql => format!( + "psql -q -U {} -d {} -v ON_ERROR_STOP=1 --single-transaction", + self.user, self.name + ), + // a cluster dump connects to each database itself, and stopping at the + // first error would stop at a role that is already there + Kind::Cluster => format!("psql -q -U {} -d postgres", self.user), + } + } +} + +enum Dump { + Directory, + Gzip, + Header(Vec), +} + +impl Dump { + fn of(path: &Path) -> Result { + if path.is_dir() { + return Ok(Self::Directory); + } + + let header = read_header(path)?; + if header.starts_with(GZIP_MAGIC) { + return Ok(Self::Gzip); + } + + Ok(Self::Header(header)) + } +} + +// how the dump has to be fed back in: pg_restore for an archive, psql into the +// database for a single database dump, psql into postgres for a whole cluster +#[derive(Clone, Copy, PartialEq)] +enum Kind { + Archive, + Sql, + Cluster, +} + +impl Kind { + fn of(header: &[u8]) -> Self { + if header.starts_with(CUSTOM_MAGIC) || header.starts_with(TAR_MAGIC) { + return Self::Archive; + } + + if String::from_utf8_lossy(header).contains(CLUSTER_MARKER) { + return Self::Cluster; + } + + Self::Sql + } + + fn tool(&self) -> &'static str { + match self { + Self::Archive => "pg_restore", + Self::Sql | Self::Cluster => "psql", + } + } +} + +fn read_header(path: &Path) -> Result> { + let mut file = File::open(path).with_context(|| format!("opening {}", path.display()))?; + let mut header = vec![0; HEADER_LEN]; + + let read = file + .read(&mut header) + .with_context(|| format!("reading {}", path.display()))?; + header.truncate(read); + + Ok(header) +} + +fn piped(db: &Database, script: &str, input: &Path) -> Result { + let file = File::open(input).with_context(|| format!("opening {}", input.display()))?; + let mut command = CommandBuilder::docker() + .args("exec -i") + .args(&db.container) + .args("sh -c") + .arg(script) + .build()?; + + command.stdin(Stdio::from(file)); + Ok(command) +} + +fn wait_until_ready(db: &Database) -> Result<()> { + debug_eprintln!("waiting until pg_isready"); + while !CommandBuilder::docker() + .args(&format!( + "exec {} pg_isready -U {} -d {}", + db.container, db.user, db.name + )) + .build()? + .stdout(Stdio::null()) + .spawn()? + .wait()? + .success() + { + thread::sleep(Duration::from_secs(1)); + } + + Ok(()) +} + +fn run_when_ready(db: &Database, command: &str) -> Result<()> { + wait_until_ready(db)?; + CommandBuilder::docker().args(command).exec() } pub fn import(file: &Path) -> Result<()> { + let dump = Dump::of(file)?; let db = Database::resolve()?; - let dump_file = file.to_string_lossy(); eprintln!("stopping all containers"); docker_compose::stop()?; @@ -53,36 +181,100 @@ pub fn import(file: &Path) -> Result<()> { eprintln!("starting db container"); docker_compose::start(Some(&db.service))?; - eprintln!("restoring database"); - let commands = [ - format!("cp -L {dump_file} {}:/tmp/dbdump", db.container), - format!("exec {} dropdb -U {} {}", db.container, db.user, db.name), - format!( - "exec {} createdb -U {} -E utf8 -T template0 {}", - db.container, db.user, db.name - ), - format!( - "exec {} pg_restore -U {} --dbname={} /tmp/dbdump", - db.container, db.user, db.name - ), - ]; + let remote = remote_dump(); - for command in commands { - debug_eprintln!("waiting until pg_isready"); - while !CommandBuilder::docker() + // a directory cannot be streamed, so it is the one shape that gets copied in + if matches!(dump, Dump::Directory) { + run_when_ready( + &db, + &format!("cp -L {} {}:{remote}", file.display(), db.container), + )?; + } + + let (kind, restore) = match &dump { + Dump::Directory => (Kind::Archive, None), + Dump::Header(header) => { + let kind = Kind::of(header); + + (kind, Some(db.restore_with(kind))) + } + Dump::Gzip => { + wait_until_ready(&db)?; + let out = piped(&db, &format!("gunzip -c | head -c {HEADER_LEN}"), file)? + .output() + .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 + } + ); + } + + let kind = Kind::of(&out.stdout); + + (kind, Some(format!("gunzip -c | {}", db.restore_with(kind)))) + } + }; + + let tool = kind.tool(); + if kind == Kind::Cluster { + eprintln!( + "restoring a whole cluster dump with psql; roles that already exist \ + will report errors" + ); + } else { + eprintln!("restoring database with {tool}"); + } + + run_when_ready( + &db, + &format!("exec {} dropdb -U {} {}", db.container, db.user, db.name), + )?; + + // a cluster dump creates the database itself, and would trip over one that + // is already there + if kind != Kind::Cluster { + run_when_ready( + &db, + &format!( + "exec {} createdb -U {} -E utf8 -T template0 {}", + db.container, db.user, db.name + ), + )?; + } + + wait_until_ready(&db)?; + let status = match &restore { + Some(script) => piped(&db, script, file)?.spawn()?.wait()?, + None => CommandBuilder::docker() .args(&format!( - "exec {} pg_isready -U {} -d {}", + "exec {} pg_restore -U {} --dbname={} {remote}", db.container, db.user, db.name )) .build()? - .stdout(Stdio::null()) .spawn()? - .wait()? - .success() - { - thread::sleep(Duration::from_secs(1)); - } - CommandBuilder::docker().args(&command).exec()?; + .wait()?, + }; + + if !status.success() { + bail!("{tool} failed, the database is left empty"); + } + + if matches!(dump, Dump::Directory) { + let _ = CommandBuilder::docker() + .args(&format!("exec {} rm -rf {remote}", db.container)) + .exec(); } eprintln!("restarting containers"); @@ -92,7 +284,7 @@ pub fn import(file: &Path) -> Result<()> { Ok(()) } -pub fn dump(file: &PathBuf) -> Result<()> { +pub fn dump(file: &PathBuf, format: Format) -> Result<()> { let db = Database::resolve()?; eprintln!("dumping to local file {}", file.to_string_lossy()); @@ -101,8 +293,11 @@ pub fn dump(file: &PathBuf) -> Result<()> { let stdout = Stdio::from(file); let command = format!( - "exec {} pg_dump -U {} --format=c {}", - db.container, db.user, db.name + "exec {} pg_dump -U {} --format={} {}", + db.container, + db.user, + format.flag(), + db.name ); CommandBuilder::docker() .args(&command) @@ -110,3 +305,39 @@ pub fn dump(file: &PathBuf) -> Result<()> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::{GZIP_MAGIC, Kind}; + + #[test] + fn an_archive_is_recognised_by_its_magic() { + assert!(matches!(Kind::of(b"PGDMP\x01\x0f"), Kind::Archive)); + assert!(matches!(Kind::of(b"toc.dat\x00\x00"), Kind::Archive)); + } + + #[test] + fn a_cluster_dump_is_recognised_by_its_header() { + let header = b"--\n-- PostgreSQL database cluster dump\n--\n\n\\restrict abc\n"; + assert!(matches!(Kind::of(header), Kind::Cluster)); + } + + #[test] + fn anything_else_is_a_single_database_dump() { + for header in [ + &b"--\n-- PostgreSQL database dump\n"[..], + &b"BEGIN;"[..], + &b""[..], + &b"PGD"[..], + &b"toc.da"[..], + ] { + assert!(matches!(Kind::of(header), Kind::Sql), "{header:?}"); + } + } + + #[test] + fn gzip_is_recognised_by_its_magic() { + assert!(b"\x1f\x8b\x08\x00rest".starts_with(GZIP_MAGIC)); + assert!(!b"PGDMP".starts_with(GZIP_MAGIC)); + } +}