Python: adds DBSIZE command

CHANGELOG.md

@@ -51,6 +51,7 @@
 * Node: Added STRLEN command ([#993](https://github.com/aws/glide-for-redis/pull/993))
 * Node: Added LINDEX command ([#999](https://github.com/aws/glide-for-redis/pull/999))
 * Python, Node: Added ZPOPMAX command ([#996](https://github.com/aws/glide-for-redis/pull/996), [#1009](https://github.com/aws/glide-for-redis/pull/1009))
+* Python: Added DBSIZE command ([#1040](https://github.com/aws/glide-for-redis/pull/1040))
 
 #### Features
 * Python, Node: Added support in Lua Scripts ([#775](https://github.com/aws/glide-for-redis/pull/775), [#860](https://github.com/aws/glide-for-redis/pull/860))

glide-core/src/protobuf/redis_request.proto

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

glide-core/src/socket_listener.rs

@@ -363,6 +363,7 @@ fn get_command(request: &Command) -> Option<Cmd> {
         RequestType::Time => Some(cmd("TIME")),
         RequestType::Zrank => Some(cmd("ZRANK")),
         RequestType::Rename => Some(cmd("RENAME")),
+        RequestType::DBSize => Some(cmd("DBSIZE")),
     }
 }
 

python/python/glide/async_commands/cluster_commands.py

@@ -237,6 +237,7 @@ async def client_getname(
         """
         Get the name of the connection to which the request is routed.
         See https://redis.io/commands/client-getname/ for more details.
+
         Args:
             route (Optional[Route]): The command will be routed to a random node, unless `route` is provided,
             in which case the client will route the command to the nodes defined by `route`.
@@ -257,3 +258,22 @@ async def client_getname(
             TClusterResponse[Optional[str]],
             await self._execute_command(RequestType.ClientGetName, [], route),
         )
+
+    async def dbsize(self, route: Optional[Route] = None) -> int:
+        """
+        Returns the number of keys in the database.
+        See https://redis.io/commands/dbsize for more details.
+
+        Args:
+            route (Optional[Route]): The command will be routed to all primaries, unless `route` is provided,
+            in which case the client will route the command to the nodes defined by `route`.
+
+        Returns:
+            int: The number of keys in the database.
+            In the case of routing the query to multiple nodes, returns the aggregated number of keys across the different nodes.
+
+        Examples:
+            >>> await client.dbsize()
+                10  # Indicates there are 10 keys in the cluster.
+        """
+        return cast(int, await self._execute_command(RequestType.DBSize, [], route))

python/python/glide/async_commands/standalone_commands.py

@@ -192,3 +192,17 @@ async def client_getname(self) -> Optional[str]:
         return cast(
             Optional[str], await self._execute_command(RequestType.ClientGetName, [])
         )
+
+    async def dbsize(self) -> int:
+        """
+        Returns the number of keys in the currently selected database.
+        See https://redis.io/commands/dbsize for more details.
+
+        Returns:
+            int: The number of keys in the currently selected database.
+
+        Examples:
+            >>> await client.dbsize()
+                10  # Indicates there are 10 keys in the current database.
+        """
+        return cast(int, await self._execute_command(RequestType.DBSize, []))

python/python/glide/async_commands/transaction.py

@@ -1376,6 +1376,16 @@ def zscore(self: TTransaction, key: str, member: str) -> TTransaction:
         """
         return self.append_command(RequestType.ZScore, [key, member])
 
+    def dbsize(self: TTransaction) -> TTransaction:
+        """
+        Returns the number of keys in the currently selected database.
+        See https://redis.io/commands/dbsize for more details.
+
+        Commands response:
+            int: The number of keys in the database.
+        """
+        return self.append_command(RequestType.DBSize, [])
+
 
 class Transaction(BaseTransaction):
     """

python/python/tests/test_async_client.py

@@ -1559,6 +1559,27 @@ async def test_echo(self, redis_client: TRedisClient):
         message = get_random_string(5)
         assert await redis_client.echo(message) == message
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_dbsize(self, redis_client: TRedisClient):
+        assert await redis_client.custom_command(["FLUSHALL"]) == OK
+
+        assert await redis_client.dbsize() == 0
+        key_value_pairs = [(get_random_string(10), "foo") for _ in range(10)]
+
+        for key, value in key_value_pairs:
+            assert await redis_client.set(key, value) == OK
+        assert await redis_client.dbsize() == 10
+
+        if isinstance(redis_client, RedisClusterClient):
+            assert await redis_client.custom_command(["FLUSHALL"]) == OK
+            key = get_random_string(5)
+            assert await redis_client.set(key, value) == OK
+            assert await redis_client.dbsize(SlotKeyRoute(SlotType.PRIMARY, key)) == 1
+        else:
+            assert await redis_client.select(1) == OK
+            assert await redis_client.dbsize() == 0
+
 
 class TestCommandsUnitTests:
     def test_expiry_cmd_args(self):

python/python/tests/test_transaction.py

@@ -36,6 +36,9 @@ async def transaction_test(
     value2 = get_random_string(5)
     args: List[TResult] = []
 
+    transaction.dbsize()
+    args.append(0)
+
     transaction.set(key, value)
     args.append(OK)
     transaction.get(key)
@@ -256,6 +259,7 @@ async def test_transaction_exec_abort(self, redis_client: TRedisClient):
     @pytest.mark.parametrize("cluster_mode", [True])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_cluster_transaction(self, redis_client: RedisClusterClient):
+        assert await redis_client.custom_command(["FLUSHALL"]) == OK
         keyslot = get_random_string(3)
         transaction = ClusterTransaction()
         transaction.info()
@@ -293,6 +297,7 @@ async def test_can_return_null_on_watch_transaction_failures(
     @pytest.mark.parametrize("cluster_mode", [False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_standalone_transaction(self, redis_client: RedisClient):
+        assert await redis_client.custom_command(["FLUSHALL"]) == OK
         keyslot = get_random_string(3)
         key = "{{{}}}:{}".format(keyslot, get_random_string(3))  # to get the same slot
         value = get_random_string(5)