mirror of https://github.com/jumpserver/jumpserver
perf: 优化用户授权树工具类存放目录 user_perm_tree 文件中
Bai
parent
5e5061a825
commit
26d9cdc50d
|
|
@ -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']
|
||||
|
|
|
|||
|
|
@ -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__)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from .permission import *
|
||||
from .user_permission import *
|
||||
from .account import *
|
||||
from .user_perm_tree import *
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue