From f826f43495136eefc75ecd42b4948d500d4122ab Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 14 May 2025 15:05:12 +0800
Subject: [PATCH] perf: simplify db using

---
 apps/common/db/utils.py   | 10 ++++++++--
 apps/common/decorators.py |  4 ++--
 apps/notifications/ws.py  |  2 +-
 apps/terminal/ws.py       |  2 +-
 4 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/apps/common/db/utils.py b/apps/common/db/utils.py
index fd05016da..a3099c2cf 100644
--- a/apps/common/db/utils.py
+++ b/apps/common/db/utils.py
@@ -56,7 +56,14 @@ def close_old_connections():
 
 
 @contextmanager
-def safe_db_connection(auto_close=False):
+def safe_db_connection(auto_close=True):
+    close_old_connections()
+    yield
+    close_old_connections()
+
+
+@contextmanager
+def safe_atomic_db_connection(auto_close=False):
     in_atomic_block = connection.in_atomic_block  # 当前是否处于事务中
     autocommit = transaction.get_autocommit()  # 是否启用了自动提交
     created = False
@@ -70,7 +77,6 @@ def safe_db_connection(auto_close=False):
     finally:
         # 如果不是事务中（API 请求中可能需要提交事务），则关闭连接
         if auto_close or (created and not in_atomic_block and autocommit):
-            print("close connection in safe_db_connection")
             close_old_connections()
 
 
diff --git a/apps/common/decorators.py b/apps/common/decorators.py
index 6e9cc0926..963c5c481 100644
--- a/apps/common/decorators.py
+++ b/apps/common/decorators.py
@@ -11,7 +11,7 @@ from functools import wraps
 
 from django.db import transaction
 
-from .db.utils import open_db_connection, safe_db_connection
+from .db.utils import open_db_connection, safe_atomic_db_connection
 from .utils import logger
 
 
@@ -318,7 +318,7 @@ def bulk_handle(handler, batch_size=50, timeout=0.5):
                 if not cache:
                     return
                 with tmp_to_org(org_id):
-                    with safe_db_connection():
+                    with safe_atomic_db_connection():
                         handler(cache)
                 cache.clear()
 
diff --git a/apps/notifications/ws.py b/apps/notifications/ws.py
index e81da8542..5616694de 100644
--- a/apps/notifications/ws.py
+++ b/apps/notifications/ws.py
@@ -50,7 +50,7 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         user_id = str(self.scope["user"].id)
 
         # 先发一个消息再说
-        with safe_db_connection(auto_close=True):
+        with safe_db_connection():
             self.send_unread_msg_count()
 
         def handle_new_site_msg_recv(msg):
diff --git a/apps/terminal/ws.py b/apps/terminal/ws.py
index 24bb5eb11..372647e7e 100644
--- a/apps/terminal/ws.py
+++ b/apps/terminal/ws.py
@@ -52,7 +52,7 @@ class TerminalTaskWebsocket(JsonWebsocketConsumer):
         self.send(bytes_data=content)
 
     def get_terminal_tasks(self, task_id=None):
-        with safe_db_connection(auto_close=True):
+        with safe_db_connection():
             critical_time = timezone.now() - datetime.timedelta(minutes=10)
             tasks = self.terminal.task_set.filter(is_finished=False, date_created__gte=critical_time)
             if task_id: