Merge pull request #1165 from exograph/query-typed-no-data

For `query_typed`, deal with the no-data case.
sfackler · Jul 23, 2024 · 3afcc3d · 3afcc3d
2 parents 9f196e7 + 0fc4005
commit 3afcc3d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git 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
@@ -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);
 }