sqlx_core/
acquire.rs

1use crate::database::Database;
2use crate::error::Error;
3use crate::pool::{MaybePoolConnection, Pool, PoolConnection};
4
5use crate::transaction::Transaction;
6use futures_core::future::BoxFuture;
7use std::ops::{Deref, DerefMut};
8
9/// Acquire connections or transactions from a database in a generic way.
10///
11/// If you want to accept generic database connections that implement
12/// [`Acquire`] which then allows you to [`acquire`][`Acquire::acquire`] a
13/// connection or [`begin`][`Acquire::begin`] a transaction, then you can do it
14/// like that:
15///
16/// ```rust
17/// # use sqlx::{Acquire, postgres::Postgres, error::BoxDynError};
18/// # #[cfg(any(postgres_9_6, postgres_15))]
19/// async fn run_query<'a, A>(conn: A) -> Result<(), BoxDynError>
20/// where
21///     A: Acquire<'a, Database = Postgres>,
22/// {
23///     let mut conn = conn.acquire().await?;
24///
25///     sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
26///     sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
27///
28///     Ok(())
29/// }
30/// ```
31///
32/// If you run into a lifetime error about "implementation of `sqlx::Acquire` is
33/// not general enough", the [workaround] looks like this:
34///
35/// ```rust
36/// # use std::future::Future;
37/// # use sqlx::{Acquire, postgres::Postgres, error::BoxDynError};
38/// # #[cfg(any(postgres_9_6, postgres_15))]
39/// fn run_query<'a, 'c, A>(conn: A) -> impl Future<Output = Result<(), BoxDynError>> + Send + 'a
40/// where
41///     A: Acquire<'c, Database = Postgres> + Send + 'a,
42/// {
43///     async move {
44///         let mut conn = conn.acquire().await?;
45///
46///         sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
47///         sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
48///
49///         Ok(())
50///     }
51/// }
52/// ```
53///
54/// However, if you really just want to accept both, a transaction or a
55/// connection as an argument to a function, then it's easier to just accept a
56/// mutable reference to a database connection like so:
57///
58/// ```rust
59/// # use sqlx::{postgres::PgConnection, error::BoxDynError};
60/// # #[cfg(any(postgres_9_6, postgres_15))]
61/// async fn run_query(conn: &mut PgConnection) -> Result<(), BoxDynError> {
62///     sqlx::query!("SELECT 1 as v").fetch_one(&mut *conn).await?;
63///     sqlx::query!("SELECT 2 as v").fetch_one(&mut *conn).await?;
64///
65///     Ok(())
66/// }
67/// ```
68///
69/// The downside of this approach is that you have to `acquire` a connection
70/// from a pool first and can't directly pass the pool as argument.
71///
72/// [workaround]: https://github.com/launchbadge/sqlx/issues/1015#issuecomment-767787777
73pub trait Acquire<'c> {
74    type Database: Database;
75
76    type Connection: Deref<Target = <Self::Database as Database>::Connection> + DerefMut + Send;
77
78    fn acquire(self) -> BoxFuture<'c, Result<Self::Connection, Error>>;
79
80    fn begin(self) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>;
81}
82
83impl<'a, DB: Database> Acquire<'a> for &'_ Pool<DB> {
84    type Database = DB;
85
86    type Connection = PoolConnection<DB>;
87
88    fn acquire(self) -> BoxFuture<'static, Result<Self::Connection, Error>> {
89        Box::pin(self.acquire())
90    }
91
92    fn begin(self) -> BoxFuture<'static, Result<Transaction<'a, DB>, Error>> {
93        let conn = self.acquire();
94
95        Box::pin(async move {
96            Transaction::begin(MaybePoolConnection::PoolConnection(conn.await?)).await
97        })
98    }
99}
100
101#[macro_export]
102macro_rules! impl_acquire {
103    ($DB:ident, $C:ident) => {
104        impl<'c> $crate::acquire::Acquire<'c> for &'c mut $C {
105            type Database = $DB;
106
107            type Connection = &'c mut <$DB as $crate::database::Database>::Connection;
108
109            #[inline]
110            fn acquire(
111                self,
112            ) -> futures_core::future::BoxFuture<'c, Result<Self::Connection, $crate::error::Error>>
113            {
114                Box::pin(futures_util::future::ok(self))
115            }
116
117            #[inline]
118            fn begin(
119                self,
120            ) -> futures_core::future::BoxFuture<
121                'c,
122                Result<$crate::transaction::Transaction<'c, $DB>, $crate::error::Error>,
123            > {
124                $crate::transaction::Transaction::begin(self)
125            }
126        }
127    };
128}