Skip to content

Commit 09497b0

Browse files
committed
feat: add support for quick active health checks on connections
1 parent 20f0aae commit 09497b0

File tree

3 files changed

+22
-0
lines changed

3 files changed

+22
-0
lines changed

tokio-postgres/src/client.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -489,6 +489,12 @@ impl Client {
489489
simple_query::batch_execute(self.inner(), query).await
490490
}
491491

492+
/// Check the connection is alive and wait for the confirmation.
493+
pub async fn check_connection(&self) -> Result<(), Error> {
494+
// sync is a very quick message to test the connection health.
495+
query::sync(self.inner()).await
496+
}
497+
492498
/// Begins a new database transaction.
493499
///
494500
/// The transaction will roll back by default - use the `commit` method to commit it.

tokio-postgres/src/query.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,3 +242,13 @@ impl RowStream {
242242
self.rows_affected
243243
}
244244
}
245+
246+
pub async fn sync(client: &InnerClient) -> Result<(), Error> {
247+
let buf = Bytes::from_static(b"S\0\0\0\x04");
248+
let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
249+
250+
match responses.next().await? {
251+
Message::ReadyForQuery(_) => Ok(()),
252+
_ => Err(Error::unexpected_message()),
253+
}
254+
}

tokio-postgres/tests/test/main.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,12 @@ async fn scram_password_ok() {
149149
connect("user=scram_user password=password dbname=postgres").await;
150150
}
151151

152+
#[tokio::test]
153+
async fn sync() {
154+
let client = connect("user=postgres").await;
155+
client.check_connection().await.unwrap();
156+
}
157+
152158
#[tokio::test]
153159
async fn pipelined_prepare() {
154160
let client = connect("user=postgres").await;

0 commit comments

Comments
 (0)