perf: account backup (#9013)

Co-authored-by: feng <1304903146@qq.com>

apps/assets/api/account/backup.py

@@ -4,6 +4,7 @@ from rest_framework import status, viewsets
 from rest_framework.response import Response
 
 from orgs.mixins.api import OrgBulkModelViewSet
+from common.const.choices import Trigger
 from assets import serializers
 from assets.tasks import execute_account_backup_plan
 from assets.models import (
@@ -38,9 +39,7 @@ 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(
+        task = execute_account_backup_plan.delay(pid=pid, trigger=Trigger.manual)
-            pid=pid, trigger=AccountBackupPlanExecution.Trigger.manual
-        )
         return Response({'task': task.id}, status=status.HTTP_201_CREATED)
 
     def filter_queryset(self, queryset):

apps/assets/automations/backup_account/handlers.py

@@ -82,15 +82,16 @@ class AssetAccountHandler(BaseAccountHandler):
 
         # TODO 可以优化一下查询 在账号上做 category 的缓存 避免数据量大时连表操作
         qs = Account.objects.filter(
-            asset__platform__category__in=categories
+            asset__platform__type__in=categories
-        ).annotate(category=F('asset__platform__category'))
+        ).annotate(category=F('asset__platform__type'))
+        print(qs, categories)
         if not qs.exists():
             return data_map
 
         category_dict = {}
         for i in AllTypes.grouped_choices_to_objs():
             for j in i['children']:
-                category_dict[j['value']] = j['label']
+                category_dict[j['value']] = j['display_name']
 
         header_fields = cls.get_header_fields(AccountSecretSerializer(qs.first()))
         account_category_map = defaultdict(list)

apps/assets/automations/backup_account/manager.py

@@ -12,7 +12,7 @@ from .handlers import AccountBackupHandler
 logger = get_logger(__name__)
 
 
-class AccountBackupExecutionManager:
+class AccountBackupManager:
     def __init__(self, execution):
         self.execution = execution
         self.date_start = timezone.now()

apps/assets/automations/endpoint.py

@@ -3,6 +3,7 @@ from .gather_facts.manager import GatherFactsManager
 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 ..const import AutomationTypes
 
 
@@ -13,6 +14,8 @@ class ExecutionManager:
         AutomationTypes.gather_accounts: GatherAccountsManager,
         AutomationTypes.verify_account: VerifyAccountManager,
         AutomationTypes.push_account: PushAccountManager,
+        # TODO 后期迁移到自动化策略中
+        'backup_account': AccountBackupManager,
     }
 
     def __init__(self, execution):
@@ -21,4 +24,3 @@ class ExecutionManager:
 
     def run(self, *args, **kwargs):
         return self._runner.run(*args, **kwargs)
-

apps/assets/migrations/0109_rename_categories_to_types.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.14 on 2022-11-03 08:44
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0108_auto_20221027_1053'),
+    ]
+
+    operations = [
+        migrations.RenameField(
+            model_name='accountbackupplan',
+            old_name='categories',
+            new_name='types',
+        ),
+    ]

apps/assets/models/backup.py

@@ -2,7 +2,6 @@
 # -*- coding: utf-8 -*-
 #
 import uuid
-from functools import reduce
 
 from celery import current_task
 from django.db import models
@@ -11,9 +10,9 @@ from django.utils.translation import ugettext_lazy as _
 from orgs.mixins.models import OrgModelMixin
 from ops.mixin import PeriodTaskModelMixin
 from common.utils import get_logger
+from common.const.choices import Trigger
 from common.db.encoder import ModelJSONFieldEncoder
 from common.mixins.models import CommonModelMixin
-from common.const.choices import Trigger
 
 __all__ = ['AccountBackupPlan', 'AccountBackupPlanExecution']
 
@@ -22,7 +21,7 @@ logger = get_logger(__file__)
 
 class AccountBackupPlan(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
-    categories = models.JSONField(default=list)
+    types = models.JSONField(default=list)
     recipients = models.ManyToManyField(
         'users.User', related_name='recipient_escape_route_plans', blank=True,
         verbose_name=_("Recipient")
@@ -53,7 +52,7 @@ class AccountBackupPlan(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
             'crontab': self.crontab,
             'org_id': self.org_id,
             'created_by': self.created_by,
-            'categories': self.categories,
+            'types': self.types,
             'recipients': {
                 str(recipient.id): (str(recipient), bool(recipient.secret_key))
                 for recipient in self.recipients.all()
@@ -100,9 +99,9 @@ class AccountBackupPlanExecution(OrgModelMixin):
         verbose_name = _('Account backup execution')
 
     @property
-    def categories(self):
+    def types(self):
-        categories = self.plan_snapshot.get('categories')
+        types = self.plan_snapshot.get('types')
-        return categories
+        return types
 
     @property
     def recipients(self):
@@ -111,7 +110,11 @@ class AccountBackupPlanExecution(OrgModelMixin):
             return []
         return recipients.values()
 
+    @property
+    def manager_type(self):
+        return 'backup_account'
+
     def start(self):
-        from ..task_handlers import ExecutionManager
+        from assets.automations.endpoint import ExecutionManager
         manager = ExecutionManager(execution=self)
         return manager.run()