feat: open psql in the database container

This commit is contained in:
2026-09-08 14:16:18 +00:00
parent 61f35f5ae4
commit 6389f3f1f5
6 changed files with 65 additions and 12 deletions

View File

@@ -7,6 +7,13 @@ pub enum Postgres {
/// Import a dump, in any format ahab can produce /// Import a dump, in any format ahab can produce
Import { path: PathBuf }, Import { path: PathBuf },
/// Open psql in the database container
Psql {
/// Arguments for psql itself
#[arg(trailing_var_arg = true, allow_hyphen_values = true)]
rest: Vec<String>,
},
/// Dump via pg_dump, or pg_dumpall for a whole cluster /// Dump via pg_dump, or pg_dumpall for a whole cluster
Dump { Dump {
path: PathBuf, path: PathBuf,

View File

@@ -4,6 +4,7 @@ use super::{Argv, Cmd};
pub struct Exec { pub struct Exec {
container: String, container: String,
interactive: bool, interactive: bool,
tty: bool,
inner: Argv, inner: Argv,
} }
@@ -12,10 +13,17 @@ impl Exec {
Self { Self {
container: container.to_string(), container: container.to_string(),
interactive: false, interactive: false,
tty: false,
inner, inner,
} }
} }
// a terminal, for a command that is a shell rather than a pipe stage
pub fn tty(mut self, tty: bool) -> Self {
self.tty = tty;
self
}
// keep stdin open, for a command that is fed a dump // keep stdin open, for a command that is fed a dump
pub fn interactive(mut self) -> Self { pub fn interactive(mut self) -> Self {
self.interactive = true; self.interactive = true;
@@ -31,6 +39,10 @@ impl Cmd for Exec {
argv = argv.arg("--interactive"); argv = argv.arg("--interactive");
} }
if self.tty {
argv = argv.arg("--tty");
}
argv.arg(&self.container).args(self.inner.words()) argv.arg(&self.container).args(self.inner.words())
} }
} }

View File

@@ -82,11 +82,13 @@ impl Cmd for PgRestore<'_> {
} }
} }
// reads a plain sql dump, of one database or of a whole cluster // reads a plain sql dump, or runs as the interactive shell it is
pub struct Psql<'a> { pub struct Psql<'a> {
username: &'a str, username: &'a str,
dbname: &'a str, dbname: &'a str,
quiet: bool,
atomic: bool, atomic: bool,
rest: &'a [String],
} }
impl<'a> Psql<'a> { impl<'a> Psql<'a> {
@@ -94,33 +96,51 @@ impl<'a> Psql<'a> {
Self { Self {
username, username,
dbname, dbname,
quiet: false,
atomic: false, atomic: false,
rest: &[],
} }
} }
// say nothing and print no result rows, for a restore whose output is noise
pub fn quiet(mut self) -> Self {
self.quiet = true;
self
}
// stop at the first error and undo the rest, which a cluster dump cannot do: // stop at the first error and undo the rest, which a cluster dump cannot do:
// it connects to each database itself, and trips over roles already there // it connects to each database itself, and trips over roles already there
pub fn atomic(mut self) -> Self { pub fn atomic(mut self) -> Self {
self.atomic = true; self.atomic = true;
self self
} }
// whatever the caller typed, for psql's own flags
pub fn args(mut self, rest: &'a [String]) -> Self {
self.rest = rest;
self
}
} }
impl Cmd for Psql<'_> { impl Cmd for Psql<'_> {
fn argv(&self) -> Argv { fn argv(&self) -> Argv {
let argv = Argv::new("psql") let mut argv = Argv::new("psql");
.arg("--quiet")
.flag("--output", "/dev/null") if self.quiet {
argv = argv.arg("--quiet").flag("--output", "/dev/null");
}
argv = argv
.flag("--username", self.username) .flag("--username", self.username)
.flag("--dbname", self.dbname); .flag("--dbname", self.dbname);
if self.atomic { if self.atomic {
return argv argv = argv
.flag("--variable", "ON_ERROR_STOP=1") .flag("--variable", "ON_ERROR_STOP=1")
.arg("--single-transaction"); .arg("--single-transaction");
} }
argv argv.args(self.rest)
} }
} }
@@ -208,13 +228,14 @@ mod tests {
fn only_a_single_database_restore_stops_at_the_first_error() { fn only_a_single_database_restore_stops_at_the_first_error() {
assert!( assert!(
Psql::new("u", "db") Psql::new("u", "db")
.quiet()
.atomic() .atomic()
.argv() .argv()
.quoted() .quoted()
.ends_with("--variable 'ON_ERROR_STOP=1' --single-transaction") .ends_with("--variable 'ON_ERROR_STOP=1' --single-transaction")
); );
assert_eq!( assert_eq!(
Psql::new("u", "postgres").argv().quoted(), Psql::new("u", "postgres").quiet().argv().quoted(),
"psql --quiet --output /dev/null --username u --dbname postgres" "psql --quiet --output /dev/null --username u --dbname postgres"
); );
} }

View File

@@ -2,7 +2,7 @@ mod server;
mod shape; mod shape;
use fs_err::File; use fs_err::File;
use std::io::{self, Write}; use std::io::{self, IsTerminal, Write};
use std::path::Path; use std::path::Path;
use std::process::Stdio; use std::process::Stdio;
@@ -12,8 +12,8 @@ use self::server::{Database, wait_until_ready, when_ready};
use self::shape::{Dump, HEADER_LEN, Kind}; use self::shape::{Dump, HEADER_LEN, Kind};
use crate::cli::Format; use crate::cli::Format;
use crate::cmd::{ use crate::cmd::{
Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Rm, Start, Stop, Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Psql, Rm, Start,
Up, Stop, Up,
}; };
use crate::ctx::Ctx; use crate::ctx::Ctx;
use crate::fsops::{remove_file, rename, suffixed}; use crate::fsops::{remove_file, rename, suffixed};
@@ -202,6 +202,18 @@ pub fn import(ctx: &Ctx, file: &Path) -> Result<()> {
Ok(()) Ok(())
} }
pub fn psql(ctx: &Ctx, rest: &[String]) -> Result<()> {
let db = Database::resolve(ctx)?;
// a terminal only if this one has one, so `psql -c ... | cat` still works
Psql::new(&db.user, &db.name)
.args(rest)
.in_container(&db.container)
.interactive()
.tty(io::stdin().is_terminal())
.replace(ctx)
}
pub fn dump(ctx: &Ctx, file: &Path, format: Format, gzip: bool) -> Result<()> { pub fn dump(ctx: &Ctx, file: &Path, format: Format, gzip: bool) -> Result<()> {
let db = Database::resolve(ctx)?; let db = Database::resolve(ctx)?;

View File

@@ -42,8 +42,8 @@ impl Database {
pub(super) fn restore_with(&self, kind: Kind) -> Box<dyn Cmd + '_> { pub(super) fn restore_with(&self, kind: Kind) -> Box<dyn Cmd + '_> {
match kind { match kind {
Kind::Archive => Box::new(PgRestore::new(&self.user, &self.name)), Kind::Archive => Box::new(PgRestore::new(&self.user, &self.name)),
Kind::Sql => Box::new(Psql::new(&self.user, &self.name).atomic()), Kind::Sql => Box::new(Psql::new(&self.user, &self.name).quiet().atomic()),
Kind::Cluster => Box::new(Psql::new(&self.user, "postgres")), Kind::Cluster => Box::new(Psql::new(&self.user, "postgres").quiet()),
} }
} }
} }

View File

@@ -61,6 +61,7 @@ fn run(ctx: &Ctx, command: cli::Commands) -> Result<ExitCode> {
cli::Commands::Postgres { command } => { cli::Commands::Postgres { command } => {
match command { match command {
cli::Postgres::Import { path } => commands::postgres::import(ctx, &path), cli::Postgres::Import { path } => commands::postgres::import(ctx, &path),
cli::Postgres::Psql { rest } => commands::postgres::psql(ctx, &rest),
cli::Postgres::Dump { path, format, gzip } => { cli::Postgres::Dump { path, format, gzip } => {
commands::postgres::dump(ctx, &path, format, gzip) commands::postgres::dump(ctx, &path, format, gzip)
} }