From e1f03a194bac6b90f8592e0d28527af019e6c697 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Thu, 15 May 2025 15:15:32 +0800
Subject: [PATCH] fix: Asset list gather account raise connection already
 closed

---
 apps/accounts/automations/base/manager.py | 6 +++---
 apps/assets/automations/base/manager.py   | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/apps/accounts/automations/base/manager.py b/apps/accounts/automations/base/manager.py
index bb94a2213..4c8263e5d 100644
--- a/apps/accounts/automations/base/manager.py
+++ b/apps/accounts/automations/base/manager.py
@@ -10,7 +10,7 @@ from accounts.models import BaseAccountQuerySet
 from accounts.utils import SecretGenerator
 from assets.automations.base.manager import BasePlaybookManager
 from assets.const import HostTypes
-from common.db.utils import safe_db_connection
+from common.db.utils import safe_atomic_db_connection
 from common.utils import get_logger
 
 logger = get_logger(__name__)
@@ -170,7 +170,7 @@ class BaseChangeSecretPushManager(AccountBasePlaybookManager):
         )
         super().on_host_success(host, result)
 
-        with safe_db_connection():
+        with safe_atomic_db_connection():
             account.save(update_fields=['secret', 'date_updated', 'date_change_secret', 'change_secret_status'])
             self.save_record(recorder)
 
@@ -198,6 +198,6 @@ class BaseChangeSecretPushManager(AccountBasePlaybookManager):
         )
         super().on_host_error(host, error, result)
 
-        with safe_db_connection():
+        with safe_atomic_db_connection():
             account.save(update_fields=['change_secret_status', 'date_change_secret', 'date_updated'])
             self.save_record(recorder)
diff --git a/apps/assets/automations/base/manager.py b/apps/assets/automations/base/manager.py
index 8446dbeb4..e4999d363 100644
--- a/apps/assets/automations/base/manager.py
+++ b/apps/assets/automations/base/manager.py
@@ -17,7 +17,7 @@ from sshtunnel import SSHTunnelForwarder
 
 from assets.automations.methods import platform_automation_methods
 from common.const import Status
-from common.db.utils import safe_db_connection
+from common.db.utils import safe_atomic_db_connection
 from common.tasks import send_mail_async
 from common.utils import get_logger, lazyproperty, is_openssh_format_key, ssh_pubkey_gen
 from ops.ansible import JMSInventory, DefaultCallback, SuperPlaybookRunner
@@ -123,7 +123,7 @@ class BaseManager:
         self.execution.result = self.result
         self.execution.status = self.status
 
-        with safe_db_connection():
+        with safe_atomic_db_connection():
             self.execution.save()
 
     def print_summary(self):