brass/cli/src/db.rs

use anyhow::Context;
use sqlx::migrate::Migrate;
use sqlx::{migrate::Migrator, Executor};
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,
}

#[async_std::main]
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.");
        }
    }
}

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 = PathBuf::from(
        std::env::var("CARGO_MANIFEST_DIR").expect("This command needs to be invoked using cargo"),
    )
    .join("..")
    .join("migrations")
    .canonicalize()
    .unwrap();

    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(())
}