add query domain handler

2026-10-09 11:03:25 +08:00
parent 3cd836c929
commit 354d4e3a8d
9 changed files with 476 additions and 15 deletions
@@ -16,7 +16,7 @@ DB_POOL_PRE_PING=true
 # --- Query limits ---
 QUERY_MAX_ROWS=1000
 QUERY_MAX_BYTES=262144
-QUERY_MAX_CHARS=4096
+QUERY_MAX_CELL_CHARS=4096
 QUERY_TIMEOUT_SECONDS=30
 QUERY_ALLOW_WRITE=False

@@ -0,0 +1,250 @@
+# mini-db-mcp
+
+一个面向关系型数据库的轻量级 MCP（Model Context Protocol）服务。项目基于 Python MCP SDK 和 SQLAlchemy，当前支持 MySQL，并通过**渐进式披露**减少一次性注入模型上下文的工具定义数量。
+
+## 特性
+
+- 支持 `stdio` 与无状态 Streamable HTTP 两种传输方式
+- 采用“领域工具 -> 子工具”的两阶段发现与执行流程
+- 提供数据库对象浏览、表结构查看和字段采样能力
+- 查询默认只读，并提供行数、响应大小、单元格长度和超时限制
+- 同时返回 `structuredContent` 与兼容旧客户端的 JSON 文本内容
+- 支持工具列表和元数据结果分页
+
+## 工作原理
+
+项目不会在初始化时把所有子工具定义一次性暴露给模型，而是按需加载：
+
+1. MCP 客户端通过 `tools/list` 获取顶层领域工具，例如 `metadata`。
+2. 第一次调用领域工具时传入空参数，服务返回该领域的子工具清单及其输入、输出 Schema。
+3. 模型选择子工具后，再通过 `child_tool` 和 `child_arguments` 发起第二次调用。
+4. 服务校验参数并将请求路由到对应处理器，最终返回结构化结果。
+
+```mermaid
+sequenceDiagram
+    participant Client as MCP Client
+    participant Server as mini-db-mcp
+    participant DB as MySQL
+    Client->>Server: tools/list
+    Server-->>Client: metadata
+    Client->>Server: tools/call metadata, arguments={}
+    Server-->>Client: child tool manifest
+    Client->>Server: tools/call metadata, child_tool + child_arguments
+    Server->>DB: execute metadata query
+    DB-->>Server: result
+    Server-->>Client: structuredContent + JSON text
+```
+
+当前提供以下领域和子工具：
+
+| 领域 | 子工具 | 用途 |
+| --- | --- | --- |
+| `metadata` | `list_tables` | 分页列出当前数据库中的表和视图 |
+| `metadata` | `describe_table` | 分页查看指定表或视图的字段结构 |
+| `metadata` | `sample_values` | 获取指定字段的去重样例值 |
+| `query` | `execute_query` | 执行单条 SQL；根据 `QUERY_ALLOW_WRITE` 校验读写权限，并受查询超时和结果大小限制 |
+
+主要入口位于 [`src/mini_db_mcp/server.py`](src/mini_db_mcp/server.py)，渐进式工具发现与调用逻辑位于 [`src/mini_db_mcp/tools.py`](src/mini_db_mcp/tools.py)。
+
+## 环境要求
+
+- Python 3.12+
+- 可访问的 MySQL 数据库
+- 推荐使用 [uv](https://docs.astral.sh/uv/) 管理依赖
+
+## 快速开始
+
+### 1. 安装依赖
+
+```powershell
+uv sync
+```
+
+### 2. 创建本地配置
+
+启动前必须先将示例配置复制为 `.env`。
+
+Windows PowerShell：
+
+```powershell
+Copy-Item .env.example .env
+```
+
+macOS / Linux：
+
+```bash
+cp .env.example .env
+```
+
+随后编辑 `.env`，至少填写数据库连接信息：
+
+```dotenv
+DB_HOST=127.0.0.1
+DB_PORT=3306
+DB_USER=root
+DB_PASSWORD=your_password
+DB_SCHEMA=your_database
+```
+
+不要提交 `.env`。仓库仅跟踪不含真实凭据的 `.env.example`。
+
+### 3. 启动服务
+
+默认使用 `stdio`，适合由本地 MCP 客户端拉起：
+
+```powershell
+uv run mcp
+```
+
+也可以通过模块入口启动：
+
+```powershell
+uv run python -m mini_db_mcp
+```
+
+如需使用 Streamable HTTP，将 `.env` 中的传输方式改为：
+
+```dotenv
+MCP_TRANSPORT=http
+MCP_HOST=127.0.0.1
+MCP_PORT=8000
+```
+
+然后运行：
+
+```powershell
+uv run mcp
+```
+
+MCP 端点为 `http://127.0.0.1:8000/mcp`。
+
+## MCP 客户端配置
+
+使用 `stdio` 的客户端可参考以下配置。请将项目路径替换为本机的绝对路径：
+
+```json
+{
+  "mcpServers": {
+    "mini-db-mcp": {
+      "command": "uv",
+      "args": [
+        "--directory",
+        "C:\\path\\to\\mini-db-mcp",
+        "run",
+        "mcp"
+      ]
+    }
+  }
+}
+```
+
+服务从项目工作目录读取 `.env`，因此客户端配置必须将工作目录指向仓库根目录。
+
+## CallTool 示例
+
+第一次调用 `metadata` 时不传业务参数，用于获取当前可用的子工具定义：
+
+```json
+{
+  "name": "metadata",
+  "arguments": {}
+}
+```
+
+返回结果的 `mode` 为 `manifest`，其中 `children` 包含子工具名称、说明及 JSON Schema。选择子工具后进行第二次调用：
+
+```json
+{
+  "name": "metadata",
+  "arguments": {
+    "child_tool": "list_tables",
+    "child_arguments": {
+      "keyword": "user",
+      "limit": 20
+    }
+  }
+}
+```
+
+成功结果的 `mode` 为 `result`；参数错误、未知子工具或执行失败时，`mode` 为 `error`，同时 MCP 响应的 `isError` 会被设置为 `true`。
+
+`query` 领域同样采用两阶段调用。先发现 SQL 执行子工具：
+
+```json
+{
+  "name": "query",
+  "arguments": {}
+}
+```
+
+然后选择 `execute_query` 并提供 SQL：
+
+```json
+{
+  "name": "query",
+  "arguments": {
+    "child_tool": "execute_query",
+    "child_arguments": {
+      "sql": "SELECT id, name FROM users"
+    }
+  }
+}
+```
+
+## 配置
+
+配置由 `.env` 加载，完整默认值见 [`.env.example`](.env.example)。
+
+| 配置项 | 默认值 | 说明 |
+| --- | --- | --- |
+| `DB_HOST` | `127.0.0.1` | 数据库主机 |
+| `DB_PORT` | `3306` | 数据库端口 |
+| `DB_USER` | `root` | 数据库用户 |
+| `DB_PASSWORD` | 空 | 数据库密码 |
+| `DB_SCHEMA` | 空 | 默认数据库 Schema |
+| `QUERY_MAX_ROWS` | `1000` | 单次查询最大返回行数 |
+| `QUERY_MAX_BYTES` | `262144` | 单次响应最大字节数 |
+| `QUERY_MAX_CELL_CHARS` | `4096` | 单个字段值最大字符数 |
+| `QUERY_TIMEOUT_SECONDS` | `30` | 查询超时时间（秒） |
+| `QUERY_ALLOW_WRITE` | `false` | 是否允许写操作，建议保持关闭 |
+| `MCP_TRANSPORT` | `stdio` | 传输方式：`stdio` 或 `http` |
+| `MCP_HOST` | `127.0.0.1` | HTTP 监听地址 |
+| `MCP_PORT` | `8000` | HTTP 监听端口 |
+| `LOG_LEVEL` | `INFO` | 日志级别 |
+
+## 开发与测试
+
+安装开发依赖并运行测试：
+
+```powershell
+uv sync --dev
+uv run pytest
+```
+
+使用 MCP Inspector 调试 `stdio` 服务：
+
+```powershell
+npx @modelcontextprotocol/inspector uv run mcp
+```
+
+## 项目结构
+
+```text
+src/mini_db_mcp/
+|-- main.py                 # 命令行入口与传输方式选择
+|-- server.py               # MCP Server 注册与 CallTool 入口
+|-- tools.py                # 渐进式工具发现、校验与调用
+|-- config.py               # 环境变量配置
+|-- db.py                   # 异步数据库访问
+`-- domain/
+    |-- domain_catalog.py   # 领域与子工具 Schema
+    |-- domain_handler.py   # 处理器注册表
+    `-- handler/            # 子工具具体实现
+```
+
+## 安全建议
+
+- 使用权限最小化的专用数据库账号。
+- 保持 `QUERY_ALLOW_WRITE=false`，除非明确需要写入能力并已完成额外审计。
+- 不要在 MCP 客户端配置或版本库中直接保存数据库密码。
+- 对 HTTP 部署增加反向代理、身份认证、TLS 和网络访问控制；当前服务默认仅监听本机地址。
@@ -1,8 +1,11 @@
-from typing import Any
+import asyncio
+import base64
+import json
+from time import monotonic
+from typing import Any, Iterable, Sequence
 from datetime import datetime, date, time
 from decimal import Decimal
 from uuid import UUID
-import base64
 from sqlalchemy import URL, inspect, Table, MetaData, select
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from sqlalchemy.exc import SQLAlchemyError
@@ -12,7 +15,17 @@ from mini_db_mcp.utils.pagination_util import paginate_list


 class QueryExecutionError(RuntimeError):
-    pass
+    def __init__(
+        self,
+        code: str,
+        message: str | None = None,
+        *,
+        retryable: bool = False,
+    ) -> None:
+        self.code = code
+        self.message = message or code
+        self.retryable = retryable
+        super().__init__(self.message)

 class Database:
    def __init__(self, settings: Config):
@@ -218,6 +231,76 @@ class Database:
            for value in values[:limit]
        ], truncated

+    async def execute_guarded(self, sql: str) -> dict[str, Any]:
+        start = monotonic()
+        try:
+            async with asyncio.timeout(self.settings.query.timeout_seconds):
+                async with self.engine.connect() as connection:
+                    result = await connection.exec_driver_sql(sql)
+                    columns = list(result.keys())
+                    raw_rows: Iterable[Any] = result
+                    rows, truncated = bound_rows(
+                        columns,
+                        raw_rows,
+                        max_rows=self.settings.query.max_rows,
+                        max_bytes=self.settings.query.max_bytes,
+                        max_cell_chars=self.settings.query.max_cell_chars,
+                    )
+        except SQLAlchemyError as exc:
+            raise QueryExecutionError(
+                "query_failed",
+                "查询执行失败，请检查 SQL、权限及数据库状态。",
+            ) from exc
+        except TimeoutError as exc:
+            raise QueryExecutionError(
+                "query_timeout",
+                "查询超过执行时间上限，已中断。",
+                retryable=True,
+            ) from exc
+
+        return {
+            "columns": columns,
+            "rows": rows,
+            "row_count": len(rows),
+            "truncated": truncated,
+            "elapsed_ms": int((monotonic() - start) * 1000),
+        }
+
+
+def bound_rows(
+    columns: Sequence[str],
+    source_rows: Iterable[Sequence[Any]],
+    *,
+    max_rows: int,
+    max_bytes: int,
+    max_cell_chars: int,
+) -> tuple[list[list[Any]], bool]:
+    """Bound model-facing output by rows and serialized byte size.
+
+    This only limits what is returned to the model. It does not limit the amount of
+    data the database may scan; query timeouts and database workload governance are
+    separate controls.
+    """
+
+    rows: list[list[Any]] = []
+    truncated = False
+    base_size = len(json.dumps({"columns": list(columns), "rows": []}, ensure_ascii=False).encode("utf-8"))
+    current_size = base_size
+
+    for raw_row in source_rows:
+        if len(rows) >= max_rows:
+            truncated = True
+            break
+        row = [_serialize_scalar(value, max_chars=max_cell_chars) for value in raw_row]
+        row_size = len(json.dumps(row, ensure_ascii=False, default=str).encode("utf-8"))
+        if current_size + row_size > max_bytes:
+            truncated = True
+            break
+        rows.append(row)
+        current_size += row_size
+
+    return rows, truncated
+

 def _serialize_scalar(value: Any, *, max_chars: int) -> Any:
    if value is None or isinstance(value, (bool, int, float)):
@@ -117,6 +117,29 @@ _SAMPLE_VALUES_OUTPUT = {
    "additionalProperties": False,
 }

+_QUERY_OUTPUT = {
+    "type": "object",
+    "properties": {
+        "columns": {"type": "array", "items": {"type": "string"}},
+        "rows": {"type": "array", "items": {"type": "array", "items": {}}},
+        "row_count": {"type": "integer", "minimum": 0},
+        "truncated": {"type": "boolean"},
+        "elapsed_ms": {"type": "integer", "minimum": 0},
+        "executed_sql": {"type": "string"},
+        "referenced_objects": {"type": "array", "items": {"type": "string"}},
+    },
+    "required": [
+        "columns",
+        "rows",
+        "row_count",
+        "truncated",
+        "elapsed_ms",
+        "executed_sql",
+        "referenced_objects",
+    ],
+    "additionalProperties": False,
+}
+

 def _domain(
    *,
@@ -262,6 +285,30 @@ DOMAIN_DEFINITIONS = (
            ),
        ),
    ),
+
+    _domain(
+        name="query",
+        title="Query",
+        description="Validates and executes SQL according to the configured database query policy.",
+        children=(
+            _child(
+                domain="query",
+                name="execute_query",
+                title="Execute Query",
+                description="Executes one SQL statement subject to the configured read/write policy, timeout, and result limits.",
+                input_schema=_input_schema(
+                    {
+                        "sql": _string(
+                            "A single SQL statement to execute. Read-only mode is enforced by default.",
+                            max_length=200_000,
+                        ),
+                    },
+                    required=("sql",),
+                ),
+                output_schema=_QUERY_OUTPUT,
+            ),
+        ),
+    )
 )


@@ -1,15 +1,26 @@
-from mini_db_mcp.domain.domain_handler import DomainHandlerRegistry
+from mini_db_mcp.domain.domain_handler import DomainHandlerRegistry, HandlerProvider
 from mini_db_mcp.domain.handler.metadata_handler import MetadataHandler
 from mini_db_mcp.domain.runtime.metadata_runtime import MetadataRuntime, MetadataStore
+from mini_db_mcp.domain.handler.query_handler import QueryHandler
+from mini_db_mcp.domain.runtime.query_runtime import QueryRuntime, QueryStore


-def create_domain_registry(*, 
-        metadata_store: MetadataStore
+def create_domain_registry(
+    *,
+    metadata_store: MetadataStore,
+    query_store: QueryStore,
+    query_allow_write: bool = False,
+    query_max_rows: int = 1000,
 ) -> DomainHandlerRegistry:
    """Build the domain runtimes, handlers, and registry."""
    metadata_runtime = MetadataRuntime(metadata_store)
-    return DomainHandlerRegistry(
-        [
-            MetadataHandler(metadata_runtime)
-        ]
+    query_runtime = QueryRuntime(
+        query_store,
+        allow_write=query_allow_write,
+        max_rows=query_max_rows,
    )
+    handlers: list[HandlerProvider] = [
+        MetadataHandler(metadata_runtime), 
+        QueryHandler(query_runtime)
+    ]
+    return DomainHandlerRegistry(handlers)
@@ -0,0 +1,21 @@
+from types import MappingProxyType
+from typing import Mapping
+
+from mini_db_mcp.domain.domain_handler import ToolHandler
+from mini_db_mcp.domain.runtime.query_runtime import QueryRuntime
+
+
+class QueryHandler:
+    domain_name = "query"
+
+    def __init__(self, runtime: QueryRuntime) -> None:
+        self._handlers: Mapping[str, ToolHandler] = MappingProxyType(
+            {"execute_query": runtime.execute_query}
+        )
+
+    @property
+    def handler_names(self) -> frozenset[str]:
+        return frozenset(self._handlers)
+
+    def resolve_handler(self, name: str) -> ToolHandler | None:
+        return self._handlers.get(name)
@@ -0,0 +1,41 @@
+from typing import Any, Protocol
+
+from mini_db_mcp.security import SecurityError, SqlGuard, extract_tables
+
+
+class QueryStore(Protocol):
+    """Query execution capability required by the query runtime."""
+
+    async def execute_guarded(self, sql: str) -> dict[str, Any]: ...
+
+
+class QueryRuntime:
+    """Runtime for query domain."""
+
+    def __init__(
+        self,
+        store: QueryStore,
+        *,
+        allow_write: bool = False,
+        max_rows: int = 1000,
+    ) -> None:
+        self._store = store
+        self._sql_guard = SqlGuard(allow_write=allow_write, max_rows=max_rows)
+
+    async def execute_query(self, arguments: dict[str, Any]) -> dict[str, Any]:
+        sql = arguments.get("sql")
+        if not isinstance(sql, str) or not sql.strip():
+            raise ValueError("sql must be a non-empty string")
+
+        try:
+            guarded_sql = self._sql_guard.validate(sql)
+        except SecurityError as exc:
+            raise ValueError(str(exc)) from exc
+
+        result = await self._store.execute_guarded(guarded_sql)
+        return {
+            **result,
+            "executed_sql": guarded_sql,
+            "referenced_objects": sorted(extract_tables(guarded_sql)),
+        }
+
@@ -59,6 +59,9 @@ class MiniServer():
        self.database = Database(config)
        self.handler_registry = create_domain_registry(
            metadata_store=self.database,
+            query_store=self.database,
+            query_allow_write=config.query.allow_write,
+            query_max_rows=config.query.max_rows,
        )
        self.tools_manager = ToolsManager(
            catalog=DOMAIN_CATALOG,
@@ -96,12 +96,17 @@ class ToolsManager:
                str(exc),
            )
        except Exception as exc:
+            error_code = (
+                DomainErrorCode.CHILD_EXECUTION_TIMEOUT
+                if getattr(exc, "code", None) == "query_timeout"
+                else DomainErrorCode.CHILD_EXECUTION_FAILED
+            )
            return self._error_result(
                name,
                child_tool,
-                DomainErrorCode.CHILD_EXECUTION_FAILED,
+                error_code,
                str(exc),
-                retryable=True,
+                retryable=getattr(exc, "retryable", True),
            )

        return {