Browse Source

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>
tags/1.8.1
耐小心 2 months ago
parent
commit
acd209a890
No account linked to committer's email address

+ 6
- 0
api/core/app/apps/advanced_chat/app_generator.py View File



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,

+ 2
- 1
api/core/app/apps/advanced_chat/app_runner.py View File

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



+ 7
- 2
api/core/app/apps/message_based_app_generator.py View File

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

+ 6
- 2
api/core/app/task_pipeline/message_cycle_manager.py View File

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

Loading…
Cancel
Save