Overhaul copy_out

Author: sfackler

postgres/src/client.rs

@@ -275,15 +275,15 @@ impl Client {
     /// # Ok(())
     /// # }
     /// ```
-    pub fn copy_out<'a, T>(
-        &'a mut self,
-        query: &'a T,
-        params: &'a [&(dyn ToSql + Sync)],
-    ) -> Result<impl BufRead + 'a, Error>
+    pub fn copy_out<T>(
+        &mut self,
+        query: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<impl BufRead, Error>
     where
         T: ?Sized + ToStatement,
     {
-        let stream = self.0.copy_out(query, params);
+        let stream = executor::block_on(self.0.copy_out(query, params))?;
         CopyOutReader::new(stream)
     }
 

postgres/src/transaction.rs

@@ -123,15 +123,15 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::copy_out`.
-    pub fn copy_out<'b, T>(
-        &'b mut self,
-        query: &'b T,
-        params: &'b [&(dyn ToSql + Sync)],
-    ) -> Result<impl BufRead + 'b, Error>
+    pub fn copy_out<T>(
+        &mut self,
+        query: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<impl BufRead, Error>
     where
         T: ?Sized + ToStatement,
     {
-        let stream = self.0.copy_out(query, params);
+        let stream = executor::block_on(self.0.copy_out(query, params))?;
         CopyOutReader::new(stream)
     }
 

tokio-postgres/src/client.rs

@@ -3,6 +3,7 @@ use crate::cancel_query;
 use crate::codec::BackendMessages;
 use crate::config::{Host, SslMode};
 use crate::connection::{Request, RequestMessages};
+use crate::copy_out::CopyStream;
 use crate::query::RowStream;
 use crate::slice_iter;
 #[cfg(feature = "runtime")]
@@ -16,10 +17,10 @@ use crate::{cancel_query_raw, copy_in, copy_out, query, Transaction};
 use crate::{prepare, SimpleQueryMessage};
 use crate::{simple_query, Row};
 use crate::{Error, Statement};
-use bytes::{Bytes, IntoBuf};
+use bytes::IntoBuf;
 use fallible_iterator::FallibleIterator;
 use futures::channel::mpsc;
-use futures::{future, Stream, TryFutureExt, TryStream, TryStreamExt};
+use futures::{future, Stream, TryStream, TryStreamExt};
 use futures::{ready, StreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::backend::Message;
@@ -293,20 +294,17 @@ impl Client {
     /// # Panics
     ///
     /// Panics if the number of parameters provided does not match the number expected.
-    pub fn copy_out<'a, T>(
-        &'a self,
-        statement: &'a T,
-        params: &'a [&(dyn ToSql + Sync)],
-    ) -> impl Stream<Item = Result<Bytes, Error>> + 'a
+    pub async fn copy_out<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<CopyStream, Error>
     where
         T: ?Sized + ToStatement,
     {
-        let f = async move {
-            let statement = statement.__convert().into_statement(self).await?;
-            let params = slice_iter(params);
-            Ok(copy_out::copy_out(self.inner(), statement, params))
-        };
-        f.try_flatten_stream()
+        let statement = statement.__convert().into_statement(self).await?;
+        let params = slice_iter(params);
+        copy_out::copy_out(self.inner(), statement, params).await
     }
 
     /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows.

tokio-postgres/src/copy_out.rs

@@ -4,26 +4,23 @@ use crate::connection::RequestMessages;
 use crate::types::ToSql;
 use crate::{query, Error, Statement};
 use bytes::Bytes;
-use futures::{ready, Stream, TryFutureExt};
+use futures::{ready, Stream};
 use postgres_protocol::message::backend::Message;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-pub fn copy_out<'a, I>(
-    client: &'a InnerClient,
+pub async fn copy_out<'a, I>(
+    client: &InnerClient,
     statement: Statement,
     params: I,
-) -> impl Stream<Item = Result<Bytes, Error>> + 'a
+) -> Result<CopyStream, Error>
 where
-    I: IntoIterator<Item = &'a dyn ToSql> + 'a,
+    I: IntoIterator<Item = &'a dyn ToSql>,
     I::IntoIter: ExactSizeIterator,
 {
-    let f = async move {
-        let buf = query::encode(&statement, params)?;
-        let responses = start(client, buf).await?;
-        Ok(CopyOut { responses })
-    };
-    f.try_flatten_stream()
+    let buf = query::encode(&statement, params)?;
+    let responses = start(client, buf).await?;
+    Ok(CopyStream { responses })
 }
 
 async fn start(client: &InnerClient, buf: Vec<u8>) -> Result<Responses, Error> {
@@ -42,11 +39,11 @@ async fn start(client: &InnerClient, buf: Vec<u8>) -> Result<Responses, Error> {
     Ok(responses)
 }
 
-struct CopyOut {
+pub struct CopyStream {
     responses: Responses,
 }
 
-impl Stream for CopyOut {
+impl Stream for CopyStream {
     type Item = Result<Bytes, Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

tokio-postgres/src/transaction.rs

@@ -1,5 +1,6 @@
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
+use crate::copy_out::CopyStream;
 use crate::query::RowStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
@@ -10,7 +11,7 @@ use crate::Socket;
 use crate::{
     bind, query, slice_iter, Client, Error, Portal, Row, SimpleQueryMessage, Statement, ToStatement,
 };
-use bytes::{Bytes, IntoBuf};
+use bytes::IntoBuf;
 use futures::{Stream, TryStream, TryStreamExt};
 use postgres_protocol::message::frontend;
 use std::error;
@@ -211,15 +212,15 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::copy_out`.
-    pub fn copy_out<'b, T>(
-        &'b self,
-        statement: &'b T,
-        params: &'b [&(dyn ToSql + Sync)],
-    ) -> impl Stream<Item = Result<Bytes, Error>> + 'b
+    pub async fn copy_out<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<CopyStream, Error>
     where
         T: ?Sized + ToStatement,
     {
-        self.client.copy_out(statement, params)
+        self.client.copy_out(statement, params).await
     }
 
     /// Like `Client::simple_query`.

tokio-postgres/tests/test/main.rs

@@ -506,7 +506,13 @@ async fn copy_out() {
         .unwrap();
 
     let stmt = client.prepare("COPY foo TO STDOUT").await.unwrap();
-    let data = client.copy_out(&stmt, &[]).try_concat().await.unwrap();
+    let data = client
+        .copy_out(&stmt, &[])
+        .await
+        .unwrap()
+        .try_concat()
+        .await
+        .unwrap();
     assert_eq!(&data[..], b"1\tjim\n2\tjoe\n");
 }