fix: prevent database connection leaks in chatflow mode by using Session-managed queries (#24656)

Co-authored-by: 王锶奇 <wangsiqi2@tal.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
耐小心
2025-09-01 18:22:42 +08:00
committed by GitHub
parent bd482eb8ef
commit acd209a890
4 changed files with 21 additions and 5 deletions

View File

@@ -450,6 +450,12 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
worker_thread.start() worker_thread.start()
# release database connection, because the following new thread operations may take a long time
db.session.refresh(workflow)
db.session.refresh(message)
db.session.refresh(user)
db.session.close()
# return response or stream generator # return response or stream generator
response = self._handle_advanced_chat_response( response = self._handle_advanced_chat_response(
application_generate_entity=application_generate_entity, application_generate_entity=application_generate_entity,

View File

@@ -72,7 +72,8 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
app_config = self.application_generate_entity.app_config app_config = self.application_generate_entity.app_config
app_config = cast(AdvancedChatAppConfig, app_config) app_config = cast(AdvancedChatAppConfig, app_config)
app_record = db.session.query(App).where(App.id == app_config.app_id).first() with Session(db.engine, expire_on_commit=False) as session:
app_record = session.scalar(select(App).where(App.id == app_config.app_id))
if not app_record: if not app_record:
raise ValueError("App not found") raise ValueError("App not found")

View File

@@ -3,6 +3,9 @@ import logging
from collections.abc import Generator from collections.abc import Generator
from typing import Optional, Union, cast from typing import Optional, Union, cast
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom
from core.app.apps.base_app_generator import BaseAppGenerator from core.app.apps.base_app_generator import BaseAppGenerator
from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.apps.base_app_queue_manager import AppQueueManager
@@ -253,7 +256,8 @@ class MessageBasedAppGenerator(BaseAppGenerator):
:param conversation_id: conversation id :param conversation_id: conversation id
:return: conversation :return: conversation
""" """
conversation = db.session.query(Conversation).where(Conversation.id == conversation_id).first() with Session(db.engine, expire_on_commit=False) as session:
conversation = session.scalar(select(Conversation).where(Conversation.id == conversation_id))
if not conversation: if not conversation:
raise ConversationNotExistsError("Conversation not exists") raise ConversationNotExistsError("Conversation not exists")
@@ -266,7 +270,8 @@ class MessageBasedAppGenerator(BaseAppGenerator):
:param message_id: message id :param message_id: message id
:return: message :return: message
""" """
message = db.session.query(Message).where(Message.id == message_id).first() with Session(db.engine, expire_on_commit=False) as session:
message = session.scalar(select(Message).where(Message.id == message_id))
if message is None: if message is None:
raise MessageNotExistsError("Message not exists") raise MessageNotExistsError("Message not exists")

View File

@@ -3,6 +3,8 @@ from threading import Thread
from typing import Optional, Union from typing import Optional, Union
from flask import Flask, current_app from flask import Flask, current_app
from sqlalchemy import select
from sqlalchemy.orm import Session
from configs import dify_config from configs import dify_config
from core.app.entities.app_invoke_entities import ( from core.app.entities.app_invoke_entities import (
@@ -143,7 +145,8 @@ class MessageCycleManager:
:param event: event :param event: event
:return: :return:
""" """
message_file = db.session.query(MessageFile).where(MessageFile.id == event.message_file_id).first() with Session(db.engine, expire_on_commit=False) as session:
message_file = session.scalar(select(MessageFile).where(MessageFile.id == event.message_file_id))
if message_file and message_file.url is not None: if message_file and message_file.url is not None:
# get tool file id # get tool file id
@@ -183,7 +186,8 @@ class MessageCycleManager:
:param message_id: message id :param message_id: message id
:return: :return:
""" """
message_file = db.session.query(MessageFile).where(MessageFile.id == message_id).first() with Session(db.engine, expire_on_commit=False) as session:
message_file = session.scalar(select(MessageFile).where(MessageFile.id == message_id))
event_type = StreamEvent.MESSAGE_FILE if message_file else StreamEvent.MESSAGE event_type = StreamEvent.MESSAGE_FILE if message_file else StreamEvent.MESSAGE
return MessageStreamResponse( return MessageStreamResponse(