Catch command tag

Author: kelvich

tokio-postgres/src/query.rs

@@ -52,6 +52,7 @@ where
     Ok(RowStream {
         statement,
         responses,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -72,6 +73,7 @@ pub async fn query_portal(
     Ok(RowStream {
         statement: portal.statement().clone(),
         responses,
+        command_tag: None,
         _p: PhantomPinned,
     })
 }
@@ -202,6 +204,7 @@ pin_project! {
     pub struct RowStream {
         statement: Statement,
         responses: Responses,
+        command_tag: Option<String>,
         #[pin]
         _p: PhantomPinned,
     }
@@ -213,6 +216,7 @@ impl RowStream {
         RowStream {
             statement,
             responses,
+            command_tag: None,
             _p: PhantomPinned,
         }
     }
@@ -228,12 +232,24 @@ impl Stream for RowStream {
                 Message::DataRow(body) => {
                     return Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
                 }
-                Message::EmptyQueryResponse
-                | Message::CommandComplete(_)
-                | Message::PortalSuspended => {}
+                Message::EmptyQueryResponse | Message::PortalSuspended => {}
+                Message::CommandComplete(body) => {
+                    if let Ok(tag) = body.tag() {
+                        *this.command_tag = Some(tag.to_string());
+                    }
+                }
                 Message::ReadyForQuery(_) => return Poll::Ready(None),
                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
             }
         }
     }
 }
+
+impl RowStream {
+    /// Returns the command tag of this query.
+    ///
+    /// This is only available after the stream has been exhausted.
+    pub fn command_tag(&self) -> Option<String> {
+        self.command_tag.clone()
+    }
+}

tokio-postgres/tests/test/main.rs

@@ -264,7 +264,7 @@ async fn query_raw_txt() {
         .unwrap();
 
     assert_eq!(rows.len(), 1);
-    let res: i32 = rows[0].as_text(0).unwrap().parse::<i32>().unwrap();
+    let res: i32 = rows[0].as_text(0).unwrap().unwrap().parse::<i32>().unwrap();
     assert_eq!(res, 55 * 42);
 
     let rows: Vec<tokio_postgres::Row> = client
@@ -279,6 +279,25 @@ async fn query_raw_txt() {
     assert_eq!(rows[0].get::<_, &str>(0), "42");
 }
 
+#[tokio::test]
+async fn command_tag() {
+    let client = connect("user=postgres").await;
+
+    let row_stream = client
+        .query_raw_txt("select unnest('{1,2,3}'::int[]);", [])
+        .await
+        .unwrap();
+
+    pin_mut!(row_stream);
+
+    let mut rows: Vec<tokio_postgres::Row> = Vec::new();
+    while let Some(row) = row_stream.next().await {
+        rows.push(row.unwrap());
+    }
+
+    assert_eq!(row_stream.command_tag(), Some("SELECT 3".to_string()));
+}
+
 #[tokio::test]
 async fn custom_composite() {
     let client = connect("user=postgres").await;