From dc88dcdabdf32846c4d6374e40b1745228321862 Mon Sep 17 00:00:00 2001 From: Bai Date: Thu, 1 Jan 2026 21:26:03 +0800 Subject: [PATCH] perf: add Tree/TreeNode/TreeLeaf and AssetNodeTree/NodeTreeNode/AssetTreeLeaf, and TreeNodeSerializer --- apps/assets/tree/x/tree.py | 83 ++++++++ apps/assets/tree/x/tree_node.py | 68 +++++++ apps/assets/tree/x/tree_serializer.py | 55 +++++ apps/assets/tree/x1/node_tree.py | 128 ++++++++++++ apps/assets/tree/x1/tree.py | 283 ++++++++++++++++++++++++++ 5 files changed, 617 insertions(+) create mode 100644 apps/assets/tree/x/tree.py create mode 100644 apps/assets/tree/x/tree_node.py create mode 100644 apps/assets/tree/x/tree_serializer.py create mode 100644 apps/assets/tree/x1/node_tree.py create mode 100644 apps/assets/tree/x1/tree.py diff --git a/apps/assets/tree/x/tree.py b/apps/assets/tree/x/tree.py new file mode 100644 index 000000000..925c7c7a6 --- /dev/null +++ b/apps/assets/tree/x/tree.py @@ -0,0 +1,83 @@ +from django.db.models import QuerySet +from assets.models import Asset, Node, Platform, Zone +from labels.models import Label +from .tree_node import NodeAssetTreeNode + + +class Tree: + + def __init__(self): + self.root = None + self.tree_nodes = {} + + +class AssetsTreeMixin: + + asset_filter_fields = [ 'id', 'node_id', 'platform_id', 'zone_id', 'is_active', 'org_id' ] + asset_search_fields = [ 'name__icontains', 'address__icontains' ] + + def __init__(self): + super().__init__() + self.assets_queryset = None + self.tree_nodes_assets_amount_mapper = {} + self.tree_nodes_assets_mapper = [] + +class UsersTreeMixin: + + user_filter_fields = [ 'id', 'user', 'username', 'is_active' ] + user_search_fields = [ 'username__icontains', 'email__icontains' ] + + def __init__(self): + super().__init__() + self.users_queryset = None + self.tree_nodes_users_amount_mapper = {} + self.tree_nodes_users_mapper = [] + + +class NodeAssetTree(Tree, AssetsTreeMixin): + + tree_node_class = NodeAssetTreeNode + + node_filter_fields = [ 'id', 'key' ] + node_search_fields = [ 'value__icontains' ] + + def __init__(self): + super().__init__() + self.nodes: QuerySet[Node] = [] + +class PlatformAssetTree(Tree, AssetsTreeMixin): + + platform_filter_fields = [ 'id', 'category', 'type' ] + platform_search_fields = [ 'name__icontains', 'category__icontains', 'type__icontains' ] + + def __init__(self): + super().__init__() + self.platforms: QuerySet[Platform] = [] + + +# --- + +class ZoneAssetTree(Tree, AssetsTreeMixin): + + zone_filter_fields = [ 'id' ] + zone_search_fields = [ 'name__icontains' ] + + def __init__(self): + super().__init__() + self.zones: QuerySet[Zone] = [] + +class LabelTree(Tree): + + label_filter_fields = [ 'id' ] + label_search_fields = [ 'name__icontains', 'value__icontains' ] + + def __init__(self): + super().__init__() + self.labels: QuerySet[Label] = [] + +class LabelAssetTree(LabelTree, AssetsTreeMixin): + pass + + +class LabelUserTree(LabelTree, UsersTreeMixin): + pass diff --git a/apps/assets/tree/x/tree_node.py b/apps/assets/tree/x/tree_node.py new file mode 100644 index 000000000..36403efaa --- /dev/null +++ b/apps/assets/tree/x/tree_node.py @@ -0,0 +1,68 @@ + +class TreeNode: + def __init__(self): + self.key = None + self.name = None + self.parent = None + self.children = [] + + +class AssetsMixin: + fetch_asset_fields = [ + 'id', 'name', 'address', 'node_id', + 'platform_id', 'platform__category', 'platform__type', + 'is_active', 'comment', 'zone_id', 'org_id', + ] + + def __init__(self): + self.assets_amount = 0 + self.assets_amount_total = 0 + self.assets = [] + + +class UsersMixin: + + def __init__(self): + self.users_amount = 0 + self.users_amount_total = 0 + self.users = [] + + +class NodeAssetTreeNode(TreeNode, AssetsMixin): + def __init__(self): + super().__init__() + self.node = None + + +class PlatformAssetTreeNode(TreeNode, AssetsMixin): + + def __init__(self): + super().__init__() + self.type = 'category | type | platform' + self.platform_category = None + self.platform_type = None + self.platform = None + +# --- + + +class ZoneAssetTreeNode(TreeNode, AssetsMixin): + + def __init__(self): + super().__init__() + self.zone = None + + +class LabelTreeNode(TreeNode): + + def __init__(self): + super().__init__() + self.label = None + + +class LabelAssetTreeNode(LabelTreeNode, AssetsMixin): + pass + + +class LabelUserTreeNode(LabelTreeNode, UsersMixin): + pass diff --git a/apps/assets/tree/x/tree_serializer.py b/apps/assets/tree/x/tree_serializer.py new file mode 100644 index 000000000..2f0230579 --- /dev/null +++ b/apps/assets/tree/x/tree_serializer.py @@ -0,0 +1,55 @@ +from rest_framework import serializers + + +class AssetDataSerializer(serializers.Serializer): + platform_type = serializers.CharField() + org_name = serializers.CharField() + sftp = serializers.CharField() + platform = serializers.CharField() + name = serializers.CharField() + + +class NodeDataSerializer(serializers.Serializer): + id = serializers.CharField() + key = serializers.CharField() + value = serializers.CharField() + assets_amount = serializers.IntegerField() + assets_amount_total = serializers.IntegerField() + children_count_total = serializers.IntegerField() + + +META_DATA_SERIALIZER_MAP = { + 'node': NodeDataSerializer, + 'asset': AssetDataSerializer, +} + + +class TreeNodeMetaSerializer(serializers.Serializer): + type = serializers.CharField() + data = serializers.SerializerMethodField() + + def get_data(self, obj): + # 同时支持 dict / object + meta_type = getattr(obj, 'type', None) + raw_data = getattr(obj, 'data', None) + + serializer_cls = META_DATA_SERIALIZER_MAP.get(meta_type) + if not serializer_cls: + return raw_data + + return serializer_cls(raw_data).data + + +class TreeNodeSerializer(serializers.Serializer): + id = serializers.CharField() + name = serializers.CharField() + title = serializers.CharField() + pId = serializers.CharField() + isParent = serializers.BooleanField() + open = serializers.BooleanField() + meta = TreeNodeMetaSerializer() + + +# Example usage: +tree_nodes = [] +serializer = TreeNodeSerializer(tree_nodes, many=True) \ No newline at end of file diff --git a/apps/assets/tree/x1/node_tree.py b/apps/assets/tree/x1/node_tree.py new file mode 100644 index 000000000..e3c0b93a2 --- /dev/null +++ b/apps/assets/tree/x1/node_tree.py @@ -0,0 +1,128 @@ +from django.db.models import Q, Count + +from assets.models import Asset, Node, node +from .tree import Tree, TreeNode, TreeLeaf + + +class NodeTreeNode(TreeNode): + model_only_fields = [ + 'id', 'key', 'value' + ] + + def __init__(self, instance, **kwargs): + super().__init__(**kwargs) + self.instance = instance + + @property + def meta(self): + return { + 'type': 'node', + 'data': {} + } + + +class AssetTreeLeaf(TreeLeaf): + model_only_fields = [ + 'node_id', 'id', 'name', 'address', 'category', 'type' + ] + + def __init__(self, instance, **kwargs): + super().__init__(**kwargs) + self.instance = instance + + @property + def meta(self): + return { + 'type': 'asset', + 'data': {} + } + + +class AssetNodeTree(Tree): + node_model = Node + leaf_model = Asset + + def __init__(self, assets_scope_q: Q = None, org=None, **kwargs): + super().__init__(**kwargs) + self._assets_scope_q = assets_scope_q + self._org = org + self._node_id_key_mapper = {} + self._node_key_id_mapper = {} + self._node_id_node_mapper = {} + self._load_complete_nodes() + + def _load_complete_nodes(self): + nodes = Node.objects.filter(org_id=self._org).only(*NodeTreeNode.model_only_fields) + id_key_mapper = {} + key_id_mapper = {} + id_node_mapper = {} + for n in nodes: + nid = str(n.id) + id_key_mapper[nid] = n.key + key_id_mapper[n.key] = nid + id_node_mapper[nid] = n + self._node_id_key_mapper = id_key_mapper + self._node_key_id_mapper = key_id_mapper + self._node_id_node_mapper = id_node_mapper + + @property + def assets_queryset(self): + q = self._assets_scope_q or Q() + q &= Q(org_id=self._org.id) + if self.search_leaf_keyword: + k = self.search_leaf_keyword + q &= Q(name__icontains=k) | Q(address__icontains=k) + queryset = Asset.objects.filter(q) + return queryset + + def construct_nodes(self): + nid_assets_amount = self.assets_queryset.values('node_id').annotate( + assets_amount=Count('id') + ).values_list('node_id', 'assets_amount') + nid_assets_amount = { + str(nid): amount for nid, amount in nid_assets_amount + } + tree_nodes = [] + for nid, node in self._node_id_node_mapper.items(): + node: Node + assets_amount = nid_assets_amount.get(nid, 0) + key = self._node_id_key_mapper.get(nid) + assert isinstance(key, str), "TreeNode key should is str" + if key.isdigit(): + parent_key = None + else: + parent_key = ':'.join(key.split(':')[:-1]) if ':' in key else None + tree_node = NodeTreeNode( + instance=node, + key=key, + parent_key=parent_key, + name=node.name, + leaves_amount=assets_amount, + ) + tree_nodes.append(tree_node) + return tree_nodes + + def construct_leaves(self, node_key=None): + tree_leaves = [] + + node_id = None + if node_key: + node_id = self._node_key_id_mapper.get(node_key) + + q = Q() + if node_id: + q &= Q(node_id=node_id) + + assets = self.assets_queryset.filter(q).only(*AssetTreeLeaf.model_only_fields) + for asset in assets: + node_id = str(asset.node_id) + parent_key = self._node_id_key_mapper.get(node_id) + assert parent_key is not None, "TreeLeaf parent key should not be None" + tree_leaf = AssetTreeLeaf( + instance=asset, + key=None, + parent_key=parent_key, + name=asset.name, + ) + tree_leaves.append(tree_leaf) + return tree_leaves diff --git a/apps/assets/tree/x1/tree.py b/apps/assets/tree/x1/tree.py new file mode 100644 index 000000000..6e2cbfed9 --- /dev/null +++ b/apps/assets/tree/x1/tree.py @@ -0,0 +1,283 @@ +from xpack.plugins.interface import meta +from abc import abstractmethod +from assets.models import Asset, Node +from collections import defaultdict + +from common.utils import lazyproperty + +__all__ = [ 'TreeNode', 'TreeLeaf', 'Tree' ] + +class TreeNode: + + def __init__(self, key, parent_key, name, leaves_amount=0): + ## need to serialize fields + self.key = key # leaf key is None, parent_key + self.name = name + self.open = True + self.parent = None + self.parent_key = parent_key + self.is_leaf = False + self.children = [] + self.leaves = [] + self.leaves_amount = leaves_amount + + @property + def is_root(self): + return self.parent_key is None + + def add_child(self, child: 'TreeNode'): + child.parent = self + self.children.append(child) + + def add_leaf(self, leaf: 'TreeLeaf'): + leaf.parent = self + self.leaves.append(leaf) + + def matched(self, keyword): + if not keyword: + return True + keyword = str(keyword).strip().lower() + value = str(self.name).strip().lower() + return keyword in value + + def is_parent(self): + return len(self.children) > 0 or len(self.leaves) > 0 + + @property + def meta(self) -> dict: + ''' 元数据 + { + 'type': None, + 'data': {} + } + ''' + raise NotImplementedError + + @lazyproperty + def leaves_amount_total(self): + count = self.leaves_amount + for child in self.children: + child: TreeNode + count += child.leaves_amount_total + return count + + +class TreeLeaf(TreeNode): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.key = None + self.is_leaf = True + self.open = False + + +class Tree: + + def __init__(self, search_node_keyword=None, search_leaf_keyword=None, with_leaves=False): + self.root = None + self.tree = dict() + self.nodes = dict() + self.leaves = [] + + self.search_node_keyword = search_node_keyword + self.search_leaf_keyword = search_leaf_keyword + self.with_leaves = with_leaves + self.with_leaves_amount = True + + def build(self): + self.build_nodes() + if self.with_leaves: + self.build_leaves() + + def build_nodes(self): + nodes = self.construct_nodes() + self.add_nodes(nodes) + + def build_leaves(self, node_key=None): + leaves = self.construct_leaves(node_key=node_key) + self.add_leaves(leaves) + + @abstractmethod + def construct_nodes(self): + # Node.objects.all() + raise NotImplementedError + + @abstractmethod + def construct_leaves(self, node_key=None): + # Asset.objects.all() + raise NotImplementedError + + def expand_node(self, node_key, with_leaves=False): + node = self.get_node(node_key=node_key) + node: TreeNode + if with_leaves: + # 包含资产 + if not self.with_leaves: + # 树中没有资产，构建资产 + self.build_leaves(node_key=node_key) + return node.children + node.leaves + + def search_nodes(self, keyword): + keyword = self.search_node_keyword + matched_nodes = {} + for node in self.tree.values(): + node: TreeNode + if not node.matched(keyword): + continue + matched_nodes[node.key] = node + + to_remove_keys = set() + matched_keys = list(matched_nodes.keys()) + for mk in matched_keys: + descendants_keys = [k for k in matched_keys if k.startswith(f"{mk}:")] + to_remove_keys.update(descendants_keys) + + for ck in to_remove_keys: + matched_nodes.pop(ck, None) + top_nodes = matched_nodes.values() + return top_nodes + + def get_nodes(self): + if not self.search_node_keyword: + return list(self.tree.values()) + searched_nodes = self.search_nodes(keyword=self.search_node_keyword) + nodes = self.get_nodes_ancestors(searched_nodes) + return nodes + searched_nodes + + def get_nodes_ancestors(self, nodes: list[TreeNode]): + ancestors = set() + for node in nodes: + n_ancestors = self.get_node_ancestors(node) + ancestors.update(n_ancestors) + return list(ancestors) + + def get_node_ancestors(self, node: TreeNode): + if node.is_root: + return [] + ancestors = [] + parent = node.parent + while parent: + parent: TreeNode + ancestors.append(parent) + if parent.is_root: + break + parent = parent.parent + return ancestors + + def add_nodes(self, nodes: list[TreeNode]): + self.nodes = {n.key: n for n in nodes} + for node in nodes: + self.add_node(node) + + def add_node(self, node: TreeNode): + parent = self.nodes.get(node.parent_key) + if not parent: + self.root = node + self._add_node(node) + return + + parent: TreeNode + if self.exist(parent.key): + self.add_node(parent) + + self._add_node(node) + + def exist(self, node_key): + return node_key in self.tree + + def _add_node(self, node: TreeNode): + if node.is_root: + self.tree[node.key] = node + return + + parent = self.get_node(node_key=node.parent_key) + assert isinstance(parent, TreeNode) , f"Parent node {node.parent_key} not found for node {node.key}" + parent.add_child(node) + self.tree[node.key] = node + + def get_node(self, node_key=None): + return self.tree.get(node_key) + + def add_leaves(self, leaves): + for leaf in leaves: + self.add_leaf(leaf) + + def add_leaf(self, leaf: TreeLeaf): + parent = self.get_node(leaf.parent_key) + assert isinstance(parent, TreeNode) , f"Parent node {leaf.parent_key} not found for leaf {leaf.key}" + parent.add_leaf(leaf) + + +# 初始化树 +# 展开某个树节点 +# 搜索树中的节点 +# 搜索树中的叶子 + +""" + +树的节点，是根据叶子来限定范围的 + +不同的树，叶子类型不同，查询叶子的 queryset 也不同 + +叶子是从数据库中来的 +节点也是从数据库中来的 +在数据库模型中，本质上节点是叶子的一个属性 +我们获取树的出发点，首先是符合条件的叶子 + +初始化一颗完整树时，节点的范围是由叶子的范围决定的 + +只搜索叶子时，树需要重新初始化，因为叶子的范围变化了 +只搜索节点时，树不需要重新初始化，因为叶子的范围没有变化 +既搜索节点:又搜索叶子时，首先查询出符合条件的节点，然后根据符合条件的节点，查询出符合条件的叶子，然后根据叶子生成树 + +展开一个节点时，本质上展开的是一颗搜索树上的一个节点，返回节点的子节点和叶子 + +初始化一颗树有2种情况： +1. 初始化只包含节点的树 +2. 初始化既包含节点，又包含某些节点的叶子的树 + + +因为每一颗树的节点都必须包含它的叶子数量和叶子总数量 + +查询每个节点下叶子数量，必须用一条 SQL 语句完成，性能才能接受 + +因为一颗树中的所有叶子是不计其数的， + 所以在初始化一颗完整树时，不能获取出所有叶子，只能获取某个节点下的叶子 + 在初始化一颗搜索叶子的树时，虽然理论上应该要返回所有符合条件的叶子，但是为了避免叶子数量巨大，所以必须限制返回的叶子数量 + + +核心问题： +如何根据一些扁平化数据的简单父子关系，构建出一颗树中的完整父子关系 +如何定义叶子的查询范围 +如何定义节点的查询范围 + + +初始化一颗树时，首先是确定节点的范围，然后在节点范围的基础上，确定叶子的范围。 +一定是两者的交集，来构建树 +如：node_ids, Asset.objects.filter(node_id__in=node_ids).filter(id__in=[]) +其中查询叶子时，一定有节点下不包含叶子的情况，所以要区分是否包含叶子总数量为0的节点，即所谓空节点 + +# Tips: 查询每个节点下叶子总数量比较慢（所以才改成一个资产只能属于一个节点，如果属于多个节点时，就不是树的数据结构了，而是图的数据结构，复杂度会高很多） + +树初始化完之后，一定是一颗包含所有节点完整树 + +# ... +queryset 分三部分 + +完整树 +nodes = nodes_queryset +assets = nodes & assets_queryset + + +搜索树 +nodes = tree:nodes and search_nodes_queryset +assets = nodes and search_assets_queryset + +展开完整树的某个节点 +nodes = tree:nodes:children +assets = tree:assets:expand_nodes_queryset + +展开搜索树的某个节点 +nodes = search_tree:nodes:children +assets = search_tree:assets:expand_nodes_queryset +""" \ No newline at end of file