YANPM/apps/cli/src/cmd/db_migrate_and_generate.rs

use clap::{Arg, Command};
use container::containers::{
    ConfigInfoType,
    db::{DBInfo, sqlite::SQLiteContainer},
};
use migration::{generate_entity, migrate_database};
use shared::db_type::DBType;

use crate::cmd::CliCommand;

const MAX_DB_READY_ATTEMPTS: u8 = 10;
const DB_READY_CHECK_INTERVAL_SECS: u64 = 2;
const DB_READY_STRING: [&str; 1] = ["ready to accept connections"];

pub fn get_cli_command() -> CliCommand {
    CliCommand {
        command: command(),
        action,
    }
}

fn command() -> Command {
    Command::new("db:migrate_and_generate")
        .arg(
            Arg::new("output_path")
                .short('o')
                .long("output-path")
                .value_name("PATH")
                .help("Path to output the generated entity schema")
                .required(true),
        )
        .about("Migrate database and generate entity schema")
}

fn action(
    _matches: &clap::ArgMatches,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> {
    let output_path = _matches.get_one::<String>("output_path");
    let output_path = output_path.unwrap().to_string();
    Box::pin(async move {
        let mut error_occurred = false;
        let database_configs = vec![
            SQLiteContainer::new(None)
                .await
                .get_db_container_config_info()
                .await,
            // only sqlite is required to generate entities when using Seaorm
            // PostgreSQLContainer::new(None)
            //     .await
            //     .get_db_container_config_info()
            //     .await,
        ];

        for db_config in database_configs {
            let config = container::Config {
                database: db_config,
                agent: None,
            };
            let mut detached_handler = container::start_detached(&config).await;
            match migrate_and_generate_entity(&config, &output_path).await {
                Ok(_) => println!("Migration and entity generation succeeded."),
                Err(_) => {
                    eprintln!("Migration and entity generation failed.");
                    error_occurred = true;
                    break;
                }
            }
            detached_handler.stop().await;
        }

        if error_occurred {
            std::process::exit(1);
        } else {
            std::process::exit(0);
        }
    })
}

async fn migrate_and_generate_entity(
    config: &container::Config,
    output_path: &str,
) -> Result<(), ()> {
    let ready_result = await_database_ready(config).await;
    if ready_result.is_err() {
        eprintln!("Database did not become ready in time.");
        return Err(());
    }

    let db_url = match &config.database {
        ConfigInfoType::Containerized(container_info) => &container_info.url,
        ConfigInfoType::PreExisting(pre_existing_info) => &pre_existing_info.url,
    };

    let db_type = get_database_type(config);
    match migrate_database(db_url).await {
        Ok(_) => {
            println!("Database migrated successfully for {:?}", db_type);
        }
        Err(e) => {
            eprintln!("Failed to migrate database for {}: {:#?}", db_type, e);
            return Err(());
        }
    }

    match generate_entity(db_url, output_path).await {
        Ok(_) => {
            println!(
                "Database entity schema generated successfully for {:?}",
                db_type
            );
        }
        Err(e) => {
            eprintln!(
                "Failed to generate database entity schema for {}: {:#?}",
                db_type, e
            );
            return Err(());
        }
    }

    Ok(())
}

fn get_database_type(config: &container::Config) -> DBType {
    match config.database {
        ConfigInfoType::Containerized(ref container_info) => container_info.db_type.clone(),
        ConfigInfoType::PreExisting(ref pre_existing_info) => pre_existing_info.db_type.clone(),
    }
}

async fn await_database_ready(config: &container::Config) -> Result<(), ()> {
    match config.database {
        ConfigInfoType::Containerized(ref container_info) => {
            let container_type = &container_info.db_type;
            for attempt in 1..=MAX_DB_READY_ATTEMPTS {
                println!(
                    "Checking if database container {} is ready (attempt {}/{})...",
                    container_info.db_type, attempt, MAX_DB_READY_ATTEMPTS
                );
                let logs = match container_info.container.stdout_to_vec().await {
                    Ok(logs) => logs,
                    Err(e) => {
                        eprintln!("Failed to get container logs: {}", e);
                        return Err(());
                    }
                };
                let log_output = String::from_utf8_lossy(&logs);
                if DB_READY_STRING.iter().any(|&s| log_output.contains(s)) {
                    println!("Database container {} is ready.", container_info.db_type);
                    return Ok(());
                }
                tokio::time::sleep(std::time::Duration::from_secs(DB_READY_CHECK_INTERVAL_SECS))
                    .await;
            }
            eprintln!(
                "Database container {} did not become ready after {} attempts.",
                container_type, MAX_DB_READY_ATTEMPTS
            );
            Err(())
        }

        ConfigInfoType::PreExisting(ref pre_existing_info) => {
            println!(
                "Pre-existing database of type {} assumed to be ready.",
                pre_existing_info.db_type
            );
            Ok(())
        }
    }
}