From 26d9cdc50d8eef7d34e0666118bb59a2ef4e2189 Mon Sep 17 00:00:00 2001 From: Bai Date: Fri, 9 Dec 2022 14:03:00 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E7=94=A8=E6=88=B7?= =?UTF-8?q?=E6=8E=88=E6=9D=83=E6=A0=91=E5=B7=A5=E5=85=B7=E7=B1=BB=E5=AD=98?= =?UTF-8?q?=E6=94=BE=E7=9B=AE=E5=BD=95=20user=5Fperm=5Ftree=20=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E4=B8=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/perms/api/user_permission/tree/mixin.py | 2 +- apps/perms/signal_handlers/refresh_perms.py | 2 +- apps/perms/tasks.py | 2 +- apps/perms/utils/__init__.py | 1 + apps/perms/utils/user_perm_tree.py | 145 +++++++++++++++++++ apps/perms/utils/user_permission.py | 125 ---------------- apps/settings/signal_handlers.py | 2 +- 7 files changed, 150 insertions(+), 129 deletions(-) create mode 100644 apps/perms/utils/user_perm_tree.py diff --git a/apps/perms/api/user_permission/tree/mixin.py b/apps/perms/api/user_permission/tree/mixin.py index b5f9cc5b6..d78a29b04 100644 --- a/apps/perms/api/user_permission/tree/mixin.py +++ b/apps/perms/api/user_permission/tree/mixin.py @@ -3,7 +3,7 @@ from rest_framework.request import Request from users.models import User from common.http import is_true -from perms.utils.user_permission import UserPermTreeRefreshUtil +from perms.utils import UserPermTreeRefreshUtil __all__ = ['RebuildTreeMixin'] diff --git a/apps/perms/signal_handlers/refresh_perms.py b/apps/perms/signal_handlers/refresh_perms.py index 0f37aee79..88a5fe674 100644 --- a/apps/perms/signal_handlers/refresh_perms.py +++ b/apps/perms/signal_handlers/refresh_perms.py @@ -10,7 +10,7 @@ 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 UserPermTreeExpireUtil +from perms.utils import UserPermTreeExpireUtil logger = get_logger(__file__) diff --git a/apps/perms/tasks.py b/apps/perms/tasks.py index 781feffbd..3df7a1728 100644 --- a/apps/perms/tasks.py +++ b/apps/perms/tasks.py @@ -14,7 +14,7 @@ from common.utils.timezone import local_now, dt_parser from common.const.crontab import CRONTAB_AT_AM_TEN from perms.models import AssetPermission -from perms.utils.user_permission import UserPermTreeExpireUtil +from perms.utils import UserPermTreeExpireUtil from perms.notifications import ( PermedAssetsWillExpireUserMsg, AssetPermsWillExpireForOrgAdminMsg, diff --git a/apps/perms/utils/__init__.py b/apps/perms/utils/__init__.py index fc2a94e88..8563b21e9 100644 --- a/apps/perms/utils/__init__.py +++ b/apps/perms/utils/__init__.py @@ -1,3 +1,4 @@ from .permission import * from .user_permission import * from .account import * +from .user_perm_tree import * diff --git a/apps/perms/utils/user_perm_tree.py b/apps/perms/utils/user_perm_tree.py new file mode 100644 index 000000000..d45b6a325 --- /dev/null +++ b/apps/perms/utils/user_perm_tree.py @@ -0,0 +1,145 @@ +import time +from collections import defaultdict + +from django.core.cache import cache + +from users.models import User +from orgs.models import Organization +from orgs.utils import ( + 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 perms.locks import UserGrantedTreeRebuildLock +from perms.models import ( + AssetPermission, + UserAssetGrantedTreeNodeRelation +) +from .permission import AssetPermissionUtil + + +logger = get_logger(__name__) + +__all__ = ['UserPermTreeRefreshUtil', 'UserPermTreeExpireUtil'] + + +class _UserPermTreeCacheMixin: + """ 缓存数据 users: {org_id, org_id }, 记录用户授权树已经构建完成的组织集合 """ + cache_key_template = 'perms.user.node_tree.built_orgs.user_id:{user_id}' + + def get_cache_key(self, user_id): + return self.cache_key_template.format(user_id=user_id) + + @lazyproperty + def client(self): + return cache.client.get_client(write=True) + + +class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin): + """ 用户授权树刷新工具, 针对某一个用户授权树的刷新 """ + + def __init__(self, user): + self.user = user + self.orgs = self.user.orgs.distinct() + self.org_ids = [str(o.id) for o in self.orgs] + + @lazyproperty + def cache_key_user(self): + return self.get_cache_key(self.user.id) + + @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: + 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): + 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) + ) + + 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) + + +class UserPermTreeExpireUtil(_UserPermTreeCacheMixin): + """ 用户授权树过期工具 """ + + @lazyproperty + def cache_key_all_user(self): + return self.get_cache_key('*') + + 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) + self.expire_perm_tree_for_perms(perm_ids) + + @tmp_to_root_org() + def expire_perm_tree_for_perms(self, perm_ids): + org_perm_ids = AssetPermission.objects.filter(id__in=perm_ids).values_list('org_id', 'id') + org_perms_mapper = defaultdict(set) + for org_id, perm_id in org_perm_ids: + org_perms_mapper[org_id].add(perm_id) + for org_id, perms_id in org_perms_mapper.items(): + org_ids = [org_id] + user_ids = AssetPermission.get_all_users_for_perms(perm_ids, flat=True) + self.expire_perm_tree_for_users_orgs(user_ids, org_ids) + + def expire_perm_tree_for_user_group(self, user_group): + group_ids = [user_group.id] + org_ids = [user_group.org_id] + self.expire_perm_tree_for_user_groups_orgs(group_ids, org_ids) + + def expire_perm_tree_for_user_groups_orgs(self, group_ids, org_ids): + user_ids = User.groups.through.objects.filter(usergroup_id__in=group_ids) \ + .values_list('user_id', flat=True).distinct() + self.expire_perm_tree_for_users_orgs(user_ids, org_ids) + + @on_transaction_commit + 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_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() + logger.info('Expire all user perm tree') + diff --git a/apps/perms/utils/user_permission.py b/apps/perms/utils/user_permission.py index 9ba517125..81b4fc00e 100644 --- a/apps/perms/utils/user_permission.py +++ b/apps/perms/utils/user_permission.py @@ -1,9 +1,7 @@ -import time from collections import defaultdict from typing import List, Tuple from django.conf import settings -from django.core.cache import cache from django.db.models import Q, QuerySet from django.utils.translation import gettext as _ @@ -15,19 +13,15 @@ from assets.models import ( AssetQuerySet, NodeQuerySet ) -from orgs.models import Organization from orgs.utils import ( tmp_to_org, current_org, ensure_in_real_or_default_org, - tmp_to_root_org ) from common.db.models import output_as_string, UnionQuerySet -from common.decorator import on_transaction_commit from common.utils import get_logger from common.utils.common import lazyproperty, timeit -from perms.locks import UserGrantedTreeRebuildLock from perms.models import ( AssetPermission, PermNode, @@ -41,125 +35,6 @@ NODE_ONLY_FIELDS = ('id', 'key', 'parent_key', 'org_id') logger = get_logger(__name__) -class UserPermTreeCacheMixin: - """ 缓存数据 users: {org_id, org_id }, 记录用户授权树已经构建完成的组织集合 """ - - cache_key_template = 'perms.user.node_tree.built_orgs.user_id:{user_id}' - - def get_cache_key(self, user_id): - return self.cache_key_template.format(user_id=user_id) - - @lazyproperty - def client(self): - return cache.client.get_client(write=True) - - -class UserPermTreeRefreshUtil(UserPermTreeCacheMixin): - """ 用户授权树刷新工具, 针对某一个用户授权树的刷新 """ - - def __init__(self, user): - self.user = user - self.orgs = self.user.orgs.distinct() - self.org_ids = [str(o.id) for o in self.orgs] - - @lazyproperty - def cache_key_user(self): - return self.get_cache_key(self.user.id) - - @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: - 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): - 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) - ) - - 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) - - -class UserPermTreeExpireUtil(UserPermTreeCacheMixin): - """ 用户授权树过期工具 """ - - @lazyproperty - def cache_key_all_user(self): - return self.get_cache_key('*') - - 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) - self.expire_perm_tree_for_perms(perm_ids) - - @tmp_to_root_org() - def expire_perm_tree_for_perms(self, perm_ids): - org_perm_ids = AssetPermission.objects.filter(id__in=perm_ids).values_list('org_id', 'id') - org_perms_mapper = defaultdict(set) - for org_id, perm_id in org_perm_ids: - org_perms_mapper[org_id].add(perm_id) - for org_id, perms_id in org_perms_mapper.items(): - org_ids = [org_id] - user_ids = AssetPermission.get_all_users_for_perms(perm_ids, flat=True) - self.expire_perm_tree_for_users_orgs(user_ids, org_ids) - - def expire_perm_tree_for_user_group(self, user_group): - group_ids = [user_group.id] - org_ids = [user_group.org_id] - self.expire_perm_tree_for_user_groups_orgs(group_ids, org_ids) - - def expire_perm_tree_for_user_groups_orgs(self, group_ids, org_ids): - user_ids = User.groups.through.objects.filter(usergroup_id__in=group_ids)\ - .values_list('user_id', flat=True).distinct() - self.expire_perm_tree_for_users_orgs(user_ids, org_ids) - - @on_transaction_commit - 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_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() - logger.info('Expire all user perm tree') - - class UserGrantedUtilsBase: user: User diff --git a/apps/settings/signal_handlers.py b/apps/settings/signal_handlers.py index 58447ba79..7568de368 100644 --- a/apps/settings/signal_handlers.py +++ b/apps/settings/signal_handlers.py @@ -37,7 +37,7 @@ def refresh_settings_on_changed(sender, instance=None, **kwargs): if instance.is_name('PERM_SINGLE_ASSET_TO_UNGROUP_NODE'): """ 过期所有用户授权树 """ logger.debug('Expire all user perm tree') - from perms.utils.user_permission import UserPermTreeExpireUtil + from perms.utils import UserPermTreeExpireUtil UserPermTreeExpireUtil().expire_perm_tree_for_all_user()