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