add ready_status on simple queries

Author: conradludgate

postgres/src/client.rs

@@ -439,7 +439,9 @@ impl Client {
     /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
     /// them to this method!
     pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
-        self.connection.block_on(self.client.batch_execute(query))
+        self.connection
+            .block_on(self.client.batch_execute(query))
+            .map(|_| ())
     }
 
     /// Begins a new database transaction.

postgres/src/transaction.rs

@@ -35,6 +35,7 @@ impl<'a> Transaction<'a> {
     pub fn commit(mut self) -> Result<(), Error> {
         self.connection
             .block_on(self.transaction.take().unwrap().commit())
+            .map(|_| ())
     }
 
     /// Rolls the transaction back, discarding all changes made within it.
@@ -43,6 +44,7 @@ impl<'a> Transaction<'a> {
     pub fn rollback(mut self) -> Result<(), Error> {
         self.connection
             .block_on(self.transaction.take().unwrap().rollback())
+            .map(|_| ())
     }
 
     /// Like `Client::prepare`.
@@ -193,6 +195,7 @@ impl<'a> Transaction<'a> {
     pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
         self.connection
             .block_on(self.transaction.as_ref().unwrap().batch_execute(query))
+            .map(|_| ())
     }
 
     /// Like `Client::cancel_token`.

tokio-postgres/src/client.rs

@@ -17,8 +17,8 @@ use crate::types::{Oid, ToSql, Type};
 use crate::Socket;
 use crate::{
     copy_both, copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken,
-    CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
-    TransactionBuilder,
+    CopyInSink, Error, ReadyForQueryStatus, Row, SimpleQueryMessage, Statement, ToStatement,
+    Transaction, TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -526,7 +526,7 @@ 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 async fn batch_execute(&self, query: &str) -> Result<(), Error> {
+    pub async fn batch_execute(&self, query: &str) -> Result<ReadyForQueryStatus, Error> {
         simple_query::batch_execute(self.inner(), query).await
     }
 

tokio-postgres/src/lib.rs

@@ -119,6 +119,8 @@
 #![doc(html_root_url = "https://docs.rs/tokio-postgres/0.7")]
 #![warn(rust_2018_idioms, clippy::all, missing_docs)]
 
+use postgres_protocol::message::backend::ReadyForQueryBody;
+
 pub use crate::cancel_token::CancelToken;
 pub use crate::client::Client;
 pub use crate::config::Config;
@@ -143,6 +145,31 @@ pub use crate::transaction::Transaction;
 pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder};
 use crate::types::ToSql;
 
+/// After executing a query, the connection will be in one of these states
+#[derive(Clone, Copy, Debug, PartialEq)]
+#[repr(u8)]
+pub enum ReadyForQueryStatus {
+    /// Connection state is unknown
+    Unknown,
+    /// Connection is idle (no transactions)
+    Idle = b'I',
+    /// Connection is in a transaction block
+    Transaction = b'T',
+    /// Connection is in a failed transaction block
+    FailedTransaction = b'E',
+}
+
+impl From<ReadyForQueryBody> for ReadyForQueryStatus {
+    fn from(value: ReadyForQueryBody) -> Self {
+        match value.status() {
+            b'I' => Self::Idle,
+            b'T' => Self::Transaction,
+            b'E' => Self::FailedTransaction,
+            _ => Self::Unknown,
+        }
+    }
+}
+
 pub mod binary_copy;
 mod bind;
 #[cfg(feature = "runtime")]

tokio-postgres/src/query.rs

@@ -2,7 +2,7 @@ use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::types::{BorrowToSql, IsNull};
-use crate::{Error, Portal, Row, Statement};
+use crate::{Error, Portal, ReadyForQueryStatus, Row, Statement};
 use bytes::{BufMut, Bytes, BytesMut};
 use futures_util::{ready, Stream};
 use log::{debug, log_enabled, Level};
@@ -55,7 +55,7 @@ where
         statement,
         responses,
         command_tag: None,
-        status: None,
+        status: ReadyForQueryStatus::Unknown,
         output_format: Format::Binary,
         _p: PhantomPinned,
     })
@@ -109,7 +109,7 @@ where
         statement,
         responses,
         command_tag: None,
-        status: None,
+        status: ReadyForQueryStatus::Unknown,
         output_format: Format::Text,
         _p: PhantomPinned,
     })
@@ -132,7 +132,7 @@ pub async fn query_portal(
         statement: portal.statement().clone(),
         responses,
         command_tag: None,
-        status: None,
+        status: ReadyForQueryStatus::Unknown,
         output_format: Format::Binary,
         _p: PhantomPinned,
     })
@@ -266,7 +266,7 @@ pin_project! {
         responses: Responses,
         command_tag: Option<String>,
         output_format: Format,
-        status: Option<u8>,
+        status: ReadyForQueryStatus,
         #[pin]
         _p: PhantomPinned,
     }
@@ -293,7 +293,7 @@ impl Stream for RowStream {
                     }
                 }
                 Message::ReadyForQuery(status) => {
-                    *this.status = Some(status.status());
+                    *this.status = status.into();
                     return Poll::Ready(None);
                 }
                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
@@ -313,7 +313,7 @@ impl RowStream {
     /// Returns if the connection is ready for querying, with the status of the connection.
     ///
     /// This might be available only after the stream has been exhausted.
-    pub fn ready_status(&self) -> Option<u8> {
+    pub fn ready_status(&self) -> ReadyForQueryStatus {
         self.status
     }
 }

tokio-postgres/src/simple_query.rs

@@ -1,7 +1,7 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
-use crate::{Error, SimpleQueryMessage, SimpleQueryRow};
+use crate::{Error, ReadyForQueryStatus, SimpleQueryMessage, SimpleQueryRow};
 use bytes::Bytes;
 use fallible_iterator::FallibleIterator;
 use futures_util::{ready, Stream};
@@ -40,19 +40,23 @@ pub async fn simple_query(client: &InnerClient, query: &str) -> Result<SimpleQue
     Ok(SimpleQueryStream {
         responses,
         columns: None,
+        status: ReadyForQueryStatus::Unknown,
         _p: PhantomPinned,
     })
 }
 
-pub async fn batch_execute(client: &InnerClient, query: &str) -> Result<(), Error> {
+pub async fn batch_execute(
+    client: &InnerClient,
+    query: &str,
+) -> Result<ReadyForQueryStatus, Error> {
     debug!("executing statement batch: {}", query);
 
     let buf = encode(client, query)?;
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
     loop {
         match responses.next().await? {
-            Message::ReadyForQuery(_) => return Ok(()),
+            Message::ReadyForQuery(status) => return Ok(status.into()),
             Message::CommandComplete(_)
             | Message::EmptyQueryResponse
             | Message::RowDescription(_)
@@ -74,11 +78,21 @@ pin_project! {
     pub struct SimpleQueryStream {
         responses: Responses,
         columns: Option<Arc<[SimpleColumn]>>,
+        status: ReadyForQueryStatus,
         #[pin]
         _p: PhantomPinned,
     }
 }
 
+impl SimpleQueryStream {
+    /// Returns if the connection is ready for querying, with the status of the connection.
+    ///
+    /// This might be available only after the stream has been exhausted.
+    pub fn ready_status(&self) -> ReadyForQueryStatus {
+        self.status
+    }
+}
+
 impl Stream for SimpleQueryStream {
     type Item = Result<SimpleQueryMessage, Error>;
 
@@ -117,7 +131,10 @@ impl Stream for SimpleQueryStream {
                     };
                     return Poll::Ready(Some(Ok(SimpleQueryMessage::Row(row))));
                 }
-                Message::ReadyForQuery(_) => return Poll::Ready(None),
+                Message::ReadyForQuery(s) => {
+                    *this.status = s.into();
+                    return Poll::Ready(None);
+                }
                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
         }

tokio-postgres/src/transaction.rs

@@ -9,8 +9,8 @@ use crate::types::{BorrowToSql, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    bind, query, slice_iter, CancelToken, Client, CopyInSink, Error, Portal, Row,
-    SimpleQueryMessage, Statement, ToStatement,
+    bind, query, slice_iter, CancelToken, Client, CopyInSink, Error, Portal, ReadyForQueryStatus,
+    Row, SimpleQueryMessage, Statement, ToStatement,
 };
 use bytes::Buf;
 use futures_util::TryStreamExt;
@@ -65,7 +65,7 @@ impl<'a> Transaction<'a> {
     }
 
     /// Consumes the transaction, committing all changes made within it.
-    pub async fn commit(mut self) -> Result<(), Error> {
+    pub async fn commit(mut self) -> Result<ReadyForQueryStatus, Error> {
         self.done = true;
         let query = if let Some(sp) = self.savepoint.as_ref() {
             format!("RELEASE {}", sp.name)
@@ -78,7 +78,7 @@ impl<'a> Transaction<'a> {
     /// Rolls the transaction back, discarding all changes made within it.
     ///
     /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
-    pub async fn rollback(mut self) -> Result<(), Error> {
+    pub async fn rollback(mut self) -> Result<ReadyForQueryStatus, Error> {
         self.done = true;
         let query = if let Some(sp) = self.savepoint.as_ref() {
             format!("ROLLBACK TO {}", sp.name)
@@ -261,7 +261,7 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::batch_execute`.
-    pub async fn batch_execute(&self, query: &str) -> Result<(), Error> {
+    pub async fn batch_execute(&self, query: &str) -> Result<ReadyForQueryStatus, Error> {
         self.client.batch_execute(query).await
     }
 

tokio-postgres/tests/test/main.rs

@@ -16,7 +16,8 @@ use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::{NoTls, NoTlsStream};
 use tokio_postgres::types::{Kind, Type};
 use tokio_postgres::{
-    AsyncMessage, Client, Config, Connection, Error, IsolationLevel, SimpleQueryMessage,
+    AsyncMessage, Client, Config, Connection, Error, IsolationLevel, ReadyForQueryStatus,
+    SimpleQueryMessage,
 };
 
 mod binary_copy;
@@ -189,6 +190,25 @@ async fn insert_select() {
     assert_eq!(rows[1].get::<_, &str>(1), "bob");
 }
 
+#[tokio::test]
+async fn foo() {
+    let client = connect("user=postgres").await;
+    client.batch_execute(
+        "
+        create or replace function test2 (i int) returns table (j int) language plpgsql as $$
+        declare
+            counter int;
+        begin
+            for counter in 1..i loop
+                select delay.counter into j from (select case when counter > 5 then pg_sleep(5.0) end sleep, counter) as delay;
+                raise notice '%d', counter;
+                return next;
+            end loop;
+        end; $$;
+        "
+    ).await.unwrap();
+}
+
 #[tokio::test]
 async fn custom_enum() {
     let client = connect("user=postgres").await;
@@ -365,7 +385,7 @@ async fn ready_for_query() {
     pin_mut!(row_stream);
     while row_stream.next().await.is_none() {}
 
-    assert_eq!(row_stream.ready_status(), Some(b'T'));
+    assert_eq!(row_stream.ready_status(), ReadyForQueryStatus::Transaction);
 
     let row_stream = client
         .query_raw_txt("ROLLBACK", [] as [Option<&str>; 0])
@@ -375,7 +395,7 @@ async fn ready_for_query() {
     pin_mut!(row_stream);
     while row_stream.next().await.is_none() {}
 
-    assert_eq!(row_stream.ready_status(), Some(b'I'));
+    assert_eq!(row_stream.ready_status(), ReadyForQueryStatus::Idle);
 }
 
 #[tokio::test]