Skip to main content

kevy_client/
scan.rs

1//! Key-space iteration: `KEYS`, `SCAN`, `RANDOMKEY`.
2//!
3//! Embedded backend uses `kevy_embedded::Store::with(|inner| ...)` to
4//! reach `kevy_store::Store::collect_keys`. `SCAN` on embed is a single
5//! shot — cursor always advances 0 → 0 (i.e., all keys returned at once).
6//! Use `KEYS` directly if you want the same data; `SCAN` exists so the
7//! same code can also run against a Redis server, where iteration matters.
8
9use std::io;
10
11use kevy_resp::Reply;
12
13use crate::{Connection, array_to_bulks, string, unexpected};
14
15impl Connection {
16    /// `KEYS pattern` — every key matching `pattern` (glob: `*`, `?`,
17    /// `[abc]`). Use sparingly: O(N) over the whole keyspace.
18    pub fn keys(&mut self, pattern: &[u8]) -> io::Result<Vec<Vec<u8>>> {
19        match self {
20            Self::Embedded(s) => Ok(s.with(|inner| inner.collect_keys(Some(pattern), None))),
21            Self::Remote(c) => match c.request(&[b"KEYS".to_vec(), pattern.to_vec()])? {
22                Reply::Array(items) => array_to_bulks(items),
23                Reply::Error(e) => Err(io::Error::other(string(e))),
24                other => Err(unexpected(other)),
25            },
26        }
27    }
28
29    /// `SCAN cursor [MATCH pattern] [COUNT n]`. Returns `(next_cursor,
30    /// batch)`; iterate by re-calling with the returned cursor until
31    /// `next_cursor == 0`.
32    ///
33    /// On the embedded backend the iteration always finishes in one
34    /// call — `next_cursor` is `0` and `batch` holds every matching key.
35    /// `count` is taken as a hint only.
36    pub fn scan(
37        &mut self,
38        cursor: u64,
39        pattern: Option<&[u8]>,
40        count: Option<usize>,
41    ) -> io::Result<(u64, Vec<Vec<u8>>)> {
42        match self {
43            Self::Embedded(s) => {
44                // Embed has no real cursor: any non-zero cursor means the caller
45                // already drained on a previous call.
46                if cursor != 0 {
47                    return Ok((0, Vec::new()));
48                }
49                let batch = s.with(|inner| inner.collect_keys(pattern, count));
50                Ok((0, batch))
51            }
52            Self::Remote(c) => {
53                let mut args: Vec<Vec<u8>> =
54                    vec![b"SCAN".to_vec(), cursor.to_string().into_bytes()];
55                if let Some(pat) = pattern {
56                    args.push(b"MATCH".to_vec());
57                    args.push(pat.to_vec());
58                }
59                if let Some(n) = count {
60                    args.push(b"COUNT".to_vec());
61                    args.push(n.to_string().into_bytes());
62                }
63                match c.request(&args)? {
64                    Reply::Array(items) if items.len() == 2 => {
65                        let mut it = items.into_iter();
66                        let cursor_bulk = it.next().unwrap();
67                        let keys_arr = it.next().unwrap();
68                        let next_cursor = match cursor_bulk {
69                            Reply::Bulk(b) => std::str::from_utf8(&b)
70                                .map_err(|_| io::Error::other("non-utf8 SCAN cursor"))?
71                                .parse()
72                                .map_err(|_| io::Error::other("bad SCAN cursor"))?,
73                            other => return Err(unexpected(other)),
74                        };
75                        let keys = match keys_arr {
76                            Reply::Array(items) => array_to_bulks(items)?,
77                            other => return Err(unexpected(other)),
78                        };
79                        Ok((next_cursor, keys))
80                    }
81                    Reply::Error(e) => Err(io::Error::other(string(e))),
82                    other => Err(unexpected(other)),
83                }
84            }
85        }
86    }
87
88    /// `RANDOMKEY` — sample one key, or `None` if the keyspace is empty.
89    ///
90    /// Embed returns the lexicographically-first key (deterministic, no
91    /// RNG); the server returns a truly random key. Both honour empty.
92    pub fn randomkey(&mut self) -> io::Result<Option<Vec<u8>>> {
93        match self {
94            Self::Embedded(s) => Ok(s.with(|inner| {
95                inner.collect_keys(None, Some(1)).into_iter().next()
96            })),
97            Self::Remote(c) => match c.request(&[b"RANDOMKEY".to_vec()])? {
98                Reply::Bulk(v) => Ok(Some(v)),
99                Reply::Nil => Ok(None),
100                Reply::Error(e) => Err(io::Error::other(string(e))),
101                other => Err(unexpected(other)),
102            },
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn embedded_keys_matches_glob() {
113        let mut c = Connection::open("mem://").unwrap();
114        c.set(b"user:1", b"a").unwrap();
115        c.set(b"user:2", b"b").unwrap();
116        c.set(b"other", b"c").unwrap();
117        let mut keys = c.keys(b"user:*").unwrap();
118        keys.sort();
119        assert_eq!(keys, vec![b"user:1".to_vec(), b"user:2".to_vec()]);
120    }
121
122    #[test]
123    fn embedded_scan_returns_all_in_one_round() {
124        let mut c = Connection::open("mem://").unwrap();
125        for i in 0..5 {
126            c.set(format!("k{i}").as_bytes(), b"v").unwrap();
127        }
128        let (next, batch) = c.scan(0, None, None).unwrap();
129        assert_eq!(next, 0);
130        assert_eq!(batch.len(), 5);
131        // Any non-zero cursor means "we already finished" on embed.
132        let (next2, batch2) = c.scan(123, None, None).unwrap();
133        assert_eq!(next2, 0);
134        assert!(batch2.is_empty());
135    }
136
137    #[test]
138    fn embedded_randomkey_empty_and_present() {
139        let mut c = Connection::open("mem://").unwrap();
140        assert!(c.randomkey().unwrap().is_none());
141        c.set(b"only", b"x").unwrap();
142        assert_eq!(c.randomkey().unwrap(), Some(b"only".to_vec()));
143    }
144}