Allow pipelining with composed futures for Postgres

Author: stormshield-kg

src/pg/mod.rs

@@ -114,6 +114,48 @@ const FAKE_OID: u32 = 0;
 /// # }
 /// ```
 ///
+/// For more complex cases, an immutable reference to the connection need to be used:
+/// ```rust
+/// # include!("../doctest_setup.rs");
+/// use diesel_async::RunQueryDsl;
+///
+/// #
+/// # #[tokio::main(flavor = "current_thread")]
+/// # async fn main() {
+/// #     run_test().await.unwrap();
+/// # }
+/// #
+/// # async fn run_test() -> QueryResult<()> {
+/// #     use diesel::sql_types::{Text, Integer};
+/// #     let conn = &mut establish_connection().await;
+/// #
+///       async fn fn12(mut conn: &AsyncPgConnection) -> QueryResult<(i32, i32)> {
+///           let f1 = diesel::select(1_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+///           let f2 = diesel::select(2_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+///
+///           futures_util::try_join!(f1, f2)
+///       }
+///
+///       async fn fn34(mut conn: &AsyncPgConnection) -> QueryResult<(i32, i32)> {
+///           let f3 = diesel::select(3_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+///           let f4 = diesel::select(4_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+///
+///           futures_util::try_join!(f3, f4)
+///       }
+///
+///       let f12 = fn12(&conn);
+///       let f34 = fn34(&conn);
+///
+///       let ((r1, r2), (r3, r4)) = futures_util::try_join!(f12, f34).unwrap();
+///
+///       assert_eq!(r1, 1);
+///       assert_eq!(r2, 2);
+///       assert_eq!(r3, 3);
+///       assert_eq!(r4, 4);
+///       # Ok(())
+/// # }
+/// ```
+///
 /// ## TLS
 ///
 /// Connections created by [`AsyncPgConnection::establish`] do not support TLS.
@@ -136,6 +178,12 @@ pub struct AsyncPgConnection {
 }
 
 impl SimpleAsyncConnection for AsyncPgConnection {
+    async fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
+        SimpleAsyncConnection::batch_execute(&mut &*self, query).await
+    }
+}
+
+impl SimpleAsyncConnection for &AsyncPgConnection {
     async fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
         self.record_instrumentation(InstrumentationEvent::start_query(&StrQueryHelper::new(
             query,
@@ -167,6 +215,38 @@ impl AsyncConnectionCore for AsyncPgConnection {
     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,
+    {
+        AsyncConnectionCore::load(&mut &*self, source)
+    }
+
+    fn execute_returning_count<'conn, 'query, T>(
+        &'conn mut self,
+        source: T,
+    ) -> Self::ExecuteFuture<'conn, 'query>
+    where
+        T: QueryFragment<Self::Backend> + QueryId + 'query,
+    {
+        AsyncConnectionCore::execute_returning_count(&mut &*self, source)
+    }
+}
+
+impl AsyncConnectionCore for &AsyncPgConnection {
+    type LoadFuture<'conn, 'query> =
+        <AsyncPgConnection as AsyncConnectionCore>::LoadFuture<'conn, 'query>;
+
+    type ExecuteFuture<'conn, 'query> =
+        <AsyncPgConnection as AsyncConnectionCore>::ExecuteFuture<'conn, 'query>;
+
+    type Stream<'conn, 'query> = <AsyncPgConnection as AsyncConnectionCore>::Stream<'conn, 'query>;
+
+    type Row<'conn, 'query> = <AsyncPgConnection as AsyncConnectionCore>::Row<'conn, 'query>;
+
+    type Backend = <AsyncPgConnection as AsyncConnectionCore>::Backend;
+
     fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
     where
         T: AsQuery + 'query,
@@ -962,4 +1042,37 @@ mod tests {
         assert_eq!(r1, 1);
         assert_eq!(r2, 2);
     }
+
+    #[tokio::test]
+    async fn pipelining_with_composed_futures() {
+        let database_url =
+            std::env::var("DATABASE_URL").expect("DATABASE_URL must be set in order to run tests");
+        let conn = crate::AsyncPgConnection::establish(&database_url)
+            .await
+            .unwrap();
+
+        async fn fn12(mut conn: &AsyncPgConnection) -> QueryResult<(i32, i32)> {
+            let f1 = diesel::select(1_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+            let f2 = diesel::select(2_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+
+            futures_util::try_join!(f1, f2)
+        }
+
+        async fn fn34(mut conn: &AsyncPgConnection) -> QueryResult<(i32, i32)> {
+            let f3 = diesel::select(3_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+            let f4 = diesel::select(4_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+
+            futures_util::try_join!(f3, f4)
+        }
+
+        let f12 = fn12(&conn);
+        let f34 = fn34(&conn);
+
+        let ((r1, r2), (r3, r4)) = futures_util::try_join!(f12, f34).unwrap();
+
+        assert_eq!(r1, 1);
+        assert_eq!(r2, 2);
+        assert_eq!(r3, 3);
+        assert_eq!(r4, 4);
+    }
 }