src/infrastructure/db/mod.rs

pub mod healing;
pub mod pool_router;
pub mod utils;

pub use pool_router::DbRouter;

use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
use std::env;

pub type DbPool = PgPool;

pub async fn init_db() -> DbPool {
    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set in environment");
    let cleaned_url = utils::clean_database_url(&database_url);
    let schema = utils::database_schema();

    if !cleaned_url.starts_with("postgres://") && !cleaned_url.starts_with("postgresql://") {
        tracing::error!("FATAL: Invalid DATABASE_URL detected. URL must start with 'postgres://' or 'postgresql://'. Final cleaned URL: '{}'", cleaned_url);
        panic!("INVALID_DATABASE_URL_SCHEME: Please ensure your environment variable doesn't contain psql wrappers or quotes.");
    }

    let mut base_options = cleaned_url
        .parse::<PgConnectOptions>()
        .expect("DATABASE_URL could not be parsed as a Postgres connection URL");
    base_options = base_options.statement_cache_capacity(0);

    let mut retry_count = 0;
    let max_retries = 5;
    let mut bootstrap_pool: Option<PgPool> = None;

    while retry_count < max_retries {
        tracing::info!(
            "Database connection pool initializing (Attempt {}/{})...",
            retry_count + 1,
            max_retries
        );

        match PgPoolOptions::new()
            .max_connections(1)
            .acquire_timeout(std::time::Duration::from_secs(30))
            .connect_with(base_options.clone())
            .await
        {
            Ok(pool) => {
                bootstrap_pool = Some(pool);
                break;
            }
            Err(e) => {
                retry_count += 1;
                if retry_count >= max_retries {
                    tracing::error!(
                        "FATAL: Failed to create bootstrap connection pool after {} attempts: {}",
                        max_retries,
                        e
                    );
                    panic!("DATABASE_CONNECTION_FAILURE: {}", e);
                }
                let delay = std::time::Duration::from_secs(2u64.pow(retry_count as u32));
                tracing::warn!(
                    "Bootstrap connection failed: {}. Retrying in {:?}...",
                    e,
                    delay
                );
                tokio::time::sleep(delay).await;
            }
        }
    }

    let bootstrap_pool = bootstrap_pool.unwrap();

    match sqlx::query(&format!("CREATE SCHEMA IF NOT EXISTS {schema}"))
        .execute(&bootstrap_pool)
        .await
    {
        Ok(_) => tracing::info!("Schema '{}' ensured.", schema),
        Err(e) => {
            tracing::warn!(
                "Failed to create schema '{}' (likely permission issue): {}. Proceeding anyway...",
                schema,
                e
            );
        }
    }
    bootstrap_pool.close().await;

    let mut retry_count = 0;
    let pool = loop {
        let schema_clone_for_closure = schema.clone();
        match PgPoolOptions::new()
            .max_connections(50)
            .min_connections(10)
            .acquire_timeout(std::time::Duration::from_secs(30))
            .idle_timeout(std::time::Duration::from_secs(600))
            .max_lifetime(std::time::Duration::from_secs(1800))
            .after_connect(move |conn, _meta| {
                let schema_clone = schema_clone_for_closure.clone();
                Box::pin(async move {
                    sqlx::query(&format!("SET search_path TO {},public", schema_clone))
                        .execute(conn)
                        .await?;
                    Ok(())
                })
            })
            .connect_with(base_options.clone())
            .await
        {
            Ok(p) => break p,
            Err(e) => {
                retry_count += 1;
                if retry_count >= max_retries {
                    tracing::error!("FATAL: Connection pool creation failed after {} attempts: {}. Ensure your database is accessible.", max_retries, e);
                    panic!("DATABASE_POOL_CREATION_FAILURE: {}", e);
                }
                let delay = std::time::Duration::from_secs(2u64.pow(retry_count as u32));
                tracing::warn!("Pool connection failed: {}. Retrying in {:?}...", e, delay);
                tokio::time::sleep(delay).await;
            }
        }
    };

    // Run database migrations using standard SQLX migration system (with SRE self-healing integration enabled)
    let migrate_result = sqlx::migrate!("./migrations").run(&pool).await;
    if let Err(e) = migrate_result {
        let err_msg = e.to_string();
        if err_msg.contains("was previously applied but is missing") {
            tracing::warn!("Detected dirty/missing future migrations in _sqlx_migrations table. Self-healing database migrations table...");
            // Clean up any migration version >= 33 which is not present in our current codebase
            match sqlx::query("DELETE FROM _sqlx_migrations WHERE version >= 33")
                .execute(&pool)
                .await
            {
                Ok(rows) => {
                    tracing::info!("Cleaned up {} dirty migration rows from _sqlx_migrations. Retrying migrations...", rows.rows_affected());
                    if let Err(retry_err) = sqlx::migrate!("./migrations").run(&pool).await {
                        tracing::error!("Database migration retry failed: {}", retry_err);
                        panic!("DATABASE_MIGRATION_FAILURE: {}", retry_err);
                    }
                }
                Err(heal_err) => {
                    tracing::error!("Failed to clean _sqlx_migrations table: {}", heal_err);
                    panic!("DATABASE_MIGRATION_FAILURE: {}", e);
                }
            }
        } else {
            tracing::error!("Database migration failed: {}", e);
            panic!("DATABASE_MIGRATION_FAILURE: {}", e);
        }
    }
    // Run database self-healing dynamically to guarantee schema parity
    if let Err(e) = healing::self_heal_schema(&pool).await {
        tracing::error!("Database self-healing failed: {}", e);
    }

    pool
}

/// Initialise both the primary R/W pool and an optional read-replica pool,
/// then wrap them in a `DbRouter` for automatic query routing.
pub async fn init_db_router() -> DbRouter {
    let primary = init_db().await;
    let schema = utils::database_schema();
    let read_opt = pool_router::init_read_replica(&schema).await;
    let has_replica = read_opt.is_some();
    let read = read_opt.unwrap_or_else(|| primary.clone());
    tracing::info!(
        "🗄️  DbRouter ready (read-replica: {})",
        if has_replica {
            "enabled"
        } else {
            "using primary"
        }
    );
    DbRouter::new(primary, read, has_replica)
}