diff --git a/README.md b/README.md index e86031f..51d31af 100644 --- a/README.md +++ b/README.md @@ -83,10 +83,27 @@ 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 dump -F tar # pg_dump, tar +ahab postgres dump -F directory # pg_dump, a directory of files +ahab postgres dump -F cluster # pg_dumpall, roles and all databases +ahab postgres dump -F plain -z # any of them, gzipped +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..e0a90b3 100644 --- a/src/cli/postgres.rs +++ b/src/cli/postgres.rs @@ -1,12 +1,59 @@ 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, in any format ahab can produce Import { path: PathBuf }, - /// Dump via pg_dump with format=c - Dump { path: PathBuf }, + /// Dump via pg_dump, or pg_dumpall for a whole cluster + Dump { + path: PathBuf, + + /// Dump format + #[arg(short = 'F', long, value_enum, default_value_t = Format::Custom)] + format: Format, + + /// Compress the dump with gzip + #[arg(short = 'z', long)] + gzip: bool, + }, +} + +#[derive(ValueEnum, Clone, Copy, Debug, Default, PartialEq)] +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, + + /// pg_dump --format=t, restored with pg_restore + #[value(alias = "t")] + Tar, + + /// pg_dump --format=d, a directory rather than a file + #[value(alias = "d")] + Directory, + + /// pg_dumpall, the whole cluster including roles + #[value(alias = "all")] + Cluster, +} + +impl Format { + // pg_dump's -F letter, which a cluster dump does not have: that is pg_dumpall + pub fn flag(&self) -> Option<&'static str> { + match self { + Self::Custom => Some("c"), + Self::Plain => Some("p"), + Self::Tar => Some("t"), + Self::Directory => Some("d"), + Self::Cluster => None, + } + } } 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..f70e2e4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,9 @@ 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, gzip } => { + scripts::postgres::dump(&path, format, gzip) + } }, 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..b0fcbcd 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::{self, Read, Write}, 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,159 @@ 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 -o /dev/null -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 -o /dev/null -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", + } + } +} + +// pg_dumpall recreates roles the cluster already has, so this is the expected +// shape of a working restore rather than a problem +fn is_existing_role_error(line: &str) -> bool { + line.starts_with("ERROR:") && line.contains("role \"") && line.ends_with("already exists") +} + +fn restore_cluster(db: &Database, script: &str, file: &Path) -> Result<()> { + let out = piped(db, script, file)?.output().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; + } + + eprintln!("{line}"); + } + + if existing > 0 { + eprintln!( + "left {existing} existing role{} alone", + if existing == 1 { "" } else { "s" } + ); + } + + if !out.status.success() { + bail!("psql failed, the database is left empty"); + } + + Ok(()) +} + +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 +216,99 @@ 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() - .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)); + // 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))) } - CommandBuilder::docker().args(&command).exec()?; + 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(); + 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)?; + match (kind, restore.as_deref()) { + // psql's output is read rather than streamed here, to keep the expected + // role errors out of the way + (Kind::Cluster, Some(script)) => restore_cluster(&db, script, file)?, + (_, restore) => { + let status = match restore { + Some(script) => piped(&db, script, file)?.spawn()?.wait()?, + None => CommandBuilder::docker() + .args(&format!( + "exec {} pg_restore -U {} --dbname={} {remote}", + db.container, db.user, db.name + )) + .build()? + .spawn()? + .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,21 +318,129 @@ pub fn import(file: &Path) -> Result<()> { Ok(()) } -pub fn dump(file: &PathBuf) -> Result<()> { +pub fn dump(file: &PathBuf, format: Format, gzip: bool) -> Result<()> { let db = Database::resolve()?; + if format == Format::Directory { + if gzip { + bail!("a directory dump is a directory of already compressed files, not a stream"); + } + + return dump_directory(&db, file); + } + eprintln!("dumping to local file {}", file.to_string_lossy()); - let file = File::create(file)?; - let stdout = Stdio::from(file); + let stdout = Stdio::from(File::create(file)?); + let dumping = dump_command(&db, format); - let command = format!( - "exec {} pg_dump -U {} --format=c {}", - db.container, db.user, db.name - ); - CommandBuilder::docker() - .args(&command) - .exec_redirect_stdout(stdout)?; + if gzip { + CommandBuilder::docker() + .args("exec") + .args(&db.container) + .args("sh -c") + .arg(format!("{dumping} | gzip")) + .exec_redirect_stdout(stdout)?; + } else { + CommandBuilder::docker() + .args(&format!("exec {} {dumping}", db.container)) + .exec_redirect_stdout(stdout)?; + } Ok(()) } + +// what to run inside the container, without the docker exec in front of it +fn dump_command(db: &Database, format: Format) -> String { + match format.flag() { + Some(flag) => format!("pg_dump -U {} --format={flag} {}", db.user, db.name), + None => format!("pg_dumpall -U {}", 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(db: &Database, target: &Path) -> Result<()> { + if target.exists() { + bail!( + "{} already exists; a directory dump will not be written over it", + target.display() + ); + } + + eprintln!("dumping to local directory {}", target.display()); + let remote = remote_dump(); + + CommandBuilder::docker() + .args(&format!( + "exec {} pg_dump -U {} --format=d -f {remote} {}", + db.container, db.user, db.name + )) + .exec()?; + + let copied = CommandBuilder::docker() + .args(&format!( + "cp {}:{remote} {}", + db.container, + target.display() + )) + .exec(); + + let _ = CommandBuilder::docker() + .args(&format!("exec {} rm -rf {remote}", db.container)) + .exec(); + + copied +} + +#[cfg(test)] +mod tests { + use super::{GZIP_MAGIC, Kind, 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}"); + } + } + + #[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)); + } +}