From f655c3b74c59144aed35859fecf830e5ce5960db Mon Sep 17 00:00:00 2001
From: zach-com <54674467+zach-com@users.noreply.github.com>
Date: Wed, 16 Dec 2020 15:26:06 -0500
Subject: [PATCH] Address pull request comments

---
 postgres/src/client.rs          | 11 +++++++++--
 tokio-postgres/src/bind.rs      |  3 ++-
 tokio-postgres/src/client.rs    |  9 ---------
 tokio-postgres/src/error/mod.rs |  5 ++++-
 tokio-postgres/src/query.rs     |  3 ++-
 5 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/postgres/src/client.rs b/postgres/src/client.rs
index 00114165..fcd7f772 100644
--- a/postgres/src/client.rs
+++ b/postgres/src/client.rs
@@ -4,6 +4,7 @@ use crate::{
     ToStatement, Transaction, TransactionBuilder,
 };
 use std::task::Poll;
+use std::time::Duration;
 use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
 use tokio_postgres::types::{BorrowToSql, ToSql, Type};
 use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket};
@@ -414,8 +415,14 @@ impl Client {
     }
 
     /// 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))
+    pub fn is_valid(&mut self, timeout: Duration) -> Result<(), Error> {
+        let is_valid = Client::is_valid_inner(&self.client, timeout);
+        self.connection.block_on(is_valid)
+    }
+
+    async fn is_valid_inner(client: &tokio_postgres::Client, timeout: Duration) -> Result<(), Error> {
+        let trivial_query = client.simple_query("");
+        tokio::time::timeout(timeout, trivial_query).await?.map(|_| ())
     }
 
     /// Executes a sequence of SQL statements using the simple query protocol.
diff --git a/tokio-postgres/src/bind.rs b/tokio-postgres/src/bind.rs
index 9c5c4921..75bd938d 100644
--- a/tokio-postgres/src/bind.rs
+++ b/tokio-postgres/src/bind.rs
@@ -20,8 +20,9 @@ where
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
 {
+    type BytesResult = Result<bytes::Bytes, Error>;
     let name = format!("p{}", NEXT_ID.fetch_add(1, Ordering::SeqCst));
-    let buf = client.with_buf(|buf| {
+    let buf = client.with_buf::<_, BytesResult>(|buf| {
         query::encode_bind(&statement, params, &name, buf)?;
         frontend::sync(buf);
         Ok(buf.split().freeze())
diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs
index 090b4148..359a7cd1 100644
--- a/tokio-postgres/src/client.rs
+++ b/tokio-postgres/src/client.rs
@@ -450,15 +450,6 @@ 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
     }
diff --git a/tokio-postgres/src/error/mod.rs b/tokio-postgres/src/error/mod.rs
index 52909560..bb0dec9a 100644
--- a/tokio-postgres/src/error/mod.rs
+++ b/tokio-postgres/src/error/mod.rs
@@ -5,6 +5,7 @@ use postgres_protocol::message::backend::{ErrorFields, ErrorResponseBody};
 use std::error::{self, Error as _Error};
 use std::fmt;
 use std::io;
+use tokio::time::error::Elapsed;
 
 pub use self::sqlstate::*;
 
@@ -493,8 +494,10 @@ impl Error {
     pub(crate) fn connect(e: io::Error) -> Error {
         Error::new(Kind::Connect, Some(Box::new(e)))
     }
+}
 
-    pub(crate) fn timeout() -> Error {
+impl From<Elapsed> for Error {
+    fn from(_e: Elapsed) -> Error {
         Error::new(Kind::Timeout, None)
     }
 }
diff --git a/tokio-postgres/src/query.rs b/tokio-postgres/src/query.rs
index f139ed91..d6179de4 100644
--- a/tokio-postgres/src/query.rs
+++ b/tokio-postgres/src/query.rs
@@ -61,7 +61,8 @@ pub async fn query_portal(
     portal: &Portal,
     max_rows: i32,
 ) -> Result<RowStream, Error> {
-    let buf = client.with_buf(|buf| {
+    type BytesResult = Result<bytes::Bytes, Error>;
+    let buf = client.with_buf::<_, BytesResult>(|buf| {
         frontend::execute(portal.name(), max_rows, buf).map_err(Error::encode)?;
         frontend::sync(buf);
         Ok(buf.split().freeze())