From 424ef4d9a542919655f2a02d79a18b2294538106 Mon Sep 17 00:00:00 2001 From: feng <1304903146@qq.com> Date: Mon, 17 Mar 2025 17:26:09 +0800 Subject: [PATCH] perf: Automation execution --- .../automations/change_secret_dashboard.py | 45 ++++++++++++------- .../automations/backup_account/handlers.py | 18 +++++--- .../models/automations/backup_account.py | 3 +- apps/accounts/serializers/account/account.py | 1 + apps/accounts/tasks/common.py | 2 +- apps/assets/models/automations/base.py | 1 + apps/assets/tasks/common.py | 2 +- apps/i18n/lina/en.json | 4 +- apps/i18n/lina/zh.json | 4 +- 9 files changed, 53 insertions(+), 27 deletions(-) diff --git a/apps/accounts/api/automations/change_secret_dashboard.py b/apps/accounts/api/automations/change_secret_dashboard.py index 4dc73cfd6..a395d4002 100644 --- a/apps/accounts/api/automations/change_secret_dashboard.py +++ b/apps/accounts/api/automations/change_secret_dashboard.py @@ -2,6 +2,7 @@ # from collections import defaultdict +from django.core.cache import cache from django.http.response import JsonResponse from django.utils import timezone from rest_framework.views import APIView @@ -25,6 +26,7 @@ class ChangeSecretDashboardApi(APIView): tp = AutomationTypes.change_secret task_name = 'accounts.tasks.automation.execute_account_automation_task' + ongoing_change_secret_cache_key = "ongoing_change_secret_cache_key" @lazyproperty def days(self): @@ -146,24 +148,33 @@ class ChangeSecretDashboardApi(APIView): }) if _all or query_params.get('total_count_ongoing_change_secret'): - execution_ids = [] - inspect = app.control.inspect() - active_tasks = inspect.active() - if active_tasks: - for tasks in active_tasks.values(): - for task in tasks: - _id = task.get('id') - name = task.get('name') - tp = task.get('kwargs', {}).get('tp') - if name == self.task_name and tp == self.tp: - execution_ids.append(_id) + ongoing_counts = cache.get(self.ongoing_change_secret_cache_key) + if ongoing_counts is None: + execution_ids = [] + inspect = app.control.inspect() + active_tasks = inspect.active() + if active_tasks: + for tasks in active_tasks.values(): + for task in tasks: + _id = task.get('id') + name = task.get('name') + tp = task.get('kwargs', {}).get('tp') + if name == self.task_name and tp == self.tp: + execution_ids.append(_id) - snapshots = AutomationExecution.objects.filter(id__in=execution_ids).values_list('snapshot', flat=True) + snapshots = AutomationExecution.objects.filter(id__in=execution_ids).values_list('snapshot', flat=True) - asset_ids = {asset for i in snapshots for asset in i.get('assets', [])} - account_ids = {account for i in snapshots for account in i.get('accounts', [])} - data['total_count_ongoing_change_secret'] = len(execution_ids) - data['total_count_ongoing_change_secret_assets'] = len(asset_ids) - data['total_count_ongoing_change_secret_accounts'] = len(account_ids) + asset_ids = {asset for i in snapshots for asset in i.get('assets', [])} + account_ids = {account for i in snapshots for account in i.get('accounts', [])} + + ongoing_counts = (len(execution_ids), len(asset_ids), len(account_ids)) + data['total_count_ongoing_change_secret'] = ongoing_counts[0] + data['total_count_ongoing_change_secret_assets'] = ongoing_counts[1] + data['total_count_ongoing_change_secret_accounts'] = ongoing_counts[2] + cache.set(self.ongoing_change_secret_cache_key, ongoing_counts, 60) + else: + data['total_count_ongoing_change_secret'] = ongoing_counts[0] + data['total_count_ongoing_change_secret_assets'] = ongoing_counts[1] + data['total_count_ongoing_change_secret_accounts'] = ongoing_counts[2] return JsonResponse(data, status=200) diff --git a/apps/accounts/automations/backup_account/handlers.py b/apps/accounts/automations/backup_account/handlers.py index 1f1f01cc4..737c427a2 100644 --- a/apps/accounts/automations/backup_account/handlers.py +++ b/apps/accounts/automations/backup_account/handlers.py @@ -36,17 +36,24 @@ class BaseAccountHandler: if isinstance(v, OrderedDict): cls.unpack_data(v, data) else: + if isinstance(v, dict): + v = v.get('label') data[k] = v return data @classmethod def get_header_fields(cls, serializer: serializers.Serializer): try: - backup_fields = getattr(serializer, 'Meta').fields_backup + exclude_backup_fields = getattr(serializer, 'Meta').exclude_backup_fields except AttributeError: - backup_fields = serializer.fields.keys() + exclude_backup_fields = [] + backup_fields = serializer.fields.keys() + header_fields = {} for field in backup_fields: + if field in exclude_backup_fields: + continue + v = serializer.fields[field] if isinstance(v, serializers.Serializer): _fields = cls.get_header_fields(v) @@ -189,8 +196,8 @@ class AccountBackupHandler: attachment_list = [attachment] AccountBackupExecutionTaskMsg(name, user).publish(attachment_list) - for file in files: - os.remove(file) + # for file in files: + # os.remove(file) def send_backup_obj_storage(self, files, recipients, password): if not files: @@ -275,7 +282,8 @@ class AccountBackupHandler: else: recipients = recipients_part_one or recipients_part_two files = self.create_excel() - self.send_backup_mail(files, recipients) + print(files) + # self.send_backup_mail(files, recipients) def run(self): print('{}: {}'.format(_('Plan start'), local_now_display())) diff --git a/apps/accounts/models/automations/backup_account.py b/apps/accounts/models/automations/backup_account.py index 7f44cd226..e617d1b1e 100644 --- a/apps/accounts/models/automations/backup_account.py +++ b/apps/accounts/models/automations/backup_account.py @@ -7,8 +7,8 @@ from django.utils.translation import gettext_lazy as _ from accounts.const import AccountBackupType, AutomationTypes from common.db import fields -from orgs.mixins.models import OrgManager from common.utils import get_logger +from orgs.mixins.models import OrgManager from .base import AccountBaseAutomation __all__ = ['BackupAccountAutomation'] @@ -59,6 +59,7 @@ class BackupAccountAutomation(AccountBaseAutomation): def to_attr_json(self): attr_json = super().to_attr_json() attr_json.update({ + 'id': str(self.id), 'types': self.types, 'backup_type': self.backup_type, 'is_password_divided_by_email': self.is_password_divided_by_email, diff --git a/apps/accounts/serializers/account/account.py b/apps/accounts/serializers/account/account.py index 844abce9c..438fb0244 100644 --- a/apps/accounts/serializers/account/account.py +++ b/apps/accounts/serializers/account/account.py @@ -460,6 +460,7 @@ class AccountSecretSerializer(SecretReadableMixin, AccountSerializer): 'secret': {'write_only': False}, 'spec_info': {'label': _('Spec info')}, } + exclude_backup_fields = ['passphrase', 'push_now', 'params'] class AccountHistorySerializer(serializers.ModelSerializer): diff --git a/apps/accounts/tasks/common.py b/apps/accounts/tasks/common.py index 4582322c3..69cd8f404 100644 --- a/apps/accounts/tasks/common.py +++ b/apps/accounts/tasks/common.py @@ -13,6 +13,6 @@ def quickstart_automation_by_snapshot(task_name, tp, task_snapshot=None): data['id'] = str(uuid.uuid4()) execution = AutomationExecution.objects.create( - trigger=Trigger.manual, **data + type=tp, trigger=Trigger.manual, **data ) execution.start() diff --git a/apps/assets/models/automations/base.py b/apps/assets/models/automations/base.py index bf0caf50b..cbeadea57 100644 --- a/apps/assets/models/automations/base.py +++ b/apps/assets/models/automations/base.py @@ -109,6 +109,7 @@ class BaseAutomation(PeriodTaskModelMixin, JMSOrgBaseModel): execution = self.execution_model.objects.create( id=eid, + type=self.type, trigger=trigger, automation=self, snapshot=self.to_attr_json(), diff --git a/apps/assets/tasks/common.py b/apps/assets/tasks/common.py index 977347000..08c57c972 100644 --- a/apps/assets/tasks/common.py +++ b/apps/assets/tasks/common.py @@ -42,7 +42,7 @@ def quickstart_automation(task_name, tp, task_snapshot=None): break execution = AutomationExecution.objects.create( - trigger=Trigger.manual, **data + type=tp, trigger=Trigger.manual, **data ) execution.start() return execution diff --git a/apps/i18n/lina/en.json b/apps/i18n/lina/en.json index a003cbede..59f558d5c 100644 --- a/apps/i18n/lina/en.json +++ b/apps/i18n/lina/en.json @@ -1520,5 +1520,7 @@ "Organization": "Organization", "IpGroupHelpText": "* indicates match all. for example: 192.168.10.1, 192.168.1.0/24, 10.1.1.1-10.1.1.20, 2001:db8:2de::e13, 2001:db8:1a:1110::/64", "SelectRisk": "Select risk", - "IpGroup": "IP group" + "IpGroup": "IP group", + "PublicIP": "Public IP", + "PrivateIP": "Private IP" } \ No newline at end of file diff --git a/apps/i18n/lina/zh.json b/apps/i18n/lina/zh.json index 31a7327fb..249fa996f 100644 --- a/apps/i18n/lina/zh.json +++ b/apps/i18n/lina/zh.json @@ -1522,5 +1522,7 @@ "Reopen": "重新打开", "Review": "审查", "SelectRisk": "选择风险", - "IpGroup": "IP 组" + "IpGroup": "IP 组", + "PublicIP": "公有 IP", + "PrivateIP": "私有 IP" } \ No newline at end of file