admin/aiagent

Fork 0

Files

rjb 9d3198f6bc 知你客服

2026-03-06 22:31:41 +08:00

4.5 KiB

Raw Blame History

数据库查询工具实现总结

✅ 完成状态

任务: 数据库查询工具实现
状态: ✅ 已完成
完成时间: 2026-01-23

📋 实现功能

1. 核心功能

✅ 默认数据库查询
- 使用SQLAlchemy连接默认数据库
- 支持所有SELECT查询操作
- 自动格式化查询结果为JSON
✅ 指定数据源查询
- 支持通过data_source_id参数查询外部数据库
- 自动使用数据源连接器（MySQL、PostgreSQL等）
- 支持多种数据库类型
✅ SQL注入防护
- 只允许SELECT查询
- 禁止INSERT、UPDATE、DELETE、DROP等危险操作
- 检查危险关键字（INSERT、UPDATE、DELETE、DROP、TRUNCATE、ALTER等）
- 禁止多语句查询
- 自动清理SQL注释
✅ 查询超时控制
- 默认超时时间：30秒
- 最大超时时间：300秒
- 使用asyncio实现异步超时控制
✅ 结果格式化
- 返回JSON格式结果
- 包含查询状态、行数、数据内容
- 自动处理日期时间等特殊类型

🔧 技术实现

文件修改

backend/app/services/builtin_tools.py
- 实现 database_query_tool 函数
- 实现 _validate_sql_query SQL验证函数
- 实现 _execute_default_db_query 默认数据库查询
- 实现 _execute_data_source_query 数据源查询
- 更新 DATABASE_QUERY_SCHEMA 工具定义
backend/app/main.py
- 注册 database_query_tool 到工具注册表

核心代码

async def database_query_tool(
    query: str, 
    database: str = "default",
    data_source_id: Optional[str] = None,
    timeout: int = 30
) -> str:
    """
    数据库查询工具
    
    - 只允许SELECT查询
    - 支持默认数据库和指定数据源
    - 包含SQL注入防护和超时控制
    """

🧪 测试结果

测试用例

✅ SQL验证测试
- 正常SELECT查询 ✅
- INSERT/UPDATE/DELETE/DROP查询被拒绝 ✅
- 多语句查询被拒绝 ✅
- 带WHERE和别名的复杂查询 ✅
✅ 数据库查询测试
- 默认数据库查询成功 ✅
- SQL注入防护生效 ✅
- 复杂SELECT查询成功 ✅
- 超时控制生效 ✅

测试脚本

文件: test_database_query_tool.py
所有测试用例通过 ✅

📊 工具Schema

{
  "type": "function",
  "function": {
    "name": "database_query",
    "description": "执行数据库查询（只允许SELECT查询，支持默认数据库和指定数据源）",
    "parameters": {
      "type": "object",
      "properties": {
        "query": {
          "type": "string",
          "description": "SQL查询语句（只允许SELECT查询）"
        },
        "data_source_id": {
          "type": "string",
          "description": "数据源ID（可选，如果提供则使用指定的数据源）"
        },
        "timeout": {
          "type": "integer",
          "description": "查询超时时间（秒，默认30秒，最大300秒）",
          "default": 30,
          "minimum": 1,
          "maximum": 300
        }
      },
      "required": ["query"]
    }
  }
}

🔒 安全特性

SQL注入防护
- 只允许SELECT查询
- 关键字黑名单检查
- 禁止多语句执行
超时控制
- 防止长时间运行的查询
- 可配置超时时间
错误处理
- 详细的错误信息
- 不泄露敏感信息

📝 使用示例

示例1: 查询默认数据库

result = await database_query_tool(
    query="SELECT id, username, email FROM users LIMIT 10",
    timeout=30
)

示例2: 查询指定数据源

result = await database_query_tool(
    query="SELECT * FROM products WHERE price > 100",
    data_source_id="your-data-source-id",
    timeout=60
)

示例3: 在LLM节点中使用

在LLM节点的工具配置中启用 database_query 工具，LLM可以自动调用：

用户: "查询一下有多少个用户"
LLM: 调用 database_query(query="SELECT COUNT(*) as count FROM users")

🎯 下一步

根据待完善功能清单，接下来可以：

工具调用可视化 - 显示工具调用过程
监控和告警前端界面 - 后端API已完成，需要前端实现
工具动态注册机制 - 支持HTTP工具、工作流工具的动态注册

📊 统计

代码行数: 约200行
测试用例: 8个
测试通过率: 100%
工具总数: 9个（新增1个）

最后更新: 2026-01-23
文档版本: v1.0

4.5 KiB Raw Blame History Unescape Escape