cache.py 2.45 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import aioredis
import asyncio
import json

from django.conf import settings


class Backend:
    class NotFound(Exception):
        pass

    async def get(self, key):
        raise NotImplementedError

Eliot Berriot's avatar
Eliot Berriot committed
15
    async def set(self, key, expire=None):
16
17
        raise NotImplementedError

18
19
20
    async def close(self):
        return

21
22
23
24
25
26
27
28
29
30
31

class Dummy(Backend):
    def __init__(self):
        self._cache = {}

    async def get(self, key):
        try:
            return self._cache[key]
        except KeyError:
            raise self.NotFound(key)

Eliot Berriot's avatar
Eliot Berriot committed
32
    async def set(self, key, value, expire=None):
33
34
35
        self._cache[key] = value


36
37
38
39
40
41
42
43
44
45
46
47
class Noop(Backend):
    """
    A backend that caches nothing
    """

    async def get(self, key):
        raise self.NotFound(key)

    async def set(self, key, value, expire=None):
        return


48
class Redis(Backend):
Eliot Berriot's avatar
Eliot Berriot committed
49
    def __init__(self, params, write_only=False):
50
51
        self.params = params
        self._redis = None
Eliot Berriot's avatar
Eliot Berriot committed
52
        self.write_only = write_only
53
54
55
56
57
58
59
60
61
62

    async def redis(self):
        if self._redis:
            return self._redis
        self._redis = await aioredis.create_redis(
            loop=asyncio.get_event_loop(), **self.params
        )
        return self._redis

    async def get(self, key):
Eliot Berriot's avatar
Eliot Berriot committed
63
64
        if self.write_only:
            raise self.NotFound(key)
65
66
67
68
69
        r = await self.redis()
        try:
            v = await r.get(key)
        except KeyError:
            raise self.NotFound(key)
70
71
        if v is None:
            raise self.NotFound(key)
Eliot Berriot's avatar
Eliot Berriot committed
72
73
74
75
76
        try:
            return json.loads(v)
        except TypeError:
            # Null, empty string, etc.
            return v
77

78
79
80
    async def set(self, key, value, expire=None):
        if expire is None:
            expire = settings.CACHE_DEFAULT_EXPIRATION
81
        r = await self.redis()
82
83
84
        if expire:
            return await r.setex(key, expire, json.dumps(value))
        return await r.set(key, json.dumps(value))
85

86
87
88
89
90
    async def close(self):
        if not self._redis:
            return
        await self._redis.close()

91

92
93
94
_DEFAULT = None


Eliot Berriot's avatar
Eliot Berriot committed
95
96
97
98
99
100
101
102
103
104
105
106
_WRITE_ONLY_DEFAULT = None


def get_write_only_default():
    global _WRITE_ONLY_DEFAULT
    if _WRITE_ONLY_DEFAULT:
        return _WRITE_ONLY_DEFAULT
    _WRITE_ONLY_DEFAULT = Redis(settings.ASYNC_REDIS_PARAMS, write_only=True)
    return _WRITE_ONLY_DEFAULT


def get_default(**kwargs):
107
108
109
    global _DEFAULT
    if _DEFAULT:
        return _DEFAULT
Eliot Berriot's avatar
Eliot Berriot committed
110
    _DEFAULT = Redis(settings.ASYNC_REDIS_PARAMS, **kwargs)
111
    return _DEFAULT