mirror of https://github.com/jumpserver/jumpserver
* Bai reactor tree ( 重构获取完整资产树中节点下资产总数的逻辑) (#5548) * tree: v0.1 * tree: v0.2 * tree: v0.3 * tree: v0.4 * tree: 添加并发锁未请求到时的debug日志 * 以空间换时间的方式优化资产树 * Reactor tree togther v2 (#5576) * Bai reactor tree ( 重构获取完整资产树中节点下资产总数的逻辑) (#5548) * tree: v0.1 * tree: v0.2 * tree: v0.3 * tree: v0.4 * tree: 添加并发锁未请求到时的debug日志 * 以空间换时间的方式优化资产树 * 修改授权适配新方案 * 添加树处理工具 * 完成新的用户授权树计算以及修改一些信号 * 重构了获取资产的一些 api * 重构了一些节点的api * 整理了一些代码 * 完成了api 的重构 * 重构检查节点数量功能 * 完成重构授权树工具类 * api 添加强制刷新参数 * 整理一些信号 * 处理一些信号的问题 * 完成了信号的处理 * 重构了资产树相关的锁机制 * RebuildUserTreeTask 还得添加回来 * 优化下不能在root组织的检查函数 * 优化资产树变化时锁的使用 * 修改一些算法的小工具 * 资产树锁不再校验是否在具体组织里 * 整理了一些信号的位置 * 修复资产与节点关系维护的bug * 去掉一些调试代码 * 修复资产授权过期检查刷新授权树的 bug * 添加了可重入锁 * 添加一些计时，优化一些sql * 增加 union 查询的支持 * 尝试用 sql 解决节点资产数量问题 * 开始优化计算授权树节点资产数量不用冗余表 * 新代码能跑起来了，修复一下bug * 去掉 UserGrantedMappingNode 换成 UserAssetGrantedTreeNodeRelation * 修了些bug，做了些优化 * 优化QuerySetStage 执行逻辑 * 与小白的内存结合了 * 删掉老的表，迁移新的 assets_amount 字段 * 优化用户授权页面资产列表 count 慢 * 修复批量命令数量不对 * 修改获取非直接授权节点的 children 的逻辑 * 获取整棵树的节点 * 回退锁 * 整理迁移脚本 * 改变更新树策略 * perf: 修改一波缩进 * fix: 修改handler名称 * 修复授权树获取资产sql 泛滥 * 修复授权规则有效bug * 修复一些bug * 修复一些bug * 又修了一些小bug * 去掉了老的 get_nodes_all_assets * 修改一些写法 * Reactor tree togther b2 (#5570) * fix: 修改handler名称 * perf: 优化生成树 * perf: 去掉注释 * 优化了一些 * 重新生成迁移脚本 * 去掉周期检查节点资产数量的任务 * Pr@reactor tree togther guang@perf mapping (#5573) * fix: 修改handler名称 * perf: mapping 拆分出来 * 修改名称 * perf: 修改锁名 * perf: 去掉检查节点任务 * perf: 修改一下名称 * perf: 优化一波 Co-authored-by: Jiangjie.Bai <32935519+BaiJiangJie@users.noreply.github.com> Co-authored-by: Bai <bugatti_it@163.com> Co-authored-by: xinwen <coderWen@126.com> Co-authored-by: xinwen <coderWen@126.com> Co-authored-by: 老广 <ibuler@qq.com>pull/5610/head
Jiangjie.Bai
GitHub
parent
709e7af953
commit
7cf6e54f01
|
|
@ -3,10 +3,10 @@
|
|||
from assets.api import FilterAssetByNodeMixin
|
||||
from rest_framework.viewsets import ModelViewSet
|
||||
from rest_framework.generics import RetrieveAPIView
|
||||
from rest_framework.response import Response
|
||||
from rest_framework import status
|
||||
from django.shortcuts import get_object_or_404
|
||||
from django.utils.decorators import method_decorator
|
||||
|
||||
from assets.locks import NodeTreeUpdateLock
|
||||
from common.utils import get_logger, get_object_or_none
|
||||
from common.permissions import IsOrgAdmin, IsOrgAdminOrAppUser, IsSuperUser
|
||||
from orgs.mixins.api import OrgBulkModelViewSet
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import List
|
||||
|
||||
from common.utils.common import timeit
|
||||
from assets.models import Node, Asset
|
||||
from assets.pagination import AssetLimitOffsetPagination
|
||||
from common.utils import lazyproperty
|
||||
|
|
@ -7,6 +8,8 @@ from assets.utils import get_node, is_query_node_all_assets
|
|||
|
||||
|
||||
class SerializeToTreeNodeMixin:
|
||||
|
||||
@timeit
|
||||
def serialize_nodes(self, nodes: List[Node], with_asset_amount=False):
|
||||
if with_asset_amount:
|
||||
def _name(node: Node):
|
||||
|
|
@ -43,6 +46,7 @@ class SerializeToTreeNodeMixin:
|
|||
return platform
|
||||
return default
|
||||
|
||||
@timeit
|
||||
def serialize_assets(self, assets, node_key=None):
|
||||
if node_key is None:
|
||||
get_pid = lambda asset: getattr(asset, 'parent_key', '')
|
||||
|
|
|
|||
|
|
@ -17,12 +17,9 @@ from common.const.signals import PRE_REMOVE, POST_REMOVE
|
|||
from assets.models import Asset
|
||||
from common.utils import get_logger, get_object_or_none
|
||||
from common.tree import TreeNodeSerializer
|
||||
from common.const.distributed_lock_key import UPDATE_NODE_TREE_LOCK_KEY
|
||||
from orgs.mixins.api import OrgModelViewSet
|
||||
from orgs.mixins import generics
|
||||
from orgs.lock import org_level_transaction_lock
|
||||
from orgs.utils import current_org
|
||||
from assets.tasks import check_node_assets_amount_task
|
||||
from ..hands import IsOrgAdmin
|
||||
from ..models import Node
|
||||
from ..tasks import (
|
||||
|
|
@ -31,6 +28,7 @@ from ..tasks import (
|
|||
)
|
||||
from .. import serializers
|
||||
from .mixin import SerializeToTreeNodeMixin
|
||||
from assets.locks import NodeTreeUpdateLock
|
||||
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
|
@ -50,11 +48,6 @@ class NodeViewSet(OrgModelViewSet):
|
|||
permission_classes = (IsOrgAdmin,)
|
||||
serializer_class = serializers.NodeSerializer
|
||||
|
||||
@action(methods=[POST], detail=False, url_name='launch-check-assets-amount-task')
|
||||
def launch_check_assets_amount_task(self, request):
|
||||
task = check_node_assets_amount_task.delay(current_org.id)
|
||||
return Response(data={'task': task.id})
|
||||
|
||||
# 仅支持根节点指直接创建，子节点下的节点需要通过children接口创建
|
||||
def perform_create(self, serializer):
|
||||
child_key = Node.org_root().get_next_child_key()
|
||||
|
|
@ -184,9 +177,9 @@ class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
|
|||
if not include_assets:
|
||||
return []
|
||||
assets = self.instance.get_assets().only(
|
||||
"id", "hostname", "ip", "os",
|
||||
"org_id", "protocols", "is_active"
|
||||
)
|
||||
"id", "hostname", "ip", "os", "platform_id",
|
||||
"org_id", "protocols", "is_active",
|
||||
).prefetch_related('platform')
|
||||
return self.serialize_assets(assets, self.instance.key)
|
||||
|
||||
|
||||
|
|
@ -219,8 +212,6 @@ class NodeAddChildrenApi(generics.UpdateAPIView):
|
|||
return Response("OK")
|
||||
|
||||
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='patch')
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='put')
|
||||
class NodeAddAssetsApi(generics.UpdateAPIView):
|
||||
model = Node
|
||||
serializer_class = serializers.NodeAssetsSerializer
|
||||
|
|
@ -233,8 +224,6 @@ class NodeAddAssetsApi(generics.UpdateAPIView):
|
|||
instance.assets.add(*tuple(assets))
|
||||
|
||||
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='patch')
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='put')
|
||||
class NodeRemoveAssetsApi(generics.UpdateAPIView):
|
||||
model = Node
|
||||
serializer_class = serializers.NodeAssetsSerializer
|
||||
|
|
@ -251,8 +240,6 @@ class NodeRemoveAssetsApi(generics.UpdateAPIView):
|
|||
Node.org_root().assets.add(*orphan_assets)
|
||||
|
||||
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='patch')
|
||||
@method_decorator(org_level_transaction_lock(UPDATE_NODE_TREE_LOCK_KEY), name='put')
|
||||
class MoveAssetsToNodeApi(generics.UpdateAPIView):
|
||||
model = Node
|
||||
serializer_class = serializers.NodeAssetsSerializer
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
from orgs.utils import current_org
|
||||
from common.utils.lock import DistributedLock
|
||||
|
||||
|
||||
class NodeTreeUpdateLock(DistributedLock):
|
||||
name_template = 'assets.node.tree.update.<org_id:{org_id}>'
|
||||
|
||||
def get_name(self):
|
||||
if current_org:
|
||||
org_id = current_org.id
|
||||
else:
|
||||
org_id = 'current_org_is_null'
|
||||
name = self.name_template.format(
|
||||
org_id=org_id
|
||||
)
|
||||
return name
|
||||
|
||||
def __init__(self, blocking=True):
|
||||
name = self.get_name()
|
||||
super().__init__(name=name, blocking=blocking,
|
||||
release_lock_on_transaction_commit=True)
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
# Generated by Django 3.1 on 2021-02-04 09:49
|
||||
|
||||
from django.db import migrations
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('assets', '0065_auto_20210121_1549'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.RemoveField(
|
||||
model_name='node',
|
||||
name='assets_amount',
|
||||
),
|
||||
]
|
||||
|
|
@ -17,7 +17,7 @@ from orgs.mixins.models import OrgModelMixin, OrgManager
|
|||
from .base import ConnectivityMixin
|
||||
from .utils import Connectivity
|
||||
|
||||
__all__ = ['Asset', 'ProtocolsMixin', 'Platform']
|
||||
__all__ = ['Asset', 'ProtocolsMixin', 'Platform', 'AssetQuerySet']
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
|
@ -41,13 +41,6 @@ def default_node():
|
|||
|
||||
|
||||
class AssetManager(OrgManager):
|
||||
def get_queryset(self):
|
||||
return super().get_queryset().annotate(
|
||||
platform_base=models.F('platform__base')
|
||||
)
|
||||
|
||||
|
||||
class AssetOrgManager(OrgManager):
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -230,7 +223,6 @@ class Asset(ProtocolsMixin, NodesRelationMixin, OrgModelMixin):
|
|||
comment = models.TextField(default='', blank=True, verbose_name=_('Comment'))
|
||||
|
||||
objects = AssetManager.from_queryset(AssetQuerySet)()
|
||||
org_objects = AssetOrgManager.from_queryset(AssetQuerySet)()
|
||||
_connectivity = None
|
||||
|
||||
def __str__(self):
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from django.db import models
|
|||
from django.utils.translation import ugettext_lazy as _
|
||||
from django.conf import settings
|
||||
|
||||
from common.utils.common import timeit
|
||||
from common.utils import (
|
||||
ssh_key_string_to_obj, ssh_key_gen, get_logger, lazyproperty
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,15 +18,3 @@ class FavoriteAsset(CommonModelMixin):
|
|||
@classmethod
|
||||
def get_user_favorite_assets_id(cls, user):
|
||||
return cls.objects.filter(user=user).values_list('asset', flat=True)
|
||||
|
||||
@classmethod
|
||||
def get_user_favorite_assets(cls, user, asset_perms_id=None):
|
||||
from assets.models import Asset
|
||||
from perms.utils.asset.user_permission import get_user_granted_all_assets
|
||||
asset_ids = get_user_granted_all_assets(
|
||||
user,
|
||||
via_mapping_node=False,
|
||||
asset_perms_id=asset_perms_id
|
||||
).values_list('id', flat=True)
|
||||
query_name = cls.asset.field.related_query_name()
|
||||
return Asset.org_objects.filter(**{f'{query_name}__user_id': user.id}, id__in=asset_ids).distinct()
|
||||
|
|
|
|||
|
|
@ -1,23 +1,32 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import uuid
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
import threading
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from collections import defaultdict
|
||||
from django.db import models, transaction
|
||||
from django.db.models import Q
|
||||
from django.db.models import Q, Manager
|
||||
from django.db.utils import IntegrityError
|
||||
from django.utils.translation import ugettext_lazy as _
|
||||
from django.utils.translation import ugettext
|
||||
from django.db.transaction import atomic
|
||||
from django.core.cache import cache
|
||||
|
||||
from common.utils.lock import DistributedLock
|
||||
from common.utils.common import timeit
|
||||
from common.db.models import output_as_string
|
||||
from common.utils import get_logger
|
||||
from common.utils.common import lazyproperty
|
||||
from orgs.mixins.models import OrgModelMixin, OrgManager
|
||||
from orgs.utils import get_current_org, tmp_to_org
|
||||
from orgs.models import Organization
|
||||
|
||||
|
||||
__all__ = ['Node', 'FamilyMixin', 'compute_parent_key']
|
||||
__all__ = ['Node', 'FamilyMixin', 'compute_parent_key', 'NodeQuerySet']
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
|
|
@ -247,9 +256,125 @@ class FamilyMixin:
|
|||
return [*tuple(ancestors), self, *tuple(children)]
|
||||
|
||||
|
||||
class NodeAssetsMixin:
|
||||
class NodeAllAssetsMappingMixin:
|
||||
# Use a new plan
|
||||
|
||||
# { org_id: { node_key: [ asset1_id, asset2_id ] } }
|
||||
orgid_nodekey_assetsid_mapping = defaultdict(dict)
|
||||
|
||||
@classmethod
|
||||
def get_node_all_assets_id_mapping(cls, org_id):
|
||||
_mapping = cls.get_node_all_assets_id_mapping_from_memory(org_id)
|
||||
if _mapping:
|
||||
return _mapping
|
||||
|
||||
_mapping = cls.get_node_all_assets_id_mapping_from_cache_or_generate_to_cache(org_id)
|
||||
cls.set_node_all_assets_id_mapping_to_memory(org_id, mapping=_mapping)
|
||||
return _mapping
|
||||
|
||||
# from memory
|
||||
@classmethod
|
||||
def get_node_all_assets_id_mapping_from_memory(cls, org_id):
|
||||
mapping = cls.orgid_nodekey_assetsid_mapping.get(org_id, {})
|
||||
return mapping
|
||||
|
||||
@classmethod
|
||||
def set_node_all_assets_id_mapping_to_memory(cls, org_id, mapping):
|
||||
cls.orgid_nodekey_assetsid_mapping[org_id] = mapping
|
||||
|
||||
@classmethod
|
||||
def expire_node_all_assets_id_mapping_from_memory(cls, org_id):
|
||||
org_id = str(org_id)
|
||||
cls.orgid_nodekey_assetsid_mapping.pop(org_id, None)
|
||||
|
||||
# get order: from memory -> (from cache -> to generate)
|
||||
@classmethod
|
||||
def get_node_all_assets_id_mapping_from_cache_or_generate_to_cache(cls, org_id):
|
||||
mapping = cls.get_node_all_assets_id_mapping_from_cache(org_id)
|
||||
if mapping:
|
||||
return mapping
|
||||
|
||||
lock_key = f'KEY_LOCK_GENERATE_ORG_{org_id}_NODE_ALL_ASSETS_ID_MAPPING'
|
||||
logger.info(f'Thread[{threading.get_ident()}] acquiring lock[{lock_key}] ...')
|
||||
with DistributedLock(lock_key):
|
||||
logger.info(f'Thread[{threading.get_ident()}] acquire lock[{lock_key}] ok')
|
||||
# 这里使用无限期锁，原因是如果这里卡住了，就卡在数据库了，说明
|
||||
# 数据库繁忙，所以不应该再有线程执行这个操作，使数据库忙上加忙
|
||||
|
||||
# 这里最好先判断内存中有没有，防止同一进程的多个线程重复从 cache 中获取数据，
|
||||
# 但逻辑过于繁琐，直接判断 cache 吧
|
||||
_mapping = cls.get_node_all_assets_id_mapping_from_cache(org_id)
|
||||
if _mapping:
|
||||
return _mapping
|
||||
|
||||
_mapping = cls.generate_node_all_assets_id_mapping(org_id)
|
||||
cls.set_node_all_assets_id_mapping_to_cache(org_id=org_id, mapping=_mapping)
|
||||
return _mapping
|
||||
|
||||
@classmethod
|
||||
def get_node_all_assets_id_mapping_from_cache(cls, org_id):
|
||||
cache_key = cls._get_cache_key_for_node_all_assets_id_mapping(org_id)
|
||||
mapping = cache.get(cache_key)
|
||||
return mapping
|
||||
|
||||
@classmethod
|
||||
def set_node_all_assets_id_mapping_to_cache(cls, org_id, mapping):
|
||||
cache_key = cls._get_cache_key_for_node_all_assets_id_mapping(org_id)
|
||||
cache.set(cache_key, mapping, timeout=None)
|
||||
|
||||
@classmethod
|
||||
def expire_node_all_assets_id_mapping_from_cache(cls, org_id):
|
||||
cache_key = cls._get_cache_key_for_node_all_assets_id_mapping(org_id)
|
||||
cache.delete(cache_key)
|
||||
|
||||
@staticmethod
|
||||
def _get_cache_key_for_node_all_assets_id_mapping(org_id):
|
||||
return 'ASSETS_ORG_NODE_ALL_ASSETS_ID_MAPPING_{}'.format(org_id)
|
||||
|
||||
@classmethod
|
||||
def generate_node_all_assets_id_mapping(cls, org_id):
|
||||
from .asset import Asset
|
||||
|
||||
t1 = time.time()
|
||||
with tmp_to_org(org_id):
|
||||
nodes_id_key = Node.objects.filter(org_id=org_id) \
|
||||
.annotate(char_id=output_as_string('id')) \
|
||||
.values_list('char_id', 'key')
|
||||
|
||||
# * 直接取出全部. filter(node__org_id=org_id)(大规模下会更慢)
|
||||
nodes_assets_id = Asset.nodes.through.objects.all() \
|
||||
.annotate(char_node_id=output_as_string('node_id')) \
|
||||
.annotate(char_asset_id=output_as_string('asset_id')) \
|
||||
.values_list('char_node_id', 'char_asset_id')
|
||||
|
||||
node_id_ancestor_keys_mapping = {
|
||||
node_id: cls.get_node_ancestor_keys(node_key, with_self=True)
|
||||
for node_id, node_key in nodes_id_key
|
||||
}
|
||||
|
||||
nodeid_assetsid_mapping = defaultdict(set)
|
||||
for node_id, asset_id in nodes_assets_id:
|
||||
nodeid_assetsid_mapping[node_id].add(asset_id)
|
||||
|
||||
t2 = time.time()
|
||||
|
||||
mapping = defaultdict(set)
|
||||
for node_id, node_key in nodes_id_key:
|
||||
assets_id = nodeid_assetsid_mapping[node_id]
|
||||
node_ancestor_keys = node_id_ancestor_keys_mapping[node_id]
|
||||
for ancestor_key in node_ancestor_keys:
|
||||
mapping[ancestor_key].update(assets_id)
|
||||
|
||||
t3 = time.time()
|
||||
logger.debug('t1-t2(DB Query): {} s, t3-t2(Generate mapping): {} s'.format(t2-t1, t3-t2))
|
||||
return mapping
|
||||
|
||||
|
||||
class NodeAssetsMixin(NodeAllAssetsMappingMixin):
|
||||
org_id: str
|
||||
key = ''
|
||||
id = None
|
||||
objects: Manager
|
||||
|
||||
def get_all_assets(self):
|
||||
from .asset import Asset
|
||||
|
|
@ -263,8 +388,7 @@ class NodeAssetsMixin:
|
|||
# 可是 startswith 会导致表关联时 Asset 索引失效
|
||||
from .asset import Asset
|
||||
node_ids = cls.objects.filter(
|
||||
Q(key__startswith=f'{key}:') |
|
||||
Q(key=key)
|
||||
Q(key__startswith=f'{key}:') | Q(key=key)
|
||||
).values_list('id', flat=True).distinct()
|
||||
assets = Asset.objects.filter(
|
||||
nodes__id__in=list(node_ids)
|
||||
|
|
@ -283,29 +407,39 @@ class NodeAssetsMixin:
|
|||
return self.get_all_assets().valid()
|
||||
|
||||
@classmethod
|
||||
def get_nodes_all_assets_ids(cls, nodes_keys):
|
||||
assets_ids = cls.get_nodes_all_assets(nodes_keys).values_list('id', flat=True)
|
||||
def get_nodes_all_assets_ids_by_keys(cls, nodes_keys):
|
||||
nodes = Node.objects.filter(key__in=nodes_keys)
|
||||
assets_ids = cls.get_nodes_all_assets(*nodes).values_list('id', flat=True)
|
||||
return assets_ids
|
||||
|
||||
@classmethod
|
||||
def get_nodes_all_assets(cls, nodes_keys, extra_assets_ids=None):
|
||||
def get_nodes_all_assets(cls, *nodes):
|
||||
from .asset import Asset
|
||||
nodes_keys = cls.clean_children_keys(nodes_keys)
|
||||
q = Q()
|
||||
node_ids = ()
|
||||
for key in nodes_keys:
|
||||
q |= Q(key__startswith=f'{key}:')
|
||||
q |= Q(key=key)
|
||||
if q:
|
||||
node_ids = Node.objects.filter(q).distinct().values_list('id', flat=True)
|
||||
node_ids = set()
|
||||
descendant_node_query = Q()
|
||||
for n in nodes:
|
||||
node_ids.add(n.id)
|
||||
descendant_node_query |= Q(key__istartswith=f'{n.key}:')
|
||||
if descendant_node_query:
|
||||
_ids = Node.objects.order_by().filter(descendant_node_query).values_list('id', flat=True)
|
||||
node_ids.update(_ids)
|
||||
return Asset.objects.order_by().filter(nodes__id__in=node_ids).distinct()
|
||||
|
||||
q = Q(nodes__id__in=list(node_ids))
|
||||
if extra_assets_ids:
|
||||
q |= Q(id__in=extra_assets_ids)
|
||||
if q:
|
||||
return Asset.org_objects.filter(q).distinct()
|
||||
else:
|
||||
return Asset.objects.none()
|
||||
@property
|
||||
def assets_amount(self):
|
||||
assets_id = self.get_all_assets_id()
|
||||
return len(assets_id)
|
||||
|
||||
def get_all_assets_id(self):
|
||||
assets_id = self.get_all_assets_id_by_node_key(org_id=self.org_id, node_key=self.key)
|
||||
return set(assets_id)
|
||||
|
||||
@classmethod
|
||||
def get_all_assets_id_by_node_key(cls, org_id, node_key):
|
||||
org_id = str(org_id)
|
||||
nodekey_assetsid_mapping = cls.get_node_all_assets_id_mapping(org_id)
|
||||
assets_id = nodekey_assetsid_mapping.get(node_key, [])
|
||||
return set(assets_id)
|
||||
|
||||
|
||||
class SomeNodesMixin:
|
||||
|
|
@ -416,7 +550,6 @@ class Node(OrgModelMixin, SomeNodesMixin, FamilyMixin, NodeAssetsMixin):
|
|||
date_create = models.DateTimeField(auto_now_add=True)
|
||||
parent_key = models.CharField(max_length=64, verbose_name=_("Parent key"),
|
||||
db_index=True, default='')
|
||||
assets_amount = models.IntegerField(default=0)
|
||||
|
||||
objects = OrgManager.from_queryset(NodeQuerySet)()
|
||||
is_node = True
|
||||
|
|
|
|||
|
|
@ -199,7 +199,7 @@ class SystemUser(BaseUser):
|
|||
from assets.models import Node
|
||||
nodes_keys = self.nodes.all().values_list('key', flat=True)
|
||||
assets_ids = set(self.assets.all().values_list('id', flat=True))
|
||||
nodes_assets_ids = Node.get_nodes_all_assets_ids(nodes_keys)
|
||||
nodes_assets_ids = Node.get_nodes_all_assets_ids_by_keys(nodes_keys)
|
||||
assets_ids.update(nodes_assets_ids)
|
||||
assets = Asset.objects.filter(id__in=assets_ids)
|
||||
return assets
|
||||
|
|
|
|||
|
|
@ -111,7 +111,7 @@ class AssetSerializer(BulkOrgResourceModelSerializer):
|
|||
@classmethod
|
||||
def setup_eager_loading(cls, queryset):
|
||||
""" Perform necessary eager loading of data. """
|
||||
queryset = queryset.select_related('admin_user', 'domain', 'platform')
|
||||
queryset = queryset.prefetch_related('admin_user', 'domain', 'platform')
|
||||
queryset = queryset.prefetch_related('nodes', 'labels')
|
||||
return queryset
|
||||
|
||||
|
|
@ -166,13 +166,6 @@ class AssetDisplaySerializer(AssetSerializer):
|
|||
'connectivity',
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def setup_eager_loading(cls, queryset):
|
||||
queryset = super().setup_eager_loading(queryset)
|
||||
queryset = queryset\
|
||||
.annotate(admin_user_username=F('admin_user__username'))
|
||||
return queryset
|
||||
|
||||
|
||||
class PlatformSerializer(serializers.ModelSerializer):
|
||||
meta = serializers.DictField(required=False, allow_null=True)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
from .common import *
|
||||
from .maintain_nodes_tree import *
|
||||
|
|
@ -1,21 +1,17 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
from operator import add, sub
|
||||
|
||||
from assets.utils import is_asset_exists_in_node
|
||||
from django.db.models.signals import (
|
||||
post_save, m2m_changed, pre_delete, post_delete, pre_save
|
||||
)
|
||||
from django.db.models import Q, F
|
||||
from django.dispatch import receiver
|
||||
|
||||
from common.exceptions import M2MReverseNotAllowed
|
||||
from common.const.signals import PRE_ADD, POST_ADD, POST_REMOVE, PRE_CLEAR, PRE_REMOVE
|
||||
from common.const.signals import POST_ADD, POST_REMOVE, PRE_REMOVE
|
||||
from common.utils import get_logger
|
||||
from common.decorator import on_transaction_commit
|
||||
from .models import Asset, SystemUser, Node, compute_parent_key
|
||||
from assets.models import Asset, SystemUser, Node
|
||||
from users.models import User
|
||||
from .tasks import (
|
||||
from assets.tasks import (
|
||||
update_assets_hardware_info_util,
|
||||
test_asset_connectivity_util,
|
||||
push_system_user_to_assets_manual,
|
||||
|
|
@ -23,7 +19,6 @@ from .tasks import (
|
|||
add_nodes_assets_to_system_users
|
||||
)
|
||||
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
|
|
@ -202,134 +197,6 @@ def on_asset_nodes_add(instance, action, reverse, pk_set, **kwargs):
|
|||
m2m_model.objects.bulk_create(to_create)
|
||||
|
||||
|
||||
def _update_node_assets_amount(node: Node, asset_pk_set: set, operator=add):
|
||||
"""
|
||||
一个节点与多个资产关系变化时，更新计数
|
||||
|
||||
:param node: 节点实例
|
||||
:param asset_pk_set: 资产的`id`集合, 内部不会修改该值
|
||||
:param operator: 操作
|
||||
* -> Node
|
||||
# -> Asset
|
||||
|
||||
* [3]
|
||||
/ \
|
||||
* * [2]
|
||||
/ \
|
||||
* * [1]
|
||||
/ / \
|
||||
* [a] # # [b]
|
||||
|
||||
"""
|
||||
# 获取节点[1]祖先节点的 `key` 含自己，也就是[1, 2, 3]节点的`key`
|
||||
ancestor_keys = node.get_ancestor_keys(with_self=True)
|
||||
ancestors = Node.objects.filter(key__in=ancestor_keys).order_by('-key')
|
||||
to_update = []
|
||||
for ancestor in ancestors:
|
||||
# 迭代祖先节点的`key`，顺序是 [1] -> [2] -> [3]
|
||||
# 查询该节点及其后代节点是否包含要操作的资产，将包含的从要操作的
|
||||
# 资产集合中去掉，他们是重复节点，无论增加或删除都不会影响节点的资产数量
|
||||
|
||||
asset_pk_set -= set(Asset.objects.filter(
|
||||
id__in=asset_pk_set
|
||||
).filter(
|
||||
Q(nodes__key__istartswith=f'{ancestor.key}:') |
|
||||
Q(nodes__key=ancestor.key)
|
||||
).distinct().values_list('id', flat=True))
|
||||
if not asset_pk_set:
|
||||
# 要操作的资产集合为空，说明都是重复资产，不用改变节点资产数量
|
||||
# 而且既然它包含了，它的祖先节点肯定也包含了，所以祖先节点都不用
|
||||
# 处理了
|
||||
break
|
||||
ancestor.assets_amount = operator(F('assets_amount'), len(asset_pk_set))
|
||||
to_update.append(ancestor)
|
||||
Node.objects.bulk_update(to_update, fields=('assets_amount', 'parent_key'))
|
||||
|
||||
|
||||
def _remove_ancestor_keys(ancestor_key, tree_set):
|
||||
# 这里判断 `ancestor_key` 不能是空，防止数据错误导致的死循环
|
||||
# 判断是否在集合里，来区分是否已被处理过
|
||||
while ancestor_key and ancestor_key in tree_set:
|
||||
tree_set.remove(ancestor_key)
|
||||
ancestor_key = compute_parent_key(ancestor_key)
|
||||
|
||||
|
||||
def _update_nodes_asset_amount(node_keys, asset_pk, operator):
|
||||
"""
|
||||
一个资产与多个节点关系变化时，更新计数
|
||||
|
||||
:param node_keys: 节点 id 的集合
|
||||
:param asset_pk: 资产 id
|
||||
:param operator: 操作
|
||||
"""
|
||||
|
||||
# 所有相关节点的祖先节点，组成一棵局部树
|
||||
ancestor_keys = set()
|
||||
for key in node_keys:
|
||||
ancestor_keys.update(Node.get_node_ancestor_keys(key))
|
||||
|
||||
# 相关节点可能是其他相关节点的祖先节点，如果是从相关节点里干掉
|
||||
node_keys -= ancestor_keys
|
||||
|
||||
to_update_keys = []
|
||||
for key in node_keys:
|
||||
# 遍历相关节点，处理它及其祖先节点
|
||||
# 查询该节点是否包含待处理资产
|
||||
exists = is_asset_exists_in_node(asset_pk, key)
|
||||
parent_key = compute_parent_key(key)
|
||||
|
||||
if exists:
|
||||
# 如果资产在该节点，那么他及其祖先节点都不用处理
|
||||
_remove_ancestor_keys(parent_key, ancestor_keys)
|
||||
continue
|
||||
else:
|
||||
# 不存在，要更新本节点
|
||||
to_update_keys.append(key)
|
||||
# 这里判断 `parent_key` 不能是空，防止数据错误导致的死循环
|
||||
# 判断是否在集合里，来区分是否已被处理过
|
||||
while parent_key and parent_key in ancestor_keys:
|
||||
exists = is_asset_exists_in_node(asset_pk, parent_key)
|
||||
if exists:
|
||||
_remove_ancestor_keys(parent_key, ancestor_keys)
|
||||
break
|
||||
else:
|
||||
to_update_keys.append(parent_key)
|
||||
ancestor_keys.remove(parent_key)
|
||||
parent_key = compute_parent_key(parent_key)
|
||||
|
||||
Node.objects.filter(key__in=to_update_keys).update(
|
||||
assets_amount=operator(F('assets_amount'), 1)
|
||||
)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=Asset.nodes.through)
|
||||
def update_nodes_assets_amount(action, instance, reverse, pk_set, **kwargs):
|
||||
# 不允许 `pre_clear` ，因为该信号没有 `pk_set`
|
||||
# [官网](https://docs.djangoproject.com/en/3.1/ref/signals/#m2m-changed)
|
||||
refused = (PRE_CLEAR,)
|
||||
if action in refused:
|
||||
raise ValueError
|
||||
|
||||
mapper = {
|
||||
PRE_ADD: add,
|
||||
POST_REMOVE: sub
|
||||
}
|
||||
if action not in mapper:
|
||||
return
|
||||
|
||||
operator = mapper[action]
|
||||
|
||||
if reverse:
|
||||
node: Node = instance
|
||||
asset_pk_set = set(pk_set)
|
||||
_update_node_assets_amount(node, asset_pk_set, operator)
|
||||
else:
|
||||
asset_pk = instance.id
|
||||
# 与资产直接关联的节点
|
||||
node_keys = set(Node.objects.filter(id__in=pk_set).values_list('key', flat=True))
|
||||
_update_nodes_asset_amount(node_keys, asset_pk, operator)
|
||||
|
||||
|
||||
RELATED_NODE_IDS = '_related_node_ids'
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import os
|
||||
import threading
|
||||
|
||||
from django.db.models.signals import (
|
||||
m2m_changed, post_save, post_delete
|
||||
)
|
||||
from django.dispatch import receiver
|
||||
from django.utils.functional import LazyObject
|
||||
|
||||
from common.signals import django_ready
|
||||
from common.utils.connection import RedisPubSub
|
||||
from common.utils import get_logger
|
||||
from assets.models import Asset, Node
|
||||
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
# clear node assets mapping for memory
|
||||
# ------------------------------------
|
||||
|
||||
|
||||
def get_node_assets_mapping_for_memory_pub_sub():
|
||||
return RedisPubSub('fm.node_all_assets_id_memory_mapping')
|
||||
|
||||
|
||||
class NodeAssetsMappingForMemoryPubSub(LazyObject):
|
||||
def _setup(self):
|
||||
self._wrapped = get_node_assets_mapping_for_memory_pub_sub()
|
||||
|
||||
|
||||
node_assets_mapping_for_memory_pub_sub = NodeAssetsMappingForMemoryPubSub()
|
||||
|
||||
|
||||
def expire_node_assets_mapping_for_memory(org_id):
|
||||
# 所有进程清除(自己的 memory 数据)
|
||||
org_id = str(org_id)
|
||||
node_assets_mapping_for_memory_pub_sub.publish(org_id)
|
||||
# 当前进程清除(cache 数据)
|
||||
logger.debug(
|
||||
"Expire node assets id mapping from cache of org={}, pid={}"
|
||||
"".format(org_id, os.getpid())
|
||||
)
|
||||
Node.expire_node_all_assets_id_mapping_from_cache(org_id)
|
||||
|
||||
|
||||
@receiver(post_save, sender=Node)
|
||||
def on_node_post_create(sender, instance, created, update_fields, **kwargs):
|
||||
if created:
|
||||
need_expire = True
|
||||
elif update_fields and 'key' in update_fields:
|
||||
need_expire = True
|
||||
else:
|
||||
need_expire = False
|
||||
|
||||
if need_expire:
|
||||
expire_node_assets_mapping_for_memory(instance.org_id)
|
||||
|
||||
|
||||
@receiver(post_delete, sender=Node)
|
||||
def on_node_post_delete(sender, instance, **kwargs):
|
||||
expire_node_assets_mapping_for_memory(instance.org_id)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=Asset.nodes.through)
|
||||
def on_node_asset_change(sender, instance, **kwargs):
|
||||
expire_node_assets_mapping_for_memory(instance.org_id)
|
||||
|
||||
|
||||
@receiver(django_ready)
|
||||
def subscribe_node_assets_mapping_expire(sender, **kwargs):
|
||||
logger.debug("Start subscribe for expire node assets id mapping from memory")
|
||||
|
||||
def keep_subscribe():
|
||||
subscribe = node_assets_mapping_for_memory_pub_sub.subscribe()
|
||||
for message in subscribe.listen():
|
||||
if message["type"] != "message":
|
||||
continue
|
||||
org_id = message['data'].decode()
|
||||
Node.expire_node_all_assets_id_mapping_from_memory(org_id)
|
||||
logger.debug(
|
||||
"Expire node assets id mapping from memory of org={}, pid={}"
|
||||
"".format(str(org_id), os.getpid())
|
||||
)
|
||||
t = threading.Thread(target=keep_subscribe)
|
||||
t.daemon = True
|
||||
t.start()
|
||||
|
|
@ -12,6 +12,7 @@ __all__ = ['add_nodes_assets_to_system_users']
|
|||
@tmp_to_root_org()
|
||||
def add_nodes_assets_to_system_users(nodes_keys, system_users):
|
||||
from ..models import Node
|
||||
assets = Node.get_nodes_all_assets(nodes_keys).values_list('id', flat=True)
|
||||
nodes = Node.objects.filter(key__in=nodes_keys)
|
||||
assets = Node.get_nodes_all_assets(*nodes)
|
||||
for system_user in system_users:
|
||||
system_user.assets.add(*tuple(assets))
|
||||
|
|
|
|||
|
|
@ -141,7 +141,8 @@ def gather_asset_users(assets, task_name=None):
|
|||
|
||||
@shared_task(queue="ansible")
|
||||
def gather_nodes_asset_users(nodes_key):
|
||||
assets = Node.get_nodes_all_assets(nodes_key)
|
||||
nodes = Node.objects.filter(key__in=nodes_key)
|
||||
assets = Node.get_nodes_all_assets(*nodes)
|
||||
assets_groups_by_100 = [assets[i:i+100] for i in range(0, len(assets), 100)]
|
||||
for _assets in assets_groups_by_100:
|
||||
gather_asset_users(_assets)
|
||||
|
|
|
|||
|
|
@ -1,27 +0,0 @@
|
|||
from celery import shared_task
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
|
||||
from orgs.models import Organization
|
||||
from orgs.utils import tmp_to_org
|
||||
from ops.celery.decorator import register_as_period_task
|
||||
from assets.utils import check_node_assets_amount
|
||||
|
||||
from common.utils.lock import AcquireFailed
|
||||
from common.utils import get_logger
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
@shared_task(queue='celery_heavy_tasks')
|
||||
def check_node_assets_amount_task(org_id=Organization.ROOT_ID):
|
||||
try:
|
||||
with tmp_to_org(Organization.get_instance(org_id)):
|
||||
check_node_assets_amount()
|
||||
except AcquireFailed:
|
||||
logger.error(_('The task of self-checking is already running and cannot be started repeatedly'))
|
||||
|
||||
|
||||
@register_as_period_task(crontab='0 2 * * *')
|
||||
@shared_task(queue='celery_heavy_tasks')
|
||||
def check_node_assets_amount_period_task():
|
||||
check_node_assets_amount_task()
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
from assets.tree import Tree
|
||||
|
||||
|
||||
def test():
|
||||
from orgs.models import Organization
|
||||
from assets.models import Node, Asset
|
||||
import time
|
||||
Organization.objects.get(id='1863cf22-f666-474e-94aa-935fe175203c').change_to()
|
||||
|
||||
t1 = time.time()
|
||||
nodes = list(Node.objects.exclude(key__startswith='-').only('id', 'key', 'parent_key'))
|
||||
node_asset_id_pairs = Asset.nodes.through.objects.all().values_list('node_id', 'asset_id')
|
||||
t2 = time.time()
|
||||
node_asset_id_pairs = list(node_asset_id_pairs)
|
||||
tree = Tree(nodes, node_asset_id_pairs)
|
||||
tree.build_tree()
|
||||
tree.nodes = None
|
||||
tree.node_asset_id_pairs = None
|
||||
import pickle
|
||||
d = pickle.dumps(tree)
|
||||
print('------------', len(d))
|
||||
return tree
|
||||
tree.compute_tree_node_assets_amount()
|
||||
|
||||
print(f'校对算法准确性 ......')
|
||||
for node in nodes:
|
||||
tree_node = tree.key_tree_node_mapper[node.key]
|
||||
if tree_node.assets_amount != node.assets_amount:
|
||||
print(f'ERROR: {tree_node.assets_amount} {node.assets_amount}')
|
||||
# print(f'OK {tree_node.asset_amount} {node.assets_amount}')
|
||||
|
||||
print(f'数据库时间: {t2 - t1}')
|
||||
return tree
|
||||
|
|
@ -2,7 +2,6 @@
|
|||
from django.urls import path, re_path
|
||||
from rest_framework_nested import routers
|
||||
from rest_framework_bulk.routes import BulkRouter
|
||||
from django.db.transaction import non_atomic_requests
|
||||
|
||||
from common import api as capi
|
||||
|
||||
|
|
@ -57,9 +56,9 @@ urlpatterns = [
|
|||
path('nodes/children/', api.NodeChildrenApi.as_view(), name='node-children-2'),
|
||||
path('nodes/<uuid:pk>/children/add/', api.NodeAddChildrenApi.as_view(), name='node-add-children'),
|
||||
path('nodes/<uuid:pk>/assets/', api.NodeAssetsApi.as_view(), name='node-assets'),
|
||||
path('nodes/<uuid:pk>/assets/add/', non_atomic_requests(api.NodeAddAssetsApi.as_view()), name='node-add-assets'),
|
||||
path('nodes/<uuid:pk>/assets/replace/', non_atomic_requests(api.MoveAssetsToNodeApi.as_view()), name='node-replace-assets'),
|
||||
path('nodes/<uuid:pk>/assets/remove/', non_atomic_requests(api.NodeRemoveAssetsApi.as_view()), name='node-remove-assets'),
|
||||
path('nodes/<uuid:pk>/assets/add/', api.NodeAddAssetsApi.as_view(), name='node-add-assets'),
|
||||
path('nodes/<uuid:pk>/assets/replace/', api.MoveAssetsToNodeApi.as_view(), name='node-replace-assets'),
|
||||
path('nodes/<uuid:pk>/assets/remove/', api.NodeRemoveAssetsApi.as_view(), name='node-remove-assets'),
|
||||
path('nodes/<uuid:pk>/tasks/', api.NodeTaskCreateApi.as_view(), name='node-task-create'),
|
||||
|
||||
path('gateways/<uuid:pk>/test-connective/', api.GatewayTestConnectionApi.as_view(), name='test-gateway-connective'),
|
||||
|
|
|
|||
|
|
@ -1,43 +1,16 @@
|
|||
# ~*~ coding: utf-8 ~*~
|
||||
#
|
||||
import time
|
||||
|
||||
from django.db.models import Q
|
||||
|
||||
from common.utils import get_logger, dict_get_any, is_uuid, get_object_or_none
|
||||
from common.utils.lock import DistributedLock
|
||||
from collections import defaultdict
|
||||
from common.utils import get_logger, dict_get_any, is_uuid, get_object_or_none, timeit
|
||||
from common.http import is_true
|
||||
from .models import Asset, Node
|
||||
from common.struct import Stack
|
||||
from common.db.models import output_as_string
|
||||
|
||||
from .models import Node
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
@DistributedLock(name="assets.node.check_node_assets_amount", blocking=False)
|
||||
def check_node_assets_amount():
|
||||
for node in Node.objects.all():
|
||||
logger.info(f'Check node assets amount: {node}')
|
||||
assets_amount = Asset.objects.filter(
|
||||
Q(nodes__key__istartswith=f'{node.key}:') | Q(nodes=node)
|
||||
).distinct().count()
|
||||
|
||||
if node.assets_amount != assets_amount:
|
||||
logger.warn(f'Node wrong assets amount <Node:{node.key}> '
|
||||
f'{node.assets_amount} right is {assets_amount}')
|
||||
node.assets_amount = assets_amount
|
||||
node.save()
|
||||
# 防止自检程序给数据库的压力太大
|
||||
time.sleep(0.1)
|
||||
|
||||
|
||||
def is_asset_exists_in_node(asset_pk, node_key):
|
||||
return Asset.objects.filter(
|
||||
id=asset_pk
|
||||
).filter(
|
||||
Q(nodes__key__istartswith=f'{node_key}:') | Q(nodes__key=node_key)
|
||||
).exists()
|
||||
|
||||
|
||||
def is_query_node_all_assets(request):
|
||||
request = request
|
||||
query_all_arg = request.query_params.get('all', 'true')
|
||||
|
|
@ -57,3 +30,79 @@ def get_node(request):
|
|||
else:
|
||||
node = get_object_or_none(Node, key=node_id)
|
||||
return node
|
||||
|
||||
|
||||
class NodeAssetsInfo:
|
||||
__slots__ = ('key', 'assets_amount', 'assets')
|
||||
|
||||
def __init__(self, key, assets_amount, assets):
|
||||
self.key = key
|
||||
self.assets_amount = assets_amount
|
||||
self.assets = assets
|
||||
|
||||
def __str__(self):
|
||||
return self.key
|
||||
|
||||
|
||||
class NodeAssetsUtil:
|
||||
def __init__(self, nodes, nodekey_assetsid_mapper):
|
||||
"""
|
||||
:param nodes: 节点
|
||||
:param nodekey_assetsid_mapper: 节点直接资产id的映射 {"key1": set(), "key2": set()}
|
||||
"""
|
||||
self.nodes = nodes
|
||||
# node_id --> set(asset_id1, asset_id2)
|
||||
self.nodekey_assetsid_mapper = nodekey_assetsid_mapper
|
||||
self.nodekey_assetsinfo_mapper = {}
|
||||
|
||||
@timeit
|
||||
def generate(self):
|
||||
# 准备排序好的资产信息数据
|
||||
infos = []
|
||||
for node in self.nodes:
|
||||
assets = self.nodekey_assetsid_mapper.get(node.key, set())
|
||||
info = NodeAssetsInfo(key=node.key, assets_amount=0, assets=assets)
|
||||
infos.append(info)
|
||||
infos = sorted(infos, key=lambda i: [int(i) for i in i.key.split(':')])
|
||||
# 这个守卫需要添加一下，避免最后一个无法出栈
|
||||
guarder = NodeAssetsInfo(key='', assets_amount=0, assets=set())
|
||||
infos.append(guarder)
|
||||
|
||||
stack = Stack()
|
||||
for info in infos:
|
||||
# 如果栈顶的不是这个节点的父祖节点，那么可以出栈了，可以计算资产数量了
|
||||
while stack.top and not info.key.startswith(f'{stack.top.key}:'):
|
||||
pop_info = stack.pop()
|
||||
pop_info.assets_amount = len(pop_info.assets)
|
||||
self.nodekey_assetsinfo_mapper[pop_info.key] = pop_info
|
||||
if not stack.top:
|
||||
continue
|
||||
stack.top.assets.update(pop_info.assets)
|
||||
stack.push(info)
|
||||
|
||||
def get_assets_by_key(self, key):
|
||||
info = self.nodekey_assetsinfo_mapper[key]
|
||||
return info['assets']
|
||||
|
||||
def get_assets_amount(self, key):
|
||||
info = self.nodekey_assetsinfo_mapper[key]
|
||||
return info.assets_amount
|
||||
|
||||
@classmethod
|
||||
def test_it(cls):
|
||||
from assets.models import Node, Asset
|
||||
|
||||
nodes = list(Node.objects.all())
|
||||
nodes_assets = Asset.nodes.through.objects.all()\
|
||||
.annotate(aid=output_as_string('asset_id'))\
|
||||
.values_list('node__key', 'aid')
|
||||
|
||||
mapping = defaultdict(set)
|
||||
for key, asset_id in nodes_assets:
|
||||
mapping[key].add(asset_id)
|
||||
|
||||
util = cls(nodes, mapping)
|
||||
util.generate()
|
||||
return util
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,2 +0,0 @@
|
|||
UPDATE_NODE_TREE_LOCK_KEY = 'org_level_transaction_lock_{org_id}_assets_update_node_tree'
|
||||
UPDATE_MAPPING_NODE_TASK_LOCK_KEY = 'org_level_transaction_lock_{user_id}_update_mapping_node_task'
|
||||
|
|
@ -82,3 +82,7 @@ class JMSModel(JMSBaseModel):
|
|||
|
||||
def concated_display(name1, name2):
|
||||
return Concat(F(name1), Value('('), F(name2), Value(')'))
|
||||
|
||||
|
||||
def output_as_string(field_name):
|
||||
return ExpressionWrapper(F(field_name), output_field=CharField())
|
||||
|
|
|
|||
|
|
@ -254,3 +254,22 @@ def get_disk_usage():
|
|||
mount_points = [p.mountpoint for p in partitions]
|
||||
usages = {p: psutil.disk_usage(p) for p in mount_points}
|
||||
return usages
|
||||
|
||||
|
||||
class Time:
|
||||
def __init__(self):
|
||||
self._timestamps = []
|
||||
self._msgs = []
|
||||
|
||||
def begin(self):
|
||||
self._timestamps.append(time.time())
|
||||
|
||||
def time(self, msg):
|
||||
self._timestamps.append(time.time())
|
||||
self._msgs.append(msg)
|
||||
|
||||
def print(self):
|
||||
last, *timestamps = self._timestamps
|
||||
for timestamp, msg in zip(timestamps, self._msgs):
|
||||
logger.debug(f'TIME_IT: {msg} {timestamp-last}')
|
||||
last = timestamp
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
from functools import wraps
|
||||
import threading
|
||||
|
||||
from redis_lock import Lock as RedisLock
|
||||
from redis_lock import Lock as RedisLock, NotAcquired
|
||||
from redis import Redis
|
||||
from django.db import transaction
|
||||
|
||||
from common.utils import get_logger
|
||||
from common.utils.inspect import copy_function_args
|
||||
|
|
@ -16,7 +17,8 @@ class AcquireFailed(RuntimeError):
|
|||
|
||||
|
||||
class DistributedLock(RedisLock):
|
||||
def __init__(self, name, blocking=True, expire=60*2, auto_renewal=True):
|
||||
def __init__(self, name, blocking=True, expire=None, release_lock_on_transaction_commit=False,
|
||||
release_raise_exc=False, auto_renewal_seconds=60*2):
|
||||
"""
|
||||
使用 redis 构造的分布式锁
|
||||
|
||||
|
|
@ -25,31 +27,46 @@ class DistributedLock(RedisLock):
|
|||
:param blocking:
|
||||
该参数只在锁作为装饰器或者 `with` 时有效。
|
||||
:param expire:
|
||||
锁的过期时间，注意不一定是锁到这个时间就释放了，分两种情况
|
||||
当 `auto_renewal=False` 时，锁会释放
|
||||
当 `auto_renewal=True` 时，如果过期之前程序还没释放锁，我们会延长锁的存活时间。
|
||||
这里的作用是防止程序意外终止没有释放锁，导致死锁。
|
||||
锁的过期时间
|
||||
:param release_lock_on_transaction_commit:
|
||||
是否在当前事务结束后再释放锁
|
||||
:param release_raise_exc:
|
||||
释放锁时，如果没有持有锁是否抛异常或静默
|
||||
:param auto_renewal_seconds:
|
||||
当持有一个无限期锁的时候，刷新锁的时间，具体参考 `redis_lock.Lock#auto_renewal`
|
||||
"""
|
||||
self.kwargs_copy = copy_function_args(self.__init__, locals())
|
||||
redis = Redis(host=CONFIG.REDIS_HOST, port=CONFIG.REDIS_PORT, password=CONFIG.REDIS_PASSWORD)
|
||||
|
||||
if expire is None:
|
||||
expire = auto_renewal_seconds
|
||||
auto_renewal = True
|
||||
else:
|
||||
auto_renewal = False
|
||||
|
||||
super().__init__(redis_client=redis, name=name, expire=expire, auto_renewal=auto_renewal)
|
||||
self._blocking = blocking
|
||||
self._release_lock_on_transaction_commit = release_lock_on_transaction_commit
|
||||
self._release_raise_exc = release_raise_exc
|
||||
|
||||
def __enter__(self):
|
||||
thread_id = threading.current_thread().ident
|
||||
logger.debug(f'DISTRIBUTED_LOCK: <thread_id:{thread_id}> attempt to acquire <lock:{self._name}> ...')
|
||||
logger.debug(f'Attempt to acquire global lock: thread {thread_id} lock {self._name}')
|
||||
acquired = self.acquire(blocking=self._blocking)
|
||||
if self._blocking and not acquired:
|
||||
logger.debug(f'DISTRIBUTED_LOCK: <thread_id:{thread_id}> was not acquired <lock:{self._name}>, but blocking=True')
|
||||
logger.debug(f'Not acquired lock, but blocking=True, thread {thread_id} lock {self._name}')
|
||||
raise EnvironmentError("Lock wasn't acquired, but blocking=True")
|
||||
if not acquired:
|
||||
logger.debug(f'DISTRIBUTED_LOCK: <thread_id:{thread_id}> acquire <lock:{self._name}> failed')
|
||||
logger.debug(f'Not acquired the lock, thread {thread_id} lock {self._name}')
|
||||
raise AcquireFailed
|
||||
logger.debug(f'DISTRIBUTED_LOCK: <thread_id:{thread_id}> acquire <lock:{self._name}> ok')
|
||||
logger.debug(f'Acquire lock success, thread {thread_id} lock {self._name}')
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type=None, exc_value=None, traceback=None):
|
||||
self.release()
|
||||
if self._release_lock_on_transaction_commit:
|
||||
transaction.on_commit(self.release)
|
||||
else:
|
||||
self.release()
|
||||
|
||||
def __call__(self, func):
|
||||
@wraps(func)
|
||||
|
|
@ -57,5 +74,17 @@ class DistributedLock(RedisLock):
|
|||
# 要创建一个新的锁对象
|
||||
with self.__class__(**self.kwargs_copy):
|
||||
return func(*args, **kwds)
|
||||
|
||||
return inner
|
||||
|
||||
def locked_by_me(self):
|
||||
if self.locked():
|
||||
if self.get_owner_id() == self.id:
|
||||
return True
|
||||
return False
|
||||
|
||||
def release(self):
|
||||
try:
|
||||
super().release()
|
||||
except AcquireFailed as e:
|
||||
if self._release_raise_exc:
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -1,131 +0,0 @@
|
|||
from uuid import uuid4
|
||||
from functools import wraps
|
||||
|
||||
from django.core.cache import cache
|
||||
from django.db.transaction import atomic
|
||||
from rest_framework.request import Request
|
||||
from rest_framework.exceptions import NotAuthenticated
|
||||
|
||||
from orgs.utils import current_org
|
||||
from common.exceptions import SomeoneIsDoingThis, Timeout
|
||||
from common.utils.timezone import dt_formater, now
|
||||
|
||||
# Redis 中锁值得模板，该模板提供了很强的可读性，方便调试与排错
|
||||
VALUE_TEMPLATE = '{stage}:{username}:{user_id}:{now}:{rand_str}'
|
||||
|
||||
# 锁的状态
|
||||
DOING = 'doing' # 处理中，此状态的锁可以被干掉
|
||||
COMMITING = 'commiting' # 提交事务中，此状态很重要，要确保事务在锁消失之前返回了，不要轻易删除该锁
|
||||
|
||||
client = cache.client.get_client(write=True)
|
||||
|
||||
|
||||
"""
|
||||
将锁的状态从 `doing` 切换到 `commiting`
|
||||
KEYS[1]： key
|
||||
ARGV[1]: doingvalue
|
||||
ARGV[2]: commitingvalue
|
||||
ARGV[3]: timeout
|
||||
"""
|
||||
change_lock_state_to_commiting_lua = '''
|
||||
if (redis.call("get", KEYS[1]) == ARGV[1])
|
||||
then
|
||||
return redis.call("set", KEYS[1], ARGV[2], "EX", ARGV[3], "XX")
|
||||
else
|
||||
return 0
|
||||
end
|
||||
'''
|
||||
change_lock_state_to_commiting_lua_obj = client.register_script(change_lock_state_to_commiting_lua)
|
||||
|
||||
|
||||
"""
|
||||
释放锁，两种`value`都要检查`doing`和`commiting`
|
||||
KEYS[1]： key
|
||||
ARGV[1]: 两个 `value` 中的其中一个
|
||||
ARGV[2]: 两个 `value` 中的其中一个
|
||||
"""
|
||||
release_lua = '''
|
||||
if (redis.call("get",KEYS[1]) == ARGV[1] or redis.call("get",KEYS[1]) == ARGV[2])
|
||||
then
|
||||
return redis.call("del",KEYS[1])
|
||||
else
|
||||
return 0
|
||||
end
|
||||
'''
|
||||
release_lua_obj = client.register_script(release_lua)
|
||||
|
||||
|
||||
def acquire(key, value, timeout):
|
||||
return client.set(key, value, ex=timeout, nx=True)
|
||||
|
||||
|
||||
def get(key):
|
||||
return client.get(key)
|
||||
|
||||
|
||||
def change_lock_state_to_commiting(key, doingvalue, commitingvalue, timeout=600):
|
||||
# 将锁的状态从 `doing` 切换到 `commiting`
|
||||
return bool(change_lock_state_to_commiting_lua_obj(keys=(key,), args=(doingvalue, commitingvalue, timeout)))
|
||||
|
||||
|
||||
def release(key, value1, value2):
|
||||
# 释放锁，两种`value` `doing`和`commiting` 都要检查
|
||||
return release_lua_obj(keys=(key,), args=(value1, value2))
|
||||
|
||||
|
||||
def _generate_value(request: Request, stage=DOING):
|
||||
# 不支持匿名用户
|
||||
user = request.user
|
||||
if user.is_anonymous:
|
||||
raise NotAuthenticated
|
||||
|
||||
return VALUE_TEMPLATE.format(
|
||||
stage=stage, username=user.username, user_id=user.id,
|
||||
now=dt_formater(now()), rand_str=uuid4()
|
||||
)
|
||||
|
||||
|
||||
default_wait_msg = SomeoneIsDoingThis.default_detail
|
||||
|
||||
|
||||
def org_level_transaction_lock(key, timeout=300, wait_msg=default_wait_msg):
|
||||
"""
|
||||
被装饰的 `View` 必须取消自身的 `ATOMIC_REQUESTS`，因为该装饰器要有事务的完全控制权
|
||||
[官网](https://docs.djangoproject.com/en/3.1/topics/db/transactions/#tying-transactions-to-http-requests)
|
||||
|
||||
1. 获取锁：只有当锁对应的 `key` 不存在时成功获取，`value` 设置为 `doing`
|
||||
2. 开启事务：本次请求的事务必须确保在这里开启
|
||||
3. 执行 `View` 体
|
||||
4. `View` 体执行结束未异常，此时事务还未提交
|
||||
5. 检查锁是否过时，过时事务回滚，不过时，重新设置`key`延长`key`有效期，已确保足够时间提交事务，同时把`key`的状态改为`commiting`
|
||||
6. 提交事务
|
||||
7. 释放锁，释放的时候会检查`doing`与`commiting`的值，因为删除或者更改锁必须提供与当前锁的`value`相同的值，确保不误删
|
||||
[锁参考文章](http://doc.redisfans.com/string/set.html#id2)
|
||||
"""
|
||||
|
||||
def decorator(fun):
|
||||
@wraps(fun)
|
||||
def wrapper(request, *args, **kwargs):
|
||||
# `key`可能是组织相关的，如果是把组织`id`加上
|
||||
_key = key.format(org_id=current_org.id)
|
||||
doing_value = _generate_value(request)
|
||||
commiting_value = _generate_value(request, stage=COMMITING)
|
||||
try:
|
||||
lock = acquire(_key, doing_value, timeout)
|
||||
if not lock:
|
||||
raise SomeoneIsDoingThis(detail=wait_msg)
|
||||
with atomic(savepoint=False):
|
||||
ret = fun(request, *args, **kwargs)
|
||||
# 提交事务前，检查一下锁是否还在
|
||||
# 锁在的话，更新锁的状态为 `commiting`，延长锁时间，确保事务提交
|
||||
# 锁不在的话回滚
|
||||
ok = change_lock_state_to_commiting(_key, doing_value, commiting_value)
|
||||
if not ok:
|
||||
# 超时或者被中断了
|
||||
raise Timeout
|
||||
return ret
|
||||
finally:
|
||||
# 释放锁，锁的两个值都要尝试，不确定异常是从什么位置抛出的
|
||||
release(_key, commiting_value, doing_value)
|
||||
return wrapper
|
||||
return decorator
|
||||
|
|
@ -184,3 +184,8 @@ def org_aware_func(org_arg_name):
|
|||
|
||||
|
||||
current_org = LocalProxy(get_current_org)
|
||||
|
||||
|
||||
def ensure_in_real_or_default_org():
|
||||
if not current_org or current_org.is_root():
|
||||
raise ValueError('You must in a real or default org!')
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from applications.models import Application
|
|||
from perms.utils.application.permission import (
|
||||
get_application_system_users_id
|
||||
)
|
||||
from perms.api.asset.user_permission.mixin import ForAdminMixin, ForUserMixin
|
||||
from perms.api.asset.user_permission.mixin import RoleAdminMixin, RoleUserMixin
|
||||
from common.permissions import IsOrgAdminOrAppUser
|
||||
from perms.hands import User, SystemUser
|
||||
from perms import serializers
|
||||
|
|
@ -43,11 +43,11 @@ class GrantedApplicationSystemUsersMixin(ListAPIView):
|
|||
return system_users
|
||||
|
||||
|
||||
class UserGrantedApplicationSystemUsersApi(ForAdminMixin, GrantedApplicationSystemUsersMixin):
|
||||
class UserGrantedApplicationSystemUsersApi(RoleAdminMixin, GrantedApplicationSystemUsersMixin):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedApplicationSystemUsersApi(ForUserMixin, GrantedApplicationSystemUsersMixin):
|
||||
class MyGrantedApplicationSystemUsersApi(RoleUserMixin, GrantedApplicationSystemUsersMixin):
|
||||
pass
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from applications.api.mixin import (
|
|||
SerializeApplicationToTreeNodeMixin
|
||||
)
|
||||
from perms import serializers
|
||||
from perms.api.asset.user_permission.mixin import ForAdminMixin, ForUserMixin
|
||||
from perms.api.asset.user_permission.mixin import RoleAdminMixin, RoleUserMixin
|
||||
from perms.utils.application.user_permission import (
|
||||
get_user_granted_all_applications
|
||||
)
|
||||
|
|
@ -34,11 +34,11 @@ class AllGrantedApplicationsMixin(CommonApiMixin, ListAPIView):
|
|||
return queryset.only(*self.only_fields)
|
||||
|
||||
|
||||
class UserAllGrantedApplicationsApi(ForAdminMixin, AllGrantedApplicationsMixin):
|
||||
class UserAllGrantedApplicationsApi(RoleAdminMixin, AllGrantedApplicationsMixin):
|
||||
pass
|
||||
|
||||
|
||||
class MyAllGrantedApplicationsApi(ForUserMixin, AllGrantedApplicationsMixin):
|
||||
class MyAllGrantedApplicationsApi(RoleUserMixin, AllGrantedApplicationsMixin):
|
||||
pass
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -4,37 +4,23 @@ from rest_framework.request import Request
|
|||
|
||||
from common.permissions import IsOrgAdminOrAppUser, IsValidUser
|
||||
from common.utils import lazyproperty
|
||||
from common.http import is_true
|
||||
from orgs.utils import tmp_to_root_org
|
||||
from users.models import User
|
||||
from perms.models import UserGrantedMappingNode
|
||||
from perms.utils.asset.user_permission import UserGrantedTreeRefreshController
|
||||
|
||||
|
||||
class UserNodeGrantStatusDispatchMixin:
|
||||
class PermBaseMixin:
|
||||
user: User
|
||||
|
||||
@staticmethod
|
||||
def get_mapping_node_by_key(key, user):
|
||||
return UserGrantedMappingNode.objects.get(key=key, user=user)
|
||||
|
||||
def dispatch_get_data(self, key, user):
|
||||
status = UserGrantedMappingNode.get_node_granted_status(key, user)
|
||||
if status == UserGrantedMappingNode.GRANTED_DIRECT:
|
||||
return self.get_data_on_node_direct_granted(key)
|
||||
elif status == UserGrantedMappingNode.GRANTED_INDIRECT:
|
||||
return self.get_data_on_node_indirect_granted(key)
|
||||
else:
|
||||
return self.get_data_on_node_not_granted(key)
|
||||
|
||||
def get_data_on_node_direct_granted(self, key):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_data_on_node_indirect_granted(self, key):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_data_on_node_not_granted(self, key):
|
||||
raise NotImplementedError
|
||||
def get(self, request, *args, **kwargs):
|
||||
force = is_true(request.query_params.get('rebuild_tree'))
|
||||
controller = UserGrantedTreeRefreshController(self.user)
|
||||
controller.refresh_if_need(force)
|
||||
return super().get(request, *args, **kwargs)
|
||||
|
||||
|
||||
class ForAdminMixin:
|
||||
class RoleAdminMixin(PermBaseMixin):
|
||||
permission_classes = (IsOrgAdminOrAppUser,)
|
||||
kwargs: dict
|
||||
|
||||
|
|
@ -44,7 +30,7 @@ class ForAdminMixin:
|
|||
return User.objects.get(id=user_id)
|
||||
|
||||
|
||||
class ForUserMixin:
|
||||
class RoleUserMixin(PermBaseMixin):
|
||||
permission_classes = (IsValidUser,)
|
||||
request: Request
|
||||
|
||||
|
|
|
|||
|
|
@ -1,156 +0,0 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
from perms.api.asset.user_permission.mixin import UserNodeGrantStatusDispatchMixin
|
||||
from rest_framework.generics import ListAPIView
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.request import Request
|
||||
from django.conf import settings
|
||||
|
||||
from assets.api.mixin import SerializeToTreeNodeMixin
|
||||
from common.utils import get_logger
|
||||
from perms.pagination import GrantedAssetLimitOffsetPagination
|
||||
from assets.models import Asset, Node, FavoriteAsset
|
||||
from perms import serializers
|
||||
from perms.utils.asset.user_permission import (
|
||||
get_node_all_granted_assets, get_user_direct_granted_assets,
|
||||
get_user_granted_all_assets
|
||||
)
|
||||
from .mixin import ForAdminMixin, ForUserMixin
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class UserDirectGrantedAssetsApi(ListAPIView):
|
||||
"""
|
||||
用户直接授权的资产的列表，也就是授权规则上直接授权的资产，并非是来自节点的
|
||||
"""
|
||||
serializer_class = serializers.AssetGrantedSerializer
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
filterset_fields = ['hostname', 'ip', 'id', 'comment']
|
||||
search_fields = ['hostname', 'ip', 'comment']
|
||||
|
||||
def get_queryset(self):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
user = self.user
|
||||
assets = get_user_direct_granted_assets(user)\
|
||||
.prefetch_related('platform')\
|
||||
.only(*self.only_fields)
|
||||
return assets
|
||||
|
||||
|
||||
class UserFavoriteGrantedAssetsApi(ListAPIView):
|
||||
serializer_class = serializers.AssetGrantedSerializer
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
filterset_fields = ['hostname', 'ip', 'id', 'comment']
|
||||
search_fields = ['hostname', 'ip', 'comment']
|
||||
|
||||
def get_queryset(self):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
user = self.user
|
||||
assets = FavoriteAsset.get_user_favorite_assets(user)\
|
||||
.prefetch_related('platform')\
|
||||
.only(*self.only_fields)
|
||||
return assets
|
||||
|
||||
|
||||
class AssetsAsTreeMixin(SerializeToTreeNodeMixin):
|
||||
"""
|
||||
将 资产 序列化成树的结构返回
|
||||
"""
|
||||
def list(self, request: Request, *args, **kwargs):
|
||||
queryset = self.filter_queryset(self.get_queryset())
|
||||
if request.query_params.get('search'):
|
||||
# 如果用户搜索的条件不精准，会导致返回大量的无意义数据。
|
||||
# 这里限制一下返回数据的最大条数
|
||||
queryset = queryset[:999]
|
||||
data = self.serialize_assets(queryset, None)
|
||||
return Response(data=data)
|
||||
|
||||
|
||||
class UserDirectGrantedAssetsForAdminApi(ForAdminMixin, UserDirectGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyDirectGrantedAssetsApi(ForUserMixin, UserDirectGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class UserFavoriteGrantedAssetsForAdminApi(ForAdminMixin, UserFavoriteGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyFavoriteGrantedAssetsApi(ForUserMixin, UserFavoriteGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class UserDirectGrantedAssetsAsTreeForAdminApi(ForAdminMixin, AssetsAsTreeMixin, UserDirectGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyUngroupAssetsAsTreeApi(ForUserMixin, AssetsAsTreeMixin, UserDirectGrantedAssetsApi):
|
||||
def get_queryset(self):
|
||||
queryset = super().get_queryset()
|
||||
if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
queryset = queryset.none()
|
||||
return queryset
|
||||
|
||||
|
||||
class UserAllGrantedAssetsApi(ForAdminMixin, ListAPIView):
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
serializer_class = serializers.AssetGrantedSerializer
|
||||
filterset_fields = ['hostname', 'ip', 'id', 'comment']
|
||||
search_fields = ['hostname', 'ip', 'comment']
|
||||
|
||||
def get_queryset(self):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
queryset = get_user_granted_all_assets(self.user)
|
||||
queryset = queryset.prefetch_related('platform')
|
||||
return queryset.only(*self.only_fields)
|
||||
|
||||
|
||||
class MyAllGrantedAssetsApi(ForUserMixin, UserAllGrantedAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyAllAssetsAsTreeApi(ForUserMixin, AssetsAsTreeMixin, UserAllGrantedAssetsApi):
|
||||
search_fields = ['hostname', 'ip']
|
||||
|
||||
|
||||
class UserGrantedNodeAssetsApi(UserNodeGrantStatusDispatchMixin, ListAPIView):
|
||||
serializer_class = serializers.AssetGrantedSerializer
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
filterset_fields = ['hostname', 'ip', 'id', 'comment']
|
||||
search_fields = ['hostname', 'ip', 'comment']
|
||||
pagination_class = GrantedAssetLimitOffsetPagination
|
||||
pagination_node: Node
|
||||
|
||||
def get_queryset(self):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
node_id = self.kwargs.get("node_id")
|
||||
node = Node.objects.get(id=node_id)
|
||||
self.pagination_node = node
|
||||
return self.dispatch_get_data(node.key, self.user)
|
||||
|
||||
def get_data_on_node_direct_granted(self, key):
|
||||
# 如果这个节点是直接授权的(或者说祖先节点直接授权的), 获取下面的所有资产
|
||||
return Node.get_node_all_assets_by_key_v2(key)
|
||||
|
||||
def get_data_on_node_indirect_granted(self, key):
|
||||
self.pagination_node = self.get_mapping_node_by_key(key, self.user)
|
||||
return get_node_all_granted_assets(self.user, key)
|
||||
|
||||
def get_data_on_node_not_granted(self, key):
|
||||
return Asset.objects.none()
|
||||
|
||||
|
||||
class UserGrantedNodeAssetsForAdminApi(ForAdminMixin, UserGrantedNodeAssetsApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodeAssetsApi(ForUserMixin, UserGrantedNodeAssetsApi):
|
||||
pass
|
||||
|
|
@ -0,0 +1 @@
|
|||
from .views import *
|
||||
|
|
@ -0,0 +1,127 @@
|
|||
from rest_framework.response import Response
|
||||
from rest_framework.request import Request
|
||||
|
||||
from users.models import User
|
||||
from assets.api.mixin import SerializeToTreeNodeMixin
|
||||
from common.utils import get_logger
|
||||
from perms.pagination import NodeGrantedAssetPagination, AllGrantedAssetPagination
|
||||
from assets.models import Asset, Node
|
||||
from perms import serializers
|
||||
from perms.utils.asset.user_permission import UserGrantedAssetsQueryUtils, QuerySetStage
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# 获取数据的 ------------------------------------------------------------
|
||||
|
||||
class UserDirectGrantedAssetsQuerysetMixin:
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
user: User
|
||||
|
||||
def get_queryset(self):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
user = self.user
|
||||
assets = UserGrantedAssetsQueryUtils(user) \
|
||||
.get_direct_granted_assets() \
|
||||
.prefetch_related('platform') \
|
||||
.only(*self.only_fields)
|
||||
return assets
|
||||
|
||||
|
||||
class UserAllGrantedAssetsQuerysetMixin:
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
pagination_class = AllGrantedAssetPagination
|
||||
user: User
|
||||
|
||||
def get_union_queryset(self, qs_stage: QuerySetStage):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
qs_stage.prefetch_related('platform').only(*self.only_fields)
|
||||
queryset = UserGrantedAssetsQueryUtils(self.user) \
|
||||
.get_all_granted_assets(qs_stage)
|
||||
return queryset
|
||||
|
||||
|
||||
class UserFavoriteGrantedAssetsMixin:
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
user: User
|
||||
|
||||
def get_union_queryset(self, qs_stage: QuerySetStage):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
user = self.user
|
||||
qs_stage.prefetch_related('platform').only(*self.only_fields)
|
||||
utils = UserGrantedAssetsQueryUtils(user)
|
||||
assets = utils.get_favorite_assets(qs_stage=qs_stage)
|
||||
return assets
|
||||
|
||||
|
||||
class UserGrantedNodeAssetsMixin:
|
||||
only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
|
||||
pagination_class = NodeGrantedAssetPagination
|
||||
pagination_node: Node
|
||||
user: User
|
||||
|
||||
def get_union_queryset(self, qs_stage: QuerySetStage):
|
||||
if getattr(self, 'swagger_fake_view', False):
|
||||
return Asset.objects.none()
|
||||
node_id = self.kwargs.get("node_id")
|
||||
qs_stage.prefetch_related('platform').only(*self.only_fields)
|
||||
node, assets = UserGrantedAssetsQueryUtils(self.user).get_node_all_assets(
|
||||
node_id, qs_stage=qs_stage
|
||||
)
|
||||
self.pagination_node = node
|
||||
return assets
|
||||
|
||||
|
||||
# 控制格式的 ----------------------------------------------------
|
||||
|
||||
class AssetsUnionQuerysetMixin:
|
||||
def get_queryset_union_prefer(self):
|
||||
if hasattr(self, 'get_union_queryset'):
|
||||
# 为了支持 union 查询
|
||||
queryset = Asset.objects.all().distinct()
|
||||
queryset = self.filter_queryset(queryset)
|
||||
qs_stage = QuerySetStage()
|
||||
qs_stage.and_with_queryset(queryset)
|
||||
queryset = self.get_union_queryset(qs_stage)
|
||||
else:
|
||||
queryset = self.filter_queryset(self.get_queryset())
|
||||
return queryset
|
||||
|
||||
|
||||
class AssetsSerializerFormatMixin(AssetsUnionQuerysetMixin):
|
||||
serializer_class = serializers.AssetGrantedSerializer
|
||||
filterset_fields = ['hostname', 'ip', 'id', 'comment']
|
||||
search_fields = ['hostname', 'ip', 'comment']
|
||||
|
||||
def list(self, request, *args, **kwargs):
|
||||
queryset = self.get_queryset_union_prefer()
|
||||
|
||||
page = self.paginate_queryset(queryset)
|
||||
if page is not None:
|
||||
serializer = self.get_serializer(page, many=True)
|
||||
return self.get_paginated_response(serializer.data)
|
||||
|
||||
serializer = self.get_serializer(queryset, many=True)
|
||||
return Response(serializer.data)
|
||||
|
||||
|
||||
class AssetsTreeFormatMixin(AssetsUnionQuerysetMixin, SerializeToTreeNodeMixin):
|
||||
"""
|
||||
将 资产 序列化成树的结构返回
|
||||
"""
|
||||
|
||||
def list(self, request: Request, *args, **kwargs):
|
||||
queryset = self.get_queryset_union_prefer()
|
||||
|
||||
if request.query_params.get('search'):
|
||||
# 如果用户搜索的条件不精准，会导致返回大量的无意义数据。
|
||||
# 这里限制一下返回数据的最大条数
|
||||
queryset = queryset[:999]
|
||||
data = self.serialize_assets(queryset, None)
|
||||
return Response(data=data)
|
||||
|
||||
# def get_serializer_class(self):
|
||||
# return EmptySerializer
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
from rest_framework.generics import ListAPIView
|
||||
from django.conf import settings
|
||||
|
||||
from common.utils import get_logger
|
||||
from ..mixin import RoleAdminMixin, RoleUserMixin
|
||||
from .mixin import (
|
||||
UserAllGrantedAssetsQuerysetMixin, UserDirectGrantedAssetsQuerysetMixin, UserFavoriteGrantedAssetsMixin,
|
||||
UserGrantedNodeAssetsMixin, AssetsSerializerFormatMixin, AssetsTreeFormatMixin,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
'UserDirectGrantedAssetsForAdminApi', 'MyDirectGrantedAssetsApi', 'UserFavoriteGrantedAssetsForAdminApi',
|
||||
'MyFavoriteGrantedAssetsApi', 'UserDirectGrantedAssetsAsTreeForAdminApi', 'MyUngroupAssetsAsTreeApi',
|
||||
'UserAllGrantedAssetsApi', 'MyAllGrantedAssetsApi', 'MyAllAssetsAsTreeApi', 'UserGrantedNodeAssetsForAdminApi',
|
||||
'MyGrantedNodeAssetsApi',
|
||||
]
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class UserDirectGrantedAssetsForAdminApi(UserDirectGrantedAssetsQuerysetMixin,
|
||||
RoleAdminMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyDirectGrantedAssetsApi(UserDirectGrantedAssetsQuerysetMixin,
|
||||
RoleUserMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class UserFavoriteGrantedAssetsForAdminApi(UserFavoriteGrantedAssetsMixin,
|
||||
RoleAdminMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyFavoriteGrantedAssetsApi(UserFavoriteGrantedAssetsMixin,
|
||||
RoleUserMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class UserDirectGrantedAssetsAsTreeForAdminApi(UserDirectGrantedAssetsQuerysetMixin,
|
||||
RoleAdminMixin,
|
||||
AssetsTreeFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyUngroupAssetsAsTreeApi(UserDirectGrantedAssetsQuerysetMixin,
|
||||
RoleUserMixin,
|
||||
AssetsTreeFormatMixin,
|
||||
ListAPIView):
|
||||
def get_queryset(self):
|
||||
queryset = super().get_queryset()
|
||||
if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
queryset = queryset.none()
|
||||
return queryset
|
||||
|
||||
|
||||
class UserAllGrantedAssetsApi(UserAllGrantedAssetsQuerysetMixin,
|
||||
RoleAdminMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyAllGrantedAssetsApi(UserAllGrantedAssetsQuerysetMixin,
|
||||
RoleUserMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyAllAssetsAsTreeApi(UserAllGrantedAssetsQuerysetMixin,
|
||||
RoleUserMixin,
|
||||
AssetsTreeFormatMixin,
|
||||
ListAPIView):
|
||||
search_fields = ['hostname', 'ip']
|
||||
|
||||
|
||||
class UserGrantedNodeAssetsForAdminApi(UserGrantedNodeAssetsMixin,
|
||||
RoleAdminMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodeAssetsApi(UserGrantedNodeAssetsMixin,
|
||||
RoleUserMixin,
|
||||
AssetsSerializerFormatMixin,
|
||||
ListAPIView):
|
||||
pass
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import abc
|
||||
from django.conf import settings
|
||||
from rest_framework.generics import (
|
||||
ListAPIView
|
||||
)
|
||||
|
|
@ -10,16 +9,11 @@ from rest_framework.request import Request
|
|||
|
||||
from assets.api.mixin import SerializeToTreeNodeMixin
|
||||
from common.utils import get_logger
|
||||
from .mixin import ForAdminMixin, ForUserMixin, UserNodeGrantStatusDispatchMixin
|
||||
from perms.hands import Node, User
|
||||
from .mixin import RoleAdminMixin, RoleUserMixin
|
||||
from perms.hands import User
|
||||
from perms import serializers
|
||||
from perms.utils.asset.user_permission import (
|
||||
get_indirect_granted_node_children,
|
||||
get_user_granted_nodes_list_via_mapping_node,
|
||||
get_top_level_granted_nodes,
|
||||
rebuild_user_tree_if_need, get_favorite_node,
|
||||
get_ungrouped_node
|
||||
)
|
||||
|
||||
from perms.utils.asset.user_permission import UserGrantedNodesQueryUtils
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
@ -61,7 +55,6 @@ class BaseGrantedNodeApi(_GrantedNodeStructApi, metaclass=abc.ABCMeta):
|
|||
serializer_class = serializers.NodeGrantedSerializer
|
||||
|
||||
def list(self, request, *args, **kwargs):
|
||||
rebuild_user_tree_if_need(request, self.user)
|
||||
nodes = self.get_nodes()
|
||||
serializer = self.get_serializer(nodes, many=True)
|
||||
return Response(serializer.data)
|
||||
|
|
@ -73,7 +66,6 @@ class BaseNodeChildrenApi(NodeChildrenMixin, BaseGrantedNodeApi, metaclass=abc.A
|
|||
|
||||
class BaseGrantedNodeAsTreeApi(SerializeToTreeNodeMixin, _GrantedNodeStructApi, metaclass=abc.ABCMeta):
|
||||
def list(self, request: Request, *args, **kwargs):
|
||||
rebuild_user_tree_if_need(request, self.user)
|
||||
nodes = self.get_nodes()
|
||||
nodes = self.serialize_nodes(nodes, with_asset_amount=True)
|
||||
return Response(data=nodes)
|
||||
|
|
@ -83,30 +75,16 @@ class BaseNodeChildrenAsTreeApi(NodeChildrenMixin, BaseGrantedNodeAsTreeApi, met
|
|||
pass
|
||||
|
||||
|
||||
class UserGrantedNodeChildrenMixin(UserNodeGrantStatusDispatchMixin):
|
||||
class UserGrantedNodeChildrenMixin:
|
||||
user: User
|
||||
request: Request
|
||||
|
||||
def get_children(self):
|
||||
user = self.user
|
||||
key = self.request.query_params.get('key')
|
||||
|
||||
if not key:
|
||||
nodes = list(get_top_level_granted_nodes(user))
|
||||
else:
|
||||
nodes = self.dispatch_get_data(key, user)
|
||||
nodes = UserGrantedNodesQueryUtils(user).get_node_children(key)
|
||||
return nodes
|
||||
|
||||
def get_data_on_node_direct_granted(self, key):
|
||||
return Node.objects.filter(parent_key=key)
|
||||
|
||||
def get_data_on_node_indirect_granted(self, key):
|
||||
nodes = get_indirect_granted_node_children(self.user, key)
|
||||
return nodes
|
||||
|
||||
def get_data_on_node_not_granted(self, key):
|
||||
return Node.objects.none()
|
||||
|
||||
|
||||
class UserGrantedNodesMixin:
|
||||
"""
|
||||
|
|
@ -115,41 +93,38 @@ class UserGrantedNodesMixin:
|
|||
user: User
|
||||
|
||||
def get_nodes(self):
|
||||
nodes = []
|
||||
if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
nodes.append(get_ungrouped_node(self.user))
|
||||
nodes.append(get_favorite_node(self.user))
|
||||
nodes.extend(get_user_granted_nodes_list_via_mapping_node(self.user))
|
||||
utils = UserGrantedNodesQueryUtils(self.user)
|
||||
nodes = utils.get_whole_tree_nodes()
|
||||
return nodes
|
||||
|
||||
|
||||
# ------------------------------------------
|
||||
# 最终的 api
|
||||
class UserGrantedNodeChildrenForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi):
|
||||
class UserGrantedNodeChildrenForAdminApi(RoleAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodeChildrenApi(ForUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi):
|
||||
class MyGrantedNodeChildrenApi(RoleUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi):
|
||||
pass
|
||||
|
||||
|
||||
class UserGrantedNodeChildrenAsTreeForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi):
|
||||
class UserGrantedNodeChildrenAsTreeForAdminApi(RoleAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodeChildrenAsTreeApi(ForUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi):
|
||||
class MyGrantedNodeChildrenAsTreeApi(RoleUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi):
|
||||
pass
|
||||
|
||||
|
||||
class UserGrantedNodesForAdminApi(ForAdminMixin, UserGrantedNodesMixin, BaseGrantedNodeApi):
|
||||
class UserGrantedNodesForAdminApi(RoleAdminMixin, UserGrantedNodesMixin, BaseGrantedNodeApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodesApi(ForUserMixin, UserGrantedNodesMixin, BaseGrantedNodeApi):
|
||||
class MyGrantedNodesApi(RoleUserMixin, UserGrantedNodesMixin, BaseGrantedNodeApi):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodesAsTreeApi(ForUserMixin, UserGrantedNodesMixin, BaseGrantedNodeAsTreeApi):
|
||||
class MyGrantedNodesAsTreeApi(RoleUserMixin, UserGrantedNodesMixin, BaseGrantedNodeAsTreeApi):
|
||||
pass
|
||||
|
||||
# ------------------------------------------
|
||||
|
|
|
|||
|
|
@ -1,29 +1,23 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
from itertools import chain
|
||||
|
||||
from rest_framework.generics import ListAPIView
|
||||
from rest_framework.request import Request
|
||||
from rest_framework.response import Response
|
||||
from django.db.models import F, Value, CharField, Q
|
||||
from django.db.models import F, Value, CharField
|
||||
from django.conf import settings
|
||||
|
||||
from common.utils.common import timeit
|
||||
from orgs.utils import tmp_to_root_org
|
||||
from common.permissions import IsValidUser
|
||||
from common.utils import get_logger, get_object_or_none
|
||||
from .mixin import UserNodeGrantStatusDispatchMixin, ForUserMixin, ForAdminMixin
|
||||
from .mixin import RoleUserMixin, RoleAdminMixin
|
||||
from perms.utils.asset.user_permission import (
|
||||
get_indirect_granted_node_children, UNGROUPED_NODE_KEY, FAVORITE_NODE_KEY,
|
||||
get_user_direct_granted_assets, get_top_level_granted_nodes,
|
||||
get_user_granted_nodes_list_via_mapping_node,
|
||||
get_user_granted_all_assets, rebuild_user_tree_if_need,
|
||||
get_user_all_assetpermissions_id, get_favorite_node,
|
||||
get_ungrouped_node, compute_tmp_mapping_node_from_perm,
|
||||
TMP_GRANTED_FIELD, count_direct_granted_node_assets,
|
||||
count_node_all_granted_assets
|
||||
UserGrantedTreeBuildUtils, get_user_all_asset_perm_ids,
|
||||
UserGrantedNodesQueryUtils, UserGrantedAssetsQueryUtils,
|
||||
QuerySetStage,
|
||||
)
|
||||
from perms.models import AssetPermission
|
||||
from assets.models import Asset, FavoriteAsset
|
||||
from perms.models import AssetPermission, PermNode
|
||||
from assets.models import Asset
|
||||
from assets.api import SerializeToTreeNodeMixin
|
||||
from perms.hands import Node
|
||||
|
||||
|
|
@ -33,76 +27,45 @@ logger = get_logger(__name__)
|
|||
class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView):
|
||||
permission_classes = (IsValidUser,)
|
||||
|
||||
def add_ungrouped_resource(self, data: list, user, asset_perms_id):
|
||||
@timeit
|
||||
def add_ungrouped_resource(self, data: list, nodes_query_utils, assets_query_utils):
|
||||
if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
return
|
||||
ungrouped_node = nodes_query_utils.get_ungrouped_node()
|
||||
|
||||
ungrouped_node = get_ungrouped_node(user, asset_perms_id=asset_perms_id)
|
||||
direct_granted_assets = get_user_direct_granted_assets(
|
||||
user, asset_perms_id=asset_perms_id
|
||||
).annotate(
|
||||
direct_granted_assets = assets_query_utils.get_direct_granted_assets().annotate(
|
||||
parent_key=Value(ungrouped_node.key, output_field=CharField())
|
||||
).prefetch_related('platform')
|
||||
|
||||
data.extend(self.serialize_nodes([ungrouped_node], with_asset_amount=True))
|
||||
data.extend(self.serialize_assets(direct_granted_assets))
|
||||
|
||||
def add_favorite_resource(self, data: list, user, asset_perms_id):
|
||||
favorite_node = get_favorite_node(user, asset_perms_id)
|
||||
favorite_assets = FavoriteAsset.get_user_favorite_assets(
|
||||
user, asset_perms_id=asset_perms_id
|
||||
).annotate(
|
||||
@timeit
|
||||
def add_favorite_resource(self, data: list, nodes_query_utils, assets_query_utils):
|
||||
favorite_node = nodes_query_utils.get_favorite_node()
|
||||
|
||||
qs_state = QuerySetStage().annotate(
|
||||
parent_key=Value(favorite_node.key, output_field=CharField())
|
||||
).prefetch_related('platform')
|
||||
favorite_assets = assets_query_utils.get_favorite_assets(qs_stage=qs_state, only=())
|
||||
|
||||
data.extend(self.serialize_nodes([favorite_node], with_asset_amount=True))
|
||||
data.extend(self.serialize_assets(favorite_assets))
|
||||
|
||||
@timeit
|
||||
def add_node_filtered_by_system_user(self, data: list, user, asset_perms_id):
|
||||
tmp_nodes = compute_tmp_mapping_node_from_perm(user, asset_perms_id=asset_perms_id)
|
||||
granted_nodes_key = []
|
||||
for _node in tmp_nodes:
|
||||
_granted = getattr(_node, TMP_GRANTED_FIELD, False)
|
||||
if not _granted:
|
||||
if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
assets_amount = count_direct_granted_node_assets(user, _node.key, asset_perms_id)
|
||||
else:
|
||||
assets_amount = count_node_all_granted_assets(user, _node.key, asset_perms_id)
|
||||
_node.assets_amount = assets_amount
|
||||
else:
|
||||
granted_nodes_key.append(_node.key)
|
||||
utils = UserGrantedTreeBuildUtils(user, asset_perms_id)
|
||||
nodes = utils.get_whole_tree_nodes()
|
||||
data.extend(self.serialize_nodes(nodes, with_asset_amount=True))
|
||||
|
||||
# 查询他们的子节点
|
||||
q = Q()
|
||||
for _key in granted_nodes_key:
|
||||
q |= Q(key__startswith=f'{_key}:')
|
||||
def add_assets(self, data: list, assets_query_utils: UserGrantedAssetsQueryUtils):
|
||||
qs_stage = QuerySetStage().annotate(parent_key=F('nodes__key')).prefetch_related('platform')
|
||||
|
||||
if q:
|
||||
descendant_nodes = Node.objects.filter(q).distinct()
|
||||
else:
|
||||
descendant_nodes = Node.objects.none()
|
||||
|
||||
data.extend(self.serialize_nodes(chain(tmp_nodes, descendant_nodes), with_asset_amount=True))
|
||||
|
||||
def add_assets(self, data: list, user, asset_perms_id):
|
||||
if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
|
||||
all_assets = get_user_granted_all_assets(
|
||||
user,
|
||||
via_mapping_node=False,
|
||||
include_direct_granted_assets=False,
|
||||
asset_perms_id=asset_perms_id
|
||||
)
|
||||
all_assets = assets_query_utils.get_direct_granted_nodes_assets(qs_stage=qs_stage)
|
||||
else:
|
||||
all_assets = get_user_granted_all_assets(
|
||||
user,
|
||||
via_mapping_node=False,
|
||||
include_direct_granted_assets=True,
|
||||
asset_perms_id=asset_perms_id
|
||||
)
|
||||
all_assets = assets_query_utils.get_all_granted_assets(qs_stage=qs_stage)
|
||||
|
||||
all_assets = all_assets.annotate(
|
||||
parent_key=F('nodes__key')
|
||||
).prefetch_related('platform')
|
||||
data.extend(self.serialize_assets(all_assets))
|
||||
|
||||
@tmp_to_root_org()
|
||||
|
|
@ -117,7 +80,7 @@ class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView):
|
|||
|
||||
user = request.user
|
||||
data = []
|
||||
asset_perms_id = get_user_all_assetpermissions_id(user)
|
||||
asset_perms_id = get_user_all_asset_perm_ids(user)
|
||||
|
||||
system_user_id = request.query_params.get('system_user')
|
||||
if system_user_id:
|
||||
|
|
@ -125,89 +88,72 @@ class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView):
|
|||
id__in=asset_perms_id, system_users__id=system_user_id, actions__gt=0
|
||||
).values_list('id', flat=True).distinct())
|
||||
|
||||
self.add_ungrouped_resource(data, user, asset_perms_id)
|
||||
self.add_favorite_resource(data, user, asset_perms_id)
|
||||
nodes_query_utils = UserGrantedNodesQueryUtils(user, asset_perms_id)
|
||||
assets_query_utils = UserGrantedAssetsQueryUtils(user, asset_perms_id)
|
||||
|
||||
self.add_ungrouped_resource(data, nodes_query_utils, assets_query_utils)
|
||||
self.add_favorite_resource(data, nodes_query_utils, assets_query_utils)
|
||||
|
||||
if system_user_id:
|
||||
# 有系统用户筛选的需要重新计算树结构
|
||||
self.add_node_filtered_by_system_user(data, user, asset_perms_id)
|
||||
else:
|
||||
rebuild_user_tree_if_need(request, user)
|
||||
all_nodes = get_user_granted_nodes_list_via_mapping_node(user)
|
||||
all_nodes = nodes_query_utils.get_whole_tree_nodes(with_special=False)
|
||||
data.extend(self.serialize_nodes(all_nodes, with_asset_amount=True))
|
||||
|
||||
self.add_assets(data, user, asset_perms_id)
|
||||
self.add_assets(data, assets_query_utils)
|
||||
return Response(data=data)
|
||||
|
||||
|
||||
class GrantedNodeChildrenWithAssetsAsTreeApiMixin(UserNodeGrantStatusDispatchMixin,
|
||||
SerializeToTreeNodeMixin,
|
||||
class GrantedNodeChildrenWithAssetsAsTreeApiMixin(SerializeToTreeNodeMixin,
|
||||
ListAPIView):
|
||||
"""
|
||||
带资产的授权树
|
||||
"""
|
||||
user: None
|
||||
|
||||
def get_data_on_node_direct_granted(self, key):
|
||||
nodes = Node.objects.filter(parent_key=key)
|
||||
assets = Asset.org_objects.filter(nodes__key=key).distinct()
|
||||
assets = assets.prefetch_related('platform')
|
||||
return nodes, assets
|
||||
def ensure_key(self):
|
||||
key = self.request.query_params.get('key', None)
|
||||
id = self.request.query_params.get('id', None)
|
||||
|
||||
def get_data_on_node_indirect_granted(self, key):
|
||||
user = self.user
|
||||
asset_perms_id = get_user_all_assetpermissions_id(user)
|
||||
if key is not None:
|
||||
return key
|
||||
|
||||
nodes = get_indirect_granted_node_children(user, key)
|
||||
|
||||
assets = Asset.org_objects.filter(
|
||||
nodes__key=key,
|
||||
).filter(
|
||||
granted_by_permissions__id__in=asset_perms_id
|
||||
).distinct()
|
||||
assets = assets.prefetch_related('platform')
|
||||
return nodes, assets
|
||||
|
||||
def get_data_on_node_not_granted(self, key):
|
||||
return Node.objects.none(), Asset.objects.none()
|
||||
|
||||
def get_data(self, key, user):
|
||||
assets, nodes = [], []
|
||||
if not key:
|
||||
root_nodes = get_top_level_granted_nodes(user)
|
||||
nodes.extend(root_nodes)
|
||||
elif key == UNGROUPED_NODE_KEY:
|
||||
assets = get_user_direct_granted_assets(user)
|
||||
assets = assets.prefetch_related('platform')
|
||||
elif key == FAVORITE_NODE_KEY:
|
||||
assets = FavoriteAsset.get_user_favorite_assets(user)
|
||||
else:
|
||||
nodes, assets = self.dispatch_get_data(key, user)
|
||||
return nodes, assets
|
||||
|
||||
def id2key_if_have(self):
|
||||
id = self.request.query_params.get('id')
|
||||
if id is not None:
|
||||
node = get_object_or_none(Node, id=id)
|
||||
if node:
|
||||
return node.key
|
||||
node = get_object_or_none(Node, id=id)
|
||||
if node:
|
||||
return node.key
|
||||
|
||||
def list(self, request: Request, *args, **kwargs):
|
||||
key = self.request.query_params.get('key')
|
||||
if key is None:
|
||||
key = self.id2key_if_have()
|
||||
user = self.user
|
||||
key = self.ensure_key()
|
||||
|
||||
nodes_query_utils = UserGrantedNodesQueryUtils(user)
|
||||
assets_query_utils = UserGrantedAssetsQueryUtils(user)
|
||||
|
||||
nodes = PermNode.objects.none()
|
||||
assets = Asset.objects.none()
|
||||
|
||||
if not key:
|
||||
nodes = nodes_query_utils.get_top_level_nodes()
|
||||
elif key == PermNode.UNGROUPED_NODE_KEY:
|
||||
assets = assets_query_utils.get_ungroup_assets()
|
||||
elif key == PermNode.FAVORITE_NODE_KEY:
|
||||
assets = assets_query_utils.get_favorite_assets()
|
||||
else:
|
||||
nodes = nodes_query_utils.get_node_children(key)
|
||||
assets = assets_query_utils.get_node_assets(key)
|
||||
assets = assets.prefetch_related('platform')
|
||||
|
||||
user = self.user
|
||||
rebuild_user_tree_if_need(request, user)
|
||||
nodes, assets = self.get_data(key, user)
|
||||
|
||||
tree_nodes = self.serialize_nodes(nodes, with_asset_amount=True)
|
||||
tree_assets = self.serialize_assets(assets, key)
|
||||
return Response(data=[*tree_nodes, *tree_assets])
|
||||
|
||||
|
||||
class UserGrantedNodeChildrenWithAssetsAsTreeApi(ForAdminMixin, GrantedNodeChildrenWithAssetsAsTreeApiMixin):
|
||||
class UserGrantedNodeChildrenWithAssetsAsTreeApi(RoleAdminMixin, GrantedNodeChildrenWithAssetsAsTreeApiMixin):
|
||||
pass
|
||||
|
||||
|
||||
class MyGrantedNodeChildrenWithAssetsAsTreeApi(ForUserMixin, GrantedNodeChildrenWithAssetsAsTreeApiMixin):
|
||||
class MyGrantedNodeChildrenWithAssetsAsTreeApi(RoleUserMixin, GrantedNodeChildrenWithAssetsAsTreeApiMixin):
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
from rest_framework import generics
|
||||
from django.db.models import Q
|
||||
from django.utils.decorators import method_decorator
|
||||
|
||||
from assets.models import SystemUser
|
||||
from common.permissions import IsValidUser
|
||||
from orgs.utils import tmp_to_root_org
|
||||
from perms.utils.asset.user_permission import get_user_all_asset_perm_ids
|
||||
from .. import serializers
|
||||
|
||||
|
||||
|
|
@ -16,9 +16,9 @@ class SystemUserPermission(generics.ListAPIView):
|
|||
def get_queryset(self):
|
||||
user = self.request.user
|
||||
|
||||
asset_perms_id = get_user_all_asset_perm_ids(user)
|
||||
queryset = SystemUser.objects.filter(
|
||||
Q(granted_by_permissions__users=user) |
|
||||
Q(granted_by_permissions__user_groups__users=user)
|
||||
granted_by_permissions__id__in=asset_perms_id
|
||||
).distinct()
|
||||
|
||||
return queryset
|
||||
|
|
|
|||
|
|
@ -1,47 +0,0 @@
|
|||
from django.utils.crypto import get_random_string
|
||||
from perms.utils import rebuild_user_mapping_nodes_if_need_with_lock
|
||||
|
||||
from common.thread_pools import SingletonThreadPoolExecutor
|
||||
from common.utils import get_logger
|
||||
from perms.models import RebuildUserTreeTask
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class Executor(SingletonThreadPoolExecutor):
|
||||
pass
|
||||
|
||||
|
||||
executor = Executor()
|
||||
|
||||
|
||||
def run_mapping_node_tasks():
|
||||
failed_user_ids = []
|
||||
|
||||
ident = get_random_string()
|
||||
logger.debug(f'[{ident}]mapping_node_tasks running')
|
||||
|
||||
while True:
|
||||
task = RebuildUserTreeTask.objects.exclude(
|
||||
user_id__in=failed_user_ids
|
||||
).first()
|
||||
|
||||
if task is None:
|
||||
break
|
||||
|
||||
user = task.user
|
||||
try:
|
||||
rebuild_user_mapping_nodes_if_need_with_lock(user)
|
||||
except:
|
||||
logger.exception(f'[{ident}]mapping_node_tasks_exception')
|
||||
failed_user_ids.append(user.id)
|
||||
|
||||
logger.debug(f'[{ident}]mapping_node_tasks finished')
|
||||
|
||||
|
||||
def submit_update_mapping_node_task():
|
||||
executor.submit(run_mapping_node_tasks)
|
||||
|
||||
|
||||
def submit_update_mapping_node_task_for_user(user):
|
||||
executor.submit(rebuild_user_mapping_nodes_if_need_with_lock, user)
|
||||
|
|
@ -0,0 +1,11 @@
|
|||
from common.utils.lock import DistributedLock
|
||||
|
||||
|
||||
class UserGrantedTreeRebuildLock(DistributedLock):
|
||||
name_template = 'perms.user.asset.node.tree.rebuid.<org_id:{org_id}>.<user_id:{user_id}>'
|
||||
|
||||
def __init__(self, org_id, user_id):
|
||||
name = self.name_template.format(
|
||||
org_id=org_id, user_id=user_id
|
||||
)
|
||||
super().__init__(name=name)
|
||||
|
|
@ -1,19 +1,6 @@
|
|||
# Generated by Django 2.2.13 on 2020-08-21 08:20
|
||||
|
||||
from django.db import migrations
|
||||
from perms.tasks import dispatch_mapping_node_tasks
|
||||
|
||||
|
||||
def start_build_users_perm_tree_task(apps, schema_editor):
|
||||
User = apps.get_model('users', 'User')
|
||||
RebuildUserTreeTask = apps.get_model('perms', 'RebuildUserTreeTask')
|
||||
|
||||
user_ids = User.objects.all().values_list('id', flat=True).distinct()
|
||||
RebuildUserTreeTask.objects.bulk_create(
|
||||
[RebuildUserTreeTask(user_id=i) for i in user_ids]
|
||||
)
|
||||
|
||||
dispatch_mapping_node_tasks.delay()
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
|
@ -23,5 +10,4 @@ class Migration(migrations.Migration):
|
|||
]
|
||||
|
||||
operations = [
|
||||
migrations.RunPython(start_build_users_perm_tree_task)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,65 @@
|
|||
# Generated by Django 3.1 on 2021-02-04 09:49
|
||||
|
||||
import assets.models.node
|
||||
from django.conf import settings
|
||||
from django.db import migrations, models
|
||||
import django.db.models.deletion
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('assets', '0066_remove_node_assets_amount'),
|
||||
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
|
||||
('perms', '0017_auto_20210104_0435'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name='UserAssetGrantedTreeNodeRelation',
|
||||
fields=[
|
||||
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('created_by', models.CharField(blank=True, max_length=32, null=True, verbose_name='Created by')),
|
||||
('updated_by', models.CharField(blank=True, max_length=32, null=True, verbose_name='Updated by')),
|
||||
('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),
|
||||
('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')),
|
||||
('org_id', models.CharField(blank=True, db_index=True, default='', max_length=36, verbose_name='Organization')),
|
||||
('node_key', models.CharField(db_index=True, max_length=64, verbose_name='Key')),
|
||||
('node_parent_key', models.CharField(db_index=True, default='', max_length=64, verbose_name='Parent key')),
|
||||
('node_from', models.CharField(choices=[('granted', 'Direct node granted'), ('child', 'Have children node'), ('asset', 'Direct asset granted')], db_index=True, max_length=16)),
|
||||
('node_assets_amount', models.IntegerField(default=0)),
|
||||
('node', models.ForeignKey(db_constraint=False, default=None, on_delete=django.db.models.deletion.CASCADE, related_name='granted_node_rels', to='assets.node')),
|
||||
('user', models.ForeignKey(db_constraint=False, on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
|
||||
],
|
||||
options={
|
||||
'abstract': False,
|
||||
},
|
||||
bases=(assets.models.node.FamilyMixin, models.Model),
|
||||
),
|
||||
migrations.RemoveField(
|
||||
model_name='usergrantedmappingnode',
|
||||
name='node',
|
||||
),
|
||||
migrations.RemoveField(
|
||||
model_name='usergrantedmappingnode',
|
||||
name='user',
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name='PermNode',
|
||||
fields=[
|
||||
],
|
||||
options={
|
||||
'ordering': [],
|
||||
'proxy': True,
|
||||
'indexes': [],
|
||||
'constraints': [],
|
||||
},
|
||||
bases=('assets.node',),
|
||||
),
|
||||
migrations.DeleteModel(
|
||||
name='RebuildUserTreeTask',
|
||||
),
|
||||
migrations.DeleteModel(
|
||||
name='UserGrantedMappingNode',
|
||||
),
|
||||
]
|
||||
|
|
@ -2,7 +2,10 @@ import logging
|
|||
from functools import reduce
|
||||
|
||||
from django.utils.translation import ugettext_lazy as _
|
||||
from django.db.models import F
|
||||
|
||||
from common.db.models import ChoiceSet
|
||||
from orgs.mixins.models import OrgModelMixin
|
||||
from common.db import models
|
||||
from common.utils import lazyproperty
|
||||
from assets.models import Asset, SystemUser, Node, FamilyMixin
|
||||
|
|
@ -11,7 +14,7 @@ from .base import BasePermission
|
|||
|
||||
|
||||
__all__ = [
|
||||
'AssetPermission', 'Action', 'UserGrantedMappingNode', 'RebuildUserTreeTask',
|
||||
'AssetPermission', 'Action', 'PermNode', 'UserAssetGrantedTreeNodeRelation',
|
||||
]
|
||||
|
||||
# 使用场景
|
||||
|
|
@ -135,39 +138,109 @@ class AssetPermission(BasePermission):
|
|||
from assets.models import Node
|
||||
nodes_keys = self.nodes.all().values_list('key', flat=True)
|
||||
assets_ids = set(self.assets.all().values_list('id', flat=True))
|
||||
nodes_assets_ids = Node.get_nodes_all_assets_ids(nodes_keys)
|
||||
nodes_assets_ids = Node.get_nodes_all_assets_ids_by_keys(nodes_keys)
|
||||
assets_ids.update(nodes_assets_ids)
|
||||
assets = Asset.objects.filter(id__in=assets_ids)
|
||||
return assets
|
||||
|
||||
|
||||
class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, models.JMSBaseModel):
|
||||
class NodeFrom(ChoiceSet):
|
||||
granted = 'granted', 'Direct node granted'
|
||||
child = 'child', 'Have children node'
|
||||
asset = 'asset', 'Direct asset granted'
|
||||
|
||||
class UserGrantedMappingNode(FamilyMixin, models.JMSBaseModel):
|
||||
node = models.ForeignKey('assets.Node', default=None, on_delete=models.CASCADE,
|
||||
db_constraint=False, null=True, related_name='mapping_nodes')
|
||||
key = models.CharField(max_length=64, verbose_name=_("Key"), db_index=True) # '1:1:1:1'
|
||||
user = models.ForeignKey('users.User', db_constraint=False, on_delete=models.CASCADE)
|
||||
granted = models.BooleanField(default=False, db_index=True)
|
||||
asset_granted = models.BooleanField(default=False, db_index=True)
|
||||
parent_key = models.CharField(max_length=64, default='', verbose_name=_('Parent key'), db_index=True) # '1:1:1:1'
|
||||
assets_amount = models.IntegerField(default=0)
|
||||
node = models.ForeignKey('assets.Node', default=None, on_delete=models.CASCADE,
|
||||
db_constraint=False, null=False, related_name='granted_node_rels')
|
||||
node_key = models.CharField(max_length=64, verbose_name=_("Key"), db_index=True)
|
||||
node_parent_key = models.CharField(max_length=64, default='', verbose_name=_('Parent key'), db_index=True)
|
||||
node_from = models.CharField(choices=NodeFrom.choices, max_length=16, db_index=True)
|
||||
node_assets_amount = models.IntegerField(default=0)
|
||||
|
||||
GRANTED_DIRECT = 1
|
||||
GRANTED_INDIRECT = 2
|
||||
GRANTED_NONE = 0
|
||||
@property
|
||||
def key(self):
|
||||
return self.node_key
|
||||
|
||||
@property
|
||||
def parent_key(self):
|
||||
return self.node_parent_key
|
||||
|
||||
@classmethod
|
||||
def get_node_granted_status(cls, key, user):
|
||||
ancestor_keys = Node.get_node_ancestor_keys(key, with_self=True)
|
||||
has_granted = UserGrantedMappingNode.objects.filter(
|
||||
key__in=ancestor_keys, user=user
|
||||
).values_list('granted', flat=True)
|
||||
if not has_granted:
|
||||
return cls.GRANTED_NONE
|
||||
if any(list(has_granted)):
|
||||
return cls.GRANTED_DIRECT
|
||||
return cls.GRANTED_INDIRECT
|
||||
def get_node_granted_status(cls, user, key):
|
||||
ancestor_keys = set(cls.get_node_ancestor_keys(key, with_self=True))
|
||||
ancestor_rel_nodes = cls.objects.filter(user=user, node_key__in=ancestor_keys)
|
||||
|
||||
for rel_node in ancestor_rel_nodes:
|
||||
if rel_node.key == key:
|
||||
return rel_node.node_from, rel_node
|
||||
if rel_node.node_from == cls.NodeFrom.granted:
|
||||
return cls.NodeFrom.granted, None
|
||||
return '', None
|
||||
|
||||
|
||||
class RebuildUserTreeTask(models.JMSBaseModel):
|
||||
user = models.ForeignKey('users.User', on_delete=models.CASCADE, verbose_name=_('User'))
|
||||
class PermNode(Node):
|
||||
class Meta:
|
||||
proxy = True
|
||||
ordering = []
|
||||
|
||||
# 特殊节点
|
||||
UNGROUPED_NODE_KEY = 'ungrouped'
|
||||
UNGROUPED_NODE_VALUE = _('Ungrouped')
|
||||
FAVORITE_NODE_KEY = 'favorite'
|
||||
FAVORITE_NODE_VALUE = _('Favorite')
|
||||
|
||||
node_from = ''
|
||||
granted_assets_amount = 0
|
||||
|
||||
# 提供可以设置 资产数量的字段
|
||||
_assets_amount = None
|
||||
|
||||
annotate_granted_node_rel_fields = {
|
||||
'granted_assets_amount': F('granted_node_rels__node_assets_amount'),
|
||||
'node_from': F('granted_node_rels__node_from')
|
||||
}
|
||||
|
||||
@property
|
||||
def assets_amount(self):
|
||||
_assets_amount = getattr(self, '_assets_amount')
|
||||
if isinstance(_assets_amount, int):
|
||||
return _assets_amount
|
||||
return super().assets_amount
|
||||
|
||||
@assets_amount.setter
|
||||
def assets_amount(self, value):
|
||||
self._assets_amount = value
|
||||
|
||||
def use_granted_assets_amount(self):
|
||||
self.assets_amount = self.granted_assets_amount
|
||||
|
||||
@classmethod
|
||||
def get_ungrouped_node(cls, assets_amount):
|
||||
return cls(
|
||||
id=cls.UNGROUPED_NODE_KEY,
|
||||
key=cls.UNGROUPED_NODE_KEY,
|
||||
value=cls.UNGROUPED_NODE_VALUE,
|
||||
assets_amount=assets_amount
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_favorite_node(cls, assets_amount):
|
||||
node = cls(
|
||||
id=cls.FAVORITE_NODE_KEY,
|
||||
key=cls.FAVORITE_NODE_KEY,
|
||||
value=cls.FAVORITE_NODE_VALUE,
|
||||
)
|
||||
node.assets_amount = assets_amount
|
||||
return node
|
||||
|
||||
def get_granted_status(self, user):
|
||||
status, rel_node = UserAssetGrantedTreeNodeRelation.get_node_granted_status(user, self.key)
|
||||
self.node_from = status
|
||||
if rel_node:
|
||||
self.granted_assets_amount = rel_node.node_assets_amount
|
||||
return status
|
||||
|
||||
def save(self):
|
||||
# 这是个只读 Model
|
||||
raise NotImplementedError
|
||||
|
|
|
|||
|
|
@ -1,30 +1,54 @@
|
|||
from rest_framework.pagination import LimitOffsetPagination
|
||||
from rest_framework.request import Request
|
||||
from django.db.models import Sum
|
||||
|
||||
from perms.models import UserAssetGrantedTreeNodeRelation
|
||||
from common.utils import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class GrantedAssetLimitOffsetPagination(LimitOffsetPagination):
|
||||
class GrantedAssetPaginationBase(LimitOffsetPagination):
|
||||
|
||||
def paginate_queryset(self, queryset, request: Request, view=None):
|
||||
self._request = request
|
||||
self._view = view
|
||||
self._user = request.user
|
||||
return super().paginate_queryset(queryset, request, view=None)
|
||||
|
||||
def get_count(self, queryset):
|
||||
exclude_query_params = {
|
||||
self.limit_query_param,
|
||||
self.offset_query_param,
|
||||
'key', 'all', 'show_current_asset',
|
||||
'cache_policy', 'display', 'draw'
|
||||
'cache_policy', 'display', 'draw',
|
||||
'order',
|
||||
}
|
||||
for k, v in self._request.query_params.items():
|
||||
if k not in exclude_query_params and v is not None:
|
||||
logger.warn(f'Not hit node.assets_amount because find a unknow query_param `{k}` -> {self._request.get_full_path()}')
|
||||
return super().get_count(queryset)
|
||||
return self.get_count_from_nodes(queryset)
|
||||
|
||||
def get_count_from_nodes(self, queryset):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class NodeGrantedAssetPagination(GrantedAssetPaginationBase):
|
||||
def get_count_from_nodes(self, queryset):
|
||||
node = getattr(self._view, 'pagination_node', None)
|
||||
if node:
|
||||
logger.debug(f'{self._request.get_full_path()} hit node.assets_amount[{node.assets_amount}]')
|
||||
logger.debug(f'Hit node.assets_amount[{node.assets_amount}] -> {self._request.get_full_path()}')
|
||||
return node.assets_amount
|
||||
else:
|
||||
logger.warn(f'Not hit node.assets_amount[{node}] because {self._view} not has `pagination_node` -> {self._request.get_full_path()}')
|
||||
return super().get_count(queryset)
|
||||
|
||||
def paginate_queryset(self, queryset, request: Request, view=None):
|
||||
self._request = request
|
||||
self._view = view
|
||||
return super().paginate_queryset(queryset, request, view=None)
|
||||
|
||||
class AllGrantedAssetPagination(GrantedAssetPaginationBase):
|
||||
def get_count_from_nodes(self, queryset):
|
||||
assets_amount = sum(UserAssetGrantedTreeNodeRelation.objects.filter(
|
||||
user=self._user, node_parent_key=''
|
||||
).values_list('node_assets_amount', flat=True))
|
||||
logger.debug(f'Hit all assets amount {assets_amount} -> {self._request.get_full_path()}')
|
||||
return assets_amount
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
from . import common
|
||||
from . import refresh_perms
|
||||
|
|
@ -1,31 +1,22 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
from django.db.models.signals import m2m_changed, pre_delete, pre_save
|
||||
from django.db.models.signals import m2m_changed
|
||||
from django.dispatch import receiver
|
||||
|
||||
from perms.tasks import create_rebuild_user_tree_task, \
|
||||
create_rebuild_user_tree_task_by_related_nodes_or_assets
|
||||
from users.models import User, UserGroup
|
||||
from assets.models import Asset, SystemUser
|
||||
from assets.models import SystemUser
|
||||
from applications.models import Application
|
||||
from common.utils import get_logger
|
||||
from common.exceptions import M2MReverseNotAllowed
|
||||
from common.const.signals import POST_ADD, POST_REMOVE, POST_CLEAR
|
||||
from .models import AssetPermission, ApplicationPermission
|
||||
from common.const.signals import POST_ADD
|
||||
from perms.models import AssetPermission, ApplicationPermission
|
||||
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
def handle_rebuild_user_tree(instance, action, reverse, pk_set, **kwargs):
|
||||
if action.startswith('post'):
|
||||
if reverse:
|
||||
create_rebuild_user_tree_task(pk_set)
|
||||
else:
|
||||
create_rebuild_user_tree_task([instance.id])
|
||||
|
||||
|
||||
def handle_bind_groups_systemuser(instance, action, reverse, pk_set, **kwargs):
|
||||
@receiver(m2m_changed, sender=User.groups.through)
|
||||
def on_user_groups_change(sender, instance, action, reverse, pk_set, **kwargs):
|
||||
"""
|
||||
UserGroup 增加 User 时，增加的 User 需要与 UserGroup 关联的动态系统用户相关联
|
||||
"""
|
||||
|
|
@ -47,53 +38,11 @@ def handle_bind_groups_systemuser(instance, action, reverse, pk_set, **kwargs):
|
|||
system_user.users.add(*users_id)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=User.groups.through)
|
||||
def on_user_groups_change(**kwargs):
|
||||
handle_rebuild_user_tree(**kwargs)
|
||||
handle_bind_groups_systemuser(**kwargs)
|
||||
|
||||
|
||||
@receiver([pre_save], sender=AssetPermission)
|
||||
def on_asset_perm_deactive(instance: AssetPermission, **kwargs):
|
||||
try:
|
||||
old = AssetPermission.objects.only('is_active').get(id=instance.id)
|
||||
if instance.is_active != old.is_active:
|
||||
create_rebuild_user_tree_task_by_asset_perm(instance)
|
||||
except AssetPermission.DoesNotExist:
|
||||
pass
|
||||
|
||||
|
||||
@receiver([pre_delete], sender=AssetPermission)
|
||||
def on_asset_permission_delete(instance, **kwargs):
|
||||
# 授权删除之前，查出所有相关用户
|
||||
create_rebuild_user_tree_task_by_asset_perm(instance)
|
||||
|
||||
|
||||
def create_rebuild_user_tree_task_by_asset_perm(asset_perm: AssetPermission):
|
||||
user_ids = set()
|
||||
user_ids.update(
|
||||
UserGroup.objects.filter(
|
||||
assetpermissions=asset_perm, users__id__isnull=False
|
||||
).distinct().values_list('users__id', flat=True)
|
||||
)
|
||||
user_ids.update(
|
||||
User.objects.filter(assetpermissions=asset_perm).distinct().values_list('id', flat=True)
|
||||
)
|
||||
create_rebuild_user_tree_task(user_ids)
|
||||
|
||||
|
||||
def need_rebuild_mapping_node(action):
|
||||
return action in (POST_REMOVE, POST_ADD, POST_CLEAR)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=AssetPermission.nodes.through)
|
||||
def on_permission_nodes_changed(instance, action, reverse, pk_set, model, **kwargs):
|
||||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
create_rebuild_user_tree_task_by_asset_perm(instance)
|
||||
|
||||
if action != POST_ADD:
|
||||
return
|
||||
logger.debug("Asset permission nodes change signal received")
|
||||
|
|
@ -110,9 +59,6 @@ def on_permission_assets_changed(instance, action, reverse, pk_set, model, **kwa
|
|||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
create_rebuild_user_tree_task_by_asset_perm(instance)
|
||||
|
||||
if action != POST_ADD:
|
||||
return
|
||||
logger.debug("Asset permission assets change signal received")
|
||||
|
|
@ -150,9 +96,6 @@ def on_asset_permission_users_changed(instance, action, reverse, pk_set, model,
|
|||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
create_rebuild_user_tree_task(pk_set)
|
||||
|
||||
if action != POST_ADD:
|
||||
return
|
||||
logger.debug("Asset permission users change signal received")
|
||||
|
|
@ -171,10 +114,6 @@ def on_asset_permission_user_groups_changed(instance, action, pk_set, model,
|
|||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
user_ids = User.objects.filter(groups__id__in=pk_set).distinct().values_list('id', flat=True)
|
||||
create_rebuild_user_tree_task(user_ids)
|
||||
|
||||
if action != POST_ADD:
|
||||
return
|
||||
logger.debug("Asset permission user groups change signal received")
|
||||
|
|
@ -187,21 +126,6 @@ def on_asset_permission_user_groups_changed(instance, action, pk_set, model,
|
|||
system_user.groups.add(*tuple(groups))
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=Asset.nodes.through)
|
||||
def on_node_asset_change(action, instance, reverse, pk_set, **kwargs):
|
||||
if not need_rebuild_mapping_node(action):
|
||||
return
|
||||
|
||||
if reverse:
|
||||
asset_pk_set = pk_set
|
||||
node_pk_set = [instance.id]
|
||||
else:
|
||||
asset_pk_set = [instance.id]
|
||||
node_pk_set = pk_set
|
||||
|
||||
create_rebuild_user_tree_task_by_related_nodes_or_assets.delay(node_pk_set, asset_pk_set)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=ApplicationPermission.system_users.through)
|
||||
def on_application_permission_system_users_changed(sender, instance: ApplicationPermission, action, reverse, pk_set, **kwargs):
|
||||
if not instance.category_remote_app:
|
||||
|
|
@ -0,0 +1,115 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
from django.db.models.signals import m2m_changed, pre_delete, pre_save, post_save
|
||||
from django.dispatch import receiver
|
||||
|
||||
from users.models import User
|
||||
from assets.models import Asset
|
||||
from orgs.utils import current_org
|
||||
from common.utils import get_logger
|
||||
from common.exceptions import M2MReverseNotAllowed
|
||||
from common.const.signals import POST_ADD, POST_REMOVE, POST_CLEAR
|
||||
from perms.models import AssetPermission
|
||||
from perms.utils.asset.user_permission import UserGrantedTreeRefreshController
|
||||
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=User.groups.through)
|
||||
def on_user_groups_change(sender, instance, action, reverse, pk_set, **kwargs):
|
||||
if action.startswith('post'):
|
||||
if reverse:
|
||||
group_ids = [instance.id]
|
||||
user_ids = pk_set
|
||||
else:
|
||||
group_ids = pk_set
|
||||
user_ids = [instance.id]
|
||||
|
||||
exists = AssetPermission.user_groups.through.objects.filter(usergroup_id__in=group_ids).exists()
|
||||
if exists:
|
||||
org_ids = [current_org.id]
|
||||
UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users(org_ids, user_ids)
|
||||
|
||||
|
||||
@receiver([pre_delete], sender=AssetPermission)
|
||||
def on_asset_perm_pre_delete(sender, instance, **kwargs):
|
||||
# 授权删除之前，查出所有相关用户
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
|
||||
|
||||
|
||||
@receiver([pre_save], sender=AssetPermission)
|
||||
def on_asset_perm_pre_save(sender, instance, **kwargs):
|
||||
try:
|
||||
old = AssetPermission.objects.get(id=instance.id)
|
||||
|
||||
if old.is_valid != instance.is_valid:
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
|
||||
except AssetPermission.DoesNotExist:
|
||||
pass
|
||||
|
||||
|
||||
@receiver([post_save], sender=AssetPermission)
|
||||
def on_asset_perm_post_save(sender, instance, created, **kwargs):
|
||||
if created:
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
|
||||
|
||||
|
||||
def need_rebuild_mapping_node(action):
|
||||
return action in (POST_REMOVE, POST_ADD, POST_CLEAR)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=AssetPermission.nodes.through)
|
||||
def on_permission_nodes_changed(sender, instance, action, reverse, **kwargs):
|
||||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=AssetPermission.assets.through)
|
||||
def on_permission_assets_changed(sender, instance, action, reverse, pk_set, model, **kwargs):
|
||||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=AssetPermission.users.through)
|
||||
def on_asset_permission_users_changed(sender, action, reverse, pk_set, **kwargs):
|
||||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users(
|
||||
[current_org.id], pk_set
|
||||
)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=AssetPermission.user_groups.through)
|
||||
def on_asset_permission_user_groups_changed(sender, action, pk_set, reverse, **kwargs):
|
||||
if reverse:
|
||||
raise M2MReverseNotAllowed
|
||||
|
||||
if need_rebuild_mapping_node(action):
|
||||
user_ids = User.groups.through.objects.filter(usergroup_id__in=pk_set).distinct().values_list('user_id', flat=True)
|
||||
UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users(
|
||||
[current_org.id], user_ids
|
||||
)
|
||||
|
||||
|
||||
@receiver(m2m_changed, sender=Asset.nodes.through)
|
||||
def on_node_asset_change(action, instance, reverse, pk_set, **kwargs):
|
||||
if not need_rebuild_mapping_node(action):
|
||||
return
|
||||
|
||||
if reverse:
|
||||
asset_pk_set = pk_set
|
||||
node_pk_set = [instance.id]
|
||||
else:
|
||||
asset_pk_set = [instance.id]
|
||||
node_pk_set = pk_set
|
||||
|
||||
UserGrantedTreeRefreshController.add_need_refresh_on_nodes_assets_relate_change(node_pk_set, asset_pk_set)
|
||||
|
|
@ -2,39 +2,18 @@
|
|||
from __future__ import absolute_import, unicode_literals
|
||||
from datetime import timedelta
|
||||
|
||||
from django.db import transaction
|
||||
from django.db.models import Q
|
||||
from django.db.transaction import atomic
|
||||
from django.conf import settings
|
||||
from celery import shared_task
|
||||
from common.utils import get_logger
|
||||
from common.utils.timezone import now, dt_formater, dt_parser
|
||||
from users.models import User
|
||||
from ops.celery.decorator import register_as_period_task
|
||||
from assets.models import Node
|
||||
from perms.models import RebuildUserTreeTask, AssetPermission
|
||||
from perms.utils.asset.user_permission import rebuild_user_mapping_nodes_if_need_with_lock, lock
|
||||
from perms.models import AssetPermission
|
||||
from perms.utils.asset.user_permission import UserGrantedTreeRefreshController
|
||||
|
||||
logger = get_logger(__file__)
|
||||
|
||||
|
||||
@shared_task(queue='node_tree')
|
||||
def rebuild_user_mapping_nodes_celery_task(user_id):
|
||||
user = User.objects.get(id=user_id)
|
||||
try:
|
||||
rebuild_user_mapping_nodes_if_need_with_lock(user)
|
||||
except lock.SomeoneIsDoingThis:
|
||||
pass
|
||||
|
||||
|
||||
@shared_task(queue='node_tree')
|
||||
def dispatch_mapping_node_tasks():
|
||||
user_ids = RebuildUserTreeTask.objects.all().values_list('user_id', flat=True).distinct()
|
||||
logger.info(f'>>> dispatch_mapping_node_tasks for users {list(user_ids)}')
|
||||
for id in user_ids:
|
||||
rebuild_user_mapping_nodes_celery_task.delay(id)
|
||||
|
||||
|
||||
@register_as_period_task(interval=settings.PERM_EXPIRED_CHECK_PERIODIC)
|
||||
@shared_task(queue='celery_check_asset_perm_expired')
|
||||
@atomic()
|
||||
|
|
@ -60,66 +39,9 @@ def check_asset_permission_expired():
|
|||
setting.value = dt_formater(end)
|
||||
setting.save()
|
||||
|
||||
ids = AssetPermission.objects.filter(
|
||||
asset_perm_ids = AssetPermission.objects.filter(
|
||||
date_expired__gte=start, date_expired__lte=end
|
||||
).distinct().values_list('id', flat=True)
|
||||
logger.info(f'>>> checking {start} to {end} have {ids} expired')
|
||||
dispatch_process_expired_asset_permission.delay(list(ids))
|
||||
|
||||
|
||||
@shared_task(queue='node_tree')
|
||||
def dispatch_process_expired_asset_permission(asset_perms_id):
|
||||
user_ids = User.objects.filter(
|
||||
Q(assetpermissions__id__in=asset_perms_id) |
|
||||
Q(groups__assetpermissions__id__in=asset_perms_id)
|
||||
).distinct().values_list('id', flat=True)
|
||||
RebuildUserTreeTask.objects.bulk_create(
|
||||
[RebuildUserTreeTask(user_id=user_id) for user_id in user_ids]
|
||||
)
|
||||
|
||||
dispatch_mapping_node_tasks.delay()
|
||||
|
||||
|
||||
def create_rebuild_user_tree_task(user_ids):
|
||||
RebuildUserTreeTask.objects.bulk_create(
|
||||
[RebuildUserTreeTask(user_id=i) for i in user_ids]
|
||||
)
|
||||
transaction.on_commit(dispatch_mapping_node_tasks.delay)
|
||||
|
||||
|
||||
@shared_task(queue='node_tree')
|
||||
def create_rebuild_user_tree_task_by_related_nodes_or_assets(node_ids, asset_ids):
|
||||
node_ids = set(node_ids)
|
||||
node_keys = set()
|
||||
nodes = Node.objects.filter(id__in=node_ids)
|
||||
for _node in nodes:
|
||||
node_keys.update(_node.get_ancestor_keys())
|
||||
node_ids.update(
|
||||
Node.objects.filter(key__in=node_keys).values_list('id', flat=True)
|
||||
)
|
||||
|
||||
asset_perms_id = set()
|
||||
asset_perms_id.update(
|
||||
AssetPermission.objects.filter(
|
||||
assets__id__in=asset_ids
|
||||
).values_list('id', flat=True).distinct()
|
||||
)
|
||||
asset_perms_id.update(
|
||||
AssetPermission.objects.filter(
|
||||
nodes__id__in=node_ids
|
||||
).values_list('id', flat=True).distinct()
|
||||
)
|
||||
|
||||
user_ids = set()
|
||||
user_ids.update(
|
||||
User.objects.filter(
|
||||
assetpermissions__id__in=asset_perms_id
|
||||
).distinct().values_list('id', flat=True)
|
||||
)
|
||||
user_ids.update(
|
||||
User.objects.filter(
|
||||
groups__assetpermissions__id__in=asset_perms_id
|
||||
).distinct().values_list('id', flat=True)
|
||||
)
|
||||
|
||||
create_rebuild_user_tree_task(user_ids)
|
||||
asset_perm_ids = list(asset_perm_ids)
|
||||
logger.info(f'>>> checking {start} to {end} have {asset_perm_ids} expired')
|
||||
UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids_cross_orgs(asset_perm_ids)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -5,7 +5,6 @@ import forgery_py
|
|||
from .base import FakeDataGenerator
|
||||
|
||||
from assets.models import *
|
||||
from assets.utils import check_node_assets_amount
|
||||
|
||||
|
||||
class AdminUsersGenerator(FakeDataGenerator):
|
||||
|
|
@ -93,4 +92,4 @@ class AssetsGenerator(FakeDataGenerator):
|
|||
self.set_assets_nodes(creates)
|
||||
|
||||
def after_generate(self):
|
||||
check_node_assets_amount()
|
||||
pass
|
||||
|
|
|
|||
Loading…
Reference in New Issue