Change the copy_in interface

Rather than taking in a Stream and advancing it internally, return a
Sink that can be advanced by the calling code. This significantly
simplifies encoding logic for things like tokio-postgres-binary-copy.

Similarly, the blocking interface returns a Writer.

Closes sfackler#489

Author: sfackler

postgres/src/client.rs

@@ -1,19 +1,15 @@
+use crate::iter::Iter;
+#[cfg(feature = "runtime")]
+use crate::Config;
+use crate::{CopyInWriter, CopyOutReader, Statement, ToStatement, Transaction};
 use fallible_iterator::FallibleIterator;
 use futures::executor;
-use std::io::{BufRead, Read};
 use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
 use tokio_postgres::types::{ToSql, Type};
 #[cfg(feature = "runtime")]
 use tokio_postgres::Socket;
 use tokio_postgres::{Error, Row, SimpleQueryMessage};
 
-use crate::copy_in_stream::CopyInStream;
-use crate::copy_out_reader::CopyOutReader;
-use crate::iter::Iter;
-#[cfg(feature = "runtime")]
-use crate::Config;
-use crate::{Statement, ToStatement, Transaction};
-
 /// A synchronous PostgreSQL client.
 ///
 /// This is a lightweight wrapper over the asynchronous tokio_postgres `Client`.
@@ -264,29 +260,33 @@ impl Client {
     /// The `query` argument can either be a `Statement`, or a raw query string. The data in the provided reader is
     /// passed along to the server verbatim; it is the caller's responsibility to ensure it uses the proper format.
     ///
+    /// The copy *must* be explicitly completed via the `finish` method. If it is not, the copy will be aborted.
+    ///
     /// # Examples
     ///
     /// ```no_run
     /// use postgres::{Client, NoTls};
+    /// use std::io::Write;
     ///
-    /// # fn main() -> Result<(), postgres::Error> {
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
     /// let mut client = Client::connect("host=localhost user=postgres", NoTls)?;
     ///
-    /// client.copy_in("COPY people FROM stdin", &[], &mut "1\tjohn\n2\tjane\n".as_bytes())?;
+    /// let mut writer = client.copy_in("COPY people FROM stdin", &[])?;
+    /// writer.write_all(b"1\tjohn\n2\tjane\n")?;
+    /// writer.finish()?;
     /// # Ok(())
     /// # }
     /// ```
-    pub fn copy_in<T, R>(
+    pub fn copy_in<T>(
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-        reader: R,
-    ) -> Result<u64, Error>
+    ) -> Result<CopyInWriter<'_>, Error>
     where
         T: ?Sized + ToStatement,
-        R: Read + Unpin,
     {
-        executor::block_on(self.0.copy_in(query, params, CopyInStream(reader)))
+        let sink = executor::block_on(self.0.copy_in(query, params))?;
+        Ok(CopyInWriter::new(sink))
     }
 
     /// Executes a `COPY TO STDOUT` statement, returning a reader of the resulting data.
@@ -312,7 +312,7 @@ impl Client {
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-    ) -> Result<impl BufRead, Error>
+    ) -> Result<CopyOutReader<'_>, Error>
     where
         T: ?Sized + ToStatement,
     {

postgres/src/copy_in_stream.rs

@@ -0,0 +1,63 @@
+use bytes::{Bytes, BytesMut};
+use futures::{executor, SinkExt};
+use std::io;
+use std::io::Write;
+use std::marker::PhantomData;
+use std::pin::Pin;
+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> {
+    sink: Pin<Box<CopyInSink<Bytes>>>,
+    buf: BytesMut,
+    _p: PhantomData<&'a mut ()>,
+}
+
+// no-op impl to extend borrow until drop
+impl Drop for CopyInWriter<'_> {
+    fn drop(&mut self) {}
+}
+
+impl<'a> CopyInWriter<'a> {
+    pub(crate) fn new(sink: CopyInSink<Bytes>) -> CopyInWriter<'a> {
+        CopyInWriter {
+            sink: Box::pin(sink),
+            buf: BytesMut::new(),
+            _p: PhantomData,
+        }
+    }
+
+    /// Completes the copy, returning the number of rows written.
+    ///
+    /// If this is not called, the copy will be aborted.
+    pub fn finish(mut self) -> Result<u64, Error> {
+        self.flush_inner()?;
+        executor::block_on(self.sink.as_mut().finish())
+    }
+
+    fn flush_inner(&mut self) -> Result<(), Error> {
+        if self.buf.is_empty() {
+            return Ok(());
+        }
+
+        executor::block_on(self.sink.as_mut().send(self.buf.split().freeze()))
+    }
+}
+
+impl Write for CopyInWriter<'_> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if self.buf.len() > 4096 {
+            self.flush()?;
+        }
+
+        self.buf.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.flush_inner()
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+    }
+}

postgres/src/copy_in_writer.rs

@@ -1,33 +1,24 @@
 use bytes::{Buf, Bytes};
-use futures::{executor, Stream};
+use futures::executor;
 use std::io::{self, BufRead, Cursor, Read};
 use std::marker::PhantomData;
 use std::pin::Pin;
-use tokio_postgres::Error;
+use tokio_postgres::{CopyStream, Error};
 
 /// The reader returned by the `copy_out` method.
-pub struct CopyOutReader<'a, S>
-where
-    S: Stream,
-{
-    it: executor::BlockingStream<Pin<Box<S>>>,
+pub struct CopyOutReader<'a> {
+    it: executor::BlockingStream<Pin<Box<CopyStream>>>,
     cur: Cursor<Bytes>,
     _p: PhantomData<&'a mut ()>,
 }
 
 // no-op impl to extend borrow until drop
-impl<'a, S> Drop for CopyOutReader<'a, S>
-where
-    S: Stream,
-{
+impl Drop for CopyOutReader<'_> {
     fn drop(&mut self) {}
 }
 
-impl<'a, S> CopyOutReader<'a, S>
-where
-    S: Stream<Item = Result<Bytes, Error>>,
-{
-    pub(crate) fn new(stream: S) -> Result<CopyOutReader<'a, S>, Error> {
+impl<'a> CopyOutReader<'a> {
+    pub(crate) fn new(stream: CopyStream) -> Result<CopyOutReader<'a>, Error> {
         let mut it = executor::block_on_stream(Box::pin(stream));
         let cur = match it.next() {
             Some(Ok(cur)) => cur,
@@ -43,10 +34,7 @@ where
     }
 }
 
-impl<'a, S> Read for CopyOutReader<'a, S>
-where
-    S: Stream<Item = Result<Bytes, Error>>,
-{
+impl Read for CopyOutReader<'_> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         let b = self.fill_buf()?;
         let len = usize::min(buf.len(), b.len());
@@ -56,10 +44,7 @@ where
     }
 }
 
-impl<'a, S> BufRead for CopyOutReader<'a, S>
-where
-    S: Stream<Item = Result<Bytes, Error>>,
-{
+impl BufRead for CopyOutReader<'_> {
     fn fill_buf(&mut self) -> io::Result<&[u8]> {
         if self.cur.remaining() == 0 {
             match self.it.next() {

postgres/src/copy_out_reader.rs

@@ -69,6 +69,8 @@ pub use tokio_postgres::{
 pub use crate::client::*;
 #[cfg(feature = "runtime")]
 pub use crate::config::Config;
+pub use crate::copy_in_writer::CopyInWriter;
+pub use crate::copy_out_reader::CopyOutReader;
 #[doc(no_inline)]
 pub use crate::error::Error;
 #[doc(no_inline)]
@@ -80,7 +82,7 @@ pub use crate::transaction::*;
 mod client;
 #[cfg(feature = "runtime")]
 pub mod config;
-mod copy_in_stream;
+mod copy_in_writer;
 mod copy_out_reader;
 mod iter;
 mod transaction;

postgres/src/lib.rs

@@ -1,4 +1,4 @@
-use std::io::Read;
+use std::io::{Read, Write};
 use tokio_postgres::types::Type;
 use tokio_postgres::NoTls;
 
@@ -154,13 +154,9 @@ fn copy_in() {
         .simple_query("CREATE TEMPORARY TABLE foo (id INT, name TEXT)")
         .unwrap();
 
-    client
-        .copy_in(
-            "COPY foo FROM stdin",
-            &[],
-            &mut &b"1\tsteven\n2\ttimothy"[..],
-        )
-        .unwrap();
+    let mut writer = client.copy_in("COPY foo FROM stdin", &[]).unwrap();
+    writer.write_all(b"1\tsteven\n2\ttimothy").unwrap();
+    writer.finish().unwrap();
 
     let rows = client
         .query("SELECT id, name FROM foo ORDER BY id", &[])
@@ -173,6 +169,25 @@ fn copy_in() {
     assert_eq!(rows[1].get::<_, &str>(1), "timothy");
 }
 
+#[test]
+fn copy_in_abort() {
+    let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+    client
+        .simple_query("CREATE TEMPORARY TABLE foo (id INT, name TEXT)")
+        .unwrap();
+
+    let mut writer = client.copy_in("COPY foo FROM stdin", &[]).unwrap();
+    writer.write_all(b"1\tsteven\n2\ttimothy").unwrap();
+    drop(writer);
+
+    let rows = client
+        .query("SELECT id, name FROM foo ORDER BY id", &[])
+        .unwrap();
+
+    assert_eq!(rows.len(), 0);
+}
+
 #[test]
 fn copy_out() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();

postgres/src/test.rs

@@ -1,14 +1,10 @@
+use crate::iter::Iter;
+use crate::{CopyInWriter, CopyOutReader, Portal, Statement, ToStatement};
 use fallible_iterator::FallibleIterator;
 use futures::executor;
-use std::io::{BufRead, Read};
 use tokio_postgres::types::{ToSql, Type};
 use tokio_postgres::{Error, Row, SimpleQueryMessage};
 
-use crate::copy_in_stream::CopyInStream;
-use crate::copy_out_reader::CopyOutReader;
-use crate::iter::Iter;
-use crate::{Portal, Statement, ToStatement};
-
 /// A representation of a PostgreSQL database transaction.
 ///
 /// Transactions will implicitly roll back by default when dropped. Use the `commit` method to commit the changes made
@@ -117,25 +113,24 @@ impl<'a> Transaction<'a> {
     }
 
     /// Like `Client::copy_in`.
-    pub fn copy_in<T, R>(
+    pub fn copy_in<T>(
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-        reader: R,
-    ) -> Result<u64, Error>
+    ) -> Result<CopyInWriter<'_>, Error>
     where
         T: ?Sized + ToStatement,
-        R: Read + Unpin,
     {
-        executor::block_on(self.0.copy_in(query, params, CopyInStream(reader)))
+        let sink = executor::block_on(self.0.copy_in(query, params))?;
+        Ok(CopyInWriter::new(sink))
     }
 
     /// Like `Client::copy_out`.
     pub fn copy_out<T>(
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-    ) -> Result<impl BufRead, Error>
+    ) -> Result<CopyOutReader<'_>, Error>
     where
         T: ?Sized + ToStatement,
     {

postgres/src/transaction.rs

@@ -8,7 +8,6 @@ edition = "2018"
 byteorder = "1.0"
 bytes = "0.5"
 futures = "0.3"
-parking_lot = "0.10"
 pin-project-lite = "0.1"
 tokio-postgres = { version = "=0.5.0-alpha.2", default-features = false, path = "../tokio-postgres" }