Skip to main content

kevy_client/
collections.rs

1//! Connection methods for the four collection-typed Redis data types:
2//! hash, list, set, sorted set. Plus the small `LPUSH`/`SADD`-style
3//! request builders shared between them.
4//!
5//! Lives in its own module so `lib.rs` stays focused on the `Connection`
6//! enum + open + the generic + string ops. Behaviour and API are
7//! unchanged from the original single-file layout.
8
9use 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    // ===== Hash =====
18
19    /// `HSET key field value [field value ...]`. Returns the number of
20    /// fields that were newly added (not overwrites).
21    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    /// `HGET key field`. `None` when the key or field is absent.
42    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    /// `HDEL key field [field ...]`. Returns the number of fields actually
55    /// removed.
56    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    /// `HLEN key`. Number of fields in the hash (0 if absent).
74    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    /// `HGETALL key`. Returns a flat `[f0, v0, f1, v1, ...]` matching the
86    /// Redis wire shape; empty when the key is absent.
87    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    /// `HKEYS key`. Returns the hash's field names (empty if absent).
99    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    /// `HVALS key`. Returns the hash's values (empty if absent).
111    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    // ===== List =====
123
124    /// `LPUSH key value [value ...]`. Returns the new list length.
125    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    /// `RPUSH key value [value ...]`. Returns the new list length.
133    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    /// `LPOP key count`. Returns up to `count` values from the head; empty
141    /// when the key is absent or already drained.
142    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    /// `RPOP key count`. Symmetric to [`Self::lpop`] from the tail.
150    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    /// `LLEN key`. 0 when the key is absent.
158    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    /// `LRANGE key start stop`. Redis-style indexing — negative offsets
170    /// count from the tail (`-1` = last element).
171    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    // ===== Set =====
189
190    /// `SADD key member [member ...]`. Returns count of newly added members.
191    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    /// `SREM key member [member ...]`. Returns count of removed members.
199    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    /// `SMEMBERS key`. Order is implementation-defined; empty if absent.
207    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    /// `SCARD key`. 0 when the key is absent.
219    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    /// `SISMEMBER key member`. `false` when key or member absent.
231    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    /// `SINTER key [key ...]` — intersection of all sets.
246    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    /// `SUNION key [key ...]` — union of all sets.
254    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    /// `SDIFF key [key ...]` — members of the first set absent from the rest.
262    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    // ===== Sorted set =====
270
271    /// `ZADD key score member [score member ...]`. Returns count of newly
272    /// added members (overwrites don't count).
273    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                // Scores are formatted floats — owned in `scores`; members borrow.
278                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    /// `ZREM key member [member ...]`. Returns count of removed members.
296    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    /// `ZSCORE key member`. `None` if absent.
304    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    /// `ZCARD key`. Number of members; 0 if absent.
324    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    /// `ZRANGE key start stop`. Ascending-score order; negative indices
336    /// count from the tail.
337    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
356// ─────────────────────────────────────────────────────────────────────────
357// Shared request builders. Both backends accept a slice of byte-slices,
358// but the RESP path needs to splat them into a single argv vector.
359// ─────────────────────────────────────────────────────────────────────────
360
361pub(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// Set-combine plumbing: each backend's path computes the intersection /
412// union / difference of N sets identified by `keys`.
413
414#[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;