merge: postgres dump formats

README.md

@@ -84,9 +84,26 @@ ahab django make-command <app> <name>
 
 ```bash
 ahab postgres dump <path>            # pg_dump, custom format
-ahab postgres import <path>         # drop, create, pg_restore
+ahab postgres dump -F plain <path>   # pg_dump, plain sql
+ahab postgres dump -F tar <path>     # pg_dump, tar
+ahab postgres dump -F directory <d>  # pg_dump, a directory of files
+ahab postgres dump -F cluster <path> # pg_dumpall, roles and all databases
+ahab postgres dump -F plain -z <path> # any of them, gzipped
+ahab postgres import <path>          # 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

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};

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,
+        }
+    }
 }

src/command_builder.rs

@@ -46,6 +46,11 @@ impl CommandBuilder {
         Self::default().args("docker compose")
     }
 
+    pub fn arg(mut self, arg: impl Into<String>) -> Self {
+        self.args.push(arg.into());
+        self
+    }
+
     pub fn args<T>(mut self, args: T) -> Self
     where
         Args: From<T>,

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),

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,33 +56,133 @@ 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),
+        }
+    }
 }
 
-pub fn import(file: &Path) -> Result<()> {
-    let db = Database::resolve()?;
-    let dump_file = file.to_string_lossy();
+enum Dump {
+    Directory,
+    Gzip,
+    Header(Vec<u8>),
+}
 
-    eprintln!("stopping all containers");
-    docker_compose::stop()?;
+impl Dump {
+    fn of(path: &Path) -> Result<Self> {
+        if path.is_dir() {
+            return Ok(Self::Directory);
+        }
 
-    eprintln!("starting db container");
-    docker_compose::start(Some(&db.service))?;
+        let header = read_header(path)?;
+        if header.starts_with(GZIP_MAGIC) {
+            return Ok(Self::Gzip);
+        }
 
-    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
-        ),
-    ];
+        Ok(Self::Header(header))
+    }
+}
 
-    for command in commands {
+// 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<Vec<u8>> {
+    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<std::process::Command> {
+    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!(
@@ -82,7 +197,118 @@ pub fn import(file: &Path) -> Result<()> {
     {
         thread::sleep(Duration::from_secs(1));
     }
-        CommandBuilder::docker().args(&command).exec()?;
+
+    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()?;
+
+    eprintln!("stopping all containers");
+    docker_compose::stop()?;
+
+    eprintln!("starting db container");
+    docker_compose::start(Some(&db.service))?;
+
+    let remote = remote_dump();
+
+    // 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();
+    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
-    );
+    if gzip {
         CommandBuilder::docker()
-        .args(&command)
+            .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));
+    }
+}