Support connection validation with timeout

Author: zach-com

postgres/src/client.rs

@@ -413,6 +413,11 @@ impl Client {
         self.connection.block_on(self.client.simple_query(query))
     }
 
+    /// Validates connection, timing out after specified duration.
+    pub fn is_valid(&mut self, timeout: std::time::Duration) -> Result<(), Error> {
+        self.connection.block_on(self.client.is_valid(timeout))
+    }
+
     /// Executes a sequence of SQL statements using the simple query protocol.
     ///
     /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that

tokio-postgres/src/client.rs

@@ -450,6 +450,15 @@ impl Client {
         self.simple_query_raw(query).await?.try_collect().await
     }
 
+    /// Validates connection, timing out after specified duration.
+    pub async fn is_valid(&self, timeout: Duration) -> Result<(), Error> {
+        type SqmResult = Result<Vec<SimpleQueryMessage>, Error>;
+        type SqmTimeout = Result<SqmResult, tokio::time::error::Elapsed>;
+        let sqm_future = self.simple_query_raw("").await?.try_collect();
+        let sqm_timeout: SqmTimeout = tokio::time::timeout(timeout, sqm_future).await;
+        sqm_timeout.map_err(|_| Error::timeout())?.map(|_| ())
+    }
+
     pub(crate) async fn simple_query_raw(&self, query: &str) -> Result<SimpleQueryStream, Error> {
         simple_query::simple_query(self.inner(), query).await
     }

tokio-postgres/src/error/mod.rs

@@ -354,6 +354,7 @@ enum Kind {
     RowCount,
     #[cfg(feature = "runtime")]
     Connect,
+    Timeout,
 }
 
 struct ErrorInner {
@@ -392,6 +393,7 @@ impl fmt::Display for Error {
             Kind::RowCount => fmt.write_str("query returned an unexpected number of rows")?,
             #[cfg(feature = "runtime")]
             Kind::Connect => fmt.write_str("error connecting to server")?,
+            Kind::Timeout => fmt.write_str("timeout waiting for server")?,
         };
         if let Some(ref cause) = self.0.cause {
             write!(fmt, ": {}", cause)?;
@@ -491,4 +493,8 @@ impl Error {
     pub(crate) fn connect(e: io::Error) -> Error {
         Error::new(Kind::Connect, Some(Box::new(e)))
     }
+
+    pub(crate) fn timeout() -> Error {
+        Error::new(Kind::Timeout, None)
+    }
 }