parse message in clear way, pass channel param to SubscriptionItem

Author: estin

examples/pubsub.rs

@@ -20,11 +20,15 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
         loop {
             match subscriber.recv().await {
-                Some(Ok(cmd::SubscribeItem::Subscribed)) => println!("sub: subscribed"),
-                Some(Ok(cmd::SubscribeItem::Message(payload))) => {
-                    println!("sub: {:?}", payload)
+                Some(Ok(cmd::SubscribeItem::Subscribed(channel))) => {
+                    println!("sub: subscribed to {:?}", channel)
+                }
+                Some(Ok(cmd::SubscribeItem::Message { channel, payload })) => {
+                    println!("sub: {:?} from {:?}", payload, channel)
+                }
+                Some(Ok(cmd::SubscribeItem::UnSubscribed(channel))) => {
+                    println!("sub: unsubscribed from {:?}", channel)
                 }
-                Some(Ok(cmd::SubscribeItem::UnSubscribed)) => println!("sub: unsubscribed"),
                 Some(Err(e)) => {
                     println!("sub: {}", e);
                     return;

src/cmd/pubsub.rs

@@ -1,5 +1,5 @@
 use super::{utils, Command, CommandError};
-use ntex::util::Bytes;
+use ntex::util::{Bytes, Either};
 use std::convert::TryFrom;
 
 use crate::codec::{BulkString, Request, Response};
@@ -8,76 +8,58 @@ const TYPE_SUBSCRIBE: Bytes = Bytes::from_static(b"subscribe");
 const TYPE_UNSUBSCRIBE: Bytes = Bytes::from_static(b"unsubscribe");
 const TYPE_MESSAGE: Bytes = Bytes::from_static(b"message");
 
-/// Publish redis command
-pub fn Publish<T, V>(key: T, value: V) -> utils::IntOutputCommand
-where
-    BulkString: From<T> + From<V>,
-{
-    utils::IntOutputCommand(Request::Array(vec![
-        Request::from_static("PUBLISH"),
-        Request::BulkString(key.into()),
-        Request::BulkString(value.into()),
-    ]))
-}
-
 #[derive(Debug, PartialEq)]
 pub enum SubscribeItem {
-    Subscribed,
-    UnSubscribed,
-    Message(Bytes),
+    Subscribed(Bytes),
+    UnSubscribed(Bytes),
+    Message { channel: Bytes, payload: Bytes },
 }
 
-impl TryFrom<Response> for SubscribeItem {
-    type Error = CommandError;
+struct MessagePayload(Either<Bytes, i64>);
 
-    fn try_from(val: Response) -> Result<Self, Self::Error> {
-        let parts = if let Response::Array(ref parts) = val {
-            parts
-        } else {
-            return Err(CommandError::Output("Cannot parse response", val));
-        };
+impl TryFrom<Response> for MessagePayload {
+    type Error = (&'static str, Response);
 
-        if parts.len() != 3 {
-            return Err(CommandError::Output(
-                "Subscription message has invalid items number",
-                val,
-            ));
+    fn try_from(val: Response) -> Result<Self, Self::Error> {
+        match val {
+            Response::Bytes(bytes) => Ok(MessagePayload(Either::Left(bytes))),
+            Response::Integer(number) => Ok(MessagePayload(Either::Right(number))),
+            _ => Err(("Not a bytes object or integer", val)),
         }
+    }
+}
 
-        let (mtype, payload) = (&parts[0], &parts[2]);
+impl TryFrom<Response> for SubscribeItem {
+    type Error = CommandError;
 
-        let mtype = if let Response::Bytes(mtype) = mtype {
-            mtype
-        } else {
-            return Err(CommandError::Output(
-                "Subscription message type unknown",
-                val,
-            ));
-        };
+    fn try_from(val: Response) -> Result<Self, Self::Error> {
+        let (mtype, channel, payload) = <(Bytes, Bytes, MessagePayload)>::try_from(val)?;
 
         if mtype == &TYPE_SUBSCRIBE {
-            return Ok(SubscribeItem::Subscribed);
+            return Ok(SubscribeItem::Subscribed(channel));
         }
 
         if mtype == &TYPE_UNSUBSCRIBE {
-            return Ok(SubscribeItem::UnSubscribed);
+            return Ok(SubscribeItem::UnSubscribed(channel));
         }
 
         if mtype != &TYPE_MESSAGE {
             return Err(CommandError::Output(
                 "Subscription message type unknown",
-                val,
+                Response::Bytes(mtype),
             ));
         }
 
-        if let Response::Bytes(payload) = payload {
-            return Ok(SubscribeItem::Message(payload.clone()));
+        let payload = if let Some(payload) = payload.0.left() {
+            payload
         } else {
             return Err(CommandError::Output(
-                "Subscription message has empty payload",
-                val,
+                "Subscription message payload is not bytes",
+                Response::Nil,
             ));
-        }
+        };
+
+        Ok(SubscribeItem::Message { channel, payload })
     }
 }
 
@@ -95,6 +77,18 @@ impl Command for SubscribeOutputCommand {
     }
 }
 
+/// Publish redis command
+pub fn Publish<T, V>(key: T, value: V) -> utils::IntOutputCommand
+where
+    BulkString: From<T> + From<V>,
+{
+    utils::IntOutputCommand(Request::Array(vec![
+        Request::from_static("PUBLISH"),
+        Request::BulkString(key.into()),
+        Request::BulkString(value.into()),
+    ]))
+}
+
 /// Subscribe redis command
 pub fn Subscribe<T>(key: T) -> SubscribeOutputCommand
 where

tests/test_redis.rs

@@ -204,33 +204,43 @@ async fn test_connection() {
 #[ntex::test]
 async fn test_pubsub() {
     let key = new_key();
+    let channel = Bytes::from(key);
 
     let subscriber = RedisConnector::new("127.0.0.1:6379")
         .connect_simple()
         .await
         .unwrap();
 
-    let stream = subscriber.stream(cmd::Subscribe(&key)).unwrap();
+    let stream = subscriber.stream(cmd::Subscribe(&channel)).unwrap();
 
     let message = stream.recv().await;
-    assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::Subscribed);
+    assert_eq!(
+        message.unwrap().unwrap(),
+        cmd::SubscribeItem::Subscribed(channel.clone())
+    );
 
     let publisher = connect().await;
 
     // pub
-    let result = publisher.exec(cmd::Publish(&key, "1")).await.unwrap();
+    let result = publisher.exec(cmd::Publish(&channel, "1")).await.unwrap();
     assert_eq!(result, 1);
 
     // sub
     let message = stream.recv().await;
     assert_eq!(
         message.unwrap().unwrap(),
-        cmd::SubscribeItem::Message(Bytes::from_static(b"1"))
+        cmd::SubscribeItem::Message {
+            channel: channel.clone(),
+            payload: Bytes::from_static(b"1")
+        }
     );
 
     // unsub
-    subscriber.send(cmd::UnSubscribe(&key)).unwrap();
+    subscriber.send(cmd::UnSubscribe(&channel)).unwrap();
 
     let message = stream.recv().await;
-    assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::UnSubscribed);
+    assert_eq!(
+        message.unwrap().unwrap(),
+        cmd::SubscribeItem::UnSubscribed(channel.clone())
+    );
 }