RedisBackend.py

import os, uuid
from typing import Any, Dict, List, Optional, Set, Tuple
from redis.client import Redis
from .._Backend import _Backend
from ..Tools import key_to_str
from .Mutex import Mutex

DEFAULT_SERVICE_HOST = '127.0.0.1'
DEFAULT_SERVICE_PORT = 6379
DEFAULT_DATABASE_ID  = 0

def get_setting(settings : Dict[str, Any], name : str, default : Any) -> Any:
    value = settings.get(name, os.environ.get(name))
    return default if value is None else value

class RedisBackend(_Backend):
    def __init__(self, **settings) -> None:
        host = get_setting(settings, 'REDIS_SERVICE_HOST', DEFAULT_SERVICE_HOST)
        port = get_setting(settings, 'REDIS_SERVICE_PORT', DEFAULT_SERVICE_PORT)
        dbid = get_setting(settings, 'REDIS_DATABASE_ID',  DEFAULT_DATABASE_ID )
        self._client = Redis.from_url('redis://{host}:{port}/{dbid}'.format(host=host, port=port, dbid=dbid))
        self._mutex = Mutex(self._client)

    def lock(self, keys : List[List[str]], owner_key : Optional[str] = None) -> Tuple[bool, str]:
        str_keys = {key_to_str(key) for key in keys}
        owner_key = str(uuid.uuid4()) if owner_key is None else owner_key
        return self._mutex.acquire(str_keys, owner_key=owner_key, blocking=True)

    def unlock(self, keys : List[List[str]], owner_key : str) -> bool:
        str_keys = {key_to_str(key) for key in keys}
        return self._mutex.release(str_keys, owner_key)

    def keys(self) -> list:
        return [k.decode('UTF-8') for k in self._client.keys()]

    def exists(self, key : List[str]) -> bool:
        str_key = key_to_str(key)
        return self._client.exists(str_key) == 1

    def delete(self, key : List[str]) -> bool:
        str_key = key_to_str(key)
        return self._client.delete(str_key) == 1

    def dict_get(self, key : List[str], fields : List[str] = []) -> Dict[str, str]:
        str_key = key_to_str(key)
        if len(fields) == 0:
            keys_values = self._client.hgetall(str_key).items()
        else:
            fields = list(fields)
            keys_values = zip(fields, self._client.hmget(str_key, fields))

        attributes = {}
        for key,value in keys_values:
            str_key = key.decode('UTF-8') if isinstance(key, bytes) else key
            attributes[str_key] = value.decode('UTF-8') if isinstance(value, bytes) else value
        return attributes

    def dict_update(self, key : List[str], fields : Dict[str, str] = {}) -> None:
        str_key = key_to_str(key)
        if len(fields) > 0:
            self._client.hset(str_key, mapping=fields)

    def dict_delete(self, key : List[str], fields : List[str] = []) -> None:
        str_key = key_to_str(key)
        if len(fields) == 0:
            self._client.delete(str_key)
        else:
            self._client.hdel(str_key, set(fields))

    def list_get_all(self, key : List[str]) -> List[str]:
        str_key = key_to_str(key)
        return list(map(lambda m: m.decode('UTF-8'), self._client.lrange(str_key, 0, -1)))

    def list_push_last(self, key : List[str], item : str) -> None:
        str_key = key_to_str(key)
        self._client.rpush(str_key, item)

    def list_remove_first_occurrence(self, key : List[str], item: str) -> None:
        str_key = key_to_str(key)
        self._client.lrem(str_key, 1, item)

    def set_add(self, key : List[str], item : str) -> None:
        str_key = key_to_str(key)
        self._client.sadd(str_key, item)

    def set_has(self, key : List[str], item : str) -> bool:
        str_key = key_to_str(key)
        return self._client.sismember(str_key, item) == 1

    def set_get_all(self, key : List[str]) -> Set[str]:
        str_key = key_to_str(key)
        return set(map(lambda m: m.decode('UTF-8'), self._client.smembers(str_key)))

    def set_remove(self, key : List[str], item : str) -> None:
        str_key = key_to_str(key)
        self._client.srem(str_key, item)

    def dump(self) -> List[Tuple[str, str, str]]:
        entries = []
        for str_key in self._client.keys():
            str_key = str_key.decode('UTF-8')
            key_type = self._client.type(str_key)
            if key_type is not None: key_type = key_type.decode('UTF-8')
            key_type = {
                'hash'  : 'dict',
                'list'  : 'list',
                'set'   : 'set',
                'string': 'str',
            }.get(key_type)
            key_content = {
                'dict': lambda key: {k.decode('UTF-8'):v.decode('UTF-8') for k,v in self._client.hgetall(key).items()},
                'list': lambda key: [m.decode('UTF-8') for m in self._client.lrange(key, 0, -1)],
                'set' : lambda key: {m.decode('UTF-8') for m in self._client.smembers(key)},
                'str' : lambda key: self._client.get(key).decode('UTF-8'),
            }.get(key_type, lambda key: 'UNSUPPORTED_TYPE')
            entries.append((str_key, key_type, key_content(str_key)))
        return entries