diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index 81693252d..c191dafe4 100644 --- a/apps/assets/models/asset/common.py +++ b/apps/assets/models/asset/common.py @@ -63,18 +63,29 @@ class NodesRelationMixin: nodes = Node.objects.filter(id=Node.org_root().id) return nodes - def get_all_nodes(self, flat=False): + def get_all_nodes(self, flat=False, only_keys=False): from ..node import Node node_keys = set() for node in self.get_nodes(): ancestor_keys = node.get_ancestor_keys(with_self=True) node_keys.update(ancestor_keys) + if only_keys: + return node_keys nodes = Node.objects.filter(key__in=node_keys).distinct() - if flat: - node_ids = set(nodes.values_list('id', flat=True)) - return node_ids - else: + if not flat: return nodes + node_ids = set(nodes.values_list('id', flat=True)) + return node_ids + + @classmethod + def get_all_nodes_for_assets(cls, assets): + from ..node import Node + node_keys = set() + for asset in assets: + asset_node_keys = asset.get_all_nodes(only_keys=True) + node_keys.update(asset_node_keys) + nodes = Node.objects.filter(key__in=node_keys) + return nodes class Protocol(models.Model): diff --git a/apps/perms/api/user_permission/tree/mixin.py b/apps/perms/api/user_permission/tree/mixin.py index cbb405994..87c4eba3c 100644 --- a/apps/perms/api/user_permission/tree/mixin.py +++ b/apps/perms/api/user_permission/tree/mixin.py @@ -1,7 +1,7 @@ from rest_framework.request import Request from users.models import User -from perms.utils.user_permission import UserGrantedTreeRefreshController +from perms.utils.user_permission import UserPermTreeUtil from common.http import is_true @@ -13,5 +13,5 @@ class RebuildTreeMixin: def get(self, request: Request, *args, **kwargs): force = is_true(request.query_params.get('rebuild_tree')) - UserGrantedTreeRefreshController(self.user).refresh_if_need(force) + UserPermTreeUtil(self.user).refresh_if_need(force) return super().get(request, *args, **kwargs) diff --git a/apps/perms/locks.py b/apps/perms/locks.py index 96c766fb8..77babf7f8 100644 --- a/apps/perms/locks.py +++ b/apps/perms/locks.py @@ -2,7 +2,7 @@ from common.utils.lock import DistributedLock class UserGrantedTreeRebuildLock(DistributedLock): - name_template = 'perms.user.asset.node.tree.rebuid.' + name_template = 'perms.user.asset.node.tree.rebuild.' def __init__(self, user_id): name = self.name_template.format(user_id=user_id) diff --git a/apps/perms/signal_handlers/refresh_perms.py b/apps/perms/signal_handlers/refresh_perms.py index 2e66c4475..2d5617b34 100644 --- a/apps/perms/signal_handlers/refresh_perms.py +++ b/apps/perms/signal_handlers/refresh_perms.py @@ -10,7 +10,7 @@ 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.user_permission import UserGrantedTreeRefreshController +from perms.utils.user_permission import UserPermTreeUtil logger = get_logger(__file__) @@ -24,7 +24,7 @@ def on_user_group_delete(sender, instance: UserGroup, using, **kwargs): org_id = instance.org_id user_ids = UserGroup.users.through.objects.filter(usergroup_id=instance.id).values_list('user_id', flat=True) - UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users([org_id], list(user_ids)) + UserPermTreeUtil.add_need_refresh_orgs_for_users([org_id], list(user_ids)) @receiver(m2m_changed, sender=User.groups.through) @@ -46,14 +46,14 @@ def on_user_groups_change(sender, instance, action, reverse, pk_set, **kwargs): return org_ids = [org_id] - UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users(org_ids, user_ids) + UserPermTreeUtil.add_need_refresh_orgs_for_users(org_ids, user_ids) @receiver([pre_delete], sender=AssetPermission) def on_asset_perm_pre_delete(sender, instance, **kwargs): # 授权删除之前，查出所有相关用户 with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids([instance.id]) @receiver([pre_save], sender=AssetPermission) @@ -63,7 +63,7 @@ def on_asset_perm_pre_save(sender, instance, **kwargs): if old.is_valid != instance.is_valid: with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids([instance.id]) except AssetPermission.DoesNotExist: pass @@ -73,7 +73,7 @@ def on_asset_perm_post_save(sender, instance, created, **kwargs): if not created: return with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids([instance.id]) def need_rebuild_mapping_node(action): @@ -89,7 +89,7 @@ def on_permission_nodes_changed(sender, instance, action, reverse, **kwargs): return with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids([instance.id]) @receiver(m2m_changed, sender=AssetPermission.assets.through) @@ -100,7 +100,7 @@ def on_permission_assets_changed(sender, instance, action, reverse, pk_set, mode if not need_rebuild_mapping_node(action): return with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids([instance.id]) @receiver(m2m_changed, sender=AssetPermission.users.through) @@ -112,7 +112,7 @@ def on_asset_permission_users_changed(sender, action, reverse, instance, pk_set, return with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users( + UserPermTreeUtil.add_need_refresh_orgs_for_users( [current_org.id], pk_set ) @@ -129,7 +129,7 @@ def on_asset_permission_user_groups_changed(sender, instance, action, pk_set, re .values_list('user_id', flat=True) \ .distinct() with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users( + UserPermTreeUtil.add_need_refresh_orgs_for_users( [current_org.id], user_ids ) @@ -147,4 +147,4 @@ def on_node_asset_change(action, instance, reverse, pk_set, **kwargs): node_pk_set = pk_set with tmp_to_org(instance.org): - UserGrantedTreeRefreshController.add_need_refresh_on_nodes_assets_relate_change(node_pk_set, asset_pk_set) + UserPermTreeUtil.add_need_refresh_on_nodes_assets_relate_change(node_pk_set, asset_pk_set) diff --git a/apps/perms/tasks.py b/apps/perms/tasks.py index 564e5657e..e739f4bcd 100644 --- a/apps/perms/tasks.py +++ b/apps/perms/tasks.py @@ -16,7 +16,7 @@ from perms.notifications import ( PermedAssetsWillExpireUserMsg, AssetPermsWillExpireForOrgAdminMsg, ) from perms.models import AssetPermission -from perms.utils.user_permission import UserGrantedTreeRefreshController +from perms.utils.user_permission import UserPermTreeUtil logger = get_logger(__file__) @@ -52,7 +52,7 @@ def check_asset_permission_expired(): ).distinct().values_list('id', flat=True) 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) + UserPermTreeUtil.add_need_refresh_by_asset_perm_ids_cross_orgs(asset_perm_ids) @register_as_period_task(crontab=CRONTAB_AT_AM_TEN) diff --git a/apps/perms/utils/permission.py b/apps/perms/utils/permission.py index 4358e2075..89407be14 100644 --- a/apps/perms/utils/permission.py +++ b/apps/perms/utils/permission.py @@ -1,5 +1,8 @@ - +from django.db.models import QuerySet, Model +from collections.abc import Iterable +from assets.models import Node, Asset from common.utils import get_logger + from perms.models import AssetPermission logger = get_logger(__file__) @@ -41,14 +44,16 @@ class AssetPermissionUtil(object): perms = self.get_permissions(ids=group_perm_ids) return perms - def get_permissions_for_asset(self, asset, with_node=True, flat=False): + def get_permissions_for_assets(self, assets, with_node=True, flat=False): """ 获取资产的授权规则""" perm_ids = set() - asset_perm_ids = AssetPermission.assets.through.objects.filter(asset_id=asset.id) \ - .values_list('assetpermission_id', flat=True).distinct() + assets = self.transform_to_queryset(assets, Asset) + asset_ids = [str(a.id) for a in assets] + relations = AssetPermission.assets.through.objects.filter(asset_id__in=asset_ids) + asset_perm_ids = relations.values_list('assetpermission_id', flat=True).distinct() perm_ids.update(asset_perm_ids) if with_node: - nodes = asset.get_all_nodes() + nodes = Asset.get_all_nodes_for_assets(assets) node_perm_ids = self.get_permissions_for_nodes(nodes, flat=True) perm_ids.update(node_perm_ids) if flat: @@ -58,16 +63,12 @@ class AssetPermissionUtil(object): def get_permissions_for_nodes(self, nodes, with_ancestor=False, flat=False): """ 获取节点的授权规则 """ + nodes = self.transform_to_queryset(nodes, Node) if with_ancestor: - node_ids = set() - for node in nodes: - _nodes = node.get_ancestors(with_self=True) - _node_ids = _nodes.values_list('id', flat=True).distinct() - node_ids.update(_node_ids) - else: - node_ids = nodes.values_list('id', flat=True).distinct() - perm_ids = AssetPermission.nodes.through.objects.filter(node_id__in=node_ids) \ - .values_list('assetpermission_id', flat=True).distinct() + nodes = Node.get_ancestor_queryset(nodes) + node_ids = nodes.values_list('id', flat=True).distinct() + relations = AssetPermission.nodes.through.objects.filter(node_id__in=node_ids) + perm_ids = relations.values_list('assetpermission_id', flat=True).distinct() if flat: return perm_ids perms = self.get_permissions(ids=perm_ids) @@ -76,14 +77,14 @@ class AssetPermissionUtil(object): def get_permissions_for_user_asset(self, user, asset): """ 获取同时包含用户、资产的授权规则 """ user_perm_ids = self.get_permissions_for_user(user, flat=True) - asset_perm_ids = self.get_permissions_for_asset(asset, flat=True) + asset_perm_ids = self.get_permissions_for_assets([asset], flat=True) perm_ids = set(user_perm_ids) & set(asset_perm_ids) perms = self.get_permissions(ids=perm_ids) return perms def get_permissions_for_user_group_asset(self, user_group, asset): user_perm_ids = self.get_permissions_for_user_groups([user_group], flat=True) - asset_perm_ids = self.get_permissions_for_asset(asset, flat=True) + asset_perm_ids = self.get_permissions_for_assets([asset], flat=True) perm_ids = set(user_perm_ids) & set(asset_perm_ids) perms = self.get_permissions(ids=perm_ids) return perms @@ -92,3 +93,15 @@ class AssetPermissionUtil(object): def get_permissions(ids): perms = AssetPermission.objects.filter(id__in=ids).order_by('-date_expired') return perms + + @staticmethod + def transform_to_queryset(objs_or_ids, model): + if not objs_or_ids: + return objs_or_ids + if isinstance(objs_or_ids, QuerySet): + return objs_or_ids + ids = [str(o.id) if isinstance(o, model) else o for o in objs_or_ids] + return model.objects.filter(id__in=ids) + + + diff --git a/apps/perms/utils/user_permission.py b/apps/perms/utils/user_permission.py index 92a42a401..dc90e29bf 100644 --- a/apps/perms/utils/user_permission.py +++ b/apps/perms/utils/user_permission.py @@ -25,6 +25,7 @@ from perms.models import ( AssetPermission, PermNode, UserAssetGrantedTreeNodeRelation ) from users.models import User +from .permission import AssetPermissionUtil NodeFrom = UserAssetGrantedTreeNodeRelation.NodeFrom NODE_ONLY_FIELDS = ('id', 'key', 'parent_key', 'org_id') @@ -32,32 +33,92 @@ NODE_ONLY_FIELDS = ('id', 'key', 'parent_key', 'org_id') logger = get_logger(__name__) -def get_user_all_asset_perm_ids(user) -> set: - asset_perm_ids = set() - user_perm_id = AssetPermission.users.through.objects \ - .filter(user_id=user.id) \ - .values_list('assetpermission_id', flat=True) \ - .distinct() - asset_perm_ids.update(user_perm_id) +class UserPermTreeUtil2: + cache_key_template = 'perms.user.node_tree.built_orgs.user_id:{user_id}' - group_ids = user.groups.through.objects \ - .filter(user_id=user.id) \ - .values_list('usergroup_id', flat=True) \ - .distinct() - group_ids = list(group_ids) - groups_perm_id = AssetPermission.user_groups.through.objects \ - .filter(usergroup_id__in=group_ids) \ - .values_list('assetpermission_id', flat=True) \ - .distinct() - asset_perm_ids.update(groups_perm_id) + def __init__(self, user): + self.user = user + self.orgs = self.user.orgs.distinct() + self.org_ids = [str(o.id) for o in self.orgs] - asset_perm_ids = AssetPermission.objects.filter( - id__in=asset_perm_ids).valid().values_list('id', flat=True) - asset_perm_ids = set(asset_perm_ids) - return asset_perm_ids + def get_cache_key(self, user_id): + return self.cache_key_template.format(user_id=user_id) + + @lazyproperty + def cache_key_user(self): + return self.get_cache_key(self.user.id) + + @lazyproperty + def cache_key_all_user(self): + return self.get_cache_key('*') + + @lazyproperty + def client(self): + return cache.client.get_client(write=True) + + @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() + if not to_refresh_orgs: + logger.info('Not have to refresh orgs') + return + + with UserGrantedTreeRebuildLock(self.user.id): + for org in to_refresh_orgs: + with tmp_to_org(org): + start = time.time() + UserGrantedTreeBuildUtils(self.user).rebuild_user_granted_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) + ) + self.mark_user_orgs_refresh_finished(to_refresh_orgs) + + def clean_user_perm_tree_nodes_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): + 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 + to_refresh_orgs = Organization.objects.filter(id__in=to_refresh_org_ids) + logger.info(f'Need to refresh orgs: {to_refresh_orgs}') + return to_refresh_orgs + + def mark_user_orgs_refresh_finished(self, org_ids): + self.client.sadd(self.cache_key_user, *org_ids) + + # cls + def expire_perm_tree_for_all_user(self): + keys = self.client.keys(self.cache_key_all_user) + with self.client.pipline() as p: + for k in keys: + p.delete(k) + p.execute() + + def expire_perm_tree_for_users_orgs(self, user_ids, org_ids): + org_ids = [str(oid) for oid in org_ids] + with self.client.pipline() as p: + for uid in user_ids: + cache_key = self.get_cache_key(uid) + p.srem(cache_key, *org_ids) + p.execute() + logger.info('Expire perm tree for users: [{}], orgs: [{}]'.format(user_ids, org_ids)) + + def expire_perm_tree_for_nodes_assets(self, node_ids, asset_ids): + node_perm_ids = AssetPermissionUtil().get_permissions_for_nodes(node_ids, flat=True) + asset_perm_ids = AssetPermissionUtil().get_permissions_for_assets(asset_ids, flat=True) + perm_ids = set(node_perm_ids) | set(asset_perm_ids) -class UserGrantedTreeRefreshController: +class UserPermTreeUtil: key_template = 'perms.user.node_tree.built_orgs.user_id:{user_id}' def __init__(self, user): @@ -65,24 +126,49 @@ class UserGrantedTreeRefreshController: self.key = self.key_template.format(user_id=user.id) self.client = self.get_redis_client() - @classmethod - def clean_all_user_tree_built_mark(cls): - """ 清除所有用户已构建树的标记 """ - client = cls.get_redis_client() - key_match = cls.key_template.format(user_id='*') - keys = client.keys(key_match) - with client.pipeline() as p: - for key in keys: - p.delete(key) - p.execute() + @timeit + def refresh_if_need(self, force=False): + user = self.user + orgs = user.orgs.all().distinct() + org_ids = [str(o.id) for o in orgs] - @classmethod - def get_redis_client(cls): - return cache.client.get_client(write=True) + with tmp_to_root_org(): + user_relations = UserAssetGrantedTreeNodeRelation.objects.filter(user=user) + user_legacy_org_relations = user_relations.exclude(org_id__in=org_ids) + user_legacy_org_relations.delete() - def get_need_refresh_org_ids(self): - org_ids = self.client.smembers(self.key) - return {org_id.decode() for org_id in org_ids} + need_refresh_orgs = [] + + if not force and not self.have_need_refresh_orgs(): + return + + with UserGrantedTreeRebuildLock(user_id=user.id): + if force: + orgs = self.orgs + self.set_all_orgs_as_built() + else: + orgs = self.get_need_refresh_orgs_and_fill_up() + + for org in orgs: + with tmp_to_org(org): + t_start = time.time() + logger.info(f'Rebuild user tree: user={self.user} org={current_org}') + utils = UserGrantedTreeBuildUtils(user) + utils.rebuild_user_granted_tree() + logger.info( + f'Rebuild user tree ok: cost={time.time() - t_start} ' + f'user={self.user} org={current_org}' + ) + + @lazyproperty + def org_ids(self): + ret = {str(org.id) for org in self.orgs} + return ret + + @lazyproperty + def orgs(self): + orgs = {*self.user.orgs.all().distinct()} + return orgs def set_all_orgs_as_built(self): self.client.sadd(self.key, *self.org_ids) @@ -99,15 +185,28 @@ class UserGrantedTreeRefreshController: with self.client.pipeline() as p: p.smembers(self.key) p.sadd(self.key, *org_ids) - ret = p.execute() - built_org_ids = {org_id.decode() for org_id in ret[0]} - ids = org_ids - built_org_ids - orgs = {*Organization.objects.filter(id__in=ids)} - logger.info( - f'Need rebuild orgs are {orgs}, built orgs are {ret[0]}, ' - f'all orgs are {org_ids}' - ) - return orgs + old_org_ids, new_orgs_count = p.execute() + old_org_ids = {oid.decode() for oid in old_org_ids} + need_refresh_org_ids = org_ids - old_org_ids + need_refresh_orgs = Organization.objects.filter(id__in=need_refresh_org_ids) + logger.info(f'Need refresh orgs: {need_refresh_orgs}') + return need_refresh_orgs + + # cls + @classmethod + def get_redis_client(cls): + return cache.client.get_client(write=True) + + @classmethod + def clean_all_user_tree_built_mark(cls): + """ 清除所有用户已构建树的标记 """ + client = cls.get_redis_client() + key_match = cls.key_template.format(user_id='*') + keys = client.keys(key_match) + with client.pipeline() as p: + for key in keys: + p.delete(key) + p.execute() @classmethod @on_transaction_commit @@ -190,46 +289,6 @@ class UserGrantedTreeRefreshController: [current_org.id], user_ids ) - @lazyproperty - def org_ids(self): - ret = {str(org.id) for org in self.orgs} - return ret - - @lazyproperty - def orgs(self): - orgs = {*self.user.orgs.all().distinct()} - return orgs - - @timeit - def refresh_if_need(self, force=False): - user = self.user - - with tmp_to_root_org(): - UserAssetGrantedTreeNodeRelation.objects.filter(user=user) \ - .exclude(org_id__in=self.org_ids) \ - .delete() - - if not force and not self.have_need_refresh_orgs(): - return - - with UserGrantedTreeRebuildLock(user_id=user.id): - if force: - orgs = self.orgs - self.set_all_orgs_as_built() - else: - orgs = self.get_need_refresh_orgs_and_fill_up() - - for org in orgs: - with tmp_to_org(org): - t_start = time.time() - logger.info(f'Rebuild user tree: user={self.user} org={current_org}') - utils = UserGrantedTreeBuildUtils(user) - utils.rebuild_user_granted_tree() - logger.info( - f'Rebuild user tree ok: cost={time.time() - t_start} ' - f'user={self.user} org={current_org}' - ) - class UserGrantedUtilsBase: user: User @@ -243,7 +302,7 @@ class UserGrantedUtilsBase: if self._asset_perm_ids: return self._asset_perm_ids - asset_perm_ids = get_user_all_asset_perm_ids(self.user) + asset_perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True) return asset_perm_ids diff --git a/apps/settings/signal_handlers.py b/apps/settings/signal_handlers.py index b04ae4f5e..0f4025c0a 100644 --- a/apps/settings/signal_handlers.py +++ b/apps/settings/signal_handlers.py @@ -38,8 +38,8 @@ def refresh_settings_on_changed(sender, instance=None, **kwargs): if instance.name == 'PERM_SINGLE_ASSET_TO_UNGROUP_NODE': # 清除所有用户授权树已构建的标记，下次访问重新生成 logger.debug('Clean ALL User perm tree built mark') - from perms.utils.asset import UserGrantedTreeRefreshController - UserGrantedTreeRefreshController.clean_all_user_tree_built_mark() + from perms.utils.user_permission import UserPermTreeUtil + UserPermTreeUtil.clean_all_user_tree_built_mark() @receiver(django_ready)