Read full stream to ensure CannotCache statements do not leak if stream is dropped prematurely

src/run_query_dsl/mod.rs

@@ -146,7 +146,7 @@ pub mod return_futures {
     ///
     /// This is essentially `impl Future<Output = 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>>>>>;
+        utils::AndThen<Q::LoadFuture<'conn>, utils::LoadNext<Pin<Box<Q::Stream<'conn>>>, U>>;
 }
 
 /// Methods used to execute queries.

src/run_query_dsl/utils.rs

@@ -3,7 +3,8 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use diesel::QueryResult;
-use futures_core::{ready, TryFuture, TryStream};
+use futures_core::{TryFuture, TryStream};
+use futures_util::stream::TryCollect;
 use futures_util::{TryFutureExt, TryStreamExt};
 
 // We use a custom future implementation here to erase some lifetimes
@@ -80,33 +81,48 @@ where
 
 /// Converts a stream into a future, only yielding the first element.
 /// Based on [`futures_util::stream::StreamFuture`].
-pub struct LoadNext<St> {
-    stream: Option<St>,
+///
+/// Consumes the entire stream to ensure proper cleanup before returning which is
+/// required to fix: https://github.com/weiznich/diesel_async/issues/269
+#[repr(transparent)]
+pub struct LoadNext<F, T>
+where
+    F: TryStream<Ok = T, Error = diesel::result::Error>,
+{
+    future: TryCollect<F, Vec<T>>,
 }
 
-impl<St> LoadNext<St> {
-    pub(crate) fn new(stream: St) -> Self {
+impl<F, T> LoadNext<F, T>
+where
+    F: TryStream<Ok = T, Error = diesel::result::Error>,
+{
+    pub(crate) fn new(stream: F) -> Self {
         Self {
-            stream: Some(stream),
+            future: stream.try_collect(),
         }
     }
 }
 
-impl<St> Future for LoadNext<St>
+impl<F, T> Future for LoadNext<F, T>
 where
-    St: TryStream<Error = diesel::result::Error> + Unpin,
+    F: TryStream<Ok = T, Error = diesel::result::Error>,
+    TryCollect<F, Vec<T>>: Future<Output = Result<Vec<T>, diesel::result::Error>>,
 {
-    type Output = QueryResult<St::Ok>;
+    type Output = QueryResult<T>;
 
-    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)),
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match unsafe {
+            // SAFETY: This projects pinning to the only inner field
+            self.map_unchecked_mut(|s| &mut s.future)
+        }
+        .poll(cx)
+        {
+            Poll::Ready(Ok(results)) => match results.into_iter().next() {
+                Some(first) => Poll::Ready(Ok(first)),
+                None => Poll::Ready(Err(diesel::result::Error::NotFound)),
+            },
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
         }
     }
 }