sqlx/sqlx-core/src/acquire.rs

use crate::database::Database;
use crate::error::Error;
use crate::pool::{MaybePoolConnection, Pool, PoolConnection};

use crate::transaction::Transaction;
use futures_core::future::BoxFuture;
use std::ops::{Deref, DerefMut};

/// Acquire connections or transactions from a database in a generic way.
///
/// If you want to accept generic database connections that implement
/// [`Acquire`] which then allows you to [`acquire`][`Acquire::acquire`] a
/// connection or [`begin`][`Acquire::begin`] a transaction, then you can do it
/// like that:
///
/// ```rust
/// # use sqlx::{Acquire, postgres::Postgres, error::BoxDynError};
/// # #[cfg(any(postgres_9_6, postgres_15))]
/// async fn run_query<'a, A>(conn: A) -> Result<(), BoxDynError>
/// where
///     A: Acquire<'a, Database = Postgres>,
/// {
///     let mut conn = conn.acquire().await?;
///
///     sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
///     sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
///
///     Ok(())
/// }
/// ```
///
/// If you run into a lifetime error about "implementation of `sqlx::Acquire` is
/// not general enough", the [workaround] looks like this:
///
/// ```rust
/// # use std::future::Future;
/// # use sqlx::{Acquire, postgres::Postgres, error::BoxDynError};
/// # #[cfg(any(postgres_9_6, postgres_15))]
/// fn run_query<'a, 'c, A>(conn: A) -> impl Future<Output = Result<(), BoxDynError>> + Send + 'a
/// where
///     A: Acquire<'c, Database = Postgres> + Send + 'a,
/// {
///     async move {
///         let mut conn = conn.acquire().await?;
///
///         sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
///         sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
///
///         Ok(())
///     }
/// }
/// ```
///
/// However, if you really just want to accept both, a transaction or a
/// connection as an argument to a function, then it's easier to just accept a
/// mutable reference to a database connection like so:
///
/// ```rust
/// # use sqlx::{postgres::PgConnection, error::BoxDynError};
/// # #[cfg(any(postgres_9_6, postgres_15))]
/// async fn run_query(conn: &mut PgConnection) -> Result<(), BoxDynError> {
///     sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
///     sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
///
///     Ok(())
/// }
/// ```
///
/// The downside of this approach is that you have to `acquire` a connection
/// from a pool first and can't directly pass the pool as argument.
///
/// [workaround]: https://github.com/launchbadge/sqlx/issues/1015#issuecomment-767787777
pub trait Acquire<'c> {
    type Database: Database;

    type Connection: Deref<Target = <Self::Database as Database>::Connection> + DerefMut + Send;

    fn acquire(self) -> BoxFuture<'c, Result<Self::Connection, Error>>;

    fn begin(self) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>;
}

impl<'a, DB: Database> Acquire<'a> for &'_ Pool<DB> {
    type Database = DB;

    type Connection = PoolConnection<DB>;

    fn acquire(self) -> BoxFuture<'static, Result<Self::Connection, Error>> {
        Box::pin(self.acquire())
    }

    fn begin(self) -> BoxFuture<'static, Result<Transaction<'a, DB>, Error>> {
        let conn = self.acquire();

        Box::pin(async move {
            Transaction::begin(MaybePoolConnection::PoolConnection(conn.await?)).await
        })
    }
}

#[macro_export]
macro_rules! impl_acquire {
    ($DB:ident, $C:ident) => {
        impl<'c> $crate::acquire::Acquire<'c> for &'c mut $C {
            type Database = $DB;

            type Connection = &'c mut <$DB as $crate::database::Database>::Connection;

            #[inline]
            fn acquire(
                self,
            ) -> futures_core::future::BoxFuture<'c, Result<Self::Connection, $crate::error::Error>>
            {
                Box::pin(futures_util::future::ok(self))
            }

            #[inline]
            fn begin(
                self,
            ) -> futures_core::future::BoxFuture<
                'c,
                Result<$crate::transaction::Transaction<'c, $DB>, $crate::error::Error>,
            > {
                $crate::transaction::Transaction::begin(self)
            }
        }
    };
}