[OSS-26] Queues improvements (#32)

stefandanaita · Jul 30, 2024 · 4084a1f · 4084a1f
1 parent 552763f
commit 4084a1f
Show file tree

Hide file tree

Showing 11 changed files with 479 additions and 54 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,6 +16,7 @@ chrono = { version = "0.4", features = ["serde"] }
 http = "1"
 reqwest = { version = "0.12", features = ["json"] }
 reqwest-middleware = { version = "0.3.1", features = ["json"] }
+rust_decimal = "1.35"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 thiserror = "1"

diff --git a/src/api/_generic.rs b/src/api/_generic.rs
@@ -3,6 +3,7 @@ use http::StatusCode;
 use reqwest::Response;
 use serde::de::DeserializeOwned;
 
+#[tracing::instrument(skip(response))]
 pub async fn handle_response<T>(response: Response) -> Result<T, RabbitMqClientError>
 where
     T: DeserializeOwned,
@@ -24,6 +25,7 @@ where
     }
 }
 
+#[tracing::instrument(skip(response))]
 pub async fn handle_empty_response(response: Response) -> Result<(), RabbitMqClientError> {
     let status = response.status();
 

diff --git a/src/api/exchange.rs b/src/api/exchange.rs
@@ -4,7 +4,6 @@ use crate::errors::RabbitMqClientError;
 use crate::RabbitMqClient;
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 
 #[async_trait]
 pub trait ExchangeApi {
@@ -39,13 +38,6 @@ pub trait ExchangeApi {
         exchange: String,
     ) -> Result<(), RabbitMqClientError>;
 
-    async fn publish_message(
-        &self,
-        vhost: String,
-        exchange: String,
-        request: RabbitMqExchangeMessagePublishRequest,
-    ) -> Result<RabbitMqExchangeMessagePublishResponse, RabbitMqClientError>;
-
     async fn list_source_bindings(
         &self,
         vhost: String,
@@ -151,28 +143,6 @@ impl ExchangeApi for RabbitMqClient {
         handle_empty_response(response).await
     }
 
-    async fn publish_message(
-        &self,
-        vhost: String,
-        exchange: String,
-        request: RabbitMqExchangeMessagePublishRequest,
-    ) -> Result<RabbitMqExchangeMessagePublishResponse, RabbitMqClientError> {
-        let response = self
-            .client
-            .request(
-                reqwest::Method::POST,
-                format!(
-                    "{}/api/exchanges/{}/{}/publish",
-                    self.api_url, vhost, exchange
-                ),
-            )
-            .json(&request)
-            .send()
-            .await?;
-
-        handle_response(response).await
-    }
-
     async fn list_source_bindings(
         &self,
         vhost: String,
@@ -229,8 +199,8 @@ pub struct RabbitMqExchange {
 
 #[derive(Debug, Deserialize)]
 pub struct RabbitMqExchangeMessageStats {
-    pub publish_in: i64,
-    pub publish_out: i64,
+    pub publish_in: Option<i64>,
+    pub publish_out: Option<i64>,
 }
 
 #[derive(Debug, Serialize)]
@@ -241,24 +211,3 @@ pub struct RabbitMqExchangeRequest {
     pub durable: bool,
     pub internal: bool,
 }
-
-#[derive(Debug, Serialize)]
-pub struct RabbitMqExchangeMessagePublishRequest {
-    pub properties: HashMap<String, String>,
-    pub routing_key: String,
-    pub payload: String,
-    pub payload_encoding: RabbitMqMessageEncoding,
-}
-
-#[derive(Debug, Serialize)]
-pub enum RabbitMqMessageEncoding {
-    #[serde(rename = "string")]
-    String,
-    #[serde(rename = "base64")]
-    Base64,
-}
-
-#[derive(Debug, Deserialize)]
-pub struct RabbitMqExchangeMessagePublishResponse {
-    pub routed: bool,
-}
diff --git a/src/api/message.rs b/src/api/message.rs
@@ -0,0 +1,127 @@
+use crate::api::_generic::handle_response;
+use crate::errors::RabbitMqClientError;
+use crate::RabbitMqClient;
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+#[async_trait]
+pub trait MessageApi {
+    async fn publish_message(
+        &self,
+        vhost: String,
+        exchange: String,
+        request: RabbitMqPublishMessageRequest,
+    ) -> Result<RabbitMqPublishMessageResponse, RabbitMqClientError>;
+
+    async fn get_messages(
+        &self,
+        vhost: String,
+        queue: String,
+        options: RabbitMqGetMessagesOptions,
+    ) -> Result<Vec<RabbitMqMessage>, RabbitMqClientError>;
+}
+
+#[async_trait]
+impl MessageApi for RabbitMqClient {
+    async fn publish_message(
+        &self,
+        vhost: String,
+        exchange: String,
+        request: RabbitMqPublishMessageRequest,
+    ) -> Result<RabbitMqPublishMessageResponse, RabbitMqClientError> {
+        let response = self
+            .client
+            .request(
+                reqwest::Method::POST,
+                format!(
+                    "{}/api/exchanges/{}/{}/publish",
+                    self.api_url, vhost, exchange
+                ),
+            )
+            .json(&request)
+            .send()
+            .await?;
+
+        handle_response(response).await
+    }
+
+    #[tracing::instrument(skip(self))]
+    async fn get_messages(
+        &self,
+        vhost: String,
+        queue: String,
+        options: RabbitMqGetMessagesOptions,
+    ) -> Result<Vec<RabbitMqMessage>, RabbitMqClientError> {
+        let response = self
+            .client
+            .request(
+                reqwest::Method::POST,
+                format!("{}/api/queues/{}/{}/get", self.api_url, vhost, queue),
+            )
+            .json(&options)
+            .send()
+            .await?;
+
+        handle_response(response).await
+    }
+}
+
+#[derive(Debug, Serialize)]
+pub struct RabbitMqPublishMessageRequest {
+    pub properties: HashMap<String, String>,
+    pub routing_key: String,
+    pub payload: String,
+    pub payload_encoding: RabbitMqMessageEncoding,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub enum RabbitMqMessageEncoding {
+    #[serde(rename = "string")]
+    String,
+    #[serde(rename = "base64")]
+    Base64,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct RabbitMqPublishMessageResponse {
+    pub routed: bool,
+}
+
+#[derive(Debug, Serialize)]
+pub struct RabbitMqGetMessagesOptions {
+    pub count: u32,
+    #[serde(rename = "ackmode")]
+    pub ack_mode: RabbitMqGetMessagesAckMode,
+    pub encoding: RabbitMqGetMessagesEncoding,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub truncate: Option<u64>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum RabbitMqGetMessagesAckMode {
+    AckRequeueTrue,
+    AckRequeueFalse,
+    RejectRequeueTrue,
+    RejectRequeueFalse,
+}
+
+#[derive(Debug, Serialize)]
+pub enum RabbitMqGetMessagesEncoding {
+    #[serde(rename = "auto")]
+    Auto,
+    #[serde(rename = "base64")]
+    Base64,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct RabbitMqMessage {
+    pub payload_bytes: u64,
+    pub redelivered: bool,
+    pub exchange: String,
+    pub routing_key: String,
+    pub message_count: u64,
+    pub payload: String,
+    pub payload_encoding: RabbitMqMessageEncoding,
+}
diff --git a/src/api/mod.rs b/src/api/mod.rs
@@ -2,6 +2,7 @@ mod _generic;
 pub mod binding;
 pub mod connection;
 pub mod exchange;
+pub mod message;
 pub mod node;
 pub mod overview;
 pub mod permission;

diff --git a/src/api/queue.rs b/src/api/queue.rs
@@ -3,6 +3,7 @@ use crate::api::binding::RabbitMqBinding;
 use crate::errors::RabbitMqClientError;
 use crate::RabbitMqClient;
 use async_trait::async_trait;
+use rust_decimal::Decimal;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
@@ -53,6 +54,7 @@ pub trait QueueApi {
 
 #[async_trait]
 impl QueueApi for RabbitMqClient {
+    #[tracing::instrument(skip(self))]
     async fn list_queues(
         &self,
         vhost: Option<String>,
@@ -69,6 +71,7 @@ impl QueueApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn get_queue(
         &self,
         vhost: String,
@@ -86,6 +89,7 @@ impl QueueApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn get_queue_bindings(
         &self,
         vhost: String,
@@ -103,6 +107,7 @@ impl QueueApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn create_queue(
         &self,
         vhost: String,
@@ -121,6 +126,7 @@ impl QueueApi for RabbitMqClient {
         }
     }
 
+    #[tracing::instrument(skip(self))]
     async fn update_queue(
         &self,
         vhost: String,
@@ -140,6 +146,7 @@ impl QueueApi for RabbitMqClient {
         handle_empty_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn delete_queue(&self, vhost: String, name: String) -> Result<(), RabbitMqClientError> {
         let response = self
             .client
@@ -153,6 +160,7 @@ impl QueueApi for RabbitMqClient {
         handle_empty_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn purge_queue(&self, vhost: String, name: String) -> Result<(), RabbitMqClientError> {
         let response = self
             .client
@@ -166,6 +174,7 @@ impl QueueApi for RabbitMqClient {
         handle_empty_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn set_queue_actions(
         &self,
         vhost: String,
@@ -190,7 +199,7 @@ impl QueueApi for RabbitMqClient {
 pub struct RabbitMqQueue {
     pub name: String,
     pub node: String,
-    pub arguments: HashMap<String, String>,
+    pub arguments: HashMap<String, RabbitMqArgument>,
     pub state: String,
     #[serde(rename = "type")]
     pub kind: String,
@@ -208,6 +217,13 @@ pub struct RabbitMqQueue {
     pub message_stats: Option<RabbitMqQueueMessageStats>,
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(untagged)]
+pub enum RabbitMqArgument {
+    String(String),
+    Decimal(Decimal),
+}
+
 #[derive(Debug, Deserialize)]
 pub struct RabbitMqQueueMessageStats {
     #[serde(default)]

diff --git a/src/api/user.rs b/src/api/user.rs
@@ -33,6 +33,7 @@ pub trait UserApi {
 
 #[async_trait]
 impl UserApi for RabbitMqClient {
+    #[tracing::instrument(skip(self))]
     async fn who_am_i(&self) -> Result<RabbitMqWhoAmI, RabbitMqClientError> {
         let response = self
             .client
@@ -43,6 +44,7 @@ impl UserApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn list_users(&self) -> Result<Vec<RabbitMqUser>, RabbitMqClientError> {
         let response = self
             .client
@@ -53,6 +55,7 @@ impl UserApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn list_users_without_permissions(
         &self,
     ) -> Result<Vec<RabbitMqUser>, RabbitMqClientError> {
@@ -68,6 +71,7 @@ impl UserApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn bulk_delete_users(
         &self,
         users: RabbitMqUsersBulkDeleteRequest,
@@ -85,6 +89,7 @@ impl UserApi for RabbitMqClient {
         handle_empty_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn list_user_permissions(
         &self,
         user: String,
@@ -101,6 +106,7 @@ impl UserApi for RabbitMqClient {
         handle_response(response).await
     }
 
+    #[tracing::instrument(skip(self))]
     async fn list_user_topic_permissions(
         &self,
         user: String,