From c134c4b8cfc763654b8b3dd621918e1e339049b4 Mon Sep 17 00:00:00 2001 From: Chris Tsang Date: Thu, 23 Jul 2026 11:52:50 +0100 Subject: [PATCH] Add before_acquire passthrough and idle-ping shorthand to ConnectOptions Expose SQLx's `before_acquire` pool hook, which SeaORM previously only reachable through the raw `map_sqlx__pool_opts` escape hatch. - `test_before_acquire_if_idle_for(Duration)`: DB-agnostic shorthand that pings a pooled connection before handing it out only once it has been idle for at least the given duration, instead of on every acquire. Sets `test_before_acquire(false)`. - `map_sqlx_{postgres,mysql,sqlite}_before_acquire`: per-backend passthrough of SQLx's `before_acquire` callback. - `get_test_before_acquire` / `get_test_before_acquire_if_idle_for` getters. Because SQLx's single `before_acquire` slot replaces rather than composes and has no getter, a generic `apply_before_acquire` helper composes the idle-ping shorthand with any user-provided callback into one closure (idle-ping first, then the user callback). The helper returns the pool options untouched when neither option is set, so existing behavior is unchanged. Purely additive: no existing signature, default, or behavior is changed. --- CHANGELOG.md | 8 ++ src/database/mod.rs | 170 ++++++++++++++++++++++++++++++++++++ src/driver/sqlx_common.rs | 117 +++++++++++++++++++++++++ src/driver/sqlx_mysql.rs | 7 ++ src/driver/sqlx_postgres.rs | 7 ++ src/driver/sqlx_sqlite.rs | 7 ++ 6 files changed, 316 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6426d5d627..9e875790b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/) and this project adheres to [Semantic Versioning](http://semver.org/). +## [Unreleased] + +### Added + +- `ConnectOptions::test_before_acquire_if_idle_for(Duration)` — ping a pooled connection before it is handed out only once it has been idle for at least the given duration, instead of on every acquire (`test_before_acquire`). Setting it disables `test_before_acquire`. +- `ConnectOptions::map_sqlx_postgres_before_acquire` / `map_sqlx_mysql_before_acquire` / `map_sqlx_sqlite_before_acquire` — install a per-backend SQLx `before_acquire` callback. Composes with the idle-ping shorthand above: the idle-ping runs first, then the callback. +- `ConnectOptions::get_test_before_acquire` / `get_test_before_acquire_if_idle_for` getters. + ## 2.0.0 - 2026-07-19 ### Release Candidates diff --git a/src/database/mod.rs b/src/database/mod.rs index 057c9525c5..97cc3fd23a 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -127,6 +127,9 @@ pub struct ConnectOptions { /// Statement timeout (PostgreSQL only) pub(crate) statement_timeout: Option, pub(crate) test_before_acquire: bool, + /// If set, a pooled connection is pinged before being handed out only when it has been + /// idle for at least this long (see [`ConnectOptions::test_before_acquire_if_idle_for`]). + pub(crate) test_before_acquire_if_idle_for: Option, /// Only establish connections to the DB as needed. If set to `true`, the db connection will /// be created using SQLx's [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy) /// method. @@ -155,6 +158,16 @@ pub struct ConnectOptions { #[debug(skip)] pub(crate) sqlite_opts_fn: Option SqliteConnectOptions + Send + Sync>>, + + #[cfg(feature = "sqlx-mysql")] + #[debug(skip)] + pub(crate) mysql_before_acquire_fn: Option>, + #[cfg(feature = "sqlx-postgres")] + #[debug(skip)] + pub(crate) pg_before_acquire_fn: Option>, + #[cfg(feature = "sqlx-sqlite")] + #[debug(skip)] + pub(crate) sqlite_before_acquire_fn: Option>, } impl Database { @@ -264,6 +277,7 @@ impl ConnectOptions { application_name: None, statement_timeout: None, test_before_acquire: true, + test_before_acquire_if_idle_for: None, connect_lazy: false, after_connect: None, #[cfg(feature = "sqlx-mysql")] @@ -278,6 +292,12 @@ impl ConnectOptions { pg_opts_fn: None, #[cfg(feature = "sqlx-sqlite")] sqlite_opts_fn: None, + #[cfg(feature = "sqlx-mysql")] + mysql_before_acquire_fn: None, + #[cfg(feature = "sqlx-postgres")] + pg_before_acquire_fn: None, + #[cfg(feature = "sqlx-sqlite")] + sqlite_before_acquire_fn: None, } } @@ -457,11 +477,61 @@ impl ConnectOptions { } /// If true, the connection will be pinged upon acquiring from the pool (default true). + /// + /// See [`test_before_acquire_if_idle_for`](Self::test_before_acquire_if_idle_for) for the + /// cheaper "only ping stale connections" variant. pub fn test_before_acquire(&mut self, value: bool) -> &mut Self { self.test_before_acquire = value; self } + /// Get whether a pooled connection is pinged on every acquire (default true). + pub fn get_test_before_acquire(&self) -> bool { + self.test_before_acquire + } + + /// Ping a pooled connection before handing it out, but only once it has been idle for at + /// least `idle`. + /// + /// [`test_before_acquire`](Self::test_before_acquire) pings on *every* acquire, which adds + /// a round-trip to each checkout. This shorthand instead pings only connections that have + /// been idle long enough to plausibly have been dropped by the server, a proxy, or a + /// firewall — the common failure mode — while letting hot connections through untouched. + /// + /// Calling this sets [`test_before_acquire`](Self::test_before_acquire) to `false` (SQLx + /// runs the per-acquire ping *and* this hook, so leaving it enabled would ping on every + /// acquire regardless and defeat the threshold). + /// + /// It composes with a per-backend [`map_sqlx_postgres_before_acquire`] callback (and its + /// MySQL / SQLite counterparts): the idle-ping runs first, then your callback. + /// + /// Applies only to pools built through [`Database::connect`]. Pools adopted via + /// `SqlxPostgresConnector::from_sqlx_postgres_pool` (and the MySQL / SQLite equivalents) + /// bypass [`ConnectOptions`] entirely — configure `before_acquire` on your own + /// [`sqlx::pool::PoolOptions`] in that case. + /// + /// # Example + /// ``` + /// # use sea_orm::ConnectOptions; + /// # use std::time::Duration; + /// let mut opt = ConnectOptions::new("postgres://localhost/db"); + /// opt.test_before_acquire_if_idle_for(Duration::from_secs(30)); + /// assert_eq!(opt.get_test_before_acquire(), false); + /// ``` + /// + /// [`map_sqlx_postgres_before_acquire`]: Self::map_sqlx_postgres_before_acquire + pub fn test_before_acquire_if_idle_for(&mut self, idle: Duration) -> &mut Self { + self.test_before_acquire = false; + self.test_before_acquire_if_idle_for = Some(idle); + self + } + + /// Get the idle threshold set by + /// [`test_before_acquire_if_idle_for`](Self::test_before_acquire_if_idle_for), if any. + pub fn get_test_before_acquire_if_idle_for(&self) -> Option { + self.test_before_acquire_if_idle_for + } + /// If set to `true`, the db connection pool will be created using SQLx's /// [connect_lazy](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html#method.connect_lazy) method. pub fn connect_lazy(&mut self, value: bool) -> &mut Self { @@ -564,4 +634,104 @@ impl ConnectOptions { self.sqlite_pool_opts_fn = Some(Arc::new(f)); self } + + #[cfg(feature = "sqlx-mysql")] + #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-mysql")))] + /// Set a `before_acquire` callback run on an idle pooled MySQL connection before it is + /// handed out. + /// + /// Return `Ok(true)` to use the connection, or `Ok(false)`/`Err(_)` to discard it and let + /// the pool try another (opening a fresh one if needed). Composes after the idle-ping + /// installed by + /// [`test_before_acquire_if_idle_for`](Self::test_before_acquire_if_idle_for) — that runs + /// first, then this callback. + /// + /// Applies only to pools built through [`Database::connect`], not to pools adopted via + /// `SqlxMySqlConnector::from_sqlx_mysql_pool`. + pub fn map_sqlx_mysql_before_acquire(&mut self, f: F) -> &mut Self + where + F: for<'c> Fn( + &'c mut sqlx::mysql::MySqlConnection, + sqlx::pool::PoolConnectionMetadata, + ) + -> futures_util::future::BoxFuture<'c, Result> + + Send + + Sync + + 'static, + { + self.mysql_before_acquire_fn = Some(Arc::new(f)); + self + } + + #[cfg(feature = "sqlx-postgres")] + #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-postgres")))] + /// Set a `before_acquire` callback run on an idle pooled PostgreSQL connection before it + /// is handed out. + /// + /// Return `Ok(true)` to use the connection, or `Ok(false)`/`Err(_)` to discard it and let + /// the pool try another (opening a fresh one if needed). Composes after the idle-ping + /// installed by + /// [`test_before_acquire_if_idle_for`](Self::test_before_acquire_if_idle_for) — that runs + /// first, then this callback. + /// + /// Applies only to pools built through [`Database::connect`], not to pools adopted via + /// `SqlxPostgresConnector::from_sqlx_postgres_pool`. + /// + /// # Example + /// ``` + /// # use sea_orm::ConnectOptions; + /// # use std::time::Duration; + /// use sea_orm::sqlx::Connection; + /// + /// let mut opt = ConnectOptions::new("postgres://localhost/db"); + /// opt.test_before_acquire_if_idle_for(Duration::from_secs(30)) + /// .map_sqlx_postgres_before_acquire(|conn, _meta| { + /// Box::pin(async move { + /// conn.ping().await?; + /// Ok(true) + /// }) + /// }); + /// ``` + pub fn map_sqlx_postgres_before_acquire(&mut self, f: F) -> &mut Self + where + F: for<'c> Fn( + &'c mut sqlx::postgres::PgConnection, + sqlx::pool::PoolConnectionMetadata, + ) + -> futures_util::future::BoxFuture<'c, Result> + + Send + + Sync + + 'static, + { + self.pg_before_acquire_fn = Some(Arc::new(f)); + self + } + + #[cfg(feature = "sqlx-sqlite")] + #[cfg_attr(docsrs, doc(cfg(feature = "sqlx-sqlite")))] + /// Set a `before_acquire` callback run on an idle pooled SQLite connection before it is + /// handed out. + /// + /// Return `Ok(true)` to use the connection, or `Ok(false)`/`Err(_)` to discard it and let + /// the pool try another (opening a fresh one if needed). Composes after the idle-ping + /// installed by + /// [`test_before_acquire_if_idle_for`](Self::test_before_acquire_if_idle_for) — that runs + /// first, then this callback. + /// + /// Applies only to pools built through [`Database::connect`], not to pools adopted via + /// `SqlxSqliteConnector::from_sqlx_sqlite_pool`. + pub fn map_sqlx_sqlite_before_acquire(&mut self, f: F) -> &mut Self + where + F: for<'c> Fn( + &'c mut sqlx::sqlite::SqliteConnection, + sqlx::pool::PoolConnectionMetadata, + ) + -> futures_util::future::BoxFuture<'c, Result> + + Send + + Sync + + 'static, + { + self.sqlite_before_acquire_fn = Some(Arc::new(f)); + self + } } diff --git a/src/driver/sqlx_common.rs b/src/driver/sqlx_common.rs index 6090dfd644..ab6d5330b3 100644 --- a/src/driver/sqlx_common.rs +++ b/src/driver/sqlx_common.rs @@ -1,4 +1,21 @@ use crate::{ConnAcquireErr, ConnectOptions, DbErr, RuntimeErr}; +use std::{sync::Arc, time::Duration}; + +/// Callback stored for a `before_acquire` hook on [`ConnectOptions`]. +/// +/// This mirrors the signature accepted by SQLx's +/// [`PoolOptions::before_acquire`][sqlx::pool::PoolOptions::before_acquire] for the given +/// database backend `DB`. Note the fully-qualified [`futures_util::future::BoxFuture`]: the +/// crate-local `BoxFuture` alias collapses to `T` under the `sync` feature, which would not +/// match SQLx's signature. +pub(crate) type BeforeAcquireFn = Arc< + dyn for<'c> Fn( + &'c mut ::Connection, + sqlx::pool::PoolConnectionMetadata, + ) -> futures_util::future::BoxFuture<'c, Result> + + Send + + Sync, +>; /// Converts an [sqlx::error] execution error to a [DbErr] pub fn sqlx_error_to_exec_err(err: sqlx::Error) -> DbErr { @@ -63,4 +80,104 @@ impl ConnectOptions { opt = opt.test_before_acquire(self.test_before_acquire); opt } + + /// Install the composed `before_acquire` hook onto a [`sqlx::pool::PoolOptions`]. + /// + /// SQLx exposes a single `before_acquire` slot whose setter *replaces* rather than + /// composes, and offers no getter to read it back. To let the idle-ping shorthand + /// ([`ConnectOptions::test_before_acquire_if_idle_for`]) coexist with a user-provided + /// per-backend callback ([`ConnectOptions::map_sqlx_postgres_before_acquire`] and + /// friends), this composes both into one closure: the idle-ping runs first, then the + /// user callback. When a ping threshold is set, `test_before_acquire` is forced off so + /// the connection is pinged only past the threshold rather than on every acquire. + /// + /// Returns `opt` untouched when neither option is configured, so callers that opt into + /// nothing get byte-for-byte the previous behavior. + pub(crate) fn apply_before_acquire( + mut opt: sqlx::pool::PoolOptions, + ping_after_idle: Option, + user_cb: Option>, + ) -> sqlx::pool::PoolOptions + where + DB: sqlx::Database, + { + use sqlx::Connection; + + if ping_after_idle.is_none() && user_cb.is_none() { + return opt; + } + if ping_after_idle.is_some() { + opt = opt.test_before_acquire(false); + } + opt.before_acquire(move |conn, meta| { + let user_cb = user_cb.clone(); + Box::pin(async move { + if let Some(threshold) = ping_after_idle { + // `idle_for` is `Copy`; read it before `meta` is moved into the user callback. + // `>=` matches the "idle for at least `threshold`" contract documented on + // `ConnectOptions::test_before_acquire_if_idle_for`. + if meta.idle_for >= threshold { + conn.ping().await?; + } + } + match user_cb { + Some(user_cb) => user_cb(conn, meta).await, + None => Ok(true), + } + }) + }) + } +} + +#[cfg(all(test, feature = "sqlx-postgres"))] +mod tests { + use crate::ConnectOptions; + use sqlx::Connection; + use std::time::Duration; + + #[test] + fn idle_shorthand_disables_test_before_acquire() { + let mut opt = ConnectOptions::new("postgres://localhost/db"); + assert!(opt.get_test_before_acquire()); + assert_eq!(opt.get_test_before_acquire_if_idle_for(), None); + + opt.test_before_acquire_if_idle_for(Duration::from_secs(30)); + assert!(!opt.get_test_before_acquire()); + assert_eq!( + opt.get_test_before_acquire_if_idle_for(), + Some(Duration::from_secs(30)) + ); + } + + #[test] + fn compose_shorthand_and_user_callback() { + let mut opt = ConnectOptions::new("postgres://localhost/db"); + opt.test_before_acquire_if_idle_for(Duration::from_secs(30)) + .map_sqlx_postgres_before_acquire(|conn, _meta| { + Box::pin(async move { + conn.ping().await?; + Ok(true) + }) + }); + + // Composing both into SQLx's single `before_acquire` slot type-checks and returns a + // usable `PoolOptions`. Behavioral ping timing requires a live pool, covered elsewhere. + let pool_opts = ConnectOptions::apply_before_acquire::( + sqlx::pool::PoolOptions::new(), + opt.get_test_before_acquire_if_idle_for(), + opt.pg_before_acquire_fn.clone(), + ); + let _ = pool_opts; + } + + #[test] + fn apply_before_acquire_noop_when_unset() { + // With neither option set, the helper must return the options untouched. + let opts = ConnectOptions::apply_before_acquire::( + sqlx::pool::PoolOptions::new(), + None, + None, + ); + let _ = opts; + } } diff --git a/src/driver/sqlx_mysql.rs b/src/driver/sqlx_mysql.rs index 4145bb6a90..25b3603e0a 100644 --- a/src/driver/sqlx_mysql.rs +++ b/src/driver/sqlx_mysql.rs @@ -89,7 +89,14 @@ impl SqlxMySqlConnector { let after_connect = options.after_connect.clone(); let connect_lazy = options.connect_lazy; let mysql_pool_opts_fn = options.mysql_pool_opts_fn.clone(); + let mysql_before_acquire = options.mysql_before_acquire_fn.clone(); + let ping_after_idle = options.test_before_acquire_if_idle_for; let mut pool_options = options.sqlx_pool_options(); + pool_options = crate::ConnectOptions::apply_before_acquire::( + pool_options, + ping_after_idle, + mysql_before_acquire, + ); if let Some(f) = &mysql_pool_opts_fn { pool_options = f(pool_options); } diff --git a/src/driver/sqlx_postgres.rs b/src/driver/sqlx_postgres.rs index 65b0c2e0a3..938a59edeb 100644 --- a/src/driver/sqlx_postgres.rs +++ b/src/driver/sqlx_postgres.rs @@ -117,7 +117,14 @@ impl SqlxPostgresConnector { let lazy = options.connect_lazy; let after_connect = options.after_connect.clone(); let pg_pool_opts_fn = options.pg_pool_opts_fn.clone(); + let pg_before_acquire = options.pg_before_acquire_fn.clone(); + let ping_after_idle = options.test_before_acquire_if_idle_for; let mut pool_options = options.sqlx_pool_options(); + pool_options = crate::ConnectOptions::apply_before_acquire::( + pool_options, + ping_after_idle, + pg_before_acquire, + ); if let Some(sql) = set_search_path_sql { pool_options = pool_options.after_connect(move |conn, _| { diff --git a/src/driver/sqlx_sqlite.rs b/src/driver/sqlx_sqlite.rs index 1a4cb540b8..0df465c977 100644 --- a/src/driver/sqlx_sqlite.rs +++ b/src/driver/sqlx_sqlite.rs @@ -99,7 +99,14 @@ impl SqlxSqliteConnector { let after_conn = options.after_connect.clone(); let connect_lazy = options.connect_lazy; let sqlite_pool_opts_fn = options.sqlite_pool_opts_fn.clone(); + let sqlite_before_acquire = options.sqlite_before_acquire_fn.clone(); + let ping_after_idle = options.test_before_acquire_if_idle_for; let mut pool_options = options.sqlx_pool_options(); + pool_options = crate::ConnectOptions::apply_before_acquire::( + pool_options, + ping_after_idle, + sqlite_before_acquire, + ); if let Some(f) = &sqlite_pool_opts_fn { pool_options = f(pool_options);