From c9534bb9c4df7bf920dab52b674f417fdea82152 Mon Sep 17 00:00:00 2001 From: ibuler Date: Mon, 20 Feb 2023 23:50:33 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E6=8E=A8=E9=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/accounts/api/account/account.py | 2 +- .../models/automations/change_secret.py | 2 +- .../models/automations/push_account.py | 1 - apps/accounts/serializers/account/account.py | 3 ++ .../serializers/automations/push_account.py | 54 ++----------------- apps/accounts/tasks/common.py | 13 ++--- apps/accounts/tasks/gather_accounts.py | 4 +- apps/accounts/tasks/push_account.py | 31 ++++++----- apps/accounts/tasks/verify_account.py | 4 +- 9 files changed, 33 insertions(+), 81 deletions(-) diff --git a/apps/accounts/api/account/account.py b/apps/accounts/api/account/account.py index 3a5e6b45e..dd57761c2 100644 --- a/apps/accounts/api/account/account.py +++ b/apps/accounts/api/account/account.py @@ -92,7 +92,7 @@ class AccountsTaskCreateAPI(CreateAPIView): def perform_create(self, serializer): data = serializer.validated_data - accounts = self.get_accounts() + accounts = data.get('accounts') account_ids = accounts.values_list('id', flat=True) asset_ids = [account.asset_id for account in accounts] diff --git a/apps/accounts/models/automations/change_secret.py b/apps/accounts/models/automations/change_secret.py index d52fc2cae..76ae9b4f2 100644 --- a/apps/accounts/models/automations/change_secret.py +++ b/apps/accounts/models/automations/change_secret.py @@ -16,11 +16,11 @@ class ChangeSecretMixin(models.Model): choices=SecretType.choices, max_length=16, default=SecretType.PASSWORD, verbose_name=_('Secret type') ) + secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Secret')) secret_strategy = models.CharField( choices=SecretStrategy.choices, max_length=16, default=SecretStrategy.custom, verbose_name=_('Secret strategy') ) - secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Secret')) password_rules = models.JSONField(default=dict, verbose_name=_('Password rules')) ssh_key_change_strategy = models.CharField( choices=SSHKeyStrategy.choices, max_length=16, diff --git a/apps/accounts/models/automations/push_account.py b/apps/accounts/models/automations/push_account.py index 4e4d95d9d..411e8cb12 100644 --- a/apps/accounts/models/automations/push_account.py +++ b/apps/accounts/models/automations/push_account.py @@ -9,7 +9,6 @@ __all__ = ['PushAccountAutomation'] class PushAccountAutomation(ChangeSecretMixin, AccountBaseAutomation): - accounts = None triggers = models.JSONField(max_length=16, default=list, verbose_name=_('Triggers')) username = models.CharField(max_length=128, verbose_name=_('Username')) action = models.CharField(max_length=16, verbose_name=_('Action')) diff --git a/apps/accounts/serializers/account/account.py b/apps/accounts/serializers/account/account.py index c8d281be3..1f9c143bd 100644 --- a/apps/accounts/serializers/account/account.py +++ b/apps/accounts/serializers/account/account.py @@ -144,4 +144,7 @@ class AccountTaskSerializer(serializers.Serializer): ('push', 'push'), ) action = serializers.ChoiceField(choices=ACTION_CHOICES, write_only=True) + accounts = serializers.PrimaryKeyRelatedField( + queryset=Account.objects, required=False, allow_empty=True, many=True + ) task = serializers.CharField(read_only=True) diff --git a/apps/accounts/serializers/automations/push_account.py b/apps/accounts/serializers/automations/push_account.py index aad6041ba..07ebaa2d0 100644 --- a/apps/accounts/serializers/automations/push_account.py +++ b/apps/accounts/serializers/automations/push_account.py @@ -7,58 +7,12 @@ from .change_secret import ( class PushAccountAutomationSerializer(ChangeSecretAutomationSerializer): - # dynamic_username = serializers.BooleanField(label=_('Dynamic username'), default=False) - # triggers = TreeChoicesField( - # choice_cls=TriggerChoice, label=_('Triggers'), - # default=TriggerChoice.all(), - # ) - # action = LabeledChoiceField( - # choices=PushAccountActionChoice.choices, label=_('Action'), - # default=PushAccountActionChoice.create_and_push - # ) - class Meta(ChangeSecretAutomationSerializer.Meta): model = PushAccountAutomation - fields = copy.copy(ChangeSecretAutomationSerializer.Meta.fields) - fields.remove('recipients') - - # fields = ChangeSecretAutomationSerializer.Meta.fields + [ - # 'dynamic_username', 'triggers', 'action' - # ] - - # def validate_username(self, value): - # if self.initial_data.get('dynamic_username'): - # value = '@USER' - # queryset = self.Meta.model.objects.filter(username=value) - # if self.instance: - # queryset = queryset.exclude(id=self.instance.id) - # if queryset.exists(): - # raise serializers.ValidationError(_('Username already exists')) - # return value - # - # def validate_dynamic_username(self, value): - # if not value: - # return value - # queryset = self.Meta.model.objects.filter(username='@USER') - # if self.instance: - # queryset = queryset.exclude(id=self.instance.id) - # if queryset.exists(): - # raise serializers.ValidationError(_('Dynamic username already exists')) - # return value - # - # def validate_triggers(self, value): - # # Now triggers readonly, set all - # return TriggerChoice.all() - # - # def get_field_names(self, declared_fields, info): - # fields = super().get_field_names(declared_fields, info) - # excludes = [ - # 'recipients', 'is_periodic', 'interval', 'crontab', - # 'periodic_display', 'assets', 'nodes' - # ] - # fields = [f for f in fields if f not in excludes] - # fields[fields.index('accounts')] = 'username' - # return fields + fields = [ + n for n in ChangeSecretAutomationSerializer.Meta.fields + if n not in ['recipients'] + ] class PushAccountUpdateAssetSerializer(ChangeSecretUpdateAssetSerializer): diff --git a/apps/accounts/tasks/common.py b/apps/accounts/tasks/common.py index 5c8aeedf5..52f2d13b4 100644 --- a/apps/accounts/tasks/common.py +++ b/apps/accounts/tasks/common.py @@ -4,17 +4,14 @@ from assets.tasks.common import generate_automation_execution_data from common.const.choices import Trigger -def automation_execute_start(task_name, tp, task_snapshot=None): +def quickstart_automation_by_snapshot(task_name, tp, task_snapshot=None): from accounts.models import AutomationExecution data = generate_automation_execution_data(task_name, tp, task_snapshot) - while True: - try: - _id = data['id'] - AutomationExecution.objects.get(id=_id) - data['id'] = str(uuid.uuid4()) - except AutomationExecution.DoesNotExist: - break + pk = data['id'] + if AutomationExecution.objects.exists(id=pk): + data['id'] = str(uuid.uuid4()) + execution = AutomationExecution.objects.create( trigger=Trigger.manual, **data ) diff --git a/apps/accounts/tasks/gather_accounts.py b/apps/accounts/tasks/gather_accounts.py index cf2283e6f..ceead3f9d 100644 --- a/apps/accounts/tasks/gather_accounts.py +++ b/apps/accounts/tasks/gather_accounts.py @@ -4,7 +4,7 @@ from django.utils.translation import gettext_lazy as _ from django.utils.translation import gettext_noop from accounts.const import AutomationTypes -from accounts.tasks.common import automation_execute_start +from accounts.tasks.common import quickstart_automation_by_snapshot from assets.models import Node from common.utils import get_logger from orgs.utils import org_aware_func @@ -22,7 +22,7 @@ def gather_asset_accounts_util(nodes, task_name): 'nodes': [str(node.id) for node in nodes], } tp = AutomationTypes.verify_account - automation_execute_start(task_name, tp, task_snapshot) + quickstart_automation_by_snapshot(task_name, tp, task_snapshot) @shared_task( diff --git a/apps/accounts/tasks/push_account.py b/apps/accounts/tasks/push_account.py index a1c4b6f7c..ab7832b2f 100644 --- a/apps/accounts/tasks/push_account.py +++ b/apps/accounts/tasks/push_account.py @@ -1,8 +1,9 @@ from celery import shared_task +from collections import defaultdict from django.utils.translation import gettext_noop, ugettext_lazy as _ from accounts.const import AutomationTypes -from accounts.tasks.common import automation_execute_start +from accounts.tasks.common import quickstart_automation_by_snapshot from common.utils import get_logger logger = get_logger(__file__) @@ -11,31 +12,29 @@ __all__ = [ ] -def push_util(account, asset_ids, task_name): - task_snapshot = { - 'secret': account.secret, - 'secret_type': account.secret_type, - 'accounts': [account.username], - 'assets': asset_ids, - } - tp = AutomationTypes.push_account - automation_execute_start(task_name, tp, task_snapshot) - - @shared_task( queue="ansible", verbose_name=_('Push accounts to assets'), activity_callback=lambda self, account_ids, asset_ids: (account_ids, None) ) -def push_accounts_to_assets_task(account_ids, asset_ids): +def push_accounts_to_assets_task(account_ids): from accounts.models import PushAccountAutomation - from assets.models import Asset from accounts.models import Account - assets = Asset.objects.filter(id__in=asset_ids) accounts = Account.objects.filter(id__in=account_ids) task_name = gettext_noop("Push accounts to assets") task_name = PushAccountAutomation.generate_unique_name(task_name) + account_asset_mapper = defaultdict(set) for account in accounts: - push_util(account, assets, task_name) + account_asset_mapper[account.username].add(account.asset) + + for username, assets in account_asset_mapper.items(): + task_snapshot = { + 'secret': account.secret, + 'secret_type': account.secret_type, + 'accounts': [account.username], + 'assets': asset_ids, + } + tp = AutomationTypes.push_account + quickstart_automation_by_snapshot(task_name, tp, task_snapshot) \ No newline at end of file diff --git a/apps/accounts/tasks/verify_account.py b/apps/accounts/tasks/verify_account.py index f5221d975..8bb656e29 100644 --- a/apps/accounts/tasks/verify_account.py +++ b/apps/accounts/tasks/verify_account.py @@ -3,7 +3,7 @@ from django.utils.translation import gettext_lazy as _ from django.utils.translation import gettext_noop from accounts.const import AutomationTypes -from accounts.tasks.common import automation_execute_start +from accounts.tasks.common import quickstart_automation_by_snapshot from assets.const import GATEWAY_NAME from common.utils import get_logger from orgs.utils import org_aware_func @@ -22,7 +22,7 @@ def verify_connectivity_util(assets, tp, accounts, task_name): 'accounts': account_usernames, 'assets': [str(asset.id) for asset in assets], } - automation_execute_start(task_name, tp, task_snapshot) + quickstart_automation_by_snapshot(task_name, tp, task_snapshot) @org_aware_func("assets")