use clap::{Arg, Command}; use container::{ db::{DBInfo, sqlite::SQLiteContainer}, types::ConfigInfoType, }; 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 + Send>> { let output_path = _matches.get_one::("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, }; 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(()) } } }