Python: adds ZRANK command (#1065)

CHANGELOG.md

@@ -21,6 +21,7 @@
 * Python: Added json module and JSON.SET JSON.GET commands  ([#1056](https://github.com/aws/glide-for-redis/pull/1056))
 * Node: Added Time command. ([#1114](https://github.com/aws/glide-for-redis/pull/1114))
 * Python, Node: Added LINDEX command ([#1058](https://github.com/aws/glide-for-redis/pull/1058), [#999](https://github.com/aws/glide-for-redis/pull/999))
+* Python: Added ZRANK command ([#1065](https://github.com/aws/glide-for-redis/pull/1065))
 
 #### Features
 

glide-core/src/protobuf/redis_request.proto

@@ -133,6 +133,7 @@ enum RequestType {
     JsonGet = 89;
     ZRemRangeByScore = 90;
     Time = 91;
+    Zrank = 92;
 }
 
 message Command {

glide-core/src/socket_listener.rs

@@ -363,6 +363,7 @@ fn get_command(request: &Command) -> Option<Cmd> {
         RequestType::JsonGet => Some(cmd("JSON.GET")),
         RequestType::ZRemRangeByScore => Some(cmd("ZREMRANGEBYSCORE")),
         RequestType::Time => Some(cmd("TIME")),
+        RequestType::Zrank => Some(cmd("ZRANK")),
     }
 }
 

python/python/glide/async_commands/core.py

@@ -1634,6 +1634,65 @@ async def zrange_withscores(
             Mapping[str, float], await self._execute_command(RequestType.Zrange, args)
         )
 
+    async def zrank(
+        self,
+        key: str,
+        member: str,
+    ) -> Optional[int]:
+        """
+        Returns the rank of `member` in the sorted set stored at `key`, with scores ordered from low to high.
+
+        See https://redis.io/commands/zrank for more details.
+
+        To get the rank of `member` with it's score, see `zrank_withscore`.
+
+        Args:
+            key (str): The key of the sorted set.
+            member (str): The member whose rank is to be retrieved.
+
+        Returns:
+            Optional[int]: The rank of `member` in the sorted set.
+            If `key` doesn't exist, or if `member` is not present in the set, None will be returned.
+
+            Examples:
+            >>> await client.zrank("my_sorted_set", "member2")
+                1  # Indicates that "member2" has the second-lowest score in the sorted set "my_sorted_set".
+            >>> await client.zrank("my_sorted_set", "non_existing_member")
+                None  # Indicates that "non_existing_member" is not present in the sorted set "my_sorted_set".
+        """
+        return cast(
+            Optional[int], await self._execute_command(RequestType.Zrank, [key, member])
+        )
+
+    async def zrank_withscore(
+        self,
+        key: str,
+        member: str,
+    ) -> Optional[List[Union[int, float]]]:
+        """
+        Returns the rank of `member` in the sorted set stored at `key` with it's score, where scores are ordered from the lowest to highest.
+
+        See https://redis.io/commands/zrank for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            member (str): The member whose rank is to be retrieved.
+
+        Returns:
+            Optional[List[Union[int, float]]]: A list containing the rank and score of `member` in the sorted set.
+            If `key` doesn't exist, or if `member` is not present in the set, None will be returned.
+
+        Examples:
+            >>> await client.zrank_withscore("my_sorted_set", "member2")
+                [1 , 6.0]  # Indicates that "member2" with score 6.0 has the second-lowest score in the sorted set "my_sorted_set".
+            >>> await client.zrank_withscore("my_sorted_set", "non_existing_member")
+                None  # Indicates that "non_existing_member" is not present in the sorted set "my_sorted_set".
+        """
+        return cast(
+            Optional[List[Union[int, float]]],
+            await self._execute_command(RequestType.Zrank, [key, member, "WITHSCORE"]),
+        )
+
     async def zrem(
         self,
         key: str,

python/python/glide/async_commands/transaction.py

@@ -1296,6 +1296,48 @@ def zrange_withscores(
 
         return self.append_command(RequestType.Zrange, args)
 
+    def zrank(
+        self: TTransaction,
+        key: str,
+        member: str,
+    ) -> TTransaction:
+        """
+        Returns the rank of `member` in the sorted set stored at `key`, with scores ordered from low to high.
+
+        See https://redis.io/commands/zrank for more details.
+
+        To get the rank of `member` with it's score, see `zrank_withscore`.
+
+        Args:
+            key (str): The key of the sorted set.
+            member (str): The member whose rank is to be retrieved.
+
+        Commands response:
+            Optional[int]: The rank of `member` in the sorted set.
+            If `key` doesn't exist, or if `member` is not present in the set, None will be returned.
+        """
+        return self.append_command(RequestType.Zrank, [key, member])
+
+    def zrank_withscore(
+        self: TTransaction,
+        key: str,
+        member: str,
+    ) -> TTransaction:
+        """
+        Returns the rank of `member` in the sorted set stored at `key` with it's score, where scores are ordered from the lowest to highest.
+
+        See https://redis.io/commands/zrank for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            member (str): The member whose rank is to be retrieved.
+
+        Commands response:
+            Optional[List[Union[int, float]]]: A list containing the rank and score of `member` in the sorted set.
+            If `key` doesn't exist, or if `member` is not present in the set, None will be returned.
+        """
+        return self.append_command(RequestType.Zrank, [key, member, "WITHSCORE"])
+
     def zrem(
         self: TTransaction,
         key: str,

python/python/glide/constants.py

@@ -14,13 +14,13 @@
 TResult = Union[
     TOK,
     str,
-    List[str],
     List[List[str]],
     int,
     None,
     Dict[str, T],
     float,
     Set[T],
+    List[T],
 ]
 TRequest = Union[RedisRequest, ConnectionRequest]
 # When routing to a single node, response will be T

python/python/tests/test_async_client.py

@@ -1502,6 +1502,27 @@ async def test_zrange_different_types_of_keys(self, redis_client: TRedisClient):
         with pytest.raises(RequestError):
             await redis_client.zrange_withscores(key, RangeByIndex(start=0, stop=1))
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_zrank(self, redis_client: TRedisClient):
+        key = get_random_string(10)
+        members_scores = {"one": 1.5, "two": 2, "three": 3}
+        assert await redis_client.zadd(key, members_scores) == 3
+        assert await redis_client.zrank(key, "one") == 0
+        if not await check_if_server_version_lt(redis_client, "7.2.0"):
+            assert await redis_client.zrank_withscore(key, "one") == [0, 1.5]
+            assert await redis_client.zrank_withscore(key, "non_existing_field") is None
+            assert (
+                await redis_client.zrank_withscore("non_existing_key", "field") is None
+            )
+
+        assert await redis_client.zrank(key, "non_existing_field") is None
+        assert await redis_client.zrank("non_existing_key", "field") is None
+
+        assert await redis_client.set(key, "value") == OK
+        with pytest.raises(RequestError):
+            await redis_client.zrank(key, "one")
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_type(self, redis_client: TRedisClient):

python/python/tests/test_transaction.py

@@ -15,11 +15,13 @@
 from glide.constants import OK, TResult
 from glide.redis_client import RedisClient, RedisClusterClient, TRedisClient
 from tests.conftest import create_client
-from tests.test_async_client import get_random_string
+from tests.test_async_client import check_if_server_version_lt, get_random_string
 
 
-def transaction_test(
-    transaction: Union[Transaction, ClusterTransaction], keyslot: str
+async def transaction_test(
+    transaction: Union[Transaction, ClusterTransaction],
+    keyslot: str,
+    redis_client: TRedisClient,
 ) -> List[TResult]:
     key = "{{{}}}:{}".format(keyslot, get_random_string(3))  # to get the same slot
     key2 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # to get the same slot
@@ -147,6 +149,11 @@ def transaction_test(
 
     transaction.zadd(key8, {"one": 1, "two": 2, "three": 3})
     args.append(3)
+    transaction.zrank(key8, "one")
+    args.append(0)
+    if not await check_if_server_version_lt(redis_client, "7.2.0"):
+        transaction.zrank_withscore(key8, "one")
+        args.append([0, 1])
     transaction.zadd_incr(key8, "one", 3)
     args.append(4)
     transaction.zrem(key8, ["one"])
@@ -252,7 +259,7 @@ async def test_cluster_transaction(self, redis_client: RedisClusterClient):
         keyslot = get_random_string(3)
         transaction = ClusterTransaction()
         transaction.info()
-        expected = transaction_test(transaction, keyslot)
+        expected = await transaction_test(transaction, keyslot, redis_client)
         result = await redis_client.exec(transaction)
         assert isinstance(result, list)
         assert isinstance(result[0], str)
@@ -296,7 +303,7 @@ async def test_standalone_transaction(self, redis_client: RedisClient):
         transaction.get(key)
         transaction.select(0)
         transaction.get(key)
-        expected = transaction_test(transaction, keyslot)
+        expected = await transaction_test(transaction, keyslot, redis_client)
         result = await redis_client.exec(transaction)
         assert isinstance(result, list)
         assert isinstance(result[0], str)