Overhaul simple_query

Author: sfackler

postgres/src/client.rs

@@ -303,22 +303,7 @@ impl Client {
     /// functionality to safely imbed that data in the request. Do not form statements via string concatenation and pass
     /// them to this method!
     pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
-        self.simple_query_iter(query)?.collect()
-    }
-
-    /// Like `simple_query`, except that it returns a fallible iterator over the resulting values rather than buffering
-    /// the response in memory.
-    ///
-    /// # Warning
-    ///
-    /// Prepared statements should be use for any query which contains user-specified data, as they provided the
-    /// functionality to safely imbed that data in the request. Do not form statements via string concatenation and pass
-    /// them to this method!
-    pub fn simple_query_iter<'a>(
-        &'a mut self,
-        query: &'a str,
-    ) -> Result<impl FallibleIterator<Item = SimpleQueryMessage, Error = Error> + 'a, Error> {
-        Ok(Iter::new(self.0.simple_query(query)))
+        executor::block_on(self.0.simple_query(query))
     }
 
     /// Executes a sequence of SQL statements using the simple query protocol.

postgres/src/transaction.rs

@@ -137,15 +137,7 @@ impl<'a> Transaction<'a> {
 
     /// Like `Client::simple_query`.
     pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
-        self.simple_query_iter(query)?.collect()
-    }
-
-    /// Like `Client::simple_query_iter`.
-    pub fn simple_query_iter<'b>(
-        &'b mut self,
-        query: &'b str,
-    ) -> Result<impl FallibleIterator<Item = SimpleQueryMessage, Error = Error> + 'b, Error> {
-        Ok(Iter::new(self.0.simple_query(query)))
+        executor::block_on(self.0.simple_query(query))
     }
 
     /// Like `Client::batch_execute`.

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::simple_query::SimpleQueryStream;
 use crate::copy_out::CopyStream;
 use crate::query::RowStream;
 use crate::slice_iter;
@@ -20,7 +21,7 @@ use crate::{Error, Statement};
 use bytes::IntoBuf;
 use fallible_iterator::FallibleIterator;
 use futures::channel::mpsc;
-use futures::{future, Stream, TryStream, TryStreamExt};
+use futures::{future, TryStream, TryStreamExt};
 use futures::{ready, StreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::backend::Message;
@@ -320,11 +321,15 @@ impl Client {
     /// Prepared statements should be use for any query which contains user-specified data, as they provided the
     /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
     /// them to this method!
-    pub fn simple_query<'a>(
-        &'a self,
-        query: &'a str,
-    ) -> impl Stream<Item = Result<SimpleQueryMessage, Error>> + 'a {
-        simple_query::simple_query(self.inner(), query)
+    pub async fn simple_query(
+        &self,
+        query: &str,
+    ) -> Result<Vec<SimpleQueryMessage>, Error> {
+        self.simple_query_raw(query).await?.try_collect().await
+    }
+
+    pub(crate) async fn simple_query_raw(&self, query: &str) -> Result<SimpleQueryStream, Error> {
+        simple_query::simple_query(self.inner(), query).await
     }
 
     /// Executes a sequence of SQL statements using the simple query protocol.

tokio-postgres/src/connect.rs

@@ -4,7 +4,7 @@ use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
 use crate::tls::{MakeTlsConnect, TlsConnect};
 use crate::{Client, Config, Connection, Error, SimpleQueryMessage, Socket};
-use futures::future;
+use futures::{future, Future};
 use futures::{FutureExt, Stream};
 use pin_utils::pin_mut;
 use std::io;
@@ -50,7 +50,7 @@ where
         }
     }
 
-    return Err(error.unwrap());
+    Err(error.unwrap())
 }
 
 async fn connect_once<T>(
@@ -73,7 +73,16 @@ where
     let (mut client, mut connection) = connect_raw(socket, tls, config).await?;
 
     if let TargetSessionAttrs::ReadWrite = config.target_session_attrs {
-        let rows = client.simple_query("SHOW transaction_read_only");
+        let rows = client.simple_query_raw("SHOW transaction_read_only");
+        pin_mut!(rows);
+
+        let rows = future::poll_fn(|cx| {
+            if connection.poll_unpin(cx)?.is_ready() {
+                return Poll::Ready(Err(Error::closed()));
+            }
+
+            rows.as_mut().poll(cx)
+        }).await?;
         pin_mut!(rows);
 
         loop {

tokio-postgres/src/simple_query.rs

@@ -3,27 +3,24 @@ use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::{Error, SimpleQueryMessage, SimpleQueryRow};
 use fallible_iterator::FallibleIterator;
-use futures::{ready, Stream, TryFutureExt};
+use futures::{ready, Stream};
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
-pub fn simple_query<'a>(
-    client: &'a InnerClient,
-    query: &'a str,
-) -> impl Stream<Item = Result<SimpleQueryMessage, Error>> + 'a {
-    let f = async move {
-        let buf = encode(query)?;
-        let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+pub async fn simple_query(
+    client: &InnerClient,
+    query: &str,
+) -> Result<SimpleQueryStream, Error> {
+    let buf = encode(query)?;
+    let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
-        Ok(SimpleQuery {
-            responses,
-            columns: None,
-        })
-    };
-    f.try_flatten_stream()
+    Ok(SimpleQueryStream {
+        responses,
+        columns: None,
+    })
 }
 
 pub async fn batch_execute(client: &InnerClient, query: &str) -> Result<(), Error> {
@@ -48,12 +45,12 @@ fn encode(query: &str) -> Result<Vec<u8>, Error> {
     Ok(buf)
 }
 
-struct SimpleQuery {
+pub struct SimpleQueryStream {
     responses: Responses,
     columns: Option<Arc<[String]>>,
 }
 
-impl Stream for SimpleQuery {
+impl Stream for SimpleQueryStream {
     type Item = Result<SimpleQueryMessage, Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

tokio-postgres/src/transaction.rs

@@ -12,7 +12,7 @@ use crate::{
     bind, query, slice_iter, Client, Error, Portal, Row, SimpleQueryMessage, Statement, ToStatement,
 };
 use bytes::IntoBuf;
-use futures::{Stream, TryStream, TryStreamExt};
+use futures::{TryStream, TryStreamExt};
 use postgres_protocol::message::frontend;
 use std::error;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -224,11 +224,11 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::simple_query`.
-    pub fn simple_query<'b>(
-        &'b self,
-        query: &'b str,
-    ) -> impl Stream<Item = Result<SimpleQueryMessage, Error>> + 'b {
-        self.client.simple_query(query)
+    pub async fn simple_query(
+        &self,
+        query: &str,
+    ) -> Result<Vec<SimpleQueryMessage>, Error> {
+        self.client.simple_query(query).await
     }
 
     /// Like `Client::batch_execute`.

tokio-postgres/tests/test/main.rs

@@ -264,7 +264,6 @@ async fn simple_query() {
             INSERT INTO foo (name) VALUES ('steven'), ('joe');
             SELECT * FROM foo ORDER BY id;",
         )
-        .try_collect::<Vec<_>>()
         .await
         .unwrap();