ahab/src/commands/postgres.rs

mod server;
mod shape;

use std::io::{self, IsTerminal, Write};
use std::path::{Path, PathBuf};
use std::process::Stdio;

use anyhow::{Context, Result, bail};

use self::server::{Database, wait_until_ready, when_ready};
use self::shape::{Dump, HEADER_LEN, Kind};
use crate::cli::Format;
use crate::cli::postgres as cli;
use crate::cmd::{
    Cmd, Cp, CreateDb, DropDb, Gunzip, Gzip, Head, PgDump, PgDumpAll, PgRestore, Psql, Rm,
};
use crate::ctx::Ctx;
use crate::fsops::{create_private_new, suffixed};
use crate::output::{note, plural, warning};

// 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 {
    // the container's /tmp is shared with whatever else runs in it, and a name
    // that can be worked out in advance is one a symlink can be planted at
    let spun = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |since| since.subsec_nanos());

    format!("/tmp/ahab-dump-{}-{spun:09}", std::process::id())
}

// beside the target but unique to this run: a fixed name is a path a second
// dump would share, and one a symlink can be planted at ahead of time
fn partial_path(file: &Path) -> PathBuf {
    suffixed(file, &format!(".{}.partial", std::process::id()))
}

// the database was dropped before the restore was attempted, so a restore that
// failed leaves nothing behind: said the same way wherever it happens, since
// what the user has to do about it does not depend on which tool it was
fn left_empty(tool: &str) -> anyhow::Error {
    anyhow::anyhow!("{tool} failed, the database is left empty")
}

// pg_dumpall recreates roles the cluster already has, so this is the expected
// complaint rather than a failure
fn is_existing_role_error(line: &str) -> bool {
    line.starts_with("ERROR:") && line.contains("role \"") && line.ends_with("already exists")
}

fn restore_cluster(ctx: &Ctx, db: &Database, restore: &dyn Cmd, file: &Path) -> Result<()> {
    let out = restore
        .in_container(&db.container)
        .interactive()
        .stdin_from_captured(ctx, file)
        .context("running psql")?;

    io::stdout().write_all(&out.stdout).ok();

    let mut existing = 0;
    let mut failed = 0;

    for line in String::from_utf8_lossy(&out.stderr).lines() {
        if is_existing_role_error(line) {
            existing += 1;
            continue;
        }

        // this psql runs without ON_ERROR_STOP, so it exits 0 whatever the sql
        // did: these lines are the only account of what happened, which makes
        // them the result rather than progress, and --quiet has to keep them
        if line.starts_with("ERROR:") {
            failed += 1;
            warning!("{line}");
        } else {
            note!(ctx, "{line}");
        }
    }

    if existing > 0 {
        note!(
            ctx,
            "left {existing} existing role{} alone",
            plural(existing)
        );
    }

    if !out.status.success() {
        return Err(left_empty("psql"));
    }
    if failed > 0 {
        bail!(
            "psql reported {failed} error{}, so the cluster restored only in part",
            plural(failed)
        );
    }

    Ok(())
}

// which shape is inside the compression, read through the container's own gunzip
// rather than assuming the host has one. this only reads, so it is a probe and
// runs even in a dry run, where the answer decides what the plan says
fn gzip_kind(ctx: &Ctx, db: &Database, file: &Path) -> Result<Kind> {
    wait_until_ready(ctx, db)?;

    let out = Gunzip
        .pipe(&Head::bytes(HEADER_LEN))?
        // head closes the pipe once it has its bytes, which kills gunzip
        .allow_early_close()
        .in_container(&db.container)
        .interactive()
        .probe_with_stdin(ctx, file)
        .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
            }
        );
    }

    Ok(Kind::of(&out.stdout))
}

pub fn import(ctx: &Ctx, file: &Path) -> Result<()> {
    let dump = Dump::of(file)?;
    let db = Database::resolve(ctx)?;

    if matches!(dump, Dump::Directory) {
        local_path_for_docker(file)?;
    }

    note!(ctx, "stopping all containers");
    ctx.compose().stop().run(ctx)?;

    // everything from here runs with the project down, so an error leaves it
    // that way and the user has no reason to guess as much
    imported(ctx, &db, &dump, file)
        .map_err(|e| e.context("the project is left stopped, `docker compose up` starts it again"))
}

// docker cp reads `container:path`, splitting on the first colon, so a local
// path holding one is read as a container name and something else entirely
fn local_path_for_docker(file: &Path) -> Result<()> {
    match file.to_string_lossy().contains(':') {
        true => bail!(
            "{} has a colon in it, which docker cp reads as a container name; \
             rename it to copy it in or out",
            file.display()
        ),
        false => Ok(()),
    }
}

fn imported(ctx: &Ctx, db: &Database, dump: &Dump, file: &Path) -> Result<()> {
    note!(ctx, "starting db container");
    ctx.compose().start(&db.service).run(ctx)?;

    let remote = remote_dump();

    // a directory cannot be streamed, so it is the one shape that gets copied in
    if matches!(dump, Dump::Directory) {
        when_ready(ctx, db, &Cp::into_container(file, &db.container, &remote))?;
    }

    let kind = match &dump {
        Dump::Directory => Kind::Archive,
        Dump::Header(header) => Kind::of(header),
        Dump::Gzip => match gzip_kind(ctx, db, file) {
            Ok(kind) => kind,
            // looking inside needs a container to decompress with, and a dry run
            // is worth printing with the project down, which is when it is most
            // likely to be asked for
            Err(e) if ctx.dry_run => {
                note!(ctx, "planning for a single database dump: {e:#}");
                Kind::Sql
            }
            Err(e) => return Err(e),
        },
    };

    let restore = db.restore_with(kind);
    // the name of what actually runs, so the message cannot drift from it
    let tool = restore.argv().program().to_string();
    note!(ctx, "restoring database with {tool}");

    when_ready(
        ctx,
        db,
        &DropDb {
            username: &db.user,
            dbname: &db.name,
        }
        .in_container(&db.container),
    )?;

    // a cluster dump creates the database itself, and would trip over one that
    // is already there
    if kind != Kind::Cluster {
        when_ready(
            ctx,
            db,
            &CreateDb {
                username: &db.user,
                dbname: &db.name,
            }
            .in_container(&db.container),
        )?;
    }

    wait_until_ready(ctx, db)?;
    let restored = restore_dump(ctx, db, dump, kind, file, &remote, &tool, restore);

    // the copy inside the container holds the whole database, and its /tmp
    // outlives the command: it goes whether or not the restore worked, which is
    // exactly when it used to be left behind
    if matches!(dump, Dump::Directory) {
        let _ = Rm::recursive(&remote).in_container(&db.container).run(ctx);
    }
    restored?;

    note!(ctx, "restarting containers");
    ctx.compose().stop().run(ctx)?;
    ctx.compose().up().run(ctx)?;

    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn restore_dump(
    ctx: &Ctx,
    db: &Database,
    dump: &Dump,
    kind: Kind,
    file: &Path,
    remote: &str,
    tool: &str,
    restore: Box<dyn Cmd + '_>,
) -> Result<()> {
    if ctx.dry_run {
        note!(ctx, "would restore with {tool}");
        return Ok(());
    }

    // a directory dump was copied in whole, so pg_restore reads it from the
    // container; every other shape is fed in on stdin, through gunzip when it
    // arrives compressed
    if matches!(dump, Dump::Directory) {
        let status = PgRestore::new(&db.user, &db.name)
            .from(remote)
            .in_container(&db.container)
            .status(ctx)?;

        if !status.success() {
            return Err(left_empty(tool));
        }

        return Ok(());
    }

    let restore: Box<dyn Cmd> = match dump {
        Dump::Gzip => Box::new(Gunzip.pipe(&*restore)?),
        _ => restore,
    };

    if kind == Kind::Cluster {
        // psql's output is read rather than streamed here, to keep the
        // expected role errors out of the way
        return restore_cluster(ctx, db, &*restore, file);
    }

    let status = restore
        .in_container(&db.container)
        .interactive()
        .stdin_from(ctx, file)?;

    if !status.success() {
        return Err(left_empty(tool));
    }

    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, args: &cli::Dump) -> Result<()> {
    let (file, format, gzip) = (args.path.as_path(), args.format, args.gzip);
    let db = Database::resolve(ctx)?;

    if format == Format::Directory {
        // main says this as an argument error before getting here; kept as the
        // last word in case anything else ever calls dump
        if gzip {
            bail!("a directory dump is a directory of already compressed files, not a stream");
        }

        return dump_directory(ctx, &db, file);
    }

    note!(ctx, "dumping to local file {}", file.to_string_lossy());

    // written beside the target and renamed once the dump succeeds, so a failure
    // cannot destroy the dump that is already there
    let partial = partial_path(file);
    // a dry run produces no dump, so it must not lay a hand on the target either
    let stdout = if ctx.dry_run {
        Stdio::null()
    } else {
        Stdio::from(create_private_new(&partial)?)
    };

    let dumping = dump_command(&db, format);

    let dumped = if gzip {
        dumping
            .pipe(&Gzip)?
            .in_container(&db.container)
            .stream_to(ctx, stdout)
    } else {
        dumping.in_container(&db.container).stream_to(ctx, stdout)
    };

    if let Err(e) = dumped {
        let _ = ctx.fs().remove_file(&partial);
        return Err(e);
    }

    ctx.fs().rename(&partial, file)?;

    Ok(())
}

// pg_dump for one database, pg_dumpall for a cluster, which has no format letter
fn dump_command(db: &Database, format: Format) -> Box<dyn Cmd + '_> {
    match PgDump::of(&db.user, &db.name, format) {
        Some(dump) => Box::new(dump),
        None => Box::new(PgDumpAll { username: &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(ctx: &Ctx, db: &Database, target: &Path) -> Result<()> {
    // symlink_metadata rather than exists(), which follows the link and so is
    // false for a dangling one: docker cp would then write through it
    if std::fs::symlink_metadata(target).is_ok() {
        bail!(
            "{} already exists; a directory dump will not be written over it",
            target.display()
        );
    }

    note!(ctx, "dumping to local directory {}", target.display());
    let remote = remote_dump();

    PgDump::of(&db.user, &db.name, Format::Directory)
        .expect("a directory dump has a pg_dump letter")
        .to(&remote)
        .in_container(&db.container)
        .run(ctx)?;

    let copied = Cp::out_of_container(&db.container, &remote, target).run(ctx);

    let _ = Rm::recursive(&remote).in_container(&db.container).run(ctx);

    copied
}

#[cfg(test)]
mod tests {
    use super::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}");
        }
    }
}