change to enum CoLinkMQType

Author: stneng

src/application.rs

@@ -548,8 +548,12 @@ pub struct CoLinkInfo {
     pub version: String,
 }
 
+pub enum CoLinkMQType {
+    RabbitMQ,
+    RedisStream,
+}
 pub struct CoLinkSubscriber {
-    mq_type: i32, // 0 for rabbitmq, 1 for redis stream
+    mq_type: CoLinkMQType,
     queue_name: String,
     rabbitmq_consumer: Option<lapin::Consumer>,
     redis_connection: Option<redis::aio::Connection>,
@@ -562,7 +566,7 @@ impl CoLinkSubscriber {
             let client = redis::Client::open(mq_uri)?;
             let con = client.get_async_connection().await?;
             Ok(Self {
-                mq_type: 1,
+                mq_type: CoLinkMQType::RedisStream,
                 queue_name: queue_name.to_string(),
                 rabbitmq_consumer: None,
                 redis_connection: Some(con),
@@ -579,7 +583,7 @@ impl CoLinkSubscriber {
                 )
                 .await?;
             Ok(Self {
-                mq_type: 0,
+                mq_type: CoLinkMQType::RabbitMQ,
                 queue_name: queue_name.to_string(),
                 rabbitmq_consumer: Some(consumer),
                 redis_connection: None,
@@ -588,44 +592,46 @@ impl CoLinkSubscriber {
     }
 
     pub async fn get_next(&mut self) -> Result<Vec<u8>, Error> {
-        if self.mq_type == 0 {
-            let delivery = self
-                .rabbitmq_consumer
-                .as_mut()
-                .unwrap()
-                .next()
-                .await
-                .expect("error in consumer");
-            let delivery = delivery.expect("error in consumer");
-            delivery.ack(BasicAckOptions::default()).await?;
-            Ok(delivery.data)
-        } else if self.mq_type == 1 {
-            let opts = StreamReadOptions::default()
-                .group(&self.queue_name, uuid::Uuid::new_v4().to_string())
-                .block(0)
-                .count(1);
-            let res: StreamReadReply = self
-                .redis_connection
-                .as_mut()
-                .unwrap()
-                .xread_options(&[&self.queue_name], &[">"], &opts)
-                .await?;
-            let id = &res.keys[0].ids[0].id;
-            let data: Vec<u8> =
-                FromRedisValue::from_redis_value(res.keys[0].ids[0].map.get("payload").unwrap())?;
-            self.redis_connection
-                .as_mut()
-                .unwrap()
-                .xack(&self.queue_name, &self.queue_name, &[id])
-                .await?;
-            self.redis_connection
-                .as_mut()
-                .unwrap()
-                .xdel(&self.queue_name, &[id])
-                .await?;
-            Ok(data)
-        } else {
-            Err("mq type is not supported.")?
+        match self.mq_type {
+            CoLinkMQType::RabbitMQ => {
+                let delivery = self
+                    .rabbitmq_consumer
+                    .as_mut()
+                    .unwrap()
+                    .next()
+                    .await
+                    .expect("error in consumer");
+                let delivery = delivery.expect("error in consumer");
+                delivery.ack(BasicAckOptions::default()).await?;
+                Ok(delivery.data)
+            }
+            CoLinkMQType::RedisStream => {
+                let opts = StreamReadOptions::default()
+                    .group(&self.queue_name, uuid::Uuid::new_v4().to_string())
+                    .block(0)
+                    .count(1);
+                let res: StreamReadReply = self
+                    .redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xread_options(&[&self.queue_name], &[">"], &opts)
+                    .await?;
+                let id = &res.keys[0].ids[0].id;
+                let data: Vec<u8> = FromRedisValue::from_redis_value(
+                    res.keys[0].ids[0].map.get("payload").unwrap(),
+                )?;
+                self.redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xack(&self.queue_name, &self.queue_name, &[id])
+                    .await?;
+                self.redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xdel(&self.queue_name, &[id])
+                    .await?;
+                Ok(data)
+            }
         }
     }
 }