perf: add Tree/TreeNode/TreeLeaf and AssetNodeTree/NodeTreeNode/AssetTreeLeaf, and TreeNodeSerializer

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

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

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)

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

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
+"""