dj-embe/framework/redis.py

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()