Skip to content

Commit 6b67bf6

Browse files
committed
feat: 极大的提升 SQL 工具调用的稳定性,移除 LIMINT 并将默认超时时间修改为 60s,添加MySQL连接管理器和安全检查功能,优化查询执行和错误处理逻辑
- 重构 MySQLConnectionManager,在使用前验证游标,向使用者提供辅助函数,并使用基于线程池的超时机制替代 SIGALRM,使查询可运行长达 600 秒,同时安全地失效不良连接(src/agents/common/toolkits/mysql/connection.py) - 扩大安全防护中的允许超时窗口,移除未使用的限制验证器,并更新查询模型默认值(src/agents/common/toolkits/mysql/security.py,src/agents/common/toolkits/mysql/tools.py) - 使 mysql_query 依赖 execute_query_with_timeout,移除隐式 LIMIT 注入,确保失败时连接保持健康,扩展结果格式化和故障排除提示;记录 list-table 调用,并在避免按表 COUNT 扫描的同时,丰富 mysql_describe_table 的列注释(src/agents/common/toolkits/mysql/tools.py) - 将高级代理指南重命名为 agents-config,更新 VitePress 导航,并在文档中增加 LangGraph/MCP/MySQL 配置指导(docs/.vitepress/config.mts,docs/advanced/agents-config.md) - 调整路线图错误列表,将 LangGraph 工具问题标记为已解决,并跟踪 LightRAG 查看器的缺失功能(docs/changelog/roadmap.md) - 添加 CLAUDE.md,引导贡献者参阅 AGENTS.md 获取更多信息
1 parent 093084e commit 6b67bf6

7 files changed

Lines changed: 196 additions & 137 deletions

File tree

CLAUDE.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Read AGENTS.md

docs/.vitepress/config.mts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ export default defineConfig({
3535
items: [
3636
{ text: '配置系统详解', link: '/advanced/configuration' },
3737
{ text: '文档解析', link: '/advanced/document-processing' },
38-
{ text: '智能体', link: '/advanced/agents' },
38+
{ text: '智能体', link: '/advanced/agents-config' },
3939
{ text: '品牌自定义', link: '/advanced/branding' },
4040
{ text: '其他配置', link: '/advanced/misc' }
4141
]
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ MYSQL_CHARSET=utf8mb4
101101

102102
-**只读操作**: 仅允许 SELECT、SHOW、DESCRIBE、EXPLAIN 操作
103103
-**SQL 注入防护**: 严格的表名参数验证
104-
-**超时控制**: 默认 10 秒,最大 60
104+
-**超时控制**: 默认 60 秒,最大 600
105105
-**结果限制**: 默认 10000 字符,100 行,最大 1000 行
106106

107107
### 可视化图表-MCP-Server

docs/changelog/roadmap.md

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77

88
## Bugs
99

10-
- [ ] v1 版本的 LangGraph 的工具渲染有问题
1110
- [ ] upload 接口会阻塞主进程
11+
- [ ] LightRAG 知识库查看不了解析后的文本
1212

1313
## Next
1414

@@ -39,4 +39,5 @@
3939
- [x] 修改现有的智能体Demo,并尽量将默认助手的特性兼容到 LangGraph 的 [`create_agent`](https://docs.langchain.com/oss/python/langchain/agents)
4040
- [x] 基于 create_agent 创建 SQL Viewer 智能体 <Badge type="info" text="0.3.5" />
4141
- [x] 优化 MCP 逻辑,支持 common + special 创建方式 <Badge type="info" text="0.3.5" />
42-
- [x] 修复本地知识库的 metadata 和 向量数据库中不一致的情况。
42+
- [x] 修复本地知识库的 metadata 和 向量数据库中不一致的情况。
43+
- [x] v1 版本的 LangGraph 的工具渲染有问题

src/agents/common/toolkits/mysql/connection.py

Lines changed: 83 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import signal
1+
import concurrent.futures
22
import threading
33
import time
44
from contextlib import contextmanager
@@ -93,55 +93,85 @@ def test_connection(self) -> bool:
9393
pass
9494
return False
9595

96+
def _invalidate_connection(self, connection: pymysql.Connection | None = None):
97+
"""关闭并清理失效的连接"""
98+
try:
99+
if connection:
100+
connection.close()
101+
except Exception:
102+
pass
103+
finally:
104+
self.connection = None
105+
96106
@contextmanager
97107
def get_cursor(self):
98108
"""获取数据库游标的上下文管理器"""
99109
max_retries = 2
110+
cursor = None
111+
connection = None
112+
last_error: Exception | None = None
113+
114+
# 优先确保成功获取游标再交给调用方执行查询
100115
for attempt in range(max_retries):
101116
try:
102117
connection = self._get_connection()
103118
cursor = connection.cursor()
104-
105-
try:
106-
yield cursor
107-
connection.commit()
108-
break # 成功,退出重试循环
109-
110-
except Exception as e:
111-
connection.rollback()
112-
113-
# 如果是连接错误,尝试重新连接
114-
if "MySQL" in str(e) or "connection" in str(e).lower():
115-
if attempt < max_retries - 1:
116-
logger.warning(f"Connection error, retrying (attempt {attempt + 1}): {e}")
117-
# 强制重新连接
118-
if self.connection:
119-
try:
120-
self.connection.close()
121-
except Exception as _:
122-
pass
123-
self.connection = None
124-
time.sleep(1)
125-
continue
126-
127-
raise e # 其他错误直接抛出
128-
129-
finally:
130-
if cursor:
131-
cursor.close()
132-
119+
break
133120
except Exception as e:
121+
last_error = e
122+
logger.warning(f"Failed to acquire cursor (attempt {attempt + 1}): {e}")
123+
self._invalidate_connection(connection)
124+
cursor = None
125+
connection = None
134126
if attempt == max_retries - 1:
135-
raise e # 最后一次尝试失败,抛出异常
127+
raise e
136128
time.sleep(1)
137129

130+
if cursor is None or connection is None:
131+
raise last_error or ConnectionError("Unable to acquire MySQL cursor")
132+
133+
try:
134+
yield cursor
135+
connection.commit()
136+
except Exception as e:
137+
try:
138+
connection.rollback()
139+
except Exception:
140+
pass
141+
142+
# 标记连接失效,等待下一次获取时重建
143+
if "MySQL" in str(e) or "connection" in str(e).lower():
144+
logger.warning(f"MySQL connection error encountered, invalidating connection: {e}")
145+
self._invalidate_connection(connection)
146+
147+
raise
148+
finally:
149+
if cursor:
150+
try:
151+
cursor.close()
152+
except Exception:
153+
pass
154+
138155
def close(self):
139156
"""关闭数据库连接"""
140157
if self.connection:
141158
self.connection.close()
142159
self.connection = None
143160
logger.info("MySQL connection closed")
144161

162+
def get_connection(self) -> pymysql.Connection:
163+
"""对外暴露的连接获取方法"""
164+
return self._get_connection()
165+
166+
def invalidate_connection(self):
167+
"""手动标记连接失效"""
168+
self._invalidate_connection(self.connection)
169+
170+
@property
171+
def database_name(self) -> str:
172+
"""返回当前配置的数据库名称"""
173+
return self.config["database"]
174+
145175

146176
class QueryTimeoutError(Exception):
147177
"""查询超时异常"""
@@ -156,25 +186,30 @@ class QueryResultTooLargeError(Exception):
156186

157187

158188
def execute_query_with_timeout(connection: pymysql.Connection, sql: str, params: tuple = None, timeout: int = 10):
159-
"""带超时的查询执行"""
189+
"""使用线程池实现超时控制,避免信号导致的生成器问题"""
160190

161-
def timeout_handler(_signum, _frame):
162-
raise QueryTimeoutError(f"Query timeout after {timeout} seconds")
163-
164-
# 设置信号处理
165-
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
166-
signal.alarm(timeout)
167-
168-
try:
191+
def query_worker():
192+
"""查询工作函数,在单独线程中执行"""
169193
cursor = connection.cursor(DictCursor)
170-
cursor.execute(sql, params or ())
171-
result = cursor.fetchall()
172-
cursor.close()
173-
return result
174-
finally:
175-
# 恢复原始信号处理
176-
signal.alarm(0)
177-
signal.signal(signal.SIGALRM, old_handler)
194+
try:
195+
if params is None:
196+
cursor.execute(sql)
197+
else:
198+
cursor.execute(sql, params)
199+
result = cursor.fetchall()
200+
return result
201+
finally:
202+
cursor.close()
203+
204+
# 使用线程池执行查询,设置超时
205+
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
206+
future = executor.submit(query_worker)
207+
try:
208+
return future.result(timeout=timeout)
209+
except concurrent.futures.TimeoutError:
210+
# 尝试取消任务
211+
future.cancel()
212+
raise QueryTimeoutError(f"Query timeout after {timeout} seconds")
178213

179214

180215
def limit_result_size(result: list, max_chars: int = 10000) -> list:

src/agents/common/toolkits/mysql/security.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,7 @@ def validate_table_name(cls, table_name: str) -> bool:
7676
# 检查表名只包含字母、数字、下划线
7777
return bool(re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table_name))
7878

79-
@classmethod
80-
def validate_limit(cls, limit: int) -> bool:
81-
"""验证limit参数"""
82-
return isinstance(limit, int) and 0 < limit <= 1000
83-
8479
@classmethod
8580
def validate_timeout(cls, timeout: int) -> bool:
8681
"""验证timeout参数"""
87-
return isinstance(timeout, int) and 1 <= timeout <= 60
82+
return isinstance(timeout, int) and 1 <= timeout <= 600

0 commit comments

Comments
 (0)