fix: ensures recover from fail with PgCopyIn

launchbadge · Oct 29, 2022 · a406d3b · a406d3b
1 parent 76ae286
commit a406d3b
Showing 1 changed file with 15 additions and 11 deletions.
diff --git a/sqlx-core/src/postgres/copy.rs b/sqlx-core/src/postgres/copy.rs
@@ -120,10 +120,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),
@@ -278,14 +281,15 @@ impl<C: DerefMut<Target = PgConnection>> PgCopyIn<C> {
             .expect("CopyWriter::finish: conn taken illegally");
 
         conn.stream.send(CopyDone).await?;
-        let cc: CommandComplete = conn
-            .stream
-            .recv_expect(MessageFormat::CommandComplete)
-            .await?;
+        let cc: CommandComplete = match conn.stream.recv_expect(MessageFormat::CommandComplete).await {
+            Ok(cc) => cc,
+            Err(e) => {
+                conn.stream.recv().await?;
+                return Err(e);
+            }
+        };
 
-        conn.stream
-            .recv_expect(MessageFormat::ReadyForQuery)
-            .await?;
+        conn.stream.recv_expect(MessageFormat::ReadyForQuery).await?;
 
         Ok(cc.rows_affected())
     }