perf: 优化账号推送

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]
 

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,

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

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)

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

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
     )

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(

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)

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