Blocking binary copy support

Author: sfackler

postgres/src/binary_copy.rs

@@ -0,0 +1,92 @@
+//! Utilities for working with the PostgreSQL binary copy format.
+
+use crate::types::{ToSql, Type};
+use crate::{CopyInWriter, CopyOutReader, Error, Rt};
+use fallible_iterator::FallibleIterator;
+use futures::StreamExt;
+use std::pin::Pin;
+#[doc(inline)]
+pub use tokio_postgres::binary_copy::BinaryCopyOutRow;
+use tokio_postgres::binary_copy::{self, BinaryCopyOutStream};
+
+/// A type which serializes rows into the PostgreSQL binary copy format.
+///
+/// The copy *must* be explicitly completed via the `finish` method. If it is not, the copy will be aborted.
+pub struct BinaryCopyInWriter<'a> {
+    runtime: Rt<'a>,
+    sink: Pin<Box<binary_copy::BinaryCopyInWriter>>,
+}
+
+impl<'a> BinaryCopyInWriter<'a> {
+    /// Creates a new writer which will write rows of the provided types.
+    pub fn new(writer: CopyInWriter<'a>, types: &[Type]) -> BinaryCopyInWriter<'a> {
+        let stream = writer
+            .sink
+            .into_unpinned()
+            .expect("writer has already been written to");
+
+        BinaryCopyInWriter {
+            runtime: writer.runtime,
+            sink: Box::pin(binary_copy::BinaryCopyInWriter::new(stream, types)),
+        }
+    }
+
+    /// Writes a single row.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of values provided does not match the number expected.
+    pub fn write(&mut self, values: &[&(dyn ToSql + Sync)]) -> Result<(), Error> {
+        self.runtime.block_on(self.sink.as_mut().write(values))
+    }
+
+    /// A maximally-flexible version of `write`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of values provided does not match the number expected.
+    pub fn write_raw<'b, I>(&mut self, values: I) -> Result<(), Error>
+    where
+        I: IntoIterator<Item = &'b dyn ToSql>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.runtime.block_on(self.sink.as_mut().write_raw(values))
+    }
+
+    /// Completes the copy, returning the number of rows added.
+    ///
+    /// This method *must* be used to complete the copy process. If it is not, the copy will be aborted.
+    pub fn finish(mut self) -> Result<u64, Error> {
+        self.runtime.block_on(self.sink.as_mut().finish())
+    }
+}
+
+/// An iterator of rows deserialized from the PostgreSQL binary copy format.
+pub struct BinaryCopyOutIter<'a> {
+    runtime: Rt<'a>,
+    stream: Pin<Box<BinaryCopyOutStream>>,
+}
+
+impl<'a> BinaryCopyOutIter<'a> {
+    /// Creates a new iterator from a raw copy out reader and the types of the columns being returned.
+    pub fn new(reader: CopyOutReader<'a>, types: &[Type]) -> BinaryCopyOutIter<'a> {
+        let stream = reader
+            .stream
+            .into_unpinned()
+            .expect("reader has already been read from");
+
+        BinaryCopyOutIter {
+            runtime: reader.runtime,
+            stream: Box::pin(BinaryCopyOutStream::new(stream, types)),
+        }
+    }
+}
+
+impl FallibleIterator for BinaryCopyOutIter<'_> {
+    type Item = BinaryCopyOutRow;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<BinaryCopyOutRow>, Error> {
+        self.runtime.block_on(self.stream.next()).transpose()
+    }
+}

postgres/src/client.rs

@@ -382,7 +382,7 @@ impl Client {
         T: ?Sized + ToStatement,
     {
         let stream = self.runtime.block_on(self.client.copy_out(query))?;
-        CopyOutReader::new(self.rt(), stream)
+        Ok(CopyOutReader::new(self.rt(), stream))
     }
 
     /// Executes a sequence of SQL statements using the simple query protocol.

postgres/src/copy_in_writer.rs

@@ -1,25 +1,25 @@
+use crate::lazy_pin::LazyPin;
 use crate::Rt;
 use bytes::{Bytes, BytesMut};
 use futures::SinkExt;
 use std::io;
 use std::io::Write;
-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> {
-    runtime: Rt<'a>,
-    sink: Pin<Box<CopyInSink<Bytes>>>,
+    pub(crate) runtime: Rt<'a>,
+    pub(crate) sink: LazyPin<CopyInSink<Bytes>>,
     buf: BytesMut,
 }
 
 impl<'a> CopyInWriter<'a> {
     pub(crate) fn new(runtime: Rt<'a>, sink: CopyInSink<Bytes>) -> CopyInWriter<'a> {
         CopyInWriter {
             runtime,
-            sink: Box::pin(sink),
+            sink: LazyPin::new(sink),
             buf: BytesMut::new(),
         }
     }
@@ -29,7 +29,7 @@ impl<'a> CopyInWriter<'a> {
     /// If this is not called, the copy will be aborted.
     pub fn finish(mut self) -> Result<u64, Error> {
         self.flush_inner()?;
-        self.runtime.block_on(self.sink.as_mut().finish())
+        self.runtime.block_on(self.sink.pinned().finish())
     }
 
     fn flush_inner(&mut self) -> Result<(), Error> {
@@ -38,7 +38,7 @@ impl<'a> CopyInWriter<'a> {
         }
 
         self.runtime
-            .block_on(self.sink.as_mut().send(self.buf.split().freeze()))
+            .block_on(self.sink.pinned().send(self.buf.split().freeze()))
     }
 }
 

postgres/src/copy_out_reader.rs

@@ -1,34 +1,24 @@
+use crate::lazy_pin::LazyPin;
 use crate::Rt;
 use bytes::{Buf, Bytes};
 use futures::StreamExt;
-use std::io::{self, BufRead, Cursor, Read};
-use std::pin::Pin;
-use tokio_postgres::{CopyOutStream, Error};
+use std::io::{self, BufRead, Read};
+use tokio_postgres::CopyOutStream;
 
 /// The reader returned by the `copy_out` method.
 pub struct CopyOutReader<'a> {
-    runtime: Rt<'a>,
-    stream: Pin<Box<CopyOutStream>>,
-    cur: Cursor<Bytes>,
+    pub(crate) runtime: Rt<'a>,
+    pub(crate) stream: LazyPin<CopyOutStream>,
+    cur: Bytes,
 }
 
 impl<'a> CopyOutReader<'a> {
-    pub(crate) fn new(
-        mut runtime: Rt<'a>,
-        stream: CopyOutStream,
-    ) -> Result<CopyOutReader<'a>, Error> {
-        let mut stream = Box::pin(stream);
-        let cur = match runtime.block_on(stream.next()) {
-            Some(Ok(cur)) => cur,
-            Some(Err(e)) => return Err(e),
-            None => Bytes::new(),
-        };
-
-        Ok(CopyOutReader {
+    pub(crate) fn new(runtime: Rt<'a>, stream: CopyOutStream) -> CopyOutReader<'a> {
+        CopyOutReader {
             runtime,
-            stream,
-            cur: Cursor::new(cur),
-        })
+            stream: LazyPin::new(stream),
+            cur: Bytes::new(),
+        }
     }
 }
 
@@ -44,15 +34,15 @@ impl Read for CopyOutReader<'_> {
 
 impl BufRead for CopyOutReader<'_> {
     fn fill_buf(&mut self) -> io::Result<&[u8]> {
-        if self.cur.remaining() == 0 {
-            match self.runtime.block_on(self.stream.next()) {
-                Some(Ok(cur)) => self.cur = Cursor::new(cur),
+        if !self.cur.has_remaining() {
+            match self.runtime.block_on(self.stream.pinned().next()) {
+                Some(Ok(cur)) => self.cur = cur,
                 Some(Err(e)) => return Err(io::Error::new(io::ErrorKind::Other, e)),
                 None => {}
             };
         }
 
-        Ok(Buf::bytes(&self.cur))
+        Ok(self.cur.bytes())
     }
 
     fn consume(&mut self, amt: usize) {

postgres/src/lazy_pin.rs

@@ -0,0 +1,28 @@
+use std::pin::Pin;
+
+pub(crate) struct LazyPin<T> {
+    value: Box<T>,
+    pinned: bool,
+}
+
+impl<T> LazyPin<T> {
+    pub fn new(value: T) -> LazyPin<T> {
+        LazyPin {
+            value: Box::new(value),
+            pinned: false,
+        }
+    }
+
+    pub fn pinned(&mut self) -> Pin<&mut T> {
+        self.pinned = true;
+        unsafe { Pin::new_unchecked(&mut *self.value) }
+    }
+
+    pub fn into_unpinned(self) -> Option<T> {
+        if self.pinned {
+            None
+        } else {
+            Some(*self.value)
+        }
+    }
+}

postgres/src/lib.rs

@@ -76,6 +76,8 @@ mod client;
 pub mod config;
 mod copy_in_writer;
 mod copy_out_reader;
+mod lazy_pin;
+pub mod binary_copy;
 mod row_iter;
 mod transaction;
 

postgres/src/test.rs

@@ -3,6 +3,8 @@ use tokio_postgres::types::Type;
 use tokio_postgres::NoTls;
 
 use super::*;
+use crate::binary_copy::{BinaryCopyInWriter, BinaryCopyOutIter};
+use fallible_iterator::FallibleIterator;
 
 #[test]
 fn prepare() {
@@ -188,6 +190,31 @@ fn copy_in_abort() {
     assert_eq!(rows.len(), 0);
 }
 
+#[test]
+fn binary_copy_in() {
+    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 writer = client.copy_in("COPY foo FROM stdin BINARY").unwrap();
+    let mut writer = BinaryCopyInWriter::new(writer, &[Type::INT4, Type::TEXT]);
+    writer.write(&[&1i32, &"steven"]).unwrap();
+    writer.write(&[&2i32, &"timothy"]).unwrap();
+    writer.finish().unwrap();
+
+    let rows = client
+        .query("SELECT id, name FROM foo ORDER BY id", &[])
+        .unwrap();
+
+    assert_eq!(rows.len(), 2);
+    assert_eq!(rows[0].get::<_, i32>(0), 1);
+    assert_eq!(rows[0].get::<_, &str>(1), "steven");
+    assert_eq!(rows[1].get::<_, i32>(0), 2);
+    assert_eq!(rows[1].get::<_, &str>(1), "timothy");
+}
+
 #[test]
 fn copy_out() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
@@ -209,6 +236,32 @@ fn copy_out() {
     client.simple_query("SELECT 1").unwrap();
 }
 
+#[test]
+fn binary_copy_out() {
+    let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
+
+    client
+        .simple_query(
+            "CREATE TEMPORARY TABLE foo (id INT, name TEXT);
+             INSERT INTO foo (id, name) VALUES (1, 'steven'), (2, 'timothy');",
+        )
+        .unwrap();
+
+    let reader = client
+        .copy_out("COPY foo (id, name) TO STDOUT BINARY")
+        .unwrap();
+    let rows = BinaryCopyOutIter::new(reader, &[Type::INT4, Type::TEXT])
+        .collect::<Vec<_>>()
+        .unwrap();
+    assert_eq!(rows.len(), 2);
+    assert_eq!(rows[0].get::<i32>(0), 1);
+    assert_eq!(rows[0].get::<&str>(1), "steven");
+    assert_eq!(rows[1].get::<i32>(0), 2);
+    assert_eq!(rows[1].get::<&str>(1), "timothy");
+
+    client.simple_query("SELECT 1").unwrap();
+}
+
 #[test]
 fn portal() {
     let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();

postgres/src/transaction.rs

@@ -155,7 +155,7 @@ impl<'a> Transaction<'a> {
         T: ?Sized + ToStatement,
     {
         let stream = self.runtime.block_on(self.transaction.copy_out(query))?;
-        CopyOutReader::new(self.rt(), stream)
+        Ok(CopyOutReader::new(self.rt(), stream))
     }
 
     /// Like `Client::simple_query`.

tokio-postgres/src/copy_out.rs

@@ -22,11 +22,13 @@ pub async fn copy_out(client: &InnerClient, statement: Statement) -> Result<Copy
 async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
+    println!("a");
     match responses.next().await? {
         Message::BindComplete => {}
         _ => return Err(Error::unexpected_message()),
     }
 
+    println!("b");
     match responses.next().await? {
         Message::CopyOutResponse(_) => {}
         _ => return Err(Error::unexpected_message()),
@@ -50,6 +52,7 @@ impl Stream for CopyOutStream {
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.project();
 
+        println!("c");
         match ready!(this.responses.poll_next(cx)?) {
             Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
             Message::CopyDone => Poll::Ready(None),