simple query ready for query (sfackler#22)

* add ready_status on simple queries

* add correct socket2 features

Author: conradludgate

postgres/src/client.rs

@@ -419,7 +419,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/Cargo.toml

@@ -58,6 +58,7 @@ postgres-protocol = { version = "0.6.6", path = "../postgres-protocol" }
 postgres-types = { version = "0.2.5", path = "../postgres-types" }
 tokio = { version = "1.27", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
+socket2 = { version = "0.5", features = ["all"] }
 rand = "0.8.5"
 whoami = "1.4.1"
 

tokio-postgres/src/client.rs

@@ -15,8 +15,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;
@@ -506,7 +506,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/generic_client.rs

@@ -175,7 +175,8 @@ impl GenericClient for Client {
     }
 
     async fn batch_execute(&self, query: &str) -> Result<(), Error> {
-        self.batch_execute(query).await
+        self.batch_execute(query).await?;
+        Ok(())
     }
 
     fn client(&self) -> &Client {
@@ -272,7 +273,8 @@ impl GenericClient for Transaction<'_> {
     }
 
     async fn batch_execute(&self, query: &str) -> Result<(), Error> {
-        self.batch_execute(query).await
+        self.batch_execute(query).await?;
+        Ok(())
     }
 
     fn client(&self) -> &Client {

tokio-postgres/src/lib.rs

@@ -118,6 +118,8 @@
 //! | `with-time-0_3` | Enable support for the 0.3 version of the `time` crate. | [time](https://crates.io/crates/time/0.3.0) 0.3 | no |
 #![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;
@@ -142,6 +144,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,
     })
@@ -269,7 +269,7 @@ pin_project! {
         responses: Responses,
         command_tag: Option<String>,
         output_format: Format,
-        status: Option<u8>,
+        status: ReadyForQueryStatus,
         #[pin]
         _p: PhantomPinned,
     }
@@ -296,7 +296,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()))),
@@ -316,7 +316,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

@@ -2,7 +2,7 @@ use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::query::extract_row_affected;
-use crate::{Error, SimpleQueryMessage, SimpleQueryRow};
+use crate::{Error, ReadyForQueryStatus, SimpleQueryMessage, SimpleQueryRow};
 use bytes::Bytes;
 use fallible_iterator::FallibleIterator;
 use futures_util::{ready, Stream};
@@ -41,19 +41,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(_)
@@ -75,11 +79,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>;
 
@@ -111,7 +125,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;
@@ -365,7 +366,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 +376,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]