1use std::io;
9
10use kevy_resp::Reply;
11
12use crate::cluster::ClusterClient;
13use crate::collections::{list_pop, list_push, remote_set_combine, set_multi};
14use crate::{array_to_bulks, string, unexpected, vec2, vec3};
15
16fn count(reply: Reply) -> io::Result<usize> {
19 match reply {
20 Reply::Int(n) if n >= 0 => Ok(n as usize),
21 Reply::Error(e) => Err(io::Error::other(string(e))),
22 other => Err(unexpected(other)),
23 }
24}
25
26fn bulks(reply: Reply) -> io::Result<Vec<Vec<u8>>> {
28 match reply {
29 Reply::Array(items) => array_to_bulks(items),
30 Reply::Error(e) => Err(io::Error::other(string(e))),
31 other => Err(unexpected(other)),
32 }
33}
34
35impl ClusterClient {
36 pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> io::Result<usize> {
40 let mut args = Vec::with_capacity(2 + pairs.len() * 2);
41 args.push(b"HSET".to_vec());
42 args.push(key.to_vec());
43 for (f, v) in pairs {
44 args.push(f.to_vec());
45 args.push(v.to_vec());
46 }
47 count(self.request_keyed(key, &args)?)
48 }
49
50 pub fn hget(&mut self, key: &[u8], field: &[u8]) -> io::Result<Option<Vec<u8>>> {
52 match self.request_keyed(key, &vec3(b"HGET", key, field))? {
53 Reply::Bulk(v) => Ok(Some(v)),
54 Reply::Nil => Ok(None),
55 Reply::Error(e) => Err(io::Error::other(string(e))),
56 other => Err(unexpected(other)),
57 }
58 }
59
60 pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> io::Result<usize> {
62 let mut args = Vec::with_capacity(fields.len() + 2);
63 args.push(b"HDEL".to_vec());
64 args.push(key.to_vec());
65 args.extend(fields.iter().map(|f| f.to_vec()));
66 count(self.request_keyed(key, &args)?)
67 }
68
69 pub fn hlen(&mut self, key: &[u8]) -> io::Result<usize> {
71 count(self.request_keyed(key, &vec2(b"HLEN", key))?)
72 }
73
74 pub fn hgetall(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
76 bulks(self.request_keyed(key, &vec2(b"HGETALL", key))?)
77 }
78
79 pub fn hkeys(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
81 bulks(self.request_keyed(key, &vec2(b"HKEYS", key))?)
82 }
83
84 pub fn hvals(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
86 bulks(self.request_keyed(key, &vec2(b"HVALS", key))?)
87 }
88
89 pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
93 list_push(self.route_mut(key), b"LPUSH", key, values)
94 }
95
96 pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> io::Result<usize> {
98 list_push(self.route_mut(key), b"RPUSH", key, values)
99 }
100
101 pub fn lpop(&mut self, key: &[u8], n: usize) -> io::Result<Vec<Vec<u8>>> {
103 list_pop(self.route_mut(key), b"LPOP", key, n)
104 }
105
106 pub fn rpop(&mut self, key: &[u8], n: usize) -> io::Result<Vec<Vec<u8>>> {
108 list_pop(self.route_mut(key), b"RPOP", key, n)
109 }
110
111 pub fn llen(&mut self, key: &[u8]) -> io::Result<usize> {
113 count(self.request_keyed(key, &vec2(b"LLEN", key))?)
114 }
115
116 pub fn lrange(&mut self, key: &[u8], start: i64, stop: i64) -> io::Result<Vec<Vec<u8>>> {
118 let args = vec![
119 b"LRANGE".to_vec(),
120 key.to_vec(),
121 start.to_string().into_bytes(),
122 stop.to_string().into_bytes(),
123 ];
124 bulks(self.request_keyed(key, &args)?)
125 }
126
127 pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
131 set_multi(self.route_mut(key), b"SADD", key, members)
132 }
133
134 pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
136 set_multi(self.route_mut(key), b"SREM", key, members)
137 }
138
139 pub fn smembers(&mut self, key: &[u8]) -> io::Result<Vec<Vec<u8>>> {
141 bulks(self.request_keyed(key, &vec2(b"SMEMBERS", key))?)
142 }
143
144 pub fn scard(&mut self, key: &[u8]) -> io::Result<usize> {
146 count(self.request_keyed(key, &vec2(b"SCARD", key))?)
147 }
148
149 pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> io::Result<bool> {
151 match self.request_keyed(key, &vec3(b"SISMEMBER", key, member))? {
152 Reply::Int(1) => Ok(true),
153 Reply::Int(0) => Ok(false),
154 Reply::Error(e) => Err(io::Error::other(string(e))),
155 other => Err(unexpected(other)),
156 }
157 }
158
159 pub fn sinter(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
161 self.set_combine(b"SINTER", keys)
162 }
163
164 pub fn sunion(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
166 self.set_combine(b"SUNION", keys)
167 }
168
169 pub fn sdiff(&mut self, keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
171 self.set_combine(b"SDIFF", keys)
172 }
173
174 fn set_combine(&mut self, verb: &[u8], keys: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
175 let Some(first) = keys.first() else {
176 return Ok(Vec::new());
177 };
178 remote_set_combine(self.route_mut(first), verb, keys)
179 }
180
181 pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> io::Result<usize> {
185 let mut args = Vec::with_capacity(2 + pairs.len() * 2);
186 args.push(b"ZADD".to_vec());
187 args.push(key.to_vec());
188 for (score, m) in pairs {
189 args.push(score.to_string().into_bytes());
190 args.push(m.to_vec());
191 }
192 count(self.request_keyed(key, &args)?)
193 }
194
195 pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> io::Result<usize> {
197 set_multi(self.route_mut(key), b"ZREM", key, members)
198 }
199
200 pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> io::Result<Option<f64>> {
202 match self.request_keyed(key, &vec3(b"ZSCORE", key, member))? {
203 Reply::Bulk(v) => {
204 let s = std::str::from_utf8(&v)
205 .map_err(|_| io::Error::other("non-utf8 score reply"))?;
206 let n: f64 = s
207 .parse()
208 .map_err(|_| io::Error::other(format!("bad score float: {s}")))?;
209 Ok(Some(n))
210 }
211 Reply::Nil => Ok(None),
212 Reply::Error(e) => Err(io::Error::other(string(e))),
213 other => Err(unexpected(other)),
214 }
215 }
216
217 pub fn zcard(&mut self, key: &[u8]) -> io::Result<usize> {
219 count(self.request_keyed(key, &vec2(b"ZCARD", key))?)
220 }
221
222 pub fn zrange(&mut self, key: &[u8], start: i64, stop: i64) -> io::Result<Vec<Vec<u8>>> {
224 let args = vec![
225 b"ZRANGE".to_vec(),
226 key.to_vec(),
227 start.to_string().into_bytes(),
228 stop.to_string().into_bytes(),
229 ];
230 bulks(self.request_keyed(key, &args)?)
231 }
232}