diff --git a/apps/common/apps.py b/apps/common/apps.py index 18ed8a9a9..23913256d 100644 --- a/apps/common/apps.py +++ b/apps/common/apps.py @@ -4,6 +4,7 @@ import os import sys from django.apps import AppConfig +from django.db import close_old_connections class CommonConfig(AppConfig): @@ -21,3 +22,4 @@ class CommonConfig(AppConfig): if not os.environ.get('DJANGO_DEBUG_SHELL'): django_ready.send(CommonConfig) + close_old_connections() diff --git a/apps/common/db/utils.py b/apps/common/db/utils.py index c127e84aa..fd05016da 100644 --- a/apps/common/db/utils.py +++ b/apps/common/db/utils.py @@ -56,7 +56,7 @@ def close_old_connections(): @contextmanager -def safe_db_connection(): +def safe_db_connection(auto_close=False): in_atomic_block = connection.in_atomic_block # 当前是否处于事务中 autocommit = transaction.get_autocommit() # 是否启用了自动提交 created = False @@ -69,7 +69,7 @@ def safe_db_connection(): yield finally: # 如果不是事务中(API 请求中可能需要提交事务),则关闭连接 - if created and not in_atomic_block and autocommit: + 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/notifications/ws.py b/apps/notifications/ws.py index 5616694de..e81da8542 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(): + with safe_db_connection(auto_close=True): self.send_unread_msg_count() def handle_new_site_msg_recv(msg): diff --git a/apps/ops/signal_handlers.py b/apps/ops/signal_handlers.py index a5afdd2ee..fb5b89851 100644 --- a/apps/ops/signal_handlers.py +++ b/apps/ops/signal_handlers.py @@ -115,11 +115,10 @@ def on_celery_task_pre_run(task_id='', kwargs=None, **others): @signals.task_postrun.connect def on_celery_task_post_run(task_id='', state='', **kwargs): - close_old_connections() - CeleryTaskExecution.objects.filter(id=task_id).update( state=state, date_finished=timezone.now(), is_finished=True ) + close_old_connections() @signals.after_task_publish.connect diff --git a/apps/terminal/ws.py b/apps/terminal/ws.py index 372647e7e..24bb5eb11 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(): + with safe_db_connection(auto_close=True): 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: