Support batches on pgmq_delete

core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgmq"
-version = "0.16.3"
+version = "0.17.0"
 edition = "2021"
 authors = ["Tembo.io"]
 description = "A distributed message queue for Rust applications, on Postgres."

core/sqlx-data.json

@@ -1,5 +1,26 @@
 {
   "db": "PostgreSQL",
+  "0c81d8ee07fac75c18d9bdebd9b2ac5e473b4ef43a5d332f5d58ae6a0e01e600": {
+    "describe": {
+      "columns": [
+        {
+          "name": "pgmq_delete",
+          "ordinal": 0,
+          "type_info": "Bool"
+        }
+      ],
+      "nullable": [
+        null
+      ],
+      "parameters": {
+        "Left": [
+          "Text",
+          "Int8"
+        ]
+      }
+    },
+    "query": "SELECT * from pgmq_delete($1::text, $2::bigint)"
+  },
   "0e13957cc5c1f96a07ec242ce393f0738d62d79d34182e3c9333e9f9712d4cbb": {
     "describe": {
       "columns": [
@@ -119,6 +140,27 @@
     },
     "query": "SELECT * from pgmq_archive($1::text, $2)"
   },
+  "8e4f6635dc4cfb5ed42ddb87930b985793e81243b8a5609dec4ab566aaab4e9c": {
+    "describe": {
+      "columns": [
+        {
+          "name": "pgmq_delete",
+          "ordinal": 0,
+          "type_info": "Bool"
+        }
+      ],
+      "nullable": [
+        null
+      ],
+      "parameters": {
+        "Left": [
+          "Text",
+          "Int8Array"
+        ]
+      }
+    },
+    "query": "SELECT * from pgmq_delete($1::text, $2::bigint[])"
+  },
   "9919286cee87946b387f69e67df94f94eb0acdd3b5f4848faf092c55d484b61a": {
     "describe": {
       "columns": [
@@ -185,27 +227,6 @@
     },
     "query": "SELECT * from pgmq_set_vt($1::text, $2::bigint, $3::integer);"
   },
-  "c38b0cc48a6744bdb28ad1fa637223b6427eda19c0a52dc8881750d6a50ac0ef": {
-    "describe": {
-      "columns": [
-        {
-          "name": "pgmq_delete",
-          "ordinal": 0,
-          "type_info": "Bool"
-        }
-      ],
-      "nullable": [
-        null
-      ],
-      "parameters": {
-        "Left": [
-          "Text",
-          "Int8"
-        ]
-      }
-    },
-    "query": "SELECT * from pgmq_delete($1::text, $2)"
-  },
   "c89ad5584222a1f9c0d505fe58bf035d68d033047176b106e289ad6842f49972": {
     "describe": {
       "columns": [

core/src/pg_ext.rs

@@ -295,7 +295,7 @@ impl PGMQueueExt {
     // Delete a message by message id.
     pub async fn delete(&self, queue_name: &str, msg_id: i64) -> Result<bool, PgmqError> {
         let row = sqlx::query!(
-            "SELECT * from pgmq_delete($1::text, $2)",
+            "SELECT * from pgmq_delete($1::text, $2::bigint)",
             queue_name,
             msg_id
         )
@@ -304,5 +304,15 @@ impl PGMQueueExt {
         Ok(row.pgmq_delete.expect("no delete result"))
     }
 
-    //
+    // Delete with a slice of message ids
+    pub async fn delete_batch(&self, queue_name: &str, msg_id: &[i64]) -> Result<bool, PgmqError> {
+        let row = sqlx::query!(
+            "SELECT * from pgmq_delete($1::text, $2::bigint[])",
+            queue_name,
+            msg_id
+        )
+        .fetch_one(&self.connection)
+        .await?;
+        Ok(row.pgmq_delete.expect("no delete result"))
+    }
 }

core/tests/integration_test.rs

@@ -870,6 +870,18 @@ async fn test_extension_api() {
         .await
         .expect("failed to delete");
     assert!(!deleted);
+
+    // delete a batch of messages
+    let m1 = queue.send(&test_queue, &del_msg).await.unwrap();
+    let m2 = queue.send(&test_queue, &del_msg).await.unwrap();
+    let m3 = queue.send(&test_queue, &del_msg).await.unwrap();
+    let delete_result = queue
+        .delete_batch(&test_queue, &[m1, m2, m3])
+        .await
+        .expect("delete batch error");
+    let rowcount = rowcount(&test_queue, &queue.connection).await;
+    assert_eq!(rowcount, 0);
+    assert_eq!(delete_result, true);
 }
 
 #[tokio::test]

src/lib.rs

@@ -12,7 +12,7 @@ pub mod partition;
 
 use pgmq_crate::errors::PgmqError;
 use pgmq_crate::query::{
-    archive, check_input, delete, init_queue, pop, read, PGMQ_SCHEMA, TABLE_PREFIX,
+    archive, check_input, delete, delete_batch, init_queue, pop, read, PGMQ_SCHEMA, TABLE_PREFIX,
 };
 
 use errors::PgmqExtError;
@@ -215,6 +215,22 @@ fn pgmq_delete(queue_name: &str, msg_id: i64) -> Result<Option<bool>, PgmqExtErr
     }
 }
 
+#[pg_extern(name = "pgmq_delete")]
+fn pgmq_delete_batch(queue_name: &str, msg_ids: Vec<i64>) -> Result<Option<bool>, PgmqExtError> {
+    let mut num_deleted = 0;
+    let query = delete_batch(queue_name, &msg_ids)?;
+    Spi::connect(|mut client| {
+        let tup_table = client.update(&query, None, None);
+        match tup_table {
+            Ok(tup_table) => num_deleted = tup_table.len(),
+            Err(e) => {
+                error!("error deleting message: {}", e);
+            }
+        }
+    });
+    Ok(Some(true))
+}
+
 /// archive a message forever instead of deleting it
 #[pg_extern]
 fn pgmq_archive(queue_name: &str, msg_id: i64) -> Result<Option<bool>, PgmqExtError> {