perf: Push account recorder

apps/accounts/api/account/account.py

@@ -187,7 +187,7 @@ class AccountHistoriesSecretAPI(ExtraFilterFieldsMixin, AccountRecordViewLogMixi
 
     @property
     def latest_change_secret_record(self) -> ChangeSecretRecord:
-        return self.account.change_secret_records.filter(
+        return self.account.changesecretrecords.filter(
             status=ChangeSecretRecordStatusChoice.pending
         ).order_by('-date_created').first()
 

apps/accounts/automations/base/manager.py

@@ -1,13 +1,15 @@
 from copy import deepcopy
 
 from django.conf import settings
+from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 
 from accounts.automations.methods import platform_automation_methods
-from accounts.const import SSHKeyStrategy, SecretStrategy, SecretType
+from accounts.const import SSHKeyStrategy, SecretStrategy, SecretType, ChangeSecretRecordStatusChoice
 from accounts.models import BaseAccountQuerySet
 from assets.automations.base.manager import BasePlaybookManager
 from assets.const import HostTypes
+from common.db.utils import safe_db_connection
 from common.utils import get_logger
 
 logger = get_logger(__name__)
@@ -32,6 +34,8 @@ class BaseChangeSecretPushManager(AccountBasePlaybookManager):
             'ssh_key_change_strategy', SSHKeyStrategy.set_jms
         )
         self.account_ids = self.execution.snapshot['accounts']
+        self.record_map = self.execution.snapshot.get('record_map', {})  # 这个是某个失败的记录重试
+        self.name_recorder_mapper = {}  # 做个映射，方便后面处理
 
     def gen_account_inventory(self, account, asset, h, path_dir):
         raise NotImplementedError
@@ -119,3 +123,51 @@ class BaseChangeSecretPushManager(AccountBasePlaybookManager):
             inventory_hosts.append(h)
 
         return inventory_hosts
+
+    def on_host_success(self, host, result):
+        recorder = self.name_recorder_mapper.get(host)
+        if not recorder:
+            return
+        recorder.status = ChangeSecretRecordStatusChoice.success.value
+        recorder.date_finished = timezone.now()
+
+        account = recorder.account
+        if not account:
+            print("Account not found, deleted ?")
+            return
+
+        account.secret = getattr(recorder, 'new_secret', account.secret)
+        account.date_updated = timezone.now()
+
+        with safe_db_connection():
+            recorder.save(update_fields=['status', 'date_finished'])
+            account.save(update_fields=['secret', 'date_updated'])
+
+        self.summary['ok_accounts'] += 1
+        self.result['ok_accounts'].append(
+            {
+                "asset": str(account.asset),
+                "username": account.username,
+            }
+        )
+        super().on_host_success(host, result)
+
+    def on_host_error(self, host, error, result):
+        recorder = self.name_recorder_mapper.get(host)
+        if not recorder:
+            return
+        recorder.status = ChangeSecretRecordStatusChoice.failed.value
+        recorder.date_finished = timezone.now()
+        recorder.error = error
+        try:
+            recorder.save()
+        except Exception as e:
+            print(f"\033[31m Save {host} recorder error: {e} \033[0m\n")
+        self.summary['fail_accounts'] += 1
+        self.result['fail_accounts'].append(
+            {
+                "asset": str(recorder.asset),
+                "username": recorder.account.username,
+            }
+        )
+        super().on_host_error(host, error, result)

apps/accounts/automations/change_secret/manager.py

@@ -2,7 +2,6 @@ import os
 import time
 
 from django.conf import settings
-from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from xlsxwriter import Workbook
 
@@ -12,7 +11,6 @@ from accounts.const import (
 from accounts.models import ChangeSecretRecord
 from accounts.notifications import ChangeSecretExecutionTaskMsg, ChangeSecretReportMsg
 from accounts.serializers import ChangeSecretRecordBackUpSerializer
-from common.db.utils import safe_db_connection
 from common.decorators import bulk_create_decorator
 from common.utils import get_logger
 from common.utils.file import encrypt_and_compress_zip_file
@@ -26,11 +24,6 @@ logger = get_logger(__name__)
 class ChangeSecretManager(BaseChangeSecretPushManager):
     ansible_account_prefer = ''
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.record_map = self.execution.snapshot.get('record_map', {})  # 这个是某个失败的记录重试
-        self.name_recorder_mapper = {}  # 做个映射，方便后面处理
-
     @classmethod
     def method_type(cls):
         return AutomationTypes.change_secret
@@ -74,54 +67,6 @@ class ChangeSecretManager(BaseChangeSecretPushManager):
         )
         return recorder
 
-    def on_host_success(self, host, result):
-        recorder = self.name_recorder_mapper.get(host)
-        if not recorder:
-            return
-        recorder.status = ChangeSecretRecordStatusChoice.success.value
-        recorder.date_finished = timezone.now()
-
-        account = recorder.account
-        if not account:
-            print("Account not found, deleted ?")
-            return
-
-        account.secret = recorder.new_secret
-        account.date_updated = timezone.now()
-
-        with safe_db_connection():
-            recorder.save(update_fields=['status', 'date_finished'])
-            account.save(update_fields=['secret', 'date_updated'])
-
-        self.summary['ok_accounts'] += 1
-        self.result['ok_accounts'].append(
-            {
-                "asset": str(account.asset),
-                "username": account.username,
-            }
-        )
-        super().on_host_success(host, result)
-
-    def on_host_error(self, host, error, result):
-        recorder = self.name_recorder_mapper.get(host)
-        if not recorder:
-            return
-        recorder.status = ChangeSecretRecordStatusChoice.failed.value
-        recorder.date_finished = timezone.now()
-        recorder.error = error
-        try:
-            recorder.save()
-        except Exception as e:
-            print(f"\033[31m Save {host} recorder error: {e} \033[0m\n")
-        self.summary['fail_accounts'] += 1
-        self.result['fail_accounts'].append(
-            {
-                "asset": str(recorder.asset),
-                "username": recorder.account.username,
-            }
-        )
-        super().on_host_success(host, result)
-
     def check_secret(self):
         if self.secret_strategy == SecretStrategy.custom \
                 and not self.execution.snapshot['secret']:

apps/accounts/automations/push_account/manager.py

@@ -1,9 +1,11 @@
 from django.utils.translation import gettext_lazy as _
 
 from accounts.const import AutomationTypes
+from common.decorators import bulk_create_decorator
 from common.utils import get_logger
 from common.utils.timezone import local_now_filename
 from ..base.manager import BaseChangeSecretPushManager
+from ...models import PushSecretRecord
 
 logger = get_logger(__name__)
 
@@ -17,12 +19,33 @@ class PushAccountManager(BaseChangeSecretPushManager):
         return account.secret
 
     def gen_account_inventory(self, account, asset, h, path_dir):
+        self.get_or_create_record(asset, account, h['name'])
         secret = self.get_secret(account)
         secret_type = account.secret_type
         new_secret, private_key_path = self.handle_ssh_secret(secret_type, secret, path_dir)
         h = self.gen_inventory(h, account, new_secret, private_key_path, asset)
         return h
 
+    def get_or_create_record(self, asset, account, name):
+        asset_account_id = f'{asset.id}-{account.id}'
+
+        if asset_account_id in self.record_map:
+            record_id = self.record_map[asset_account_id]
+            recorder = PushSecretRecord.objects.filter(id=record_id).first()
+        else:
+            recorder = self.create_record(asset, account)
+
+        self.name_recorder_mapper[name] = recorder
+        return recorder
+    
+    @bulk_create_decorator(PushSecretRecord)
+    def create_record(self, asset, account):
+        recorder = PushSecretRecord(
+            asset=asset, account=account, execution=self.execution,
+            comment=f'{account.username}@{asset.address}'
+        )
+        return recorder
+
     def print_summary(self):
         print('\n\n' + '-' * 80)
         plan_execution_end = _('Plan execution end')

apps/accounts/migrations/0029_alter_changesecretrecord_account_and_more.py

@@ -0,0 +1,51 @@
+# Generated by Django 4.1.13 on 2025-01-23 07:22
+
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0011_auto_20241204_1516'),
+        ('accounts', '0028_remove_checkaccountengine_is_active_and_more'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='changesecretrecord',
+            name='account',
+            field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='%(class)ss', to='accounts.account'),
+        ),
+        migrations.AlterField(
+            model_name='changesecretrecord',
+            name='asset',
+            field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='asset_%(class)ss', to='assets.asset'),
+        ),
+        migrations.AlterField(
+            model_name='changesecretrecord',
+            name='execution',
+            field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='execution_%(class)ss', to='accounts.automationexecution'),
+        ),
+        migrations.CreateModel(
+            name='PushSecretRecord',
+            fields=[
+                ('created_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Created by')),
+                ('updated_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Updated by')),
+                ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),
+                ('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')),
+                ('comment', models.TextField(blank=True, default='', verbose_name='Comment')),
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('date_finished', models.DateTimeField(blank=True, db_index=True, null=True, verbose_name='Date finished')),
+                ('status', models.CharField(default='pending', max_length=16, verbose_name='Status')),
+                ('error', models.TextField(blank=True, null=True, verbose_name='Error')),
+                ('account', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='%(class)ss', to='accounts.account')),
+                ('asset', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='asset_%(class)ss', to='assets.asset')),
+                ('execution', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='execution_%(class)ss', to='accounts.automationexecution')),
+            ],
+            options={
+                'verbose_name': 'Push secret record',
+            },
+        ),
+    ]

apps/accounts/models/automations/change_secret.py

@@ -9,7 +9,7 @@ from common.db import fields
 from common.db.models import JMSBaseModel
 from .base import AccountBaseAutomation, ChangeSecretMixin
 
-__all__ = ['ChangeSecretAutomation', 'ChangeSecretRecord', ]
+__all__ = ['ChangeSecretAutomation', 'ChangeSecretRecord', 'BaseSecretRecord']
 
 
 class ChangeSecretAutomation(ChangeSecretMixin, AccountBaseAutomation):
@@ -30,36 +30,42 @@ class ChangeSecretAutomation(ChangeSecretMixin, AccountBaseAutomation):
         return attr_json
 
 
-class ChangeSecretRecord(JMSBaseModel):
+class BaseSecretRecord(JMSBaseModel):
     account = models.ForeignKey(
         'accounts.Account', on_delete=models.SET_NULL,
-        null=True, related_name='change_secret_records'
+        null=True, related_name='%(class)ss'
     )
     asset = models.ForeignKey(
         'assets.Asset', on_delete=models.SET_NULL,
-        null=True, related_name='asset_change_secret_records'
+        null=True, related_name='asset_%(class)ss'
     )
     execution = models.ForeignKey(
         'accounts.AutomationExecution', on_delete=models.SET_NULL,
-        null=True, related_name='execution_change_secret_records',
+        null=True, related_name='execution_%(class)ss',
     )
-    old_secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Old secret'))
-    new_secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('New secret'))
     date_finished = models.DateTimeField(blank=True, null=True, verbose_name=_('Date finished'), db_index=True)
-    ignore_fail = models.BooleanField(default=False, verbose_name=_('Ignore fail'))
     status = models.CharField(
         max_length=16, verbose_name=_('Status'), default=ChangeSecretRecordStatusChoice.pending.value
     )
     error = models.TextField(blank=True, null=True, verbose_name=_('Error'))
 
     class Meta:
-        verbose_name = _("Change secret record")
+        abstract = True
 
     def __str__(self):
         return f'{self.account.username}@{self.asset}'
 
-    @staticmethod
+    @classmethod
-    def get_valid_records():
+    def get_valid_records(cls):
-        return ChangeSecretRecord.objects.exclude(
+        return cls.objects.exclude(
             Q(execution__isnull=True) | Q(asset__isnull=True) | Q(account__isnull=True)
         )
+
+
+class ChangeSecretRecord(BaseSecretRecord):
+    old_secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Old secret'))
+    new_secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('New secret'))
+    ignore_fail = models.BooleanField(default=False, verbose_name=_('Ignore fail'))
+
+    class Meta:
+        verbose_name = _("Change secret record")

apps/accounts/models/automations/push_account.py

@@ -3,10 +3,10 @@ from django.utils.translation import gettext_lazy as _
 
 from accounts.const import AutomationTypes
 from accounts.models import Account
-from .base import AccountBaseAutomation
+from .base import AccountBaseAutomation, ChangeSecretMixin
-from .change_secret import ChangeSecretMixin
+from .change_secret import BaseSecretRecord
 
-__all__ = ['PushAccountAutomation']
+__all__ = ['PushAccountAutomation', 'PushSecretRecord']
 
 
 class PushAccountAutomation(ChangeSecretMixin, AccountBaseAutomation):
@@ -36,3 +36,8 @@ class PushAccountAutomation(ChangeSecretMixin, AccountBaseAutomation):
 
     class Meta:
         verbose_name = _("Push asset account")
+
+
+class PushSecretRecord(BaseSecretRecord):
+    class Meta:
+        verbose_name = _("Push secret record")