From 757e688ab36e545e451b8f4047929f8865936831 Mon Sep 17 00:00:00 2001
From: feng626 <1304903146@qq.com>
Date: Sun, 4 Sep 2022 16:33:36 +0800
Subject: [PATCH] =?UTF-8?q?=E8=B4=A6=E5=8F=B7=E5=A4=87=E4=BB=BD=E4=BC=98?=
 =?UTF-8?q?=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/task_handlers/backup/handlers.py | 44 ++++++++++++--------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/apps/assets/task_handlers/backup/handlers.py b/apps/assets/task_handlers/backup/handlers.py
index 42d21484d..3b8eac87d 100644
--- a/apps/assets/task_handlers/backup/handlers.py
+++ b/apps/assets/task_handlers/backup/handlers.py
@@ -49,16 +49,23 @@ class BaseAccountHandler:
         return header_fields
 
     @classmethod
-    def create_row(cls, account, serializer_cls, header_fields=None):
-        serializer = serializer_cls(account)
-        if not header_fields:
-            header_fields = cls.get_header_fields(serializer)
-        data = cls.unpack_data(serializer.data)
+    def create_row(cls, data, header_fields):
+        data = cls.unpack_data(data)
         row_dict = {}
         for field, header_name in header_fields.items():
-            row_dict[header_name] = str(data[field])
+            row_dict[header_name] = str(data.get(field, field))
         return row_dict
 
+    @classmethod
+    def add_rows(cls, data, header_fields, sheet):
+        data_map = defaultdict(list)
+        for i in data:
+            row = cls.create_row(i, header_fields)
+            if sheet not in data_map:
+                data_map[sheet].append(list(row.keys()))
+            data_map[sheet].append(list(row.values()))
+        return data_map
+
 
 class AssetAccountHandler(BaseAccountHandler):
     @staticmethod
@@ -73,22 +80,25 @@ class AssetAccountHandler(BaseAccountHandler):
         data_map = defaultdict(list)
 
         # TODO 可以优化一下查询 在账号上做type的缓存 避免数据量大时连表操作
-        accounts = Account.objects.filter(
+        qs = Account.objects.filter(
             asset__platform__type__in=types
         ).annotate(type=F('asset__platform__type'))
-        if not accounts.first():
+        if not qs.exists():
             return data_map
 
         type_dict = dict(Type.CHOICES)
-        header_fields = cls.get_header_fields(AccountSecretSerializer(accounts.first()))
-        for account in accounts:
-            sheet_name = type_dict[account.type]
-            row = cls.create_row(account, AccountSecretSerializer, header_fields)
-            if sheet_name not in data_map:
-                data_map[sheet_name].append(list(row.keys()))
-            data_map[sheet_name].append(list(row.values()))
-
-        logger.info('\n\033[33m- 共收集 {} 条账号\033[0m'.format(accounts.count()))
+        header_fields = cls.get_header_fields(AccountSecretSerializer(qs.first()))
+        account_type_map = defaultdict(list)
+        for account in qs:
+            account_type_map[account.type].append(account)
+
+        data_map = {}
+        for tp, accounts in account_type_map.items():
+            sheet_name = type_dict[tp]
+            data = AccountSecretSerializer(accounts, many=True).data
+            data_map.update(cls.add_rows(data, header_fields, sheet_name))
+
+        logger.info('\n\033[33m- 共收集 {} 条账号\033[0m'.format(qs.count()))
         return data_map