lambda-forge 0.1.7

use crate::Error;

use lambda_http::tracing::warn;
use or_panic::prelude::*;
use sqlx::{PgPool, postgres::PgPoolOptions};
use std::{env, time::Duration};
use tokio::sync::OnceCell;

/// Reusable Postgres connection pool between lambda invocations.
static DB_POOL: OnceCell<PgPool> = OnceCell::const_new();

/// High-level database handle backed by a globally shared Postgres connection pool.
///
/// This provides access to a lazily initialized [`PgPool`] which is stored in a process wide
/// [`OnceCell`]. The underlying pool is created at most once per process (in a cold start
/// worst case), and then reused accross all subsequent invocations.
#[derive(Clone, Copy)]
pub struct Database {
    /// Postgres Connection Pool.
    pool: &'static PgPool,
}
impl Database {
    /// Try to initialize a new [`Database`] instance.
    pub async fn try_new() -> Result<Database, Error> {
        let pool = Database::get_or_init_cnx_pool().await?;

        Ok(Database { pool })
    }

    /// Get the cached database connection pool, or initialize a new one.
    pub async fn get_or_init_cnx_pool() -> Result<&'static PgPool, Error> {
        if let Some(pool) = DB_POOL.get() {
            return Ok(pool);
        }

        warn!("initiating new database connection pool (`DB_POOL` was empty).");

        let db_url = env::var("DB_URL").map_err(|_| {
            Error::InvalidEnvironmentConfig("missing the required `DB_URL` environment variable.")
        })?;
        let db_max_cnx = env::var("DB_MAX_CONNS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(10);

        let pool = PgPoolOptions::new()
            .max_connections(db_max_cnx)
            .acquire_timeout(Duration::from_secs(5))
            .connect(&db_url)
            .await?;

        // GUARD(race conditions): recheck the DB_POOL value.
        if let Some(pool) = DB_POOL.get() {
            return Ok(pool);
        }

        let _ = DB_POOL.set(pool);
        let db_pool = DB_POOL
            .get()
            .or_panic("`DB_POOL` should never NOT be initialized here");

        Ok(db_pool)
    }

    /// Get the database connection pool.
    pub fn pool(&self) -> &PgPool {
        self.pool
    }

    /// Health check for the database, runs simple `SELECT 1` query.
    pub async fn ping(&self) -> Result<(), Error> {
        sqlx::query("SELECT 1").execute(self.pool()).await?;
        Ok(())
    }

    /// Begin a SQL transaction.
    pub async fn begin(&self) -> Result<sqlx::Transaction<'_, sqlx::Postgres>, Error> {
        Ok(self.pool().begin().await?)
    }
}