support subscribe/unsubscribe to multimple channels, added ssubscribe/sunsubscribe and psubsribe/punsubscribe

Author: estin

examples/pubsub.rs

@@ -5,20 +5,26 @@ use std::error::Error;
 async fn main() -> Result<(), Box<dyn Error>> {
     env_logger::init();
 
+    let channel = "pubsub_channel";
+
     // subscriber
     let client = RedisConnector::new("127.0.0.1:6379")
         .connect_simple()
         .await?;
 
-    let pubsub = client.subscribe(cmd::Subscribe("pubsub")).unwrap();
+    let pubsub = client.subscribe(cmd::Subscribe(vec![channel])).unwrap();
 
     ntex::rt::spawn(async move {
         loop {
             match pubsub.recv().await {
                 Some(Ok(cmd::SubscribeItem::Subscribed(channel))) => {
                     println!("sub: subscribed to {:?}", channel)
                 }
-                Some(Ok(cmd::SubscribeItem::Message { channel, payload })) => {
+                Some(Ok(cmd::SubscribeItem::Message {
+                    pattern: _,
+                    channel,
+                    payload,
+                })) => {
                     println!("sub: {:?} from {:?}", payload, channel)
                 }
                 Some(Ok(cmd::SubscribeItem::UnSubscribed(channel))) => {
@@ -39,7 +45,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     for i in 0..5 {
         let value = i.to_string();
         println!("pub: {}", value);
-        redis.exec(cmd::Publish("pubsub", &value)).await?;
+        redis.exec(cmd::Publish(channel, &value)).await?;
     }
 
     Ok(())

src/cmd/mod.rs

@@ -18,7 +18,10 @@ pub use self::connection::{Ping, Reset, Select};
 pub use self::hashes::{HDel, HGet, HGetAll, HIncrBy, HLen, HSet};
 pub use self::keys::{Del, Exists, Expire, ExpireAt, Keys, Ttl, TtlResult};
 pub use self::lists::{LIndex, LPop, LPush, RPop, RPush};
-pub use self::pubsub::{Publish, Subscribe, SubscribeItem, UnSubscribe};
+pub use self::pubsub::{
+    PSubscribe, PUnSubscribe, Publish, SPublish, SSubscribe, SUnSubscribe, Subscribe,
+    SubscribeItem, UnSubscribe,
+};
 pub use self::strings::{Get, IncrBy, Set};
 
 /// Trait implemented by types that can be used as redis commands

src/cmd/pubsub.rs

@@ -6,15 +6,31 @@ use crate::codec::{BulkString, Request, Response};
 
 const TYPE_SUBSCRIBE: Bytes = Bytes::from_static(b"subscribe");
 const TYPE_UNSUBSCRIBE: Bytes = Bytes::from_static(b"unsubscribe");
+const TYPE_SSUBSCRIBE: Bytes = Bytes::from_static(b"ssubscribe");
+const TYPE_SUNSUBSCRIBE: Bytes = Bytes::from_static(b"sunsubscribe");
+const TYPE_PSUBSCRIBE: Bytes = Bytes::from_static(b"psubscribe");
+const TYPE_PUNSUBSCRIBE: Bytes = Bytes::from_static(b"punsubscribe");
 const TYPE_MESSAGE: Bytes = Bytes::from_static(b"message");
+const TYPE_SMESSAGE: Bytes = Bytes::from_static(b"smessage");
+const TYPE_PMESSAGE: Bytes = Bytes::from_static(b"pmessage");
 
 pub trait PubSubCommand {}
 
+// #[derive(Debug, PartialEq)]
+// pub struct Channel {
+//     pub name: Bytes,
+//     pub pattern: Option<Bytes>,
+// }
+
 #[derive(Debug, PartialEq)]
 pub enum SubscribeItem {
     Subscribed(Bytes),
     UnSubscribed(Bytes),
-    Message { channel: Bytes, payload: Bytes },
+    Message {
+        pattern: Option<Bytes>,
+        channel: Bytes,
+        payload: Bytes,
+    },
 }
 
 struct MessagePayload(Either<Bytes, i64>);
@@ -35,33 +51,68 @@ impl TryFrom<Response> for SubscribeItem {
     type Error = CommandError;
 
     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(channel));
-        }
-
-        if mtype == &TYPE_UNSUBSCRIBE {
-            return Ok(SubscribeItem::UnSubscribed(channel));
-        }
-
-        if mtype != &TYPE_MESSAGE {
-            return Err(CommandError::Output(
-                "Subscription message type unknown",
-                Response::Bytes(mtype),
-            ));
-        }
-
-        let payload = if let Some(payload) = payload.0.left() {
-            payload
-        } else {
-            return Err(CommandError::Output(
-                "Subscription message payload is not bytes",
-                Response::Nil,
-            ));
+        let (mtype, pattern, channel, payload) = match val {
+            Response::Array(ary) => match ary.len() {
+                // subscribe or ssubscribe message
+                3 => {
+                    let mut ary_iter = ary.into_iter();
+                    (
+                        Bytes::try_from(ary_iter.next().expect("No value"))?,
+                        None,
+                        Bytes::try_from(ary_iter.next().expect("No value"))?,
+                        MessagePayload::try_from(ary_iter.next().expect("No value"))?,
+                    )
+                }
+                // psubscribe message
+                4 => {
+                    let mut ary_iter = ary.into_iter();
+                    (
+                        Bytes::try_from(ary_iter.next().expect("No value"))?,
+                        Some(Bytes::try_from(ary_iter.next().expect("No value"))?),
+                        Bytes::try_from(ary_iter.next().expect("No value"))?,
+                        MessagePayload::try_from(ary_iter.next().expect("No value"))?,
+                    )
+                }
+                _ => {
+                    return Err(CommandError::Output(
+                        "Array needs to be 3 or 4 elements",
+                        Response::Array(ary),
+                    ))
+                }
+            },
+            _ => return Err(CommandError::Output("Unexpected value", val)),
         };
 
-        Ok(SubscribeItem::Message { channel, payload })
+        match &mtype {
+            s if s == &TYPE_SUBSCRIBE || s == &TYPE_SSUBSCRIBE || s == &TYPE_PSUBSCRIBE => {
+                return Ok(SubscribeItem::Subscribed(channel));
+            }
+            s if s == &TYPE_UNSUBSCRIBE || s == &TYPE_SUNSUBSCRIBE || s == &TYPE_PUNSUBSCRIBE => {
+                return Ok(SubscribeItem::UnSubscribed(channel));
+            }
+            s if s == &TYPE_MESSAGE || s == &TYPE_SMESSAGE || s == &TYPE_PMESSAGE => {
+                let payload = if let Some(payload) = payload.0.left() {
+                    payload
+                } else {
+                    return Err(CommandError::Output(
+                        "Subscription message payload is not bytes",
+                        Response::Nil,
+                    ));
+                };
+
+                return Ok(SubscribeItem::Message {
+                    pattern,
+                    channel,
+                    payload,
+                });
+            }
+            _ => {
+                return Err(CommandError::Output(
+                    "Subscription message type unknown",
+                    Response::Bytes(mtype),
+                ));
+            }
+        };
     }
 }
 
@@ -93,7 +144,7 @@ impl Command for UnSubscribeOutputCommand {
     }
 }
 
-/// Publish redis command
+/// PUBLISH redis command
 pub fn Publish<T, V>(key: T, value: V) -> utils::IntOutputCommand
 where
     BulkString: From<T> + From<V>,
@@ -105,26 +156,94 @@ where
     ]))
 }
 
-/// Subscribe redis command
-pub fn Subscribe<T>(key: T) -> SubscribeOutputCommand
+/// SPUBLISH redis command
+pub fn SPublish<T, V>(key: T, value: V) -> utils::IntOutputCommand
 where
-    BulkString: From<T>,
+    BulkString: From<T> + From<V>,
 {
-    SubscribeOutputCommand(Request::Array(vec![
-        Request::from_static("SUBSCRIBE"),
+    utils::IntOutputCommand(Request::Array(vec![
+        Request::from_static("SPUBLISH"),
         Request::BulkString(key.into()),
+        Request::BulkString(value.into()),
     ]))
 }
 
-/// Unsubscribe redis command
-pub fn UnSubscribe<T>(key: T) -> UnSubscribeOutputCommand
+/// SUBSCRIBE redis command
+pub fn Subscribe<T>(channels: Vec<T>) -> SubscribeOutputCommand
 where
     BulkString: From<T>,
 {
-    UnSubscribeOutputCommand(Request::Array(vec![
-        Request::from_static("UNSUBSCRIBE"),
-        Request::BulkString(key.into()),
-    ]))
+    let mut req = Request::from_static("SUBSCRIBE");
+    for channel in channels {
+        req = req.add(Request::BulkString(channel.into()));
+    }
+    SubscribeOutputCommand(req)
+}
+
+/// UNSUBSCRIBE redis command
+pub fn UnSubscribe<T>(channels: Option<Vec<T>>) -> UnSubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    let mut req = Request::from_static("UNSUBSCRIBE");
+    if let Some(channels) = channels {
+        for channel in channels {
+            req = req.add(Request::BulkString(channel.into()));
+        }
+    }
+    UnSubscribeOutputCommand(req)
+}
+
+/// SSUBSCRIBE redis command
+pub fn SSubscribe<T>(channels: Vec<T>) -> SubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    let mut req = Request::from_static("SSUBSCRIBE");
+    for channel in channels {
+        req = req.add(Request::BulkString(channel.into()));
+    }
+    SubscribeOutputCommand(req)
+}
+
+/// SUNSUBSCRIBE redis command
+pub fn SUnSubscribe<T>(channels: Option<Vec<T>>) -> UnSubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    let mut req = Request::from_static("SUNSUBSCRIBE");
+    if let Some(channels) = channels {
+        for channel in channels {
+            req = req.add(Request::BulkString(channel.into()));
+        }
+    }
+    UnSubscribeOutputCommand(req)
+}
+
+/// PSUBSCRIBE redis command
+pub fn PSubscribe<T>(channels: Vec<T>) -> SubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    let mut req = Request::from_static("PSUBSCRIBE");
+    for channel in channels {
+        req = req.add(Request::BulkString(channel.into()));
+    }
+    SubscribeOutputCommand(req)
+}
+
+/// PUNSUBSCRIBE redis command
+pub fn PUnSubscribe<T>(channels: Option<Vec<T>>) -> UnSubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    let mut req = Request::from_static("PUNSUBSCRIBE");
+    if let Some(channels) = channels {
+        for channel in channels {
+            req = req.add(Request::BulkString(channel.into()));
+        }
+    }
+    UnSubscribeOutputCommand(req)
 }
 
 impl PubSubCommand for SubscribeOutputCommand {}

src/simple.rs

@@ -115,7 +115,7 @@ impl<U: Command + PubSubCommand> SubscriptionClient<U> {
     /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
     ///     let redis = RedisConnector::new("127.0.0.1:6379").connect_simple().await?;
     ///    
-    ///     let subscriber = redis.subscribe(cmd::Subscribe("test"))?;
+    ///     let subscriber = redis.subscribe(cmd::Subscribe(vec!["test"]))?;
     ///     // do some work
     ///
     ///     // go back to normal client