From 1cc983b2eb3e8b9a19029e5ccc9c06f20f21bda4 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Mon, 7 Nov 2022 16:10:26 +0800
Subject: [PATCH] perf: automation button (#9023)

Co-authored-by: feng <1304903146@qq.com>
---
 apps/assets/api/asset/asset.py          |  5 ++--
 apps/assets/automations/base/manager.py |  3 +++
 apps/assets/automations/endpoint.py     | 10 +++++---
 apps/assets/serializers/asset/common.py |  2 +-
 apps/assets/tasks/__init__.py           |  1 +
 apps/assets/tasks/gather_accounts.py    | 34 +++++++++++++++++++++++++
 6 files changed, 48 insertions(+), 7 deletions(-)
 create mode 100644 apps/assets/tasks/gather_accounts.py

diff --git a/apps/assets/api/asset/asset.py b/apps/assets/api/asset/asset.py
index 43461730f..4e1176d17 100644
--- a/apps/assets/api/asset/asset.py
+++ b/apps/assets/api/asset/asset.py
@@ -82,10 +82,11 @@ class AssetsTaskMixin:
     def perform_assets_task(self, serializer):
         data = serializer.validated_data
         assets = data.get('assets', [])
+        asset_ids = [asset.id for asset in assets]
         if data['action'] == "refresh":
-            task = update_assets_hardware_info_manual.delay(assets)
+            task = update_assets_hardware_info_manual.delay(asset_ids)
         else:
-            task = test_assets_connectivity_manual.delay(assets)
+            task = test_assets_connectivity_manual.delay(asset_ids)
         return task
 
     def perform_create(self, serializer):
diff --git a/apps/assets/automations/base/manager.py b/apps/assets/automations/base/manager.py
index c5f334167..ab8ea01bf 100644
--- a/apps/assets/automations/base/manager.py
+++ b/apps/assets/automations/base/manager.py
@@ -221,6 +221,7 @@ class BasePlaybookManager:
         else:
             print(">>> 开始执行任务\n")
 
+        self.execution.date_start = timezone.now()
         for i, runner in enumerate(runners, start=1):
             if len(runners) > 1:
                 print(">>> 开始执行第 {} 批任务".format(i))
@@ -231,3 +232,5 @@ class BasePlaybookManager:
             except Exception as e:
                 self.on_runner_failed(runner, e)
             print('\n')
+        self.execution.date_finished = timezone.now()
+        self.execution.save()
diff --git a/apps/assets/automations/endpoint.py b/apps/assets/automations/endpoint.py
index eb7b2f4ca..2548e9184 100644
--- a/apps/assets/automations/endpoint.py
+++ b/apps/assets/automations/endpoint.py
@@ -4,16 +4,18 @@ from .gather_accounts.manager import GatherAccountsManager
 from .verify_account.manager import VerifyAccountManager
 from .push_account.manager import PushAccountManager
 from .backup_account.manager import AccountBackupManager
+from .ping.manager import PingManager
 from ..const import AutomationTypes
 
 
 class ExecutionManager:
     manager_type_mapper = {
-        AutomationTypes.change_secret: ChangeSecretManager,
-        AutomationTypes.gather_facts: GatherFactsManager,
-        AutomationTypes.gather_accounts: GatherAccountsManager,
-        AutomationTypes.verify_account: VerifyAccountManager,
+        AutomationTypes.ping: PingManager,
         AutomationTypes.push_account: PushAccountManager,
+        AutomationTypes.gather_facts: GatherFactsManager,
+        AutomationTypes.change_secret: ChangeSecretManager,
+        AutomationTypes.verify_account: VerifyAccountManager,
+        AutomationTypes.gather_accounts: GatherAccountsManager,
         # TODO 后期迁移到自动化策略中
         'backup_account': AccountBackupManager,
     }
diff --git a/apps/assets/serializers/asset/common.py b/apps/assets/serializers/asset/common.py
index c44be52ce..a6c4b4fe2 100644
--- a/apps/assets/serializers/asset/common.py
+++ b/apps/assets/serializers/asset/common.py
@@ -77,7 +77,7 @@ class AssetSerializer(OrgResourceSerializerMixin, WritableNestedModelSerializer)
             'nodes', 'labels', 'protocols', 'accounts', 'nodes_display',
         ]
         read_only_fields = [
-            'category', 'type', 'specific',
+            'category', 'type', 'specific', 'info',
             'connectivity', 'date_verified',
             'created_by', 'date_created',
         ]
diff --git a/apps/assets/tasks/__init__.py b/apps/assets/tasks/__init__.py
index c4d56528c..060f4d2d9 100644
--- a/apps/assets/tasks/__init__.py
+++ b/apps/assets/tasks/__init__.py
@@ -9,3 +9,4 @@ from .gather_facts import *
 from .nodes_amount import *
 from .push_account import *
 from .verify_account import *
+from .gather_accounts import *
diff --git a/apps/assets/tasks/gather_accounts.py b/apps/assets/tasks/gather_accounts.py
new file mode 100644
index 000000000..4e372aca7
--- /dev/null
+++ b/apps/assets/tasks/gather_accounts.py
@@ -0,0 +1,34 @@
+# ~*~ coding: utf-8 ~*~
+from celery import shared_task
+from django.utils.translation import gettext_noop
+
+from orgs.utils import tmp_to_root_org, org_aware_func
+from common.utils import get_logger
+from assets.models import Node
+
+__all__ = ['gather_asset_accounts']
+logger = get_logger(__name__)
+
+
+@org_aware_func("nodes")
+def gather_asset_accounts_util(nodes, task_name):
+    from assets.models import GatherAccountsAutomation
+    task_name = GatherAccountsAutomation.generate_unique_name(task_name)
+
+    data = {
+        'name': task_name,
+        'comment': ', '.join([str(i) for i in nodes])
+    }
+    instance = GatherAccountsAutomation.objects.create(**data)
+    instance.nodes.add(*nodes)
+    instance.execute()
+
+
+@shared_task(queue="ansible")
+def gather_asset_accounts(node_ids, task_name=None):
+    if task_name is None:
+        task_name = gettext_noop("Gather assets accounts")
+
+    with tmp_to_root_org():
+        nodes = Node.objects.filter(id__in=node_ids)
+    gather_asset_accounts_util(nodes=nodes, task_name=task_name)