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