2016-08-09 09:27:37 +00:00
|
|
|
# ~*~ coding: utf-8 ~*~
|
|
|
|
#
|
2019-08-21 12:27:21 +00:00
|
|
|
from treelib import Tree
|
2019-09-26 07:09:10 +00:00
|
|
|
from treelib.exceptions import NodeIDAbsentError
|
2019-08-21 12:27:21 +00:00
|
|
|
from collections import defaultdict
|
|
|
|
from copy import deepcopy
|
2018-01-26 06:47:10 +00:00
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
from common.utils import get_object_or_none, get_logger, timeit
|
2019-09-26 07:09:10 +00:00
|
|
|
from .models import SystemUser, Asset
|
2016-09-04 09:43:03 +00:00
|
|
|
|
2017-12-06 10:31:51 +00:00
|
|
|
|
2019-06-27 13:43:10 +00:00
|
|
|
logger = get_logger(__file__)
|
2017-12-07 05:01:33 +00:00
|
|
|
|
|
|
|
|
2017-12-10 16:29:25 +00:00
|
|
|
def get_system_user_by_name(name):
|
|
|
|
system_user = get_object_or_none(SystemUser, name=name)
|
|
|
|
return system_user
|
2017-12-12 04:19:45 +00:00
|
|
|
|
|
|
|
|
2018-12-10 02:11:54 +00:00
|
|
|
def get_system_user_by_id(id):
|
|
|
|
system_user = get_object_or_none(SystemUser, id=id)
|
|
|
|
return system_user
|
|
|
|
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
class TreeService(Tree):
|
|
|
|
tag_sep = ' / '
|
2019-06-27 13:43:10 +00:00
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.nodes_assets_map = defaultdict(set)
|
|
|
|
self.all_nodes_assets_map = {}
|
2019-09-26 07:09:10 +00:00
|
|
|
self._invalid_assets = frozenset()
|
2019-06-27 13:43:10 +00:00
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
@classmethod
|
|
|
|
@timeit
|
|
|
|
def new(cls):
|
|
|
|
from .models import Node
|
2019-08-31 04:07:03 +00:00
|
|
|
from orgs.utils import tmp_to_root_org
|
|
|
|
|
|
|
|
with tmp_to_root_org():
|
2019-09-25 08:31:18 +00:00
|
|
|
all_nodes = list(Node.objects.all().values("key", "value"))
|
|
|
|
all_nodes.sort(key=lambda x: len(x["key"].split(":")))
|
2019-08-31 04:07:03 +00:00
|
|
|
tree = cls()
|
|
|
|
tree.create_node(tag='', identifier='')
|
|
|
|
for node in all_nodes:
|
2019-09-25 08:31:18 +00:00
|
|
|
key = node["key"]
|
|
|
|
value = node["value"]
|
|
|
|
parent_key = ":".join(key.split(":")[:-1])
|
2019-09-26 07:09:10 +00:00
|
|
|
tree.safe_create_node(
|
2019-09-25 08:31:18 +00:00
|
|
|
tag=value, identifier=key,
|
|
|
|
parent=parent_key,
|
2019-08-31 04:07:03 +00:00
|
|
|
)
|
|
|
|
tree.init_assets()
|
2019-08-21 12:27:21 +00:00
|
|
|
return tree
|
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
@timeit
|
2019-08-21 12:27:21 +00:00
|
|
|
def init_assets(self):
|
2019-09-11 13:22:25 +00:00
|
|
|
from orgs.utils import tmp_to_root_org
|
|
|
|
self.all_nodes_assets_map = {}
|
2019-09-11 13:35:48 +00:00
|
|
|
self.nodes_assets_map = defaultdict(set)
|
2019-09-11 13:22:25 +00:00
|
|
|
with tmp_to_root_org():
|
|
|
|
queryset = Asset.objects.all().values_list('id', 'nodes__key')
|
2019-09-26 07:09:10 +00:00
|
|
|
invalid_assets = Asset.objects.filter(is_active=False)\
|
|
|
|
.values_list('id', flat=True)
|
|
|
|
self._invalid_assets = frozenset(invalid_assets)
|
2019-09-11 13:22:25 +00:00
|
|
|
for asset_id, key in queryset:
|
|
|
|
if not key:
|
|
|
|
continue
|
|
|
|
self.nodes_assets_map[key].add(asset_id)
|
2019-08-21 12:27:21 +00:00
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def safe_create_node(self, **kwargs):
|
|
|
|
parent = kwargs.get("parent")
|
|
|
|
if not self.contains(parent):
|
|
|
|
kwargs['parent'] = self.root
|
|
|
|
self.create_node(**kwargs)
|
|
|
|
|
2019-09-24 07:00:32 +00:00
|
|
|
def all_children_ids(self, nid, with_self=True):
|
2019-08-21 12:27:21 +00:00
|
|
|
children_ids = self.expand_tree(nid)
|
|
|
|
if not with_self:
|
|
|
|
next(children_ids)
|
2019-09-24 07:00:32 +00:00
|
|
|
return list(children_ids)
|
|
|
|
|
|
|
|
def all_children(self, nid, with_self=True, deep=False):
|
|
|
|
children_ids = self.all_children_ids(nid, with_self=with_self)
|
2019-08-21 12:27:21 +00:00
|
|
|
return [self.get_node(i, deep=deep) for i in children_ids]
|
|
|
|
|
|
|
|
def ancestors(self, nid, with_self=False, deep=False):
|
|
|
|
ancestor_ids = list(self.rsearch(nid))
|
|
|
|
ancestor_ids.pop()
|
|
|
|
if not with_self:
|
|
|
|
ancestor_ids.pop(0)
|
|
|
|
return [self.get_node(i, deep=deep) for i in ancestor_ids]
|
|
|
|
|
|
|
|
def get_node_full_tag(self, nid):
|
2019-08-28 03:43:55 +00:00
|
|
|
ancestors = self.ancestors(nid, with_self=True)
|
2019-08-21 12:27:21 +00:00
|
|
|
ancestors.reverse()
|
2019-08-28 03:43:55 +00:00
|
|
|
return self.tag_sep.join([n.tag for n in ancestors])
|
2019-08-21 12:27:21 +00:00
|
|
|
|
|
|
|
def get_family(self, nid, deep=False):
|
|
|
|
ancestors = self.ancestors(nid, with_self=False, deep=deep)
|
|
|
|
children = self.all_children(nid, with_self=False)
|
|
|
|
return ancestors + [self[nid]] + children
|
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
@staticmethod
|
|
|
|
def is_parent(child, parent):
|
|
|
|
parent_id = child.bpointer
|
|
|
|
return parent_id == parent.identifier
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
def root_node(self):
|
|
|
|
return self.get_node(self.root)
|
|
|
|
|
|
|
|
def get_node(self, nid, deep=False):
|
|
|
|
node = super().get_node(nid)
|
|
|
|
if deep:
|
|
|
|
node = self.copy_node(node)
|
|
|
|
return node
|
|
|
|
|
|
|
|
def parent(self, nid, deep=False):
|
|
|
|
parent = super().parent(nid)
|
|
|
|
if deep:
|
|
|
|
parent = self.copy_node(parent)
|
|
|
|
return parent
|
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def set_assets(self, nid, assets):
|
|
|
|
self.nodes_assets_map[nid] = set(assets)
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
def assets(self, nid):
|
2019-08-31 04:07:03 +00:00
|
|
|
assets = self.nodes_assets_map[nid]
|
|
|
|
return assets
|
2019-08-21 12:27:21 +00:00
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def valid_assets(self, nid):
|
|
|
|
return set(self.assets(nid)) - set(self._invalid_assets)
|
2019-08-21 12:27:21 +00:00
|
|
|
|
|
|
|
def all_assets(self, nid):
|
|
|
|
assets = self.all_nodes_assets_map.get(nid)
|
|
|
|
if assets:
|
|
|
|
return assets
|
|
|
|
assets = set(self.assets(nid))
|
|
|
|
children = self.children(nid)
|
|
|
|
for child in children:
|
|
|
|
assets.update(self.all_assets(child.identifier))
|
2019-09-26 07:09:10 +00:00
|
|
|
self.all_nodes_assets_map[nid] = assets
|
2019-08-21 12:27:21 +00:00
|
|
|
return assets
|
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def all_valid_assets(self, nid):
|
|
|
|
return set(self.all_assets(nid)) - set(self._invalid_assets)
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
def assets_amount(self, nid):
|
|
|
|
return len(self.all_assets(nid))
|
2019-06-27 13:43:10 +00:00
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def valid_assets_amount(self, nid):
|
|
|
|
return len(self.all_valid_assets(nid))
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
@staticmethod
|
|
|
|
def copy_node(node):
|
|
|
|
new_node = deepcopy(node)
|
|
|
|
new_node.fpointer = None
|
|
|
|
return new_node
|
|
|
|
|
2019-09-26 07:09:10 +00:00
|
|
|
def safe_add_ancestors(self, node, ancestors):
|
|
|
|
# 如果没有祖先节点,那么添加该节点, 父节点是root node
|
|
|
|
if len(ancestors) == 0:
|
2019-08-21 12:27:21 +00:00
|
|
|
parent = self.root_node()
|
|
|
|
else:
|
2019-09-26 07:09:10 +00:00
|
|
|
parent = ancestors[0]
|
2019-08-21 12:27:21 +00:00
|
|
|
|
|
|
|
# 如果当前节点已再树中,则移动当前节点到父节点中
|
|
|
|
# 这个是由于 当前节点放到了二级节点中
|
2019-09-26 07:09:10 +00:00
|
|
|
if not self.contains(parent.identifier):
|
|
|
|
# logger.debug('Add parent: {}'.format(parent.identifier))
|
|
|
|
self.safe_add_ancestors(parent, ancestors[1:])
|
|
|
|
|
2019-08-21 12:27:21 +00:00
|
|
|
if self.contains(node.identifier):
|
2019-09-26 07:09:10 +00:00
|
|
|
# msg = 'Move node to parent: {} => {}'.format(
|
|
|
|
# node.identifier, parent.identifier
|
|
|
|
# )
|
|
|
|
# logger.debug(msg)
|
2019-08-21 12:27:21 +00:00
|
|
|
self.move_node(node.identifier, parent.identifier)
|
|
|
|
else:
|
2019-09-26 07:09:10 +00:00
|
|
|
# logger.debug('Add node: {}'.format(node.identifier))
|
2019-08-21 12:27:21 +00:00
|
|
|
self.add_node(node, parent)
|
2019-09-26 07:09:10 +00:00
|
|
|
|
2019-08-31 04:07:03 +00:00
|
|
|
#
|
|
|
|
# def __getstate__(self):
|
|
|
|
# self.mutex = None
|
|
|
|
# return self.__dict__
|
|
|
|
#
|
2019-09-26 07:09:10 +00:00
|
|
|
|
|
|
|
def __setstate__(self, state):
|
|
|
|
self.__dict__ = state
|
|
|
|
if '_invalid_assets' not in state:
|
|
|
|
self._invalid_assets = frozenset()
|
|
|
|
# self.mutex = threading.Lock()
|