add copy_both_simple method

Signed-off-by: Petros Angelatos <petrosagg@gmail.com>

Author: petrosagg

tokio-postgres/src/client.rs

@@ -1,6 +1,7 @@
 use crate::codec::BackendMessages;
 use crate::config::{Host, SslMode};
 use crate::connection::{Request, RequestMessages};
+use crate::copy_both::CopyBothDuplex;
 use crate::copy_out::CopyOutStream;
 use crate::query::RowStream;
 use crate::simple_query::SimpleQueryStream;
@@ -11,8 +12,9 @@ use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
-    Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    copy_both, copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken,
+    CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
+    TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -433,6 +435,15 @@ impl Client {
         copy_out::copy_out(self.inner(), statement).await
     }
 
+    /// Executes a CopyBoth query, returning a combined Stream+Sink type to read and write copy
+    /// data.
+    pub async fn copy_both_simple<T>(&self, query: &str) -> Result<CopyBothDuplex<T>, Error>
+    where
+        T: Buf + 'static + Send,
+    {
+        copy_both::copy_both_simple(self.inner(), query).await
+    }
+
     /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows.
     ///
     /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that

tokio-postgres/src/connection.rs

@@ -1,4 +1,5 @@
 use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
+use crate::copy_both::CopyBothReceiver;
 use crate::copy_in::CopyInReceiver;
 use crate::error::DbError;
 use crate::maybe_tls_stream::MaybeTlsStream;
@@ -21,6 +22,7 @@ use tokio_util::codec::Framed;
 pub enum RequestMessages {
     Single(FrontendMessage),
     CopyIn(CopyInReceiver),
+    CopyBoth(CopyBothReceiver),
 }
 
 pub struct Request {
@@ -259,6 +261,24 @@ where
                         .map_err(Error::io)?;
                     self.pending_request = Some(RequestMessages::CopyIn(receiver));
                 }
+                RequestMessages::CopyBoth(mut receiver) => {
+                    let message = match receiver.poll_next_unpin(cx) {
+                        Poll::Ready(Some(message)) => message,
+                        Poll::Ready(None) => {
+                            trace!("poll_write: finished copy_both request");
+                            continue;
+                        }
+                        Poll::Pending => {
+                            trace!("poll_write: waiting on copy_both stream");
+                            self.pending_request = Some(RequestMessages::CopyBoth(receiver));
+                            return Ok(true);
+                        }
+                    };
+                    Pin::new(&mut self.stream)
+                        .start_send(message)
+                        .map_err(Error::io)?;
+                    self.pending_request = Some(RequestMessages::CopyBoth(receiver));
+                }
             }
         }
     }

tokio-postgres/src/copy_both.rs

@@ -0,0 +1,249 @@
+use crate::client::{InnerClient, Responses};
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::{simple_query, Error};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
+use futures::channel::mpsc;
+use futures::future;
+use futures::{ready, Sink, SinkExt, Stream, StreamExt};
+use log::debug;
+use pin_project_lite::pin_project;
+use postgres_protocol::message::backend::Message;
+use postgres_protocol::message::frontend;
+use postgres_protocol::message::frontend::CopyData;
+use std::marker::{PhantomData, PhantomPinned};
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pub(crate) enum CopyBothMessage {
+    Message(FrontendMessage),
+    Done,
+}
+
+pub struct CopyBothReceiver {
+    receiver: mpsc::Receiver<CopyBothMessage>,
+    done: bool,
+}
+
+impl CopyBothReceiver {
+    pub(crate) fn new(receiver: mpsc::Receiver<CopyBothMessage>) -> CopyBothReceiver {
+        CopyBothReceiver {
+            receiver,
+            done: false,
+        }
+    }
+}
+
+impl Stream for CopyBothReceiver {
+    type Item = FrontendMessage;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<FrontendMessage>> {
+        if self.done {
+            return Poll::Ready(None);
+        }
+
+        match ready!(self.receiver.poll_next_unpin(cx)) {
+            Some(CopyBothMessage::Message(message)) => Poll::Ready(Some(message)),
+            Some(CopyBothMessage::Done) => {
+                self.done = true;
+                let mut buf = BytesMut::new();
+                frontend::copy_done(&mut buf);
+                frontend::sync(&mut buf);
+                Poll::Ready(Some(FrontendMessage::Raw(buf.freeze())))
+            }
+            None => {
+                self.done = true;
+                let mut buf = BytesMut::new();
+                frontend::copy_fail("", &mut buf).unwrap();
+                frontend::sync(&mut buf);
+                Poll::Ready(Some(FrontendMessage::Raw(buf.freeze())))
+            }
+        }
+    }
+}
+
+enum SinkState {
+    Active,
+    Closing,
+    Reading,
+}
+
+pin_project! {
+    /// A sink for `COPY ... FROM STDIN` query data.
+    ///
+    /// The copy *must* be explicitly completed via the `Sink::close` or `finish` methods. If it is
+    /// not, the copy will be aborted.
+    pub struct CopyBothDuplex<T> {
+        #[pin]
+        sender: mpsc::Sender<CopyBothMessage>,
+        responses: Responses,
+        buf: BytesMut,
+        state: SinkState,
+        #[pin]
+        _p: PhantomPinned,
+        _p2: PhantomData<T>,
+    }
+}
+
+impl<T> CopyBothDuplex<T>
+where
+    T: Buf + 'static + Send,
+{
+    pub(crate) fn new(sender: mpsc::Sender<CopyBothMessage>, responses: Responses) -> Self {
+        Self {
+            sender,
+            responses,
+            buf: BytesMut::new(),
+            state: SinkState::Active,
+            _p: PhantomPinned,
+            _p2: PhantomData,
+        }
+    }
+
+    /// A poll-based version of `finish`.
+    pub fn poll_finish(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<u64, Error>> {
+        loop {
+            match self.state {
+                SinkState::Active => {
+                    ready!(self.as_mut().poll_flush(cx))?;
+                    let mut this = self.as_mut().project();
+                    ready!(this.sender.as_mut().poll_ready(cx)).map_err(|_| Error::closed())?;
+                    this.sender
+                        .start_send(CopyBothMessage::Done)
+                        .map_err(|_| Error::closed())?;
+                    *this.state = SinkState::Closing;
+                }
+                SinkState::Closing => {
+                    let this = self.as_mut().project();
+                    ready!(this.sender.poll_close(cx)).map_err(|_| Error::closed())?;
+                    *this.state = SinkState::Reading;
+                }
+                SinkState::Reading => {
+                    let this = self.as_mut().project();
+                    match ready!(this.responses.poll_next(cx))? {
+                        Message::CommandComplete(body) => {
+                            let rows = body
+                                .tag()
+                                .map_err(Error::parse)?
+                                .rsplit(' ')
+                                .next()
+                                .unwrap()
+                                .parse()
+                                .unwrap_or(0);
+                            return Poll::Ready(Ok(rows));
+                        }
+                        _ => return Poll::Ready(Err(Error::unexpected_message())),
+                    }
+                }
+            }
+        }
+    }
+
+    /// Completes the copy, returning the number of rows inserted.
+    ///
+    /// The `Sink::close` method is equivalent to `finish`, except that it does not return the
+    /// number of rows.
+    pub async fn finish(mut self: Pin<&mut Self>) -> Result<u64, Error> {
+        future::poll_fn(|cx| self.as_mut().poll_finish(cx)).await
+    }
+}
+
+impl<T> Stream for CopyBothDuplex<T> {
+    type Item = Result<Bytes, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+
+        match ready!(this.responses.poll_next(cx)?) {
+            Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
+            Message::CopyDone => Poll::Ready(None),
+            _ => Poll::Ready(Some(Err(Error::unexpected_message()))),
+        }
+    }
+}
+
+impl<T> Sink<T> for CopyBothDuplex<T>
+where
+    T: Buf + 'static + Send,
+{
+    type Error = Error;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        self.project()
+            .sender
+            .poll_ready(cx)
+            .map_err(|_| Error::closed())
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Error> {
+        let this = self.project();
+
+        let data: Box<dyn Buf + Send> = if item.remaining() > 4096 {
+            if this.buf.is_empty() {
+                Box::new(item)
+            } else {
+                Box::new(this.buf.split().freeze().chain(item))
+            }
+        } else {
+            this.buf.put(item);
+            if this.buf.len() > 4096 {
+                Box::new(this.buf.split().freeze())
+            } else {
+                return Ok(());
+            }
+        };
+
+        let data = CopyData::new(data).map_err(Error::encode)?;
+        this.sender
+            .start_send(CopyBothMessage::Message(FrontendMessage::CopyData(data)))
+            .map_err(|_| Error::closed())
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        let mut this = self.project();
+
+        if !this.buf.is_empty() {
+            ready!(this.sender.as_mut().poll_ready(cx)).map_err(|_| Error::closed())?;
+            let data: Box<dyn Buf + Send> = Box::new(this.buf.split().freeze());
+            let data = CopyData::new(data).map_err(Error::encode)?;
+            this.sender
+                .as_mut()
+                .start_send(CopyBothMessage::Message(FrontendMessage::CopyData(data)))
+                .map_err(|_| Error::closed())?;
+        }
+
+        this.sender.poll_flush(cx).map_err(|_| Error::closed())
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        self.poll_finish(cx).map_ok(|_| ())
+    }
+}
+
+pub async fn copy_both_simple<T>(
+    client: &InnerClient,
+    query: &str,
+) -> Result<CopyBothDuplex<T>, Error>
+where
+    T: Buf + 'static + Send,
+{
+    debug!("executing copy both query {}", query);
+
+    let buf = simple_query::encode(client, query)?;
+
+    let (mut sender, receiver) = mpsc::channel(1);
+    let receiver = CopyBothReceiver::new(receiver);
+    let mut responses = client.send(RequestMessages::CopyBoth(receiver))?;
+
+    sender
+        .send(CopyBothMessage::Message(FrontendMessage::Raw(buf)))
+        .await
+        .map_err(|_| Error::closed())?;
+
+    match responses.next().await? {
+        Message::CopyBothResponse(_) => {}
+        _ => return Err(Error::unexpected_message()),
+    }
+
+    Ok(CopyBothDuplex::new(sender, responses))
+}

tokio-postgres/src/lib.rs

@@ -155,6 +155,7 @@ mod connect_raw;
 mod connect_socket;
 mod connect_tls;
 mod connection;
+mod copy_both;
 mod copy_in;
 mod copy_out;
 pub mod error;