perf: 优化用户详情页授权列表加载速度&添加可重入锁

apps/perms/api/asset/asset_permission.py

@@ -26,12 +26,6 @@ class AssetPermissionViewSet(BasePermissionViewSet):
         'node_id', 'node', 'asset_id', 'hostname', 'ip'
     ]
 
-    def get_queryset(self):
-        queryset = super().get_queryset().prefetch_related(
-            "nodes", "assets", "users", "user_groups", "system_users"
-        )
-        return queryset
-
     def filter_node(self, queryset):
         node_id = self.request.query_params.get('node_id')
         node_name = self.request.query_params.get('node')

apps/perms/api/asset/user_permission/user_permission_nodes_with_assets.py

@@ -14,7 +14,6 @@ from .mixin import RoleUserMixin, RoleAdminMixin
 from perms.utils.asset.user_permission import (
     UserGrantedTreeBuildUtils, get_user_all_asset_perm_ids,
     UserGrantedNodesQueryUtils, UserGrantedAssetsQueryUtils,
-    QuerySetStage,
 )
 from perms.models import AssetPermission, PermNode
 from assets.models import Asset
@@ -44,10 +43,10 @@ class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView):
     def add_favorite_resource(self, data: list, nodes_query_utils, assets_query_utils):
         favorite_node = nodes_query_utils.get_favorite_node()
 
-        qs_state = QuerySetStage().annotate(
+        favorite_assets = assets_query_utils.get_favorite_assets()
+        favorite_assets = favorite_assets.annotate(
             parent_key=Value(favorite_node.key, output_field=CharField())
         ).prefetch_related('platform')
-        favorite_assets = assets_query_utils.get_favorite_assets(qs_stage=qs_state, only=())
 
         data.extend(self.serialize_nodes([favorite_node], with_asset_amount=True))
         data.extend(self.serialize_assets(favorite_assets))
@@ -59,13 +58,11 @@ class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView):
         data.extend(self.serialize_nodes(nodes, with_asset_amount=True))
 
     def add_assets(self, data: list, assets_query_utils: UserGrantedAssetsQueryUtils):
-        qs_stage = QuerySetStage().annotate(parent_key=F('nodes__key')).prefetch_related('platform')
-
         if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
-            all_assets = assets_query_utils.get_direct_granted_nodes_assets(qs_stage=qs_stage)
+            all_assets = assets_query_utils.get_direct_granted_nodes_assets()
         else:
-            all_assets = assets_query_utils.get_all_granted_assets(qs_stage=qs_stage)
-
+            all_assets = assets_query_utils.get_all_granted_assets()
+        all_assets = all_assets.annotate(parent_key=F('nodes__key')).prefetch_related('platform')
         data.extend(self.serialize_assets(all_assets))
 
     @tmp_to_root_org()
@@ -144,8 +141,6 @@ class GrantedNodeChildrenWithAssetsAsTreeApiMixin(SerializeToTreeNodeMixin,
             assets = assets_query_utils.get_node_assets(key)
         assets = assets.prefetch_related('platform')
 
-        user = self.user
-
         tree_nodes = self.serialize_nodes(nodes, with_asset_amount=True)
         tree_assets = self.serialize_assets(assets, key)
         return Response(data=[*tree_nodes, *tree_assets])

apps/perms/api/base.py

@@ -45,7 +45,7 @@ class BasePermissionViewSet(OrgBulkModelViewSet):
         if not self.is_query_all():
             queryset = queryset.filter(users=user)
             return queryset
-        groups = user.groups.all()
+        groups = list(user.groups.all().values_list('id', flat=True))
         queryset = queryset.filter(
             Q(users=user) | Q(user_groups__in=groups)
         ).distinct()

apps/perms/migrations/0018_auto_20210208_1515.py

@@ -1,4 +1,4 @@
-# Generated by Django 3.1 on 2021-02-04 09:49
+# Generated by Django 3.1 on 2021-02-08 07:15
 
 import assets.models.node
 from django.conf import settings
@@ -9,8 +9,8 @@ import django.db.models.deletion
 class Migration(migrations.Migration):
 
     dependencies = [
-        ('assets', '0066_remove_node_assets_amount'),
         migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('assets', '0065_auto_20210121_1549'),
         ('perms', '0017_auto_20210104_0435'),
     ]
 

apps/perms/pagination.py

@@ -1,37 +1,17 @@
-from rest_framework.pagination import LimitOffsetPagination
+from django.conf import settings
 from rest_framework.request import Request
-from django.db.models import Sum
 
+from assets.pagination import AssetPaginationBase
 from perms.models import UserAssetGrantedTreeNodeRelation
 from common.utils import get_logger
 
 logger = get_logger(__name__)
 
 
-class GrantedAssetPaginationBase(LimitOffsetPagination):
-
-    def paginate_queryset(self, queryset, request: Request, view=None):
-        self._request = request
-        self._view = view
-        self._user = request.user
-        return super().paginate_queryset(queryset, request, view=None)
-
-    def get_count(self, queryset):
-        exclude_query_params = {
-            self.limit_query_param,
-            self.offset_query_param,
-            'key', 'all', 'show_current_asset',
-            'cache_policy', 'display', 'draw',
-            'order',
-        }
-        for k, v in self._request.query_params.items():
-            if k not in exclude_query_params and v is not None:
-                logger.warn(f'Not hit node.assets_amount because find a unknow query_param `{k}` -> {self._request.get_full_path()}')
-                return super().get_count(queryset)
-        return self.get_count_from_nodes(queryset)
-
-    def get_count_from_nodes(self, queryset):
-        raise NotImplementedError
+class GrantedAssetPaginationBase(AssetPaginationBase):
+    def init_attrs(self, queryset, request: Request, view=None):
+        super().init_attrs(queryset, request, view)
+        self._user = view.user
 
 
 class NodeGrantedAssetPagination(GrantedAssetPaginationBase):
@@ -42,11 +22,13 @@ class NodeGrantedAssetPagination(GrantedAssetPaginationBase):
             return node.assets_amount
         else:
             logger.warn(f'Not hit node.assets_amount[{node}] because {self._view} not has `pagination_node` -> {self._request.get_full_path()}')
-            return super().get_count(queryset)
+            return None
 
 
 class AllGrantedAssetPagination(GrantedAssetPaginationBase):
     def get_count_from_nodes(self, queryset):
+        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
+            return None
         assets_amount = sum(UserAssetGrantedTreeNodeRelation.objects.filter(
             user=self._user, node_parent_key=''
         ).values_list('node_assets_amount', flat=True))

apps/perms/serializers/asset/permission.py

@@ -3,9 +3,12 @@
 
 from rest_framework import serializers
 from django.utils.translation import ugettext_lazy as _
+from django.db.models import Prefetch
 
 from orgs.mixins.serializers import BulkOrgResourceModelSerializer
 from perms.models import AssetPermission, Action
+from assets.models import Asset, Node, SystemUser
+from users.models import User, UserGroup
 
 __all__ = [
     'AssetPermissionSerializer',
@@ -68,5 +71,11 @@ class AssetPermissionSerializer(BulkOrgResourceModelSerializer):
     @classmethod
     def setup_eager_loading(cls, queryset):
         """ Perform necessary eager loading of data. """
-        queryset = queryset.prefetch_related('users', 'user_groups', 'assets', 'nodes', 'system_users')
+        queryset = queryset.prefetch_related(
+            Prefetch('system_users', queryset=SystemUser.objects.only('id')),
+            Prefetch('user_groups', queryset=UserGroup.objects.only('id')),
+            Prefetch('users', queryset=User.objects.only('id')),
+            Prefetch('assets', queryset=Asset.objects.only('id')),
+            Prefetch('nodes', queryset=Node.objects.only('id'))
+        )
         return queryset

apps/perms/utils/asset/user_permission.py

@@ -115,8 +115,8 @@ class UnionQuerySet(QuerySet):
     def __getitem__(self, item):
         return self.__execute()[item]
 
-    def __next__(self):
-        return next(self.__execute())
+    def __iter__(self):
+        return iter(self.__execute())
 
     @classmethod
     def test_it(cls):
@@ -299,12 +299,12 @@ class UserGrantedTreeRefreshController:
         cls.remove_builed_orgs_from_users(orgs_id, users_id)
 
     @classmethod
+    @ensure_in_real_or_default_org
     def add_need_refresh_on_nodes_assets_relate_change(cls, node_ids, asset_ids):
         """
         1，计算与这些资产有关的授权
         2，计算与这些节点以及祖先节点有关的授权
         """
-        ensure_in_real_or_default_org()
 
         node_ids = set(node_ids)
         ancestor_node_keys = set()
@@ -340,8 +340,8 @@ class UserGrantedTreeRefreshController:
                 cls.add_need_refresh_by_asset_perm_ids(perm_ids)
 
     @classmethod
+    @ensure_in_real_or_default_org
     def add_need_refresh_by_asset_perm_ids(cls, asset_perm_ids):
-        ensure_in_real_or_default_org()
 
         group_ids = AssetPermission.user_groups.through.objects.filter(
             assetpermission_id__in=asset_perm_ids
@@ -429,8 +429,8 @@ class UserGrantedTreeBuildUtils(UserGrantedUtilsBase):
         return asset_ids
 
     @timeit
+    @ensure_in_real_or_default_org
     def rebuild_user_granted_tree(self):
-        ensure_in_real_or_default_org()
         logger.info(f'Rebuild user:{self.user} tree in org:{current_org}')
 
         user = self.user
@@ -618,13 +618,13 @@ class UserGrantedTreeBuildUtils(UserGrantedUtilsBase):
 
 class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase):
 
-    def get_favorite_assets(self, only=('id', )) -> QuerySet:
+    def get_favorite_assets(self) -> QuerySet:
         favorite_asset_ids = FavoriteAsset.objects.filter(
             user=self.user
         ).values_list('asset_id', flat=True)
         favorite_asset_ids = list(favorite_asset_ids)
         assets = self.get_all_granted_assets()
-        assets = assets.filter(id__in=favorite_asset_ids).only(*only)
+        assets = assets.filter(id__in=favorite_asset_ids)
         return assets
 
     def get_ungroup_assets(self) -> AssetQuerySet:
@@ -670,7 +670,7 @@ class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase):
         granted_status = node.get_granted_status(self.user)
 
         if granted_status == NodeFrom.granted:
-            assets = Asset.objects.order_by().filter(nodes_id=node.id)
+            assets = Asset.objects.order_by().filter(nodes__id=node.id)
             return assets
         elif granted_status == NodeFrom.asset:
             return self._get_indirect_granted_node_assets(node.id)
@@ -678,7 +678,7 @@ class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase):
             return Asset.objects.none()
 
     def _get_indirect_granted_node_assets(self, id) -> AssetQuerySet:
-        assets = Asset.objects.order_by().filter(nodes_id=id) & self.get_direct_granted_assets()
+        assets = Asset.objects.order_by().filter(nodes__id=id).distinct() & self.get_direct_granted_assets()
         return assets
 
     def _get_indirect_granted_node_all_assets(self, node) -> QuerySet: