Merge pull request #101 from weiznich/feature/more_pool_configs

weiznich · web-flow · commit fead86962d1d · 2023-09-01T08:05:06.000Z
Add a `RecyclingMethod` configuration to the connection pool
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,8 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 ## Unreleased
 
 * Add a `AsyncConnectionWrapper` type to turn a `diesel_async::AsyncConnection` into a `diesel::Connection`. This might be used to execute migrations via `diesel_migrations`. 
+* Add some connection pool configurations to specify how connections
+in the pool should be checked if they are still valid
 
 ## [0.3.2] - 2023-07-24
 
@@ -56,3 +58,5 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 [0.2.1]: https://github.com/weiznich/diesel_async/compare/v0.2.0...v0.2.1
 [0.2.2]: https://github.com/weiznich/diesel_async/compare/v0.2.1...v0.2.2
 [0.3.0]: https://github.com/weiznich/diesel_async/compare/v0.2.0...v0.3.0
+[0.3.1]: https://github.com/weiznich/diesel_async/compare/v0.3.0...v0.3.1
+[0.3.2]: https://github.com/weiznich/diesel_async/compare/v0.3.1...v0.3.2
diff --git a/examples/postgres/pooled-with-rustls/src/main.rs b/examples/postgres/pooled-with-rustls/src/main.rs
@@ -1,6 +1,7 @@
 use diesel::{ConnectionError, ConnectionResult};
 use diesel_async::pooled_connection::bb8::Pool;
 use diesel_async::pooled_connection::AsyncDieselConnectionManager;
+use diesel_async::pooled_connection::ManagerConfig;
 use diesel_async::AsyncPgConnection;
 use futures_util::future::BoxFuture;
 use futures_util::FutureExt;
@@ -10,12 +11,14 @@ use std::time::Duration;
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let db_url = std::env::var("DATABASE_URL").expect("Env var `DATABASE_URL` not set");
 
+    let config = ManagerConfig {
+        custom_setup: Box::new(establish_connection),
+        ..ManagerConfig::default()
+    };
+
     // First we have to construct a connection manager with our custom `establish_connection`
     // function
-    let mgr = AsyncDieselConnectionManager::<AsyncPgConnection>::new_with_setup(
-        db_url,
-        establish_connection,
-    );
+    let mgr = AsyncDieselConnectionManager::<AsyncPgConnection>::new_with_config(db_url, config);
     // From that connection we can then create a pool, here given with some example settings.
     //
     // This creates a TLS configuration that's equivalent to `libpq'` `sslmode=verify-full`, which
diff --git a/src/async_connection_wrapper.rs b/src/async_connection_wrapper.rs
@@ -258,16 +258,22 @@ mod implementation {
         B: BlockOn,
         Self: diesel::Connection,
         C: crate::AsyncConnection<Backend = <Self as diesel::Connection>::Backend>
-            + crate::pooled_connection::PoolableConnection,
+            + crate::pooled_connection::PoolableConnection
+            + 'static,
+        diesel::dsl::BareSelect<diesel::dsl::AsExprOf<i32, diesel::sql_types::Integer>>:
+            crate::methods::ExecuteDsl<C>,
+        diesel::query_builder::SqlQuery: crate::methods::ExecuteDsl<C>,
     {
         fn ping(&mut self) -> diesel::QueryResult<()> {
-            diesel::Connection::execute_returning_count(self, &C::make_ping_query()).map(|_| ())
+            let fut = crate::pooled_connection::PoolableConnection::ping(
+                &mut self.inner,
+                &crate::pooled_connection::RecyclingMethod::Verified,
+            );
+            self.runtime.block_on(fut)
         }
 
         fn is_broken(&mut self) -> bool {
-            <C::TransactionManager as crate::TransactionManager<_>>::is_broken_transaction_manager(
-                &mut self.inner,
-            )
+            crate::pooled_connection::PoolableConnection::is_broken(&mut self.inner)
         }
     }
 
diff --git a/src/mysql/mod.rs b/src/mysql/mod.rs
@@ -307,13 +307,7 @@ impl AsyncMysqlConnection {
     feature = "mobc",
     feature = "r2d2"
 ))]
-impl crate::pooled_connection::PoolableConnection for AsyncMysqlConnection {
-    type PingQuery = crate::pooled_connection::CheckConnectionQuery;
-
-    fn make_ping_query() -> Self::PingQuery {
-        crate::pooled_connection::CheckConnectionQuery
-    }
-}
+impl crate::pooled_connection::PoolableConnection for AsyncMysqlConnection {}
 
 #[cfg(test)]
 mod tests {
diff --git a/src/pg/mod.rs b/src/pg/mod.rs
@@ -477,10 +477,10 @@ async fn lookup_type(
     feature = "r2d2"
 ))]
 impl crate::pooled_connection::PoolableConnection for AsyncPgConnection {
-    type PingQuery = crate::pooled_connection::CheckConnectionQuery;
+    fn is_broken(&mut self) -> bool {
+        use crate::TransactionManager;
 
-    fn make_ping_query() -> Self::PingQuery {
-        crate::pooled_connection::CheckConnectionQuery
+        Self::TransactionManager::is_broken_transaction_manager(self) || self.conn.is_closed()
     }
 }
 
diff --git a/src/pooled_connection/bb8.rs b/src/pooled_connection/bb8.rs
@@ -43,6 +43,7 @@
 
 use super::{AsyncDieselConnectionManager, PoolError, PoolableConnection};
 use bb8::ManageConnection;
+use diesel::query_builder::QueryFragment;
 
 /// Type alias for using [`bb8::Pool`] with [`diesel-async`]
 pub type Pool<C> = bb8::Pool<AsyncDieselConnectionManager<C>>;
@@ -55,19 +56,24 @@ pub type RunError = bb8::RunError<super::PoolError>;
 impl<C> ManageConnection for AsyncDieselConnectionManager<C>
 where
     C: PoolableConnection + 'static,
+    diesel::dsl::BareSelect<diesel::dsl::AsExprOf<i32, diesel::sql_types::Integer>>:
+        crate::methods::ExecuteDsl<C>,
+    diesel::query_builder::SqlQuery: QueryFragment<C::Backend>,
 {
     type Connection = C;
 
     type Error = PoolError;
 
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
-        (self.setup)(&self.connection_url)
+        (self.manager_config.custom_setup)(&self.connection_url)
             .await
             .map_err(PoolError::ConnectionError)
     }
 
     async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
-        conn.ping().await.map_err(PoolError::QueryError)
+        conn.ping(&self.manager_config.recycling_method)
+            .await
+            .map_err(PoolError::QueryError)
     }
 
     fn has_broken(&self, conn: &mut Self::Connection) -> bool {
diff --git a/src/pooled_connection/deadpool.rs b/src/pooled_connection/deadpool.rs
@@ -41,6 +41,7 @@
 //! ```
 use super::{AsyncDieselConnectionManager, PoolableConnection};
 use deadpool::managed::Manager;
+use diesel::query_builder::QueryFragment;
 
 /// Type alias for using [`deadpool::managed::Pool`] with [`diesel-async`]
 pub type Pool<C> = deadpool::managed::Pool<AsyncDieselConnectionManager<C>>;
@@ -63,13 +64,16 @@ pub type HookErrorCause = deadpool::managed::HookErrorCause<super::PoolError>;
 impl<C> Manager for AsyncDieselConnectionManager<C>
 where
     C: PoolableConnection + Send + 'static,
+    diesel::dsl::BareSelect<diesel::dsl::AsExprOf<i32, diesel::sql_types::Integer>>:
+        crate::methods::ExecuteDsl<C>,
+    diesel::query_builder::SqlQuery: QueryFragment<C::Backend>,
 {
     type Type = C;
 
     type Error = super::PoolError;
 
     async fn create(&self) -> Result<Self::Type, Self::Error> {
-        (self.setup)(&self.connection_url)
+        (self.manager_config.custom_setup)(&self.connection_url)
             .await
             .map_err(super::PoolError::ConnectionError)
     }
@@ -80,7 +84,9 @@ where
                 "Broken connection",
             ));
         }
-        obj.ping().await.map_err(super::PoolError::QueryError)?;
+        obj.ping(&self.manager_config.recycling_method)
+            .await
+            .map_err(super::PoolError::QueryError)?;
         Ok(())
     }
 }
diff --git a/src/pooled_connection/mobc.rs b/src/pooled_connection/mobc.rs
@@ -40,6 +40,7 @@
 //! # }
 //! ```
 use super::{AsyncDieselConnectionManager, PoolError, PoolableConnection};
+use diesel::query_builder::QueryFragment;
 use mobc::Manager;
 
 /// Type alias for using [`mobc::Pool`] with [`diesel-async`]
@@ -52,19 +53,24 @@ pub type Builder<C> = mobc::Builder<AsyncDieselConnectionManager<C>>;
 impl<C> Manager for AsyncDieselConnectionManager<C>
 where
     C: PoolableConnection + 'static,
+    diesel::dsl::BareSelect<diesel::dsl::AsExprOf<i32, diesel::sql_types::Integer>>:
+        crate::methods::ExecuteDsl<C>,
+    diesel::query_builder::SqlQuery: QueryFragment<C::Backend>,
 {
     type Connection = C;
 
     type Error = PoolError;
 
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
-        (self.setup)(&self.connection_url)
+        (self.manager_config.custom_setup)(&self.connection_url)
             .await
             .map_err(PoolError::ConnectionError)
     }
 
     async fn check(&self, mut conn: Self::Connection) -> Result<Self::Connection, Self::Error> {
-        conn.ping().await.map_err(PoolError::QueryError)?;
+        conn.ping(&self.manager_config.recycling_method)
+            .await
+            .map_err(PoolError::QueryError)?;
         Ok(conn)
     }
 }
diff --git a/src/pooled_connection/mod.rs b/src/pooled_connection/mod.rs
@@ -8,9 +8,9 @@
 use crate::{AsyncConnection, SimpleAsyncConnection};
 use crate::{TransactionManager, UpdateAndFetchResults};
 use diesel::associations::HasTable;
-use diesel::query_builder::{QueryFragment, QueryId};
 use diesel::QueryResult;
 use futures_util::{future, FutureExt};
+use std::borrow::Cow;
 use std::fmt;
 use std::ops::DerefMut;
 
@@ -42,18 +42,83 @@ impl fmt::Display for PoolError {
 
 impl std::error::Error for PoolError {}
 
-type SetupCallback<C> =
+/// Type of the custom setup closure passed to [`ManagerConfig::custom_setup`]
+pub type SetupCallback<C> =
     Box<dyn Fn(&str) -> future::BoxFuture<diesel::ConnectionResult<C>> + Send + Sync>;
 
+/// Type of the recycle check callback for the [`RecyclingMethod::CustomFunction`] variant
+pub type RecycleCheckCallback<C> =
+    dyn Fn(&mut C) -> future::BoxFuture<QueryResult<()>> + Send + Sync;
+
+/// Possible methods of how a connection is recycled.
+#[derive(Default)]
+pub enum RecyclingMethod<C> {
+    /// Only check for open transactions when recycling existing connections
+    /// Unless you have special needs this is a safe choice.
+    ///
+    /// If the database connection is closed you will recieve an error on the first place
+    /// you actually try to use the connection
+    Fast,
+    /// In addition to checking for open transactions a test query is executed
+    ///
+    /// This is slower, but guarantees that the database connection is ready to be used.
+    #[default]
+    Verified,
+    /// Like `Verified` but with a custom query
+    CustomQuery(Cow<'static, str>),
+    /// Like `Verified` but with a custom callback that allows to perform more checks
+    ///
+    /// The connection is only recycled if the callback returns `Ok(())`
+    CustomFunction(Box<RecycleCheckCallback<C>>),
+}
+
+impl<C: fmt::Debug> fmt::Debug for RecyclingMethod<C> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Fast => write!(f, "Fast"),
+            Self::Verified => write!(f, "Verified"),
+            Self::CustomQuery(arg0) => f.debug_tuple("CustomQuery").field(arg0).finish(),
+            Self::CustomFunction(_) => f.debug_tuple("CustomFunction").finish(),
+        }
+    }
+}
+
+/// Configuration object for a Manager.
+///
+/// This currently only makes it possible to specify which [`RecyclingMethod`]
+/// should be used when retrieving existing objects from the [`Pool`].
+pub struct ManagerConfig<C> {
+    /// Method of how a connection is recycled. See [RecyclingMethod].
+    pub recycling_method: RecyclingMethod<C>,
+    /// Construct a new connection manger
+    /// with a custom setup procedure
+    ///
+    /// This can be used to for example establish a SSL secured
+    /// postgres connection
+    pub custom_setup: SetupCallback<C>,
+}
+
+impl<C> Default for ManagerConfig<C>
+where
+    C: AsyncConnection + 'static,
+{
+    fn default() -> Self {
+        Self {
+            recycling_method: Default::default(),
+            custom_setup: Box::new(|url| C::establish(url).boxed()),
+        }
+    }
+}
+
 /// An connection manager for use with diesel-async.
 ///
 /// See the concrete pool implementations for examples:
 /// * [deadpool](self::deadpool)
 /// * [bb8](self::bb8)
 /// * [mobc](self::mobc)
 pub struct AsyncDieselConnectionManager<C> {
-    setup: SetupCallback<C>,
     connection_url: String,
+    manager_config: ManagerConfig<C>,
 }
 
 impl<C> fmt::Debug for AsyncDieselConnectionManager<C> {
@@ -66,28 +131,31 @@ impl<C> fmt::Debug for AsyncDieselConnectionManager<C> {
     }
 }
 
-impl<C> AsyncDieselConnectionManager<C> {
+impl<C> AsyncDieselConnectionManager<C>
+where
+    C: AsyncConnection + 'static,
+{
     /// Returns a new connection manager,
     /// which establishes connections to the given database URL.
+    #[must_use]
     pub fn new(connection_url: impl Into<String>) -> Self
     where
         C: AsyncConnection + 'static,
     {
-        Self::new_with_setup(connection_url, |url| C::establish(url).boxed())
+        Self::new_with_config(connection_url, Default::default())
     }
 
-    /// Construct a new connection manger
-    /// with a custom setup procedure
-    ///
-    /// This can be used to for example establish a SSL secured
-    /// postgres connection
-    pub fn new_with_setup(
+    /// Returns a new connection manager,
+    /// which establishes connections with the given database URL
+    /// and that uses the specified configuration
+    #[must_use]
+    pub fn new_with_config(
         connection_url: impl Into<String>,
-        setup: impl Fn(&str) -> future::BoxFuture<diesel::ConnectionResult<C>> + Send + Sync + 'static,
+        manager_config: ManagerConfig<C>,
     ) -> Self {
         Self {
-            setup: Box::new(setup),
             connection_url: connection_url.into(),
+            manager_config,
         }
     }
 }
@@ -218,9 +286,8 @@ where
     }
 }
 
-#[doc(hidden)]
 #[derive(diesel::query_builder::QueryId)]
-pub struct CheckConnectionQuery;
+struct CheckConnectionQuery;
 
 impl<DB> diesel::query_builder::QueryFragment<DB> for CheckConnectionQuery
 where
@@ -244,19 +311,34 @@ impl<C> diesel::query_dsl::RunQueryDsl<C> for CheckConnectionQuery {}
 #[doc(hidden)]
 #[async_trait::async_trait]
 pub trait PoolableConnection: AsyncConnection {
-    type PingQuery: QueryFragment<Self::Backend> + QueryId + Send;
-
-    fn make_ping_query() -> Self::PingQuery;
-
     /// Check if a connection is still valid
     ///
-    /// The default implementation performs a `SELECT 1` query
-    async fn ping(&mut self) -> diesel::QueryResult<()>
+    /// The default implementation will perform a check based on the provided
+    /// recycling method variant
+    async fn ping(&mut self, config: &RecyclingMethod<Self>) -> diesel::QueryResult<()>
     where
         for<'a> Self: 'a,
+        diesel::dsl::BareSelect<diesel::dsl::AsExprOf<i32, diesel::sql_types::Integer>>:
+            crate::methods::ExecuteDsl<Self>,
+        diesel::query_builder::SqlQuery: crate::methods::ExecuteDsl<Self>,
     {
-        use crate::RunQueryDsl;
-        Self::make_ping_query().execute(self).await.map(|_| ())
+        use crate::run_query_dsl::RunQueryDsl;
+        use diesel::IntoSql;
+
+        match config {
+            RecyclingMethod::Fast => Ok(()),
+            RecyclingMethod::Verified => {
+                diesel::select(1_i32.into_sql::<diesel::sql_types::Integer>())
+                    .execute(self)
+                    .await
+                    .map(|_| ())
+            }
+            RecyclingMethod::CustomQuery(query) => diesel::sql_query(query.as_ref())
+                .execute(self)
+                .await
+                .map(|_| ()),
+            RecyclingMethod::CustomFunction(c) => c(self).await,
+        }
     }
 
     /// Checks if the connection is broken and should not be reused

Original file line number	Diff line number	Diff line change
`@@ -477,10 +477,10 @@ async fn lookup_type(`
`477`	`477`	`feature = "r2d2"`
`478`	`478`	`))]`
`479`	`479`	`impl crate::pooled_connection::PoolableConnection for AsyncPgConnection {`
`480`		`- type PingQuery = crate::pooled_connection::CheckConnectionQuery;`
	`480`	`+ fn is_broken(&mut self) -> bool {`
	`481`	`+ use crate::TransactionManager;`
`481`	`482`
`482`		`- fn make_ping_query() -> Self::PingQuery {`
`483`		`- crate::pooled_connection::CheckConnectionQuery`
	`483`	`+ Self::TransactionManager::is_broken_transaction_manager(self) \|\| self.conn.is_closed()`
`484`	`484`	`}`
`485`	`485`	`}`
`486`	`486`