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