From d6aea54722f43f4976429e1a9e39cee254d53e1e Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Wed, 15 Nov 2023 16:44:35 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E8=B4=A6=E5=8F=B7=E6=94=B6=E9=9B=86?=
 =?UTF-8?q?=E6=9C=AA=E5=90=8C=E6=AD=A5=E8=B5=84=E4=BA=A7=E6=97=B6=20?=
 =?UTF-8?q?=E5=8F=98=E6=9B=B4=E6=95=B0=E6=8D=AE=E9=94=99=E8=AF=AF=20(#1212?=
 =?UTF-8?q?3)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 .../automations/gather_accounts/manager.py    | 58 ++++++++++++-------
 1 file changed, 36 insertions(+), 22 deletions(-)

diff --git a/apps/accounts/automations/gather_accounts/manager.py b/apps/accounts/automations/gather_accounts/manager.py
index c3af9ffe8..0941a78fa 100644
--- a/apps/accounts/automations/gather_accounts/manager.py
+++ b/apps/accounts/automations/gather_accounts/manager.py
@@ -17,7 +17,8 @@ class GatherAccountsManager(AccountBasePlaybookManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.host_asset_mapper = {}
-        self.gathered_accounts = []
+        self.asset_account_info = {}
+
         self.asset_username_mapper = defaultdict(set)
         self.is_sync_account = self.execution.snapshot.get('is_sync_account')
 
@@ -46,49 +47,58 @@ class GatherAccountsManager(AccountBasePlaybookManager):
             data.append(d)
         return data
 
-    def update_or_create_accounts(self, asset, result):
+    def collect_asset_account_info(self, asset, result):
         data = self.generate_data(asset, result)
-        with tmp_to_org(asset.org_id):
-            GatheredAccount.objects.filter(asset=asset, present=True).update(present=False)
-            for d in data:
-                username = d['username']
-                gathered_account, __ = GatheredAccount.objects.update_or_create(
-                    defaults=d, asset=asset, username=username,
-                )
-                self.gathered_accounts.append(gathered_account)
+        self.asset_account_info[asset] = data
 
     def on_host_success(self, host, result):
         info = result.get('debug', {}).get('res', {}).get('info', {})
         asset = self.host_asset_mapper.get(host)
         if asset and info:
             result = self.filter_success_result(asset.type, info)
-            self.update_or_create_accounts(asset, result)
+            self.collect_asset_account_info(asset, result)
         else:
             logger.error(f'Not found {host} info')
 
+    def update_or_create_accounts(self):
+        for asset, data in self.asset_account_info.items():
+            with tmp_to_org(asset.org_id):
+                gathered_accounts = []
+                GatheredAccount.objects.filter(asset=asset, present=True).update(present=False)
+                for d in data:
+                    username = d['username']
+                    gathered_account, __ = GatheredAccount.objects.update_or_create(
+                        defaults=d, asset=asset, username=username,
+                    )
+                    gathered_accounts.append(gathered_account)
+                if not self.is_sync_account:
+                    return
+                GatheredAccount.sync_accounts(gathered_accounts)
+
     def run(self, *args, **kwargs):
         super().run(*args, **kwargs)
-        self.send_email_if_need()
-        if not self.is_sync_account:
-            return
-        GatheredAccount.sync_accounts(self.gathered_accounts)
+        users, change_info = self.generate_send_users_and_change_info()
+        self.update_or_create_accounts()
+        self.send_email_if_need(users, change_info)
 
-    def send_email_if_need(self):
+    def generate_send_users_and_change_info(self):
         recipients = self.execution.recipients
         if not self.asset_username_mapper or not recipients:
-            return
+            return None, None
 
         users = User.objects.filter(id__in=recipients)
         if not users:
-            return
+            return users, None
 
         asset_ids = self.asset_username_mapper.keys()
         assets = Asset.objects.filter(id__in=asset_ids)
+        gather_accounts = GatheredAccount.objects.filter(asset_id__in=asset_ids, present=True)
         asset_id_map = {str(asset.id): asset for asset in assets}
-        asset_qs = assets.values_list('id', 'accounts__username')
-        system_asset_username_mapper = defaultdict(set)
+        asset_id_username = list(assets.values_list('id', 'accounts__username'))
+        asset_id_username.extend(list(gather_accounts.values_list('asset_id', 'username')))
 
-        for asset_id, username in asset_qs:
+        system_asset_username_mapper = defaultdict(set)
+        for asset_id, username in asset_id_username:
             system_asset_username_mapper[str(asset_id)].add(username)
 
         change_info = {}
@@ -110,7 +120,11 @@ class GatherAccountsManager(AccountBasePlaybookManager):
                 'remove_usernames': ', '.join(remove_usernames),
             }
 
-        if not change_info:
+        return users, change_info
+
+    @staticmethod
+    def send_email_if_need(users, change_info):
+        if not users or not change_info:
             return
 
         for user in users: