test and complete support for postgres copy

Author: Montana Low

sqlx-core/src/postgres/connection/mod.rs

@@ -20,8 +20,6 @@ use crate::transaction::Transaction;
 
 pub use self::stream::PgStream;
 
-pub use self::stream::PgStream;
-
 pub(crate) mod describe;
 mod establish;
 mod executor;

sqlx-core/src/postgres/copy.rs

@@ -8,6 +8,7 @@ use crate::postgres::message::{
 use crate::postgres::Postgres;
 use bytes::{BufMut, Bytes};
 use futures_core::stream::BoxStream;
+use smallvec::alloc::borrow::Cow;
 use sqlx_rt::{AsyncRead, AsyncReadExt, AsyncWriteExt};
 use std::convert::TryFrom;
 use std::ops::{Deref, DerefMut};
@@ -235,8 +236,18 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
                 "fail_with: expected ErrorResponse, got: {:?}",
                 msg.format
             )),
-            // FIXME: inspect the `DatabaseError` to make sure we're not discarding another error
-            Err(Error::Database(_db)) => Ok(()),
+            Err(Error::Database(e)) => {
+                match e.code() {
+                    Some(Cow::Borrowed("57014")) => {
+                        // postgres abort received error code
+                        conn.stream
+                            .recv_expect(MessageFormat::ReadyForQuery)
+                            .await?;
+                        Ok(())
+                    }
+                    _ => Err(Error::Database(e)),
+                }
+            }
             Err(e) => Err(e),
         }
     }
@@ -256,6 +267,10 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
             .recv_expect(MessageFormat::CommandComplete)
             .await?;
 
+        conn.stream
+            .recv_expect(MessageFormat::ReadyForQuery)
+            .await?;
+
         Ok(cc.rows_affected())
     }
 }
@@ -279,7 +294,7 @@ async fn pg_begin_copy_out<'c, C: DerefMut<Target = PgConnection> + Send + 'c>(
 
     let _: CopyResponse = conn
         .stream
-        .recv_expect(MessageFormat::CopyInResponse)
+        .recv_expect(MessageFormat::CopyOutResponse)
         .await?;
 
     let stream: TryAsyncStream<'c, Bytes> = try_stream! {
@@ -289,6 +304,8 @@ async fn pg_begin_copy_out<'c, C: DerefMut<Target = PgConnection> + Send + 'c>(
                 MessageFormat::CopyData => r#yield!(msg.decode::<CopyData<Bytes>>()?.0),
                 MessageFormat::CopyDone => {
                     let _ = msg.decode::<CopyDone>()?;
+                    conn.stream.recv_expect(MessageFormat::CommandComplete).await?;
+                    conn.stream.recv_expect(MessageFormat::ReadyForQuery).await?;
                     return Ok(())
                 },
                 _ => return Err(err_protocol!("unexpected message format during copy out: {:?}", msg.format))

sqlx-core/src/postgres/message/copy.rs

@@ -3,7 +3,6 @@ use crate::io::{BufExt, BufMutExt, Decode, Encode};
 use bytes::{Buf, BufMut, Bytes};
 use std::ops::Deref;
 
-
 /// The same structure is sent for both `CopyInResponse` and `CopyOutResponse`
 pub struct CopyResponse {
     pub format: i8,

sqlx-core/src/postgres/message/mod.rs

@@ -106,6 +106,8 @@ impl MessageFormat {
             b'C' => MessageFormat::CommandComplete,
             b'd' => MessageFormat::CopyData,
             b'c' => MessageFormat::CopyDone,
+            b'G' => MessageFormat::CopyInResponse,
+            b'H' => MessageFormat::CopyOutResponse,
             b'D' => MessageFormat::DataRow,
             b'E' => MessageFormat::ErrorResponse,
             b'I' => MessageFormat::EmptyQueryResponse,

tests/postgres/postgres.rs

@@ -1,4 +1,4 @@
-use futures::TryStreamExt;
+use futures::{StreamExt, TryStreamExt};
 use sqlx::postgres::{
     PgConnectOptions, PgConnection, PgDatabaseError, PgErrorPosition, PgSeverity,
 };
@@ -1104,3 +1104,99 @@ async fn test_pg_server_num() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn it_can_copy_in() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+    conn.execute(
+        r#"
+        CREATE TEMPORARY TABLE users (id INTEGER NOT NULL);
+    "#,
+    )
+    .await?;
+
+    let mut copy = conn
+        .copy_in_raw(
+            r#"
+        COPY users (id) FROM STDIN WITH (FORMAT CSV, HEADER);
+    "#,
+        )
+        .await?;
+
+    copy.send("id\n1\n2\n".as_bytes()).await?;
+    let rows = copy.finish().await?;
+    assert_eq!(rows, 2);
+
+    // conn is safe for reuse
+    let value = sqlx::query("select 1 + 1")
+        .try_map(|row: PgRow| row.try_get::<i32, _>(0))
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(2i32, value);
+
+    Ok(())
+}
+
+#[sqlx_macros::test]
+async fn it_can_abort_copy_in() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+    conn.execute(
+        r#"
+        CREATE TEMPORARY TABLE users (id INTEGER NOT NULL);
+    "#,
+    )
+    .await?;
+
+    let mut copy = conn
+        .copy_in_raw(
+            r#"
+        COPY users (id) FROM STDIN WITH (FORMAT CSV, HEADER);
+    "#,
+        )
+        .await?;
+
+    copy.abort("this is only a test").await?;
+
+    // conn is safe for reuse
+    let value = sqlx::query("select 1 + 1")
+        .try_map(|row: PgRow| row.try_get::<i32, _>(0))
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(2i32, value);
+
+    Ok(())
+}
+
+#[sqlx_macros::test]
+async fn it_can_copy_out() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+
+    {
+        let mut copy = conn
+            .copy_out_raw(
+                "
+            COPY (SELECT generate_series(1, 2) AS id) TO STDOUT WITH (FORMAT CSV, HEADER);
+        ",
+            )
+            .await?;
+
+        assert_eq!(copy.next().await.unwrap().unwrap(), "id\n");
+        assert_eq!(copy.next().await.unwrap().unwrap(), "1\n");
+        assert_eq!(copy.next().await.unwrap().unwrap(), "2\n");
+        if copy.next().await.is_some() {
+            anyhow::bail!("Unexpected data from COPY");
+        }
+    }
+
+    // conn is safe for reuse
+    let value = sqlx::query("select 1 + 1")
+        .try_map(|row: PgRow| row.try_get::<i32, _>(0))
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(2i32, value);
+
+    Ok(())
+}

tests/sqlite/.gitignore

@@ -0,0 +1,2 @@
+sqlite.db
+