rok-cli 0.1.2

Developer CLI for rok-based Axum applications
//! `rok db:*` — database management commands.

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

// ── commands ──────────────────────────────────────────────────────────────────

/// `rok db:migrate` — apply all pending migrations.
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(())
}

/// `rok db:rollback` — rollback the last N migration batches (default 1).
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(())
}

/// `rok db:status` — print migration status table.
pub async fn status() -> anyhow::Result<()> {
    let pool = connect().await?;
    runner(pool).status().await?;
    Ok(())
}

/// `rok db:seed` — run seeders (guidance, as seeders are compiled into the app binary).
pub async fn seed(class: Option<&str>) -> anyhow::Result<()> {
    // Seeders are user-defined Rust code; the CLI can't call them without
    // linking against the app binary.  Print guidance instead.
    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(())
}

/// `rok db:fresh` — drop all tables, re-run all migrations.
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(())
}