sea_orm/driver/
sqlx_common.rs

1use crate::{ConnAcquireErr, ConnectOptions, DbErr, RuntimeErr};
2use std::{sync::Arc, time::Duration};
3
4/// Callback stored for a `before_acquire` hook on [`ConnectOptions`].
5///
6/// This mirrors the signature accepted by SQLx's
7/// [`PoolOptions::before_acquire`][sqlx::pool::PoolOptions::before_acquire] for the given
8/// database backend `DB`. Note the fully-qualified [`futures_util::future::BoxFuture`]: the
9/// crate-local `BoxFuture` alias collapses to `T` under the `sync` feature, which would not
10/// match SQLx's signature.
11pub(crate) type BeforeAcquireFn<DB> = Arc<
12    dyn for<'c> Fn(
13            &'c mut <DB as sqlx::Database>::Connection,
14            sqlx::pool::PoolConnectionMetadata,
15        ) -> futures_util::future::BoxFuture<'c, Result<bool, sqlx::Error>>
16        + Send
17        + Sync,
18>;
19
20/// Converts an [sqlx::error] execution error to a [DbErr]
21pub fn sqlx_error_to_exec_err(err: sqlx::Error) -> DbErr {
22    DbErr::Exec(RuntimeErr::SqlxError(err.into()))
23}
24
25/// Converts an [sqlx::error] query error to a [DbErr]
26pub fn sqlx_error_to_query_err(err: sqlx::Error) -> DbErr {
27    DbErr::Query(RuntimeErr::SqlxError(err.into()))
28}
29
30/// Converts an [sqlx::error] connection error to a [DbErr]
31pub fn sqlx_error_to_conn_err(err: sqlx::Error) -> DbErr {
32    DbErr::Conn(RuntimeErr::SqlxError(err.into()))
33}
34
35/// Converts an [sqlx::error] error to a [DbErr]
36pub fn sqlx_map_err_ignore_not_found<T: std::fmt::Debug>(
37    err: Result<Option<T>, sqlx::Error>,
38) -> Result<Option<T>, DbErr> {
39    if let Err(sqlx::Error::RowNotFound) = err {
40        Ok(None)
41    } else {
42        err.map_err(sqlx_error_to_query_err)
43    }
44}
45
46/// Converts an [sqlx::error] error to a [DbErr]
47pub fn sqlx_conn_acquire_err(sqlx_err: sqlx::Error) -> DbErr {
48    match sqlx_err {
49        sqlx::Error::PoolTimedOut => DbErr::ConnectionAcquire(ConnAcquireErr::Timeout),
50        sqlx::Error::PoolClosed => DbErr::ConnectionAcquire(ConnAcquireErr::ConnectionClosed),
51        _ => DbErr::Conn(RuntimeErr::SqlxError(sqlx_err.into())),
52    }
53}
54
55impl ConnectOptions {
56    /// Convert [ConnectOptions] into [sqlx::pool::PoolOptions]
57    pub fn sqlx_pool_options<DB>(self) -> sqlx::pool::PoolOptions<DB>
58    where
59        DB: sqlx::Database,
60    {
61        let mut opt = sqlx::pool::PoolOptions::new();
62        if let Some(max_connections) = self.max_connections {
63            opt = opt.max_connections(max_connections);
64        }
65        if let Some(min_connections) = self.min_connections {
66            opt = opt.min_connections(min_connections);
67        }
68        if let Some(connect_timeout) = self.connect_timeout {
69            opt = opt.acquire_timeout(connect_timeout);
70        }
71        if let Some(idle_timeout) = self.idle_timeout {
72            opt = opt.idle_timeout(idle_timeout);
73        }
74        if let Some(acquire_timeout) = self.acquire_timeout {
75            opt = opt.acquire_timeout(acquire_timeout);
76        }
77        if let Some(max_lifetime) = self.max_lifetime {
78            opt = opt.max_lifetime(max_lifetime);
79        }
80        opt = opt.test_before_acquire(self.test_before_acquire);
81        opt
82    }
83
84    /// Install the composed `before_acquire` hook onto a [`sqlx::pool::PoolOptions`].
85    ///
86    /// SQLx exposes a single `before_acquire` slot whose setter *replaces* rather than
87    /// composes, and offers no getter to read it back. To let the idle-ping shorthand
88    /// ([`ConnectOptions::test_before_acquire_if_idle_for`]) coexist with a user-provided
89    /// per-backend callback ([`ConnectOptions::map_sqlx_postgres_before_acquire`] and
90    /// friends), this composes both into one closure: the idle-ping runs first, then the
91    /// user callback. When a ping threshold is set, `test_before_acquire` is forced off so
92    /// the connection is pinged only past the threshold rather than on every acquire.
93    ///
94    /// Returns `opt` untouched when neither option is configured, so callers that opt into
95    /// nothing get byte-for-byte the previous behavior.
96    pub(crate) fn apply_before_acquire<DB>(
97        mut opt: sqlx::pool::PoolOptions<DB>,
98        ping_after_idle: Option<Duration>,
99        user_cb: Option<BeforeAcquireFn<DB>>,
100    ) -> sqlx::pool::PoolOptions<DB>
101    where
102        DB: sqlx::Database,
103    {
104        use sqlx::Connection;
105
106        if ping_after_idle.is_none() && user_cb.is_none() {
107            return opt;
108        }
109        if ping_after_idle.is_some() {
110            opt = opt.test_before_acquire(false);
111        }
112        opt.before_acquire(move |conn, meta| {
113            let user_cb = user_cb.clone();
114            Box::pin(async move {
115                if let Some(threshold) = ping_after_idle {
116                    // `idle_for` is `Copy`; read it before `meta` is moved into the user callback.
117                    // `>=` matches the "idle for at least `threshold`" contract documented on
118                    // `ConnectOptions::test_before_acquire_if_idle_for`.
119                    if meta.idle_for >= threshold {
120                        conn.ping().await?;
121                    }
122                }
123                match user_cb {
124                    Some(user_cb) => user_cb(conn, meta).await,
125                    None => Ok(true),
126                }
127            })
128        })
129    }
130}
131
132#[cfg(all(test, feature = "sqlx-postgres"))]
133mod tests {
134    use crate::ConnectOptions;
135    use sqlx::Connection;
136    use std::time::Duration;
137
138    #[test]
139    fn idle_shorthand_disables_test_before_acquire() {
140        let mut opt = ConnectOptions::new("postgres://localhost/db");
141        assert!(opt.get_test_before_acquire());
142        assert_eq!(opt.get_test_before_acquire_if_idle_for(), None);
143
144        opt.test_before_acquire_if_idle_for(Duration::from_secs(30));
145        assert!(!opt.get_test_before_acquire());
146        assert_eq!(
147            opt.get_test_before_acquire_if_idle_for(),
148            Some(Duration::from_secs(30))
149        );
150    }
151
152    #[test]
153    fn compose_shorthand_and_user_callback() {
154        let mut opt = ConnectOptions::new("postgres://localhost/db");
155        opt.test_before_acquire_if_idle_for(Duration::from_secs(30))
156            .map_sqlx_postgres_before_acquire(|conn, _meta| {
157                Box::pin(async move {
158                    conn.ping().await?;
159                    Ok(true)
160                })
161            });
162
163        // Composing both into SQLx's single `before_acquire` slot type-checks and returns a
164        // usable `PoolOptions`. Behavioral ping timing requires a live pool, covered elsewhere.
165        let pool_opts = ConnectOptions::apply_before_acquire::<sqlx::Postgres>(
166            sqlx::pool::PoolOptions::new(),
167            opt.get_test_before_acquire_if_idle_for(),
168            opt.pg_before_acquire_fn.clone(),
169        );
170        let _ = pool_opts;
171    }
172
173    #[test]
174    fn apply_before_acquire_noop_when_unset() {
175        // With neither option set, the helper must return the options untouched.
176        let opts = ConnectOptions::apply_before_acquire::<sqlx::Postgres>(
177            sqlx::pool::PoolOptions::new(),
178            None,
179            None,
180        );
181        let _ = opts;
182    }
183}
sea_orm/driver/sqlx_common.rs

sea_orm/driver/
sqlx_common.rs
