sqlx-record/src/conn_provider.rs

use sqlx::pool::PoolConnection;

#[cfg(feature = "mysql")]
use sqlx::{MySql, MySqlConnection, MySqlPool, Transaction};

#[cfg(feature = "postgres")]
use sqlx::{Postgres, PgConnection, PgPool, Transaction};

#[cfg(feature = "sqlite")]
use sqlx::{Sqlite, SqliteConnection, SqlitePool, Transaction};

// ============================================================================
// MySQL Implementation
// ============================================================================

#[cfg(feature = "mysql")]
pub enum ConnProvider<'a> {
    /// Stores a reference to an existing connection
    Borrowed {
        conn: &'a mut PoolConnection<MySql>,
    },
    /// Stores an owned connection acquired from a pool
    Owned {
        pool: MySqlPool,
        conn: Option<PoolConnection<MySql>>,
    },
    /// Stores a reference to a transaction
    Transaction {
        tx: &'a mut Transaction<'static, MySql>,
    },
}

#[cfg(feature = "mysql")]
impl<'a> ConnProvider<'a> {
    /// Create a ConnProvider from a borrowed connection reference
    pub fn from_ref(conn: &'a mut PoolConnection<MySql>) -> Self {
        ConnProvider::Borrowed { conn }
    }

    /// Create a ConnProvider that will lazily acquire a connection from the pool
    pub fn from_pool(pool: MySqlPool) -> Self {
        ConnProvider::Owned { pool, conn: None }
    }

    /// Create a ConnProvider from a borrowed transaction reference
    pub fn from_tx(tx: &'a mut Transaction<'static, MySql>) -> Self {
        ConnProvider::Transaction { tx }
    }

    /// Get a mutable reference to the underlying connection.
    /// For borrowed connections, returns the reference directly.
    /// For owned connections, lazily acquires from pool on first call.
    /// For transactions, returns the transaction's underlying connection.
    pub async fn get_conn(&mut self) -> Result<&mut MySqlConnection, sqlx::Error> {
        match self {
            ConnProvider::Borrowed { conn } => Ok(&mut **conn),
            ConnProvider::Owned { pool, conn } => {
                if conn.is_none() {
                    *conn = Some(pool.acquire().await?);
                }
                Ok(&mut **conn.as_mut().unwrap())
            }
            ConnProvider::Transaction { tx } => Ok(&mut **tx),
        }
    }
}

// ============================================================================
// PostgreSQL Implementation
// ============================================================================

#[cfg(feature = "postgres")]
pub enum ConnProvider<'a> {
    /// Stores a reference to an existing connection
    Borrowed {
        conn: &'a mut PoolConnection<Postgres>,
    },
    /// Stores an owned connection acquired from a pool
    Owned {
        pool: PgPool,
        conn: Option<PoolConnection<Postgres>>,
    },
    /// Stores a reference to a transaction
    Transaction {
        tx: &'a mut Transaction<'static, Postgres>,
    },
}

#[cfg(feature = "postgres")]
impl<'a> ConnProvider<'a> {
    /// Create a ConnProvider from a borrowed connection reference
    pub fn from_ref(conn: &'a mut PoolConnection<Postgres>) -> Self {
        ConnProvider::Borrowed { conn }
    }

    /// Create a ConnProvider that will lazily acquire a connection from the pool
    pub fn from_pool(pool: PgPool) -> Self {
        ConnProvider::Owned { pool, conn: None }
    }

    /// Create a ConnProvider from a borrowed transaction reference
    pub fn from_tx(tx: &'a mut Transaction<'static, Postgres>) -> Self {
        ConnProvider::Transaction { tx }
    }

    /// Get a mutable reference to the underlying connection.
    /// For borrowed connections, returns the reference directly.
    /// For owned connections, lazily acquires from pool on first call.
    /// For transactions, returns the transaction's underlying connection.
    pub async fn get_conn(&mut self) -> Result<&mut PgConnection, sqlx::Error> {
        match self {
            ConnProvider::Borrowed { conn } => Ok(&mut **conn),
            ConnProvider::Owned { pool, conn } => {
                if conn.is_none() {
                    *conn = Some(pool.acquire().await?);
                }
                Ok(&mut **conn.as_mut().unwrap())
            }
            ConnProvider::Transaction { tx } => Ok(&mut **tx),
        }
    }
}

// ============================================================================
// SQLite Implementation
// ============================================================================

#[cfg(feature = "sqlite")]
pub enum ConnProvider<'a> {
    /// Stores a reference to an existing connection
    Borrowed {
        conn: &'a mut PoolConnection<Sqlite>,
    },
    /// Stores an owned connection acquired from a pool
    Owned {
        pool: SqlitePool,
        conn: Option<PoolConnection<Sqlite>>,
    },
    /// Stores a reference to a transaction
    Transaction {
        tx: &'a mut Transaction<'static, Sqlite>,
    },
}

#[cfg(feature = "sqlite")]
impl<'a> ConnProvider<'a> {
    /// Create a ConnProvider from a borrowed connection reference
    pub fn from_ref(conn: &'a mut PoolConnection<Sqlite>) -> Self {
        ConnProvider::Borrowed { conn }
    }

    /// Create a ConnProvider that will lazily acquire a connection from the pool
    pub fn from_pool(pool: SqlitePool) -> Self {
        ConnProvider::Owned { pool, conn: None }
    }

    /// Create a ConnProvider from a borrowed transaction reference
    pub fn from_tx(tx: &'a mut Transaction<'static, Sqlite>) -> Self {
        ConnProvider::Transaction { tx }
    }

    /// Get a mutable reference to the underlying connection.
    /// For borrowed connections, returns the reference directly.
    /// For owned connections, lazily acquires from pool on first call.
    /// For transactions, returns the transaction's underlying connection.
    pub async fn get_conn(&mut self) -> Result<&mut SqliteConnection, sqlx::Error> {
        match self {
            ConnProvider::Borrowed { conn } => Ok(&mut **conn),
            ConnProvider::Owned { pool, conn } => {
                if conn.is_none() {
                    *conn = Some(pool.acquire().await?);
                }
                Ok(&mut **conn.as_mut().unwrap())
            }
            ConnProvider::Transaction { tx } => Ok(&mut **tx),
        }
    }
}