fix(postgres): track the ReadyForQuery owed by COPY TO STDOUT

sqlx-postgres/src/connection/mod.rs

@@ -107,7 +107,7 @@ impl PgConnection {
         Ok(())
     }
 
-    async fn recv_ready_for_query(&mut self) -> Result<(), Error> {
+    pub(crate) async fn recv_ready_for_query(&mut self) -> Result<(), Error> {
         let r: ReadyForQuery = self.inner.stream.recv_expect().await?;
 
         self.inner.pending_ready_for_query_count -= 1;

sqlx-postgres/src/copy.rs

@@ -323,23 +323,28 @@ async fn pg_begin_copy_out<'c, C: DerefMut<Target = PgConnection> + Send + 'c>(
     statement: &str,
 ) -> Result<BoxStream<'c, Result<Bytes>>> {
     conn.wait_until_ready().await?;
-    conn.inner.stream.send(Query(statement)).await?;
+    conn.inner.stream.write_msg(Query(statement))?;
+    // Counted so a cancelled COPY's owed `ReadyForQuery` is drained by `wait_until_ready`
+    // instead of being mistaken for the reply to the connection's next message.
+    // Before the flush: a cancelled flush leaves the `Query` buffered, to be sent later.
+    conn.inner.pending_ready_for_query_count += 1;
+    conn.inner.stream.flush().await?;
 
     let _: CopyOutResponse = conn.inner.stream.recv_expect().await?;
 
     let stream: TryAsyncStream<'c, Bytes> = try_stream! {
         loop {
             match conn.inner.stream.recv().await {
                 Err(e) => {
-                    conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+                    conn.recv_ready_for_query().await?;
                     return Err(e);
                 },
                 Ok(msg) => match msg.format {
                     BackendMessageFormat::CopyData => r#yield!(msg.decode::<CopyData<Bytes>>()?.0),
                     BackendMessageFormat::CopyDone => {
                         let _ = msg.decode::<CopyDone>()?;
                         conn.inner.stream.recv_expect::<CommandComplete>().await?;
-                        conn.inner.stream.recv_expect::<ReadyForQuery>().await?;
+                        conn.recv_ready_for_query().await?;
                         return Ok(())
                     },
                     _ => return Err(err_protocol!("unexpected message format during copy out: {:?}", msg.format))

tests/postgres/postgres.rs

@@ -1787,6 +1787,50 @@ async fn it_can_copy_out() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[sqlx_macros::test]
+async fn it_can_reuse_connection_after_cancelled_copy_out() -> anyhow::Result<()> {
+    use sqlx::postgres::PgPoolCopyExt;
+    use sqlx_core::rt::timeout;
+
+    let mut setup = new::<Postgres>().await?;
+    setup
+        .execute("CREATE TABLE IF NOT EXISTS _sqlx_copy_out_cancel (id INT)")
+        .await?;
+
+    let pool = PgPoolOptions::new()
+        .max_connections(1)
+        .connect(&dotenvy::var("DATABASE_URL")?)
+        .await?;
+
+    let mut lock = setup.begin().await?;
+    lock.execute("LOCK TABLE _sqlx_copy_out_cancel IN ACCESS EXCLUSIVE MODE")
+        .await?;
+
+    let copy = async {
+        let mut stream = pool
+            .copy_out_raw("COPY (SELECT id FROM _sqlx_copy_out_cancel) TO STDOUT")
+            .await?;
+        while stream.try_next().await?.is_some() {}
+        Ok::<_, sqlx::Error>(())
+    };
+    assert!(
+        timeout(Duration::from_millis(200), copy).await.is_err(),
+        "COPY should time out while the table is locked"
+    );
+
+    lock.commit().await?;
+
+    let one: i32 = sqlx::query_scalar("SELECT 1").fetch_one(&pool).await?;
+    assert_eq!(one, 1);
+
+    let mut stream = pool
+        .copy_out_raw("COPY (SELECT id FROM _sqlx_copy_out_cancel) TO STDOUT")
+        .await?;
+    while stream.try_next().await?.is_some() {}
+
+    Ok(())
+}
+
 #[sqlx_macros::test]
 async fn it_encodes_custom_array_issue_1504() -> anyhow::Result<()> {
     use sqlx::encode::IsNull;