From 30b89e5cc93d82630e5c293fceb52b422ce8b87a Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Wed, 12 Apr 2023 17:59:13 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E8=B4=A6=E5=8F=B7=E6=A8=A1=E7=89=88?= =?UTF-8?q?=E6=9B=B4=E6=96=B0=20(#10184)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: feng <1304903146@qq.com> --- apps/accounts/const/account.py | 1 + apps/accounts/serializers/account/account.py | 40 ++++++-------- apps/accounts/serializers/account/template.py | 55 +++++++++++++++---- apps/assets/models/asset/common.py | 16 ++++++ 4 files changed, 76 insertions(+), 36 deletions(-) diff --git a/apps/accounts/const/account.py b/apps/accounts/const/account.py index 29185b233..55fa02d80 100644 --- a/apps/accounts/const/account.py +++ b/apps/accounts/const/account.py @@ -18,6 +18,7 @@ class AliasAccount(TextChoices): class Source(TextChoices): LOCAL = 'local', _('Local') COLLECTED = 'collected', _('Collected') + TEMPLATE = 'template', _('Template') class AccountInvalidPolicy(TextChoices): diff --git a/apps/accounts/serializers/account/account.py b/apps/accounts/serializers/account/account.py index 0a1cf3f83..b2a982fd4 100644 --- a/apps/accounts/serializers/account/account.py +++ b/apps/accounts/serializers/account/account.py @@ -1,5 +1,4 @@ import uuid -from collections import defaultdict from django.db import IntegrityError from django.db.models import Q @@ -10,7 +9,7 @@ from rest_framework.validators import UniqueTogetherValidator from accounts.const import SecretType, Source, AccountInvalidPolicy from accounts.models import Account, AccountTemplate from accounts.tasks import push_accounts_to_assets_task -from assets.const import Category, AllTypes, Protocol +from assets.const import Category, AllTypes from assets.models import Asset from common.serializers import SecretReadableMixin from common.serializers.fields import ObjectRelatedField, LabeledChoiceField @@ -80,12 +79,12 @@ class AccountCreateUpdateSerializerMixin(serializers.Serializer): raise serializers.ValidationError({'template': 'Template not found'}) # Set initial data from template - ignore_fields = ['id', 'name', 'date_created', 'date_updated', 'org_id'] + ignore_fields = ['id', 'date_created', 'date_updated', 'org_id'] field_names = [ field.name for field in template._meta.fields if field.name not in ignore_fields ] - attrs = {'source': 'template', 'source_id': template.id} + attrs = {'source': Source.TEMPLATE, 'source_id': str(template.id)} for name in field_names: value = getattr(template, name, None) if value is None: @@ -135,6 +134,16 @@ class AccountCreateUpdateSerializerMixin(serializers.Serializer): else: raise serializers.ValidationError('Account already exists') + def validate(self, attrs): + attrs = super().validate(attrs) + if self.instance: + return attrs + + if 'source' in self.initial_data: + attrs['source'] = self.initial_data['source'] + attrs['source_id'] = self.initial_data['source_id'] + return attrs + def create(self, validated_data): push_now = validated_data.pop('push_now', None) instance, stat = self.do_create(validated_data) @@ -146,6 +155,7 @@ class AccountCreateUpdateSerializerMixin(serializers.Serializer): validated_data.pop('username', None) validated_data.pop('on_invalid', None) push_now = validated_data.pop('push_now', None) + validated_data['source_id'] = None instance = super().update(instance, validated_data) self.push_account_if_need(instance, push_now, 'updated') return instance @@ -233,25 +243,6 @@ class AssetAccountBulkSerializer(AccountCreateUpdateSerializerMixin, serializers initial_data = self.initial_data self.from_template_if_need(initial_data) - @staticmethod - def _get_valid_secret_type_assets(assets, secret_type): - if isinstance(assets, list): - asset_ids = [a.id for a in assets] - assets = Asset.objects.filter(id__in=asset_ids) - - asset_protocol = assets.prefetch_related('protocols').values_list('id', 'protocols__name') - protocol_secret_types_map = Protocol.protocol_secret_types() - asset_secret_types_mapp = defaultdict(set) - - for asset_id, protocol in asset_protocol: - secret_types = set(protocol_secret_types_map.get(protocol, [])) - asset_secret_types_mapp[asset_id].update(secret_types) - - return [ - asset for asset in assets - if secret_type in asset_secret_types_mapp.get(asset.id, []) - ] - @staticmethod def get_filter_lookup(vd): return { @@ -314,7 +305,8 @@ class AssetAccountBulkSerializer(AccountCreateUpdateSerializerMixin, serializers vd['name'] = vd.get('username') create_handler = self.get_create_handler(on_invalid) - secret_type_supports = self._get_valid_secret_type_assets(assets, secret_type) + asset_ids = [asset.id for asset in assets] + secret_type_supports = Asset.get_secret_type_assets(asset_ids, secret_type) _results = {} for asset in assets: diff --git a/apps/accounts/serializers/account/template.py b/apps/accounts/serializers/account/template.py index a72565cf5..45190f051 100644 --- a/apps/accounts/serializers/account/template.py +++ b/apps/accounts/serializers/account/template.py @@ -1,4 +1,5 @@ -from accounts.models import AccountTemplate +from accounts.models import AccountTemplate, Account +from assets.models import Asset from common.serializers import SecretReadableMixin from .base import BaseAccountSerializer @@ -7,17 +8,47 @@ class AccountTemplateSerializer(BaseAccountSerializer): class Meta(BaseAccountSerializer.Meta): model = AccountTemplate - # @classmethod - # def validate_required(cls, attrs): - # # TODO 选择模版后检查一些必填项 - # required_field_dict = {} - # error = _('This field is required.') - # for k, v in cls().fields.items(): - # if v.required and k not in attrs: - # required_field_dict[k] = error - # if not required_field_dict: - # return - # raise serializers.ValidationError(required_field_dict) + @staticmethod + def bulk_update_accounts(instance, diff): + accounts = Account.objects.filter(source_id=instance.id) + if not accounts: + return + + secret_type = diff.pop('secret_type', None) + diff.pop('secret', None) + update_accounts = [] + for account in accounts: + for field, value in diff.items(): + setattr(account, field, value) + update_accounts.append(account) + if update_accounts: + Account.objects.bulk_update(update_accounts, diff.keys()) + + if secret_type is None: + return + + update_accounts = [] + asset_ids = accounts.values_list('asset_id', flat=True) + secret_type_supports = Asset.get_secret_type_assets(asset_ids, secret_type) + asset_ids_supports = [asset.id for asset in secret_type_supports] + for account in accounts: + asset_id = account.asset_id + if asset_id not in asset_ids_supports: + continue + account.secret_type = secret_type + account.secret = instance.secret + update_accounts.append(account) + if update_accounts: + Account.objects.bulk_update(update_accounts, ['secret', 'secret_type']) + + def update(self, instance, validated_data): + diff = { + k: v for k, v in validated_data.items() + if getattr(instance, k) != v + } + instance = super().update(instance, validated_data) + self.bulk_update_accounts(instance, diff) + return instance class AccountTemplateSecretSerializer(SecretReadableMixin, AccountTemplateSerializer): diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index 05fada2e3..abfc1be8b 100644 --- a/apps/assets/models/asset/common.py +++ b/apps/assets/models/asset/common.py @@ -271,6 +271,22 @@ class Asset(NodesRelationMixin, AbsConnectivity, JMSOrgBaseModel): tree_node = TreeNode(**data) return tree_node + @staticmethod + def get_secret_type_assets(asset_ids, secret_type): + assets = Asset.objects.filter(id__in=asset_ids) + asset_protocol = assets.prefetch_related('protocols').values_list('id', 'protocols__name') + protocol_secret_types_map = const.Protocol.protocol_secret_types() + asset_secret_types_mapp = defaultdict(set) + + for asset_id, protocol in asset_protocol: + secret_types = set(protocol_secret_types_map.get(protocol, [])) + asset_secret_types_mapp[asset_id].update(secret_types) + + return [ + asset for asset in assets + if secret_type in asset_secret_types_mapp.get(asset.id, []) + ] + class Meta: unique_together = [('org_id', 'name')] verbose_name = _("Asset")