From ad60e6eea8b2f0751eead89d88f0657d7ed13e7a Mon Sep 17 00:00:00 2001 From: tt-P607 <68868379+tt-P607@users.noreply.github.com> Date: Fri, 31 Oct 2025 23:11:38 +0800 Subject: [PATCH] =?UTF-8?q?=E5=91=80=EF=BC=8C=E6=9F=92=E6=9F=92=EF=BC=81?= =?UTF-8?q?=E8=BF=99=E6=AC=A1=E7=9A=84=E4=BB=A3=E7=A0=81=E6=94=B9=E5=8A=A8?= =?UTF-8?q?=E8=99=BD=E7=84=B6=E4=B8=8D=E5=A4=A7=EF=BC=8C=E4=BD=86=E5=8D=B4?= =?UTF-8?q?=E5=83=8F=E7=BB=99=E6=88=91=E4=BB=AC=E7=9A=84=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E5=BA=93=E8=BF=9E=E6=8E=A5=E6=B1=A0=E5=8A=A0=E4=BA=86=E4=B8=80?= =?UTF-8?q?=E5=B1=82=E5=9D=9A=E5=9B=BA=E7=9A=84=E2=80=9C=E6=8A=A4=E7=9B=BE?= =?UTF-8?q?=E2=80=9D=E5=91=A2=EF=BC=8C=E7=9C=9F=E6=98=AF=E4=B8=AA=E5=BE=88?= =?UTF-8?q?=E6=A3=92=E7=9A=84=E4=BC=98=E5=8C=96=EF=BC=81=E2=99=AA~?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 我为你准备好了这次的 Commit Message，快来看看吧！ refactor(database): 优化数据库连接池的关闭与验证逻辑 - 使用 `asyncio.shield` 保护连接关闭操作，确保即使在任务被取消（如流式聊天中断）的情况下，数据库会话也能被安全地关闭，防止资源泄漏。 - 在连接验证查询中，使用 `sqlalchemy.text()` 构造SQL语句，这是SQLAlchemy 2.0推荐的最佳实践，可以避免潜在的SQL注入风险并提高代码可读性。 --- src/common/database/connection_pool_manager.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/common/database/connection_pool_manager.py b/src/common/database/connection_pool_manager.py index e295e2284..e1df6fd2e 100644 --- a/src/common/database/connection_pool_manager.py +++ b/src/common/database/connection_pool_manager.py @@ -8,6 +8,7 @@ import time from contextlib import asynccontextmanager from typing import Any +from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from src.common.logger import get_logger @@ -53,10 +54,16 @@ class ConnectionInfo: async def close(self): """关闭连接""" try: - await self.session.close() + # 使用 shield 保护 close 操作，确保即使任务被取消也能完成关闭 + # 通过 `cast` 明确告知类型检查器 `shield` 的返回类型，避免类型错误 + from typing import cast + await cast(asyncio.Future, asyncio.shield(self.session.close())) logger.debug("连接已关闭") except asyncio.CancelledError: - logger.warning("关闭连接时任务被取消") + # 这是一个预期的行为，例如在流式聊天中断时 + logger.debug("关闭连接时任务被取消") + # 重新抛出异常以确保任务状态正确 + raise except Exception as e: logger.warning(f"关闭连接时出错: {e}") @@ -172,7 +179,7 @@ class ConnectionPoolManager: # 验证连接是否仍然有效 try: # 执行一个简单的查询来验证连接 - await connection_info.session.execute("SELECT 1") + await connection_info.session.execute(text("SELECT 1")) return connection_info except Exception as e: logger.debug(f"连接验证失败，将移除: {e}")