redis cache backend import pickle import random import re from django core cache backends base import defaulttimeout basecache from django utils functional import cachedproperty from django utils moduleloading import importstring class redisserializer def initself protocolnone self protocol pickle highestprotocol if protocol is none else protocol def dumpsself obj only skip pickling for integers a int subclasses as bool should be pickled if typeobj is int return obj return pickle dumpsobj self protocol def loadsself data try return intdata except valueerror return pickle loadsdata class rediscacheclient def init self servers serializernone poolclassnone parserclassnone options import redis self lib redis self servers servers self pools self client self lib redis if isinstancepoolclass str poolclass importstringpoolclass self poolclass poolclass or self lib connectionpool if isinstanceserializer str serializer importstringserializer if callableserializer serializer serializer self serializer serializer or redisserializer if isinstanceparserclass str parserclass importstringparserclass parserclass parserclass or self lib connection defaultparser self pooloptions parserclass parserclass options def getconnectionpoolindexself write write to the first server read from other servers if there are more otherwise read from the first server if write or lenself servers 1 return 0 return random randint1 lenself servers 1 def getconnectionpoolself write index self getconnectionpoolindexwrite if index not in self pools self poolsindex self poolclass fromurl self serversindex self pooloptions return self poolsindex def getclientself keynone writefalse key is used so that the method signature remains the same and custom cache client can be implemented which might require the key to select the server e g sharding pool self getconnectionpoolwrite return self clientconnectionpoolpool def addself key value timeout client self getclientkey writetrue value self serializer dumpsvalue if timeout 0 if ret boolclient setkey value nxtrue client deletekey return ret else return boolclient setkey value extimeout nxtrue def getself key default client self getclientkey value client getkey return default if value is none else self serializer loadsvalue def setself key value timeout client self getclientkey writetrue value self serializer dumpsvalue if timeout 0 client deletekey else client setkey value extimeout def touchself key timeout client self getclientkey writetrue if timeout is none return boolclient persistkey else return boolclient expirekey timeout def deleteself key client self getclientkey writetrue return boolclient deletekey def getmanyself keys client self getclientnone ret client mgetkeys return k self serializer loadsv for k v in zipkeys ret if v is not none def haskeyself key client self getclientkey return boolclient existskey def incrself key delta client self getclientkey writetrue if not client existskey raise valueerrorkey s not found key return client incrkey delta def setmanyself data timeout client self getclientnone writetrue pipeline client pipeline pipeline msetk self serializer dumpsv for k v in data items if timeout is not none setting timeout for each key as redis does not support timeout with mset for key in data pipeline expirekey timeout pipeline execute def deletemanyself keys client self getclientnone writetrue client deletekeys def clearself client self getclientnone writetrue return boolclient flushdb class rediscachebasecache def initself server params super initparams if isinstanceserver str self servers re split server else self servers server self class rediscacheclient self options params getoptions cachedproperty def cacheself return self classself servers self options def getbackendtimeoutself timeoutdefaulttimeout if timeout defaulttimeout timeout self defaulttimeout the key will be made persistent if none used as a timeout nonpositive values will cause the key to be deleted return none if timeout is none else max0 inttimeout def addself key value timeoutdefaulttimeout versionnone key self makeandvalidatekeykey versionversion return self cache addkey value self getbackendtimeouttimeout def getself key defaultnone versionnone key self makeandvalidatekeykey versionversion return self cache getkey default def setself key value timeoutdefaulttimeout versionnone key self makeandvalidatekeykey versionversion self cache setkey value self getbackendtimeouttimeout def touchself key timeoutdefaulttimeout versionnone key self makeandvalidatekeykey versionversion return self cache touchkey self getbackendtimeouttimeout def deleteself key versionnone key self makeandvalidatekeykey versionversion return self cache deletekey def getmanyself keys versionnone keymap self makeandvalidatekeykey versionversion key for key in keys ret self cache getmanykeymap keys return keymapk v for k v in ret items def haskeyself key versionnone key self makeandvalidatekeykey versionversion return self cache haskeykey def incrself key delta1 versionnone key self makeandvalidatekeykey versionversion return self cache incrkey delta def setmanyself data timeoutdefaulttimeout versionnone if not data return safedata for key value in data items key self makeandvalidatekeykey versionversion safedatakey value self cache setmanysafedata self getbackendtimeouttimeout return def deletemanyself keys versionnone if not keys return safekeys self makeandvalidatekeykey versionversion for key in keys self cache deletemanysafekeys def clearself return self cache clear only skip pickling for integers a int subclasses as bool should be pickled write to the first server read from other servers if there are more otherwise read from the first server key is used so that the method signature remains the same and custom cache client can be implemented which might require the key to select the server e g sharding setting timeout for each key as redis does not support timeout with mset the key will be made persistent if none used as a timeout non positive values will cause the key to be deleted
import pickle
import random
import re

from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache
from django.utils.functional import cached_property
from django.utils.module_loading import import_string


class RedisSerializer:
    def __init__(self, protocol=None):
        self.protocol = pickle.HIGHEST_PROTOCOL if protocol is None else protocol

    def dumps(self, obj):
        
        
        if type(obj) is int:
            return obj
        return pickle.dumps(obj, self.protocol)

    def loads(self, data):
        try:
            return int(data)
        except ValueError:
            return pickle.loads(data)


class RedisCacheClient:
    def __init__(
        self,
        servers,
        serializer=None,
        pool_class=None,
        parser_class=None,
        **options,
    ):
        import redis

        self._lib = redis
        self._servers = servers
        self._pools = {}

        self._client = self._lib.Redis

        if isinstance(pool_class, str):
            pool_class = import_string(pool_class)
        self._pool_class = pool_class or self._lib.ConnectionPool

        if isinstance(serializer, str):
            serializer = import_string(serializer)
        if callable(serializer):
            serializer = serializer()
        self._serializer = serializer or RedisSerializer()

        if isinstance(parser_class, str):
            parser_class = import_string(parser_class)
        parser_class = parser_class or self._lib.connection.DefaultParser

        self._pool_options = {"parser_class": parser_class, **options}

    def _get_connection_pool_index(self, write):
        
        
        if write or len(self._servers) == 1:
            return 0
        return random.randint(1, len(self._servers) - 1)

    def _get_connection_pool(self, write):
        index = self._get_connection_pool_index(write)
        if index not in self._pools:
            self._pools[index] = self._pool_class.from_url(
                self._servers[index],
                **self._pool_options,
            )
        return self._pools[index]

    def get_client(self, key=None, *, write=False):
        
        
        
        pool = self._get_connection_pool(write)
        return self._client(connection_pool=pool)

    def add(self, key, value, timeout):
        client = self.get_client(key, write=True)
        value = self._serializer.dumps(value)

        if timeout == 0:
            if ret := bool(client.set(key, value, nx=True)):
                client.delete(key)
            return ret
        else:
            return bool(client.set(key, value, ex=timeout, nx=True))

    def get(self, key, default):
        client = self.get_client(key)
        value = client.get(key)
        return default if value is None else self._serializer.loads(value)

    def set(self, key, value, timeout):
        client = self.get_client(key, write=True)
        value = self._serializer.dumps(value)
        if timeout == 0:
            client.delete(key)
        else:
            client.set(key, value, ex=timeout)

    def touch(self, key, timeout):
        client = self.get_client(key, write=True)
        if timeout is None:
            return bool(client.persist(key))
        else:
            return bool(client.expire(key, timeout))

    def delete(self, key):
        client = self.get_client(key, write=True)
        return bool(client.delete(key))

    def get_many(self, keys):
        client = self.get_client(None)
        ret = client.mget(keys)
        return {
            k: self._serializer.loads(v) for k, v in zip(keys, ret) if v is not None
        }

    def has_key(self, key):
        client = self.get_client(key)
        return bool(client.exists(key))

    def incr(self, key, delta):
        client = self.get_client(key, write=True)
        if not client.exists(key):
            raise ValueError("Key '%s' not found." % key)
        return client.incr(key, delta)

    def set_many(self, data, timeout):
        client = self.get_client(None, write=True)
        pipeline = client.pipeline()
        pipeline.mset({k: self._serializer.dumps(v) for k, v in data.items()})

        if timeout is not None:
            
            
            for key in data:
                pipeline.expire(key, timeout)
        pipeline.execute()

    def delete_many(self, keys):
        client = self.get_client(None, write=True)
        client.delete(*keys)

    def clear(self):
        client = self.get_client(None, write=True)
        return bool(client.flushdb())


class RedisCache(BaseCache):
    def __init__(self, server, params):
        super().__init__(params)
        if isinstance(server, str):
            self._servers = re.split("[;,]", server)
        else:
            self._servers = server

        self._class = RedisCacheClient
        self._options = params.get("OPTIONS", {})

    @cached_property
    def _cache(self):
        return self._class(self._servers, **self._options)

    def get_backend_timeout(self, timeout=DEFAULT_TIMEOUT):
        if timeout == DEFAULT_TIMEOUT:
            timeout = self.default_timeout
        
        
        return None if timeout is None else max(0, int(timeout))

    def add(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.add(key, value, self.get_backend_timeout(timeout))

    def get(self, key, default=None, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.get(key, default)

    def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
        key = self.make_and_validate_key(key, version=version)
        self._cache.set(key, value, self.get_backend_timeout(timeout))

    def touch(self, key, timeout=DEFAULT_TIMEOUT, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.touch(key, self.get_backend_timeout(timeout))

    def delete(self, key, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.delete(key)

    def get_many(self, keys, version=None):
        key_map = {
            self.make_and_validate_key(key, version=version): key for key in keys
        }
        ret = self._cache.get_many(key_map.keys())
        return {key_map[k]: v for k, v in ret.items()}

    def has_key(self, key, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.has_key(key)

    def incr(self, key, delta=1, version=None):
        key = self.make_and_validate_key(key, version=version)
        return self._cache.incr(key, delta)

    def set_many(self, data, timeout=DEFAULT_TIMEOUT, version=None):
        if not data:
            return []
        safe_data = {}
        for key, value in data.items():
            key = self.make_and_validate_key(key, version=version)
            safe_data[key] = value
        self._cache.set_many(safe_data, self.get_backend_timeout(timeout))
        return []

    def delete_many(self, keys, version=None):
        if not keys:
            return
        safe_keys = [self.make_and_validate_key(key, version=version) for key in keys]
        self._cache.delete_many(safe_keys)

    def clear(self):
        return self._cache.clear()