From 31e541ac7a9c7d18ee2b3b91c58349e77eac28f7 Mon Sep 17 00:00:00 2001
From: Thomas de Zeeuw <thomasdezeeuw@gmail.com>
Date: Sat, 20 Jan 2024 11:36:06 +0100
Subject: [PATCH] Fix handling of deferred constraints for PostgreSQL (#2913)

---
 sqlx-postgres/src/connection/executor.rs | 15 ++++++++++-----
 tests/postgres/postgres.rs               | 24 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/sqlx-postgres/src/connection/executor.rs b/sqlx-postgres/src/connection/executor.rs
index 67ffa16ac9..b1ce29c516 100644
--- a/sqlx-postgres/src/connection/executor.rs
+++ b/sqlx-postgres/src/connection/executor.rs
@@ -402,13 +402,18 @@ impl<'c> Executor<'c> for &'c mut PgConnection {
             let s = self.run(sql, arguments, 1, persistent, metadata).await?;
             pin_mut!(s);
 
-            while let Some(s) = s.try_next().await? {
-                if let Either::Right(r) = s {
-                    return Ok(Some(r));
+            // With deferred constraints we need to check all responses as we
+            // could get a OK response (with uncommitted data), only to get an
+            // error response after (when the deferred constraint is actually
+            // checked).
+            let mut ret = None;
+            while let Some(result) = s.try_next().await? {
+                match result {
+                    Either::Right(r) if ret.is_none() => ret = Some(r),
+                    _ => {}
                 }
             }
-
-            Ok(None)
+            Ok(ret)
         })
     }
 
diff --git a/tests/postgres/postgres.rs b/tests/postgres/postgres.rs
index 477a455d74..352d3278ea 100644
--- a/tests/postgres/postgres.rs
+++ b/tests/postgres/postgres.rs
@@ -1813,3 +1813,27 @@ async fn test_shrink_buffers() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn test_error_handling_with_deferred_constraints() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+
+    sqlx::query("CREATE TABLE IF NOT EXISTS deferred_constraint ( id INTEGER PRIMARY KEY )")
+        .execute(&mut conn)
+        .await?;
+
+    sqlx::query("CREATE TABLE IF NOT EXISTS deferred_constraint_fk ( fk INTEGER CONSTRAINT deferred_fk REFERENCES deferred_constraint(id) DEFERRABLE INITIALLY DEFERRED )")
+            .execute(&mut conn)
+            .await?;
+
+    let result: sqlx::Result<i32> =
+        sqlx::query_scalar("INSERT INTO deferred_constraint_fk VALUES (1) RETURNING fk")
+            .fetch_one(&mut conn)
+            .await;
+
+    let err = result.unwrap_err();
+    let db_err = err.as_database_error().unwrap();
+    assert_eq!(db_err.constraint(), Some("deferred_fk"));
+
+    Ok(())
+}