357 lines
13 KiB
Python
357 lines
13 KiB
Python
from typing import Annotated, Any
|
||
|
||
from pydantic import BaseModel, Field
|
||
|
||
from yuxi.agents.toolkits.registry import tool
|
||
from yuxi.utils import logger
|
||
|
||
from .connection import (
|
||
MySQLConnectionManager,
|
||
QueryTimeoutError,
|
||
execute_query_with_timeout,
|
||
limit_result_size,
|
||
)
|
||
from .exceptions import MySQLConnectionError
|
||
from .security import MySQLSecurityChecker
|
||
|
||
# 全局连接管理器实例
|
||
_connection_manager: MySQLConnectionManager | None = None
|
||
|
||
|
||
def get_connection_manager() -> MySQLConnectionManager:
|
||
"""获取全局连接管理器"""
|
||
global _connection_manager
|
||
if _connection_manager is None:
|
||
import os
|
||
|
||
# 从环境变量中读取 MySQL 配置
|
||
mysql_config = {
|
||
"host": os.getenv("MYSQL_HOST"),
|
||
"user": os.getenv("MYSQL_USER"),
|
||
"password": os.getenv("MYSQL_PASSWORD"),
|
||
"database": os.getenv("MYSQL_DATABASE"),
|
||
"port": int(os.getenv("MYSQL_PORT") or "3306"),
|
||
"charset": "utf8mb4",
|
||
"description": os.getenv("MYSQL_DATABASE_DESCRIPTION") or "默认 MySQL 数据库",
|
||
}
|
||
# 验证配置完整性
|
||
required_keys = ["host", "user", "password", "database"]
|
||
for key in required_keys:
|
||
if not mysql_config[key]:
|
||
raise MySQLConnectionError(
|
||
f"MySQL configuration missing required key: {key}, please check your environment variables."
|
||
)
|
||
|
||
_connection_manager = MySQLConnectionManager(mysql_config)
|
||
return _connection_manager
|
||
|
||
|
||
@tool(
|
||
category="mysql",
|
||
tags=["数据库", "查询"],
|
||
display_name="列出MySQL表",
|
||
name_or_callable="mysql_list_tables",
|
||
)
|
||
def mysql_list_tables() -> str:
|
||
"""【查询表名及说明】获取数据库中的所有表名
|
||
|
||
这个工具用来列出当前数据库中所有的表名，帮助你了解数据库的结构。
|
||
"""
|
||
try:
|
||
conn_manager = get_connection_manager()
|
||
|
||
with conn_manager.get_cursor() as cursor:
|
||
# 获取表名
|
||
cursor.execute("SHOW TABLES")
|
||
logger.debug("Executed `SHOW TABLES` query")
|
||
tables = cursor.fetchall()
|
||
|
||
if not tables:
|
||
return "数据库中没有找到任何表"
|
||
|
||
# 提取表名
|
||
table_names = []
|
||
for table in tables:
|
||
table_name = list(table.values())[0]
|
||
table_names.append(table_name)
|
||
|
||
# 获取每个表的行数信息
|
||
# table_info = []
|
||
# for table_name in table_names:
|
||
# try:
|
||
# cursor.execute(f"SELECT COUNT(*) as count FROM `{table_name}`")
|
||
# logger.debug(f"Executed `SELECT COUNT(*) FROM {table_name}` query")
|
||
# count_result = cursor.fetchone()
|
||
# row_count = count_result["count"]
|
||
# table_info.append(f"- {table_name} (约 {row_count} 行)")
|
||
# except Exception:
|
||
# table_info.append(f"- {table_name} (无法获取行数)")
|
||
|
||
all_table_names = "\n".join(table_names)
|
||
result = f"数据库中的表:\n{all_table_names}"
|
||
if db_note := conn_manager.config.get("description"):
|
||
result = f"数据库说明: {db_note}\n\n" + result
|
||
logger.info(f"Retrieved {len(table_names)} tables from database")
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = f"获取表名失败: {str(e)}"
|
||
logger.error(error_msg)
|
||
return error_msg
|
||
|
||
|
||
class TableDescribeModel(BaseModel):
|
||
"""获取表结构的参数模型"""
|
||
|
||
table_name: str = Field(description="要查询的表名", example="users")
|
||
|
||
|
||
@tool(
|
||
category="mysql",
|
||
tags=["数据库", "结构"],
|
||
display_name="描述MySQL表结构",
|
||
name_or_callable="mysql_describe_table",
|
||
args_schema=TableDescribeModel,
|
||
)
|
||
def mysql_describe_table(table_name: Annotated[str, "要查询结构的表名"]) -> str:
|
||
"""【描述表】获取指定表的详细结构信息
|
||
|
||
这个工具用来查看表的字段信息、数据类型、是否允许NULL、默认值、键类型等。
|
||
帮助你了解表的结构，以便编写正确的SQL查询。
|
||
"""
|
||
try:
|
||
# 验证表名安全性
|
||
if not MySQLSecurityChecker.validate_table_name(table_name):
|
||
return "表名包含非法字符，请检查表名"
|
||
|
||
conn_manager = get_connection_manager()
|
||
|
||
with conn_manager.get_cursor() as cursor:
|
||
# 获取表结构
|
||
cursor.execute(f"DESCRIBE `{table_name}`")
|
||
columns = cursor.fetchall()
|
||
|
||
if not columns:
|
||
return f"表 {table_name} 不存在或没有字段"
|
||
|
||
# 获取字段备注信息
|
||
column_comments: dict[str, str] = {}
|
||
try:
|
||
cursor.execute(
|
||
"""
|
||
SELECT COLUMN_NAME, COLUMN_COMMENT
|
||
FROM information_schema.COLUMNS
|
||
WHERE TABLE_NAME = %s AND TABLE_SCHEMA = %s
|
||
""",
|
||
(table_name, conn_manager.database_name),
|
||
)
|
||
comment_rows = cursor.fetchall()
|
||
for row in comment_rows:
|
||
column_name = row.get("COLUMN_NAME")
|
||
if column_name:
|
||
column_comments[column_name] = row.get("COLUMN_COMMENT") or ""
|
||
except Exception as e:
|
||
logger.warning(f"Failed to fetch column comments for table {table_name}: {e}")
|
||
|
||
# 格式化输出
|
||
result = f"表 `{table_name}` 的结构:\n\n"
|
||
result += "字段名\t\t类型\t\tNULL\t键\t默认值\t\t额外\t备注\n"
|
||
result += "-" * 80 + "\n"
|
||
|
||
for col in columns:
|
||
field = col["Field"] or ""
|
||
type_str = col["Type"] or ""
|
||
null_str = col["Null"] or ""
|
||
key_str = col["Key"] or ""
|
||
default_str = col.get("Default") or ""
|
||
extra_str = col.get("Extra") or ""
|
||
comment_str = column_comments.get(field, "")
|
||
|
||
# 格式化输出
|
||
result += (
|
||
f"{field:<16}\t{type_str:<16}\t{null_str:<8}\t{key_str:<4}\t"
|
||
f"{default_str:<16}\t{extra_str:<16}\t{comment_str}\n"
|
||
)
|
||
|
||
# 获取索引信息
|
||
try:
|
||
cursor.execute(f"SHOW INDEX FROM `{table_name}`")
|
||
indexes = cursor.fetchall()
|
||
|
||
if indexes:
|
||
result += "\n索引信息:\n"
|
||
index_dict = {}
|
||
for idx in indexes:
|
||
key_name = idx["Key_name"]
|
||
if key_name not in index_dict:
|
||
index_dict[key_name] = []
|
||
index_dict[key_name].append(idx["Column_name"])
|
||
|
||
for key_name, columns in index_dict.items():
|
||
result += f"- {key_name}: {', '.join(columns)}\n"
|
||
except Exception as e:
|
||
logger.warning(f"Failed to get index info for table {table_name}: {e}")
|
||
|
||
logger.info(f"Retrieved structure for table {table_name}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = f"获取表 {table_name} 结构失败: {str(e)}"
|
||
logger.error(error_msg)
|
||
return error_msg
|
||
|
||
|
||
class QueryModel(BaseModel):
|
||
"""执行SQL查询的参数模型"""
|
||
|
||
sql: str = Field(description="要执行的SQL查询语句（只能是SELECT语句）", example="SELECT * FROM users WHERE id = 1")
|
||
timeout: int | None = Field(default=60, description="查询超时时间（秒），默认60秒，最大600秒", ge=1, le=600)
|
||
|
||
|
||
@tool(
|
||
category="mysql",
|
||
tags=["数据库", "SQL"],
|
||
display_name="执行MySQL查询",
|
||
name_or_callable="mysql_query",
|
||
args_schema=QueryModel,
|
||
)
|
||
def mysql_query(
|
||
sql: Annotated[str, "要执行的SQL查询语句（只能是SELECT语句）"],
|
||
timeout: Annotated[int | None, "查询超时时间（秒），默认60秒，最大600秒"] = 60,
|
||
) -> str:
|
||
"""【执行 SQL 查询】执行只读的SQL查询语句
|
||
|
||
这个工具用来执行SQL查询并返回结果。支持复杂的SELECT查询，包括JOIN、GROUP BY等。
|
||
注意：只能执行查询操作，不能修改数据。
|
||
|
||
参数:
|
||
- sql: SQL查询语句
|
||
- timeout: 查询超时时间（防止长时间运行的查询）
|
||
"""
|
||
try:
|
||
# 验证SQL安全性
|
||
if not MySQLSecurityChecker.validate_sql(sql):
|
||
return "SQL语句包含不安全的操作或可能的注入攻击，请检查SQL语句"
|
||
|
||
if not MySQLSecurityChecker.validate_timeout(timeout):
|
||
return "timeout参数必须在1-600之间"
|
||
|
||
conn_manager = get_connection_manager()
|
||
connection = conn_manager.get_connection()
|
||
|
||
effective_timeout = timeout or 60
|
||
try:
|
||
result = execute_query_with_timeout(connection, sql, timeout=effective_timeout)
|
||
except QueryTimeoutError as timeout_error:
|
||
logger.error(f"MySQL query timed out after {effective_timeout} seconds: {timeout_error}")
|
||
raise
|
||
except Exception:
|
||
conn_manager.invalidate_connection()
|
||
raise
|
||
|
||
if not result:
|
||
return "查询执行成功，但没有返回任何结果"
|
||
|
||
# 限制结果大小
|
||
limited_result = limit_result_size(result, max_chars=10000)
|
||
|
||
# 检查结果是否被截断
|
||
if len(limited_result) < len(result):
|
||
warning = f"\n\n⚠️ 警告: 查询结果过大，只显示了前 {len(limited_result)} 行（共 {len(result)} 行）。\n"
|
||
warning += "建议使用更精确的查询条件或使用LIMIT子句来减少返回的数据量。"
|
||
else:
|
||
warning = ""
|
||
|
||
# 格式化输出
|
||
if limited_result:
|
||
# 获取列名
|
||
columns = list(limited_result[0].keys())
|
||
|
||
# 计算每列的最大宽度
|
||
col_widths = {}
|
||
for col in columns:
|
||
col_widths[col] = max(len(str(col)), max(len(str(row.get(col, ""))) for row in limited_result))
|
||
col_widths[col] = min(col_widths[col], 50) # 限制最大宽度
|
||
|
||
# 构建表头
|
||
header = "| " + " | ".join(f"{col:<{col_widths[col]}}" for col in columns) + " |"
|
||
separator = "|" + "|".join("-" * (col_widths[col] + 2) for col in columns) + "|"
|
||
|
||
# 构建数据行
|
||
rows = []
|
||
for row in limited_result:
|
||
row_str = "| " + " | ".join(f"{str(row.get(col, '')):<{col_widths[col]}}" for col in columns) + " |"
|
||
rows.append(row_str)
|
||
|
||
result_str = f"查询结果（共 {len(limited_result)} 行）:\n\n"
|
||
result_str += header + "\n" + separator + "\n"
|
||
result_str += "\n".join(rows[:50]) # 最多显示50行
|
||
|
||
if len(rows) > 50:
|
||
result_str += f"\n\n... 还有 {len(rows) - 50} 行未显示 ..."
|
||
|
||
result_str += warning
|
||
|
||
logger.info(f"Query executed successfully, returned {len(limited_result)} rows")
|
||
return result_str
|
||
|
||
return "查询执行成功，但返回数据为空"
|
||
|
||
except Exception as e:
|
||
error_msg = f"SQL查询执行失败: {str(e)}\n\n{sql}"
|
||
|
||
# 提供更有用的错误信息
|
||
if "timeout" in str(e).lower():
|
||
error_msg += "\n\n💡 建议：查询超时了，请尝试以下方法：\n"
|
||
error_msg += "1. 减少查询的数据量（使用WHERE条件过滤）\n"
|
||
error_msg += "2. 使用LIMIT子句限制返回行数\n"
|
||
error_msg += "3. 增加timeout参数值（最大600秒）"
|
||
elif "table" in str(e).lower() and "doesn't exist" in str(e).lower():
|
||
error_msg += "\n\n💡 建议：表不存在，请使用 mysql_list_tables 查看可用的表名"
|
||
elif "column" in str(e).lower() and "doesn't exist" in str(e).lower():
|
||
error_msg += "\n\n💡 建议：列不存在，请使用 mysql_describe_table 查看表结构"
|
||
elif "not enough arguments for format string" in str(e).lower():
|
||
error_msg += (
|
||
"\n\n💡 建议：SQL 中的百分号 (%) 被当作参数占位符使用。"
|
||
" 如需匹配包含百分号的文本，请将百分号写成双百分号 (%%) 或使用参数化查询。"
|
||
)
|
||
|
||
logger.error(error_msg)
|
||
return error_msg
|
||
|
||
|
||
def _get_db_description() -> str:
|
||
"""获取数据库描述"""
|
||
import os
|
||
|
||
return os.getenv("MYSQL_DATABASE_DESCRIPTION") or ""
|
||
|
||
|
||
# 用于跟踪是否已注入描述，避免重复
|
||
_db_description_injected: bool = False
|
||
|
||
|
||
def _inject_db_description(tools: list[Any]) -> None:
|
||
"""将数据库描述注入到工具描述中"""
|
||
global _db_description_injected
|
||
if _db_description_injected:
|
||
return
|
||
|
||
db_desc = _get_db_description()
|
||
if not db_desc:
|
||
return
|
||
|
||
for _tool in tools:
|
||
if hasattr(_tool, "description"):
|
||
# 在描述末尾添加数据库说明
|
||
_tool.description = f"{_tool.description}\n\n当前数据库说明: {db_desc}"
|
||
|
||
_db_description_injected = True
|
||
|
||
|
||
def get_mysql_tools() -> list[Any]:
|
||
"""获取MySQL工具列表"""
|
||
tools = [mysql_list_tables, mysql_describe_table, mysql_query]
|
||
_inject_db_description(tools)
|
||
return tools
|