flask-demo-api/flask_demo_api/repository/redis.py

import pickle

import redis  # type: ignore

from flask_demo_api.protocols.models import KeyDTO
from flask_demo_api.protocols.repository import Repository


class RedisRepository(Repository):
    def __init__(self, redis_pool: redis.Redis) -> None:
        self.pool = redis_pool

    def __set(self, key, val):
        data = pickle.dumps(val)
        self.pool.set(key, data)

    def __get(self, key):
        data = self.pool.get(key)
        if data is not None:
            return pickle.loads(data)
        return None

    def get_key(self, obj: KeyDTO) -> KeyDTO | None:
        data = self.__get(obj.key)
        if not data:
            return None
        return KeyDTO(key=obj.key, val=data)

    def add_key(self, obj: KeyDTO) -> KeyDTO | None:
        data = self.__get(obj.key)
        if data:
            return None
        self.__set(obj.key, obj.val)
        return obj

    def put_key(self, obj: KeyDTO) -> KeyDTO | None:
        data = self.__get(obj.key)
        if data is None:
            return None
        self.__set(obj.key, obj.val)
        return obj

    def delete_key(self, obj: KeyDTO) -> None:
        self.pool.delete(obj.key)
        return None