diff --git a/glide-core/src/client/mod.rs b/glide-core/src/client/mod.rs index 1b3cfd2e9e..420c95a62d 100644 --- a/glide-core/src/client/mod.rs +++ b/glide-core/src/client/mod.rs @@ -122,7 +122,7 @@ impl Client { routing: Option, ) -> redis::RedisFuture<'a, Value> { let expected_type = expected_type_for_cmd(cmd); - run_with_timeout(self.request_timeout, async { + run_with_timeout(self.request_timeout, async move { match self.internal_client { ClientWrapper::Standalone(ref mut client) => client.send_command(cmd).await, diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs index 803016ef2a..c4864310d0 100644 --- a/glide-core/src/client/value_conversion.rs +++ b/glide-core/src/client/value_conversion.rs @@ -5,12 +5,16 @@ use redis::{ cluster_routing::Routable, from_owned_redis_value, Cmd, ErrorKind, RedisResult, Value, }; +#[derive(Clone, Copy)] pub(crate) enum ExpectedReturnType { Map, + MapOfStringToDouble, Double, Boolean, + BulkString, Set, DoubleOrNull, + ZrankReturnType, } pub(crate) fn convert_to_expected_type( @@ -25,25 +29,51 @@ pub(crate) fn convert_to_expected_type( ExpectedReturnType::Map => match value { Value::Nil => Ok(value), Value::Map(_) => Ok(value), - Value::Array(array) => { - let mut map = Vec::with_capacity(array.len() / 2); - let mut iterator = array.into_iter(); - while let Some(key) = iterator.next() { - let Some(value) = iterator.next() else { - return Err(( - ErrorKind::TypeError, - "Response has odd number of items, and cannot be entered into a map", - ) - .into()); - }; - map.push((key, value)); - } + Value::Array(array) => convert_array_to_map(array, None, None), + _ => Err(( + ErrorKind::TypeError, + "Response couldn't be converted to map", + format!("(response was {:?})", value), + ) + .into()), + }, + ExpectedReturnType::MapOfStringToDouble => match value { + Value::Nil => Ok(value), + Value::Map(map) => { + let map_clone = map.clone(); + let result = map + .into_iter() + .map(|(key, inner_value)| { + let key_str = match key { + Value::BulkString(_) => key, + _ => Value::BulkString(from_owned_redis_value::(key)?.into()), + }; + match inner_value { + Value::BulkString(_) => Ok(( + key_str, + Value::Double(from_owned_redis_value::(inner_value)?.into()), + )), + Value::Double(_) => Ok((key_str, inner_value)), + _ => Err(( + ErrorKind::TypeError, + "Response couldn't be converted to map of {string: double}", + format!("(response was {:?})", map_clone), + ) + .into()), + } + }) + .collect::>(); - Ok(Value::Map(map)) + result.map(Value::Map) } + Value::Array(array) => convert_array_to_map( + array, + Some(ExpectedReturnType::BulkString), + Some(ExpectedReturnType::Double), + ), _ => Err(( ErrorKind::TypeError, - "Response couldn't be converted to map", + "Response couldn't be converted to map of {string: double}", format!("(response was {:?})", value), ) .into()), @@ -67,16 +97,101 @@ pub(crate) fn convert_to_expected_type( Value::Nil => Ok(value), _ => Ok(Value::Double(from_owned_redis_value::(value)?.into())), }, + ExpectedReturnType::ZrankReturnType => match value { + Value::Nil => Ok(value), + Value::Array(mut array) => { + if array.len() != 2 { + return Err(( + ErrorKind::TypeError, + "Array must contain exactly two elements", + ) + .into()); + } + + array[1] = + convert_to_expected_type(array[1].clone(), Some(ExpectedReturnType::Double))?; + + Ok(Value::Array(array)) + } + _ => Err(( + ErrorKind::TypeError, + "Response couldn't be converted to Array (ZrankResponseType)", + format!("(response was {:?})", value), + ) + .into()), + }, + ExpectedReturnType::BulkString => Ok(Value::BulkString( + from_owned_redis_value::(value)?.into(), + )), } } +fn convert_array_to_map( + array: Vec, + key_expected_return_type: Option, + value_expected_return_type: Option, +) -> RedisResult { + let mut map = Vec::new(); + let mut iterator = array.into_iter(); + while let Some(key) = iterator.next() { + match key { + Value::Array(inner_array) => { + if inner_array.len() != 2 { + return Err(( + ErrorKind::TypeError, + "Array inside map must contain exactly two elements", + ) + .into()); + } + let mut inner_iterator = inner_array.into_iter(); + let Some(inner_key) = inner_iterator.next() else { + return Err((ErrorKind::TypeError, "Missing key inside array of map").into()); + }; + let Some(inner_value) = inner_iterator.next() else { + return Err((ErrorKind::TypeError, "Missing value inside array of map").into()); + }; + + map.push(( + convert_to_expected_type(inner_key, key_expected_return_type)?, + convert_to_expected_type(inner_value, value_expected_return_type)?, + )); + } + _ => { + let Some(value) = iterator.next() else { + return Err(( + ErrorKind::TypeError, + "Response has odd number of items, and cannot be entered into a map", + ) + .into()); + }; + map.push(( + convert_to_expected_type(key, key_expected_return_type)?, + convert_to_expected_type(value, value_expected_return_type)?, + )); + } + } + } + Ok(Value::Map(map)) +} pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option { - if cmd.arg_idx(0) == Some(b"ZADD") { - return if cmd.position(b"INCR").is_some() { - Some(ExpectedReturnType::DoubleOrNull) - } else { - None - }; + let first_arg = cmd.arg_idx(0); + match first_arg { + Some(b"ZADD") => { + return cmd + .position(b"INCR") + .map(|_| ExpectedReturnType::DoubleOrNull); + } + Some(b"ZRANGE") | Some(b"ZDIFF") => { + return cmd + .position(b"WITHSCORES") + .map(|_| ExpectedReturnType::MapOfStringToDouble); + } + Some(b"ZRANK") | Some(b"ZREVRANK") => { + return cmd + .position(b"WITHSCORE") + .map(|_| ExpectedReturnType::ZrankReturnType); + } + _ => {} } let command = cmd.command()?; @@ -91,6 +206,7 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option { } b"SMEMBERS" => Some(ExpectedReturnType::Set), b"ZSCORE" => Some(ExpectedReturnType::DoubleOrNull), + b"ZPOPMIN" | b"ZPOPMAX" => Some(ExpectedReturnType::MapOfStringToDouble), _ => None, } } @@ -119,6 +235,191 @@ mod tests { .is_none()); } + #[test] + fn convert_zrange_zdiff_only_if_withsocres_is_included() { + assert!(matches!( + expected_type_for_cmd(redis::cmd("ZRANGE").arg("0").arg("-1").arg("WITHSCORES")), + Some(ExpectedReturnType::MapOfStringToDouble) + )); + + assert!(expected_type_for_cmd(redis::cmd("ZRANGE").arg("0").arg("-1")).is_none()); + + assert!(matches!( + expected_type_for_cmd(redis::cmd("ZDIFF").arg("1").arg("WITHSCORES")), + Some(ExpectedReturnType::MapOfStringToDouble) + )); + + assert!(expected_type_for_cmd(redis::cmd("ZDIFF").arg("1")).is_none()); + } + + #[test] + fn zpopmin_zpopmax_return_type() { + assert!(matches!( + expected_type_for_cmd(redis::cmd("ZPOPMIN").arg("1")), + Some(ExpectedReturnType::MapOfStringToDouble) + )); + + assert!(matches!( + expected_type_for_cmd(redis::cmd("ZPOPMAX").arg("1")), + Some(ExpectedReturnType::MapOfStringToDouble) + )); + } + + #[test] + fn convert_zank_zrevrank_only_if_withsocres_is_included() { + assert!(matches!( + expected_type_for_cmd( + redis::cmd("ZRANK") + .arg("key") + .arg("member") + .arg("WITHSCORE") + ), + Some(ExpectedReturnType::ZrankReturnType) + )); + + assert!(expected_type_for_cmd(redis::cmd("ZRANK").arg("key").arg("member")).is_none()); + + assert!(matches!( + expected_type_for_cmd( + redis::cmd("ZREVRANK") + .arg("key") + .arg("member") + .arg("WITHSCORE") + ), + Some(ExpectedReturnType::ZrankReturnType) + )); + + assert!(expected_type_for_cmd(redis::cmd("ZREVRANK").arg("key").arg("member")).is_none()); + } + + #[test] + fn test_convert_to_map_of_string_to_double() { + assert_eq!( + convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::MapOfStringToDouble)), + Ok(Value::Nil) + ); + let redis_map = vec![ + ( + Value::BulkString(b"key1".to_vec()), + Value::BulkString(b"10.5".to_vec()), + ), + ( + Value::BulkString(b"key2".to_vec()), + Value::BulkString(b"20.8".to_vec()), + ), + ( + Value::Double(20.5.into()), + Value::BulkString(b"30.2".to_vec()), + ), + ]; + + let converted_map = convert_to_expected_type( + Value::Map(redis_map), + Some(ExpectedReturnType::MapOfStringToDouble), + ) + .unwrap(); + + let converted_map = if let Value::Map(map) = converted_map { + map + } else { + panic!("Expected a Map, but got {:?}", converted_map); + }; + + assert_eq!(converted_map.len(), 3); + + let (key, value) = &converted_map[0]; + assert_eq!(*key, Value::BulkString(b"key1".to_vec())); + assert_eq!(*value, Value::Double(10.5.into())); + + let (key, value) = &converted_map[1]; + assert_eq!(*key, Value::BulkString(b"key2".to_vec())); + assert_eq!(*value, Value::Double(20.8.into())); + + let (key, value) = &converted_map[2]; + assert_eq!(*key, Value::BulkString(b"20.5".to_vec())); + assert_eq!(*value, Value::Double(30.2.into())); + + let array_of_arrays = vec![ + Value::Array(vec![ + Value::BulkString(b"key1".to_vec()), + Value::BulkString(b"10.5".to_vec()), + ]), + Value::Array(vec![ + Value::BulkString(b"key2".to_vec()), + Value::Double(20.5.into()), + ]), + ]; + + let converted_map = convert_to_expected_type( + Value::Array(array_of_arrays), + Some(ExpectedReturnType::MapOfStringToDouble), + ) + .unwrap(); + + let converted_map = if let Value::Map(map) = converted_map { + map + } else { + panic!("Expected a Map, but got {:?}", converted_map); + }; + + assert_eq!(converted_map.len(), 2); + + let (key, value) = &converted_map[0]; + assert_eq!(*key, Value::BulkString(b"key1".to_vec())); + assert_eq!(*value, Value::Double(10.5.into())); + + let (key, value) = &converted_map[1]; + assert_eq!(*key, Value::BulkString(b"key2".to_vec())); + assert_eq!(*value, Value::Double(20.5.into())); + + let array_of_arrays_err: Vec = vec![Value::Array(vec![ + Value::BulkString(b"key".to_vec()), + Value::BulkString(b"value".to_vec()), + Value::BulkString(b"10.5".to_vec()), + ])]; + + assert!(convert_to_expected_type( + Value::Array(array_of_arrays_err), + Some(ExpectedReturnType::MapOfStringToDouble) + ) + .is_err()); + } + + #[test] + fn test_convert_to_zrank_return_type() { + assert_eq!( + convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ZrankReturnType)), + Ok(Value::Nil) + ); + + let array = vec![ + Value::BulkString(b"key".to_vec()), + Value::BulkString(b"20.5".to_vec()), + ]; + + let array_result = convert_to_expected_type( + Value::Array(array), + Some(ExpectedReturnType::ZrankReturnType), + ) + .unwrap(); + + let array_result = if let Value::Array(array) = array_result { + array + } else { + panic!("Expected an Array, but got {:?}", array_result); + }; + assert_eq!(array_result.len(), 2); + + assert_eq!(array_result[0], Value::BulkString(b"key".to_vec())); + assert_eq!(array_result[1], Value::Double(20.5.into())); + + let array_err = vec![Value::BulkString(b"key".to_vec())]; + assert!(convert_to_expected_type( + Value::Array(array_err), + Some(ExpectedReturnType::ZrankReturnType) + ) + .is_err()); + } #[test] fn pass_null_value_for_double_or_null() { assert_eq!(