diff --git a/postgres/src/client.rs b/postgres/src/client.rs index 4ccb9b5a..f4b9ef3a 100644 --- a/postgres/src/client.rs +++ b/postgres/src/client.rs @@ -1,8 +1,5 @@ -use bytes::{Buf, Bytes}; -use futures::stream; use futures::{Async, Future, Poll, Stream}; -use std::io::{self, BufRead, Cursor, Read}; -use std::marker::PhantomData; +use std::io::{self, Read}; use tokio_postgres::types::{ToSql, Type}; use tokio_postgres::Error; #[cfg(feature = "runtime")] @@ -10,7 +7,7 @@ use tokio_postgres::{MakeTlsMode, Socket, TlsMode}; #[cfg(feature = "runtime")] use crate::Builder; -use crate::{Query, Statement, ToStatement, Transaction}; +use crate::{CopyOutReader, Query, Statement, ToStatement, Transaction}; pub struct Client(tokio_postgres::Client); @@ -81,19 +78,8 @@ impl Client { T: ?Sized + ToStatement, { let statement = query.__statement(self)?; - let mut stream = self.0.copy_out(&statement.0, params).wait(); - - let cur = match stream.next() { - Some(Ok(cur)) => cur, - Some(Err(e)) => return Err(e), - None => Bytes::new(), - }; - - Ok(CopyOutReader { - stream, - cur: Cursor::new(cur), - _p: PhantomData, - }) + let stream = self.0.copy_out(&statement.0, params); + CopyOutReader::new(stream) } pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> { @@ -145,42 +131,3 @@ where } } } - -pub struct CopyOutReader<'a> { - stream: stream::Wait, - cur: Cursor, - _p: PhantomData<&'a mut ()>, -} - -// no-op impl to extend borrow until drop -impl<'a> Drop for CopyOutReader<'a> { - fn drop(&mut self) {} -} - -impl<'a> Read for CopyOutReader<'a> { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let b = self.fill_buf()?; - let len = usize::min(buf.len(), b.len()); - buf[..len].copy_from_slice(&b[..len]); - self.consume(len); - Ok(len) - } -} - -impl<'a> BufRead for CopyOutReader<'a> { - fn fill_buf(&mut self) -> io::Result<&[u8]> { - if self.cur.remaining() == 0 { - match self.stream.next() { - Some(Ok(cur)) => self.cur = Cursor::new(cur), - Some(Err(e)) => return Err(io::Error::new(io::ErrorKind::Other, e)), - None => {} - }; - } - - Ok(Buf::bytes(&self.cur)) - } - - fn consume(&mut self, amt: usize) { - self.cur.advance(amt); - } -} diff --git a/postgres/src/copy_out_reader.rs b/postgres/src/copy_out_reader.rs new file mode 100644 index 00000000..b8244f26 --- /dev/null +++ b/postgres/src/copy_out_reader.rs @@ -0,0 +1,61 @@ +use bytes::{Buf, Bytes}; +use futures::stream::{self, Stream}; +use std::io::{self, BufRead, Cursor, Read}; +use std::marker::PhantomData; +use tokio_postgres::Error; + +pub struct CopyOutReader<'a> { + it: stream::Wait, + cur: Cursor, + _p: PhantomData<&'a mut ()>, +} + +// no-op impl to extend borrow until drop +impl<'a> Drop for CopyOutReader<'a> { + fn drop(&mut self) {} +} + +impl<'a> CopyOutReader<'a> { + pub(crate) fn new(stream: tokio_postgres::CopyOut) -> Result, Error> { + let mut it = stream.wait(); + let cur = match it.next() { + Some(Ok(cur)) => cur, + Some(Err(e)) => return Err(e), + None => Bytes::new(), + }; + + Ok(CopyOutReader { + it, + cur: Cursor::new(cur), + _p: PhantomData, + }) + } +} + +impl<'a> Read for CopyOutReader<'a> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let b = self.fill_buf()?; + let len = usize::min(buf.len(), b.len()); + buf[..len].copy_from_slice(&b[..len]); + self.consume(len); + Ok(len) + } +} + +impl<'a> BufRead for CopyOutReader<'a> { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + if self.cur.remaining() == 0 { + match self.it.next() { + Some(Ok(cur)) => self.cur = Cursor::new(cur), + Some(Err(e)) => return Err(io::Error::new(io::ErrorKind::Other, e)), + None => {} + }; + } + + Ok(Buf::bytes(&self.cur)) + } + + fn consume(&mut self, amt: usize) { + self.cur.advance(amt); + } +} diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 638ef26c..ece72ca0 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -6,6 +6,7 @@ use tokio::runtime::{self, Runtime}; #[cfg(feature = "runtime")] mod builder; mod client; +mod copy_out_reader; mod portal; mod query; mod query_portal; @@ -20,6 +21,7 @@ mod test; #[cfg(feature = "runtime")] pub use crate::builder::*; pub use crate::client::*; +pub use crate::copy_out_reader::*; pub use crate::portal::*; pub use crate::query::*; pub use crate::query_portal::*;