perf: 优化授权的资产,速度快 10 倍

pull/12596/head
ibuler 2024-01-24 15:41:03 +08:00 committed by 老广
parent 34aa48d18c
commit e8bbc44647
5 changed files with 21 additions and 24 deletions

View File

@ -429,7 +429,7 @@ class NodeAssetsMixin(NodeAllAssetsMappingMixin):
@classmethod @classmethod
@timeit @timeit
def get_nodes_all_assets(cls, *nodes): def get_nodes_all_assets(cls, *nodes, distinct=True):
from .asset import Asset from .asset import Asset
node_ids = set() node_ids = set()
descendant_node_query = Q() descendant_node_query = Q()
@ -439,7 +439,10 @@ class NodeAssetsMixin(NodeAllAssetsMappingMixin):
if descendant_node_query: if descendant_node_query:
_ids = Node.objects.order_by().filter(descendant_node_query).values_list('id', flat=True) _ids = Node.objects.order_by().filter(descendant_node_query).values_list('id', flat=True)
node_ids.update(_ids) node_ids.update(_ids)
return Asset.objects.order_by().filter(nodes__id__in=node_ids).distinct() assets = Asset.objects.order_by().filter(nodes__id__in=node_ids)
if distinct:
assets = assets.distinct()
return assets
def get_all_asset_ids(self): def get_all_asset_ids(self):
asset_ids = self.get_all_asset_ids_by_node_key(org_id=self.org_id, node_key=self.key) asset_ids = self.get_all_asset_ids_by_node_key(org_id=self.org_id, node_key=self.key)

View File

@ -69,7 +69,7 @@ def digest_sql_query():
for query in queries: for query in queries:
sql = query['sql'] sql = query['sql']
print(" # {}: {}".format(query['time'], sql[:1000])) print(" # {}: {}".format(query['time'], sql[:1000]))
if len(queries) < 3: if len(queries) < 3:
continue continue
print("- Table: {}".format(table_name)) print("- Table: {}".format(table_name))

View File

@ -38,7 +38,7 @@ class UserPermedAssetRetrieveApi(SelfOrPKUserMixin, RetrieveAPIView):
class BaseUserPermedAssetsApi(SelfOrPKUserMixin, ListAPIView): class BaseUserPermedAssetsApi(SelfOrPKUserMixin, ListAPIView):
ordering = ('name',) ordering = []
search_fields = ('name', 'address', 'comment') search_fields = ('name', 'address', 'comment')
ordering_fields = ("name", "address") ordering_fields = ("name", "address")
filterset_class = AssetFilterSet filterset_class = AssetFilterSet

View File

@ -8,9 +8,9 @@ from rest_framework import serializers
from accounts.models import Account from accounts.models import Account
from assets.const import Category, AllTypes from assets.const import Category, AllTypes
from assets.models import Node, Asset, Platform from assets.models import Node, Asset, Platform
from assets.serializers.asset.common import AssetLabelSerializer, AssetProtocolsPermsSerializer from assets.serializers.asset.common import AssetProtocolsPermsSerializer
from common.serializers.fields import ObjectRelatedField, LabeledChoiceField
from common.serializers import ResourceLabelsMixin from common.serializers import ResourceLabelsMixin
from common.serializers.fields import ObjectRelatedField, LabeledChoiceField
from orgs.mixins.serializers import OrgResourceModelSerializerMixin from orgs.mixins.serializers import OrgResourceModelSerializerMixin
from perms.serializers.permission import ActionChoicesField from perms.serializers.permission import ActionChoicesField

View File

@ -7,10 +7,10 @@ from django.db.models import Q
from rest_framework.utils.encoders import JSONEncoder from rest_framework.utils.encoders import JSONEncoder
from assets.const import AllTypes from assets.const import AllTypes
from assets.models import FavoriteAsset, Asset from assets.models import FavoriteAsset, Asset, Node
from common.utils.common import timeit, get_logger from common.utils.common import timeit, get_logger
from orgs.utils import current_org, tmp_to_root_org from orgs.utils import current_org, tmp_to_root_org
from perms.models import PermNode, UserAssetGrantedTreeNodeRelation from perms.models import PermNode, UserAssetGrantedTreeNodeRelation, AssetPermission
from .permission import AssetPermissionUtil from .permission import AssetPermissionUtil
__all__ = ['AssetPermissionPermAssetUtil', 'UserPermAssetUtil', 'UserPermNodeUtil'] __all__ = ['AssetPermissionPermAssetUtil', 'UserPermAssetUtil', 'UserPermNodeUtil']
@ -21,38 +21,32 @@ logger = get_logger(__name__)
class AssetPermissionPermAssetUtil: class AssetPermissionPermAssetUtil:
def __init__(self, perm_ids): def __init__(self, perm_ids):
self.perm_ids = perm_ids self.perm_ids = set(perm_ids)
def get_all_assets(self): def get_all_assets(self):
node_assets = self.get_perm_nodes_assets() node_assets = self.get_perm_nodes_assets()
direct_assets = self.get_direct_assets() direct_assets = self.get_direct_assets()
# 比原来的查到所有 asset id 再搜索块很多,因为当资产量大的时候,搜索会很慢 # 比原来的查到所有 asset id 再搜索块很多,因为当资产量大的时候,搜索会很慢
return (node_assets | direct_assets).distinct() return (node_assets | direct_assets).order_by().distinct()
@timeit @timeit
def get_perm_nodes_assets(self, flat=False): def get_perm_nodes_assets(self):
""" 获取所有授权节点下的资产 """ """ 获取所有授权节点下的资产 """
from assets.models import Node
from ..models import AssetPermission
nodes_ids = AssetPermission.objects \ nodes_ids = AssetPermission.objects \
.filter(id__in=self.perm_ids) \ .filter(id__in=self.perm_ids) \
.values_list('nodes', flat=True) .values_list('nodes', flat=True)
nodes_ids = set(nodes_ids)
nodes = Node.objects.filter(id__in=nodes_ids).only('id', 'key') nodes = Node.objects.filter(id__in=nodes_ids).only('id', 'key')
assets = PermNode.get_nodes_all_assets(*nodes) assets = PermNode.get_nodes_all_assets(*nodes, distinct=False)
if flat:
return set(assets.values_list('id', flat=True))
return assets return assets
@timeit @timeit
def get_direct_assets(self, flat=False): def get_direct_assets(self):
""" 获取直接授权的资产 """ """ 获取直接授权的资产 """
from ..models import AssetPermission asset_ids = AssetPermission.assets.through.objects \
asset_ids = AssetPermission.objects \ .filter(assetpermission_id__in=self.perm_ids) \
.filter(id__in=self.perm_ids) \ .values_list('asset_id', flat=True)
.values_list('assets', flat=True) assets = Asset.objects.filter(id__in=asset_ids)
assets = Asset.objects.filter(id__in=asset_ids).distinct()
if flat:
return set(assets.values_list('id', flat=True))
return assets return assets