perf: simplify db using

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()
 
 

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()
 

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):

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: