1use crate::{KevyError, KevyResult};
10
11use kevy_resp::Reply;
12use kevy_resp_client::RespClient;
13
14use crate::{Connection, array_to_bulks, store_err, string, unexpected};
15
16impl Connection {
17 pub fn hset(&mut self, key: &[u8], pairs: &[(&[u8], &[u8])]) -> KevyResult<usize> {
22 match self {
23 Self::Embedded(s) => s.hset(key, pairs),
24 Self::Remote(c) => {
25 let mut args: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
26 args.push(b"HSET");
27 args.push(key);
28 for &(f, v) in pairs {
29 args.push(f);
30 args.push(v);
31 }
32 match c.request_borrowed(&args)? {
33 Reply::Int(n) if n >= 0 => Ok(n as usize),
34 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
35 other => Err(unexpected(other)),
36 }
37 }
38 }
39 }
40
41 pub fn hget(&mut self, key: &[u8], field: &[u8]) -> KevyResult<Option<Vec<u8>>> {
43 match self {
44 Self::Embedded(s) => s.hget(key, field),
45 Self::Remote(c) => match c.request_borrowed(&[b"HGET", key, field])? {
46 Reply::Bulk(v) => Ok(Some(v)),
47 Reply::Nil => Ok(None),
48 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
49 other => Err(unexpected(other)),
50 },
51 }
52 }
53
54 pub fn hdel(&mut self, key: &[u8], fields: &[&[u8]]) -> KevyResult<usize> {
57 match self {
58 Self::Embedded(s) => s.hdel(key, fields),
59 Self::Remote(c) => {
60 let mut args: Vec<&[u8]> = Vec::with_capacity(fields.len() + 2);
61 args.push(b"HDEL");
62 args.push(key);
63 args.extend_from_slice(fields);
64 match c.request_borrowed(&args)? {
65 Reply::Int(n) if n >= 0 => Ok(n as usize),
66 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
67 other => Err(unexpected(other)),
68 }
69 }
70 }
71 }
72
73 pub fn hlen(&mut self, key: &[u8]) -> KevyResult<usize> {
75 match self {
76 Self::Embedded(s) => s.with(|inner| inner.hlen(key)).map_err(store_err),
77 Self::Remote(c) => match c.request_borrowed(&[b"HLEN", key])? {
78 Reply::Int(n) if n >= 0 => Ok(n as usize),
79 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
80 other => Err(unexpected(other)),
81 },
82 }
83 }
84
85 pub fn hgetall(&mut self, key: &[u8]) -> KevyResult<Vec<Vec<u8>>> {
88 match self {
89 Self::Embedded(s) => s.with(|inner| inner.hgetall(key)).map_err(store_err),
90 Self::Remote(c) => match c.request_borrowed(&[b"HGETALL", key])? {
91 Reply::Array(items) => array_to_bulks(items),
92 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
93 other => Err(unexpected(other)),
94 },
95 }
96 }
97
98 pub fn hkeys(&mut self, key: &[u8]) -> KevyResult<Vec<Vec<u8>>> {
100 match self {
101 Self::Embedded(s) => s.with(|inner| inner.hkeys(key)).map_err(store_err),
102 Self::Remote(c) => match c.request_borrowed(&[b"HKEYS", key])? {
103 Reply::Array(items) => array_to_bulks(items),
104 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
105 other => Err(unexpected(other)),
106 },
107 }
108 }
109
110 pub fn hvals(&mut self, key: &[u8]) -> KevyResult<Vec<Vec<u8>>> {
112 match self {
113 Self::Embedded(s) => s.with(|inner| inner.hvals(key)).map_err(store_err),
114 Self::Remote(c) => match c.request_borrowed(&[b"HVALS", key])? {
115 Reply::Array(items) => array_to_bulks(items),
116 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
117 other => Err(unexpected(other)),
118 },
119 }
120 }
121
122 pub fn lpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
126 match self {
127 Self::Embedded(s) => s.lpush(key, values),
128 Self::Remote(c) => list_push(c, b"LPUSH", key, values),
129 }
130 }
131
132 pub fn rpush(&mut self, key: &[u8], values: &[&[u8]]) -> KevyResult<usize> {
134 match self {
135 Self::Embedded(s) => s.rpush(key, values),
136 Self::Remote(c) => list_push(c, b"RPUSH", key, values),
137 }
138 }
139
140 pub fn lpop(&mut self, key: &[u8], count: usize) -> KevyResult<Vec<Vec<u8>>> {
143 match self {
144 Self::Embedded(s) => s.lpop(key, count),
145 Self::Remote(c) => list_pop(c, b"LPOP", key, count),
146 }
147 }
148
149 pub fn rpop(&mut self, key: &[u8], count: usize) -> KevyResult<Vec<Vec<u8>>> {
151 match self {
152 Self::Embedded(s) => s.rpop(key, count),
153 Self::Remote(c) => list_pop(c, b"RPOP", key, count),
154 }
155 }
156
157 pub fn llen(&mut self, key: &[u8]) -> KevyResult<usize> {
159 match self {
160 Self::Embedded(s) => s.llen(key),
161 Self::Remote(c) => match c.request_borrowed(&[b"LLEN", key])? {
162 Reply::Int(n) if n >= 0 => Ok(n as usize),
163 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
164 other => Err(unexpected(other)),
165 },
166 }
167 }
168
169 pub fn lrange(&mut self, key: &[u8], start: i64, stop: i64) -> KevyResult<Vec<Vec<u8>>> {
172 match self {
173 Self::Embedded(s) => s
174 .with(|inner| inner.lrange(key, start, stop))
175 .map_err(store_err),
176 Self::Remote(c) => {
177 let start_s = start.to_string();
178 let stop_s = stop.to_string();
179 match c.request_borrowed(&[b"LRANGE", key, start_s.as_bytes(), stop_s.as_bytes()])? {
180 Reply::Array(items) => array_to_bulks(items),
181 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
182 other => Err(unexpected(other)),
183 }
184 }
185 }
186 }
187
188 pub fn sadd(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
192 match self {
193 Self::Embedded(s) => s.sadd(key, members),
194 Self::Remote(c) => set_multi(c, b"SADD", key, members),
195 }
196 }
197
198 pub fn srem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
200 match self {
201 Self::Embedded(s) => s.srem(key, members),
202 Self::Remote(c) => set_multi(c, b"SREM", key, members),
203 }
204 }
205
206 pub fn smembers(&mut self, key: &[u8]) -> KevyResult<Vec<Vec<u8>>> {
208 match self {
209 Self::Embedded(s) => s.smembers(key),
210 Self::Remote(c) => match c.request_borrowed(&[b"SMEMBERS", key])? {
211 Reply::Array(items) => array_to_bulks(items),
212 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
213 other => Err(unexpected(other)),
214 },
215 }
216 }
217
218 pub fn scard(&mut self, key: &[u8]) -> KevyResult<usize> {
220 match self {
221 Self::Embedded(s) => s.scard(key),
222 Self::Remote(c) => match c.request_borrowed(&[b"SCARD", key])? {
223 Reply::Int(n) if n >= 0 => Ok(n as usize),
224 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
225 other => Err(unexpected(other)),
226 },
227 }
228 }
229
230 pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> KevyResult<bool> {
232 match self {
233 Self::Embedded(s) => s
234 .with(|inner| inner.sismember(key, member))
235 .map_err(store_err),
236 Self::Remote(c) => match c.request_borrowed(&[b"SISMEMBER", key, member])? {
237 Reply::Int(1) => Ok(true),
238 Reply::Int(0) => Ok(false),
239 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
240 other => Err(unexpected(other)),
241 },
242 }
243 }
244
245 pub fn sinter(&mut self, keys: &[&[u8]]) -> KevyResult<Vec<Vec<u8>>> {
247 match self {
248 Self::Embedded(s) => embed_set_combine(s, keys, SetOp::Inter),
249 Self::Remote(c) => remote_set_combine(c, b"SINTER", keys),
250 }
251 }
252
253 pub fn sunion(&mut self, keys: &[&[u8]]) -> KevyResult<Vec<Vec<u8>>> {
255 match self {
256 Self::Embedded(s) => embed_set_combine(s, keys, SetOp::Union),
257 Self::Remote(c) => remote_set_combine(c, b"SUNION", keys),
258 }
259 }
260
261 pub fn sdiff(&mut self, keys: &[&[u8]]) -> KevyResult<Vec<Vec<u8>>> {
263 match self {
264 Self::Embedded(s) => embed_set_combine(s, keys, SetOp::Diff),
265 Self::Remote(c) => remote_set_combine(c, b"SDIFF", keys),
266 }
267 }
268
269 pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, &[u8])]) -> KevyResult<usize> {
274 match self {
275 Self::Embedded(s) => s.zadd(key, pairs),
276 Self::Remote(c) => {
277 let scores: Vec<String> = pairs.iter().map(|(s, _)| s.to_string()).collect();
279 let mut args: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
280 args.push(b"ZADD");
281 args.push(key);
282 for (i, &(_, m)) in pairs.iter().enumerate() {
283 args.push(scores[i].as_bytes());
284 args.push(m);
285 }
286 match c.request_borrowed(&args)? {
287 Reply::Int(n) if n >= 0 => Ok(n as usize),
288 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
289 other => Err(unexpected(other)),
290 }
291 }
292 }
293 }
294
295 pub fn zrem(&mut self, key: &[u8], members: &[&[u8]]) -> KevyResult<usize> {
297 match self {
298 Self::Embedded(s) => s.zrem(key, members),
299 Self::Remote(c) => set_multi(c, b"ZREM", key, members),
300 }
301 }
302
303 pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> KevyResult<Option<f64>> {
305 match self {
306 Self::Embedded(s) => s.zscore(key, member),
307 Self::Remote(c) => match c.request_borrowed(&[b"ZSCORE", key, member])? {
308 Reply::Bulk(v) => {
309 let s = std::str::from_utf8(&v)
310 .map_err(|_| KevyError::Protocol("non-utf8 score reply".into()))?;
311 let n: f64 = s
312 .parse()
313 .map_err(|_| KevyError::Protocol(format!("bad score float: {s}")))?;
314 Ok(Some(n))
315 }
316 Reply::Nil => Ok(None),
317 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
318 other => Err(unexpected(other)),
319 },
320 }
321 }
322
323 pub fn zcard(&mut self, key: &[u8]) -> KevyResult<usize> {
325 match self {
326 Self::Embedded(s) => s.zcard(key),
327 Self::Remote(c) => match c.request_borrowed(&[b"ZCARD", key])? {
328 Reply::Int(n) if n >= 0 => Ok(n as usize),
329 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
330 other => Err(unexpected(other)),
331 },
332 }
333 }
334
335 pub fn zrange(&mut self, key: &[u8], start: i64, stop: i64) -> KevyResult<Vec<Vec<u8>>> {
338 match self {
339 Self::Embedded(s) => s
340 .with(|inner| inner.zrange(key, start, stop))
341 .map(|pairs| pairs.into_iter().map(|(m, _score)| m).collect())
342 .map_err(store_err),
343 Self::Remote(c) => {
344 let start_s = start.to_string();
345 let stop_s = stop.to_string();
346 match c.request_borrowed(&[b"ZRANGE", key, start_s.as_bytes(), stop_s.as_bytes()])? {
347 Reply::Array(items) => array_to_bulks(items),
348 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
349 other => Err(unexpected(other)),
350 }
351 }
352 }
353 }
354}
355
356pub(crate) fn list_push(
362 c: &mut RespClient,
363 verb: &[u8],
364 key: &[u8],
365 values: &[&[u8]],
366) -> KevyResult<usize> {
367 let mut args: Vec<&[u8]> = Vec::with_capacity(values.len() + 2);
368 args.push(verb);
369 args.push(key);
370 args.extend_from_slice(values);
371 match c.request_borrowed(&args)? {
372 Reply::Int(n) if n >= 0 => Ok(n as usize),
373 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
374 other => Err(unexpected(other)),
375 }
376}
377
378pub(crate) fn list_pop(
379 c: &mut RespClient,
380 verb: &[u8],
381 key: &[u8],
382 count: usize,
383) -> KevyResult<Vec<Vec<u8>>> {
384 let count_s = count.to_string();
385 match c.request_borrowed(&[verb, key, count_s.as_bytes()])? {
386 Reply::Array(items) => array_to_bulks(items),
387 Reply::Bulk(v) => Ok(vec![v]),
388 Reply::Nil => Ok(Vec::new()),
389 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
390 other => Err(unexpected(other)),
391 }
392}
393
394pub(crate) fn set_multi(
395 c: &mut RespClient,
396 verb: &[u8],
397 key: &[u8],
398 members: &[&[u8]],
399) -> KevyResult<usize> {
400 let mut args: Vec<&[u8]> = Vec::with_capacity(members.len() + 2);
401 args.push(verb);
402 args.push(key);
403 args.extend_from_slice(members);
404 match c.request_borrowed(&args)? {
405 Reply::Int(n) if n >= 0 => Ok(n as usize),
406 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
407 other => Err(unexpected(other)),
408 }
409}
410
411#[derive(Clone, Copy)]
415enum SetOp {
416 Inter,
417 Union,
418 Diff,
419}
420
421fn embed_set_combine(
422 s: &kevy_embedded::Store,
423 keys: &[&[u8]],
424 op: SetOp,
425) -> KevyResult<Vec<Vec<u8>>> {
426 use std::collections::HashSet;
427 if keys.is_empty() {
428 return Ok(Vec::new());
429 }
430 let snapshots: Vec<Vec<Vec<u8>>> = keys
431 .iter()
432 .map(|k| s.smembers(k))
433 .collect::<KevyResult<_>>()?;
434 let mut iter = snapshots.into_iter();
435 let mut acc: HashSet<Vec<u8>> = iter.next().unwrap_or_default().into_iter().collect();
436 for rest in iter {
437 let other: HashSet<Vec<u8>> = rest.into_iter().collect();
438 acc = match op {
439 SetOp::Inter => acc.intersection(&other).cloned().collect(),
440 SetOp::Union => acc.union(&other).cloned().collect(),
441 SetOp::Diff => acc.difference(&other).cloned().collect(),
442 };
443 }
444 Ok(acc.into_iter().collect())
445}
446
447pub(crate) fn remote_set_combine(
448 c: &mut RespClient,
449 verb: &[u8],
450 keys: &[&[u8]],
451) -> KevyResult<Vec<Vec<u8>>> {
452 let mut args: Vec<&[u8]> = Vec::with_capacity(keys.len() + 1);
453 args.push(verb);
454 args.extend_from_slice(keys);
455 match c.request_borrowed(&args)? {
456 Reply::Array(items) => array_to_bulks(items),
457 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
458 other => Err(unexpected(other)),
459 }
460}
461
462#[cfg(test)]
463#[path = "collections_tests.rs"]
464mod tests;