Merge pull request #833 from sfackler/deferred-constraint

Fix commit-time error reporting
sfackler · Oct 19, 2021 · 7f38b39 · 7f38b39
2 parents d454616 + b2df115
commit 7f38b39
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 11 deletions.
diff --git a/tokio-postgres/src/query.rs b/tokio-postgres/src/query.rs
@@ -99,21 +99,22 @@ where
     };
     let mut responses = start(client, buf).await?;
 
+    let mut rows = 0;
     loop {
         match responses.next().await? {
             Message::DataRow(_) => {}
             Message::CommandComplete(body) => {
-                let rows = body
+                rows = body
                     .tag()
                     .map_err(Error::parse)?
                     .rsplit(' ')
                     .next()
                     .unwrap()
                     .parse()
                     .unwrap_or(0);
-                return Ok(rows);
             }
-            Message::EmptyQueryResponse => return Ok(0),
+            Message::EmptyQueryResponse => rows = 0,
+            Message::ReadyForQuery(_) => return Ok(rows),
             _ => return Err(Error::unexpected_message()),
         }
     }
@@ -203,15 +204,17 @@ impl Stream for RowStream {
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.project();
-        match ready!(this.responses.poll_next(cx)?) {
-            Message::DataRow(body) => {
-                Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
+        loop {
+            match ready!(this.responses.poll_next(cx)?) {
+                Message::DataRow(body) => {
+                    return Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
+                }
+                Message::EmptyQueryResponse
+                | Message::CommandComplete(_)
+                | Message::PortalSuspended => {}
+                Message::ReadyForQuery(_) => return Poll::Ready(None),
+                _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
-            Message::EmptyQueryResponse
-            | Message::CommandComplete(_)
-            | Message::PortalSuspended => Poll::Ready(None),
-            Message::ErrorResponse(body) => Poll::Ready(Some(Err(Error::db(body)))),
-            _ => Poll::Ready(Some(Err(Error::unexpected_message()))),
         }
     }
 }
diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs
@@ -805,3 +805,29 @@ async fn query_opt() {
         .err()
         .unwrap();
 }
+
+#[tokio::test]
+async fn deferred_constraint() {
+    let client = connect("user=postgres").await;
+
+    client
+        .batch_execute(
+            "
+            CREATE TEMPORARY TABLE t (
+                i INT,
+                UNIQUE (i) DEFERRABLE INITIALLY DEFERRED
+            );
+        ",
+        )
+        .await
+        .unwrap();
+
+    client
+        .execute("INSERT INTO t (i) VALUES (1)", &[])
+        .await
+        .unwrap();
+    client
+        .execute("INSERT INTO t (i) VALUES (1)", &[])
+        .await
+        .unwrap_err();
+}