Use custom Future combinators to avoid GAT errors

stormshield-kg · stormshield-kg · commit c9c23d1e5b6b · 2025-09-19T20:55:19.000+02:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -36,6 +36,7 @@ deadpool = { version = "0.12", optional = true, default-features = false, featur
 ] }
 mobc = { version = ">=0.7,<0.10", optional = true }
 scoped-futures = { version = "0.1", features = ["std"] }
+pin-project-lite = "0.2.16"
 
 [dependencies.diesel]
 version = "~2.3.0"
diff --git a/src/pg/mod.rs b/src/pg/mod.rs
@@ -1167,45 +1167,42 @@ mod tests {
             .await
             .unwrap();
 
-        fn erase<'a, T: Future + Send + 'a>(t: T) -> impl Future<Output = T::Output> + Send + 'a {
-            t
-        }
-
         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);
 
-            erase(try_join(f1, f2)).await
+            try_join(f1, f2).await
         }
 
-        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);
-
-            try_join(f3, f4).boxed().await
-        }
-
-        async fn fn56(mut conn: &AsyncPgConnection) -> QueryResult<(i32, i32)> {
+        async fn fn37(
+            mut conn: &AsyncPgConnection,
+        ) -> QueryResult<(usize, Vec<i32>, i32, Vec<i32>, i32)> {
+            let f3 = diesel::select(0_i32.into_sql::<Integer>()).execute(&mut conn);
+            let f4 = diesel::select(4_i32.into_sql::<Integer>()).load::<i32>(&mut conn);
             let f5 = diesel::select(5_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
-            let f6 = diesel::select(6_i32.into_sql::<Integer>()).get_result::<i32>(&mut conn);
+            let f6 = diesel::select(6_i32.into_sql::<Integer>()).get_results::<i32>(&mut conn);
+            let f7 = diesel::select(7_i32.into_sql::<Integer>()).first::<i32>(&mut conn);
 
-            try_join!(f5.boxed(), f6.boxed())
+            try_join!(f3, f4, f5, f6, f7)
         }
 
         conn.transaction(|conn| {
             async move {
                 let f12 = fn12(conn);
-                let f34 = fn34(conn);
-                let f56 = fn56(conn);
+                let f37 = fn37(conn);
 
-                let ((r1, r2), (r3, r4), (r5, r6)) = try_join!(f12, f34, f56).unwrap();
+                let ((r1, r2), (r3, r4, r5, r6, r7)) = try_join!(f12, f37).unwrap();
 
                 assert_eq!(r1, 1);
                 assert_eq!(r2, 2);
-                assert_eq!(r3, 3);
-                assert_eq!(r4, 4);
+                assert_eq!(r3, 1);
+                assert_eq!(r4, vec![4]);
                 assert_eq!(r5, 5);
-                assert_eq!(r6, 6);
+                assert_eq!(r6, vec![6]);
+                assert_eq!(r7, 7);
+
+                fn12(conn).await?;
+                fn37(conn).await?;
 
                 QueryResult::<_>::Ok(())
             }
diff --git a/src/run_query_dsl/mod.rs b/src/run_query_dsl/mod.rs
@@ -1,13 +1,13 @@
+mod utils;
+
 use crate::AsyncConnectionCore;
 use diesel::associations::HasTable;
 use diesel::query_builder::IntoUpdateTarget;
 use diesel::result::QueryResult;
 use diesel::AsChangeset;
 use futures_core::future::BoxFuture;
-use futures_core::Stream;
-use futures_util::{future, stream, FutureExt, StreamExt, TryFutureExt, TryStreamExt};
+use futures_util::{stream, FutureExt, StreamExt, TryStreamExt};
 use std::future::Future;
-use std::pin::Pin;
 
 /// The traits used by `QueryDsl`.
 ///
@@ -22,7 +22,7 @@ pub mod methods {
     use diesel::expression::QueryMetadata;
     use diesel::query_builder::{AsQuery, QueryFragment, QueryId};
     use diesel::query_dsl::CompatibleType;
-    use futures_util::{Future, Stream, TryFutureExt};
+    use futures_util::{Future, Stream};
 
     /// The `execute` method
     ///
@@ -74,6 +74,7 @@ pub mod methods {
         type LoadFuture<'conn>: Future<Output = QueryResult<Self::Stream<'conn>>> + Send
         where
             Conn: 'conn;
+
         /// The inner stream returned by [`LoadQuery::internal_load`]
         type Stream<'conn>: Stream<Item = QueryResult<U>> + Send
         where
@@ -96,10 +97,7 @@ pub mod methods {
         ST: 'static,
     {
         type LoadFuture<'conn>
-            = future::MapOk<
-            Conn::LoadFuture<'conn, 'query>,
-            fn(Conn::Stream<'conn, 'query>) -> Self::Stream<'conn>,
-        >
+            = utils::Map<Conn::LoadFuture<'conn, 'query>, Self::Stream<'conn>>
         where
             Conn: 'conn;
 
@@ -112,33 +110,13 @@ pub mod methods {
             Conn: 'conn;
 
         fn internal_load(self, conn: &mut Conn) -> Self::LoadFuture<'_> {
-            conn.load(self)
-                .map_ok(map_result_stream_future::<U, _, _, DB, ST>)
+            utils::Map::new(conn.load(self), |stream| {
+                Ok(stream?.map(|row| {
+                    U::build_from_row(&row?).map_err(diesel::result::Error::DeserializationError)
+                }))
+            })
         }
     }
-
-    #[allow(clippy::type_complexity)]
-    fn map_result_stream_future<'s, 'a, U, S, R, DB, ST>(
-        stream: S,
-    ) -> stream::Map<S, fn(QueryResult<R>) -> QueryResult<U>>
-    where
-        S: Stream<Item = QueryResult<R>> + Send + 's,
-        R: diesel::row::Row<'a, DB> + 's,
-        DB: Backend + 'static,
-        U: FromSqlRow<ST, DB> + 'static,
-        ST: 'static,
-    {
-        stream.map(map_row_helper::<_, DB, U, ST>)
-    }
-
-    fn map_row_helper<'a, R, DB, U, ST>(row: QueryResult<R>) -> QueryResult<U>
-    where
-        U: FromSqlRow<ST, DB>,
-        R: diesel::row::Row<'a, DB>,
-        DB: Backend,
-    {
-        U::build_from_row(&row?).map_err(diesel::result::Error::DeserializationError)
-    }
 }
 
 /// The return types produced by the various [`RunQueryDsl`] methods
@@ -149,37 +127,24 @@ pub mod methods {
 // the same connection
 #[allow(type_alias_bounds)] // we need these bounds otherwise we cannot use GAT's
 pub mod return_futures {
+    use crate::run_query_dsl::utils;
+
     use super::methods::LoadQuery;
-    use diesel::QueryResult;
-    use futures_util::{future, stream};
+    use futures_util::stream;
     use std::pin::Pin;
 
     /// The future returned by [`RunQueryDsl::load`](super::RunQueryDsl::load)
     /// and [`RunQueryDsl::get_results`](super::RunQueryDsl::get_results)
     ///
     /// This is essentially `impl Future<Output = QueryResult<Vec<U>>>`
-    pub type LoadFuture<'conn, 'query, Q: LoadQuery<'query, Conn, U>, Conn, U> = future::AndThen<
-        Q::LoadFuture<'conn>,
-        stream::TryCollect<Q::Stream<'conn>, Vec<U>>,
-        fn(Q::Stream<'conn>) -> stream::TryCollect<Q::Stream<'conn>, Vec<U>>,
-    >;
+    pub type LoadFuture<'conn, 'query, Q: LoadQuery<'query, Conn, U>, Conn, U> =
+        utils::AndThen<Q::LoadFuture<'conn>, stream::TryCollect<Q::Stream<'conn>, Vec<U>>>;
 
     /// The future returned by [`RunQueryDsl::get_result`](super::RunQueryDsl::get_result)
     ///
     /// This is essentially `impl Future<Output = QueryResult<U>>`
-    pub type GetResult<'conn, 'query, Q: LoadQuery<'query, Conn, U>, Conn, U> = future::AndThen<
-        Q::LoadFuture<'conn>,
-        future::Map<
-            stream::StreamFuture<Pin<Box<Q::Stream<'conn>>>>,
-            fn((Option<QueryResult<U>>, Pin<Box<Q::Stream<'conn>>>)) -> QueryResult<U>,
-        >,
-        fn(
-            Q::Stream<'conn>,
-        ) -> future::Map<
-            stream::StreamFuture<Pin<Box<Q::Stream<'conn>>>>,
-            fn((Option<QueryResult<U>>, Pin<Box<Q::Stream<'conn>>>)) -> QueryResult<U>,
-        >,
-    >;
+    pub type GetResult<'conn, 'query, Q: LoadQuery<'query, Conn, U>, Conn, U> =
+        utils::AndThen<Q::LoadFuture<'conn>, utils::LoadNext<Pin<Box<Q::Stream<'conn>>>>>;
 }
 
 /// Methods used to execute queries.
@@ -346,13 +311,9 @@ pub trait RunQueryDsl<Conn>: Sized {
         Conn: AsyncConnectionCore,
         Self: methods::LoadQuery<'query, Conn, U> + 'query,
     {
-        fn collect_result<U, S>(stream: S) -> stream::TryCollect<S, Vec<U>>
-        where
-            S: Stream<Item = QueryResult<U>>,
-        {
-            stream.try_collect()
-        }
-        self.internal_load(conn).and_then(collect_result::<U, _>)
+        utils::AndThen::new(self.internal_load(conn), |stream| {
+            Ok(stream?.try_collect::<Vec<_>>())
+        })
     }
 
     /// Executes the given query, returning a [`Stream`] with the returned rows.
@@ -547,29 +508,9 @@ pub trait RunQueryDsl<Conn>: Sized {
         Conn: AsyncConnectionCore,
         Self: methods::LoadQuery<'query, Conn, U> + 'query,
     {
-        #[allow(clippy::type_complexity)]
-        fn get_next_stream_element<S, U>(
-            stream: S,
-        ) -> future::Map<
-            stream::StreamFuture<Pin<Box<S>>>,
-            fn((Option<QueryResult<U>>, Pin<Box<S>>)) -> QueryResult<U>,
-        >
-        where
-            S: Stream<Item = QueryResult<U>>,
-        {
-            fn map_option_to_result<U, S>(
-                (o, _): (Option<QueryResult<U>>, Pin<Box<S>>),
-            ) -> QueryResult<U> {
-                match o {
-                    Some(s) => s,
-                    None => Err(diesel::result::Error::NotFound),
-                }
-            }
-
-            Box::pin(stream).into_future().map(map_option_to_result)
-        }
-
-        self.load_stream(conn).and_then(get_next_stream_element)
+        utils::AndThen::new(self.internal_load(conn), |stream| {
+            Ok(utils::LoadNext::new(Box::pin(stream?)))
+        })
     }
 
     /// Runs the command, returning an `Vec` with the affected rows.
diff --git a/src/run_query_dsl/utils.rs b/src/run_query_dsl/utils.rs
@@ -0,0 +1,130 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use diesel::QueryResult;
+use futures_core::{ready, TryFuture, TryStream};
+use futures_util::TryStreamExt;
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Reimplementation of [`futures_util::future::Map`] without the generic closure argument
+    #[project = MapProj]
+    #[project_replace = MapProjReplace]
+    pub enum Map<Fut: Future, T> {
+         Incomplete {
+            #[pin]
+            future: Fut,
+            f: fn(Fut::Output) -> QueryResult<T>,
+        },
+        Complete,
+    }
+}
+
+impl<Fut: Future, T> Map<Fut, T> {
+    pub(crate) fn new(future: Fut, f: fn(Fut::Output) -> QueryResult<T>) -> Self {
+        Self::Incomplete { future, f }
+    }
+}
+
+impl<Fut: Future, T> Future for Map<Fut, T> {
+    type Output = QueryResult<T>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<QueryResult<T>> {
+        match self.as_mut().project() {
+            MapProj::Incomplete { future, .. } => {
+                let output = ready!(future.poll(cx));
+                match self.as_mut().project_replace(Map::Complete) {
+                    MapProjReplace::Incomplete { f, .. } => Poll::Ready(f(output)),
+                    MapProjReplace::Complete => unreachable!(),
+                }
+            }
+            MapProj::Complete => panic!("Map polled after completion"),
+        }
+    }
+}
+
+pin_project! {
+    /// Reimplementation of [`futures_util::future::AndThen`] without the generic closure argument
+    #[project = AndThenProj]
+    pub enum AndThen<Fut1: Future, Fut2> {
+        First {
+            #[pin]
+            future1: Map<Fut1, Fut2>,
+        },
+        Second {
+            #[pin]
+            future2: Fut2,
+        },
+        Empty,
+    }
+}
+
+impl<Fut1: Future, Fut2> AndThen<Fut1, Fut2> {
+    pub(crate) fn new(fut1: Fut1, f: fn(Fut1::Output) -> QueryResult<Fut2>) -> AndThen<Fut1, Fut2> {
+        Self::First {
+            future1: Map::new(fut1, f),
+        }
+    }
+}
+
+impl<Fut1, Fut2> Future for AndThen<Fut1, Fut2>
+where
+    Fut1: TryFuture<Error = diesel::result::Error>,
+    Fut2: TryFuture<Error = diesel::result::Error>,
+{
+    type Output = QueryResult<Fut2::Ok>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        loop {
+            match self.as_mut().project() {
+                AndThenProj::First { future1 } => match ready!(future1.try_poll(cx)) {
+                    Ok(future2) => self.set(Self::Second { future2 }),
+                    Err(error) => {
+                        self.set(Self::Empty);
+                        break Poll::Ready(Err(error));
+                    }
+                },
+                AndThenProj::Second { future2 } => {
+                    let output = ready!(future2.try_poll(cx));
+                    self.set(Self::Empty);
+                    break Poll::Ready(output);
+                }
+                AndThenProj::Empty => panic!("AndThen polled after completion"),
+            }
+        }
+    }
+}
+
+/// Converts a stream into a future, only yielding the first element.
+/// Based on [`futures_util::stream::StreamFuture`].
+pub struct LoadNext<St> {
+    stream: Option<St>,
+}
+
+impl<St> LoadNext<St> {
+    pub(crate) fn new(stream: St) -> Self {
+        Self {
+            stream: Some(stream),
+        }
+    }
+}
+
+impl<St> Future for LoadNext<St>
+where
+    St: TryStream<Error = diesel::result::Error> + Unpin,
+{
+    type Output = QueryResult<St::Ok>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let first = {
+            let s = self.stream.as_mut().expect("polling LoadNext twice");
+            ready!(s.try_poll_next_unpin(cx))
+        };
+        self.stream = None;
+        match first {
+            Some(first) => Poll::Ready(first),
+            None => Poll::Ready(Err(diesel::result::Error::NotFound)),
+        }
+    }
+}
