sea-orm-migration 2.0.0-rc.42

use std::future::Future;

use clap::Parser;
use dotenvy::dotenv;
use std::{error::Error, fmt::Display, process::exit};
use tracing_subscriber::{EnvFilter, prelude::*};

use sea_orm::{ConnectOptions, Database, DbConn, DbErr};
use sea_orm_cli::{MigrateSubcommands, run_migrate_generate, run_migrate_init};

use super::MigratorTraitSelf;

const MIGRATION_DIR: &str = "./";

pub async fn run_cli<M>(migrator: M)
where
    M: MigratorTraitSelf,
{
    run_cli_with_connection(migrator, Database::connect).await;
}

/// Same as [`run_cli`] where you provide the function to create the [`DbConn`].
///
/// This allows configuring the database connection as you see fit.
/// E.g. you can change settings in [`ConnectOptions`] or you can load sqlite
/// extensions.
pub async fn run_cli_with_connection<M, F, Fut>(migrator: M, make_connection: F)
where
    M: MigratorTraitSelf,
    F: FnOnce(ConnectOptions) -> Fut,
    Fut: Future<Output = Result<DbConn, DbErr>>,
{
    dotenv().ok();
    let cli = Cli::parse();

    run_cli_with_connection_inner(migrator, cli, async move |cli| {
        let url = cli
            .database_url
            .clone()
            .expect("Environment variable 'DATABASE_URL' not set");
        let schema = cli
            .database_schema
            .clone()
            .unwrap_or_else(|| "public".to_owned());

        let connect_options = ConnectOptions::new(url)
            .set_schema_search_path(schema)
            .to_owned();

        make_connection(connect_options).await
    })
    .await
    .unwrap_or_else(handle_error);
}

/// Same as [`run_cli`] where you provide a fully customized way to create the
/// [`DbConn`].
///
/// This allows loading database configuration from sources other than CLI
/// arguments or environment variables.
pub async fn run_cli_with_custom_connection(
    migrator: impl MigratorTraitSelf,
    make_connection: impl AsyncFnOnce() -> Result<DbConn, DbErr>,
) {
    dotenv().ok();
    let cli = Cli::parse();

    run_cli_with_connection_inner(migrator, cli, async |_| make_connection().await)
        .await
        .unwrap_or_else(handle_error);
}

pub async fn run_migrate<M>(
    migrator: M,
    db: &DbConn,
    command: Option<MigrateSubcommands>,
    verbose: bool,
) -> Result<(), Box<dyn Error>>
where
    M: MigratorTraitSelf,
{
    setup_tracing(verbose);
    run_migrate_inner(migrator, db, command).await
}

async fn run_cli_with_connection_inner(
    migrator: impl MigratorTraitSelf,
    cli: Cli,
    make_connection: impl AsyncFnOnce(&Cli) -> Result<DbConn, DbErr>,
) -> Result<(), Box<dyn Error>> {
    setup_tracing(cli.verbose);

    if run_non_db_command(cli.command.as_ref())? {
        return Ok(());
    }

    let db = make_connection(&cli)
        .await
        .expect("Fail to acquire database connection");
    let command = cli.command;
    run_migrate_inner(migrator, &db, command).await?;

    Ok(())
}

async fn run_migrate_inner<M>(
    migrator: M,
    db: &DbConn,
    command: Option<MigrateSubcommands>,
) -> Result<(), Box<dyn Error>>
where
    M: MigratorTraitSelf,
{
    if run_non_db_command(command.as_ref())? {
        return Ok(());
    }

    match command {
        Some(MigrateSubcommands::Fresh) => migrator.fresh(db).await?,
        Some(MigrateSubcommands::Refresh) => migrator.refresh(db).await?,
        Some(MigrateSubcommands::Reset) => migrator.reset(db).await?,
        Some(MigrateSubcommands::Status) => migrator.status(db).await?,
        Some(MigrateSubcommands::Up { num }) => migrator.up(db, num).await?,
        Some(MigrateSubcommands::Down { num }) => migrator.down(db, Some(num)).await?,
        _ => migrator.up(db, None).await?,
    };

    Ok(())
}

fn run_non_db_command(command: Option<&MigrateSubcommands>) -> Result<bool, Box<dyn Error>> {
    match command {
        Some(MigrateSubcommands::Init) => {
            run_migrate_init(MIGRATION_DIR)?;
            Ok(true)
        }
        Some(MigrateSubcommands::Generate {
            migration_name,
            universal_time: _,
            local_time,
        }) => {
            run_migrate_generate(MIGRATION_DIR, migration_name, !*local_time)?;
            Ok(true)
        }
        _ => Ok(false),
    }
}

fn setup_tracing(verbose: bool) {
    let filter = match verbose {
        true => "debug",
        false => "sea_orm_migration=info",
    };

    let filter_layer = EnvFilter::try_new(filter).unwrap();

    if verbose {
        let fmt_layer = tracing_subscriber::fmt::layer();
        tracing_subscriber::registry()
            .with(filter_layer)
            .with(fmt_layer)
            .init()
    } else {
        let fmt_layer = tracing_subscriber::fmt::layer()
            .with_target(false)
            .with_level(false)
            .without_time();
        tracing_subscriber::registry()
            .with(filter_layer)
            .with(fmt_layer)
            .init()
    };
}

#[derive(Parser)]
#[command(version)]
pub struct Cli {
    #[arg(short = 'v', long, global = true, help = "Show debug messages")]
    verbose: bool,

    #[arg(
        global = true,
        short = 's',
        long,
        env = "DATABASE_SCHEMA",
        long_help = "Database schema\n \
                    - For MySQL and SQLite, this argument is ignored.\n \
                    - For PostgreSQL, this argument is optional with default value 'public'.\n"
    )]
    database_schema: Option<String>,

    #[arg(
        global = true,
        short = 'u',
        long,
        env = "DATABASE_URL",
        help = "Database URL"
    )]
    database_url: Option<String>,

    #[command(subcommand)]
    command: Option<MigrateSubcommands>,
}

fn handle_error<E>(error: E)
where
    E: Display,
{
    eprintln!("{error}");
    exit(1);
}