Merge pull request #248 from stormshield-kg/pg_ref_run_query

Improve pipelining for Postgres

Author: weiznich

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