sqlx-postgres: Update copy

Author: joeydewaal

sqlx-core/src/io/mod.rs

@@ -24,3 +24,19 @@ pub use futures_util::io::AsyncReadExt;
 
 #[cfg(feature = "_rt-tokio")]
 pub use tokio::io::AsyncReadExt;
+
+pub async fn read_from(
+    mut source: impl AsyncRead + Unpin,
+    data: &mut Vec<u8>,
+) -> std::io::Result<usize> {
+    match () {
+        // Tokio lets us read into the buffer without zeroing first
+        #[cfg(feature = "_rt-tokio")]
+        _ => source.read_buf(data).await,
+        #[cfg(not(feature = "_rt-tokio"))]
+        _ => {
+            data.resize(data.capacity(), 0);
+            source.read(data).await
+        }
+    }
+}

sqlx-core/src/net/socket/buffered.rs

@@ -7,7 +7,7 @@ use std::pin::Pin;
 use std::task::{ready, Context, Poll};
 use std::{cmp, io};
 
-use crate::io::{AsyncRead, AsyncReadExt, ProtocolDecode, ProtocolEncode};
+use crate::io::{read_from, AsyncRead, ProtocolDecode, ProtocolEncode};
 
 // Tokio, async-std, and std all use this as the default capacity for their buffered I/O.
 const DEFAULT_BUF_SIZE: usize = 8192;
@@ -261,14 +261,8 @@ impl WriteBuffer {
     /// Read into the buffer from `source`, returning the number of bytes read.
     ///
     /// The buffer is automatically advanced by the number of bytes read.
-    pub async fn read_from(&mut self, mut source: impl AsyncRead + Unpin) -> io::Result<usize> {
-        let read = match () {
-            // Tokio lets us read into the buffer without zeroing first
-            #[cfg(feature = "_rt-tokio")]
-            _ => source.read_buf(self.buf_mut()).await?,
-            #[cfg(not(feature = "_rt-tokio"))]
-            _ => source.read(self.init_remaining_mut()).await?,
-        };
+    pub async fn read_from(&mut self, source: impl AsyncRead + Unpin) -> io::Result<usize> {
+        let read = read_from(source, self.buf_mut()).await?;
 
         if read > 0 {
             self.advance(read);

sqlx-postgres/src/connection/mod.rs

@@ -184,19 +184,16 @@ impl PgConnection {
         self.send_request(req)
     }
 
-    pub(crate) async fn start_pipe_async<F, R>(&self, callback: F) -> sqlx_core::Result<(R, Pipe)>
+    pub(crate) async fn pipe_and_forget_async<F, R>(&self, callback: F) -> sqlx_core::Result<R>
     where
         F: AsyncFnOnce(&mut MessageBuf) -> sqlx_core::Result<R>,
     {
         let mut buffer = MessageBuf::new();
         let result = (callback)(&mut buffer).await?;
-        let mut req = buffer.finish();
-        let (tx, rx) = unbounded();
-        req.chan = Some(tx);
-
+        let req = buffer.finish();
         self.send_request(req)?;
 
-        Ok((result, Pipe::new(rx)))
+        Ok(result)
     }
 }
 

sqlx-postgres/src/connection/worker.rs

@@ -152,7 +152,7 @@ impl Worker {
                     let _ = self.notif_chan.unbounded_send(notif);
                 }
                 BackendMessageFormat::ParameterStatus => {
-                    // Asynchronous response - todo
+                    // Asynchronous response
                     //
                     let ParameterStatus { name, value } = response.decode()?;
                     self.shared.insert_parameter_status(name, value);

sqlx-postgres/src/copy.rs

@@ -12,7 +12,7 @@ use crate::ext::async_stream::TryAsyncStream;
 use crate::io::AsyncRead;
 use crate::message::{
     BackendMessageFormat, CommandComplete, CopyData, CopyDone, CopyFail, CopyInResponse,
-    CopyOutResponse, CopyResponseData, Query, ReadyForQuery,
+    CopyOutResponse, CopyResponseData, ReadyForQuery,
 };
 use crate::pool::{Pool, PoolConnection};
 use crate::Postgres;
@@ -146,13 +146,13 @@ pub struct PgCopyIn<C: DerefMut<Target = PgConnection>> {
 }
 
 impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
-    async fn begin(mut conn: C, statement: &str) -> Result<Self> {
-        conn.inner.stream.send(Query(statement)).await?;
+    async fn begin(conn: C, statement: &str) -> Result<Self> {
+        let mut pipe = conn.queue_simple_query(statement)?;
 
-        let response = match conn.inner.stream.recv_expect::<CopyInResponse>().await {
+        let response = match pipe.recv_expect::<CopyInResponse>().await {
             Ok(res) => res.0,
             Err(e) => {
-                conn.inner.stream.recv().await?;
+                pipe.recv_ready_for_query().await?;
                 return Err(e);
             }
         };
@@ -194,13 +194,11 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
     /// If you're copying data from an `AsyncRead`, maybe consider [Self::read_from] instead.
     pub async fn send(&mut self, data: impl Deref<Target = [u8]>) -> Result<&mut Self> {
         for chunk in data.deref().chunks(PG_COPY_MAX_DATA_LEN) {
+            // TODO: We should probably have some kind of back-pressure here
             self.conn
                 .as_deref_mut()
                 .expect("send_data: conn taken")
-                .inner
-                .stream
-                .send(CopyData(chunk))
-                .await?;
+                .pipe_and_forget(CopyData(chunk))?;
         }
 
         Ok(self)
@@ -223,26 +221,31 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
     pub async fn read_from(&mut self, mut source: impl AsyncRead + Unpin) -> Result<&mut Self> {
         let conn: &mut PgConnection = self.conn.as_deref_mut().expect("copy_from: conn taken");
         loop {
-            let buf = conn.inner.stream.write_buffer_mut();
+            let read = conn
+                .pipe_and_forget_async(async |buf| {
+                    let write_buf = buf.buf_mut();
 
-            // Write the CopyData format code and reserve space for the length.
-            // This may end up sending an empty `CopyData` packet if, after this point,
-            // we get canceled or read 0 bytes, but that should be fine.
-            buf.put_slice(b"d\0\0\0\x04");
+                    // Write the CopyData format code and reserve space for the length.
+                    // This may end up sending an empty `CopyData` packet if, after this point,
+                    // we get canceled or read 0 bytes, but that should be fine.
+                    write_buf.put_slice(b"d\0\0\0\x04");
 
-            let read = buf.read_from(&mut source).await?;
+                    let read = sqlx_core::io::read_from(&mut source, write_buf).await?;
 
-            if read == 0 {
-                break;
-            }
+                    // Write the length
+                    let read32 = i32::try_from(read).map_err(|_| {
+                        err_protocol!("number of bytes read exceeds 2^31 - 1: {}", read)
+                    })?;
 
-            // Write the length
-            let read32 = i32::try_from(read)
-                .map_err(|_| err_protocol!("number of bytes read exceeds 2^31 - 1: {}", read))?;
+                    (&mut write_buf[1..]).put_i32(read32 + 4);
 
-            (&mut buf.get_mut()[1..]).put_i32(read32 + 4);
+                    Ok(read32)
+                })
+                .await?;
 
-            conn.inner.stream.flush().await?;
+            if read == 0 {
+                break;
+            }
         }
 
         Ok(self)
@@ -254,14 +257,14 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
     ///
     /// The server is expected to respond with an error, so only _unexpected_ errors are returned.
     pub async fn abort(mut self, msg: impl Into<String>) -> Result<()> {
-        let mut conn = self
+        let conn = self
             .conn
             .take()
             .expect("PgCopyIn::fail_with: conn taken illegally");
 
-        conn.inner.stream.send(CopyFail::new(msg)).await?;
+        let mut pipe = conn.pipe(|buf| buf.write_msg(CopyFail::new(msg)))?;
 
-        match conn.inner.stream.recv().await {
+        match pipe.recv().await {
             Ok(msg) => Err(err_protocol!(
                 "fail_with: expected ErrorResponse, got: {:?}",
                 msg.format
@@ -270,7 +273,7 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
                 match e.code() {
                     Some(Cow::Borrowed("57014")) => {
                         // postgres abort received error code
-                        conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+                        pipe.recv_expect::<ReadyForQuery>().await?;
                         Ok(())
                     }
                     _ => Err(Error::Database(e)),
@@ -284,60 +287,58 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
     ///
     /// The number of rows affected is returned.
     pub async fn finish(mut self) -> Result<u64> {
-        let mut conn = self
+        let conn = self
             .conn
             .take()
             .expect("CopyWriter::finish: conn taken illegally");
 
-        conn.inner.stream.send(CopyDone).await?;
-        let cc: CommandComplete = match conn.inner.stream.recv_expect().await {
+        let mut pipe = conn.pipe(|buf| buf.write_msg(CopyDone))?;
+        let cc: CommandComplete = match pipe.recv_expect().await {
             Ok(cc) => cc,
             Err(e) => {
-                conn.inner.stream.recv().await?;
+                pipe.recv().await?;
                 return Err(e);
             }
         };
 
-        conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+        pipe.recv_expect::<ReadyForQuery>().await?;
 
         Ok(cc.rows_affected())
     }
 }
 
 impl<C: DerefMut<Target = PgConnection>> Drop for PgCopyIn<C> {
     fn drop(&mut self) {
-        if let Some(mut conn) = self.conn.take() {
-            conn.inner
-                .stream
-                .write_msg(CopyFail::new(
-                    "PgCopyIn dropped without calling finish() or fail()",
-                ))
-                .expect("BUG: PgCopyIn abort message should not be too large");
+        if let Some(conn) = self.conn.take() {
+            conn.pipe_and_forget(CopyFail::new(
+                "PgCopyIn dropped without calling finish() or fail()",
+            ))
+            .expect("BUG: could not send PgCopyIn to background worker");
         }
     }
 }
 
 async fn pg_begin_copy_out<'c, C: DerefMut<Target = PgConnection> + Send + 'c>(
-    mut conn: C,
+    conn: C,
     statement: &str,
 ) -> Result<BoxStream<'c, Result<Bytes>>> {
-    conn.inner.stream.send(Query(statement)).await?;
+    let mut pipe = conn.queue_simple_query(statement)?;
 
-    let _: CopyOutResponse = conn.inner.stream.recv_expect().await?;
+    let _: CopyOutResponse = pipe.recv_expect().await?;
 
     let stream: TryAsyncStream<'c, Bytes> = try_stream! {
         loop {
-            match conn.inner.stream.recv().await {
+            match pipe.recv().await {
                 Err(e) => {
-                    conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+                    pipe.recv_expect::<ReadyForQuery>().await?;
                     return Err(e);
                 },
                 Ok(msg) => match msg.format {
                     BackendMessageFormat::CopyData => r#yield!(msg.decode::<CopyData<Bytes>>()?.0),
                     BackendMessageFormat::CopyDone => {
                         let _ = msg.decode::<CopyDone>()?;
-                        conn.inner.stream.recv_expect::<CommandComplete>().await?;
-                        conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+                        pipe.recv_expect::<CommandComplete>().await?;
+                        pipe.recv_expect::<ReadyForQuery>().await?;
                         return Ok(())
                     },
                     _ => return Err(err_protocol!("unexpected message format during copy out: {:?}", msg.format))