From 9e55c9bd0cc678c73b9d31c0cd19684e54f231b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Felipe=20Quintero=20Moreano?=
 <afquinteromoreano@gmail.com>
Date: Fri, 10 Feb 2023 20:38:36 -0500
Subject: [PATCH] fix: ensures recover from fail with PgCopyIn (#2179)

* CHANGELOG: mention that users should upgrade CLI

* fix: ensures recover from fail with PgCopyIn

---------

Co-authored-by: Austin Bonander <austin.bonander@gmail.com>
---
 sqlx-postgres/src/copy.rs | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/sqlx-postgres/src/copy.rs b/sqlx-postgres/src/copy.rs
index 46bb5fe658..620784308a 100644
--- a/sqlx-postgres/src/copy.rs
+++ b/sqlx-postgres/src/copy.rs
@@ -145,10 +145,13 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
         conn.wait_until_ready().await?;
         conn.stream.send(Query(statement)).await?;
 
-        let response: CopyResponse = conn
-            .stream
-            .recv_expect(MessageFormat::CopyInResponse)
-            .await?;
+        let response = match conn.stream.recv_expect(MessageFormat::CopyInResponse).await {
+            Ok(res) => res,
+            Err(e) => {
+                conn.stream.recv().await?;
+                return Err(e);
+            }
+        };
 
         Ok(PgCopyIn {
             conn: Some(conn),
@@ -299,10 +302,17 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
             .expect("CopyWriter::finish: conn taken illegally");
 
         conn.stream.send(CopyDone).await?;
-        let cc: CommandComplete = conn
+        let cc: CommandComplete = match conn
             .stream
             .recv_expect(MessageFormat::CommandComplete)
-            .await?;
+            .await
+        {
+            Ok(cc) => cc,
+            Err(e) => {
+                conn.stream.recv().await?;
+                return Err(e);
+            }
+        };
 
         conn.stream
             .recv_expect(MessageFormat::ReadyForQuery)