rullst-orm 13.0.0-alpha.1

use super::migration_lock::MigrationLock;
use super::validation::validate_table_name;
use crate::Error;

#[async_trait::async_trait]
pub trait Migration: Send + Sync {
    fn name(&self) -> &'static str;
    async fn up(&self) -> Result<(), Error>;
    async fn down(&self) -> Result<(), Error>;

    /// Whether the runner applies `up()` (or `down()`) and records it in the
    /// `migrations` table inside one [`crate::Orm::transaction`], so a failure
    /// between the two cannot leave the change applied but unrecorded.
    ///
    /// Only statements that join the task-scoped transaction are covered:
    /// `Schema` calls, `execute_query!`/`dispatch_executor!` and generated
    /// model methods. A query sent to `Orm::pool()` directly runs on another
    /// connection outside the transaction and cannot see its uncommitted
    /// tables, so a migration that uses one must keep the default `false`.
    /// PostgreSQL and SQLite roll DDL back with the transaction; MySQL and
    /// MariaDB commit implicitly at each DDL statement, so there only data
    /// changes after the last DDL statement are atomic with the record.
    /// Statements that cannot run in a transaction block, such as PostgreSQL
    /// `CREATE INDEX CONCURRENTLY`, need `false`. Unpublished v13 API.
    fn within_transaction(&self) -> bool {
        false
    }
}

#[cfg_attr(test, mutants::skip)]
pub async fn run_artisan_with_args(
    args: &[String],
    migrations: Vec<Box<dyn Migration>>,
    seeders: Vec<Box<dyn crate::Seeder>>,
) -> Result<(), Error> {
    if args.len() < 2 {
        println!("Rullst ORM Artisan CLI");
        println!("Usage:");
        println!("  make:migration <name>   Generate a new migration");
        println!("  migrate                  Run all pending migrations");
        println!("  migrate:rollback         Rollback the last batch of migrations");
        println!("  status                   Show migrations status");
        println!("  db:seed                  Populate the database with seeders");
        println!(
            "  sail:install [--force]   Generate a local docker-compose.yml (Laravel Sail style); --force replaces an existing file"
        );
        return Ok(());
    }

    let command = &args[1];
    match command.as_str() {
        "make:migration" => {
            if args.len() < 3 {
                println!("Error: migration name is required.");
                return Ok(());
            }
            let name = &args[2];
            create_migration_files(name)?;
        }
        "migrate" | "db:migrate" => {
            run_migrations(migrations).await?;
        }
        "migrate:rollback" | "db:rollback" => {
            rollback_migrations(migrations).await?;
        }
        "status" | "db:status" => {
            status_migrations(migrations).await?;
        }
        "db:seed" => {
            println!("Seeding database...");
            crate::Orm::seed(seeders).await?;
            println!("Database seeded successfully!");
        }
        "sail:install" => {
            let force = args.iter().skip(2).any(|argument| argument == "--force");
            let path = super::sail::install(std::path::Path::new("."), force)?;
            println!(
                "{} created successfully! Run `docker compose up -d` to start.",
                path.display()
            );
        }
        _ => {
            println!("Unknown command: {}", command);
        }
    }
    Ok(())
}

#[cfg_attr(test, mutants::skip)]
pub async fn run_artisan(
    migrations: Vec<Box<dyn Migration>>,
    seeders: Vec<Box<dyn crate::Seeder>>,
) -> Result<(), Error> {
    let args: Vec<String> = std::env::args().collect();
    run_artisan_with_args(&args, migrations, seeders).await
}

#[cfg_attr(mutants, mutants::skip)]
async fn migrations_table_exists(pool: &crate::RullstPool, driver: &str) -> Result<bool, Error> {
    // Only the schema/database that unqualified `migrations` statements use
    // counts; a same-named table of another application must not.
    match driver {
        "postgres" | "mysql" => {
            let query_str = if driver == "postgres" {
                "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = 'migrations'"
            } else {
                "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = 'migrations'"
            };
            let row: (i64,) = sqlx::query_as(query_str).fetch_one(pool).await?;
            Ok(row.0 > 0)
        }
        _ => {
            let query_str =
                "SELECT COUNT(*) FROM sqlite_schema WHERE type='table' AND name='migrations'";
            let row: (i64,) = sqlx::query_as(query_str).fetch_one(pool).await?;
            Ok(row.0 > 0)
        }
    }
}

#[cfg_attr(test, mutants::skip)]
async fn status_migrations(migrations: Vec<Box<dyn Migration>>) -> Result<(), Error> {
    let pool = crate::Orm::try_pool()?;
    let driver = crate::Orm::try_driver()?;

    let table_exists = migrations_table_exists(pool, driver).await?;

    let executed_set = if table_exists {
        let executed: Vec<(String,)> = sqlx::query_as("SELECT migration FROM migrations")
            .fetch_all(pool)
            .await?;
        executed
            .into_iter()
            .map(|(m,)| m)
            .collect::<std::collections::HashSet<String>>()
    } else {
        std::collections::HashSet::new()
    };

    let name_header = "Migration Name";
    let status_header = "Status";
    println!("{name_header:<40} | {status_header}");
    println!("{}", "-".repeat(55));
    for m in migrations {
        let name = m.name();
        let status = if executed_set.contains(name) {
            "Applied"
        } else {
            "Pending"
        };
        println!("{:<40} | {}", name, status);
    }

    Ok(())
}

#[cfg_attr(test, mutants::skip)]
fn create_migration_files(name: &str) -> Result<(), Error> {
    validate_table_name(name)?;
    use std::fs;

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
        .to_string();
    let sanitized_name = name.replace(['/', '\\'], "");
    let snake_name = sanitized_name.to_lowercase().replace("-", "_");
    let file_name = format!("m{}_{}", now, snake_name);

    fs::create_dir_all("src/migrations")
        .map_err(|e| Error::Internal(format!("Failed to create migrations directory: {}", e)))?;

    let new_file_path = format!("src/migrations/{}.rs", file_name);
    let template = include_str!("../migration_template.rs.txt");
    let migration_code = template
        .replace("{timestamp}", &now)
        .replace("{name}", &snake_name);

    fs::write(&new_file_path, migration_code)
        .map_err(|e| Error::Internal(format!("Failed to write migration file: {}", e)))?;
    println!("Created migration file: {}", new_file_path);

    regenerate_migrations_mod()?;

    Ok(())
}

#[cfg_attr(test, mutants::skip)]
fn regenerate_migrations_mod() -> Result<(), Error> {
    use std::fs;
    let paths = fs::read_dir("src/migrations")
        .map_err(|e| Error::Internal(format!("Failed to read migrations dir: {}", e)))?;

    let mut modules = vec![];
    for path in paths {
        let path = path.map_err(|e| Error::Internal(e.to_string()))?.path();
        if let Some(ext) = path.extension()
            && ext == "rs"
            && let Some(stem) = path.file_stem()
        {
            let stem_str = stem.to_string_lossy().to_string();
            if stem_str != "mod" && stem_str.starts_with('m') {
                modules.push(stem_str);
            }
        }
    }
    modules.sort();

    use std::fmt::Write;
    let mut mod_content = String::new();
    mod_content.push_str("// Generated by Rullst ORM Artisan. Do not edit manually.\n\n");
    for m in &modules {
        let _ = writeln!(mod_content, "pub mod {};", m);
    }
    mod_content
        .push_str("\npub fn get_migrations() -> Vec<Box<dyn rullst_orm::schema::Migration>> {\n");
    mod_content.push_str("    vec![\n");
    for m in &modules {
        let _ = writeln!(mod_content, "        Box::new({}::MigrationImpl),", m);
    }
    mod_content.push_str("    ]\n");
    mod_content.push_str("}\n");

    fs::write("src/migrations/mod.rs", mod_content)
        .map_err(|e| Error::Internal(format!("Failed to write mod.rs: {}", e)))?;
    println!("Regenerated src/migrations/mod.rs");

    Ok(())
}

/// Applies pending migrations while holding the database's runner lock, so
/// concurrent runners (for example replicas starting together) apply each
/// migration once: the executed set is read only after the lock is held.
#[cfg_attr(test, mutants::skip)]
async fn run_migrations(migrations: Vec<Box<dyn Migration>>) -> Result<(), Error> {
    let pool = crate::Orm::try_pool()?;
    let driver = crate::Orm::try_driver()?;
    let lock = MigrationLock::acquire(pool, driver).await?;
    let result = run_pending_migrations(pool, driver, migrations).await;
    lock.release().await;
    result
}

#[cfg_attr(test, mutants::skip)]
async fn run_pending_migrations(
    pool: &crate::RullstPool,
    driver: &str,
    migrations: Vec<Box<dyn Migration>>,
) -> Result<(), Error> {
    let query_str = match driver {
        "postgres" => {
            "CREATE TABLE IF NOT EXISTS migrations (
                id SERIAL PRIMARY KEY,
                migration VARCHAR(255) NOT NULL UNIQUE,
                batch INTEGER NOT NULL
            )"
        }
        "mysql" => {
            "CREATE TABLE IF NOT EXISTS migrations (
                id INT AUTO_INCREMENT PRIMARY KEY,
                migration VARCHAR(255) NOT NULL UNIQUE,
                batch INT NOT NULL
            )"
        }
        _ => {
            "CREATE TABLE IF NOT EXISTS migrations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                migration TEXT NOT NULL UNIQUE,
                batch INTEGER NOT NULL
            )"
        }
    };

    sqlx::query(query_str).execute(pool).await?;

    let executed: Vec<(String,)> = sqlx::query_as("SELECT migration FROM migrations")
        .fetch_all(pool)
        .await?;
    let executed_set: std::collections::HashSet<String> =
        executed.into_iter().map(|(m,)| m).collect();

    let batch_row: (Option<i32>,) = sqlx::query_as("SELECT MAX(batch) FROM migrations")
        .fetch_one(pool)
        .await?;
    let next_batch = batch_row.0.unwrap_or(0) + 1;

    let tracking_sql = if driver == "postgres" {
        "INSERT INTO migrations (migration, batch) VALUES ($1, $2)"
    } else {
        "INSERT INTO migrations (migration, batch) VALUES (?, ?)"
    };
    let mut count = 0;
    for m in migrations {
        let name = m.name();
        if !executed_set.contains(name) {
            println!("Migrating: {}", name);
            if m.within_transaction() {
                // The change and its record commit (or roll back) together.
                crate::Orm::transaction(move |_| {
                    Box::pin(async move {
                        m.up().await?;
                        let record = sqlx::query(tracking_sql).bind(name).bind(next_batch);
                        crate::execute_query!(record, execute, pool)?;
                        Ok::<(), Error>(())
                    })
                })
                .await?;
            } else {
                m.up().await?;
                sqlx::query(tracking_sql)
                    .bind(name)
                    .bind(next_batch)
                    .execute(pool)
                    .await?;
            }
            println!("Migrated:  {}", name);
            count += 1;
        }
    }

    if count == 0 {
        println!("Nothing to migrate.");
    }

    Ok(())
}

/// Rolls back the last batch while holding the database's runner lock.
#[cfg_attr(test, mutants::skip)]
async fn rollback_migrations(migrations: Vec<Box<dyn Migration>>) -> Result<(), Error> {
    let pool = crate::Orm::try_pool()?;
    let driver = crate::Orm::try_driver()?;
    let lock = MigrationLock::acquire(pool, driver).await?;
    let result = rollback_last_batch(pool, driver, migrations).await;
    lock.release().await;
    result
}

#[cfg_attr(test, mutants::skip)]
async fn rollback_last_batch(
    pool: &crate::RullstPool,
    driver: &str,
    migrations: Vec<Box<dyn Migration>>,
) -> Result<(), Error> {
    let table_exists = migrations_table_exists(pool, driver).await?;

    if !table_exists {
        println!("Nothing to rollback.");
        return Ok(());
    }

    let batch_row: (Option<i32>,) = sqlx::query_as("SELECT MAX(batch) FROM migrations")
        .fetch_one(pool)
        .await?;

    let last_batch = match batch_row.0 {
        Some(b) if b > 0 => b,
        _ => {
            println!("Nothing to rollback.");
            return Ok(());
        }
    };

    let rollback_query = if driver == "postgres" {
        "SELECT migration FROM migrations WHERE batch = $1 ORDER BY id DESC"
    } else {
        "SELECT migration FROM migrations WHERE batch = ? ORDER BY id DESC"
    };
    let to_rollback: Vec<(String,)> = sqlx::query_as(rollback_query)
        .bind(last_batch)
        .fetch_all(pool)
        .await?;

    let mut rollback_map: std::collections::HashMap<String, Box<dyn Migration>> =
        std::collections::HashMap::with_capacity(migrations.len());
    for m in migrations {
        rollback_map.insert(m.name().to_string(), m);
    }

    let forget_sql = if driver == "postgres" {
        "DELETE FROM migrations WHERE migration = $1"
    } else {
        "DELETE FROM migrations WHERE migration = ?"
    };
    for (name,) in to_rollback {
        if let Some(m) = rollback_map.remove(&name) {
            println!("Rolling back: {}", name);
            // Forget each migration as soon as its down() succeeds, so a later
            // failing down() leaves only the migrations that were not reverted
            // recorded as applied; the next rollback resumes from there.
            if m.within_transaction() {
                let recorded = name.clone();
                crate::Orm::transaction(move |_| {
                    Box::pin(async move {
                        m.down().await?;
                        let forget = sqlx::query(forget_sql).bind(recorded);
                        crate::execute_query!(forget, execute, pool)?;
                        Ok::<(), Error>(())
                    })
                })
                .await?;
            } else {
                m.down().await?;
                sqlx::query(forget_sql).bind(&name).execute(pool).await?;
            }
            println!("Rolled back:  {}", name);
        } else {
            println!(
                "Warning: migration {} found in database but not in compiled binary.",
                name
            );
        }
    }

    Ok(())
}