GPQBot

1
"""Repository for work with redis."""
2

3
# flake8: noqa
4
# TODO: Fix linter errors
5
import hashlib
6
import pickle
7
from typing import Any, Hashable, Optional
8

9
import aioredis
10
from starlette.datastructures import URL
11

12

13
class RedisRepo:
14
    redis: aioredis.Redis
15
    dsn: str
16
    prefix: Optional[str]
17
    delimiter: str
18
    expire: Optional[int]
19

20
    def __init__(
21
        self, dsn: URL, prefix: Optional[str] = None, expire: Optional[int] = None
22
    ) -> None:
23
        """Init repository object."""
24

25
        self.dsn = str(dsn)
26
        self.prefix = prefix
27
        self.expire = expire
28
        self.delimiter = "_"
29

30
    @classmethod
31
    async def init(
32
        cls, dsn: URL, prefix: Optional[str] = None, expire: Optional[int] = None
33
    ) -> "RedisRepo":
34
        """Init repository object."""
35

36
        repo = cls(dsn=dsn, prefix=prefix, expire=expire)
37
        repo.redis = await aioredis.create_redis_pool(repo.dsn)
38
        return repo
39

40
    async def close(self) -> None:
41
        """Close connection to redis."""
42

43
        self.redis.close()
44
        await self.redis.wait_closed()
45

46
    def __key(self, arg: Hashable) -> str:
47
        if self.prefix is not None:
48
            _prefix = self.prefix + self.delimiter
49
        else:
50
            _prefix = ""
51

52
        return _prefix + hashlib.md5(pickle.dumps(arg)).hexdigest()
53

54
    async def get(self, key: Hashable, default: Any = None) -> Any:
55
        """Get value from redis."""
56

57
        cached_data = await self.redis.get(self.__key(key))
58
        if cached_data is None:
59
            return default
60
        else:
61
            return pickle.loads(cached_data)
62

63
    async def set(
64
        self, key: Hashable, value: Any, expire: Optional[int] = None
65
    ) -> None:
66
        """Set value into redis."""
67

68
        if expire is None:
69
            expire = self.expire
70
        await self.redis.set(self.__key(key), pickle.dumps(value), expire=expire)
71

72
    async def delete(self, key: Hashable) -> None:
73
        """Remove value from redis."""
74

75
        await self.redis.delete(self.__key(key))
76

77
    async def rget(self, key: Hashable, default: Any = None) -> Any:
78
        """Get value and remove it from redis."""
79

80
        value = await self.get(key, default)
81
        await self.delete(key)
82
        return value