copy in support

Author: sfackler

postgres-shared/src/error/mod.rs

@@ -325,6 +325,14 @@ pub fn __db(e: ErrorResponseBody) -> Error {
     }
 }
 
+#[doc(hidden)]
+pub fn __user<T>(e: T) -> Error
+where
+    T: Into<Box<error::Error + Sync + Send>>,
+{
+    Error(Box::new(ErrorKind::Conversion(e.into())))
+}
+
 #[doc(hidden)]
 pub fn io(e: io::Error) -> Error {
     Error(Box::new(ErrorKind::Io(e)))

tokio-postgres/src/lib.rs

@@ -24,6 +24,7 @@ extern crate tokio_uds;
 use bytes::Bytes;
 use futures::{Async, Future, Poll, Stream};
 use postgres_shared::rows::RowIndex;
+use std::error::Error as StdError;
 use std::fmt;
 use std::io;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -96,6 +97,14 @@ impl Client {
         Query(self.0.query(&statement.0, params))
     }
 
+    pub fn copy_in<S>(&mut self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyIn<S>
+    where
+        S: Stream<Item = Vec<u8>>,
+        S::Error: Into<Box<StdError + Sync + Send>>,
+    {
+        CopyIn(self.0.copy_in(&statement.0, params, stream))
+    }
+
     pub fn copy_out(&mut self, statement: &Statement, params: &[&ToSql]) -> CopyOut {
         CopyOut(self.0.copy_out(&statement.0, params))
     }
@@ -227,6 +236,25 @@ impl Stream for Query {
     }
 }
 
+#[must_use = "futures do nothing unless polled"]
+pub struct CopyIn<S>(proto::CopyInFuture<S>)
+where
+    S: Stream<Item = Vec<u8>>,
+    S::Error: Into<Box<StdError + Sync + Send>>;
+
+impl<S> Future for CopyIn<S>
+where
+    S: Stream<Item = Vec<u8>>,
+    S::Error: Into<Box<StdError + Sync + Send>>,
+{
+    type Item = u64;
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<u64, Error> {
+        self.0.poll()
+    }
+}
+
 #[must_use = "streams do nothing unless polled"]
 pub struct CopyOut(proto::CopyOutStream);
 

tokio-postgres/src/proto/client.rs

@@ -1,14 +1,17 @@
 use antidote::Mutex;
 use futures::sync::mpsc;
+use futures::{AsyncSink, Sink, Stream};
 use postgres_protocol;
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use std::collections::HashMap;
+use std::error::Error as StdError;
 use std::sync::{Arc, Weak};
 
 use disconnected;
 use error::{self, Error};
-use proto::connection::Request;
+use proto::connection::{Request, RequestMessages};
+use proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
 use proto::copy_out::CopyOutStream;
 use proto::execute::ExecuteFuture;
 use proto::prepare::PrepareFuture;
@@ -17,7 +20,7 @@ use proto::simple_query::SimpleQueryFuture;
 use proto::statement::Statement;
 use types::{IsNull, Oid, ToSql, Type};
 
-pub struct PendingRequest(Result<Vec<u8>, Error>);
+pub struct PendingRequest(Result<RequestMessages, Error>);
 
 pub struct WeakClient(Weak<Inner>);
 
@@ -122,17 +125,45 @@ impl Client {
     }
 
     pub fn execute(&self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {
-        let pending = self.pending_execute(statement, params);
+        let pending = PendingRequest(
+            self.excecute_message(statement, params)
+                .map(RequestMessages::Single),
+        );
         ExecuteFuture::new(self.clone(), pending, statement.clone())
     }
 
     pub fn query(&self, statement: &Statement, params: &[&ToSql]) -> QueryStream {
-        let pending = self.pending_execute(statement, params);
+        let pending = PendingRequest(
+            self.excecute_message(statement, params)
+                .map(RequestMessages::Single),
+        );
         QueryStream::new(self.clone(), pending, statement.clone())
     }
 
+    pub fn copy_in<S>(&self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyInFuture<S>
+    where
+        S: Stream<Item = Vec<u8>>,
+        S::Error: Into<Box<StdError + Sync + Send>>,
+    {
+        let (mut sender, receiver) = mpsc::channel(0);
+        let pending = PendingRequest(self.excecute_message(statement, params).map(|buf| {
+            match sender.start_send(CopyMessage::Data(buf)) {
+                Ok(AsyncSink::Ready) => {}
+                _ => unreachable!("channel should have capacity"),
+            }
+            RequestMessages::CopyIn {
+                receiver: CopyInReceiver::new(receiver),
+                pending_message: None,
+            }
+        }));
+        CopyInFuture::new(self.clone(), pending, statement.clone(), stream, sender)
+    }
+
     pub fn copy_out(&self, statement: &Statement, params: &[&ToSql]) -> CopyOutStream {
-        let pending = self.pending_execute(statement, params);
+        let pending = PendingRequest(
+            self.excecute_message(statement, params)
+                .map(RequestMessages::Single),
+        );
         CopyOutStream::new(self.clone(), pending, statement.clone())
     }
 
@@ -142,42 +173,41 @@ impl Client {
         frontend::sync(&mut buf);
         let (sender, _) = mpsc::channel(0);
         let _ = self.0.sender.unbounded_send(Request {
-            messages: buf,
+            messages: RequestMessages::Single(buf),
             sender,
         });
     }
 
-    fn pending_execute(&self, statement: &Statement, params: &[&ToSql]) -> PendingRequest {
-        self.pending(|buf| {
-            let r = frontend::bind(
-                "",
-                statement.name(),
-                Some(1),
-                params.iter().zip(statement.params()),
-                |(param, ty), buf| match param.to_sql_checked(ty, buf) {
-                    Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
-                    Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
-                    Err(e) => Err(e),
-                },
-                Some(1),
-                buf,
-            );
-            match r {
-                Ok(()) => {}
-                Err(frontend::BindError::Conversion(e)) => return Err(error::conversion(e)),
-                Err(frontend::BindError::Serialization(e)) => return Err(Error::from(e)),
-            }
-            frontend::execute("", 0, buf)?;
-            frontend::sync(buf);
-            Ok(())
-        })
+    fn excecute_message(&self, statement: &Statement, params: &[&ToSql]) -> Result<Vec<u8>, Error> {
+        let mut buf = vec![];
+        let r = frontend::bind(
+            "",
+            statement.name(),
+            Some(1),
+            params.iter().zip(statement.params()),
+            |(param, ty), buf| match param.to_sql_checked(ty, buf) {
+                Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
+                Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
+                Err(e) => Err(e),
+            },
+            Some(1),
+            &mut buf,
+        );
+        match r {
+            Ok(()) => {}
+            Err(frontend::BindError::Conversion(e)) => return Err(error::conversion(e)),
+            Err(frontend::BindError::Serialization(e)) => return Err(Error::from(e)),
+        }
+        frontend::execute("", 0, &mut buf)?;
+        frontend::sync(&mut buf);
+        Ok(buf)
     }
 
     fn pending<F>(&self, messages: F) -> PendingRequest
     where
         F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
     {
         let mut buf = vec![];
-        PendingRequest(messages(&mut buf).map(|()| buf))
+        PendingRequest(messages(&mut buf).map(|()| RequestMessages::Single(buf)))
     }
 }

tokio-postgres/src/proto/connection.rs

@@ -8,11 +8,20 @@ use tokio_codec::Framed;
 
 use error::{self, DbError, Error};
 use proto::codec::PostgresCodec;
+use proto::copy_in::CopyInReceiver;
 use tls::TlsStream;
 use {bad_response, disconnected, AsyncMessage, CancelData, Notification};
 
+pub enum RequestMessages {
+    Single(Vec<u8>),
+    CopyIn {
+        receiver: CopyInReceiver,
+        pending_message: Option<Vec<u8>>,
+    },
+}
+
 pub struct Request {
-    pub messages: Vec<u8>,
+    pub messages: RequestMessages,
     pub sender: mpsc::Sender<Message>,
 }
 
@@ -28,7 +37,7 @@ pub struct Connection {
     cancel_data: CancelData,
     parameters: HashMap<String, String>,
     receiver: mpsc::UnboundedReceiver<Request>,
-    pending_request: Option<Vec<u8>>,
+    pending_request: Option<RequestMessages>,
     pending_response: Option<Message>,
     responses: VecDeque<mpsc::Sender<Message>>,
     state: State,
@@ -140,7 +149,7 @@ impl Connection {
         }
     }
 
-    fn poll_request(&mut self) -> Poll<Option<Vec<u8>>, Error> {
+    fn poll_request(&mut self) -> Poll<Option<RequestMessages>, Error> {
         if let Some(message) = self.pending_request.take() {
             trace!("retrying pending request");
             return Ok(Async::Ready(Some(message)));
@@ -170,7 +179,7 @@ impl Connection {
                     self.state = State::Terminating;
                     let mut request = vec![];
                     frontend::terminate(&mut request);
-                    request
+                    RequestMessages::Single(request)
                 }
                 Async::Ready(None) => {
                     trace!(
@@ -185,17 +194,60 @@ impl Connection {
                 }
             };
 
-            match self.stream.start_send(request)? {
-                AsyncSink::Ready => {
-                    if self.state == State::Terminating {
-                        trace!("poll_write: sent eof, closing");
-                        self.state = State::Closing;
+            match request {
+                RequestMessages::Single(request) => match self.stream.start_send(request)? {
+                    AsyncSink::Ready => {
+                        if self.state == State::Terminating {
+                            trace!("poll_write: sent eof, closing");
+                            self.state = State::Closing;
+                        }
                     }
-                }
-                AsyncSink::NotReady(request) => {
-                    trace!("poll_write: waiting on socket");
-                    self.pending_request = Some(request);
-                    return Ok(false);
+                    AsyncSink::NotReady(request) => {
+                        trace!("poll_write: waiting on socket");
+                        self.pending_request = Some(RequestMessages::Single(request));
+                        return Ok(false);
+                    }
+                },
+                RequestMessages::CopyIn {
+                    mut receiver,
+                    mut pending_message,
+                } => {
+                    let message = match pending_message.take() {
+                        Some(message) => message,
+                        None => match receiver.poll() {
+                            Ok(Async::Ready(Some(message))) => message,
+                            Ok(Async::Ready(None)) => {
+                                trace!("poll_write: finished copy_in request");
+                                continue;
+                            }
+                            Ok(Async::NotReady) => {
+                                trace!("poll_write: waiting on copy_in stream");
+                                self.pending_request = Some(RequestMessages::CopyIn {
+                                    receiver,
+                                    pending_message,
+                                });
+                                return Ok(true);
+                            }
+                            Err(()) => unreachable!("mpsc::Receiver doesn't return errors"),
+                        },
+                    };
+
+                    match self.stream.start_send(message)? {
+                        AsyncSink::Ready => {
+                            self.pending_request = Some(RequestMessages::CopyIn {
+                                receiver,
+                                pending_message: None,
+                            });
+                        }
+                        AsyncSink::NotReady(message) => {
+                            trace!("poll_write: waiting on socket");
+                            self.pending_request = Some(RequestMessages::CopyIn {
+                                receiver,
+                                pending_message: Some(message),
+                            });
+                            return Ok(false);
+                        }
+                    };
                 }
             }
         }