perf: 优化 UserPermTreeUtil 逻辑(进行中)

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):

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)

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.<user_id:{user_id}>'
+    name_template = 'perms.user.asset.node.tree.rebuild.<user_id:{user_id}>'
 
     def __init__(self, user_id):
         name = self.name_template.format(user_id=user_id)

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)

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)

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)
+
+
+

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
 
 

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)