handle errors in initial postgres COPY statements

Author: lovasoa

sqlx-core/src/postgres/copy.rs

@@ -116,19 +116,29 @@ pub struct PgCopyIn<C: DerefMut<Target = PgConnection>> {
 }
 
 impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
-    async fn begin(mut conn: C, statement: &str) -> Result<Self> {
+    async fn begin(conn: C, statement: &str) -> Result<Self> {
+        let mut conn = Self::start_copy(conn, statement).await?;
+        match conn.stream.recv_expect(MessageFormat::CopyInResponse).await {
+            Ok(response) => Ok(PgCopyIn {
+                conn: Some(conn),
+                response,
+            }),
+            Err(e) => {
+                conn.stream
+                    .send(CopyFail::new("failed to start COPY"))
+                    .await?;
+                conn.stream
+                    .recv_expect(MessageFormat::ReadyForQuery)
+                    .await?;
+                Err(e)
+            }
+        }
+    }
+
+    async fn start_copy(mut conn: C, statement: &str) -> Result<C> {
         conn.wait_until_ready().await?;
         conn.stream.send(Query(statement)).await?;
-
-        let response: CopyResponse = conn
-            .stream
-            .recv_expect(MessageFormat::CopyInResponse)
-            .await?;
-
-        Ok(PgCopyIn {
-            conn: Some(conn),
-            response,
-        })
+        Ok(conn)
     }
 
     /// Returns `true` if Postgres is expecting data in text or CSV format.
@@ -252,18 +262,15 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
                 "fail_with: expected ErrorResponse, got: {:?}",
                 msg.format
             )),
-            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(Error::Database(e)) => match e.code() {
+                Some(Cow::Borrowed("57014")) => {
+                    conn.stream
+                        .recv_expect(MessageFormat::ReadyForQuery)
+                        .await?;
+                    Ok(())
                 }
-            }
+                _ => Err(Error::Database(e)),
+            },
             Err(e) => Err(e),
         }
     }

tests/postgres/postgres.rs

@@ -1743,3 +1743,66 @@ async fn test_postgres_bytea_hex_deserialization_errors() -> anyhow::Result<()>
     }
     Ok(())
 }
+
+async fn test_copy_in_error_case(query: &str, expected_error: &str) -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+    conn.execute("CREATE TEMPORARY TABLE IF NOT EXISTS invalid_copy_target (id int4)")
+        .await?;
+
+    // Try the COPY operation
+    match conn.copy_in_raw(query).await {
+        Ok(_) => anyhow::bail!("expected error"),
+        Err(e) => assert!(
+            e.to_string().contains(expected_error),
+            "expected error to contain: {expected_error}, got: {e:?}"
+        ),
+    }
+
+    // Verify connection is still usable
+    let value = sqlx_oldapi::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_recover_from_copy_in_to_missing_table() -> anyhow::Result<()> {
+    test_copy_in_error_case(
+        r#"
+        COPY nonexistent_table (id) FROM STDIN WITH (FORMAT CSV, HEADER);
+        "#,
+        "does not exist",
+    )
+    .await
+}
+
+#[sqlx_macros::test]
+async fn it_can_recover_from_copy_in_empty_query() -> anyhow::Result<()> {
+    test_copy_in_error_case("", "EmptyQuery").await
+}
+
+#[sqlx_macros::test]
+async fn it_can_recover_from_copy_in_syntax_error() -> anyhow::Result<()> {
+    test_copy_in_error_case(
+        r#"
+        COPY FROM STDIN WITH (FORMAT CSV);
+        "#,
+        "syntax error",
+    )
+    .await
+}
+
+#[sqlx_macros::test]
+async fn it_can_recover_from_copy_in_invalid_params() -> anyhow::Result<()> {
+    test_copy_in_error_case(
+        r#"
+        COPY invalid_copy_target FROM STDIN WITH (FORMAT CSV, INVALID_PARAM true);
+        "#,
+        "invalid_param",
+    )
+    .await
+}