VAST DB MCP Server

VAST DB MCP 服务

本项目实现了一个模型上下文协议（MCP）服务，旨在充当 AI Agent/大模型与 VAST Data 数据库之间的接口。

项目目标

使用 MCP 标准，为 AI 模型提供一种安全且结构化的方式，以查询信息（架构、数据样本）并在 VAST DB 实例上执行只读查询。

核心技术

• Python： >=3.9
• MCP SDK： modelcontextprotocol/python-sdk（mcp-sdk）
• VAST DB SDK： vast-data/vastdb_sdk（vastdb）
• ASGI 服务端： uvicorn
• MCP 实现： 来自 mcp-sdk 的 FastMCP
• 配置： python-dotenv
• 测试： pytest、pytest-asyncio、pytest-mock、httpx
• 速率限制： slowapi

项目结构

/
├── sdk/                     # Cloned SDKs (ignored by git)
│   ├── python-sdk/
│   └── vastdb_sdk/
├── src/
│   └── vast_mcp_server/     # Main Python package
│       ├── __init__.py
│       ├── server.py          # FastMCP application setup
│       ├── config.py          # Loads connection details from .env
│       ├── utils.py           # Shared utility functions (e.g., auth header extraction)
│       ├── resources/         # MCP Resource handlers
│       │   ├── __init__.py
│       │   ├── schema.py      # Handler for vast://schemas
│       │   ├── table_data.py  # Handler for vast://tables/{table_name}
│       │   └── metadata.py    # Handler for vast://metadata/tables/{table_name}
│       ├── tools/             # MCP Tool handlers
│       │   ├── __init__.py
│       │   └── query.py       # Handler for vast_sql_query tool
│       └── vast_integration/  # VAST DB interaction logic
│           ├── __init__.py
│           └── db_ops.py      # Connection & query execution (async wrappers)
├── tests/                   # Pytest unit/integration tests
│   ├── __init__.py
│   ├── test_db_ops.py     # Tests for VAST DB interaction logic
│   ├── test_resources.py  # Tests for MCP resource handlers
│   └── test_tools.py      # Tests for MCP tool handlers
├── scripts/
│   └── run_server.py        # Script to start the server via uvicorn
├── .gitignore
├── .env.example             # Example environment file
├── README.md                # This file
└── pyproject.toml           # Project metadata and dependencies

已实现的 MCP 功能

身份验证： 所有资源和工具都要求在请求中发送以下 HTTP 标头：

• X-Vast-Access-Key：您的 VAST DB 访问密钥。
• X-Vast-Secret-Key：您的 VAST DB 秘密密钥。 如果未提供这些标头或凭据无效，将返回 UNAUTHENTICATED（401）错误响应。

速率限制： 所有资源和工具处理程序都根据客户端 IP 地址实施速率限制。默认限制可通过 MCP_DEFAULT_RATE_LIMIT 环境变量进行配置（例如，“10/minute”，参见 .env.example）。超过限制将返回 429 Too Many Requests 错误。

• 资源：数据库架构
  • URI： vast://schemas
  • 描述： 返回一个格式化字符串，用于描述所有已发现的表及其列（名称和类型）。需要身份验证标头。
  • 错误处理： 返回一个带有错误状态码（UNAUTHENTICATED、SERVICE_UNAVAILABLE、INTERNAL_SERVER_ERROR）的 McpResponse，以及纯文本错误消息正文（ERROR: [ErrorType] Message）。
• 资源：列出表
  • URI： vast://tables?format=FMT
  • 描述： 返回可用表名称列表。需要身份验证标头。
  • 参数：
    • format（字符串，可选，默认值：json）：输出格式（json 或 csv/list）。csv 或 list 会返回以换行符分隔的字符串。
  • 格式： JSON 字符串数组，或以换行符分隔的列表。
  • 错误处理： 返回一个带有错误状态码（UNAUTHENTICATED、BAD_REQUEST、SERVICE_UNAVAILABLE、INTERNAL_SERVER_ERROR）的 McpResponse，以及格式化的错误正文（根据 format 使用 JSON 或纯文本）。
• 资源：表元数据
  • URI： vast://metadata/tables/{table_name}
  • 描述： 返回特定表的详细元数据。需要身份验证标头。尝试包含列名、数据类型、可空性（is_nullable：例如，'YES'/'NO'）、键信息（key：例如，'PRI'）和默认值（default）。如果数据库内省无法获取名称和类型以外的字段，这些字段可能为 null。
  • 格式： JSON 对象，包含 table_name（字符串）和 columns（对象列表，每个对象可能包含 name、type、is_nullable、key、default 键）。
  • 示例响应（增强版）：

    {
      "table_name": "my_table",
      "columns": [
        {
          "name": "id",
          "type": "INTEGER",
          "is_nullable": "NO",
          "key": "PRI",
          "default": null
        },
        {
          "name": "data_column",
          "type": "VARCHAR",
          "is_nullable": "YES",
          "key": "",
          "default": "'Default Text'"
        },
        {
          "name": "timestamp",
          "type": "TIMESTAMP",
          "is_nullable": "YES",
          "key": null,
          "default": null
        }
      ]
    }
    

  • 错误处理： 返回一个带有错误状态码（UNAUTHENTICATED、NOT_FOUND、BAD_REQUEST、SERVICE_UNAVAILABLE、INTERNAL_SERVER_ERROR）的 McpResponse，以及 JSON 错误正文（{\"error\": ...}）。
• 资源：表样本数据
  • URI： vast://tables/{table_name}?limit=N&format=FMT
  • 描述： 返回指定 table_name 中的样本数据。需要身份验证标头。
  • 参数：
    • limit（整数，可选，默认值：10）：最大行数。
    • format（字符串，可选，默认值：csv）：输出格式（csv 或 json）。
  • 格式： CSV 或 JSON 字符串（对象数组）；CSV 包含标题行。
  • 错误处理： 返回一个带有错误状态码（UNAUTHENTICATED、BAD_REQUEST、SERVICE_UNAVAILABLE、INTERNAL_SERVER_ERROR）的 McpResponse，以及格式化的错误正文（根据 format 使用 JSON 或纯文本）。
• 工具：SQL 查询执行器
  • 名称： vast_sql_query
  • 参数：
    • sql（字符串，必填）：要执行的 SQL 查询。
    • format（字符串，可选，默认值：csv）：输出格式（csv 或 json）。
    • headers（字典，必填）：包含请求标头的字典，必须包含 X-Vast-Access-Key 和 X-Vast-Secret-Key。
  • 描述： 使用 headers 参数中的凭据，对 VAST DB 执行提供的 SQL 查询。
  • 格式： 返回 CSV 或 JSON 字符串（对象数组），或错误消息字符串（根据 format 使用 JSON 或纯文本）。
  • 安全性： 允许的语句类型由 MCP_ALLOWED_SQL_TYPES 环境变量控制（默认为 SELECT）。
  • 错误处理： 失败时返回格式化的错误字符串（JSON 或纯文本）。错误包括缺少或无效的标头（AuthenticationError）、不允许的查询类型（InvalidInputError）、连接问题（DatabaseConnectionError）和查询执行问题（QueryExecutionError）。

AI Agent 交互说明

将此 MCP 服务与 AI Agent 框架（例如 LangChain、LlamaIndex、自定义 Agent）集成时，请考虑以下事项：

1. Agent 提示词/配置：

   • Agent 的系统提示词或配置必须包含可用资源和工具的说明（类似上文“已实现的 MCP 功能”部分）。这样，大模型就能根据用户请求选择正确的操作。
   • 由于 MCP 缺少标准发现机制，请明确列出 URI（vast://schemas、vast://tables?...、vast://metadata/tables/{name}、vast://tables/{name}?....）和工具名称（vast_sql_query）及其功能。

2. 身份验证处理（安全关键）：

   • 绝不要在发送给大模型的提示词中包含 VAST DB 凭据（access_key、secret_key）。
   • Agent 的编排器（运行 Agent 逻辑的代码，而非大模型）负责管理凭据。
   • 在客户端安全地加载凭据（使用环境变量、Vault/AWS/Azure/GCP Secrets Manager 等密钥管理器）。
   • 对于资源： 当大模型生成目标 URI 时，编排器必须：
     • 获取已存储的凭据。
     • 构造 headers 字典：{'X-Vast-Access-Key': '...', 'X-Vast-Secret-Key': '...'}。
     • 将 vast:// URI 方案映射到正在运行的 MCP 服务的实际 HTTP URL（例如 http://localhost:8088/）。
     • 使用 MCP 客户端或标准 HTTP 客户端发起 HTTP GET 请求，并传入构造好的 headers。
   • 对于工具： 当大模型决定使用 vast_sql_query 工具并提供 sql 和 format 参数时，编排器必须：
     • 获取已存储的凭据。
     • 如上所述构造 headers 字典。
     • 将此 headers 字典注入传递给工具执行函数的参数中。大模型不应自行生成 headers 参数。

3. 请求构造：

   • 大模型需要为资源生成正确的 URI 路径参数（{table_name}）和查询参数（?format=、?limit=）。
   • 大模型为 vast_sql_query 工具生成 sql 和 format 参数。

4. 响应处理：

   • 客户端 Agent 代码需要处理不同的响应 Content-Type（例如 text/plain、application/json、text/csv）。
   • 必须检查响应状态码（尤其是资源的状态码）以检测错误（例如 401、404、500、503）。
   • 需要解析响应正文中的错误消息（纯文本或 JSON），并酌情将其反馈给用户，或用于重试/采取替代操作。

后续可能的步骤

• 实现完善的日志记录。（已完成）
• 添加单元测试。（已完成）
• 优化输出格式（例如，除 CSV 外还提供 JSON）。（已完成）
• 改进错误处理和报告。（已完成——实现了基本的自定义异常和格式化）
• 添加更细粒度的资源/工具（例如，仅列出表、获取表元数据）。（已完成——可列出表、获取表元数据）
• 为 vast_sql_query 工具实现更完善的查询验证/沙盒机制。（已完成——使用 sqlparse）
• 使查询限制（例如，是否允许非 SELECT 语句）可配置。（已完成——通过 MCP_ALLOWED_SQL_TYPES 环境变量配置）
• 添加需要运行中的 VAST DB 实例或模拟服务的集成测试。（已完成——为所有当前资源/工具添加了 ASGI/模拟测试）
• 根据需要考虑添加身份验证/授权层。（已完成——基于 Header 的身份验证）
• 将共享代码（例如 extract_auth_headers、格式化程序）重构到 utils.py 中。（已完成）
• 实现速率限制。（已完成——通过处理器上的 slowapi 实现基于 IP 的限制）

如何运行

1. 克隆代码仓库（如果尚未克隆）。
2. 设置环境： 将 .env.example 复制为 .env，并填写 VAST DB 端点、访问密钥和密钥。

   cp .env.example .env
   # Edit .env with your details
   

3. 安装依赖： 建议使用虚拟环境。