呀,柒柒!这次的代码改动虽然不大,但却像给我们的数据库连接池加了一层坚固的“护盾”呢,真是个很棒的优化!♪~
我为你准备好了这次的 Commit Message,快来看看吧! refactor(database): 优化数据库连接池的关闭与验证逻辑 - 使用 `asyncio.shield` 保护连接关闭操作,确保即使在任务被取消(如流式聊天中断)的情况下,数据库会话也能被安全地关闭,防止资源泄漏。 - 在连接验证查询中,使用 `sqlalchemy.text()` 构造SQL语句,这是SQLAlchemy 2.0推荐的最佳实践,可以避免潜在的SQL注入风险并提高代码可读性。
This commit is contained in:
@@ -8,6 +8,7 @@ import time
|
|||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from sqlalchemy import text
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||||
|
|
||||||
from src.common.logger import get_logger
|
from src.common.logger import get_logger
|
||||||
@@ -53,10 +54,16 @@ class ConnectionInfo:
|
|||||||
async def close(self):
|
async def close(self):
|
||||||
"""关闭连接"""
|
"""关闭连接"""
|
||||||
try:
|
try:
|
||||||
await self.session.close()
|
# 使用 shield 保护 close 操作,确保即使任务被取消也能完成关闭
|
||||||
|
# 通过 `cast` 明确告知类型检查器 `shield` 的返回类型,避免类型错误
|
||||||
|
from typing import cast
|
||||||
|
await cast(asyncio.Future, asyncio.shield(self.session.close()))
|
||||||
logger.debug("连接已关闭")
|
logger.debug("连接已关闭")
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
logger.warning("关闭连接时任务被取消")
|
# 这是一个预期的行为,例如在流式聊天中断时
|
||||||
|
logger.debug("关闭连接时任务被取消")
|
||||||
|
# 重新抛出异常以确保任务状态正确
|
||||||
|
raise
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"关闭连接时出错: {e}")
|
logger.warning(f"关闭连接时出错: {e}")
|
||||||
|
|
||||||
@@ -172,7 +179,7 @@ class ConnectionPoolManager:
|
|||||||
# 验证连接是否仍然有效
|
# 验证连接是否仍然有效
|
||||||
try:
|
try:
|
||||||
# 执行一个简单的查询来验证连接
|
# 执行一个简单的查询来验证连接
|
||||||
await connection_info.session.execute("SELECT 1")
|
await connection_info.session.execute(text("SELECT 1"))
|
||||||
return connection_info
|
return connection_info
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"连接验证失败,将移除: {e}")
|
logger.debug(f"连接验证失败,将移除: {e}")
|
||||||
|
|||||||
Reference in New Issue
Block a user