Skip to content

Commit f655c3b

Browse files
author
zach-com
committed
Address pull request comments
1 parent 07d9fb2 commit f655c3b

File tree

5 files changed

+17
-14
lines changed

5 files changed

+17
-14
lines changed

postgres/src/client.rs

+9-2
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::{
44
ToStatement, Transaction, TransactionBuilder,
55
};
66
use std::task::Poll;
7+
use std::time::Duration;
78
use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
89
use tokio_postgres::types::{BorrowToSql, ToSql, Type};
910
use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket};
@@ -414,8 +415,14 @@ impl Client {
414415
}
415416

416417
/// Validates connection, timing out after specified duration.
417-
pub fn is_valid(&mut self, timeout: std::time::Duration) -> Result<(), Error> {
418-
self.connection.block_on(self.client.is_valid(timeout))
418+
pub fn is_valid(&mut self, timeout: Duration) -> Result<(), Error> {
419+
let is_valid = Client::is_valid_inner(&self.client, timeout);
420+
self.connection.block_on(is_valid)
421+
}
422+
423+
async fn is_valid_inner(client: &tokio_postgres::Client, timeout: Duration) -> Result<(), Error> {
424+
let trivial_query = client.simple_query("");
425+
tokio::time::timeout(timeout, trivial_query).await?.map(|_| ())
419426
}
420427

421428
/// Executes a sequence of SQL statements using the simple query protocol.

tokio-postgres/src/bind.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,9 @@ where
2020
I: IntoIterator<Item = P>,
2121
I::IntoIter: ExactSizeIterator,
2222
{
23+
type BytesResult = Result<bytes::Bytes, Error>;
2324
let name = format!("p{}", NEXT_ID.fetch_add(1, Ordering::SeqCst));
24-
let buf = client.with_buf(|buf| {
25+
let buf = client.with_buf::<_, BytesResult>(|buf| {
2526
query::encode_bind(&statement, params, &name, buf)?;
2627
frontend::sync(buf);
2728
Ok(buf.split().freeze())

tokio-postgres/src/client.rs

-9
Original file line numberDiff line numberDiff line change
@@ -450,15 +450,6 @@ impl Client {
450450
self.simple_query_raw(query).await?.try_collect().await
451451
}
452452

453-
/// Validates connection, timing out after specified duration.
454-
pub async fn is_valid(&self, timeout: Duration) -> Result<(), Error> {
455-
type SqmResult = Result<Vec<SimpleQueryMessage>, Error>;
456-
type SqmTimeout = Result<SqmResult, tokio::time::error::Elapsed>;
457-
let sqm_future = self.simple_query_raw("").await?.try_collect();
458-
let sqm_timeout: SqmTimeout = tokio::time::timeout(timeout, sqm_future).await;
459-
sqm_timeout.map_err(|_| Error::timeout())?.map(|_| ())
460-
}
461-
462453
pub(crate) async fn simple_query_raw(&self, query: &str) -> Result<SimpleQueryStream, Error> {
463454
simple_query::simple_query(self.inner(), query).await
464455
}

tokio-postgres/src/error/mod.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use postgres_protocol::message::backend::{ErrorFields, ErrorResponseBody};
55
use std::error::{self, Error as _Error};
66
use std::fmt;
77
use std::io;
8+
use tokio::time::error::Elapsed;
89

910
pub use self::sqlstate::*;
1011

@@ -493,8 +494,10 @@ impl Error {
493494
pub(crate) fn connect(e: io::Error) -> Error {
494495
Error::new(Kind::Connect, Some(Box::new(e)))
495496
}
497+
}
496498

497-
pub(crate) fn timeout() -> Error {
499+
impl From<Elapsed> for Error {
500+
fn from(_e: Elapsed) -> Error {
498501
Error::new(Kind::Timeout, None)
499502
}
500503
}

tokio-postgres/src/query.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,8 @@ pub async fn query_portal(
6161
portal: &Portal,
6262
max_rows: i32,
6363
) -> Result<RowStream, Error> {
64-
let buf = client.with_buf(|buf| {
64+
type BytesResult = Result<bytes::Bytes, Error>;
65+
let buf = client.with_buf::<_, BytesResult>(|buf| {
6566
frontend::execute(portal.name(), max_rows, buf).map_err(Error::encode)?;
6667
frontend::sync(buf);
6768
Ok(buf.split().freeze())

0 commit comments

Comments
 (0)