diff --git a/swh/counters/interface.py b/swh/counters/interface.py --- a/swh/counters/interface.py +++ b/swh/counters/interface.py @@ -3,7 +3,7 @@ # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information -from typing import Any, Iterable +from typing import Any, Dict, Iterable, List from swh.core.api import remote_api_endpoint @@ -27,6 +27,11 @@ """Return the number of keys for the provided collection""" ... + @remote_api_endpoint("counts") + def get_counts(self, collections: List[str]) -> Dict[str, int]: + """Return the number of keys for the provided collection""" + ... + @remote_api_endpoint("counters") def get_counters(self) -> Iterable[str]: """Return the list of managed counters""" diff --git a/swh/counters/redis.py b/swh/counters/redis.py --- a/swh/counters/redis.py +++ b/swh/counters/redis.py @@ -4,7 +4,7 @@ # See top-level LICENSE file for more information import logging -from typing import Any, Iterable +from typing import Any, Dict, Iterable, List from redis.client import Redis as RedisClient from redis.exceptions import ConnectionError @@ -56,5 +56,8 @@ def get_count(self, collection: str) -> int: return self.redis_client.pfcount(collection) + def get_counts(self, collections: List[str]) -> Dict[str, int]: + return {coll: self.get_count(coll) for coll in collections} + def get_counters(self) -> Iterable[str]: return self.redis_client.keys() diff --git a/swh/counters/tests/test_redis.py b/swh/counters/tests/test_redis.py --- a/swh/counters/tests/test_redis.py +++ b/swh/counters/tests/test_redis.py @@ -78,3 +78,21 @@ assert 2 == len(counters) assert b"counter1" in counters assert b"counter2" in counters + + +def test__redis_counts(local_redis): + client = RedisClient(host=local_redis.host, port=local_redis.port) + client.pfadd("counter1", b"k1") + client.pfadd("counter1", b"k2") + client.pfadd("counter2", b"k1") + client.pfadd("counter2", b"k2") + client.pfadd("counter2", b"k3") + client.pfadd("counter3", b"k3") + + r = Redis("%s:%d" % (local_redis.host, local_redis.port)) + + counts = r.get_counts(["counter2", "counter1"]) + + assert 2 == len(counts) + assert 2 == counts["counter1"] + assert 3 == counts["counter2"]