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 `")) .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 { Ok(PathBuf::from( std::env::var("CARGO_MANIFEST_DIR").expect("This command needs to be invoked using cargo"), ) .join("..") .join("db") .canonicalize()?) }