1use crate::ValkeyResult;
12use crate::cmd::Cmd;
13use crate::executor::CommandExecutor;
14use crate::value::FromValkeyValue;
15use crate::value::ValkeyValue;
16use crate::value::to_glide_error;
17use crate::write::ToValkeyArgs;
18use async_trait::async_trait;
19use bytes::Bytes;
20
21#[async_trait]
24pub trait PubSubCommands: CommandExecutor {
25 async fn subscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
36 self.pubsub_subscribe_impl("SUBSCRIBE", channels).await
37 }
38
39 async fn unsubscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
42 self.pubsub_subscribe_impl("UNSUBSCRIBE", channels).await
43 }
44
45 async fn psubscribe<C: ToValkeyArgs + Send + Sync>(&self, patterns: &[C]) -> ValkeyResult<()> {
48 self.pubsub_subscribe_impl("PSUBSCRIBE", patterns).await
49 }
50
51 async fn punsubscribe<C: ToValkeyArgs + Send + Sync>(
54 &self,
55 patterns: &[C],
56 ) -> ValkeyResult<()> {
57 self.pubsub_subscribe_impl("PUNSUBSCRIBE", patterns).await
58 }
59
60 async fn ssubscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
63 self.pubsub_subscribe_impl("SSUBSCRIBE", channels).await
64 }
65
66 async fn sunsubscribe<C: ToValkeyArgs + Send + Sync>(
69 &self,
70 channels: &[C],
71 ) -> ValkeyResult<()> {
72 self.pubsub_subscribe_impl("SUNSUBSCRIBE", channels).await
73 }
74
75 #[doc(hidden)]
76 async fn pubsub_subscribe_impl<C: ToValkeyArgs + Send + Sync>(
77 &self,
78 keyword: &'static str,
79 channels: &[C],
80 ) -> ValkeyResult<()> {
81 let mut cmd = Cmd::new();
82 cmd.arg(keyword);
83 for c in channels {
84 cmd.arg(c);
85 }
86 self.execute_command(cmd, None).await?;
87 Ok(())
88 }
89
90 async fn pubsub_channels(&self, pattern: Option<&[u8]>) -> ValkeyResult<Vec<Bytes>> {
95 self.pubsub_channels_impl("CHANNELS", pattern).await
96 }
97
98 async fn pubsub_shardchannels(&self, pattern: Option<&[u8]>) -> ValkeyResult<Vec<Bytes>> {
101 self.pubsub_channels_impl("SHARDCHANNELS", pattern).await
102 }
103
104 async fn pubsub_numpat(&self) -> ValkeyResult<i64> {
106 let mut cmd = Cmd::new();
107 cmd.arg("PUBSUB").arg("NUMPAT");
108 i64::from_owned_valkey_value(self.execute_command(cmd, None).await?)
109 }
110
111 async fn pubsub_numsub<C: ToValkeyArgs + Send + Sync>(
113 &self,
114 channels: &[C],
115 ) -> ValkeyResult<Vec<(Bytes, i64)>> {
116 self.pubsub_numsub_impl("NUMSUB", channels).await
117 }
118
119 async fn pubsub_shardnumsub<C: ToValkeyArgs + Send + Sync>(
121 &self,
122 channels: &[C],
123 ) -> ValkeyResult<Vec<(Bytes, i64)>> {
124 self.pubsub_numsub_impl("SHARDNUMSUB", channels).await
125 }
126
127 #[doc(hidden)]
128 async fn pubsub_channels_impl(
129 &self,
130 sub: &'static str,
131 pattern: Option<&[u8]>,
132 ) -> ValkeyResult<Vec<Bytes>> {
133 let mut cmd = Cmd::new();
134 cmd.arg("PUBSUB").arg(sub);
135 if let Some(p) = pattern {
136 cmd.arg(p);
137 }
138 match self.execute_command(cmd, None).await? {
139 ValkeyValue::Array(items) => items
140 .into_iter()
141 .map(Bytes::from_owned_valkey_value)
142 .collect(),
143 ValkeyValue::Nil => Ok(Vec::new()),
144 other => Ok(vec![Bytes::from_owned_valkey_value(other)?]),
145 }
146 }
147
148 #[doc(hidden)]
149 async fn pubsub_numsub_impl<C: ToValkeyArgs + Send + Sync>(
150 &self,
151 sub: &'static str,
152 channels: &[C],
153 ) -> ValkeyResult<Vec<(Bytes, i64)>> {
154 let mut cmd = Cmd::new();
155 cmd.arg("PUBSUB").arg(sub);
156 for c in channels {
157 cmd.arg(c);
158 }
159 parse_numsub(self.execute_command(cmd, None).await?)
160 }
161}
162
163fn parse_numsub(v: ValkeyValue) -> ValkeyResult<Vec<(Bytes, i64)>> {
165 match v {
166 ValkeyValue::Nil => Ok(Vec::new()),
167 ValkeyValue::Map(pairs) => pairs
168 .into_iter()
169 .map(|(c, n)| {
170 Ok((
171 Bytes::from_owned_valkey_value(c)?,
172 i64::from_owned_valkey_value(n)?,
173 ))
174 })
175 .collect(),
176 ValkeyValue::Array(items) => {
177 let mut out = Vec::with_capacity(items.len() / 2);
178 let mut iter = items.into_iter();
179 while let (Some(c), Some(n)) = (iter.next(), iter.next()) {
180 out.push((
181 Bytes::from_owned_valkey_value(c)?,
182 i64::from_owned_valkey_value(n)?,
183 ));
184 }
185 Ok(out)
186 }
187 other => Err(to_glide_error(other, "Unexpected PUBSUB NUMSUB reply.")),
188 }
189}
190
191impl<T: CommandExecutor + ?Sized> PubSubCommands for T {}