diff --git a/sqlx-core/src/postgres/copy.rs b/sqlx-core/src/postgres/copy.rs index 0bad775085..57a84ecc21 100644 --- a/sqlx-core/src/postgres/copy.rs +++ b/sqlx-core/src/postgres/copy.rs @@ -120,10 +120,13 @@ impl> PgCopyIn { 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> PgCopyIn { .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()) }