Implement support for pipelining for the Postgres connection

Author: weiznich

src/lib.rs

@@ -3,7 +3,7 @@ use diesel::query_builder::{AsQuery, QueryFragment, QueryId};
 use diesel::row::Row;
 use diesel::{ConnectionResult, QueryResult};
 use futures::future::BoxFuture;
-use futures::Stream;
+use futures::{Future, Stream};
 #[cfg(feature = "mysql")]
 mod mysql;
 #[cfg(feature = "postgres")]
@@ -25,32 +25,37 @@ pub trait SimpleAsyncConnection {
     async fn batch_execute(&mut self, query: &str) -> QueryResult<()>;
 }
 
-pub trait AsyncConnectionGatWorkaround<'a, DB: Backend> {
-    type Stream: Stream<Item = QueryResult<Self::Row>> + Send + 'a;
-    type Row: Row<'a, DB> + 'a;
+pub trait AsyncConnectionGatWorkaround<'conn, 'query, DB: Backend> {
+    type ExecuteFuture: Future<Output = QueryResult<usize>> + Send;
+    type LoadFuture: Future<Output = QueryResult<Self::Stream>> + Send;
+    type Stream: Stream<Item = QueryResult<Self::Row>> + Send;
+    type Row: Row<'conn, DB>;
 }
 
 #[async_trait::async_trait]
 pub trait AsyncConnection: SimpleAsyncConnection + Sized + Send
 where
-    for<'a> Self: AsyncConnectionGatWorkaround<'a, Self::Backend>,
+    for<'a, 'b> Self: AsyncConnectionGatWorkaround<'a, 'b, Self::Backend>,
 {
     type Backend: Backend;
     type TransactionManager: TransactionManager<Self>;
 
     async fn establish(database_url: &str) -> ConnectionResult<Self>;
 
-    async fn load<'a, T>(
-        &'a mut self,
+    fn load<'conn, 'query, T>(
+        &'conn mut self,
         source: T,
-    ) -> QueryResult<<Self as AsyncConnectionGatWorkaround<'a, Self::Backend>>::Stream>
+    ) -> <Self as AsyncConnectionGatWorkaround<'conn, 'query, Self::Backend>>::LoadFuture
     where
-        T: AsQuery + Send,
-        T::Query: QueryFragment<Self::Backend> + QueryId + Send;
+        T: AsQuery + Send + 'query,
+        T::Query: QueryFragment<Self::Backend> + QueryId + Send + 'query;
 
-    async fn execute_returning_count<T>(&mut self, source: T) -> QueryResult<usize>
+    fn execute_returning_count<'conn, 'query, T>(
+        &'conn mut self,
+        source: T,
+    ) -> <Self as AsyncConnectionGatWorkaround<'conn, 'query, Self::Backend>>::ExecuteFuture
     where
-        T: QueryFragment<Self::Backend> + QueryId + Send;
+        T: QueryFragment<Self::Backend> + QueryId + Send + 'query;
 
     async fn transaction<F, R, E>(&mut self, callback: F) -> Result<R, E>
     where

src/mysql/mod.rs

@@ -1,5 +1,3 @@
-use std::pin::Pin;
-
 use crate::stmt_cache::{PrepareCallback, StmtCache};
 use crate::{
     AnsiTransactionManager, AsyncConnection, AsyncConnectionGatWorkaround, SimpleAsyncConnection,
@@ -9,7 +7,9 @@ use diesel::mysql::{Mysql, MysqlType};
 use diesel::query_builder::{bind_collector::RawBytesBindCollector, QueryFragment, QueryId};
 use diesel::result::{ConnectionError, ConnectionResult};
 use diesel::QueryResult;
-use futures::{Future, Stream, StreamExt, TryStreamExt};
+use futures::future::BoxFuture;
+use futures::stream::BoxStream;
+use futures::{Future, FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use mysql_async::prelude::Queryable;
 use mysql_async::{Opts, OptsBuilder, Statement};
 
@@ -35,8 +35,10 @@ impl SimpleAsyncConnection for AsyncMysqlConnection {
     }
 }
 
-impl<'a> AsyncConnectionGatWorkaround<'a, Mysql> for AsyncMysqlConnection {
-    type Stream = Pin<Box<dyn Stream<Item = QueryResult<Self::Row>> + Send + 'a>>;
+impl<'conn, 'query> AsyncConnectionGatWorkaround<'conn, 'query, Mysql> for AsyncMysqlConnection {
+    type ExecuteFuture = BoxFuture<'conn, QueryResult<usize>>;
+    type LoadFuture = BoxFuture<'conn, QueryResult<Self::Stream>>;
+    type Stream = BoxStream<'conn, QueryResult<Self::Row>>;
 
     type Row = MysqlRow;
 }
@@ -50,13 +52,15 @@ impl AsyncConnection for AsyncMysqlConnection {
     async fn establish(database_url: &str) -> diesel::ConnectionResult<Self> {
         let opts = Opts::from_url(database_url)
             .map_err(|e| diesel::result::ConnectionError::InvalidConnectionUrl(e.to_string()))?;
-        let builder = OptsBuilder::from_opts(opts).init(vec![
-            "SET sql_mode=(SELECT CONCAT(@@sql_mode, ',PIPES_AS_CONCAT'))",
-            "SET time_zone = '+00:00';",
-            "SET character_set_client = 'utf8mb4'",
-            "SET character_set_connection = 'utf8mb4'",
-            "SET character_set_results = 'utf8mb4'",
-        ]);
+        let builder = OptsBuilder::from_opts(opts)
+            .init(vec![
+                "SET sql_mode=(SELECT CONCAT(@@sql_mode, ',PIPES_AS_CONCAT'))",
+                "SET time_zone = '+00:00';",
+                "SET character_set_client = 'utf8mb4'",
+                "SET character_set_connection = 'utf8mb4'",
+                "SET character_set_results = 'utf8mb4'",
+            ])
+            .stmt_cache_size(0); // We have our own cache
 
         let conn = mysql_async::Conn::new(builder).await.map_err(ErrorHelper)?;
 
@@ -68,15 +72,16 @@ impl AsyncConnection for AsyncMysqlConnection {
         })
     }
 
-    async fn load<'a, T>(
-        &'a mut self,
+    fn load<'conn, 'query, T>(
+        &'conn mut self,
         source: T,
-    ) -> diesel::QueryResult<<Self as AsyncConnectionGatWorkaround<'a, Self::Backend>>::Stream>
+    ) -> <Self as AsyncConnectionGatWorkaround<'conn, 'query, Self::Backend>>::LoadFuture
     where
         T: diesel::query_builder::AsQuery + Send,
         T::Query: diesel::query_builder::QueryFragment<Self::Backend>
             + diesel::query_builder::QueryId
-            + Send,
+            + Send
+            + 'query,
     {
         self.with_prepared_statement(source.as_query(), |conn, stmt, binds| async move {
             let res = conn.exec_iter(&*stmt, binds).await.map_err(ErrorHelper)?;
@@ -95,20 +100,23 @@ impl AsyncConnection for AsyncMysqlConnection {
 
             Ok(stream)
         })
-        .await
+        .boxed()
     }
 
-    async fn execute_returning_count<T>(&mut self, source: T) -> diesel::QueryResult<usize>
+    fn execute_returning_count<'conn, 'query, T>(
+        &'conn mut self,
+        source: T,
+    ) -> <Self as AsyncConnectionGatWorkaround<'conn, 'query, Self::Backend>>::ExecuteFuture
     where
         T: diesel::query_builder::QueryFragment<Self::Backend>
             + diesel::query_builder::QueryId
-            + Send,
+            + Send
+            + 'query,
     {
         self.with_prepared_statement(source, |conn, stmt, binds| async move {
             conn.exec_drop(&*stmt, binds).await.map_err(ErrorHelper)?;
             Ok(conn.affected_rows() as usize)
         })
-        .await
     }
 
     fn transaction_state(&mut self) -> &mut AnsiTransactionManager {
@@ -117,14 +125,15 @@ impl AsyncConnection for AsyncMysqlConnection {
 }
 
 #[async_trait::async_trait]
-impl PrepareCallback<Statement, MysqlType> for mysql_async::Conn {
+impl PrepareCallback<Statement, MysqlType> for &'_ mut mysql_async::Conn {
     async fn prepare(
-        &mut self,
+        self,
         sql: &str,
         _metadata: &[MysqlType],
         _is_for_cache: diesel::connection::statement_cache::PrepareForCache,
-    ) -> QueryResult<Statement> {
-        Ok(self.prep(sql).await.map_err(ErrorHelper)?)
+    ) -> QueryResult<(Statement, Self)> {
+        let s = self.prep(sql).await.map_err(ErrorHelper)?;
+        Ok((s, self))
     }
 }
 
@@ -155,17 +164,22 @@ impl AsyncMysqlConnection {
         Ok(conn)
     }
 
-    async fn with_prepared_statement<'a, T, F, R>(
-        &'a mut self,
+    fn with_prepared_statement<'conn, T, F, R>(
+        &'conn mut self,
         query: T,
-        callback: impl FnOnce(&'a mut mysql_async::Conn, &'a Statement, ToSqlHelper) -> F,
-    ) -> QueryResult<R>
+        callback: impl (FnOnce(&'conn mut mysql_async::Conn, &'conn Statement, ToSqlHelper) -> F)
+            + Send
+            + 'conn,
+    ) -> BoxFuture<'conn, QueryResult<R>>
     where
+        R: Send + 'conn,
         T: QueryFragment<Mysql> + QueryId + Send,
-        F: Future<Output = QueryResult<R>>,
+        F: Future<Output = QueryResult<R>> + Send,
     {
         let mut bind_collector = RawBytesBindCollector::<Mysql>::new();
-        query.collect_binds(&mut bind_collector, &mut (), &Mysql)?;
+        if let Err(e) = query.collect_binds(&mut bind_collector, &mut (), &Mysql) {
+            return futures::future::ready(Err(e)).boxed();
+        }
 
         let binds = bind_collector.binds;
         let metadata = bind_collector.metadata;
@@ -177,24 +191,19 @@ impl AsyncMysqlConnection {
             ..
         } = self;
 
-        let conn = &mut *conn;
+        let stmt = stmt_cache.cached_prepared_statement(query, &metadata, conn, &Mysql);
 
-        let stmt = {
-            let stmt = stmt_cache
-                .cached_prepared_statement(query, &metadata, conn, &Mysql)
-                .await?;
-            stmt
-        };
-
-        let stmt = match stmt {
-            MaybeCached::CannotCache(stmt) => {
-                *last_stmt = Some(stmt);
-                last_stmt.as_ref().unwrap()
-            }
-            MaybeCached::Cached(s) => s,
-            _ => unreachable!("We've opted into breaking diesel changes and want to know if things break because someone added a new variant here")
-        };
+        stmt.and_then(|(stmt, conn)|async  move {
 
-        callback(&mut self.conn, stmt, ToSqlHelper { metadata, binds }).await
+            let stmt = match stmt {
+                MaybeCached::CannotCache(stmt) => {
+                    *last_stmt = Some(stmt);
+                    last_stmt.as_ref().unwrap()
+                }
+                MaybeCached::Cached(s) => s,
+                _ => unreachable!("We've opted into breaking diesel changes and want to know if things break because someone added a new variant here")
+            };
+            callback(conn, stmt, ToSqlHelper{metadata, binds}).await
+        }).boxed()
     }
 }

src/mysql/row.rs

@@ -61,6 +61,9 @@ impl<'a> diesel::row::Row<'a, Mysql> for MysqlRow {
         let buffer = match dbg!(value) {
             Value::NULL => None,
             Value::Bytes(b) => {
+                dbg!(&b);
+                dbg!(b.len());
+                dbg!(std::mem::size_of::<diesel::mysql::data_types::MysqlTime>());
                 // deserialize gets the length prepended, so we just use that buffer
                 // directly
                 Some(Cow::Borrowed(b as &[_]))