Split AsyncConnection trait

Author: stormshield-kg

src/lib.rs

@@ -125,12 +125,8 @@ pub trait SimpleAsyncConnection {
     fn batch_execute(&mut self, query: &str) -> impl Future<Output = QueryResult<()>> + Send;
 }
 
-/// An async connection to a database
-///
-/// This trait represents a n async database connection. It can be used to query the database through
-/// the query dsl provided by diesel, custom extensions or raw sql queries. It essentially mirrors
-/// the sync diesel [`Connection`](diesel::connection::Connection) implementation
-pub trait AsyncConnection: SimpleAsyncConnection + Sized + Send {
+/// Core trait for an async database connection
+pub trait AsyncConnectionCore: SimpleAsyncConnection + Send {
     /// The future returned by `AsyncConnection::execute`
     type ExecuteFuture<'conn, 'query>: Future<Output = QueryResult<usize>> + Send;
     /// The future returned by `AsyncConnection::load`
@@ -143,6 +139,37 @@ pub trait AsyncConnection: SimpleAsyncConnection + Sized + Send {
     /// The backend this type connects to
     type Backend: Backend;
 
+    #[doc(hidden)]
+    fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
+    where
+        T: AsQuery + 'query,
+        T::Query: QueryFragment<Self::Backend> + QueryId + 'query;
+
+    #[doc(hidden)]
+    fn execute_returning_count<'conn, 'query, T>(
+        &'conn mut self,
+        source: T,
+    ) -> Self::ExecuteFuture<'conn, 'query>
+    where
+        T: QueryFragment<Self::Backend> + QueryId + 'query;
+
+    // These functions allow the associated types (`ExecuteFuture`, `LoadFuture`, etc.) to
+    // compile without a `where Self: '_` clause. This is needed the because bound causes
+    // lifetime issues when using `transaction()` with generic `AsyncConnection`s.
+    //
+    // See: https://github.com/rust-lang/rust/issues/87479
+    #[doc(hidden)]
+    fn _silence_lint_on_execute_future(_: Self::ExecuteFuture<'_, '_>) {}
+    #[doc(hidden)]
+    fn _silence_lint_on_load_future(_: Self::LoadFuture<'_, '_>) {}
+}
+
+/// An async connection to a database
+///
+/// This trait represents an async database connection. It can be used to query the database through
+/// the query dsl provided by diesel, custom extensions or raw sql queries. It essentially mirrors
+/// the sync diesel [`Connection`](diesel::connection::Connection) implementation
+pub trait AsyncConnection: AsyncConnectionCore + Sized {
     #[doc(hidden)]
     type TransactionManager: TransactionManager<Self>;
 
@@ -336,35 +363,11 @@ pub trait AsyncConnection: SimpleAsyncConnection + Sized + Send {
         })
     }
 
-    #[doc(hidden)]
-    fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
-    where
-        T: AsQuery + 'query,
-        T::Query: QueryFragment<Self::Backend> + QueryId + 'query;
-
-    #[doc(hidden)]
-    fn execute_returning_count<'conn, 'query, T>(
-        &'conn mut self,
-        source: T,
-    ) -> Self::ExecuteFuture<'conn, 'query>
-    where
-        T: QueryFragment<Self::Backend> + QueryId + 'query;
-
     #[doc(hidden)]
     fn transaction_state(
         &mut self,
     ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData;
 
-    // These functions allow the associated types (`ExecuteFuture`, `LoadFuture`, etc.) to
-    // compile without a `where Self: '_` clause. This is needed the because bound causes
-    // lifetime issues when using `transaction()` with generic `AsyncConnection`s.
-    //
-    // See: https://github.com/rust-lang/rust/issues/87479
-    #[doc(hidden)]
-    fn _silence_lint_on_execute_future(_: Self::ExecuteFuture<'_, '_>) {}
-    #[doc(hidden)]
-    fn _silence_lint_on_load_future(_: Self::LoadFuture<'_, '_>) {}
-
     #[doc(hidden)]
     fn instrumentation(&mut self) -> &mut dyn Instrumentation;
 

src/mysql/mod.rs

@@ -1,5 +1,5 @@
 use crate::stmt_cache::{CallbackHelper, QueryFragmentHelper};
-use crate::{AnsiTransactionManager, AsyncConnection, SimpleAsyncConnection};
+use crate::{AnsiTransactionManager, AsyncConnection, AsyncConnectionCore, SimpleAsyncConnection};
 use diesel::connection::statement_cache::{
     MaybeCached, QueryFragmentForCachedStatement, StatementCache,
 };
@@ -64,30 +64,13 @@ const CONNECTION_SETUP_QUERIES: &[&str] = &[
     "SET character_set_results = 'utf8mb4'",
 ];
 
-impl AsyncConnection for AsyncMysqlConnection {
+impl AsyncConnectionCore for AsyncMysqlConnection {
     type ExecuteFuture<'conn, 'query> = BoxFuture<'conn, QueryResult<usize>>;
     type LoadFuture<'conn, 'query> = BoxFuture<'conn, QueryResult<Self::Stream<'conn, 'query>>>;
     type Stream<'conn, 'query> = BoxStream<'conn, QueryResult<Self::Row<'conn, 'query>>>;
     type Row<'conn, 'query> = MysqlRow;
     type Backend = Mysql;
 
-    type TransactionManager = AnsiTransactionManager;
-
-    async fn establish(database_url: &str) -> diesel::ConnectionResult<Self> {
-        let mut instrumentation = DynInstrumentation::default_instrumentation();
-        instrumentation.on_connection_event(InstrumentationEvent::start_establish_connection(
-            database_url,
-        ));
-        let r = Self::establish_connection_inner(database_url).await;
-        instrumentation.on_connection_event(InstrumentationEvent::finish_establish_connection(
-            database_url,
-            r.as_ref().err(),
-        ));
-        let mut conn = r?;
-        conn.instrumentation = instrumentation;
-        Ok(conn)
-    }
-
     fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
     where
         T: diesel::query_builder::AsQuery,
@@ -173,6 +156,25 @@ impl AsyncConnection for AsyncMysqlConnection {
                 .map_err(|e| diesel::result::Error::DeserializationError(Box::new(e)))
         })
     }
+}
+
+impl AsyncConnection for AsyncMysqlConnection {
+    type TransactionManager = AnsiTransactionManager;
+
+    async fn establish(database_url: &str) -> diesel::ConnectionResult<Self> {
+        let mut instrumentation = DynInstrumentation::default_instrumentation();
+        instrumentation.on_connection_event(InstrumentationEvent::start_establish_connection(
+            database_url,
+        ));
+        let r = Self::establish_connection_inner(database_url).await;
+        instrumentation.on_connection_event(InstrumentationEvent::finish_establish_connection(
+            database_url,
+            r.as_ref().err(),
+        ));
+        let mut conn = r?;
+        conn.instrumentation = instrumentation;
+        Ok(conn)
+    }
 
     fn transaction_state(&mut self) -> &mut AnsiTransactionManager {
         &mut self.transaction_manager

src/pg/mod.rs

@@ -8,7 +8,7 @@ use self::error_helper::ErrorHelper;
 use self::row::PgRow;
 use self::serialize::ToSqlHelper;
 use crate::stmt_cache::{CallbackHelper, QueryFragmentHelper};
-use crate::{AnsiTransactionManager, AsyncConnection, SimpleAsyncConnection};
+use crate::{AnsiTransactionManager, AsyncConnection, AsyncConnectionCore, SimpleAsyncConnection};
 use diesel::connection::statement_cache::{
     PrepareForCache, QueryFragmentForCachedStatement, StatementCache,
 };
@@ -160,12 +160,37 @@ impl SimpleAsyncConnection for AsyncPgConnection {
     }
 }
 
-impl AsyncConnection for AsyncPgConnection {
+impl AsyncConnectionCore for AsyncPgConnection {
     type LoadFuture<'conn, 'query> = BoxFuture<'query, QueryResult<Self::Stream<'conn, 'query>>>;
     type ExecuteFuture<'conn, 'query> = BoxFuture<'query, QueryResult<usize>>;
     type Stream<'conn, 'query> = BoxStream<'static, QueryResult<PgRow>>;
     type Row<'conn, 'query> = PgRow;
     type Backend = diesel::pg::Pg;
+
+    fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
+    where
+        T: AsQuery + 'query,
+        T::Query: QueryFragment<Self::Backend> + QueryId + 'query,
+    {
+        let query = source.as_query();
+        let load_future = self.with_prepared_statement(query, load_prepared);
+
+        self.run_with_connection_future(load_future)
+    }
+
+    fn execute_returning_count<'conn, 'query, T>(
+        &'conn mut self,
+        source: T,
+    ) -> Self::ExecuteFuture<'conn, 'query>
+    where
+        T: QueryFragment<Self::Backend> + QueryId + 'query,
+    {
+        let execute = self.with_prepared_statement(source, execute_prepared);
+        self.run_with_connection_future(execute)
+    }
+}
+
+impl AsyncConnection for AsyncPgConnection {
     type TransactionManager = AnsiTransactionManager;
 
     async fn establish(database_url: &str) -> ConnectionResult<Self> {
@@ -198,28 +223,6 @@ impl AsyncConnection for AsyncPgConnection {
         r
     }
 
-    fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
-    where
-        T: AsQuery + 'query,
-        T::Query: QueryFragment<Self::Backend> + QueryId + 'query,
-    {
-        let query = source.as_query();
-        let load_future = self.with_prepared_statement(query, load_prepared);
-
-        self.run_with_connection_future(load_future)
-    }
-
-    fn execute_returning_count<'conn, 'query, T>(
-        &'conn mut self,
-        source: T,
-    ) -> Self::ExecuteFuture<'conn, 'query>
-    where
-        T: QueryFragment<Self::Backend> + QueryId + 'query,
-    {
-        let execute = self.with_prepared_statement(source, execute_prepared);
-        self.run_with_connection_future(execute)
-    }
-
     fn transaction_state(&mut self) -> &mut AnsiTransactionManager {
         // there should be no other pending future when this is called
         // that means there is only one instance of this arc and
@@ -467,7 +470,7 @@ impl AsyncPgConnection {
     }
 
     fn with_prepared_statement<'a, T, F, R>(
-        &mut self,
+        &self,
         query: T,
         callback: fn(Arc<tokio_postgres::Client>, Statement, Vec<ToSqlHelper>) -> F,
     ) -> BoxFuture<'a, QueryResult<R>>
@@ -502,7 +505,7 @@ impl AsyncPgConnection {
     }
 
     fn with_prepared_statement_after_sql_built<'a, F, R>(
-        &mut self,
+        &self,
         callback: fn(Arc<tokio_postgres::Client>, Statement, Vec<ToSqlHelper>) -> F,
         is_safe_to_cache_prepared: QueryResult<bool>,
         query_id: Option<std::any::TypeId>,

src/pooled_connection/mod.rs

@@ -5,7 +5,7 @@
 //! * [deadpool](self::deadpool)
 //! * [bb8](self::bb8)
 //! * [mobc](self::mobc)
-use crate::{AsyncConnection, SimpleAsyncConnection};
+use crate::{AsyncConnection, AsyncConnectionCore, SimpleAsyncConnection};
 use crate::{TransactionManager, UpdateAndFetchResults};
 use diesel::associations::HasTable;
 use diesel::connection::{CacheSize, Instrumentation};
@@ -176,27 +176,18 @@ where
     }
 }
 
-impl<C> AsyncConnection for C
+impl<C> AsyncConnectionCore for C
 where
     C: DerefMut + Send,
-    C::Target: AsyncConnection,
+    C::Target: AsyncConnectionCore,
 {
     type ExecuteFuture<'conn, 'query> =
-        <C::Target as AsyncConnection>::ExecuteFuture<'conn, 'query>;
-    type LoadFuture<'conn, 'query> = <C::Target as AsyncConnection>::LoadFuture<'conn, 'query>;
-    type Stream<'conn, 'query> = <C::Target as AsyncConnection>::Stream<'conn, 'query>;
-    type Row<'conn, 'query> = <C::Target as AsyncConnection>::Row<'conn, 'query>;
+        <C::Target as AsyncConnectionCore>::ExecuteFuture<'conn, 'query>;
+    type LoadFuture<'conn, 'query> = <C::Target as AsyncConnectionCore>::LoadFuture<'conn, 'query>;
+    type Stream<'conn, 'query> = <C::Target as AsyncConnectionCore>::Stream<'conn, 'query>;
+    type Row<'conn, 'query> = <C::Target as AsyncConnectionCore>::Row<'conn, 'query>;
 
-    type Backend = <C::Target as AsyncConnection>::Backend;
-
-    type TransactionManager =
-        PoolTransactionManager<<C::Target as AsyncConnection>::TransactionManager>;
-
-    async fn establish(_database_url: &str) -> diesel::ConnectionResult<Self> {
-        Err(diesel::result::ConnectionError::BadConnection(
-            String::from("Cannot directly establish a pooled connection"),
-        ))
-    }
+    type Backend = <C::Target as AsyncConnectionCore>::Backend;
 
     fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
     where
@@ -221,6 +212,21 @@ where
         let conn = self.deref_mut();
         conn.execute_returning_count(source)
     }
+}
+
+impl<C> AsyncConnection for C
+where
+    C: DerefMut + Send,
+    C::Target: AsyncConnection,
+{
+    type TransactionManager =
+        PoolTransactionManager<<C::Target as AsyncConnection>::TransactionManager>;
+
+    async fn establish(_database_url: &str) -> diesel::ConnectionResult<Self> {
+        Err(diesel::result::ConnectionError::BadConnection(
+            String::from("Cannot directly establish a pooled connection"),
+        ))
+    }
 
     fn transaction_state(
         &mut self,