Skip to content

Commit

Permalink
fix: ensures recover from fail with PgCopyIn
Browse files Browse the repository at this point in the history
  • Loading branch information
andyquinterom committed Feb 8, 2023
1 parent ff0b0fc commit 6473a60
Showing 1 changed file with 16 additions and 6 deletions.
22 changes: 16 additions & 6 deletions sqlx-postgres/src/copy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 6473a60

Please sign in to comment.