feat: support custom format, plain sql and cluster dumps

README.md

@@ -83,10 +83,23 @@ ahab django make-command <app> <name>
 ## postgres
 
 ```bash
-ahab postgres dump <path>           # pg_dump, custom format
+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 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,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 via pg_dump
-    Dump { path: PathBuf },
+    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",
+        }
+    }
 }

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

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<u8>),
+}
+
+impl Dump {
+    fn of(path: &Path) -> Result<Self> {
+        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<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!(
+            "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 remote = remote_dump();
-    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
-        ),
-    ];
 
-    for command in commands {
+    // a directory cannot be streamed, so it is the one shape that gets copied in
-        debug_eprintln!("waiting until pg_isready");
+    if matches!(dump, Dump::Directory) {
-        while !CommandBuilder::docker()
+        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()?
+            .wait()?,
-            .success()
+    };
-        {
+
-            thread::sleep(Duration::from_secs(1));
+    if !status.success() {
-        }
+        bail!("{tool} failed, the database is left empty");
-        CommandBuilder::docker().args(&command).exec()?;
+    }
+
+    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 {}",
+        "exec {} pg_dump -U {} --format={} {}",
-        db.container, db.user, db.name
+        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));
+    }
+}