48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
import pickle
|
|
from logging import Logger
|
|
|
|
import redis.asyncio as redis
|
|
from redis.asyncio.lock import Lock
|
|
|
|
from config import RedisConfig
|
|
|
|
|
|
class Redis:
|
|
def __init__(self, logger: Logger, config: RedisConfig, rootKeyName: str) -> None:
|
|
self._client = redis.Redis(
|
|
host=config.host,
|
|
port=config.port,
|
|
password=config.password,
|
|
auto_close_connection_pool=False,
|
|
)
|
|
self._locks: dict[str, Lock] = {}
|
|
self.rootKeyName = rootKeyName
|
|
self.logger = logger
|
|
|
|
async def _get_lock(self, key) -> Lock:
|
|
if key not in self._locks:
|
|
self._locks[key] = self._client.lock(key)
|
|
return self._locks[key]
|
|
|
|
async def acquire(self, key: str) -> None:
|
|
lock = await self._get_lock(f"{self.rootKeyName}:queue:{key}")
|
|
await lock.acquire()
|
|
|
|
async def release(self, key: str) -> None:
|
|
lock = await self._get_lock(f"{self.rootKeyName}:queue:{key}")
|
|
await lock.release()
|
|
|
|
async def get(self, key: str):
|
|
self.logger.info(f"get value {key} from redis")
|
|
value = await self._client.get(f"{self.rootKeyName}:queue:{key}")
|
|
if value:
|
|
return pickle.loads(value)
|
|
return None
|
|
|
|
async def set(self, key: str, value) -> None:
|
|
self.logger.info(f"set value {key} to redis")
|
|
await self._client.set(f"{self.rootKeyName}:queue:{key}", pickle.dumps(value))
|
|
|
|
async def close(self) -> None:
|
|
await self._client.close()
|