diff --git a/CHANGELOG.md b/CHANGELOG.md index b38788bf15..52b9c11982 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ * Python: Added BLPOP and BRPOP commands ([#1369](https://github.com/aws/glide-for-redis/pull/1369)) * Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377)) * Python: Added ZDIFFSTORE command ([#1378](https://github.com/aws/glide-for-redis/pull/1378)) +* Python: Added ZRANDMEMBER command (TODO: add PR link) #### Fixes diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs index 6afaf56e17..ca4c217cc3 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, } pub(crate) fn convert_to_expected_type( @@ -270,23 +271,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 a nil value, 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 nil, return nil + // - if the server returned an empty array, return an empty array + // - otherwise, return a two-dimensional array of key-value pairs + ExpectedReturnType::ArrayOfPairs => convert_to_array_of_pairs(value, None), + // Used by ZRANDMEMBER when the WITHSCORES arg is passed. + // The server response can be a nil value, 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 nil, return nil + // - if the server returned an empty array, return an empty array + // - otherwise, return a two-dimensional array of member-score pairs, where scores are of type double + ExpectedReturnType::ArrayOfMemberScorePairs => { + // RESP2 returns scores as strings, but we want scores as type double. + convert_to_array_of_pairs(value, Some(ExpectedReturnType::Double)) + } } } @@ -367,10 +370,53 @@ 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 nil, return nil. +/// - 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::Nil => Ok(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 member-score 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, @@ -381,7 +427,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), @@ -517,7 +565,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") @@ -525,25 +573,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()), @@ -561,38 +596,71 @@ 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(); + let converted_empty_array = + convert_to_expected_type(empty_array.clone(), Some(ExpectedReturnType::ArrayOfPairs)) + .unwrap(); assert_eq!(empty_array, converted_empty_array); let converted_nil_value = - convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfKeyValuePairs)) - .unwrap(); + convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfPairs)).unwrap(); assert_eq!(Value::Nil, converted_nil_value); - 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()); } + #[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 921535ab84..6af9295a30 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -983,7 +983,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. @@ -1010,7 +1010,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 @@ -2746,6 +2746,89 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int: ), ) + async def zrandmember(self, key: str) -> Optional[str]: + """ + Returns a random element 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 string representing a random element from the sorted set. + If the sorted set does not exist or is empty, the response will be None. + + Examples: + >>> await client.zrandmember("my_sorted_set") + "GLIDE" # "GLIDE" 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 random elements 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 elements to return. + If `count` is positive, returns unique elements. + If `count` is negative, allows for duplicates elements. + + Returns: + List[str]: A list of elements from the sorted set. + If the sorted set does not exist or is empty, the response will be an empty list. + + Examples: + >>> await client.zrandmember("my_sorted_set", -3) + ["GLIDE", "GLIDE", "PYTHON"] # "GLIDE" and "PYTHON" 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 random elements 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 elements to return. + If `count` is positive, returns unique elements. + If `count` is negative, allows for duplicates elements. + + 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.zrandmember_withscores("my_sorted_set", -3) + [["GLIDE", 1.0], ["GLIDE", 1.0], ["PYTHON", 2.0]] # "GLIDE" and "PYTHON" 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 dbf3651d8f..a49c04b934 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -636,7 +636,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. @@ -654,7 +654,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 @@ -1966,6 +1966,62 @@ def zdiffstore( RequestType.ZDiffStore, [destination, str(len(keys))] + keys ) + def zrandmember(self: TTransaction, key: str) -> TTransaction: + """ + Returns a random element 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 string representing a random element 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 random elements 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 elements to return. + If `count` is positive, returns unique elements. + If `count` is negative, allows for duplicates elements. + + Command response: + List[str]: A list of elements 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 random elements 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 elements to return. + If `count` is positive, returns unique elements. + If `count` is negative, allows for duplicates elements. + + 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 0461e8132a..d3c6c9ea88 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -2378,6 +2378,83 @@ 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 element in elements: + member = element[0] + score = element[1] + assert scores.get(str(member)) == score + + # duplicate values are expected as count is negative + elements = await redis_client.zrandmember_withscores(key, -4) + assert len(elements) == 4 + for element in elements: + member = element[0] + score = element[1] + assert scores.get(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 df45b3b8e4..0fd0651f22 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -46,6 +46,7 @@ async def transaction_test( key10 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # hyper log log key11 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # streams key12 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # geo + key13 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # sorted set value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S") value2 = get_random_string(5) @@ -238,6 +239,15 @@ async def transaction_test( transaction.zdiffstore(key8, [key8, key8]) args.append(0) + transaction.zadd(key13, {"one": 1.0}) + args.append(1) + transaction.zrandmember(key13) + args.append("one") + transaction.zrandmember_count(key13, 1) + args.append(["one"]) + transaction.zrandmember_withscores(key13, 1) + args.append([["one", 1.0]]) + transaction.pfadd(key10, ["a", "b", "c"]) args.append(1)