Skip to content

Commit

Permalink
Python: adds DBSIZE command
Browse files Browse the repository at this point in the history
  • Loading branch information
shohamazon committed Mar 20, 2024
1 parent 1961d1e commit 39002f5
Show file tree
Hide file tree
Showing 7 changed files with 75 additions and 0 deletions.
1 change: 1 addition & 0 deletions glide-core/src/protobuf/redis_request.proto
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ enum RequestType {
Time = 89;
Zrank = 90;
Rename = 91;
DBSize = 92;
}

message Command {
Expand Down
1 change: 1 addition & 0 deletions glide-core/src/socket_listener.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
}
}

Expand Down
19 changes: 19 additions & 0 deletions python/python/glide/async_commands/cluster_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand All @@ -257,3 +258,21 @@ 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 nodes, 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 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, [], route))
14 changes: 14 additions & 0 deletions python/python/glide/async_commands/standalone_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, []))
10 changes: 10 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 currently selected database.
"""
return self.append_command(RequestType.DBSize, [])


class Transaction(BaseTransaction):
"""
Expand Down
25 changes: 25 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1559,6 +1559,31 @@ 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

key_value_pairs = [(get_random_string(10), "foo") for _ in range(10)]
assert await redis_client.dbsize() == 0

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):
cluster_nodes = await redis_client.custom_command(["CLUSTER", "NODES"])
assert isinstance(cluster_nodes, (str, list))
cluster_nodes = get_first_result(cluster_nodes)
replica_count = cluster_nodes.count("slave")
master_count = cluster_nodes.count("master")
assert await redis_client.dbsize(AllNodes()) == 10 * (
replica_count / master_count + 1
)
else:
assert await redis_client.select(1) == OK
assert await redis_client.dbsize() == 0


class TestCommandsUnitTests:
def test_expiry_cmd_args(self):
Expand Down
5 changes: 5 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 39002f5

Please sign in to comment.