fletch_orm/

pool.rs

//! Connection pool wrapper around `sqlx::Pool`.

use std::time::Duration;

use sqlx::Database;

use crate::error::FletchError;

/// A connection pool wrapping [`sqlx::Pool`].
///
/// Generic over `DB: sqlx::Database` so that the same API works with
/// SQLite, PostgreSQL, and MySQL while retaining compile-time type safety.
#[derive(Debug)]
pub struct Pool<DB: Database> {
    inner: sqlx::Pool<DB>,
}

impl<DB: Database> Clone for Pool<DB> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<DB: Database> Pool<DB> {
    /// Connect to the database at the given URL.
    ///
    /// This creates a connection pool with SQLx's default pool options.
    pub async fn connect(url: &str) -> Result<Self, FletchError> {
        let pool = sqlx::Pool::<DB>::connect(url).await?;
        Ok(Self { inner: pool })
    }

    /// Returns a [`PoolBuilder`] for configuring pool options before connecting.
    pub fn builder() -> PoolBuilder<DB> {
        PoolBuilder {
            options: sqlx::pool::PoolOptions::new(),
        }
    }

    /// Close the pool, waiting for all connections to be returned.
    pub async fn close(&self) {
        self.inner.close().await;
    }

    /// Returns a reference to the underlying sqlx pool.
    ///
    /// Escape hatch for advanced or raw SQLx access. Prefer fletch CRUD and
    /// the query builder where possible.
    pub fn inner(&self) -> &sqlx::Pool<DB> {
        &self.inner
    }

    /// Returns the number of connections currently active (including idle).
    pub fn size(&self) -> u32 {
        self.inner.size()
    }

    /// Returns the number of idle connections in the pool.
    pub fn num_idle(&self) -> usize {
        self.inner.num_idle()
    }

    /// Begin a new transaction.
    pub async fn begin(&self) -> Result<Transaction<'static, DB>, FletchError> {
        let tx = self.inner.begin().await?;
        Ok(Transaction { inner: tx })
    }
}

/// Builder for configuring a [`Pool`] before connecting.
///
/// Created via [`Pool::builder()`].
pub struct PoolBuilder<DB: Database> {
    options: sqlx::pool::PoolOptions<DB>,
}

impl<DB: Database> PoolBuilder<DB> {
    /// Set the maximum number of connections the pool should maintain.
    pub fn max_connections(mut self, max: u32) -> Self {
        self.options = self.options.max_connections(max);
        self
    }

    /// Set the minimum number of connections the pool should maintain.
    pub fn min_connections(mut self, min: u32) -> Self {
        self.options = self.options.min_connections(min);
        self
    }

    /// Set the maximum time a connection can sit idle before being closed.
    pub fn idle_timeout(mut self, timeout: Duration) -> Self {
        self.options = self.options.idle_timeout(timeout);
        self
    }

    /// Set the maximum lifetime of individual connections.
    pub fn max_lifetime(mut self, lifetime: Duration) -> Self {
        self.options = self.options.max_lifetime(lifetime);
        self
    }

    /// Set the maximum time to wait when acquiring a connection from the pool.
    pub fn acquire_timeout(mut self, timeout: Duration) -> Self {
        self.options = self.options.acquire_timeout(timeout);
        self
    }

    /// Connect to the database at the given URL with the configured options.
    pub async fn connect(self, url: &str) -> Result<Pool<DB>, FletchError> {
        let pool = self.options.connect(url).await?;
        Ok(Pool { inner: pool })
    }
}

/// Declarative pool configuration.
///
/// All fields are optional; unset fields use SQLx defaults.
/// Can be converted into a [`PoolBuilder`] via [`PoolConfig::into_builder`].
///
/// # Environment variables
///
/// [`PoolConfig::from_env()`] reads:
/// - `DATABASE_MAX_CONNECTIONS`
/// - `DATABASE_MIN_CONNECTIONS`
/// - `DATABASE_ACQUIRE_TIMEOUT_SECS`
/// - `DATABASE_IDLE_TIMEOUT_SECS`
/// - `DATABASE_MAX_LIFETIME_SECS`
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
pub struct PoolConfig {
    pub max_connections: Option<u32>,
    pub min_connections: Option<u32>,
    pub acquire_timeout_secs: Option<u64>,
    pub idle_timeout_secs: Option<u64>,
    pub max_lifetime_secs: Option<u64>,
}

impl PoolConfig {
    /// Load pool configuration from environment variables.
    ///
    /// Missing or unparseable variables are silently ignored (treated as unset).
    pub fn from_env() -> Self {
        Self {
            max_connections: std::env::var("DATABASE_MAX_CONNECTIONS")
                .ok()
                .and_then(|v| v.parse().ok()),
            min_connections: std::env::var("DATABASE_MIN_CONNECTIONS")
                .ok()
                .and_then(|v| v.parse().ok()),
            acquire_timeout_secs: std::env::var("DATABASE_ACQUIRE_TIMEOUT_SECS")
                .ok()
                .and_then(|v| v.parse().ok()),
            idle_timeout_secs: std::env::var("DATABASE_IDLE_TIMEOUT_SECS")
                .ok()
                .and_then(|v| v.parse().ok()),
            max_lifetime_secs: std::env::var("DATABASE_MAX_LIFETIME_SECS")
                .ok()
                .and_then(|v| v.parse().ok()),
        }
    }

    /// Convert this config into a [`PoolBuilder`], applying all set fields.
    pub fn into_builder<DB: Database>(self) -> PoolBuilder<DB> {
        let mut builder = PoolBuilder {
            options: sqlx::pool::PoolOptions::new(),
        };
        if let Some(max) = self.max_connections {
            builder = builder.max_connections(max);
        }
        if let Some(min) = self.min_connections {
            builder = builder.min_connections(min);
        }
        if let Some(secs) = self.acquire_timeout_secs {
            builder = builder.acquire_timeout(Duration::from_secs(secs));
        }
        if let Some(secs) = self.idle_timeout_secs {
            builder = builder.idle_timeout(Duration::from_secs(secs));
        }
        if let Some(secs) = self.max_lifetime_secs {
            builder = builder.max_lifetime(Duration::from_secs(secs));
        }
        builder
    }
}

/// A database transaction with auto-rollback on drop.
///
/// If the transaction is dropped without calling [`commit`](Transaction::commit),
/// it will automatically roll back.
pub struct Transaction<'a, DB: Database> {
    inner: sqlx::Transaction<'a, DB>,
}

impl<'a, DB: Database> Transaction<'a, DB> {
    /// Commit the transaction.
    pub async fn commit(self) -> Result<(), FletchError> {
        self.inner.commit().await?;
        Ok(())
    }

    /// Explicitly roll back the transaction.
    pub async fn rollback(self) -> Result<(), FletchError> {
        self.inner.rollback().await?;
        Ok(())
    }

    /// Returns a mutable reference to the underlying sqlx transaction.
    ///
    /// Escape hatch for advanced or raw SQLx access. Prefer fletch CRUD and
    /// the query builder where possible.
    pub fn inner_mut(&mut self) -> &mut sqlx::Transaction<'a, DB> {
        &mut self.inner
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use sqlx::Sqlite;

    #[tokio::test]
    async fn pool_connect_sqlite_memory() {
        let pool = Pool::<Sqlite>::connect("sqlite::memory:").await.unwrap();
        // Verify we can acquire a connection
        {
            let _conn = pool.inner().acquire().await.unwrap();
        }
        pool.close().await;
    }

    #[tokio::test]
    async fn pool_begin_and_commit() {
        let pool = Pool::<Sqlite>::connect("sqlite::memory:").await.unwrap();

        // Create a table using the pool directly
        sqlx::query("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
            .execute(pool.inner())
            .await
            .unwrap();

        // Begin a transaction and insert
        let mut tx = pool.begin().await.unwrap();
        sqlx::query("INSERT INTO test (id, name) VALUES (1, 'Alice')")
            .execute(&mut **tx.inner_mut())
            .await
            .unwrap();
        tx.commit().await.unwrap();

        // Verify the row exists
        let row: (i64, String) = sqlx::query_as("SELECT id, name FROM test WHERE id = 1")
            .fetch_one(pool.inner())
            .await
            .unwrap();
        assert_eq!(row, (1, "Alice".to_string()));

        pool.close().await;
    }

    #[tokio::test]
    async fn transaction_auto_rollback_on_drop() {
        let pool = Pool::<Sqlite>::connect("sqlite::memory:").await.unwrap();

        sqlx::query("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
            .execute(pool.inner())
            .await
            .unwrap();

        // Begin a transaction, insert, then drop without committing
        {
            let mut tx = pool.begin().await.unwrap();
            sqlx::query("INSERT INTO test (id, name) VALUES (1, 'Alice')")
                .execute(&mut **tx.inner_mut())
                .await
                .unwrap();
            // tx dropped here without commit -> rollback
        }

        // Verify the row does NOT exist
        let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM test")
            .fetch_one(pool.inner())
            .await
            .unwrap();
        assert_eq!(count.0, 0);

        pool.close().await;
    }

    #[tokio::test]
    async fn transaction_explicit_rollback() {
        let pool = Pool::<Sqlite>::connect("sqlite::memory:").await.unwrap();

        sqlx::query("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
            .execute(pool.inner())
            .await
            .unwrap();

        let mut tx = pool.begin().await.unwrap();
        sqlx::query("INSERT INTO test (id, name) VALUES (1, 'Alice')")
            .execute(&mut **tx.inner_mut())
            .await
            .unwrap();
        tx.rollback().await.unwrap();

        let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM test")
            .fetch_one(pool.inner())
            .await
            .unwrap();
        assert_eq!(count.0, 0);

        pool.close().await;
    }

    #[tokio::test]
    async fn pool_builder_connects_with_options() {
        let pool = Pool::<Sqlite>::builder()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .unwrap();

        // Verify we can acquire a connection
        let conn = pool.inner().acquire().await.unwrap();
        drop(conn);
        pool.close().await;
    }

    #[test]
    fn pool_config_default_is_all_none() {
        let config = PoolConfig::default();
        assert!(config.max_connections.is_none());
        assert!(config.min_connections.is_none());
        assert!(config.acquire_timeout_secs.is_none());
        assert!(config.idle_timeout_secs.is_none());
        assert!(config.max_lifetime_secs.is_none());
    }

    #[test]
    fn pool_config_into_builder_applies_fields() {
        let config = PoolConfig {
            max_connections: Some(20),
            min_connections: Some(2),
            acquire_timeout_secs: Some(10),
            idle_timeout_secs: Some(300),
            max_lifetime_secs: Some(1800),
        };

        let builder = config.into_builder::<Sqlite>();
        // Verify the builder was created (we can't inspect internal options,
        // but we verify it doesn't panic)
        let _ = builder;
    }

    #[tokio::test]
    async fn pool_config_into_builder_connects() {
        let config = PoolConfig {
            max_connections: Some(1),
            min_connections: None,
            acquire_timeout_secs: None,
            idle_timeout_secs: None,
            max_lifetime_secs: None,
        };

        let pool = config
            .into_builder::<Sqlite>()
            .connect("sqlite::memory:")
            .await
            .unwrap();

        let conn = pool.inner().acquire().await.unwrap();
        drop(conn);
        pool.close().await;
    }

    #[tokio::test]
    async fn pool_introspection() {
        let pool = Pool::<Sqlite>::builder()
            .max_connections(5)
            .connect("sqlite::memory:")
            .await
            .unwrap();

        // Acquire a connection to ensure the pool has at least one active connection
        let conn = pool.inner().acquire().await.unwrap();

        assert!(pool.size() >= 1);
        // num_idle should be less than size since we hold a connection
        assert!(pool.num_idle() < usize::try_from(pool.size()).unwrap());

        drop(conn);
        pool.close().await;
    }
}