1use std::io;
10
11use kevy_resp::Reply;
12
13use crate::{Connection, array_to_bulks, string, unexpected};
14
15impl Connection {
16 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 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 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 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 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}