Last active
April 11, 2021 00:52
-
-
Save abbasidaniyal/9c3e68ba15cb83e2859efb6be8d0766d to your computer and use it in GitHub Desktop.
redis_backend
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import re | |
| import redis | |
| import random | |
| import pickle | |
| import urllib | |
| import django | |
| from django.conf import settings | |
| from django.core.cache.backends.base import BaseCache, DEFAULT_TIMEOUT | |
| from django.core.cache import cache | |
| from django.core.exceptions import ImproperlyConfigured | |
| from django.utils.functional import cached_property | |
| from django.utils.module_loading import import_string | |
| class BaseCacheSerializer: | |
| def dumps(): | |
| raise NotImplementedError | |
| def loads(): | |
| raise NotImplementedError | |
| class PickleCacheSerializer(BaseCacheSerializer): | |
| def __init__(self, options): | |
| self._pickle_version = -1 | |
| if "PICKLE_VERSION" in options: | |
| try: | |
| self._pickle_version = int(options["PICKLE_VERSION"]) | |
| except (ValueError, TypeError): | |
| raise ImproperlyConfigured("must be an integer") | |
| def dumps(self, value): | |
| return pickle.dumps(value, self._pickle_version) | |
| def loads(self, value): | |
| return pickle.loads(value) | |
| class RedisCacheClient: | |
| def __init__(self, servers, options): | |
| self._servers = servers | |
| self._pools = [None] * len(servers) | |
| self._options = options | |
| self._pool_class = redis.ConnectionPool | |
| self._client = redis.Redis | |
| self._serializer = PickleCacheSerializer(self._options) | |
| self._client_kwargs = {} | |
| username = self._options.get("USERNAME", None) | |
| if username: | |
| self._client_kwargs["username"] = username | |
| password = self._options.get("PASSWORD", None) | |
| if password: | |
| self._client_kwargs["password"] = password | |
| parser = self._options.get("PARSER", redis.connection.PythonParser) | |
| if parser: | |
| if isinstance(parser, str): | |
| parser = import_string(parser) | |
| self._client_kwargs["parser_class"] = parser | |
| 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 _parse_url(self, url): | |
| parsed_url = urllib.parse.urlparse(url) | |
| kwargs = { | |
| "host": parsed_url.hostname, | |
| "port" : parsed_url.port, | |
| "username": parsed_url.username, | |
| "password": parsed_url.password, | |
| } | |
| return kwargs | |
| def _get_connection_pool(self, write): | |
| index = self._get_connection_pool_index(write) | |
| params = self._parse_url(self._servers[index]) | |
| params.update(self._client_kwargs) | |
| if self._pools[index] is None: | |
| if self._servers[index].startswith('unix:'): | |
| self._pools[index] = self._pool_class(connection_class=redis.UnixDomainSocketConnection, path=self._servers[index][5:], **self._client_kwargs) | |
| else: | |
| self._pools[index] = self._pool_class(**params) | |
| return self._pools[index] | |
| def get_client(self, key, write=False): | |
| pool = self._get_connection_pool(write) | |
| return self._client(connection_pool=pool) | |
| def get(self, key): | |
| client = self.get_client(key) | |
| return self._serializer.loads(client.get(key)) | |
| def set(self, key, value, timeout): | |
| client = self.get_client(key, write=True) | |
| client.set(key, self._serializer.dumps(value), ex=timeout) | |
| class RedisShardCacheClient(RedisCacheClient): | |
| def __init__(self, servers, params): | |
| super(RedisShardCacheClient, self).__init__(servers, params) | |
| def _get_connection_pool_index(self, key): | |
| key_hash = hash(key) | |
| return key_hash % len(self._servers) | |
| def get_client(self, key, write=False): | |
| pool = self._get_connection_pool(key) | |
| return self._client(connection_pool=pool) | |
| 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._options = params.get('OPTIONS') or {} | |
| self._client_class = self._options.get("CLIENT", RedisCacheClient) | |
| if isinstance(self._client_class, str): | |
| self._client_class = import_string(self._client_class) | |
| def get_backend_timeout(self, timeout=DEFAULT_TIMEOUT): | |
| if timeout == DEFAULT_TIMEOUT: | |
| timeout = self.default_timeout | |
| else: | |
| timeout = int(timeout) | |
| return timeout | |
| @cached_property | |
| def _caches(self): | |
| return self._client_class(self._servers, self._options) | |
| def get(self, key, default=None, version=None): | |
| key = self.make_key(key, version=version) | |
| self.validate_key(key) | |
| value = self._caches.get(key) | |
| if value is None: | |
| return default | |
| else: | |
| return value | |
| def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None): | |
| key = self.make_key(key, version=version) | |
| self.validate_key(key) | |
| return self._caches.set(key, value, timeout=self.get_backend_timeout(timeout)) | |
| if __name__ == '__main__': | |
| CACHES = { | |
| "default" : { | |
| "BACKEND" : "redis_backend.RedisCache", | |
| "LOCATION" : ["redis://localhost:6379"], | |
| "OPTIONS" : { | |
| "CLIENT" : "redis_backend.RedisShardCacheClient", | |
| "PASSWORD" : "b32099686f854124a2362a702e33069b", | |
| "PARSER": "redis.connection.HiredisParser", | |
| } | |
| } | |
| } | |
| settings.configure(CACHES=CACHES) | |
| django.setup() | |
| cache.set("key", 1) | |
| assert cache.get("key") == 1 | |
| class x: | |
| z = 10 | |
| obj = x() | |
| cache.set("obj", obj) | |
| assert cache.get("obj").z == obj.z |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment