use anyhow::Context;
use rok_orm_migrate::{FileSource, MigrationRunner};
use sqlx::postgres::PgPoolOptions;
async fn connect() -> anyhow::Result<sqlx::PgPool> {
let _ = dotenvy::dotenv();
let url = std::env::var("DATABASE_URL")
.context("DATABASE_URL not set — add it to .env or the environment")?;
PgPoolOptions::new()
.max_connections(3)
.connect(&url)
.await
.with_context(|| format!("failed to connect to database: {url}"))
}
fn migration_dir() -> &'static str {
if std::path::Path::new("database/migrations").exists() {
"database/migrations"
} else {
"migrations"
}
}
fn runner(pool: sqlx::PgPool) -> MigrationRunner {
MigrationRunner::new(pool).source(FileSource::new(migration_dir()))
}
pub async fn migrate(dry_run: bool) -> anyhow::Result<()> {
let pool = connect().await?;
if dry_run {
println!("-- dry run: pending migrations --");
runner(pool).status().await?;
} else {
runner(pool).run().await?;
println!("All migrations applied.");
}
Ok(())
}
pub async fn rollback(step: u32) -> anyhow::Result<()> {
let pool = connect().await?;
let step = step.max(1);
for i in 0..step {
runner(pool.clone()).rollback().await?;
println!(" Rolled back step {}/{step}.", i + 1);
}
println!("Rollback complete ({step} step(s)).");
Ok(())
}
pub async fn status() -> anyhow::Result<()> {
let pool = connect().await?;
runner(pool).status().await?;
Ok(())
}
pub async fn seed(class: Option<&str>) -> anyhow::Result<()> {
if let Some(class) = class {
println!(
"rok db:seed --class {class}: add this to your main.rs or a dedicated seeder binary:"
);
println!(" {}::run(&pool).await?;", class);
} else {
println!("rok db:seed: seeders must be registered and run from your app binary.");
println!();
println!("Add this to your main.rs or a seeder binary:");
println!(" UsersSeeder::run(&pool).await?;");
}
Ok(())
}
pub async fn fresh() -> anyhow::Result<()> {
let pool = connect().await?;
println!("Dropping all user tables...");
sqlx::query(
"DO $$
DECLARE r RECORD;
BEGIN
FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP
EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
END LOOP;
END $$;",
)
.execute(&pool)
.await
.context("failed to drop tables")?;
println!("Running migrations...");
runner(pool).run().await?;
println!("Database refreshed.");
Ok(())
}