add simple query versions of copy operations

Signed-off-by: Petros Angelatos <[email protected]>

Author: petrosagg

tokio-postgres/src/client.rs

@@ -424,6 +424,14 @@ impl Client {
         copy_in::copy_in(self.inner(), statement).await
     }
 
+    /// Executes a `COPY FROM STDIN` query, returning a sink used to write the copy data.
+    pub async fn copy_in_simple<U>(&self, query: &str) -> Result<CopyInSink<U>, Error>
+    where
+        U: Buf + 'static + Send,
+    {
+        copy_in::copy_in_simple(self.inner(), query).await
+    }
+
     /// Executes a `COPY TO STDOUT` statement, returning a stream of the resulting data.
     ///
     /// PostgreSQL does not support parameters in `COPY` statements, so this method does not take any.
@@ -439,6 +447,11 @@ impl Client {
         copy_out::copy_out(self.inner(), statement).await
     }
 
+    /// Executes a `COPY TO STDOUT` query, returning a stream of the resulting data.
+    pub async fn copy_out_simple(&self, query: &str) -> Result<CopyOutStream, Error> {
+        copy_out::copy_out_simple(self.inner(), query).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>

tokio-postgres/src/copy_in.rs

@@ -1,8 +1,8 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
-use crate::{query, slice_iter, Error, Statement};
-use bytes::{Buf, BufMut, BytesMut};
+use crate::{query, simple_query, slice_iter, Error, Statement};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
 use futures::channel::mpsc;
 use futures::future;
 use futures::{ready, Sink, SinkExt, Stream, StreamExt};
@@ -195,14 +195,10 @@ where
     }
 }
 
-pub async fn copy_in<T>(client: &InnerClient, statement: Statement) -> Result<CopyInSink<T>, Error>
+async fn start<T>(client: &InnerClient, buf: Bytes, simple: bool) -> Result<CopyInSink<T>, Error>
 where
     T: Buf + 'static + Send,
 {
-    debug!("executing copy in statement {}", statement.name());
-
-    let buf = query::encode(client, &statement, slice_iter(&[]))?;
-
     let (mut sender, receiver) = mpsc::channel(1);
     let receiver = CopyInReceiver::new(receiver);
     let mut responses = client.send(RequestMessages::CopyIn(receiver))?;
@@ -212,9 +208,11 @@ where
         .await
         .map_err(|_| Error::closed())?;
 
-    match responses.next().await? {
-        Message::BindComplete => {}
-        _ => return Err(Error::unexpected_message()),
+    if !simple {
+        match responses.next().await? {
+            Message::BindComplete => {}
+            _ => return Err(Error::unexpected_message()),
+        }
     }
 
     match responses.next().await? {
@@ -231,3 +229,23 @@ where
         _p2: PhantomData,
     })
 }
+
+pub async fn copy_in<T>(client: &InnerClient, statement: Statement) -> Result<CopyInSink<T>, Error>
+where
+    T: Buf + 'static + Send,
+{
+    debug!("executing copy in statement {}", statement.name());
+
+    let buf = query::encode(client, &statement, slice_iter(&[]))?;
+    start(client, buf, false).await
+}
+
+pub async fn copy_in_simple<T>(client: &InnerClient, query: &str) -> Result<CopyInSink<T>, Error>
+where
+    T: Buf + 'static + Send,
+{
+    debug!("executing copy in query {}", query);
+
+    let buf = simple_query::encode(client, query)?;
+    start(client, buf, true).await
+}

tokio-postgres/src/copy_out.rs

@@ -1,7 +1,7 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
-use crate::{query, slice_iter, Error, Statement};
+use crate::{query, simple_query, slice_iter, Error, Statement};
 use bytes::Bytes;
 use futures::{ready, Stream};
 use log::debug;
@@ -11,23 +11,36 @@ use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
+pub async fn copy_out_simple(client: &InnerClient, query: &str) -> Result<CopyOutStream, Error> {
+    debug!("executing copy out query {}", query);
+
+    let buf = simple_query::encode(client, query)?;
+    let responses = start(client, buf, true).await?;
+    Ok(CopyOutStream {
+        responses,
+        _p: PhantomPinned,
+    })
+}
+
 pub async fn copy_out(client: &InnerClient, statement: Statement) -> Result<CopyOutStream, Error> {
     debug!("executing copy out statement {}", statement.name());
 
     let buf = query::encode(client, &statement, slice_iter(&[]))?;
-    let responses = start(client, buf).await?;
+    let responses = start(client, buf, false).await?;
     Ok(CopyOutStream {
         responses,
         _p: PhantomPinned,
     })
 }
 
-async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
+async fn start(client: &InnerClient, buf: Bytes, simple: bool) -> Result<Responses, Error> {
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
 
-    match responses.next().await? {
-        Message::BindComplete => {}
-        _ => return Err(Error::unexpected_message()),
+    if !simple {
+        match responses.next().await? {
+            Message::BindComplete => {}
+            _ => return Err(Error::unexpected_message()),
+        }
     }
 
     match responses.next().await? {