Rust: modify sorted set command types to RESP3

valkey-io · Feb 8, 2024 · 8c491ac · 8c491ac
1 parent 3a3c436
commit 8c491ac
Showing 1 changed file with 150 additions and 21 deletions.
diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs
@@ -7,10 +7,12 @@ use redis::{
 
 pub(crate) enum ExpectedReturnType {
     Map,
+    MapOfStringToDouble,
     Double,
     Boolean,
     Set,
     DoubleOrNull,
+    ArrayOfDouble,
 }
 
 pub(crate) fn convert_to_expected_type(
@@ -23,27 +25,25 @@ pub(crate) fn convert_to_expected_type(
 
     match expected {
         ExpectedReturnType::Map => match value {
+            Value::Nil => Ok(value),
+            Value::Map(_) => Ok(value),
+            Value::Array(array) => convert_array_to_map(array, ExpectedReturnType::Map),
+            _ => 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(_) => 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));
-                }
-
-                Ok(Value::Map(map))
+                convert_array_to_map(array, ExpectedReturnType::MapOfStringToDouble)
             }
             _ => 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 +67,100 @@ pub(crate) fn convert_to_expected_type(
             Value::Nil => Ok(value),
             _ => Ok(Value::Double(from_owned_redis_value::<f64>(value)?.into())),
         },
+        ExpectedReturnType::ArrayOfDouble => match value {
+            Value::Array(array) => {
+                let result: Result<Vec<Value>, _> = array
+                    .into_iter()
+                    .map(|element| match element {
+                        Value::BulkString(_) => Ok(Value::Double(
+                            from_owned_redis_value::<f64>(element)?.into(),
+                        )),
+                        _ => Ok(element),
+                    })
+                    .collect();
+
+                result.map(Value::Array)
+            }
+            _ => Err((
+                ErrorKind::TypeError,
+                "Response couldn't be converted to Array",
+                format!("(response was {:?})", value),
+            )
+                .into()),
+        },
     }
 }
 
+fn convert_array_to_map(
+    array: Vec<Value>,
+    expected_return_type: ExpectedReturnType,
+) -> RedisResult<Value> {
+    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 inner_key = inner_iterator.next().unwrap();
+                let inner_value = inner_iterator.next().unwrap();
+                match expected_return_type {
+                    ExpectedReturnType::MapOfStringToDouble => {
+                        map.push((
+                            inner_key,
+                            Value::Double(from_owned_redis_value::<f64>(inner_value)?.into()),
+                        ));
+                    }
+                    _ => {
+                        map.push((inner_key, inner_value));
+                    }
+                }
+            }
+            _ => {
+                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));
+            }
+        }
+    }
+    Ok(Value::Map(map))
+}
 pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
-    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 if cmd.position(b"INCR").is_some() {
+                Some(ExpectedReturnType::DoubleOrNull)
+            } else {
+                None
+            }
+        }
+        Some(b"ZRANGE") | Some(b"ZDIFF") => {
+            return if cmd.position(b"WITHSCORES").is_some() {
+                Some(ExpectedReturnType::MapOfStringToDouble)
+            } else {
+                None
+            }
+        }
+        Some(b"ZRANK") | Some(b"ZREVRANK") => {
+            return if cmd.position(b"WITHSCORE").is_some() {
+                Some(ExpectedReturnType::ArrayOfDouble)
+            } else {
+                None
+            }
+        }
+        _ => {}
     }
 
     let command = cmd.command()?;
@@ -91,6 +175,7 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
         }
         b"SMEMBERS" => Some(ExpectedReturnType::Set),
         b"ZSCORE" => Some(ExpectedReturnType::DoubleOrNull),
+        b"ZPOPMIN" | b"ZPOPMAX" => Some(ExpectedReturnType::MapOfStringToDouble),
         _ => None,
     }
 }
@@ -119,6 +204,50 @@ 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 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::ArrayOfDouble)
+        ));
+
+        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::ArrayOfDouble)
+        ));
+
+        assert!(expected_type_for_cmd(redis::cmd("ZREVRANK").arg("key").arg("member")).is_none());
+    }
+
     #[test]
     fn pass_null_value_for_double_or_null() {
         assert_eq!(