Skip to content

Commit bc682b3

Browse files
committed
Explicitly terminate the connection in sync API
Closes #613
1 parent f6620e6 commit bc682b3

File tree

3 files changed

+38
-0
lines changed

3 files changed

+38
-0
lines changed

postgres/src/client.rs

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::{
33
CancelToken, Config, CopyInWriter, CopyOutReader, Notifications, RowIter, Statement,
44
ToStatement, Transaction, TransactionBuilder,
55
};
6+
use std::task::Poll;
67
use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
78
use tokio_postgres::types::{ToSql, Type};
89
use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket};
@@ -13,6 +14,12 @@ pub struct Client {
1314
client: tokio_postgres::Client,
1415
}
1516

17+
impl Drop for Client {
18+
fn drop(&mut self) {
19+
let _ = self.close_inner();
20+
}
21+
}
22+
1623
impl Client {
1724
pub(crate) fn new(connection: Connection, client: tokio_postgres::Client) -> Client {
1825
Client { connection, client }
@@ -524,4 +531,24 @@ impl Client {
524531
pub fn is_closed(&self) -> bool {
525532
self.client.is_closed()
526533
}
534+
535+
/// Closes the client's connection to the server.
536+
///
537+
/// This is equivalent to `Client`'s `Drop` implementation, except that it returns any error encountered to the
538+
/// caller.
539+
pub fn close(mut self) -> Result<(), Error> {
540+
self.close_inner()
541+
}
542+
543+
fn close_inner(&mut self) -> Result<(), Error> {
544+
self.client.__private_api_close();
545+
546+
self.connection.poll_block_on(|_, _, done| {
547+
if done {
548+
Poll::Ready(Ok(()))
549+
} else {
550+
Poll::Pending
551+
}
552+
})
553+
}
527554
}

postgres/src/test.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -475,3 +475,9 @@ fn notifications_timeout_iter() {
475475
assert_eq!(notifications[0].payload(), "hello");
476476
assert_eq!(notifications[1].payload(), "world");
477477
}
478+
479+
#[test]
480+
fn explicit_close() {
481+
let client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
482+
client.close().unwrap();
483+
}

tokio-postgres/src/client.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,11 @@ impl Client {
529529
pub fn is_closed(&self) -> bool {
530530
self.inner.sender.is_closed()
531531
}
532+
533+
#[doc(hidden)]
534+
pub fn __private_api_close(&mut self) {
535+
self.inner.sender.close_channel()
536+
}
532537
}
533538

534539
impl fmt::Debug for Client {

0 commit comments

Comments
 (0)