diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py index 2277cafa6..e3b2afdca 100644 --- a/apps/perms/models/asset_permission.py +++ b/apps/perms/models/asset_permission.py @@ -15,6 +15,7 @@ from common.db.models import UnionQuerySet from common.utils import date_expired_default from perms.const import ActionChoices +from .perm_node import PermNode __all__ = ['AssetPermission', 'ActionChoices'] @@ -48,6 +49,10 @@ class AssetPermissionManager(OrgManager): def valid(self): return self.get_queryset().valid() + def get_expired_permissions(self): + now = local_now() + return self.get_queryset().filter(Q(date_start__lte=now) | Q(date_expired__gte=now)) + class AssetPermission(OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) @@ -147,10 +152,3 @@ class AssetPermission(OrgModelMixin): if flat: return user_ids return User.objects.filter(id__in=user_ids) - - @classmethod - def get_expired_permissions(cls): - now = local_now() - return cls.objects.filter(Q(date_start__lte=now) | Q(date_expired__gte=now)) - - diff --git a/apps/perms/models/perm_node.py b/apps/perms/models/perm_node.py index ce061297e..4dddae98f 100644 --- a/apps/perms/models/perm_node.py +++ b/apps/perms/models/perm_node.py @@ -3,17 +3,20 @@ from django.utils.translation import ugettext_lazy as _ from django.db import models from django.db.models import F, TextChoices -from common.utils import lazyproperty -from common.db.models import BaseCreateUpdateModel from assets.models import Asset, Node, FamilyMixin, Account from orgs.mixins.models import OrgModelMixin +from common.utils import lazyproperty +from common.db.models import BaseCreateUpdateModel + + +class NodeFrom(TextChoices): + granted = 'granted', 'Direct node granted' + child = 'child', 'Have children node' + asset = 'asset', 'Direct asset granted' class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, BaseCreateUpdateModel): - class NodeFrom(TextChoices): - granted = 'granted', 'Direct node granted' - child = 'child', 'Have children node' - asset = 'asset', 'Direct asset granted' + NodeFrom = NodeFrom user = models.ForeignKey('users.User', db_constraint=False, on_delete=models.CASCADE) node = models.ForeignKey('assets.Node', default=None, on_delete=models.CASCADE, @@ -46,6 +49,8 @@ class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, BaseCreateUpd class PermNode(Node): + NodeFrom = NodeFrom + class Meta: proxy = True ordering = [] diff --git a/apps/perms/tasks.py b/apps/perms/tasks.py index 3df7a1728..57b9f48bd 100644 --- a/apps/perms/tasks.py +++ b/apps/perms/tasks.py @@ -29,7 +29,7 @@ logger = get_logger(__file__) @tmp_to_root_org() def check_asset_permission_expired(): """ 这里的任务要足够短，不要影响周期任务 """ - perms = AssetPermission.get_expired_permissions() + perms = AssetPermission.objects.get_expired_permissions() perm_ids = list(perms.distinct().values_list('id', flat=True)) logger.info(f'Checking expired permissions: {perm_ids}') UserPermTreeExpireUtil().expire_perm_tree_for_perms(perm_ids) diff --git a/apps/perms/utils/permission.py b/apps/perms/utils/permission.py index 849bec02b..bbec44343 100644 --- a/apps/perms/utils/permission.py +++ b/apps/perms/utils/permission.py @@ -1,6 +1,5 @@ -import django -from django.db.models import QuerySet, Model -from collections.abc import Iterable +from django.db.models import QuerySet + from assets.models import Node, Asset from common.utils import get_logger @@ -90,11 +89,6 @@ class AssetPermissionUtil(object): perms = self.get_permissions(ids=perm_ids) return perms - @staticmethod - def get_permissions(ids): - perms = AssetPermission.objects.filter(id__in=ids).order_by('-date_expired') - return perms - @staticmethod def convert_to_queryset_if_need(objs_or_ids, model): if not objs_or_ids: @@ -107,5 +101,7 @@ class AssetPermissionUtil(object): ] return model.objects.filter(id__in=ids) - - + @staticmethod + def get_permissions(ids): + perms = AssetPermission.objects.filter(id__in=ids).order_by('-date_expired') + return perms diff --git a/apps/perms/utils/user_perm_tree.py b/apps/perms/utils/user_perm_tree.py index 0ae9b0c61..e40f20514 100644 --- a/apps/perms/utils/user_perm_tree.py +++ b/apps/perms/utils/user_perm_tree.py @@ -1,28 +1,38 @@ import time from collections import defaultdict +from django.conf import settings from django.core.cache import cache -from common.decorator import on_transaction_commit -from common.utils import get_logger -from common.utils.common import lazyproperty, timeit +from users.models import User +from assets.models import Asset +from assets.utils import NodeAssetsUtil from orgs.models import Organization from orgs.utils import ( + current_org, tmp_to_org, tmp_to_root_org ) +from common.decorator import on_transaction_commit +from common.utils import get_logger +from common.utils.common import lazyproperty, timeit +from common.db.models import output_as_string + from perms.locks import UserGrantedTreeRebuildLock from perms.models import ( AssetPermission, - UserAssetGrantedTreeNodeRelation + UserAssetGrantedTreeNodeRelation, + PermNode ) -from perms.utils.user_permission import UserGrantedTreeBuildUtils -from users.models import User + from .permission import AssetPermissionUtil logger = get_logger(__name__) -__all__ = ['UserPermTreeRefreshUtil', 'UserPermTreeExpireUtil'] +__all__ = [ + 'UserPermTreeRefreshUtil', + 'UserPermTreeExpireUtil' +] class _UserPermTreeCacheMixin: @@ -51,34 +61,34 @@ class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin): @timeit def refresh_if_need(self, force=False): - self.clean_user_perm_tree_nodes_for_legacy_org() - to_refresh_orgs = self.orgs if force else self.get_user_need_refresh_orgs() + self._clean_user_perm_tree_for_legacy_org() + to_refresh_orgs = self.orgs if force else self._get_user_need_refresh_orgs() if not to_refresh_orgs: logger.info('Not have to refresh orgs') return with UserGrantedTreeRebuildLock(self.user.id): for org in to_refresh_orgs: - self.rebuild_user_perm_tree_for_org(org) - self.mark_user_orgs_refresh_finished([str(org.id) for org in to_refresh_orgs]) + self._rebuild_user_perm_tree_for_org(org) + self._mark_user_orgs_refresh_finished(to_refresh_orgs) - def rebuild_user_perm_tree_for_org(self, org): + def _rebuild_user_perm_tree_for_org(self, org): with tmp_to_org(org): start = time.time() - UserGrantedTreeBuildUtils(self.user).rebuild_user_granted_tree() + UserPermTreeBuildUtil(self.user).rebuild_user_perm_tree() end = time.time() logger.info( 'Refresh user [{user}] org [{org}] perm tree, user {use_time:.2f}s' - ''.format(user=self.user, org=org, use_time=end - start) + ''.format(user=self.user, org=org, use_time=end-start) ) - def clean_user_perm_tree_nodes_for_legacy_org(self): + def _clean_user_perm_tree_for_legacy_org(self): with tmp_to_root_org(): """ Clean user legacy org node relations """ user_relations = UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user) user_legacy_org_relations = user_relations.exclude(org_id__in=self.org_ids) user_legacy_org_relations.delete() - def get_user_need_refresh_orgs(self): + def _get_user_need_refresh_orgs(self): cached_org_ids = self.client.smembers(self.cache_key_user) cached_org_ids = {oid.decode() for oid in cached_org_ids} to_refresh_org_ids = set(self.org_ids) - cached_org_ids @@ -86,7 +96,7 @@ class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin): logger.info(f'Need to refresh orgs: {to_refresh_orgs}') return to_refresh_orgs - def mark_user_orgs_refresh_finished(self, org_ids): + def _mark_user_orgs_refresh_finished(self, org_ids): self.client.sadd(self.cache_key_user, *org_ids) @@ -141,3 +151,181 @@ class UserPermTreeExpireUtil(_UserPermTreeCacheMixin): p.delete(k) p.execute() logger.info('Expire all user perm tree') + + +class UserPermTreeBuildUtil(object): + node_only_fields = ('id', 'key', 'parent_key', 'org_id') + + def __init__(self, user): + self.user = user + self.user_perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True) + # {key: node} + self._perm_nodes_key_node_mapper = {} + + def rebuild_user_perm_tree(self): + self.clean_user_perm_tree() + if not self.user_perm_ids: + logger.info('User({}) not have permissions'.format(self.user)) + return + self.compute_perm_nodes() + self.compute_perm_nodes_asset_amount() + self.create_mapping_nodes() + + def clean_user_perm_tree(self): + UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user).delete() + + def compute_perm_nodes(self): + self._compute_perm_nodes_for_direct() + self._compute_perm_nodes_for_direct_asset_if_need() + self._compute_perm_nodes_for_ancestor() + + def compute_perm_nodes_asset_amount(self): + """ 这里计算的是一个组织的授权树 """ + computed = self._only_compute_root_node_assets_amount_if_need() + if computed: + return + + nodekey_assetid_mapper = defaultdict(set) + org_id = current_org.id + for key in self.perm_node_keys_for_granted: + asset_ids = PermNode.get_all_asset_ids_by_node_key(org_id, key) + nodekey_assetid_mapper[key].update(asset_ids) + + for asset_id, node_id in self.direct_asset_id_node_id_pairs: + node_key = self.perm_nodes_id_key_mapper.get(node_id) + if not node_key: + continue + nodekey_assetid_mapper[node_key].add(asset_id) + + util = NodeAssetsUtil(self.perm_nodes, nodekey_assetid_mapper) + util.generate() + + for node in self.perm_nodes: + assets_amount = util.get_assets_amount(node.key) + node.assets_amount = assets_amount + + def create_mapping_nodes(self): + to_create = [] + for node in self.perm_nodes: + relation = UserAssetGrantedTreeNodeRelation( + user=self.user, + node=node, + node_key=node.key, + node_parent_key=node.parent_key, + node_from=node.node_from, + node_assets_amount=node.assets_amount, + org_id=node.org_id + ) + to_create.append(relation) + + UserAssetGrantedTreeNodeRelation.objects.bulk_create(to_create) + + def _compute_perm_nodes_for_direct(self): + """ 直接授权的节点（叶子节点）""" + for node in self.direct_nodes: + if self.has_any_ancestor_direct_permed(node): + continue + node.node_from = node.NodeFrom.granted + self._perm_nodes_key_node_mapper[node.key] = node + + def _compute_perm_nodes_for_direct_asset_if_need(self): + """ 直接授权的资产所在的节点（叶子节点）""" + if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: + return + for node in self.direct_asset_nodes: + if self.has_any_ancestor_direct_permed(node): + continue + if node.key in self._perm_nodes_key_node_mapper: + continue + node.node_from = node.NodeFrom.asset + self._perm_nodes_key_node_mapper[node.key] = node + + def _compute_perm_nodes_for_ancestor(self): + """ 直接授权节点 和 直接授权资产所在节点 的所有祖先节点 (构造完整树) """ + ancestor_keys = set() + for node in self._perm_nodes_key_node_mapper.values(): + ancestor_keys.update(node.get_ancestor_keys()) + ancestor_keys -= set(self._perm_nodes_key_node_mapper.keys()) + + ancestors = PermNode.objects.filter(key__in=ancestor_keys).only(*self.node_only_fields) + for node in ancestors: + node.node_from = node.NodeFrom.child + self._perm_nodes_key_node_mapper[node.key] = node + + @lazyproperty + def perm_node_keys_for_granted(self): + keys = [ + key for key, node in self._perm_nodes_key_node_mapper.items() + if node.node_from == node.NodeFrom.granted + ] + return keys + + @lazyproperty + def perm_nodes_id_key_mapper(self): + mapper = { + node.id.hex: node.key + for key, node in self._perm_nodes_key_node_mapper.items() + } + return mapper + + def _only_compute_root_node_assets_amount_if_need(self): + if len(self.perm_nodes) != 1: + return False + root_node = self.perm_nodes[0] + if not root_node.is_org_root(): + return False + if root_node.node_from != root_node.NodeFrom.granted: + return False + root_node.granted_assets_amount = len(root_node.get_all_asset_ids()) + return True + + @lazyproperty + def perm_nodes(self): + """ 授权树的所有节点 """ + return list(self._perm_nodes_key_node_mapper.values()) + + def has_any_ancestor_direct_permed(self, node): + """ 任何一个祖先节点被直接授权 """ + return bool(set(node.get_ancestor_keys()) & set(self.direct_node_keys)) + + @lazyproperty + def direct_node_keys(self): + """ 直接授权的节点 keys """ + return {n.key for n in self.direct_nodes} + + @lazyproperty + def direct_nodes(self): + """ 直接授权的节点 """ + node_ids = AssetPermission.nodes.through.objects \ + .filter(assetpermission_id__in=self.user_perm_ids) \ + .values_list('node_id', flat=True).distinct() + nodes = PermNode.objects.filter(id__in=node_ids).only(*self.node_only_fields) + return nodes + + @lazyproperty + def direct_asset_nodes(self): + """ 获取直接授权的资产所在的节点 """ + node_ids = [node_id for asset_id, node_id in self.direct_asset_id_node_id_pairs] + nodes = PermNode.objects.filter(id__in=node_ids).distinct().only(*self.node_only_fields) + return nodes + + @lazyproperty + def direct_asset_id_node_id_pairs(self): + """ 直接授权的资产 id 和 节点 id """ + asset_node_pairs = Asset.nodes.through.objects \ + .filter(asset_id__in=self.direct_asset_ids) \ + .annotate( + str_asset_id=output_as_string('asset_id'), + str_node_id=output_as_string('node_id') + ).values_list('str_asset_id', 'str_node_id') + asset_node_pairs = list(asset_node_pairs) + return asset_node_pairs + + @lazyproperty + def direct_asset_ids(self): + """ 直接授权的资产 ids """ + asset_ids = AssetPermission.assets.through.objects \ + .filter(assetpermission_id__in=self.user_perm_ids) \ + .values_list('asset_id', flat=True) \ + .distinct() + return asset_ids diff --git a/apps/perms/utils/user_permission.py b/apps/perms/utils/user_permission.py index 81b4fc00e..674f562ff 100644 --- a/apps/perms/utils/user_permission.py +++ b/apps/perms/utils/user_permission.py @@ -51,217 +51,6 @@ class UserGrantedUtilsBase: return asset_perm_ids -class UserGrantedTreeBuildUtils(UserGrantedUtilsBase): - - def get_direct_granted_nodes(self) -> NodeQuerySet: - # 查询直接授权节点 - nodes = PermNode.objects.filter( - granted_by_permissions__id__in=self.asset_perm_ids - ).distinct() - return nodes - - @lazyproperty - def direct_granted_asset_ids(self) -> list: - # 3.15 - asset_ids = AssetPermission.assets.through.objects.filter( - assetpermission_id__in=self.asset_perm_ids - ).annotate( - asset_id_str=output_as_string('asset_id') - ).values_list( - 'asset_id_str', flat=True - ).distinct() - - asset_ids = list(asset_ids) - return asset_ids - - @ensure_in_real_or_default_org - def rebuild_user_granted_tree(self): - """ - 注意：调用该方法一定要被 `UserGrantedTreeRebuildLock` 锁住 - """ - user = self.user - - # 先删除旧的授权树🌲 - UserAssetGrantedTreeNodeRelation.objects.filter(user=user).delete() - - if not self.asset_perm_ids: - # 没有授权直接返回 - return - - nodes = self.compute_perm_nodes_tree() - self.compute_node_assets_amount(nodes) - if not nodes: - return - self.create_mapping_nodes(nodes) - - @timeit - def compute_perm_nodes_tree(self, node_only_fields=NODE_ONLY_FIELDS) -> list: - - # 查询直接授权节点 - nodes = self.get_direct_granted_nodes().only(*node_only_fields) - nodes = list(nodes) - - # 授权的节点 key 集合 - granted_key_set = {_node.key for _node in nodes} - - def _has_ancestor_granted(node: PermNode): - """ - 判断一个节点是否有授权过的祖先节点 - """ - ancestor_keys = set(node.get_ancestor_keys()) - return ancestor_keys & granted_key_set - - key2leaf_nodes_mapper = {} - - # 给授权节点设置 granted 标识，同时去重 - for node in nodes: - node: PermNode - if _has_ancestor_granted(node): - continue - node.node_from = NodeFrom.granted - key2leaf_nodes_mapper[node.key] = node - - # 查询授权资产关联的节点设置 - def process_direct_granted_assets(): - # 查询直接授权资产 - node_ids = {node_id_str for node_id_str, _ in self.direct_granted_asset_id_node_id_str_pairs} - # 查询授权资产关联的节点设置 2.80 - granted_asset_nodes = PermNode.objects.filter( - id__in=node_ids - ).distinct().only(*node_only_fields) - granted_asset_nodes = list(granted_asset_nodes) - - # 给资产授权关联的节点设置 is_asset_granted 标识，同时去重 - for node in granted_asset_nodes: - if _has_ancestor_granted(node): - continue - if node.key in key2leaf_nodes_mapper: - continue - node.node_from = NodeFrom.asset - key2leaf_nodes_mapper[node.key] = node - - if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: - process_direct_granted_assets() - - leaf_nodes = key2leaf_nodes_mapper.values() - - # 计算所有祖先节点 - ancestor_keys = set() - for node in leaf_nodes: - ancestor_keys.update(node.get_ancestor_keys()) - - # 从祖先节点 key 中去掉同时也是叶子节点的 key - ancestor_keys -= key2leaf_nodes_mapper.keys() - # 查出祖先节点 - ancestors = PermNode.objects.filter(key__in=ancestor_keys).only(*node_only_fields) - ancestors = list(ancestors) - for node in ancestors: - node.node_from = NodeFrom.child - result = [*leaf_nodes, *ancestors] - return result - - @timeit - def create_mapping_nodes(self, nodes): - user = self.user - to_create = [] - - for node in nodes: - to_create.append(UserAssetGrantedTreeNodeRelation( - user=user, - node=node, - node_key=node.key, - node_parent_key=node.parent_key, - node_from=node.node_from, - node_assets_amount=node.assets_amount, - org_id=node.org_id - )) - - UserAssetGrantedTreeNodeRelation.objects.bulk_create(to_create) - - @timeit - def _fill_direct_granted_node_asset_ids_from_mem(self, nodes_key, mapper): - org_id = current_org.id - for key in nodes_key: - asset_ids = PermNode.get_all_asset_ids_by_node_key(org_id, key) - mapper[key].update(asset_ids) - - @lazyproperty - def direct_granted_asset_id_node_id_str_pairs(self): - node_asset_pairs = Asset.nodes.through.objects.filter( - asset_id__in=self.direct_granted_asset_ids - ).annotate( - asset_id_str=output_as_string('asset_id'), - node_id_str=output_as_string('node_id') - ).values_list( - 'node_id_str', 'asset_id_str' - ) - node_asset_pairs = list(node_asset_pairs) - return node_asset_pairs - - @timeit - def compute_node_assets_amount(self, nodes: List[PermNode]): - """ - 这里计算的是一个组织的 - """ - # 直接授权了根节点，直接计算 - if len(nodes) == 1: - node = nodes[0] - if node.node_from == NodeFrom.granted and node.key.isdigit(): - with tmp_to_org(node.org): - node.granted_assets_amount = len(node.get_all_asset_ids()) - return - - direct_granted_nodes_key = [] - node_id_key_mapper = {} - for node in nodes: - if node.node_from == NodeFrom.granted: - direct_granted_nodes_key.append(node.key) - node_id_key_mapper[node.id.hex] = node.key - - # 授权的节点和直接资产的映射 - nodekey_assetsid_mapper = defaultdict(set) - # 直接授权的节点，资产从完整树过来 - self._fill_direct_granted_node_asset_ids_from_mem( - direct_granted_nodes_key, nodekey_assetsid_mapper - ) - - # 处理直接授权资产 - # 直接授权资产，取节点与资产的关系 - node_asset_pairs = self.direct_granted_asset_id_node_id_str_pairs - node_asset_pairs = list(node_asset_pairs) - - for node_id, asset_id in node_asset_pairs: - if node_id not in node_id_key_mapper: - continue - node_key = node_id_key_mapper[node_id] - nodekey_assetsid_mapper[node_key].add(asset_id) - - util = NodeAssetsUtil(nodes, nodekey_assetsid_mapper) - util.generate() - - for node in nodes: - assets_amount = util.get_assets_amount(node.key) - node.assets_amount = assets_amount - - def get_whole_tree_nodes(self) -> list: - node_only_fields = NODE_ONLY_FIELDS + ('value', 'full_value') - nodes = self.compute_perm_nodes_tree(node_only_fields=node_only_fields) - self.compute_node_assets_amount(nodes) - - # 查询直接授权节点的子节点 - q = Q() - for node in self.get_direct_granted_nodes().only('key'): - q |= Q(key__startswith=f'{node.key}:') - - if q: - descendant_nodes = PermNode.objects.filter(q).distinct() - else: - descendant_nodes = PermNode.objects.none() - - nodes.extend(descendant_nodes) - return nodes - - class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase): def get_favorite_assets(self) -> QuerySet: diff --git a/apps/users/models/user.py b/apps/users/models/user.py index d6137add9..b77d1151f 100644 --- a/apps/users/models/user.py +++ b/apps/users/models/user.py @@ -742,13 +742,6 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, AbstractUser): def __str__(self): return '{0.name}({0.username})'.format(self) - @classmethod - def get_group_ids_by_user_id(cls, user_id): - group_ids = cls.groups.through.objects.filter(user_id=user_id) \ - .distinct().values_list('usergroup_id', flat=True) - group_ids = list(group_ids) - return group_ids - @property def receive_backends(self): return self.user_msg_subscription.receive_backends