GeneIT/geneit_backend/src/connections/db_connection.rs

//! # Database connection management

use crate::app_config::AppConfig;
use diesel::result::{DatabaseErrorKind, Error};
use diesel::{Connection, PgConnection};
use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
use std::cell::RefCell;

const MIGRATIONS: EmbeddedMigrations = embed_migrations!();

thread_local! {
    static POSTGRES_CONNECTION: RefCell<Option<PgConnection>> = const { RefCell::new(None) };
}

/// Execute a request on the database
pub fn execute<E, I>(cb: E) -> anyhow::Result<I>
where
    E: FnOnce(&mut PgConnection) -> diesel::QueryResult<I>,
{
    // Establish connection if required
    if POSTGRES_CONNECTION.with(|i| i.borrow().is_none()) {
        let database_url = AppConfig::get().db_connection_chain();
        let conn = PgConnection::establish(&database_url)
            .unwrap_or_else(|_| panic!("Error connecting to {}", database_url));

        POSTGRES_CONNECTION.with(|i| *i.borrow_mut() = Some(conn))
    }

    match POSTGRES_CONNECTION.with(|i| cb(i.borrow_mut().as_mut().unwrap())) {
        Ok(res) => Ok(res),
        Err(e) => {
            if matches!(
                e,
                Error::DatabaseError(DatabaseErrorKind::ClosedConnection, _)
                    | Error::BrokenTransactionManager
            ) {
                log::warn!("Connection to database closed, dropping handle!");
                POSTGRES_CONNECTION.with(|i| *i.borrow_mut() = None)
            }

            log::error!("Database query error! {:?}", e);
            Err(e.into())
        }
    }
}

/// Initialize database connection
pub fn initialize_conn() -> anyhow::Result<()> {
    // Run pending diesel migrations
    execute(|db| {
        let res = db
            .run_pending_migrations(MIGRATIONS)
            .expect("Failed to run database migration!");

        for migration in res {
            log::info!("Executed database migration {migration}")
        }

        Ok(())
    })?;
    Ok(())
}