From 0fc4005ed31e3705a04cb7e58eb220d89b922dd0 Mon Sep 17 00:00:00 2001
From: Ramnivas Laddad <ramnivas@exograph.dev>
Date: Mon, 22 Jul 2024 15:07:44 -0700
Subject: [PATCH] For `query_typed`, deal with the no-data case.

If a query returns no data, we receive `Message::NoData`, which signals
the completion of the query. However, we treated it as a no-op, leading
to processing other messages and eventual failure.

This PR fixes the issue and updates the `query_typed` tests to cover
this scenario.
---
 tokio-postgres/src/query.rs       | 13 +++++++++----
 tokio-postgres/tests/test/main.rs | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/tokio-postgres/src/query.rs b/tokio-postgres/src/query.rs
index be42d66b6..3ab002871 100644
--- a/tokio-postgres/src/query.rs
+++ b/tokio-postgres/src/query.rs
@@ -89,10 +89,15 @@ where
 
     loop {
         match responses.next().await? {
-            Message::ParseComplete
-            | Message::BindComplete
-            | Message::ParameterDescription(_)
-            | Message::NoData => {}
+            Message::ParseComplete | Message::BindComplete | Message::ParameterDescription(_) => {}
+            Message::NoData => {
+                return Ok(RowStream {
+                    statement: Statement::unnamed(vec![], vec![]),
+                    responses,
+                    rows_affected: None,
+                    _p: PhantomPinned,
+                });
+            }
             Message::RowDescription(row_description) => {
                 let mut columns: Vec<Column> = vec![];
                 let mut it = row_description.fields();
diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs
index 84c46d101..9a6aa26fe 100644
--- a/tokio-postgres/tests/test/main.rs
+++ b/tokio-postgres/tests/test/main.rs
@@ -997,6 +997,13 @@ async fn query_typed_no_transaction() {
     assert_eq!(second_row.get::<_, i32>(1), 40);
     assert_eq!(second_row.get::<_, &str>(2), "literal");
     assert_eq!(second_row.get::<_, i32>(3), 5);
+
+    // Test for UPDATE that returns no data
+    let updated_rows = client
+        .query_typed("UPDATE foo set age = 33", &[])
+        .await
+        .unwrap();
+    assert_eq!(updated_rows.len(), 0);
 }
 
 #[tokio::test]
@@ -1064,4 +1071,11 @@ async fn query_typed_with_transaction() {
     assert_eq!(second_row.get::<_, i32>(1), 40);
     assert_eq!(second_row.get::<_, &str>(2), "literal");
     assert_eq!(second_row.get::<_, i32>(3), 5);
+
+    // Test for UPDATE that returns no data
+    let updated_rows = transaction
+        .query_typed("UPDATE foo set age = 33", &[])
+        .await
+        .unwrap();
+    assert_eq!(updated_rows.len(), 0);
 }