janezicmatej/ahab
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: dump in every format import understands

Browse Source
This commit is contained in:
Matej Janezic
2026-09-07 14:03:21 +00:00
parent 77703485b9
commit ad269308cc
4 changed files with 175 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/cli/postgres.rs
View File

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

4
src/main.rs
View File

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

181
src/scripts/postgres.rs
View File

@@ -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<Vec<u8>> {
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() {