diff --git a/src/cli/postgres.rs b/src/cli/postgres.rs index e0a90b3..2f50505 100644 --- a/src/cli/postgres.rs +++ b/src/cli/postgres.rs @@ -7,6 +7,13 @@ pub enum Postgres { /// Import a dump, in any format ahab can produce 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, + }, + /// Dump via pg_dump, or pg_dumpall for a whole cluster Dump { path: PathBuf, diff --git a/src/cmd/docker.rs b/src/cmd/docker.rs index d1fbc10..b57d9b1 100644 --- a/src/cmd/docker.rs +++ b/src/cmd/docker.rs @@ -4,6 +4,7 @@ use super::{Argv, Cmd}; pub struct Exec { container: String, interactive: bool, + tty: bool, inner: Argv, } @@ -12,10 +13,17 @@ impl Exec { Self { container: container.to_string(), interactive: false, + tty: false, 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 pub fn interactive(mut self) -> Self { self.interactive = true; @@ -31,6 +39,10 @@ impl Cmd for Exec { argv = argv.arg("--interactive"); } + if self.tty { + argv = argv.arg("--tty"); + } + argv.arg(&self.container).args(self.inner.words()) } } diff --git a/src/cmd/postgres.rs b/src/cmd/postgres.rs index 4154496..74b9aea 100644 --- a/src/cmd/postgres.rs +++ b/src/cmd/postgres.rs @@ -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> { username: &'a str, dbname: &'a str, + quiet: bool, atomic: bool, + rest: &'a [String], } impl<'a> Psql<'a> { @@ -94,33 +96,51 @@ impl<'a> Psql<'a> { Self { username, dbname, + quiet: 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: // it connects to each database itself, and trips over roles already there pub fn atomic(mut self) -> Self { self.atomic = true; 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<'_> { fn argv(&self) -> Argv { - let argv = Argv::new("psql") - .arg("--quiet") - .flag("--output", "/dev/null") + let mut argv = Argv::new("psql"); + + if self.quiet { + argv = argv.arg("--quiet").flag("--output", "/dev/null"); + } + + argv = argv .flag("--username", self.username) .flag("--dbname", self.dbname); if self.atomic { - return argv + argv = argv .flag("--variable", "ON_ERROR_STOP=1") .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() { assert!( Psql::new("u", "db") + .quiet() .atomic() .argv() .quoted() .ends_with("--variable 'ON_ERROR_STOP=1' --single-transaction") ); assert_eq!( - Psql::new("u", "postgres").argv().quoted(), + Psql::new("u", "postgres").quiet().argv().quoted(), "psql --quiet --output /dev/null --username u --dbname postgres" ); } diff --git a/src/commands/postgres.rs b/src/commands/postgres.rs index 8a5fe1a..17c11a5 100644 --- a/src/commands/postgres.rs +++ b/src/commands/postgres.rs @@ -2,7 +2,7 @@ mod server; mod shape; use fs_err::File; -use std::io::{self, Write}; +use std::io::{self, IsTerminal, Write}; use std::path::Path; 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 crate::cli::Format; use crate::cmd::{ - Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Rm, Start, Stop, - Up, + Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Psql, Rm, Start, + Stop, Up, }; use crate::ctx::Ctx; use crate::fsops::{remove_file, rename, suffixed}; @@ -202,6 +202,18 @@ pub fn import(ctx: &Ctx, file: &Path) -> Result<()> { 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<()> { let db = Database::resolve(ctx)?; diff --git a/src/commands/postgres/server.rs b/src/commands/postgres/server.rs index 3d1a47d..7e840cd 100644 --- a/src/commands/postgres/server.rs +++ b/src/commands/postgres/server.rs @@ -42,8 +42,8 @@ impl Database { pub(super) fn restore_with(&self, kind: Kind) -> Box { match kind { Kind::Archive => Box::new(PgRestore::new(&self.user, &self.name)), - Kind::Sql => Box::new(Psql::new(&self.user, &self.name).atomic()), - Kind::Cluster => Box::new(Psql::new(&self.user, "postgres")), + Kind::Sql => Box::new(Psql::new(&self.user, &self.name).quiet().atomic()), + Kind::Cluster => Box::new(Psql::new(&self.user, "postgres").quiet()), } } } diff --git a/src/main.rs b/src/main.rs index 34fa775..b374b37 100644 --- a/src/main.rs +++ b/src/main.rs @@ -61,6 +61,7 @@ fn run(ctx: &Ctx, command: cli::Commands) -> Result { cli::Commands::Postgres { command } => { match command { cli::Postgres::Import { path } => commands::postgres::import(ctx, &path), + cli::Postgres::Psql { rest } => commands::postgres::psql(ctx, &rest), cli::Postgres::Dump { path, format, gzip } => { commands::postgres::dump(ctx, &path, format, gzip) }