brass/cli/src/db.rs

use anyhow::Context;
use chrono::Local;
use sqlx::migrate::Migrate;
use sqlx::{migrate::Migrator, Executor};
use std::fs::File;
use std::io::Write;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    str::FromStr,
};

use brass_config::{load_config, parse_env, Environment};
use clap::{Parser, Subcommand};
use sqlx::{postgres::PgConnectOptions, Connection, PgConnection};

#[derive(Parser)]
#[command(about = "A CLI tool for managing the projects database.", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Command,
    #[arg(short, long, global = true, help = "Choose the environment (development, test, production).", value_parser = parse_env, default_value = "development")]
    environment: Environment,
}

#[derive(Subcommand)]
enum Command {
    #[command(about = "Create the database and run all migrations")]
    Setup,
    #[command(about = "Drop and recreate the database and run all migrations")]
    Reset,
    #[command(about = "Run all pending migrations on database")]
    Migrate,
    #[command(about = "Create a new migration")]
    NewMigration { title: String },
    #[command(about = "Prepare sqlx query metadata for offline compile-time verification")]
    Prepare,
}

#[async_std::main]
#[allow(unused)]
async fn main() {
    let cli = Cli::parse();
    let config = load_config(&cli.environment).expect("Could not load config!");
    let db_config =
        PgConnectOptions::from_str(&config.database_url).expect("Invalid DATABASE_URL!");

    match cli.command {
        Command::Setup => {
            create_db(&db_config)
                .await
                .expect("Failed creating database.");
            migrate_db(&db_config)
                .await
                .expect("Failed migrating database.");
        }
        Command::Reset => {
            drop_db(&db_config)
                .await
                .expect("Failed dropping database.");
            create_db(&db_config)
                .await
                .expect("Failed creating database.");
            migrate_db(&db_config)
                .await
                .expect("Failed migrating database.");
        }
        Command::Migrate => {
            migrate_db(&db_config)
                .await
                .expect("Failed migrating database.");
        }
        Command::NewMigration { title } => {
            create_new_migration(&title)
                .await
                .expect("Failed creating new migration.");
        }
        Command::Prepare => prepare().await.expect("Failed preparing query metadata."),
    }
}

async fn drop_db(db_config: &PgConnectOptions) -> anyhow::Result<()> {
    let db_name = db_config
        .get_database()
        .context("Failed to get database name!")?;

    let root_db_config = db_config.clone().database("postgres");
    let mut connection: PgConnection = Connection::connect_with(&root_db_config)
        .await
        .context("Connection to database failed!")?;

    let query_drop = format!("DROP DATABASE {}", db_name);
    connection
        .execute(query_drop.as_str())
        .await
        .context("Failed to drop database!")?;

    Ok(())
}

async fn create_db(db_config: &PgConnectOptions) -> anyhow::Result<()> {
    let db_name = db_config
        .get_database()
        .context("Failed to get database name!")?;

    let root_db_config = db_config.clone().database("postgres");
    let mut connection: PgConnection = Connection::connect_with(&root_db_config)
        .await
        .context("Connection to database failed!")?;

    let query_create = format!("CREATE DATABASE {}", db_name);
    connection
        .execute(query_create.as_str())
        .await
        .context("Failed to create database!")?;

    Ok(())
}

async fn migrate_db(db_config: &PgConnectOptions) -> anyhow::Result<()> {
    let mut connection: PgConnection = Connection::connect_with(db_config)
        .await
        .context("Connection to database failed!")?;

    let migrations_path = db_package_root()?.join("migrations");

    let migrator = Migrator::new(Path::new(&migrations_path))
        .await
        .context("Failed to create migrator!")?;

    connection
        .ensure_migrations_table()
        .await
        .context("Failed to ensure migrations table!")?;

    let applied_migrations: HashMap<_, _> = connection
        .list_applied_migrations()
        .await
        .context("Failed to list applied migrations!")?
        .into_iter()
        .map(|m| (m.version, m))
        .collect();

    for migration in migrator.iter() {
        if !applied_migrations.contains_key(&migration.version) {
            connection
                .apply(migration)
                .await
                .context("Failed to apply migration {}!")?;
            println!("Applied migration {}.", migration.version);
        }
    }

    Ok(())
}

async fn create_new_migration(title: &str) -> anyhow::Result<()> {
    let now = Local::now();
    let timestamp = now.format("%Y%m%d%H%M%S");
    let file_name = format!("{timestamp}_{title}.sql");
    let path = db_package_root()?.join("migrations").join(&file_name);

    let mut file = File::create(&path).context(format!(r#"Could not create file "{:?}""#, path))?;
    file.write_all("".as_bytes())
        .context(format!(r#"Could not write file "{:?}""#, path))?;

    println!("Created migration {file_name}.");

    Ok(())
}

async fn prepare() -> anyhow::Result<()> {
    let cargo = std::env::var("CARGO")
        .map_err(|_| anyhow::anyhow!("Please invoke me using Cargo, e.g.: `cargo db <ARGS>`"))
        .expect("Existence of CARGO env var is asserted by calling `ensure_sqlx_cli_installed`");

    let mut sqlx_prepare_command = {
        let mut cmd = std::process::Command::new(&cargo);

        cmd.args(["sqlx", "prepare", "--", "--all-targets", "--all-features"]);

        let cmd_cwd = db_package_root().context("Error finding the root of the db package!")?;
        cmd.current_dir(cmd_cwd);

        cmd
    };

    let o = sqlx_prepare_command
        .output()
        .context("Could not run {cargo} sqlx prepare!")?;

    if !o.status.success() {
        let error = anyhow::anyhow!(String::from_utf8_lossy(&o.stdout).to_string()).context("Error generating query metadata. Are you sure the database is running and all migrations are applied?");
        return Err(error);
    }

    println!("Query data written to db/.sqlx directory; please check this into version control.");
    Ok(())
}

fn db_package_root() -> Result<PathBuf, anyhow::Error> {
    Ok(PathBuf::from(
        std::env::var("CARGO_MANIFEST_DIR").expect("This command needs to be invoked using cargo"),
    )
    .join("..")
    .join("db")
    .canonicalize()?)
}