1 use std::convert::TryFrom;
2 use std::os::raw::c_void;
3 use std::ptr;
4 use std::str::Utf8Error;
5 use std::time::Duration;
6
7 use libc::size_t;
8
9 use raw::KeyType;
10
11 use crate::from_byte_string;
12 use crate::native_types::RedisType;
13 use crate::raw;
14 use crate::redismodule::REDIS_OK;
15 use crate::RedisError;
16 use crate::RedisResult;
17 use crate::RedisString;
18
19 /// `RedisKey` is an abstraction over a Redis key that allows readonly
20 /// operations.
21 ///
22 /// Its primary function is to ensure the proper deallocation of resources when
23 /// it goes out of scope. Redis normally requires that keys be managed manually
24 /// by explicitly freeing them when you're done. This can be a risky prospect,
25 /// especially with mechanics like Rust's `?` operator, so we ensure fault-free
26 /// operation through the use of the Drop trait.
27 #[derive(Clone, Copy, Debug, PartialEq)]
28 pub enum KeyMode {
29 Read,
30 ReadWrite,
31 }
32
33 #[derive(Debug)]
34 pub struct RedisKey {
35 ctx: *mut raw::RedisModuleCtx,
36 key_inner: *mut raw::RedisModuleKey,
37 }
38
39 impl RedisKey {
open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> RedisKey40 pub fn open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> RedisKey {
41 let key_inner = raw::open_key(ctx, key.inner, to_raw_mode(KeyMode::Read));
42 RedisKey { ctx, key_inner }
43 }
44
45 /// # Panics
46 ///
47 /// Will panic if `RedisModule_ModuleTypeGetValue` is missing in redismodule.h
get_value<T>(&self, redis_type: &RedisType) -> Result<Option<&T>, RedisError>48 pub fn get_value<T>(&self, redis_type: &RedisType) -> Result<Option<&T>, RedisError> {
49 verify_type(self.key_inner, redis_type)?;
50
51 let value =
52 unsafe { raw::RedisModule_ModuleTypeGetValue.unwrap()(self.key_inner).cast::<T>() };
53
54 if value.is_null() {
55 return Ok(None);
56 }
57
58 let value = unsafe { &*value };
59
60 Ok(Some(value))
61 }
62
63 /// # Panics
64 ///
65 /// Will panic if `RedisModule_KeyType` is missing in redismodule.h
key_type(&self) -> raw::KeyType66 pub fn key_type(&self) -> raw::KeyType {
67 unsafe { raw::RedisModule_KeyType.unwrap()(self.key_inner) }.into()
68 }
69
70 /// Detects whether the key pointer given to us by Redis is null.
is_null(&self) -> bool71 pub fn is_null(&self) -> bool {
72 let null_key: *mut raw::RedisModuleKey = ptr::null_mut();
73 self.key_inner == null_key
74 }
75
read(&self) -> Result<Option<String>, RedisError>76 pub fn read(&self) -> Result<Option<String>, RedisError> {
77 let val = if self.is_null() {
78 None
79 } else {
80 Some(read_key(self.key_inner)?)
81 };
82 Ok(val)
83 }
84
hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError>85 pub fn hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError> {
86 let val = if self.is_null() {
87 None
88 } else {
89 hash_mget_key(self.ctx, self.key_inner, &[field])?
90 .pop()
91 .expect("hash_mget_key should return vector of same length as input")
92 };
93 Ok(val)
94 }
95
96 /// Returns the values associated with the specified fields in the hash stored at this key.
97 /// The result will be `None` if the key does not exist.
hash_get_multi<'a, A, B>( &self, fields: &'a [A], ) -> Result<Option<HMGetResult<'a, A, B>>, RedisError> where A: Into<Vec<u8>> + Clone, RedisString: Into<B>,98 pub fn hash_get_multi<'a, A, B>(
99 &self,
100 fields: &'a [A],
101 ) -> Result<Option<HMGetResult<'a, A, B>>, RedisError>
102 where
103 A: Into<Vec<u8>> + Clone,
104 RedisString: Into<B>,
105 {
106 let val = if self.is_null() {
107 None
108 } else {
109 Some(HMGetResult {
110 fields,
111 values: hash_mget_key(self.ctx, self.key_inner, fields)?,
112 phantom: std::marker::PhantomData,
113 })
114 };
115 Ok(val)
116 }
117 }
118
119 impl Drop for RedisKey {
120 // Frees resources appropriately as a RedisKey goes out of scope.
drop(&mut self)121 fn drop(&mut self) {
122 raw::close_key(self.key_inner);
123 }
124 }
125
126 /// `RedisKeyWritable` is an abstraction over a Redis key that allows read and
127 /// write operations.
128 pub struct RedisKeyWritable {
129 ctx: *mut raw::RedisModuleCtx,
130 key_inner: *mut raw::RedisModuleKey,
131 }
132
133 impl RedisKeyWritable {
open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> RedisKeyWritable134 pub fn open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> RedisKeyWritable {
135 let key_inner = raw::open_key(ctx, key.inner, to_raw_mode(KeyMode::ReadWrite));
136 RedisKeyWritable { ctx, key_inner }
137 }
138
139 /// Detects whether the value stored in a Redis key is empty.
140 ///
141 /// Note that an empty key can be reliably detected by looking for a null
142 /// as you open the key in read mode, but when asking for write Redis
143 /// returns a non-null pointer to allow us to write to even an empty key,
144 /// so we have to check the key's value instead.
145 /*
146 fn is_empty_old(&self) -> Result<bool, Error> {
147 match self.read()? {
148 Some(s) => match s.as_str() {
149 "" => Ok(true),
150 _ => Ok(false),
151 },
152 _ => Ok(false),
153 }
154 }
155 */
156
read(&self) -> Result<Option<String>, RedisError>157 pub fn read(&self) -> Result<Option<String>, RedisError> {
158 Ok(Some(read_key(self.key_inner)?))
159 }
160
hash_set(&self, field: &str, value: RedisString) -> raw::Status161 pub fn hash_set(&self, field: &str, value: RedisString) -> raw::Status {
162 raw::hash_set(self.key_inner, field, value.inner)
163 }
164
hash_del(&self, field: &str) -> raw::Status165 pub fn hash_del(&self, field: &str) -> raw::Status {
166 raw::hash_del(self.key_inner, field)
167 }
168
hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError>169 pub fn hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError> {
170 Ok(hash_mget_key(self.ctx, self.key_inner, &[field])?
171 .pop()
172 .expect("hash_mget_key should return vector of same length as input"))
173 }
174
175 /// Returns the values associated with the specified fields in the hash stored at this key.
hash_get_multi<'a, A, B>( &self, fields: &'a [A], ) -> Result<HMGetResult<'a, A, B>, RedisError> where A: Into<Vec<u8>> + Clone, RedisString: Into<B>,176 pub fn hash_get_multi<'a, A, B>(
177 &self,
178 fields: &'a [A],
179 ) -> Result<HMGetResult<'a, A, B>, RedisError>
180 where
181 A: Into<Vec<u8>> + Clone,
182 RedisString: Into<B>,
183 {
184 Ok(HMGetResult {
185 fields,
186 values: hash_mget_key(self.ctx, self.key_inner, fields)?,
187 phantom: std::marker::PhantomData,
188 })
189 }
190
191 // `list_push_head` inserts the specified element at the head of the list stored at this key.
list_push_head(&self, element: RedisString) -> raw::Status192 pub fn list_push_head(&self, element: RedisString) -> raw::Status {
193 raw::list_push(self.key_inner, raw::Where::ListHead, element.inner)
194 }
195
196 // `list_push_tail` inserts the specified element at the tail of the list stored at this key.
list_push_tail(&self, element: RedisString) -> raw::Status197 pub fn list_push_tail(&self, element: RedisString) -> raw::Status {
198 raw::list_push(self.key_inner, raw::Where::ListTail, element.inner)
199 }
200
201 // `list_pop_head` pops and returns the first element of the list.
202 // Returns None when:
203 // 1. The list is empty.
204 // 2. The key is not a list.
list_pop_head(&self) -> Option<RedisString>205 pub fn list_pop_head(&self) -> Option<RedisString> {
206 let ptr = raw::list_pop(self.key_inner, raw::Where::ListHead);
207
208 if ptr.is_null() {
209 return None;
210 }
211
212 Some(RedisString::new(self.ctx, ptr))
213 }
214
215 // `list_pop_head` pops and returns the last element of the list.
216 // Returns None when:
217 // 1. The list is empty.
218 // 2. The key is not a list.
list_pop_tail(&self) -> Option<RedisString>219 pub fn list_pop_tail(&self) -> Option<RedisString> {
220 let ptr = raw::list_pop(self.key_inner, raw::Where::ListTail);
221
222 if ptr.is_null() {
223 return None;
224 }
225
226 Some(RedisString::new(self.ctx, ptr))
227 }
228
set_expire(&self, expire: Duration) -> RedisResult229 pub fn set_expire(&self, expire: Duration) -> RedisResult {
230 let exp_millis = expire.as_millis();
231
232 let exp_time = i64::try_from(exp_millis).map_err(|_| {
233 RedisError::String(format!(
234 "Error expire duration {} is not allowed",
235 exp_millis
236 ))
237 })?;
238
239 match raw::set_expire(self.key_inner, exp_time) {
240 raw::Status::Ok => REDIS_OK,
241
242 // Error may occur if the key wasn't open for writing or is an
243 // empty key.
244 raw::Status::Err => Err(RedisError::Str("Error while setting key expire")),
245 }
246 }
247
write(&self, val: &str) -> RedisResult248 pub fn write(&self, val: &str) -> RedisResult {
249 let val_str = RedisString::create(self.ctx, val);
250 match raw::string_set(self.key_inner, val_str.inner) {
251 raw::Status::Ok => REDIS_OK,
252 raw::Status::Err => Err(RedisError::Str("Error while setting key")),
253 }
254 }
255
256 /// # Panics
257 ///
258 /// Will panic if `RedisModule_DeleteKey` is missing in redismodule.h
delete(&self) -> RedisResult259 pub fn delete(&self) -> RedisResult {
260 unsafe { raw::RedisModule_DeleteKey.unwrap()(self.key_inner) };
261 REDIS_OK
262 }
263
264 /// # Panics
265 ///
266 /// Will panic if `RedisModule_KeyType` is missing in redismodule.h
key_type(&self) -> raw::KeyType267 pub fn key_type(&self) -> raw::KeyType {
268 unsafe { raw::RedisModule_KeyType.unwrap()(self.key_inner) }.into()
269 }
270
is_empty(&self) -> bool271 pub fn is_empty(&self) -> bool {
272 self.key_type() == KeyType::Empty
273 }
274
open_with_redis_string( ctx: *mut raw::RedisModuleCtx, key: *mut raw::RedisModuleString, ) -> RedisKeyWritable275 pub fn open_with_redis_string(
276 ctx: *mut raw::RedisModuleCtx,
277 key: *mut raw::RedisModuleString,
278 ) -> RedisKeyWritable {
279 let key_inner = raw::open_key(ctx, key, to_raw_mode(KeyMode::ReadWrite));
280 RedisKeyWritable { ctx, key_inner }
281 }
282
283 /// # Panics
284 ///
285 /// Will panic if `RedisModule_ModuleTypeGetValue` is missing in redismodule.h
get_value<'a, 'b, T>( &'a self, redis_type: &RedisType, ) -> Result<Option<&'b mut T>, RedisError>286 pub fn get_value<'a, 'b, T>(
287 &'a self,
288 redis_type: &RedisType,
289 ) -> Result<Option<&'b mut T>, RedisError> {
290 verify_type(self.key_inner, redis_type)?;
291 let value =
292 unsafe { raw::RedisModule_ModuleTypeGetValue.unwrap()(self.key_inner).cast::<T>() };
293
294 if value.is_null() {
295 return Ok(None);
296 }
297
298 let value = unsafe { &mut *value };
299 Ok(Some(value))
300 }
301
302 /// # Panics
303 ///
304 /// Will panic if `RedisModule_ModuleTypeSetValue` is missing in redismodule.h
set_value<T>(&self, redis_type: &RedisType, value: T) -> Result<(), RedisError>305 pub fn set_value<T>(&self, redis_type: &RedisType, value: T) -> Result<(), RedisError> {
306 verify_type(self.key_inner, redis_type)?;
307 let value = Box::into_raw(Box::new(value)).cast::<c_void>();
308 let status: raw::Status = unsafe {
309 raw::RedisModule_ModuleTypeSetValue.unwrap()(
310 self.key_inner,
311 *redis_type.raw_type.borrow(),
312 value,
313 )
314 }
315 .into();
316
317 status.into()
318 }
319 }
320
321 /// Opaque type used to hold multi-get results. Use the provided methods to convert
322 /// the results into the desired type of Rust collection.
323 pub struct HMGetResult<'a, A, B>
324 where
325 A: Into<Vec<u8>> + Clone,
326 RedisString: Into<B>,
327 {
328 fields: &'a [A],
329 values: Vec<Option<RedisString>>,
330 phantom: std::marker::PhantomData<B>,
331 }
332
333 pub struct HMGetIter<'a, A, B>
334 where
335 A: Into<Vec<u8>>,
336 RedisString: Into<B>,
337 {
338 fields_iter: std::slice::Iter<'a, A>,
339 values_iter: std::vec::IntoIter<Option<RedisString>>,
340 phantom: std::marker::PhantomData<B>,
341 }
342
343 impl<'a, A, B> Iterator for HMGetIter<'a, A, B>
344 where
345 A: Into<Vec<u8>> + Clone,
346 RedisString: Into<B>,
347 {
348 type Item = (A, B);
349
next(&mut self) -> Option<Self::Item>350 fn next(&mut self) -> Option<Self::Item> {
351 loop {
352 let a = self.fields_iter.next();
353 let b = self.values_iter.next();
354 match b {
355 None => return None,
356 Some(None) => continue,
357 Some(Some(rs)) => {
358 return Some((
359 a.expect("field and value slices not of same length")
360 .clone(),
361 rs.into(),
362 ))
363 }
364 }
365 }
366 }
367 }
368
369 impl<'a, A, B> IntoIterator for HMGetResult<'a, A, B>
370 where
371 A: Into<Vec<u8>> + Clone,
372 RedisString: Into<B>,
373 {
374 type Item = (A, B);
375 type IntoIter = HMGetIter<'a, A, B>;
376
377 /// Provides an iterator over the multi-get results in the form of (field-name, field-value)
378 /// pairs. The type of field-name elements is the same as that passed to the original multi-
379 /// get call, while the field-value elements may be of any type for which a `RedisString` `Into`
380 /// conversion is implemented.
381 ///
382 /// # Examples
383 ///
384 /// Get a [`HashMap`] from the results:
385 ///
386 /// ```
387 /// use std::collections::HashMap;
388 /// use redis_module::RedisError;
389 ///
390 /// let keyname = "config";
391 /// let fields = &["username", "password", "email"];
392 /// let hm = ctx
393 /// .open_key(keyname)
394 /// .hash_get_multi(fields)?
395 /// .ok_or(RedisError::Str("ERR key not found"))?;
396 /// let response: HashMap<&str, String> = hm.into_iter().collect();
397 /// ```
398 ///
399 /// Get a [`Vec`] of only the field values from the results:
400 ///
401 /// ```
402 /// use redis_module::RedisError;
403 ///
404 /// let hm = ctx
405 /// .open_key(keyname)
406 /// .hash_get_multi(fields)?
407 /// .ok_or(RedisError::Str("ERR key not found"))?;
408 /// let response: Vec<String> = hm.into_iter().map(|(_, v)| v).collect();
409 /// ```
410 ///
411 /// [`HashMap`]: std::collections::HashMap
412 /// [`Vec`]: Vec
into_iter(self) -> Self::IntoIter413 fn into_iter(self) -> Self::IntoIter {
414 Self::IntoIter {
415 fields_iter: self.fields.iter(),
416 values_iter: self.values.into_iter(),
417 phantom: std::marker::PhantomData,
418 }
419 }
420 }
421
422 impl From<raw::Status> for Result<(), RedisError> {
from(s: raw::Status) -> Self423 fn from(s: raw::Status) -> Self {
424 match s {
425 raw::Status::Ok => Ok(()),
426 raw::Status::Err => Err(RedisError::String("Generic error".to_string())),
427 }
428 }
429 }
430
431 impl Drop for RedisKeyWritable {
432 // Frees resources appropriately as a RedisKey goes out of scope.
drop(&mut self)433 fn drop(&mut self) {
434 raw::close_key(self.key_inner);
435 }
436 }
437
read_key(key: *mut raw::RedisModuleKey) -> Result<String, Utf8Error>438 fn read_key(key: *mut raw::RedisModuleKey) -> Result<String, Utf8Error> {
439 let mut length: size_t = 0;
440 from_byte_string(
441 raw::string_dma(key, &mut length, raw::KeyMode::READ),
442 length,
443 )
444 }
445
446 /// Get an arbitrary number of hash fields from a key by batching calls
447 /// to `raw::hash_get_multi`.
hash_mget_key<T>( ctx: *mut raw::RedisModuleCtx, key: *mut raw::RedisModuleKey, fields: &[T], ) -> Result<Vec<Option<RedisString>>, RedisError> where T: Into<Vec<u8>> + Clone,448 fn hash_mget_key<T>(
449 ctx: *mut raw::RedisModuleCtx,
450 key: *mut raw::RedisModuleKey,
451 fields: &[T],
452 ) -> Result<Vec<Option<RedisString>>, RedisError>
453 where
454 T: Into<Vec<u8>> + Clone,
455 {
456 const BATCH_SIZE: usize = 12;
457
458 let mut values = Vec::with_capacity(fields.len());
459 let mut values_raw = [std::ptr::null_mut(); BATCH_SIZE];
460
461 for chunk_fields in fields.chunks(BATCH_SIZE) {
462 let mut chunk_values = &mut values_raw[..chunk_fields.len()];
463 raw::hash_get_multi(key, chunk_fields, &mut chunk_values)?;
464 values.extend(chunk_values.iter().map(|ptr| {
465 if ptr.is_null() {
466 None
467 } else {
468 Some(RedisString::new(ctx, *ptr))
469 }
470 }));
471 }
472
473 Ok(values)
474 }
475
to_raw_mode(mode: KeyMode) -> raw::KeyMode476 fn to_raw_mode(mode: KeyMode) -> raw::KeyMode {
477 match mode {
478 KeyMode::Read => raw::KeyMode::READ,
479 KeyMode::ReadWrite => raw::KeyMode::READ | raw::KeyMode::WRITE,
480 }
481 }
482
483 /// # Panics
484 ///
485 /// Will panic if `RedisModule_KeyType` or `RedisModule_ModuleTypeGetType` are missing in redismodule.h
486 #[allow(clippy::not_unsafe_ptr_arg_deref)]
verify_type(key_inner: *mut raw::RedisModuleKey, redis_type: &RedisType) -> RedisResult487 pub fn verify_type(key_inner: *mut raw::RedisModuleKey, redis_type: &RedisType) -> RedisResult {
488 let key_type: KeyType = unsafe { raw::RedisModule_KeyType.unwrap()(key_inner) }.into();
489
490 if key_type != KeyType::Empty {
491 // The key exists; check its type
492 let raw_type = unsafe { raw::RedisModule_ModuleTypeGetType.unwrap()(key_inner) };
493
494 if raw_type != *redis_type.raw_type.borrow() {
495 return Err(RedisError::Str("Existing key has wrong Redis type"));
496 }
497 }
498
499 REDIS_OK
500 }
501