perf： 优化 tasks

apps/accounts/api/account/account.py

@@ -6,7 +6,7 @@ from rest_framework.response import Response
 from accounts import serializers
 from accounts.filters import AccountFilterSet
 from accounts.models import Account
-from accounts.tasks import verify_accounts_connectivity
+from accounts.tasks import verify_accounts_connectivity_task, push_accounts_to_assets_task
 from assets.models import Asset
 from authentication.const import ConfirmType
 from common.permissions import UserConfirmation
@@ -15,7 +15,7 @@ from orgs.mixins.api import OrgBulkModelViewSet
 
 __all__ = [
     'AccountViewSet', 'AccountSecretsViewSet',
-    'AccountTaskCreateAPI', 'AccountHistoriesSecretAPI'
+    'AccountsTaskCreateAPI', 'AccountHistoriesSecretAPI'
 ]
 
 from rbac.permissions import RBACPermission
@@ -29,7 +29,6 @@ class AccountViewSet(OrgBulkModelViewSet):
         'default': serializers.AccountSerializer,
     }
     rbac_perms = {
-        'verify_account': 'accounts.test_account',
         'partial_update': ['accounts.change_account'],
         'su_from_accounts': 'accounts.view_account',
     }
@@ -49,14 +48,6 @@ class AccountViewSet(OrgBulkModelViewSet):
         serializer = serializers.AccountSerializer(accounts, many=True)
         return Response(data=serializer.data)
 
-    @action(methods=['post'], detail=True, url_path='verify')
-    def verify_account(self, request, *args, **kwargs):
-        account = super().get_object()
-        account_ids = [account.id]
-        asset_ids = [account.asset_id]
-        task = verify_accounts_connectivity.delay(account_ids, asset_ids)
-        return Response(data={'task': task.id})
-
 
 class AccountSecretsViewSet(RecordViewLogMixin, AccountViewSet):
     """
@@ -86,7 +77,7 @@ class AccountHistoriesSecretAPI(RecordViewLogMixin, ListAPIView):
         return self.model.objects.filter(id=self.kwargs.get('pk'))
 
 
-class AccountTaskCreateAPI(CreateAPIView):
+class AccountsTaskCreateAPI(CreateAPIView):
     serializer_class = serializers.AccountTaskSerializer
     search_fields = AccountViewSet.search_fields
     filterset_class = AccountViewSet.filterset_class
@@ -100,10 +91,16 @@ class AccountTaskCreateAPI(CreateAPIView):
         return queryset
 
     def perform_create(self, serializer):
+        data = serializer.validated_data
         accounts = self.get_accounts()
         account_ids = accounts.values_list('id', flat=True)
         asset_ids = [account.asset_id for account in accounts]
-        task = verify_accounts_connectivity.delay(account_ids, asset_ids)
+
+        if data['action'] == 'push':
+            task = push_accounts_to_assets_task.delay(account_ids, asset_ids)
+        else:
+            task = verify_accounts_connectivity_task.delay(account_ids, asset_ids)
+
         data = getattr(serializer, '_data', {})
         data["task"] = task.id
         setattr(serializer, '_data', data)

apps/accounts/api/account/backup.py

@@ -7,7 +7,7 @@ from accounts import serializers
 from accounts.models import (
     AccountBackupAutomation, AccountBackupExecution
 )
-from accounts.tasks import execute_account_backup_plan
+from accounts.tasks import execute_account_backup_task
 from common.const.choices import Trigger
 from orgs.mixins.api import OrgBulkModelViewSet
 
@@ -38,5 +38,5 @@ class AccountBackupPlanExecutionViewSet(viewsets.ModelViewSet):
         serializer = self.get_serializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         pid = serializer.data.get('plan')
-        task = execute_account_backup_plan.delay(pid=str(pid), trigger=Trigger.manual)
+        task = execute_account_backup_task.delay(pid=str(pid), trigger=Trigger.manual)
         return Response({'task': task.id}, status=status.HTTP_201_CREATED)

apps/accounts/api/automations/base.py

@@ -4,7 +4,7 @@ from rest_framework import status, mixins, viewsets
 from rest_framework.response import Response
 
 from accounts.models import AutomationExecution
-from accounts.tasks import execute_automation
+from accounts.tasks import execute_account_automation_task
 from assets import serializers
 from assets.models import BaseAutomation
 from common.const.choices import Trigger
@@ -109,7 +109,7 @@ class AutomationExecutionViewSet(
         serializer = self.get_serializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         automation = serializer.validated_data.get('automation')
-        task = execute_automation.delay(
+        task = execute_account_automation_task.delay(
             pid=str(automation.pk), trigger=Trigger.manual, tp=self.tp
         )
         return Response({'task': task.id}, status=status.HTTP_201_CREATED)

apps/accounts/automations/change_secret/host/posix/main.yml

@@ -3,10 +3,6 @@
   tasks:
     - name: Test privileged account
       ansible.builtin.ping:
-    #
-    #    - name: print variables
-    #      debug:
-    #        msg: "Username: {{ account.username }}, Secret: {{ account.secret }}, Secret type: {{ secret_type }}"
 
     - name: Change password
       ansible.builtin.user:

apps/accounts/models/automations/backup_account.py

@@ -7,11 +7,11 @@ from celery import current_task
 from django.db import models
 from django.db.models import F
 from django.utils.translation import ugettext_lazy as _
-from common.utils import lazyproperty
 
 from common.const.choices import Trigger
 from common.db.encoder import ModelJSONFieldEncoder
 from common.utils import get_logger
+from common.utils import lazyproperty
 from ops.mixin import PeriodTaskModelMixin
 from orgs.mixins.models import OrgModelMixin, JMSOrgBaseModel
 
@@ -36,9 +36,9 @@ class AccountBackupAutomation(PeriodTaskModelMixin, JMSOrgBaseModel):
         verbose_name = _('Account backup plan')
 
     def get_register_task(self):
-        from ...tasks import execute_account_backup_plan
+        from ...tasks import execute_account_backup_task
         name = "account_backup_plan_period_{}".format(str(self.id)[:8])
-        task = execute_account_backup_plan.name
+        task = execute_account_backup_task.name
         args = (str(self.id), Trigger.timing)
         kwargs = {}
         return name, task, args, kwargs

apps/accounts/models/automations/base.py

@@ -1,6 +1,6 @@
 from django.utils.translation import gettext_lazy as _
 
-from accounts.tasks import execute_automation
+from accounts.tasks import execute_account_automation_task
 from assets.models.automations import (
     BaseAutomation as AssetBaseAutomation,
     AutomationExecution as AssetAutomationExecution
@@ -16,7 +16,7 @@ class AccountBaseAutomation(AssetBaseAutomation):
 
     @property
     def execute_task(self):
-        return execute_automation
+        return execute_account_automation_task
 
     @property
     def execution_model(self):

apps/accounts/serializers/account/account.py

@@ -3,7 +3,7 @@ from rest_framework import serializers
 
 from accounts.const import SecretType, Source
 from accounts.models import Account, AccountTemplate
-from accounts.tasks import push_accounts_to_assets
+from accounts.tasks import push_accounts_to_assets_task
 from assets.const import Category, AllTypes
 from assets.models import Asset
 from common.serializers import SecretReadableMixin, BulkModelSerializer
@@ -43,7 +43,7 @@ class AccountSerializerCreateValidateMixin:
     def push_account(instance, push_now):
         if not push_now:
             return
-        push_accounts_to_assets.delay([instance.id], [instance.asset_id])
+        push_accounts_to_assets_task.delay([instance.id], [instance.asset_id])
 
     def create(self, validated_data):
         push_now = validated_data.pop('push_now', None)
@@ -140,6 +140,8 @@ class AccountHistorySerializer(serializers.ModelSerializer):
 class AccountTaskSerializer(serializers.Serializer):
     ACTION_CHOICES = (
         ('test', 'test'),
+        ('verify', 'verify'),
+        ('push', 'push'),
     )
     action = serializers.ChoiceField(choices=ACTION_CHOICES, write_only=True)
     task = serializers.CharField(read_only=True)

apps/accounts/tasks/automation.py

@@ -24,7 +24,7 @@ def task_activity_callback(self, pid, trigger, tp):
     queue='ansible', verbose_name=_('Account execute automation'),
     activity_callback=task_activity_callback
 )
-def execute_automation(pid, trigger, tp):
+def execute_account_automation_task(pid, trigger, tp):
     model = AutomationTypes.get_type_model(tp)
     with tmp_to_root_org():
         instance = get_object_or_none(model, pk=pid)

apps/accounts/tasks/backup_account.py

@@ -10,6 +10,7 @@ logger = get_logger(__file__)
 
 
 def task_activity_callback(self, pid, trigger):
+    from accounts.models import AccountBackupAutomation
     with tmp_to_root_org():
         plan = get_object_or_none(AccountBackupAutomation, pk=pid)
     if not plan:
@@ -22,7 +23,7 @@ def task_activity_callback(self, pid, trigger):
 
 
 @shared_task(verbose_name=_('Execute account backup plan'), activity_callback=task_activity_callback)
-def execute_account_backup_plan(pid, trigger):
+def execute_account_backup_task(pid, trigger):
     from accounts.models import AccountBackupAutomation
     with tmp_to_root_org():
         plan = get_object_or_none(AccountBackupAutomation, pk=pid)

apps/accounts/tasks/gather_accounts.py

@@ -9,7 +9,7 @@ from assets.models import Node
 from common.utils import get_logger
 from orgs.utils import org_aware_func
 
-__all__ = ['gather_asset_accounts']
+__all__ = ['gather_asset_accounts_task']
 logger = get_logger(__name__)
 
 
@@ -29,7 +29,7 @@ def gather_asset_accounts_util(nodes, task_name):
     queue="ansible", verbose_name=_('Gather asset accounts'),
     activity_callback=lambda self, node_ids, task_name=None: (node_ids, None)
 )
-def gather_asset_accounts(node_ids, task_name=None):
+def gather_asset_accounts_task(node_ids, task_name=None):
     if task_name is None:
         task_name = gettext_noop("Gather assets accounts")
 

apps/accounts/tasks/push_account.py

@@ -4,43 +4,38 @@ from django.utils.translation import gettext_noop, ugettext_lazy as _
 from accounts.const import AutomationTypes
 from accounts.tasks.common import automation_execute_start
 from common.utils import get_logger
-from orgs.utils import org_aware_func
 
 logger = get_logger(__file__)
 __all__ = [
-    'push_accounts_to_assets',
+    'push_accounts_to_assets_task',
 ]
 
 
-def push_util(account, assets, task_name):
+def push_util(account, asset_ids, task_name):
     task_snapshot = {
         'secret': account.secret,
         'secret_type': account.secret_type,
         'accounts': [account.username],
-        'assets': [str(asset.id) for asset in assets],
+        'assets': asset_ids,
     }
     tp = AutomationTypes.push_account
     automation_execute_start(task_name, tp, task_snapshot)
 
 
-@org_aware_func("assets")
-def push_accounts_to_assets_util(accounts, assets):
-    from accounts.models import PushAccountAutomation
-
-    task_name = gettext_noop("Push accounts to assets")
-    task_name = PushAccountAutomation.generate_unique_name(task_name)
-    for account in accounts:
-        push_util(account, assets, task_name)
-
-
 @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(account_ids, asset_ids):
+def push_accounts_to_assets_task(account_ids, asset_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)
-    return push_accounts_to_assets_util(accounts, assets)
+
+    task_name = gettext_noop("Push accounts to assets")
+    task_name = PushAccountAutomation.generate_unique_name(task_name)
+
+    for account in accounts:
+        push_util(account, assets, task_name)

apps/accounts/tasks/verify_account.py

@@ -1,6 +1,6 @@
 from celery import shared_task
-from django.utils.translation import gettext_noop
 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
@@ -10,7 +10,7 @@ from orgs.utils import org_aware_func
 
 logger = get_logger(__name__)
 __all__ = [
-    'verify_accounts_connectivity'
+    'verify_accounts_connectivity_task'
 ]
 
 
@@ -42,7 +42,7 @@ def verify_accounts_connectivity_util(accounts, assets, task_name):
     queue="ansible", verbose_name=_('Verify asset account availability'),
     activity_callback=lambda self, account_ids, asset_ids: (account_ids, None)
 )
-def verify_accounts_connectivity(account_ids, asset_ids):
+def verify_accounts_connectivity_task(account_ids, asset_ids):
     from assets.models import Asset
     from accounts.models import Account, VerifyAccountAutomation
     assets = Asset.objects.filter(id__in=asset_ids)

apps/accounts/urls.py

@@ -25,17 +25,22 @@ router.register(r'push-account-executions', api.PushAccountExecutionViewSet, 'pu
 router.register(r'push-account-records', api.PushAccountRecordViewSet, 'push-account-record')
 
 urlpatterns = [
-    path('accounts/tasks/', api.AccountTaskCreateAPI.as_view(), name='account-task-create'),
-    path('account-secrets/<uuid:pk>/histories/', api.AccountHistoriesSecretAPI.as_view(), name='account-secret-history'),
+    path('accounts/tasks/', api.AccountsTaskCreateAPI.as_view(), name='account-task-create'),
+    path('account-secrets/<uuid:pk>/histories/', api.AccountHistoriesSecretAPI.as_view(),
+         name='account-secret-history'),
 
-    path('change-secret/<uuid:pk>/asset/remove/', api.ChangSecretRemoveAssetApi.as_view(), name='change-secret-remove-asset'),
+    path('change-secret/<uuid:pk>/asset/remove/', api.ChangSecretRemoveAssetApi.as_view(),
+         name='change-secret-remove-asset'),
     path('change-secret/<uuid:pk>/asset/add/', api.ChangSecretAddAssetApi.as_view(), name='change-secret-add-asset'),
-    path('change-secret/<uuid:pk>/nodes/', api.ChangSecretNodeAddRemoveApi.as_view(), name='change-secret-add-or-remove-node'),
+    path('change-secret/<uuid:pk>/nodes/', api.ChangSecretNodeAddRemoveApi.as_view(),
+         name='change-secret-add-or-remove-node'),
     path('change-secret/<uuid:pk>/assets/', api.ChangSecretAssetsListApi.as_view(), name='change-secret-assets'),
 
-    path('push-account/<uuid:pk>/asset/remove/', api.PushAccountRemoveAssetApi.as_view(), name='push-account-remove-asset'),
+    path('push-account/<uuid:pk>/asset/remove/', api.PushAccountRemoveAssetApi.as_view(),
+         name='push-account-remove-asset'),
     path('push-accountt/<uuid:pk>/asset/add/', api.PushAccountAddAssetApi.as_view(), name='push-account-add-asset'),
-    path('push-account/<uuid:pk>/nodes/', api.PushAccountNodeAddRemoveApi.as_view(), name='push-account-add-or-remove-node'),
+    path('push-account/<uuid:pk>/nodes/', api.PushAccountNodeAddRemoveApi.as_view(),
+         name='push-account-add-or-remove-node'),
     path('push-account/<uuid:pk>/assets/', api.PushAccountAssetsListApi.as_view(), name='push-account-assets'),
 ]
 

apps/assets/api/asset/asset.py

@@ -6,7 +6,7 @@ from django.utils.translation import gettext as _
 from rest_framework.decorators import action
 from rest_framework.response import Response
 
-from accounts.tasks import push_accounts_to_assets, verify_accounts_connectivity
+from accounts.tasks import push_accounts_to_assets_task, verify_accounts_connectivity_task
 from assets import serializers
 from assets.filters import IpInFilterBackend, LabelFilterBackend, NodeFilterBackend
 from assets.models import Asset, Gateway
@@ -205,9 +205,9 @@ class AssetTaskCreateApi(AssetsTaskMixin, generics.CreateAPIView):
         asset_ids = [asset.id]
         account_ids = accounts.values_list("id", flat=True)
         if action == "push_account":
-            task = push_accounts_to_assets.delay(account_ids, asset_ids)
+            task = push_accounts_to_assets_task.delay(account_ids, asset_ids)
         elif action == "test_account":
-            task = verify_accounts_connectivity.delay(account_ids, asset_ids)
+            task = verify_accounts_connectivity_task.delay(account_ids, asset_ids)
         else:
             task = None
         return task

apps/assets/models/automations/base.py

@@ -6,7 +6,7 @@ from django.utils.translation import ugettext_lazy as _
 
 from assets.models.asset import Asset
 from assets.models.node import Node
-from assets.tasks import execute_automation
+from assets.tasks import execute_asset_automation_task
 from common.const.choices import Trigger
 from common.db.fields import EncryptJsonDictTextField
 from ops.mixin import PeriodTaskModelMixin
@@ -49,7 +49,7 @@ class BaseAutomation(PeriodTaskModelMixin, JMSOrgBaseModel):
 
     @property
     def execute_task(self):
-        return execute_automation
+        return execute_asset_automation_task
 
     def get_register_task(self):
         name = f"automation_{self.type}_strategy_period_{str(self.id)[:8]}"

apps/assets/tasks/automation.py

@@ -1,9 +1,9 @@
 from celery import shared_task
 from django.utils.translation import gettext_lazy as _
 
-from orgs.utils import tmp_to_root_org, tmp_to_org
-from common.utils import get_logger, get_object_or_none
 from assets.const import AutomationTypes
+from common.utils import get_logger, get_object_or_none
+from orgs.utils import tmp_to_root_org, tmp_to_org
 
 logger = get_logger(__file__)
 
@@ -24,7 +24,7 @@ def task_activity_callback(self, pid, trigger, tp):
     queue='ansible', verbose_name=_('Asset execute automation'),
     activity_callback=task_activity_callback
 )
-def execute_automation(pid, trigger, tp):
+def execute_asset_automation_task(pid, trigger, tp):
     model = AutomationTypes.get_type_model(tp)
     with tmp_to_root_org():
         instance = get_object_or_none(model, pk=pid)