administrative_units_cache.py

"""
API для работы с кэшем дерева субъетов.

"""
# coding: utf-8
from __future__ import absolute_import, unicode_literals

# third party
import ujson

# project
from cache import TREE_CLIENT as redis_client

# u_{unit_id}
UNIT_KEY_TEMPLATE = "u:{0}"

COMPANY_KEY_PATTERN = "cmp:{0}:*"


def make_unit_key(unit_id):
    """Создает ключ для юнита для кэша редиса
    :param unit_id: айди юнита

    :return созданный ключ

    """
    return UNIT_KEY_TEMPLATE.format(unit_id)


def get_unit_hash(unit_instance):
    """Возвращает хэш юнита для хранения в кэше редиса на основе инстанс юнита

    :param unit_instance: инстанс юнита

    :return хэш с данными для кэша редиса

    """
    # айди компаний-прямых чилдов записываем в кэш сразу
    companies = [cid for cid in unit_instance.companies.filter(
        parent_id__isnull=True).values_list("id", flat=True)]
    children = [uid for uid in unit_instance.children.all().values_list(
        "id", flat=True)]
    return {
        "id": unit_instance.id,
        "parent_id": unit_instance.parent_id,
        "children": children,
        "companies": companies,
        "name": unit_instance.name
    }


def set_unit_companies_cache(unit_id):
    """Обновляет список компаний в ноде учреждения в кэше редиса

    :param unit_id: айди юнита

    :return результат выполнения команды редиса hset

    """
    # избегаем циклических импортов
    from administrative_units.models import AdministrativeUnit
    unit_instance = AdministrativeUnit.objects.get(pk=unit_id)
    key = make_unit_key(unit_instance.id)
    companies_list = [cid for cid in unit_instance.companies.filter(
                parent_id__isnull=True).values_list("id", flat=True)]
    return redis_client.hset(key, "companies", companies_list)


def get_unit_dict_from_cache(cache):
    """Формирует словарь данных из кэша редиса

    :param cache - кэш редиса для юнита учреждения

    :return словарь с данными по юниту

    """
    p_id = cache["parent_id"]
    # т.к. редис возвращает списки ввиде списка строк '[1, 2 ,3]'
    # то мы его грузим через json
    children = ujson.loads(cache["children"]) if cache["children"] else []
    companies = ujson.loads(cache["companies"]) if cache["companies"] else []
    _dict = {
        "id": int(cache["id"]),
        "parent_id": int(p_id) if p_id != "None" else None,
        "children": children,
        "companies": companies,
        "name": cache["name"]
    }
    return _dict


def get_unit_node(unit_id):
    """Возвращает ноду юнита из кэша ч4

    :param unit_id айди юнита

    :return ноду юнита из кэша редиса
    """
    key = make_unit_key(unit_id)
    cache = redis_client.hgetall(key)
    return get_unit_dict_from_cache(cache)


def get_units_nodes(units_id=None):
    """Возвращает список нод юнитов

    :param units_id список айди юнитов

    :return список нод юнитов

    """
    pipe = redis_client.pipeline()
    if units_id is None:
        match = make_unit_key("*")
        keys = redis_client.scan_iter(match=match, count=10000)
    else:
        # для каждого айди юнита генерирует ключ
        keys = [make_unit_key(unit_id) for unit_id in units_id]
    for key in keys:
        pipe.hgetall(key)
    # и получаем ноды юнитов по списку ключей
    nodes = [get_unit_dict_from_cache(node) for node in pipe.execute()]
    return nodes


def get_unit_descendants(unit_id, include_self=False):
    """Возвращает список айди все потомков-юнитов юнита

    :param unit_id: айди юнита
    :param include_self: включать ли переданный unit_id айди в список

    :return спиской айди юнитов-потомков

    """
    key = make_unit_key(unit_id)
    node = redis_client.hgetall(key)
    stack = ujson.loads(node["children"])
    # микро-оптимизация доступа
    pop = stack.pop
    extend = stack.extend
    hget = redis_client.hget
    # если мы включаем свой же айди(переданный) то отдаем его сразу
    if include_self:
        yield unit_id
    # пока у нас есть что-то в стеке
    while stack:
        # берем элемет стека
        child_id = pop()
        # отдаем его
        yield child_id
        key = UNIT_KEY_TEMPLATE.format(child_id)
        # получаем список айди чилдов чилда
        child_list = ujson.loads(hget(key, "children") or "[]")
        # увеличиваем стек
        extend(child_list)


def get_unit_descendant_companies(unit_id):
    """Возвращает список айди всех учреждений-потомков всех юнитов по unit_id

    :param unit_id: айди юнита

    :return список айди всех учреждений потомков юнита unit_id включая
    учреждения переданного юнита

    """
    keys = []
    # получим все ключи учреждений по этом юнитам
    for unit_id in get_unit_descendants(unit_id, True):
        match = COMPANY_KEY_PATTERN.format(unit_id)
        iterator = redis_client.scan_iter(count=10000, match=match)
        keys.extend(iterator)
    # немного схитрим и обработам сами ключи учреждений, т.к. он сами хранят id
    _int = int
    rfind = str.rfind
    ids = [_int(k[rfind(k, ":") + 1:]) for k in keys]
    return ids