omni_orchestrator/initialization/
setup_database.rs

1use anyhow::Result;
2use colored::Colorize;
3use std::{env, sync::Arc};
4use crate::db_manager::DatabaseManager;
5use crate::logging;
6
7/// Initializes the deployment database and registers all platforms.
8///
9/// - Loads the database URL from environment variables or falls back to defaults.
10/// - Connects to the database and wraps the manager in an `Arc` for shared access.
11/// - Discovers all platforms and pre-initializes their connection pools for fast access.
12///
13/// # Errors
14/// Returns an error if the database connection or platform pool initialization fails.
15pub async fn setup_database() -> Result<Arc<DatabaseManager>> {
16    // Load database URL from environment or .env file
17    let database_url = env::var("DATABASE_URL").unwrap_or_else(|_| {
18        dotenv::dotenv().ok();
19        env::var("DEFAULT_DATABASE_URL")
20            .unwrap_or_else(|_| "mysql://root:root@localhost:4001".to_string())
21    });
22
23    log::info!("{}", format!("Database URL: {}", database_url).blue());
24    let db_manager = Arc::new(DatabaseManager::new(&database_url).await?);
25    
26    // Print a banner for platform registration
27    logging::print_banner("Platform Database Registration", |s| s.bright_blue());
28
29    // Retrieve all platforms from the database
30    let platforms = db_manager.get_all_platforms().await?;
31    log::info!("{}", format!("Found {} platforms", platforms.len()).blue());
32
33    // Pre-initialize connection pools for each platform
34    for platform in &platforms {
35        log::info!(
36            "{}",
37            format!(
38                "Pre-initializing connection for platform: {}",
39                platform.name
40            )
41            .blue()
42        );
43        db_manager.get_platform_pool(&platform.name, platform.id.unwrap_or(0)).await?;
44    }
45
46    Ok(db_manager)
47}