diff --git a/CHANGELOG.md b/CHANGELOG.md index 73c7c2146a..a8d4186e25 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ * Python: Added ZDIFF command ([#1401](https://github.com/aws/glide-for-redis/pull/1401)) * Python: Added BZPOPMIN and BZPOPMAX commands ([#1399](https://github.com/aws/glide-for-redis/pull/1399)) * Python: Added ZUNIONSTORE, ZINTERSTORE commands ([#1388](https://github.com/aws/glide-for-redis/pull/1388)) +* Python: Added ZRANDMEMBER command ([#1413](https://github.com/aws/glide-for-redis/pull/1413)) #### Fixes diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs index 0f92daa325..ec74fed35a 100644 --- a/glide-core/src/client/value_conversion.rs +++ b/glide-core/src/client/value_conversion.rs @@ -20,7 +20,8 @@ pub(crate) enum ExpectedReturnType { ArrayOfDoubleOrNull, Lolwut, ArrayOfArraysOfDoubleOrNull, - ArrayOfKeyValuePairs, + ArrayOfPairs, + ArrayOfMemberScorePairs, ZMPopReturnType, KeyWithMemberAndScore, } @@ -304,23 +305,25 @@ pub(crate) fn convert_to_expected_type( .into()), } } - ExpectedReturnType::ArrayOfKeyValuePairs => match value { - Value::Nil => Ok(value), - Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => { - Ok(value) - } - Value::Array(array) - if matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) => - { - convert_flat_array_to_key_value_pairs(array) - } - _ => Err(( - ErrorKind::TypeError, - "Response couldn't be converted to an array of key-value pairs", - format!("(response was {:?})", value), - ) - .into()), - }, + // Used by HRANDFIELD when the WITHVALUES arg is passed. + // The server response can be an empty array, a flat array of key-value pairs, or a two-dimensional array of key-value pairs. + // The conversions we do here are as follows: + // + // - if the server returned an empty array, return an empty array + // - if the server returned a flat array of key-value pairs, convert to a two-dimensional array of key-value pairs + // - if the server returned a two-dimensional array of key-value pairs, return as-is + ExpectedReturnType::ArrayOfPairs => convert_to_array_of_pairs(value, None), + // Used by ZRANDMEMBER when the WITHSCORES arg is passed. + // The server response can be an empty array, a flat array of member-score pairs, or a two-dimensional array of member-score pairs. + // The server response scores can be strings or doubles. The conversions we do here are as follows: + // + // - if the server returned an empty array, return an empty array + // - if the server returned a flat array of member-score pairs, convert to a two-dimensional array of member-score pairs. The scores are converted from type string to type double. + // - if the server returned a two-dimensional array of key-value pairs, return as-is. The scores will already be of type double since this is a RESP3 response. + ExpectedReturnType::ArrayOfMemberScorePairs => { + // RESP2 returns scores as strings, but we want scores as type double. + convert_to_array_of_pairs(value, Some(ExpectedReturnType::Double)) + } // Used by BZPOPMIN/BZPOPMAX, which return an array consisting of the key of the sorted set that was popped, the popped member, and its score. // RESP2 returns the score as a string, but RESP3 returns the score as a double. Here we convert string scores into type double. ExpectedReturnType::KeyWithMemberAndScore => match value { @@ -423,10 +426,51 @@ fn convert_array_to_map( Ok(Value::Map(map)) } -/// Converts a flat array of values to a two-dimensional array, where the inner arrays are two-length arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map. +/// Used by commands like ZRANDMEMBER and HRANDFIELD. Normally a map would be more suitable for these key-value responses, but these commands may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map. +/// +/// Converts a server response as follows: +/// - if the server returned an empty array, return an empty array. +/// - if the server returned a flat array (RESP2), convert it to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs. +/// - if the server returned a two-dimensional array (RESP3), return the response as is, since it is already in the correct format. +/// - otherwise, return an error. +/// +/// `response` is a server response that we should attempt to convert as described above. +/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs. The values will only be converted if the response was a flat array, since RESP3 already returns an array of pairs with values already of the correct type. +fn convert_to_array_of_pairs( + response: Value, + value_expected_return_type: Option, +) -> RedisResult { + match response { + Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => { + // The server response is an empty array or a RESP3 array of pairs. In RESP3, the values in the pairs are + // already of the correct type, so we do not need to convert them and `response` is in the correct format. + Ok(response) + } + Value::Array(array) + if array.len() % 2 == 0 + && matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) => + { + // The server response is a RESP2 flat array with keys at even indices and their associated values at + // odd indices. + convert_flat_array_to_array_of_pairs(array, value_expected_return_type) + } + _ => Err(( + ErrorKind::TypeError, + "Response couldn't be converted to an array of key-value pairs", + format!("(response was {:?})", response), + ) + .into()), + } +} + +/// Converts a flat array of values to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map. /// /// `array` is a flat array containing keys at even-positioned elements and their associated values at odd-positioned elements. -fn convert_flat_array_to_key_value_pairs(array: Vec) -> RedisResult { +/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs. +fn convert_flat_array_to_array_of_pairs( + array: Vec, + value_expected_return_type: Option, +) -> RedisResult { if array.len() % 2 != 0 { return Err(( ErrorKind::TypeError, @@ -437,7 +481,9 @@ fn convert_flat_array_to_key_value_pairs(array: Vec) -> RedisResult Option { b"GEOPOS" => Some(ExpectedReturnType::ArrayOfArraysOfDoubleOrNull), b"HRANDFIELD" => cmd .position(b"WITHVALUES") - .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs), + .map(|_| ExpectedReturnType::ArrayOfPairs), b"ZRANDMEMBER" => cmd .position(b"WITHSCORES") - .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs), + .map(|_| ExpectedReturnType::ArrayOfMemberScorePairs), b"ZADD" => cmd .position(b"INCR") .map(|_| ExpectedReturnType::DoubleOrNull), @@ -575,7 +621,7 @@ mod tests { } #[test] - fn convert_array_of_kv_pairs() { + fn convert_to_array_of_pairs_return_type() { assert!(matches!( expected_type_for_cmd( redis::cmd("HRANDFIELD") @@ -583,25 +629,12 @@ mod tests { .arg("1") .arg("withvalues") ), - Some(ExpectedReturnType::ArrayOfKeyValuePairs) + Some(ExpectedReturnType::ArrayOfPairs) )); assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key").arg("1")).is_none()); assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key")).is_none()); - assert!(matches!( - expected_type_for_cmd( - redis::cmd("ZRANDMEMBER") - .arg("key") - .arg("1") - .arg("withscores") - ), - Some(ExpectedReturnType::ArrayOfKeyValuePairs) - )); - - assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none()); - assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none()); - let flat_array = Value::Array(vec![ Value::BulkString(b"key1".to_vec()), Value::BulkString(b"value1".to_vec()), @@ -619,34 +652,27 @@ mod tests { ]), ]); let converted_flat_array = - convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfKeyValuePairs)) - .unwrap(); + convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfPairs)).unwrap(); assert_eq!(two_dimensional_array, converted_flat_array); let converted_two_dimensional_array = convert_to_expected_type( two_dimensional_array.clone(), - Some(ExpectedReturnType::ArrayOfKeyValuePairs), + Some(ExpectedReturnType::ArrayOfPairs), ) .unwrap(); assert_eq!(two_dimensional_array, converted_two_dimensional_array); let empty_array = Value::Array(vec![]); - let converted_empty_array = convert_to_expected_type( - empty_array.clone(), - Some(ExpectedReturnType::ArrayOfKeyValuePairs), - ) - .unwrap(); - assert_eq!(empty_array, converted_empty_array); - - let converted_nil_value = - convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfKeyValuePairs)) + let converted_empty_array = + convert_to_expected_type(empty_array.clone(), Some(ExpectedReturnType::ArrayOfPairs)) .unwrap(); - assert_eq!(Value::Nil, converted_nil_value); + assert_eq!(empty_array, converted_empty_array); - let array_of_doubles = Value::Array(vec![Value::Double(5.5)]); + let flat_array_unexpected_length = + Value::Array(vec![Value::BulkString(b"somekey".to_vec())]); assert!(convert_to_expected_type( - array_of_doubles, - Some(ExpectedReturnType::ArrayOfKeyValuePairs) + flat_array_unexpected_length, + Some(ExpectedReturnType::ArrayOfPairs) ) .is_err()); } @@ -690,6 +716,42 @@ mod tests { assert_eq!(redis_response, converted_response); } + #[test] + fn convert_to_member_score_pairs_return_type() { + assert!(matches!( + expected_type_for_cmd( + redis::cmd("ZRANDMEMBER") + .arg("key") + .arg("1") + .arg("withscores") + ), + Some(ExpectedReturnType::ArrayOfMemberScorePairs) + )); + + assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none()); + assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none()); + + // convert_to_array_of_pairs_return_type already tests most functionality since the conversion for ArrayOfPairs + // and ArrayOfMemberScorePairs is mostly the same. Here we also test that the scores are converted to double + // when the server response was a RESP2 flat array. + let flat_array = Value::Array(vec![ + Value::BulkString(b"one".to_vec()), + Value::BulkString(b"1.0".to_vec()), + Value::BulkString(b"two".to_vec()), + Value::BulkString(b"2.0".to_vec()), + ]); + let expected_response = Value::Array(vec![ + Value::Array(vec![Value::BulkString(b"one".to_vec()), Value::Double(1.0)]), + Value::Array(vec![Value::BulkString(b"two".to_vec()), Value::Double(2.0)]), + ]); + let converted_flat_array = convert_to_expected_type( + flat_array, + Some(ExpectedReturnType::ArrayOfMemberScorePairs), + ) + .unwrap(); + assert_eq!(expected_response, converted_flat_array); + } + #[test] fn convert_zadd_only_if_incr_is_included() { assert!(matches!( diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 062f05702d..516494da60 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -985,7 +985,7 @@ async def hrandfield_count(self, key: str, count: int) -> List[str]: key (str): The key of the hash. count (int): The number of field names to return. If `count` is positive, returns unique elements. - If negative, allows for duplicates. + If `count` is negative, allows for duplicates elements. Returns: List[str]: A list of random field names from the hash. @@ -1012,7 +1012,7 @@ async def hrandfield_withvalues(self, key: str, count: int) -> List[List[str]]: key (str): The key of the hash. count (int): The number of field names to return. If `count` is positive, returns unique elements. - If negative, allows for duplicates. + If `count` is negative, allows for duplicates elements. Returns: List[List[str]]: A list of `[field_name, value]` lists, where `field_name` is a random field name from the @@ -2965,6 +2965,93 @@ async def zunionstore( await self._execute_command(RequestType.ZUnionStore, args), ) + async def zrandmember(self, key: str) -> Optional[str]: + """ + Returns a random member from the sorted set stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + + Returns: + Optional[str]: A random member from the sorted set. + If the sorted set does not exist or is empty, the response will be None. + + Examples: + >>> await client.zadd("my_sorted_set", {"member1": 1.0, "member2": 2.0}) + >>> await client.zrandmember("my_sorted_set") + "member1" # "member1" is a random member of "my_sorted_set". + >>> await client.zrandmember("non_existing_sorted_set") + None # "non_existing_sorted_set" is not an existing key, so None was returned. + """ + return cast( + Optional[str], + await self._execute_command(RequestType.ZRandMember, [key]), + ) + + async def zrandmember_count(self, key: str, count: int) -> List[str]: + """ + Retrieves up to the absolute value of `count` random members from the sorted set stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + count (int): The number of members to return. + If `count` is positive, returns unique members. + If `count` is negative, allows for duplicates members. + + Returns: + List[str]: A list of members from the sorted set. + If the sorted set does not exist or is empty, the response will be an empty list. + + Examples: + >>> await client.zadd("my_sorted_set", {"member1": 1.0, "member2": 2.0}) + >>> await client.zrandmember("my_sorted_set", -3) + ["member1", "member1", "member2"] # "member1" and "member2" are random members of "my_sorted_set". + >>> await client.zrandmember("non_existing_sorted_set", 3) + [] # "non_existing_sorted_set" is not an existing key, so an empty list was returned. + """ + return cast( + List[str], + await self._execute_command(RequestType.ZRandMember, [key, str(count)]), + ) + + async def zrandmember_withscores( + self, key: str, count: int + ) -> List[List[Union[str, float]]]: + """ + Retrieves up to the absolute value of `count` random members along with their scores from the sorted set + stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + count (int): The number of members to return. + If `count` is positive, returns unique members. + If `count` is negative, allows for duplicates members. + + Returns: + List[List[Union[str, float]]]: A list of `[member, score]` lists, where `member` is a random member from + the sorted set and `score` is the associated score. + If the sorted set does not exist or is empty, the response will be an empty list. + + Examples: + >>> await client.zadd("my_sorted_set", {"member1": 1.0, "member2": 2.0}) + >>> await client.zrandmember_withscores("my_sorted_set", -3) + [["member1", 1.0], ["member1", 1.0], ["member2", 2.0]] # "member1" and "member2" are random members of "my_sorted_set", and have scores of 1.0 and 2.0, respectively. + >>> await client.zrandmember_withscores("non_existing_sorted_set", 3) + [] # "non_existing_sorted_set" is not an existing key, so an empty list was returned. + """ + return cast( + List[List[Union[str, float]]], + await self._execute_command( + RequestType.ZRandMember, [key, str(count), "WITHSCORES"] + ), + ) + async def invoke_script( self, script: Script, diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 9eb523d50a..c062f6c444 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -638,7 +638,7 @@ def hrandfield_count(self: TTransaction, key: str, count: int) -> TTransaction: key (str): The key of the hash. count (int): The number of field names to return. If `count` is positive, returns unique elements. - If negative, allows for duplicates. + If `count` is negative, allows for duplicates elements. Command response: List[str]: A list of random field names from the hash. @@ -656,7 +656,7 @@ def hrandfield_withvalues(self: TTransaction, key: str, count: int) -> TTransact key (str): The key of the hash. count (int): The number of field names to return. If `count` is positive, returns unique elements. - If negative, allows for duplicates. + If `count` is negative, allows for duplicates elements. Command response: List[List[str]]: A list of `[field_name, value]` lists, where `field_name` is a random field name from the @@ -2105,6 +2105,63 @@ def zunionstore( args = _create_z_cmd_store_args(destination, keys, aggregation_type) return self.append_command(RequestType.ZUnionStore, args) + def zrandmember(self: TTransaction, key: str) -> TTransaction: + """ + Returns a random member from the sorted set stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + + Command response: + Optional[str]: A random member from the sorted set. + If the sorted set does not exist or is empty, the response will be None. + """ + return self.append_command(RequestType.ZRandMember, [key]) + + def zrandmember_count(self: TTransaction, key: str, count: int) -> TTransaction: + """ + Retrieves up to the absolute value of `count` random members from the sorted set stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + count (int): The number of members to return. + If `count` is positive, returns unique members. + If `count` is negative, allows for duplicates members. + + Command response: + List[str]: A list of members from the sorted set. + If the sorted set does not exist or is empty, the response will be an empty list. + """ + return self.append_command(RequestType.ZRandMember, [key, str(count)]) + + def zrandmember_withscores( + self: TTransaction, key: str, count: int + ) -> TTransaction: + """ + Retrieves up to the absolute value of `count` random members along with their scores from the sorted set + stored at 'key'. + + See https://valkey.io/commands/zrandmember for more details. + + Args: + key (str): The key of the sorted set. + count (int): The number of members to return. + If `count` is positive, returns unique members. + If `count` is negative, allows for duplicates members. + + Command response: + List[List[Union[str, float]]]: A list of `[member, score]` lists, where `member` is a random member from + the sorted set and `score` is the associated score. + If the sorted set does not exist or is empty, the response will be an empty list. + """ + return self.append_command( + RequestType.ZRandMember, [key, str(count), "WITHSCORES"] + ) + def dbsize(self: TTransaction) -> TTransaction: """ Returns the number of keys in the currently selected database. diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index e08d4cf45d..3fff943ea0 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -2713,6 +2713,79 @@ async def test_zdiffstore(self, redis_client: TRedisClient): await redis_client.zdiffstore("abc", ["zxy", "lkn"]) assert "CrossSlot" in str(e) + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrandmember(self, redis_client: TRedisClient): + key = get_random_string(10) + string_key = get_random_string(10) + scores = {"one": 1, "two": 2} + assert await redis_client.zadd(key, scores) == 2 + + member = await redis_client.zrandmember(key) + assert member in scores + assert await redis_client.zrandmember("non_existing_key") is None + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrandmember(string_key) + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrandmember_count(self, redis_client: TRedisClient): + key = get_random_string(10) + string_key = get_random_string(10) + scores = {"one": 1, "two": 2} + assert await redis_client.zadd(key, scores) == 2 + + # unique values are expected as count is positive + members = await redis_client.zrandmember_count(key, 4) + assert len(members) == 2 + assert set(members) == {"one", "two"} + + # duplicate values are expected as count is negative + members = await redis_client.zrandmember_count(key, -4) + assert len(members) == 4 + for member in members: + assert member in scores + + assert await redis_client.zrandmember_count(key, 0) == [] + assert await redis_client.zrandmember_count("non_existing_key", 0) == [] + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrandmember_count(string_key, 5) + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_zrandmember_withscores(self, redis_client: TRedisClient): + key = get_random_string(10) + string_key = get_random_string(10) + scores = {"one": 1, "two": 2} + assert await redis_client.zadd(key, scores) == 2 + + # unique values are expected as count is positive + elements = await redis_client.zrandmember_withscores(key, 4) + assert len(elements) == 2 + + for member, score in elements: + assert scores[str(member)] == score + + # duplicate values are expected as count is negative + elements = await redis_client.zrandmember_withscores(key, -4) + assert len(elements) == 4 + for member, score in elements: + assert scores[str(member)] == score + + assert await redis_client.zrandmember_withscores(key, 0) == [] + assert await redis_client.zrandmember_withscores("non_existing_key", 0) == [] + + # key exists, but it is not a set + assert await redis_client.set(string_key, "value") == OK + with pytest.raises(RequestError): + await redis_client.zrandmember_withscores(string_key, 5) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_type(self, redis_client: TRedisClient): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 555e5ba4c7..795358c756 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -235,6 +235,13 @@ async def transaction_test( args.append([key8, "two", 2.0]) transaction.bzpopmax([key8], 0.5) args.append([key8, "four", 4.0]) + # key8 now only contains one member ("three") + transaction.zrandmember(key8) + args.append("three") + transaction.zrandmember_count(key8, 1) + args.append(["three"]) + transaction.zrandmember_withscores(key8, 1) + args.append([["three", 3.0]]) transaction.zpopmax(key8) args.append({"three": 3.0}) transaction.zpopmin(key8)