diff --git a/apps/assets/api/tree.py b/apps/assets/api/tree.py deleted file mode 100644 index 87b955edd..000000000 --- a/apps/assets/api/tree.py +++ /dev/null @@ -1,267 +0,0 @@ -# ~*~ coding: utf-8 ~*~ - -from django.db.models import Q, Count -from django.utils.translation import gettext_lazy as _ -from rest_framework.generics import get_object_or_404 -from rest_framework.response import Response - -from assets.locks import NodeAddChildrenLock -from common.exceptions import JMSException -from common.tree import TreeNodeSerializer -from common.utils import get_logger -from orgs.mixins import generics -from orgs.utils import current_org -from orgs.models import Organization -from .mixin import SerializeToTreeNodeMixin -from .. import serializers -from ..const import AllTypes -from ..models import Node, Platform, Asset -from assets.tree.asset_tree import AssetTree - - -logger = get_logger(__file__) - -__all__ = [ - 'NodeChildrenApi', - 'NodeChildrenAsTreeApi', - 'CategoryTreeApi', -] - - -class NodeChildrenApi(generics.ListCreateAPIView): - ''' 节点的增删改查 ''' - serializer_class = serializers.NodeSerializer - search_fields = ('value',) - - instance = None - is_initial = False - perm_model = Node - - def initial(self, request, *args, **kwargs): - super().initial(request, *args, **kwargs) - self.initial_org_root_node_if_need() - self.instance = self.get_object() - - def initial_org_root_node_if_need(self): - if current_org.is_root(): - return - Node.org_root() - - def perform_create(self, serializer): - data = serializer.validated_data - _id = data.get("id") - value = data.get("value") - if value: - children = self.instance.get_children() - if children.filter(value=value).exists(): - raise JMSException(_('The same level node name cannot be the same')) - else: - value = self.instance.get_next_child_preset_name() - with NodeAddChildrenLock(self.instance): - node = self.instance.create_child(value=value, _id=_id) - # 避免查询 full value - node._full_value = node.value - serializer.instance = node - - def get_object(self): - pk = self.kwargs.get('pk') or self.request.query_params.get('id') - key = self.request.query_params.get("key") - - if not pk and not key: - self.is_initial = True - if current_org.is_root(): - node = None - else: - node = Node.org_root() - return node - - if pk: - node = get_object_or_404(Node, pk=pk) - else: - node = get_object_or_404(Node, key=key) - return node - - -class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi): - ''' 节点子节点作为树返回， - [ - { - "id": "", - "name": "", - "pId": "", - "meta": "" - } - ] - ''' - - model = Node - - def list(self, request, *args, **kwargs): - with_assets = request.query_params.get('assets', '0') == '1' - search = request.query_params.get('search') - key = request.query_params.get('key') - - if not with_assets: - # 初始化资产树 - 不包含资产 - return self.init_asset_tree() - - if search: - # 初始化资产搜索树 - 包含资产 - return self.search_asset_tree_with_assets(search) - - if key: - # 展开资产树节点 - 包含资产 - return self.expand_asset_tree_node_with_assets(key) - - # 初始化资产树 - 包含资产 - return self.init_asset_tree_with_assets() - - def init_asset_tree(self): - ''' 初始化资产树 - 不包含资产 - 返回所有节点，前端本地展开和搜索 - 全局组织: 不展开节点 - 实体组织: 展开第1级节点 - ''' - if current_org.is_root(): - orgs = Organization.objects.all() - expand_level = 0 - else: - orgs = [current_org] - expand_level = 1 - nodes = [] - for org in orgs: - tree = AssetTree(org=org) - _nodes = tree.get_nodes() - nodes.extend(_nodes) - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level - ) - return Response(data=nodes) - - def init_asset_tree_with_assets(self): - ''' 初始化资产树 - 包含资产 - 全局组织: 返回第1级节点，不返回资产，不展开 - 实体组织: 返回第1级节点和第2级节点，返回第1级节点的资产，展开第1级节点 - ''' - if current_org.is_root(): - orgs = Organization.objects.all() - node_levels = [1] - with_assets_node_levels = None - expand_level = 0 - else: - orgs = [current_org] - node_levels = [1, 2] - with_assets_node_levels = [1] - expand_level = 1 - - nodes = [] - assets = [] - for org in orgs: - tree = AssetTree( - org=org, with_assets_node_levels=with_assets_node_levels - ) - _nodes = tree.get_nodes(levels=node_levels) - nodes.extend(_nodes) - _assets = tree.get_assets() - assets.extend(_assets) - - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level, - with_assets=True - ) - assets = self.serialize_assets(assets) - data = [*nodes, *assets] - return Response(data=data) - - def expand_asset_tree_node_with_assets(self, key): - ''' 展开资产树节点 - 包含资产 - 全局组织: 返回展开节点的直接孩子节点，返回展开节点的资产，不展开节点 - 实体组织: 同上 - ''' - node = get_object_or_404(Node, key=key) - org = node.org - with_assets_node_id = node.id - expand_level = 0 - - tree = AssetTree(with_assets_node_id=with_assets_node_id, org=org) - tree_node = tree.get_node(key=node.key) - if not tree_node: - return Response(data=[]) - _nodes = tree_node.children - nodes = self.serialize_nodes( - _nodes, with_asset_amount=True, expand_level=expand_level, - with_assets=True - ) - _assets = tree.get_assets() - assets = self.serialize_assets(_assets) - data = [*nodes, *assets] - return Response(data=data) - - def search_asset_tree_with_assets(self, search): - ''' 初始化资产搜索树 - 包含资产 - 不反回完整树，资产数量为0的节点不返回 - 全局组织: 返回所有节点，返回所有资产，展开所有节点，限制资产总数量 1000（n 个组织，每个组织分配1000/n个资产） - 实体组织：同上，限制资产总数量 1000 - ''' - # 展开所有节点 - expand_level = 10000 - with_assets_all = True - with_assets_limit = 1000 - full_tree = False - assets_q_object = Q(name__icontains=search) | Q(address__icontains=search) - if current_org.is_root(): - orgs = list(Organization.objects.all()) - with_assets_limit = max(100, with_assets_limit // max(1, orgs.count())) - else: - orgs = [current_org] - - nodes = [] - assets = [] - for org in orgs: - tree = AssetTree( - assets_q_object=assets_q_object, org=org, - with_assets_all=with_assets_all, - with_assets_limit=with_assets_limit, - full_tree=full_tree - ) - _nodes = tree.get_nodes() - nodes.extend(_nodes) - _assets = tree.get_assets() - assets.extend(_assets) - - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level, - with_assets=True - ) - assets = self.serialize_assets(assets) - data = [*nodes, *assets] - return Response(data=data) - -class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView): - serializer_class = TreeNodeSerializer - rbac_perms = { - 'GET': 'assets.view_asset', - 'list': 'assets.view_asset', - } - queryset = Node.objects.none() - - def get_assets(self): - key = self.request.query_params.get('key') - platform = Platform.objects.filter(id=key).first() - if not platform: - return [] - assets = Asset.objects.filter(platform=platform).prefetch_related('platform') - return self.serialize_assets(assets, key) - - def list(self, request, *args, **kwargs): - include_asset = self.request.query_params.get('assets', '0') == '1' - # 资源数量统计可选项 (asset, account) - count_resource = self.request.query_params.get('count_resource', 'asset') - - if not self.request.query_params.get('key'): - nodes = AllTypes.to_tree_nodes(include_asset, count_resource=count_resource) - elif include_asset: - nodes = self.get_assets() - else: - nodes = [] - return Response(data=nodes) diff --git a/apps/assets/api/tree/__init__.py b/apps/assets/api/tree/__init__.py new file mode 100644 index 000000000..53a0436f9 --- /dev/null +++ b/apps/assets/api/tree/__init__.py @@ -0,0 +1,3 @@ +from .tree import * +from .base import * +from .const import * \ No newline at end of file diff --git a/apps/assets/api/tree/base.py b/apps/assets/api/tree/base.py new file mode 100644 index 000000000..b46be46a9 --- /dev/null +++ b/apps/assets/api/tree/base.py @@ -0,0 +1,294 @@ + +from django.db.models import Q +from rest_framework import generics +from rest_framework.response import Response +from rest_framework.generics import get_object_or_404 + +from users.models import User +from common.utils import lazyproperty +from common.exceptions import APIException +from orgs.utils import current_org +from orgs.models import Organization +from rbac.permissions import RBACPermission +from assets.tree.asset_tree import AssetTree +from assets.models import Node +from .mixin import SerializeToTreeNodeMixin +from .const import RenderTreeType, RenderTreeTypeChoices + +class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView): + + permission_classes = (RBACPermission,) + + perm_model = Node + + render_tree_type: RenderTreeType + tree_user: User + + + def initial(self, request, *args, **kwargs): + super().initial(request, *args, **kwargs) + self.render_tree_type = self.initial_render_tree_type() + self.tree_user = self.get_tree_user() + + def initial_render_tree_type(self): + tree_type = self.get_query_value('tree_type') + if not tree_type: + with_assets = self.request.query_params.get('assets', '0') == '1' + tree_type = RenderTreeTypeChoices.asset if with_assets else RenderTreeTypeChoices.node + return RenderTreeType(tree_type) + + def get_tree_user(self) -> User: + raise NotImplementedError + + def list(self, request, *args, **kwargs): + asset_category = self.get_query_value('asset_category') + asset_type = self.get_query_value('asset_type') + with_asset_amount = True + + if self.render_tree_type.is_node_tree: + data = self.render_node_tree(asset_category, asset_type, with_asset_amount) + elif self.render_tree_type.is_asset_tree: + data = self.render_asset_tree(asset_category, asset_type, with_asset_amount) + else: + raise APIException(f'Invalid tree type: {self.query_tree_type}') + return Response(data=data) + + def render_node_tree(self, asset_category, asset_type, with_asset_amount): + data = self.init_node_tree(asset_category, asset_type, with_asset_amount) + return data + + def render_asset_tree(self, asset_category, asset_type, with_asset_amount): + expand_node_key = self.get_query_value('key') + search_node = self.get_query_value('search_node') + search_asset = self.get_query_value('search_asset') + data = self._render_asset_tree( + expand_node_key=expand_node_key, search_node=search_node, search_asset=search_asset, + asset_category=asset_category, asset_type=asset_type, + with_asset_amount=with_asset_amount + ) + return data + + def _render_asset_tree(self, expand_node_key=None, search_node=None, search_asset=None, + asset_category=None, asset_type=None, with_asset_amount=None): + + if not search_asset and self.render_tree_type.is_asset_tree: + search = self.get_query_value('search') or '' + if ':' not in search: + search_asset = search + + if expand_node_key: + data = self.expand_asset_tree_node( + expand_node_key, asset_category, asset_type, with_asset_amount + ) + elif search_node: + data = self.search_asset_tree_nodes( + search_node, asset_category, asset_type, with_asset_amount + ) + elif search_asset: + data = self.search_asset_tree_assets( + search_asset, asset_category, asset_type, with_asset_amount + ) + else: + data = self.init_asset_tree( + asset_category, asset_type, with_asset_amount + ) + return data + + def init_node_tree(self, asset_category, asset_type, with_asset_amount): + orgs = self.get_tree_user_orgs() + + if self.org_is_global: + expand_level = 0 + else: + expand_level = 1 + + nodes = [] + for org in orgs: + tree = self.get_org_asset_tree( + asset_category=asset_category, asset_type=asset_type, org=org + ) + _nodes = tree.get_nodes() + nodes.extend(_nodes) + + data = self.serialize_nodes( + nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + return data + + def init_asset_tree(self, asset_category, asset_type, with_asset_amount): + orgs = self.get_tree_user_orgs() + if self.org_is_global: + node_levels = [1] + with_assets_node_levels = None + expand_level = 0 + else: + node_levels = [1, 2] + with_assets_node_levels = [1] + expand_level = 1 + nodes = [] + assets = [] + for org in orgs: + tree: AssetTree = self.get_org_asset_tree( + asset_category=asset_category, asset_type=asset_type, org=org, + with_assets_node_levels=with_assets_node_levels + ) + _nodes = tree.get_nodes(levels=node_levels) + nodes.extend(_nodes) + _assets = tree.get_assets() + assets.extend(_assets) + + serialized_nodes = self.serialize_nodes( + nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + serialized_assets = self.serialize_assets(assets) + data = serialized_nodes + serialized_assets + return data + + def expand_asset_tree_node(self, node_key, asset_category, asset_type, with_asset_amount): + node = get_object_or_404(Node, key=node_key) + orgs = self.get_tree_user_orgs() + # 确保用户有权限展开该节点所在组织的树 + org = orgs.filter(id=node.org_id).first() + if not org: + raise APIException( + f'No permission to expand the node in this organization: {node.org_name}' + ) + + with_assets_node_id = str(node.id) + tree: AssetTree = self.get_org_asset_tree( + asset_category=asset_category, asset_type=asset_type, org=org, + with_assets_node_id=with_assets_node_id + ) + node_children = tree.get_node_children(node.key) + node_assets = tree.get_assets() + + # (展开节点)的孩子节点不展开 + expand_level = 0 + serialized_nodes = self.serialize_nodes( + node_children, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + serialized_assets = self.serialize_assets(node_assets) + data = serialized_nodes + serialized_assets + return data + + def search_asset_tree_nodes(self, search_node, asset_category, asset_type, with_asset_amount): + orgs = self.get_tree_user_orgs() + + matched_nodes = [] + matched_nodes_ancestors = [] + for org in orgs: + tree: AssetTree = self.get_org_asset_tree( + asset_category=asset_category, asset_type=asset_type, org=org + ) + # 如果匹配的节点中包含有父子关系的节点，只返回最上一级的父节点 + _matched_nodes = tree.search_nodes(search_node, only_top_level=True) + if not _matched_nodes: + continue + matched_nodes.extend(_matched_nodes) + _ancestors = tree.get_nodes_ancestors(_matched_nodes) + matched_nodes_ancestors.extend(_ancestors) + + if not matched_nodes: + return [] + + # 匹配的节点不展开 + expand_level = 0 + serialized_matched_nodes = self.serialize_nodes( + matched_nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + + # 匹配节点的祖先节点全部展开 + expand_level = 10000 + serialized_ancestors = self.serialize_nodes( + matched_nodes_ancestors, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + data = serialized_matched_nodes + serialized_ancestors + return data + + def search_asset_tree_assets(self, search_asset, asset_category, asset_type, with_asset_amount): + orgs = self.get_tree_user_orgs() + + # 搜索时，展开所有节点 + expand_level = 10000 + + # 资产树搜索范围 + assets_q_object = Q(name__icontains=search_asset) | Q(address__icontains=search_asset) + + # 限制每个组织搜索返回的资产数量 + with_assets_limit_max = 1000 + with_assets_limit_min = 100 + with_assets_limit = max(with_assets_limit_min, with_assets_limit_max // max(1, orgs.count())) + + # 搜索树只返回包含资产的节点 + full_tree = False + + nodes = [] + assets = [] + for org in orgs: + tree: AssetTree = self.get_org_asset_tree( + assets_q_object=assets_q_object, + asset_category=asset_category, asset_type=asset_type, org=org, + with_assets_all=True, with_assets_limit=with_assets_limit, + full_tree=full_tree + ) + _nodes = tree.get_nodes() + nodes.extend(_nodes) + _assets = tree.get_assets() + assets.extend(_assets) + serialized_nodes = self.serialize_nodes( + nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, + expand_level=expand_level + ) + serialized_assets = self.serialize_assets(assets) + data = serialized_nodes + serialized_assets + return data + + # tree 抽象方法 # + + def get_org_asset_tree(self, **kwargs) -> AssetTree: + raise NotImplementedError + + # tree 辅助方法 # + + @lazyproperty + def org_is_global(self): + return current_org.is_root() + + def get_tree_user_orgs(self): + ''' 重要: 获取用户有权限渲染树的组织列表 ''' + user = self.tree_user + if self.org_is_global: + # 如果是全局组织，返回用户所在的所有实体组织 + orgs = user.orgs.all() + else: + # 如果时实体组织，从用户所在的实体组织中返回该实体组织 + orgs = user.orgs.filter(id=current_org.id) + + if not orgs: + raise APIException('No organization available for rendering the tree') + return orgs + + # view 辅助方法 # + + def get_query_value(self, query_key): + query_value = self.request.query_params.get(query_key) + if not query_value: + query_value = self.get_query_value_from_search(query_key) + return query_value + + def get_query_value_from_search(self, query_key): + search = self.request.query_params.get('search', '') + if not search: + return None + + search_list = search.split() + for _search in search_list: + if f'{query_key}:' not in _search: + continue + query_value = _search.replace(f'{query_key}:', '').strip() + return query_value \ No newline at end of file diff --git a/apps/assets/api/tree/const.py b/apps/assets/api/tree/const.py new file mode 100644 index 000000000..b47ba73f9 --- /dev/null +++ b/apps/assets/api/tree/const.py @@ -0,0 +1,28 @@ +from django.db.models import TextChoices + + +__all__ = ['RenderTreeType', 'RenderTreeTypeChoices'] + + +class RenderTreeTypeChoices(TextChoices): + node = 'node', 'Node' + asset = 'asset', 'Asset' + + +class RenderTreeType: + + def __init__(self, _type): + if _type not in RenderTreeTypeChoices.values: + raise ValueError(f'Invalid tree type: {_type}') + self._type: RenderTreeTypeChoices = _type + + @property + def is_asset_tree(self): + return self._type == RenderTreeTypeChoices.asset + + @property + def is_node_tree(self): + return self._type == RenderTreeTypeChoices.node + + def __str__(self): + return self._type.value diff --git a/apps/assets/api/tree/mixin.py b/apps/assets/api/tree/mixin.py new file mode 100644 index 000000000..adcd74ac2 --- /dev/null +++ b/apps/assets/api/tree/mixin.py @@ -0,0 +1,125 @@ +from typing import List + +from rest_framework.request import Request + +from assets.models import Platform, Protocol, MyAsset +from common.utils import lazyproperty, timeit +from assets.tree.asset_tree import AssetTreeNode, AssetTreeNodeAsset + +from .const import RenderTreeType + + +__all__ = ['SerializeToTreeNodeMixin'] + + +class SerializeToTreeNodeMixin: + request: Request + + @timeit + def serialize_nodes(self, nodes: List[AssetTreeNode], tree_type: RenderTreeType, + with_asset_amount=False, expand_level=1): + if not nodes: + return [] + + def _name(node: AssetTreeNode): + v = node.value + if not with_asset_amount: + return v + v = f'{v} ({node.assets_amount_total})' + return v + + def is_parent(node: AssetTreeNode): + if tree_type.is_asset_tree: + return node.assets_amount > 0 or not node.is_leaf + elif tree_type.is_node_tree: + return not node.is_leaf + else: + return True + + data = [ + { + 'id': node.key, + 'name': _name(node), + 'title': _name(node), + 'pId': node.parent_key, + 'isParent': is_parent(node), + 'open': node.level <= expand_level, + 'meta': { + 'type': 'node', + 'data': { + "id": node.id, + "key": node.key, + "value": node.value, + "assets_amount": node.assets_amount, + "assets_amount_total": node.assets_amount_total, + "children_count_total": node.children_count_total, + }, + } + } + for node in nodes + ] + return data + + @lazyproperty + def support_types(self): + from assets.const import AllTypes + return AllTypes.get_types_values(exclude_custom=True) + + def get_icon(self, asset): + if asset.category == 'device': + return 'switch' + if asset.type in self.support_types: + return asset.type + else: + return 'file' + + @timeit + def serialize_assets(self, assets, node_key=None, get_pid=None): + if not assets: + return [] + + if not get_pid and not node_key: + get_pid = lambda asset, platform: getattr(asset, 'parent_key', '') + + sftp_asset_ids = Protocol.objects.filter(name='sftp') \ + .values_list('asset_id', flat=True) + sftp_asset_ids = set(sftp_asset_ids) + platform_map = {p.id: p for p in Platform.objects.all()} + + data = [] + root_assets_count = 0 + MyAsset.set_asset_custom_value(assets, self.request.user) + for asset in assets: + asset: AssetTreeNodeAsset + platform = platform_map.get(asset.platform_id) + if not platform: + continue + pid = node_key or get_pid(asset, platform) + if not pid: + continue + # 根节点最多显示 1000 个资产 + if pid.isdigit(): + if root_assets_count > 1000: + continue + root_assets_count += 1 + data.append({ + 'id': str(asset.id), + 'name': asset.name, + 'title': f'{asset.address}\n{asset.comment}'.strip(), + 'pId': pid, + 'isParent': False, + 'open': False, + 'iconSkin': self.get_icon(platform), + 'chkDisabled': not asset.is_active, + 'meta': { + 'type': 'asset', + 'data': { + 'platform_type': platform.type, + 'org_name': asset.org_name, + 'sftp': asset.id in sftp_asset_ids, + 'name': asset.name, + 'address': asset.address + }, + } + }) + return data \ No newline at end of file diff --git a/apps/assets/api/tree/tree.py b/apps/assets/api/tree/tree.py new file mode 100644 index 000000000..3fca6e6f0 --- /dev/null +++ b/apps/assets/api/tree/tree.py @@ -0,0 +1,135 @@ +# ~*~ coding: utf-8 ~*~ + +from django.utils.translation import gettext_lazy as _ +from rest_framework.generics import get_object_or_404 +from rest_framework.response import Response + +from assets.locks import NodeAddChildrenLock +from common.exceptions import JMSException +from common.tree import TreeNodeSerializer +from common.utils import get_logger +from orgs.mixins import generics +from orgs.utils import current_org, tmp_to_org +from ..mixin import SerializeToTreeNodeMixin +from ... import serializers +from ...const import AllTypes +from ...models import Node, Platform, Asset +from assets.tree.asset_tree import AssetTree +from .base import AbstractAssetTreeAPI + + +logger = get_logger(__file__) + +__all__ = [ + 'NodeChildrenApi', + 'AssetTreeAPI', + 'CategoryTreeApi', +] + + +class NodeChildrenApi(generics.ListCreateAPIView): + ''' 节点的增删改查 ''' + serializer_class = serializers.NodeSerializer + search_fields = ('value',) + + instance = None + is_initial = False + perm_model = Node + + def initial(self, request, *args, **kwargs): + super().initial(request, *args, **kwargs) + self.initial_org_root_node_if_need() + self.instance = self.get_object() + + def initial_org_root_node_if_need(self): + if current_org.is_root(): + return + Node.org_root() + + def perform_create(self, serializer): + data = serializer.validated_data + _id = data.get("id") + value = data.get("value") + if value: + children = self.instance.get_children() + if children.filter(value=value).exists(): + raise JMSException(_('The same level node name cannot be the same')) + else: + value = self.instance.get_next_child_preset_name() + with NodeAddChildrenLock(self.instance): + node = self.instance.create_child(value=value, _id=_id) + # 避免查询 full value + node._full_value = node.value + serializer.instance = node + + def get_object(self): + pk = self.kwargs.get('pk') or self.request.query_params.get('id') + key = self.request.query_params.get("key") + + if not pk and not key: + self.is_initial = True + if current_org.is_root(): + node = None + else: + node = Node.org_root() + return node + + if pk: + node = get_object_or_404(Node, pk=pk) + else: + node = get_object_or_404(Node, key=key) + return node + + +class AssetTreeAPI(AbstractAssetTreeAPI): + + def initial(self, request, *args, **kwargs): + super().initial(request, *args, **kwargs) + self.initial_org_root_node_if_need() + + def initial_org_root_node_if_need(self): + if current_org.is_root(): + orgs = self.request.user.orgs.all() + else: + orgs = [current_org] + + for org in orgs: + with tmp_to_org(org): + Node.org_root() + + def get_tree_user(self): + return self.request.user + + def get_org_asset_tree(self, **kwargs) -> AssetTree: + tree = AssetTree(**kwargs) + return tree + + +class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView): + serializer_class = TreeNodeSerializer + rbac_perms = { + 'GET': 'assets.view_asset', + 'list': 'assets.view_asset', + } + queryset = Node.objects.none() + + def get_assets(self): + key = self.request.query_params.get('key') + platform = Platform.objects.filter(id=key).first() + if not platform: + return [] + assets = Asset.objects.filter(platform=platform).prefetch_related('platform') + return self.serialize_assets(assets, key) + + def list(self, request, *args, **kwargs): + include_asset = self.request.query_params.get('assets', '0') == '1' + # 资源数量统计可选项 (asset, account) + count_resource = self.request.query_params.get('count_resource', 'asset') + + if not self.request.query_params.get('key'): + nodes = AllTypes.to_tree_nodes(include_asset, count_resource=count_resource) + elif include_asset: + nodes = self.get_assets() + else: + nodes = [] + return Response(data=nodes) diff --git a/apps/assets/tree/asset_tree.py b/apps/assets/tree/asset_tree.py index e15f8902b..4b4402627 100644 --- a/apps/assets/tree/asset_tree.py +++ b/apps/assets/tree/asset_tree.py @@ -19,7 +19,8 @@ __all__ = ['AssetTree', 'AssetTreeNode'] class AssetTreeNodeAsset: model_values = [ - 'id', 'name', 'node_id', 'platform_id', 'address', 'is_active', 'comment', 'org_id' + 'id', 'name', 'node_id', 'platform_id', 'address', + 'is_active', 'comment', 'org_id' ] def __init__(self, parent_key, **kwargs): @@ -39,18 +40,19 @@ class AssetTreeNodeAsset: class AssetTreeNode(TreeNode): - def __init__(self, _id, key, value, assets_amount=0): + def __init__(self, _id, key, value, assets_amount=0, assets_amount_total=0): super().__init__(_id, key, value) self.assets_amount = assets_amount - self.assets_amount_total = 0 + self.assets_amount_total = assets_amount_total self.assets: list[AssetTreeNodeAsset] = [] - def init_assets(self, assets): - if not assets: + def init_assets(self, assets_attrs): + if not assets_attrs: return - for asset in assets: - asset['parent_key'] = self.key - self.assets.append(AssetTreeNodeAsset(**asset)) + for asset_attrs in assets_attrs: + asset_attrs['parent_key'] = self.key + asset = AssetTreeNodeAsset(**asset_attrs) + self.assets.append(asset) return self.assets def get_assets(self): diff --git a/apps/assets/tree/tree.py b/apps/assets/tree/tree.py index 93ea37b5c..ac952f14c 100644 --- a/apps/assets/tree/tree.py +++ b/apps/assets/tree/tree.py @@ -17,6 +17,13 @@ class TreeNode(object): self.parent = None self.children_count_total = 0 + def match(self, keyword): + if not keyword: + return True + keyword = str(keyword).strip().lower() + node_value = str(self.value).strip().lower() + return keyword in node_value + @lazyproperty def parent_key(self): if self.is_root: @@ -176,12 +183,14 @@ class Tree(object): nodes = {} for node in self.nodes.values(): node: TreeNode - node_value = str(node.value).strip().lower() - if keyword in node_value: - nodes[node.key] = node + if not node.match(keyword): + continue + nodes[node.key] = node + if not only_top_level: return list(nodes.values()) + # 如果匹配的节点中包含有父子关系的节点，只返回最上一级的父节点 # TODO: 优化性能 node_keys = list(nodes.keys()) children_keys = [] diff --git a/apps/assets/urls/api_urls.py b/apps/assets/urls/api_urls.py index 7527be88e..46796d504 100644 --- a/apps/assets/urls/api_urls.py +++ b/apps/assets/urls/api_urls.py @@ -41,7 +41,7 @@ urlpatterns = [ api.AssetPermUserGroupPermissionsListApi.as_view(), name='asset-perm-user-group-permission-list'), path('nodes/category/tree/', api.CategoryTreeApi.as_view(), name='asset-category-tree'), - path('nodes/children/tree/', api.NodeChildrenAsTreeApi.as_view(), name='node-children-tree'), + path('nodes/children/tree/', api.AssetTreeAPI.as_view(), name='node-children-tree'), path('nodes//children/', api.NodeChildrenApi.as_view(), name='node-children'), path('nodes/children/', api.NodeChildrenApi.as_view(), name='node-children-2'), path('nodes//children/add/', api.NodeAddChildrenApi.as_view(), name='node-add-children'), diff --git a/apps/perms/api/user_permission/assets.py b/apps/perms/api/user_permission/assets.py index 8695b3e2c..c7cfb21ec 100644 --- a/apps/perms/api/user_permission/assets.py +++ b/apps/perms/api/user_permission/assets.py @@ -70,6 +70,7 @@ class BaseUserPermedAssetsApi(SelfOrPKUserMixin, ExtraFilterFieldsMixin, ListAPI class UserAllPermedAssetsApi(BaseUserPermedAssetsApi): + # TODO: support filter by asset_category and asset_type def get_assets(self): if self.user.is_superuser and self.request.query_params.get('id'): @@ -78,6 +79,7 @@ class UserAllPermedAssetsApi(BaseUserPermedAssetsApi): node_id = self.request.query_params.get('node_id') if node_id == UserPermAssetTreeNode.SpecialKey.FAVORITE: + # TODO: Support asset_category, asset_type return UserPermUtil.get_favorite_assets(user=self.user) if node_id == UserPermAssetTreeNode.SpecialKey.UNGROUPED: diff --git a/apps/perms/api/user_permission/tree.py b/apps/perms/api/user_permission/tree.py index f78814e24..9211fd0f3 100644 --- a/apps/perms/api/user_permission/tree.py +++ b/apps/perms/api/user_permission/tree.py @@ -2,14 +2,9 @@ from django.db.models import Q from django.utils.translation import gettext_lazy as _ from django.conf import settings -from rest_framework.generics import ListAPIView -from rest_framework.response import Response -from rest_framework.generics import get_object_or_404 from common.utils import get_logger, timeit -from orgs.utils import current_org -from assets.api import SerializeToTreeNodeMixin -from assets.models import Node, FavoriteAsset, Asset +from assets.api.tree import AbstractAssetTreeAPI from assets.tree.asset_tree import AssetTreeNodeAsset from perms.tree import UserPermAssetTree, UserPermAssetTreeNode from perms.utils.utils import UserPermUtil @@ -19,361 +14,128 @@ from .mixin import SelfOrPKUserMixin logger = get_logger(__name__) __all__ = [ - 'UserPermNodeChildrenAsTreeApi' + 'UserPermedAssetTreeAPI' ] -class RenderTreeType: - NODE = 'node' - ASSET = 'asset' +class UserPermedAssetTreeAPI(SelfOrPKUserMixin, AbstractAssetTreeAPI): - def __init__(self, tp): - if tp not in [self.NODE, self.ASSET]: - raise ValueError(f'Invalid tree type: {tp}') - self._type = tp + def get_tree_user(self): + return self.user + + def get_org_asset_tree(self, **kwargs) -> UserPermAssetTree: + return self.get_user_org_asset_tree(**kwargs) + + def get_user_org_asset_tree(self, **kwargs) -> UserPermAssetTree: + tree = UserPermAssetTree(user=self.user, **kwargs) + return tree - def is_asset_tree(self): - return self._type == self.ASSET - - def is_node_tree(self): - return self._type == self.NODE - - def __str__(self): - return self._type - - -class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, ListAPIView): - - def get_query_value(self, query_key): - query_value = self.request.query_params.get(query_key) - if query_value: - return query_value - - search = self.request.query_params.get('search', '') - if not search: - return None - - search_list = search.split() - for _search in search_list: - if f'{query_key}:' not in _search: - continue - query_value = _search.replace(f'{query_key}:', '').strip() - break - return query_value - - @property - def render_tree_type(self): - tree_type = self.get_query_value('tree_type') - - if tree_type: - return RenderTreeType(tree_type) - - if self.request.query_params.get('assets', '0') == '1': - tree_type = RenderTreeType.ASSET - else: - tree_type = RenderTreeType.NODE - - return RenderTreeType(tree_type) - - @timeit - def list(self, request, *args, **kwargs): - # node / asset - expand_node_key = self.get_query_value('key') - search_node = self.get_query_value('search_node') - search_asset = self.get_query_value('search_asset') - search = self.get_query_value('search') - asset_category = self.get_query_value('asset_category') - asset_type = self.get_query_value('asset_type') - - if self.render_tree_type.is_asset_tree(): - if expand_node_key: - return self.expand_user_perm_asset_tree_node(expand_node_key, asset_category, asset_type) - elif search_node: - # search nodes - return self.search_user_perm_node_tree(search_node, asset_category, asset_type) - elif search_asset or search: - search_asset = search_asset or search - # search assets - return self.search_user_perm_asset_tree(search_asset, asset_category, asset_type) - else: - return self.init_user_perm_asset_tree(asset_category, asset_type) - else: # if self.render_tree_type.is_node_tree(): - if expand_node_key: - raise NotImplementedError(_('Expanding node in node tree is not supported yet')) - if search_node: - raise NotImplementedError(_('Searching node in node tree is not supported yet')) - if search_asset: - raise NotImplementedError(_('Searching asset in node tree is not supported yet')) - else: - return self.init_user_perm_node_tree(asset_category, asset_type) + def _render_asset_tree(self, **kwargs): + data = super()._render_asset_tree(**kwargs) + expand_node_key = kwargs.pop('expand_node_key', None) + if expand_node_key: + # 特殊节点不需要展开，资产树中特殊节点下的资产已经随着节点一起返回 + return data - - def init_user_perm_node_tree(self, asset_category=None, asset_type=None): - ''' 初始化用户权限树 - 不包含资产 - 前端搜索 - - 全局组织: 返回所有节点，不返回资产，不展开节点 - 实体组织：返回所有节点，不返回资产，展开第1级节点 - - 返回收藏节点 - 返回未分组节点 (如果需要) - ''' - if current_org.is_root(): - orgs = self.user.orgs.all() - expand_level = 0 - else: - orgs = self.user.orgs.filter(id=current_org.id) - expand_level = 1 - - if not orgs.exists(): - return Response(data=[]) - - nodes = [] - for org in orgs: - tree = UserPermAssetTree( - user=self.user, asset_category=asset_category, - asset_type=asset_type, org=org - ) - _nodes = tree.get_nodes() - nodes.extend(_nodes) - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level, - tree_type=str(self.render_tree_type) - ) - data = nodes - data = self.add_favorites_and_ungrouped_node(data, with_assets=False) - return Response(data=data) - - def search_user_perm_node_tree(self, search, asset_category=None, asset_type=None): - ''' 搜索用户授权树 - 不包含资产 - 全局组织: 返回所有匹配节点以及祖先节点，不返回匹配节点的子孙节点，不返回资产，展开所有祖先节点，不展开匹配节点 - 实体组织: 同上 - ''' - if current_org.is_root(): - orgs = self.user.orgs.all() - else: - orgs = self.user.orgs.filter(id=current_org.id) - - search_nodes = [] - nodes_ancestors = [] - for org in orgs: - tree = UserPermAssetTree( - user=self.user, asset_category=asset_category, - asset_type=asset_type, org=org - ) - _search_nodes = tree.search_nodes(search, only_top_level=True) - # tree.remove_nodes_descendants(_search_nodes) - _nodes_ancestors = tree.get_nodes_ancestors(_search_nodes) - search_nodes.extend(_search_nodes) - nodes_ancestors.extend(_nodes_ancestors) - - # 不展开搜索节点 - expand_level = 0 - # 如果有资产，则允许展开 is_parent=True 的节点 - with_assets = True - serialized_search_nodes = self.serialize_nodes( - search_nodes, with_asset_amount=True, expand_level=expand_level, - tree_type=str(self.render_tree_type) - ) - # 展开所有祖先节点 - expand_level = 10000 - serialized_nodes_ancestors = self.serialize_nodes( - nodes_ancestors, with_asset_amount=True, expand_level=expand_level - ) - data = [*serialized_nodes_ancestors, *serialized_search_nodes] - return Response(data=data) - - def init_user_perm_asset_tree(self, asset_category=None, asset_type=None): - ''' 初始化用户权限资产树 - 包含资产 - 全局组织: 返回第1级节点，不返回资产，不展开节点 - 实体组织：返回第1级和第2级节点，返回第1级节点的资产，展开第1级节点 - - 返回收藏节点和资产 - 返回未分组节点和资产 (如果需要) - ''' - - if current_org.is_root(): - orgs = self.user.orgs.all() - nodes_level = [1] - with_assets_node_levels = None - expand_level = 0 - else: - orgs = self.user.orgs.filter(id=current_org.id) - nodes_level = [1, 2] - with_assets_node_levels = [1] - expand_level = 1 - - if not orgs.exists(): - return Response(data=[]) - - nodes = [] - assets = [] - - for org in orgs: - tree = UserPermAssetTree( - user=self.user, - asset_category=asset_category, asset_type=asset_type, org=org, - with_assets_node_levels=with_assets_node_levels - ) - _nodes = tree.get_nodes(levels=nodes_level) - nodes.extend(_nodes) - _assets = tree.get_assets() - assets.extend(_assets) - - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level, - tree_type=str(self.render_tree_type) - ) - assets = self.serialize_assets(assets) - data = [*nodes, *assets] - data = self.add_favorites_and_ungrouped_node(data, with_assets=True) - return Response(data=data) - - def expand_user_perm_asset_tree_node(self, node_key, asset_category=None, asset_type=None): - ''' 展开用户权限资产树节点 - 包含资产 - 全局组织: 返回展开节点的直接孩子节点，返回展开节点的资产，不展开其他节点 - 实体组织: 同上 - ''' - expand_level = 0 - node = get_object_or_404(Node, key=node_key) - org = self.user.orgs.filter(id=node.org_id).first() - if not org: - return Response(data=[]) - - tree = UserPermAssetTree( - user=self.user, - asset_category=asset_category, asset_type=asset_type, - org=node.org, with_assets_node_id=node.id - ) - tree_node = tree.get_node(node.key) - if not tree_node: - return Response(data=[]) - - _nodes = tree_node.children - nodes = self.serialize_nodes( - _nodes, with_asset_amount=True, expand_level=expand_level, - tree_type=str(self.render_tree_type) - ) - _assets = tree.get_assets() - assets = self.serialize_assets(_assets) - data = [*nodes, *assets] - return Response(data=data) - - def search_user_perm_asset_tree(self, search, asset_category=None, asset_type=None): - ''' 初始化用户权限资产搜索树 - 包含资产 - 全局组织: 返回所有节点，返回所有资产，展开所有节点，搜索资产 (最大 1000， n 个组织，每个组织分配1000/n个资产) - 实体组织: 同上，最大资产数 1000 - ''' - expand_level = 10000 - with_assets_all = True - with_assets_limit = 1000 - if current_org.is_root(): - orgs = self.user.orgs.all() - with_assets_limit = max(100, with_assets_limit // max(1, orgs.count())) - else: - orgs = self.user.orgs.filter(id=current_org.id) - - if not orgs.exists(): - return Response(data=[]) - - assets_q_object = Q(name__icontains=search) | Q(address__icontains=search) - nodes = [] - assets = [] - for org in orgs: - tree = UserPermAssetTree( - user=self.user, assets_q_object=assets_q_object, - org=org, asset_category=asset_category, asset_type=asset_type, - with_assets_all=with_assets_all, - with_assets_limit=with_assets_limit - ) - _nodes = tree.get_nodes() - nodes.extend(_nodes) - _assets = tree.get_assets() - assets.extend(_assets) - - nodes = self.serialize_nodes( - nodes, with_asset_amount=True, expand_level=expand_level, - tree_type=str(self.render_tree_type) - ) - assets = self.serialize_assets(assets) - data = [*nodes, *assets] - return Response(data=data) - - def add_favorites_and_ungrouped_node(self, data: list, with_assets=False): - # 未分组节点和资产 - u_node, u_assets = self.get_ungrouped_node_if_need() - if u_node: - data.insert(0, u_node) - data.extend(u_assets) - - # 收藏节点和资产 - f_node, f_assets = self.get_favorite_node() - if f_node: - data.insert(0, f_node) - data.extend(f_assets) + special_nodes = self.get_special_nodes(with_assets=True, expand_level=0, **kwargs) + data = special_nodes + data return data - def get_favorite_node(self, with_assets=False): - assets = UserPermUtil.get_favorite_assets(self.user) - assets_amount = assets.count() - node = UserPermAssetTreeNode( - _id=UserPermAssetTreeNode.SpecialKey.FAVORITE.value, - key=UserPermAssetTreeNode.SpecialKey.FAVORITE.value, - value=UserPermAssetTreeNode.SpecialKey.FAVORITE.label, - assets_amount=assets_amount + def render_node_tree(self, asset_category, asset_type, with_asset_amount): + data = super().render_node_tree(asset_category, asset_type, with_asset_amount) + special_nodes = self.get_special_nodes( + with_assets=False, expand_level=0, asset_category=asset_category, asset_type=asset_type, + with_asset_amount=with_asset_amount ) - node.assets_amount_total = assets_amount - nodes = self.serialize_nodes( - [node], with_asset_amount=True, expand_level=0, - tree_type=str(self.render_tree_type) + data = special_nodes + data + return data + + def get_special_nodes(self, search_asset=None, search_node=None, + asset_category=None, asset_type=None, + with_assets=False, expand_level=0, with_asset_amount=False): + f_node, f_assets = self.get_favorite_node( + search_asset=search_asset, search_node=search_node, + asset_category=asset_category, asset_type=asset_type, with_assets=with_assets + ) + u_node, u_assets = self.get_ungrouped_node_if_need( + search_asset=search_asset, search_node=search_node, + asset_category=asset_category, asset_type=asset_type, with_assets=with_assets ) - serialized_node = nodes[0] if nodes else None - if not with_assets: - return serialized_node, [] + nodes = [] + if f_node: + nodes.append(f_node) + if u_node: + nodes.append(u_node) - if assets_amount == 0: - return serialized_node, [] + if not nodes: + return [] - assets = list(assets.values(*AssetTreeNodeAsset.model_values)) - assets = node.init_assets(assets) - serialized_assets = self.serialize_assets(assets) - return serialized_node, serialized_assets + if search_asset: + expand_level = 1 + + serialized_nodes = self.serialize_nodes( + nodes, tree_type=self.render_tree_type, + with_asset_amount=with_asset_amount, expand_level=expand_level, + ) + if with_assets: + assets = f_assets + u_assets + serialized_assets = self.serialize_assets(assets) + else: + serialized_assets = [] + + data = serialized_nodes + serialized_assets + return data - def get_ungrouped_node_if_need(self, with_assets=False): + def get_favorite_node(self, search_asset=None, search_node=None, asset_category=None, + asset_type=None, with_assets=False): + assets = UserPermUtil.get_favorite_assets( + user=self.user, search_asset=search_asset, + asset_category=asset_category, asset_type=asset_type + ) + assets_amount = assets.count() + f_node = UserPermAssetTreeNode.favorite( + assets_amount=assets_amount, assets_amount_total=assets_amount + ) + if not f_node.match(search_node): + return None, [] + + if assets_amount == 0: + return f_node, [] + + if with_assets: + assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values)) + assets = f_node.init_assets(assets_attrs) + else: + assets = [] + return f_node, assets + + def get_ungrouped_node_if_need(self, search_asset=None, search_node=None, asset_category=None, + asset_type=None, with_assets=False): if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: return None, [] - - if current_org.is_root(): + if self.org_is_global: + # 全局组织不返回未分组节点 return None, [] - - org = self.user.orgs.filter(id=current_org.id).first() - if not org: - return None, [] - - _util = UserPermUtil(user=self.user, org=org) - assets = _util.get_ungrouped_assets() + org = self.get_tree_user_orgs().first() + util = UserPermUtil(user=self.user, org=org) + assets = util.get_ungrouped_assets( + search_asset=search_asset, + asset_category=asset_category, asset_type=asset_type + ) assets_amount = assets.count() - node = UserPermAssetTreeNode( - _id=UserPermAssetTreeNode.SpecialKey.UNGROUPED.value, - key=UserPermAssetTreeNode.SpecialKey.UNGROUPED.value, - value=UserPermAssetTreeNode.SpecialKey.UNGROUPED.label, - assets_amount=assets_amount + u_node = UserPermAssetTreeNode.ungrouped( + assets_amount=assets_amount, assets_amount_total=assets_amount ) - node.assets_amount_total = assets_amount - nodes = self.serialize_nodes( - [node], with_asset_amount=True, expand_level=0, - tree_type=str(self.render_tree_type) - ) - serialized_node = nodes[0] if nodes else None - - if not with_assets: - return serialized_node, [] - + if not u_node.match(search_node): + return None, [] + if assets_amount == 0: - return serialized_node, [] + return u_node, [] - assets = assets.values(*AssetTreeNodeAsset.model_values) - assets = node.init_assets(list(assets)) - serialized_assets = self.serialize_assets(assets) - return serialized_node, serialized_assets \ No newline at end of file + if with_assets: + assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values)) + assets = u_node.init_assets(assets_attrs) + else: + assets = [] + return u_node, assets diff --git a/apps/perms/tree.py b/apps/perms/tree.py index be549af4e..fb7a71e2d 100644 --- a/apps/perms/tree.py +++ b/apps/perms/tree.py @@ -18,28 +18,48 @@ logger = get_logger(__name__) class UserPermAssetTreeNode(AssetTreeNode): - class Type: - # Neither a permission node nor a node with direct permission assets - BRIDGE = 'bridge' - # Node with direct permission - DN = 'dn' - # Node with only direct permission assets - DA = 'da' - + class Type(TextChoices): + DN = 'direct_node', 'Direct Node' + DA = 'direct_asset', 'Direct Asset' + BRIDGE = 'bridge', 'Bridge' + class SpecialKey(TextChoices): FAVORITE = 'favorite', _('Favorite') UNGROUPED = 'ungrouped', _('Ungrouped') - - def __init__(self, tp=None, **kwargs): - self.type = tp or self.Type.BRIDGE + + def __init__(self, tp, **kwargs): + self.tp = tp super().__init__(**kwargs) + @classmethod + def favorite(cls, **kwargs): + node = cls( + tp=cls.Type.BRIDGE, + _id=cls.SpecialKey.FAVORITE.value, + key=cls.SpecialKey.FAVORITE.value, + value=cls.SpecialKey.FAVORITE.label, + **kwargs + ) + return node + + @classmethod + def ungrouped(cls, **kwargs): + node = cls( + tp=cls.Type.BRIDGE, + _id=cls.SpecialKey.UNGROUPED.value, + key=cls.SpecialKey.UNGROUPED.value, + value=cls.SpecialKey.UNGROUPED.label, + **kwargs + ) + return node + def as_dict(self, simple=True): - data = super().as_dict(simple=simple) + data = super().as_dict(simple) data.update({ - 'type': self.type, + 'type': self.tp, }) return data + class UserPermAssetTree(AssetTree): diff --git a/apps/perms/urls/user_permission.py b/apps/perms/urls/user_permission.py index 422422520..229b4ae7e 100644 --- a/apps/perms/urls/user_permission.py +++ b/apps/perms/urls/user_permission.py @@ -5,7 +5,7 @@ from .. import api user_permission_urlpatterns = [ path('/assets//', api.UserPermedAssetRetrieveApi.as_view(), name='user-permed-asset'), path('/assets/', api.UserAllPermedAssetsApi.as_view(), name='user-all-assets'), - path('/nodes/children/tree/', api.UserPermNodeChildrenAsTreeApi.as_view(), name='user-perm-node-children-tree'), + path('/nodes/children/tree/', api.UserPermedAssetTreeAPI.as_view(), name='user-perm-node-children-tree'), ] user_group_permission_urlpatterns = [ diff --git a/apps/perms/utils/utils.py b/apps/perms/utils/utils.py index 16b7d3e99..70e05bddf 100644 --- a/apps/perms/utils/utils.py +++ b/apps/perms/utils/utils.py @@ -157,13 +157,27 @@ class UserPermUtil(object): return assets @classmethod - def get_favorite_assets(cls, user: User): + def get_favorite_assets(cls, user: User, search_asset=None, asset_category=None, asset_type=None): asset_ids = FavoriteAsset.get_user_favorite_asset_ids(user) - assets = Asset.objects.filter(id__in=asset_ids).valid() + q = Q(id__in=asset_ids) + if search_asset: + q &= Q(name__icontains=search_asset) | Q(address__icontains=search_asset) + if asset_category: + q &= Q(platform__category=asset_category) + if asset_type: + q &= Q(platform__type=asset_type) + assets = Asset.objects.filter(q).valid() return assets - def get_ungrouped_assets(self): - assets = Asset.objects.filter(id__in=self._user_direct_asset_ids).valid() + def get_ungrouped_assets(self, search_asset=None, asset_category=None, asset_type=None): + q = Q(id__in=self._user_direct_asset_ids) + if search_asset: + q &= Q(name__icontains=search_asset) | Q(address__icontains=search_asset) + if asset_category: + q &= Q(platform__category=asset_category) + if asset_type: + q &= Q(platform__type=asset_type) + assets = Asset.objects.filter(q).valid() return assets def _uuids_to_string(self, uuids):