use std::convert::TryFrom;
use std::os::raw::c_void;
use std::ptr;
use std::str::Utf8Error;
use std::time::Duration;
use libc::size_t;
use raw::KeyType;
use crate::from_byte_string;
use crate::native_types::RedisType;
use crate::raw;
use crate::redismodule::REDIS_OK;
use crate::RedisError;
use crate::RedisResult;
use crate::RedisString;
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum KeyMode {
Read,
ReadWrite,
}
#[derive(Debug)]
pub struct RedisKey {
ctx: *mut raw::RedisModuleCtx,
key_inner: *mut raw::RedisModuleKey,
}
impl RedisKey {
pub fn open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> Self {
let key_inner = raw::open_key(ctx, key.inner, to_raw_mode(KeyMode::Read));
Self { ctx, key_inner }
}
pub fn get_value<T>(&self, redis_type: &RedisType) -> Result<Option<&T>, RedisError> {
verify_type(self.key_inner, redis_type)?;
let value =
unsafe { raw::RedisModule_ModuleTypeGetValue.unwrap()(self.key_inner).cast::<T>() };
if value.is_null() {
return Ok(None);
}
let value = unsafe { &*value };
Ok(Some(value))
}
#[must_use]
pub fn key_type(&self) -> raw::KeyType {
unsafe { raw::RedisModule_KeyType.unwrap()(self.key_inner) }.into()
}
#[must_use]
pub fn is_null(&self) -> bool {
let null_key: *mut raw::RedisModuleKey = ptr::null_mut();
self.key_inner == null_key
}
pub fn read(&self) -> Result<Option<String>, RedisError> {
let val = if self.is_null() {
None
} else {
Some(read_key(self.key_inner)?)
};
Ok(val)
}
pub fn hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError> {
let val = if self.is_null() {
None
} else {
hash_mget_key(self.ctx, self.key_inner, &[field])?
.pop()
.expect("hash_mget_key should return vector of same length as input")
};
Ok(val)
}
pub fn 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>,
{
let val = if self.is_null() {
None
} else {
Some(HMGetResult {
fields,
values: hash_mget_key(self.ctx, self.key_inner, fields)?,
phantom: std::marker::PhantomData,
})
};
Ok(val)
}
}
impl Drop for RedisKey {
fn drop(&mut self) {
raw::close_key(self.key_inner);
}
}
pub struct RedisKeyWritable {
ctx: *mut raw::RedisModuleCtx,
key_inner: *mut raw::RedisModuleKey,
}
impl RedisKeyWritable {
pub fn open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> Self {
let key_inner = raw::open_key(ctx, key.inner, to_raw_mode(KeyMode::ReadWrite));
Self { ctx, key_inner }
}
pub fn read(&self) -> Result<Option<String>, RedisError> {
Ok(Some(read_key(self.key_inner)?))
}
#[allow(clippy::must_use_candidate)]
pub fn hash_set(&self, field: &str, value: RedisString) -> raw::Status {
raw::hash_set(self.key_inner, field, value.inner)
}
#[allow(clippy::must_use_candidate)]
pub fn hash_del(&self, field: &str) -> raw::Status {
raw::hash_del(self.key_inner, field)
}
pub fn hash_get(&self, field: &str) -> Result<Option<RedisString>, RedisError> {
Ok(hash_mget_key(self.ctx, self.key_inner, &[field])?
.pop()
.expect("hash_mget_key should return vector of same length as input"))
}
pub fn 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>,
{
Ok(HMGetResult {
fields,
values: hash_mget_key(self.ctx, self.key_inner, fields)?,
phantom: std::marker::PhantomData,
})
}
#[allow(clippy::must_use_candidate)]
pub fn list_push_head(&self, element: RedisString) -> raw::Status {
raw::list_push(self.key_inner, raw::Where::ListHead, element.inner)
}
#[allow(clippy::must_use_candidate)]
pub fn list_push_tail(&self, element: RedisString) -> raw::Status {
raw::list_push(self.key_inner, raw::Where::ListTail, element.inner)
}
#[allow(clippy::must_use_candidate)]
pub fn list_pop_head(&self) -> Option<RedisString> {
let ptr = raw::list_pop(self.key_inner, raw::Where::ListHead);
if ptr.is_null() {
return None;
}
Some(RedisString::new(self.ctx, ptr))
}
#[must_use]
pub fn list_pop_tail(&self) -> Option<RedisString> {
let ptr = raw::list_pop(self.key_inner, raw::Where::ListTail);
if ptr.is_null() {
return None;
}
Some(RedisString::new(self.ctx, ptr))
}
pub fn set_expire(&self, expire: Duration) -> RedisResult {
let exp_millis = expire.as_millis();
let exp_time = i64::try_from(exp_millis).map_err(|_| {
RedisError::String(format!(
"Error expire duration {} is not allowed",
exp_millis
))
})?;
match raw::set_expire(self.key_inner, exp_time) {
raw::Status::Ok => REDIS_OK,
raw::Status::Err => Err(RedisError::Str("Error while setting key expire")),
}
}
pub fn write(&self, val: &str) -> RedisResult {
let val_str = RedisString::create(self.ctx, val);
match raw::string_set(self.key_inner, val_str.inner) {
raw::Status::Ok => REDIS_OK,
raw::Status::Err => Err(RedisError::Str("Error while setting key")),
}
}
pub fn delete(&self) -> RedisResult {
unsafe { raw::RedisModule_DeleteKey.unwrap()(self.key_inner) };
REDIS_OK
}
#[must_use]
pub fn key_type(&self) -> raw::KeyType {
unsafe { raw::RedisModule_KeyType.unwrap()(self.key_inner) }.into()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.key_type() == KeyType::Empty
}
pub fn open_with_redis_string(
ctx: *mut raw::RedisModuleCtx,
key: *mut raw::RedisModuleString,
) -> Self {
let key_inner = raw::open_key(ctx, key, to_raw_mode(KeyMode::ReadWrite));
Self { ctx, key_inner }
}
pub fn get_value<'a, 'b, T>(
&'a self,
redis_type: &RedisType,
) -> Result<Option<&'b mut T>, RedisError> {
verify_type(self.key_inner, redis_type)?;
let value =
unsafe { raw::RedisModule_ModuleTypeGetValue.unwrap()(self.key_inner).cast::<T>() };
if value.is_null() {
return Ok(None);
}
let value = unsafe { &mut *value };
Ok(Some(value))
}
pub fn set_value<T>(&self, redis_type: &RedisType, value: T) -> Result<(), RedisError> {
verify_type(self.key_inner, redis_type)?;
let value = Box::into_raw(Box::new(value)).cast::<c_void>();
let status: raw::Status = unsafe {
raw::RedisModule_ModuleTypeSetValue.unwrap()(
self.key_inner,
*redis_type.raw_type.borrow(),
value,
)
}
.into();
status.into()
}
}
pub struct HMGetResult<'a, A, B>
where
A: Into<Vec<u8>> + Clone,
RedisString: Into<B>,
{
fields: &'a [A],
values: Vec<Option<RedisString>>,
phantom: std::marker::PhantomData<B>,
}
pub struct HMGetIter<'a, A, B>
where
A: Into<Vec<u8>>,
RedisString: Into<B>,
{
fields_iter: std::slice::Iter<'a, A>,
values_iter: std::vec::IntoIter<Option<RedisString>>,
phantom: std::marker::PhantomData<B>,
}
impl<'a, A, B> Iterator for HMGetIter<'a, A, B>
where
A: Into<Vec<u8>> + Clone,
RedisString: Into<B>,
{
type Item = (A, B);
fn next(&mut self) -> Option<Self::Item> {
loop {
let a = self.fields_iter.next();
let b = self.values_iter.next();
match b {
None => return None,
Some(None) => continue,
Some(Some(rs)) => {
return Some((
a.expect("field and value slices not of same length")
.clone(),
rs.into(),
))
}
}
}
}
}
impl<'a, A, B> IntoIterator for HMGetResult<'a, A, B>
where
A: Into<Vec<u8>> + Clone,
RedisString: Into<B>,
{
type Item = (A, B);
type IntoIter = HMGetIter<'a, A, B>;
fn into_iter(self) -> Self::IntoIter {
Self::IntoIter {
fields_iter: self.fields.iter(),
values_iter: self.values.into_iter(),
phantom: std::marker::PhantomData,
}
}
}
impl From<raw::Status> for Result<(), RedisError> {
fn from(s: raw::Status) -> Self {
match s {
raw::Status::Ok => Ok(()),
raw::Status::Err => Err(RedisError::String("Generic error".to_string())),
}
}
}
impl Drop for RedisKeyWritable {
fn drop(&mut self) {
raw::close_key(self.key_inner);
}
}
fn read_key(key: *mut raw::RedisModuleKey) -> Result<String, Utf8Error> {
let mut length: size_t = 0;
from_byte_string(
raw::string_dma(key, &mut length, raw::KeyMode::READ),
length,
)
}
fn 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,
{
const BATCH_SIZE: usize = 12;
let mut values = Vec::with_capacity(fields.len());
let mut values_raw = [std::ptr::null_mut(); BATCH_SIZE];
for chunk_fields in fields.chunks(BATCH_SIZE) {
let chunk_values = &mut values_raw[..chunk_fields.len()];
raw::hash_get_multi(key, chunk_fields, chunk_values)?;
values.extend(chunk_values.iter().map(|ptr| {
if ptr.is_null() {
None
} else {
Some(RedisString::new(ctx, *ptr))
}
}));
}
Ok(values)
}
fn to_raw_mode(mode: KeyMode) -> raw::KeyMode {
match mode {
KeyMode::Read => raw::KeyMode::READ,
KeyMode::ReadWrite => raw::KeyMode::READ | raw::KeyMode::WRITE,
}
}
#[allow(clippy::not_unsafe_ptr_arg_deref)]
pub fn verify_type(key_inner: *mut raw::RedisModuleKey, redis_type: &RedisType) -> RedisResult {
let key_type: KeyType = unsafe { raw::RedisModule_KeyType.unwrap()(key_inner) }.into();
if key_type != KeyType::Empty {
let raw_type = unsafe { raw::RedisModule_ModuleTypeGetType.unwrap()(key_inner) };
if raw_type != *redis_type.raw_type.borrow() {
return Err(RedisError::Str("Existing key has wrong Redis type"));
}
}
REDIS_OK
}