diff --git a/postgres/src/client.rs b/postgres/src/client.rs index e3c9efaa..2621cbf7 100644 --- a/postgres/src/client.rs +++ b/postgres/src/client.rs @@ -1,9 +1,33 @@ use crate::{Config, CopyInWriter, CopyOutReader, RowIter, Statement, ToStatement, Transaction}; +use std::ops::{Deref, DerefMut}; use tokio::runtime::Runtime; use tokio_postgres::tls::{MakeTlsConnect, TlsConnect}; use tokio_postgres::types::{ToSql, Type}; use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket}; +pub(crate) struct Rt<'a>(pub &'a mut Runtime); + +// no-op impl to extend the borrow until drop +impl Drop for Rt<'_> { + fn drop(&mut self) {} +} + +impl Deref for Rt<'_> { + type Target = Runtime; + + #[inline] + fn deref(&self) -> &Runtime { + self.0 + } +} + +impl DerefMut for Rt<'_> { + #[inline] + fn deref_mut(&mut self) -> &mut Runtime { + self.0 + } +} + /// A synchronous PostgreSQL client. pub struct Client { runtime: Runtime, @@ -38,6 +62,10 @@ impl Client { Config::new() } + fn rt(&mut self) -> Rt<'_> { + Rt(&mut self.runtime) + } + /// Executes a statement, returning the number of rows modified. /// /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list @@ -236,7 +264,7 @@ impl Client { let stream = self .runtime .block_on(self.client.query_raw(query, params))?; - Ok(RowIter::new(&mut self.runtime, stream)) + Ok(RowIter::new(self.rt(), stream)) } /// Creates a new prepared statement. @@ -326,7 +354,7 @@ impl Client { T: ?Sized + ToStatement, { let sink = self.runtime.block_on(self.client.copy_in(query))?; - Ok(CopyInWriter::new(&mut self.runtime, sink)) + Ok(CopyInWriter::new(self.rt(), sink)) } /// Executes a `COPY TO STDOUT` statement, returning a reader of the resulting data. @@ -354,7 +382,7 @@ impl Client { T: ?Sized + ToStatement, { let stream = self.runtime.block_on(self.client.copy_out(query))?; - CopyOutReader::new(&mut self.runtime, stream) + CopyOutReader::new(self.rt(), stream) } /// Executes a sequence of SQL statements using the simple query protocol. diff --git a/postgres/src/copy_in_writer.rs b/postgres/src/copy_in_writer.rs index 897d8756..9a9e4899 100644 --- a/postgres/src/copy_in_writer.rs +++ b/postgres/src/copy_in_writer.rs @@ -1,27 +1,22 @@ +use crate::Rt; use bytes::{Bytes, BytesMut}; use futures::SinkExt; use std::io; use std::io::Write; use std::pin::Pin; -use tokio::runtime::Runtime; use tokio_postgres::{CopyInSink, Error}; /// The writer returned by the `copy_in` method. /// /// The copy *must* be explicitly completed via the `finish` method. If it is not, the copy will be aborted. pub struct CopyInWriter<'a> { - runtime: &'a mut Runtime, + runtime: Rt<'a>, sink: Pin>>, buf: BytesMut, } -// no-op impl to extend borrow until drop -impl Drop for CopyInWriter<'_> { - fn drop(&mut self) {} -} - impl<'a> CopyInWriter<'a> { - pub(crate) fn new(runtime: &'a mut Runtime, sink: CopyInSink) -> CopyInWriter<'a> { + pub(crate) fn new(runtime: Rt<'a>, sink: CopyInSink) -> CopyInWriter<'a> { CopyInWriter { runtime, sink: Box::pin(sink), diff --git a/postgres/src/copy_out_reader.rs b/postgres/src/copy_out_reader.rs index 14f8d630..c1bc2ff6 100644 --- a/postgres/src/copy_out_reader.rs +++ b/postgres/src/copy_out_reader.rs @@ -1,25 +1,20 @@ +use crate::Rt; use bytes::{Buf, Bytes}; use futures::StreamExt; use std::io::{self, BufRead, Cursor, Read}; use std::pin::Pin; -use tokio::runtime::Runtime; use tokio_postgres::{CopyOutStream, Error}; /// The reader returned by the `copy_out` method. pub struct CopyOutReader<'a> { - runtime: &'a mut Runtime, + runtime: Rt<'a>, stream: Pin>, cur: Cursor, } -// no-op impl to extend borrow until drop -impl Drop for CopyOutReader<'_> { - fn drop(&mut self) {} -} - impl<'a> CopyOutReader<'a> { pub(crate) fn new( - runtime: &'a mut Runtime, + mut runtime: Rt<'a>, stream: CopyOutStream, ) -> Result, Error> { let mut stream = Box::pin(stream); diff --git a/postgres/src/row_iter.rs b/postgres/src/row_iter.rs index 4bd8e4d8..4be5f347 100644 --- a/postgres/src/row_iter.rs +++ b/postgres/src/row_iter.rs @@ -1,12 +1,12 @@ +use crate::Rt; use fallible_iterator::FallibleIterator; use futures::StreamExt; use std::pin::Pin; -use tokio::runtime::Runtime; use tokio_postgres::{Error, Row, RowStream}; /// The iterator returned by `query_raw`. pub struct RowIter<'a> { - runtime: &'a mut Runtime, + runtime: Rt<'a>, it: Pin>, } @@ -16,7 +16,7 @@ impl Drop for RowIter<'_> { } impl<'a> RowIter<'a> { - pub(crate) fn new(runtime: &'a mut Runtime, stream: RowStream) -> RowIter<'a> { + pub(crate) fn new(runtime: Rt<'a>, stream: RowStream) -> RowIter<'a> { RowIter { runtime, it: Box::pin(stream), diff --git a/postgres/src/transaction.rs b/postgres/src/transaction.rs index ece8de99..2b397ca5 100644 --- a/postgres/src/transaction.rs +++ b/postgres/src/transaction.rs @@ -1,4 +1,4 @@ -use crate::{CopyInWriter, CopyOutReader, Portal, RowIter, Statement, ToStatement}; +use crate::{CopyInWriter, CopyOutReader, Portal, RowIter, Rt, Statement, ToStatement}; use tokio::runtime::Runtime; use tokio_postgres::types::{ToSql, Type}; use tokio_postgres::{Error, Row, SimpleQueryMessage}; @@ -23,6 +23,10 @@ impl<'a> Transaction<'a> { } } + fn rt(&mut self) -> Rt<'_> { + Rt(self.runtime) + } + /// Consumes the transaction, committing all changes made within it. pub fn commit(self) -> Result<(), Error> { self.runtime.block_on(self.transaction.commit()) @@ -95,7 +99,7 @@ impl<'a> Transaction<'a> { let stream = self .runtime .block_on(self.transaction.query_raw(query, params))?; - Ok(RowIter::new(self.runtime, stream)) + Ok(RowIter::new(self.rt(), stream)) } /// Binds parameters to a statement, creating a "portal". @@ -133,7 +137,7 @@ impl<'a> Transaction<'a> { let stream = self .runtime .block_on(self.transaction.query_portal_raw(portal, max_rows))?; - Ok(RowIter::new(self.runtime, stream)) + Ok(RowIter::new(self.rt(), stream)) } /// Like `Client::copy_in`. @@ -142,7 +146,7 @@ impl<'a> Transaction<'a> { T: ?Sized + ToStatement, { let sink = self.runtime.block_on(self.transaction.copy_in(query))?; - Ok(CopyInWriter::new(self.runtime, sink)) + Ok(CopyInWriter::new(self.rt(), sink)) } /// Like `Client::copy_out`. @@ -151,7 +155,7 @@ impl<'a> Transaction<'a> { T: ?Sized + ToStatement, { let stream = self.runtime.block_on(self.transaction.copy_out(query))?; - CopyOutReader::new(self.runtime, stream) + CopyOutReader::new(self.rt(), stream) } /// Like `Client::simple_query`.