diff --git a/README.md b/README.md index dc094f9..51d31af 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,10 @@ ahab django make-command ```bash 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 ``` diff --git a/src/cli/postgres.rs b/src/cli/postgres.rs index 5047b20..e0a90b3 100644 --- a/src/cli/postgres.rs +++ b/src/cli/postgres.rs @@ -4,20 +4,24 @@ use clap::{Subcommand, ValueEnum}; #[derive(Subcommand, Debug)] pub enum Postgres { - /// Import a dump, custom format or plain sql + /// Import a dump, in any format ahab can produce Import { path: PathBuf }, - /// Dump via pg_dump + /// 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)] +#[derive(ValueEnum, Clone, Copy, Debug, Default, PartialEq)] pub enum Format { /// pg_dump --format=c, restored with pg_restore #[default] @@ -27,13 +31,29 @@ pub enum Format { /// 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 { - pub fn flag(&self) -> &'static str { + // 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 => "c", - Self::Plain => "p", + Self::Custom => Some("c"), + Self::Plain => Some("p"), + Self::Tar => Some("t"), + Self::Directory => Some("d"), + Self::Cluster => None, } } } diff --git a/src/main.rs b/src/main.rs index 38d7def..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, format } => scripts::postgres::dump(&path, format), + 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 7408cfc..b0fcbcd 100644 --- a/src/scripts/postgres.rs +++ b/src/scripts/postgres.rs @@ -1,7 +1,7 @@ use anyhow::{Context, Result, anyhow, bail}; use std::{ fs::File, - io::Read, + io::{self, Read, Write}, path::{Path, PathBuf}, process::Stdio, thread, @@ -61,12 +61,12 @@ impl Database { 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", + "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 -U {} -d postgres", self.user), + Kind::Cluster => format!("psql -q -o /dev/null -U {} -d postgres", self.user), } } } @@ -122,6 +122,41 @@ impl Kind { } } +// 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]; @@ -228,15 +263,7 @@ pub fn import(file: &Path) -> Result<()> { }; 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}"); - } - + eprintln!("restoring database with {tool}"); run_when_ready( &db, &format!("exec {} dropdb -U {} {}", db.container, db.user, db.name), @@ -255,20 +282,27 @@ pub fn import(file: &Path) -> Result<()> { } wait_until_ready(&db)?; - 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()?, - }; + 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 !status.success() { + bail!("{tool} failed, the database is left empty"); + } + } } if matches!(dump, Dump::Directory) { @@ -284,31 +318,100 @@ pub fn import(file: &Path) -> Result<()> { Ok(()) } -pub fn dump(file: &PathBuf, format: Format) -> 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={} {}", - db.container, - db.user, - format.flag(), - 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}; + 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() {