From 354d4e3a8d3d87237219e0c8ff4e83bd95f6fa01 Mon Sep 17 00:00:00 2001 From: "Chen, Li" Date: Mon, 28 Sep 2026 17:41:59 +0800 Subject: [PATCH] add query domain handler --- .env.example | 2 +- README.md | 250 ++++++++++++++++++ src/mini_db_mcp/db.py | 91 ++++++- src/mini_db_mcp/domain/domain_catalog.py | 47 ++++ .../domain/domain_handler_factory.py | 27 +- .../domain/handler/query_handler.py | 21 ++ .../domain/runtime/query_runtime.py | 41 +++ src/mini_db_mcp/server.py | 3 + src/mini_db_mcp/tools.py | 9 +- 9 files changed, 476 insertions(+), 15 deletions(-) create mode 100644 README.md create mode 100644 src/mini_db_mcp/domain/runtime/query_runtime.py diff --git a/.env.example b/.env.example index 1177d28..73725f5 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..bcb2255 --- /dev/null +++ b/README.md @@ -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 和网络访问控制；当前服务默认仅监听本机地址。 \ No newline at end of file diff --git a/src/mini_db_mcp/db.py b/src/mini_db_mcp/db.py index 900b68e..45c65c6 100644 --- a/src/mini_db_mcp/db.py +++ b/src/mini_db_mcp/db.py @@ -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)): @@ -232,4 +315,4 @@ def _serialize_scalar(value: Any, *, max_chars: int) -> Any: text = str(value) if len(text) > max_chars: return text[:max_chars] + "…[truncated]" - return text \ No newline at end of file + return text diff --git a/src/mini_db_mcp/domain/domain_catalog.py b/src/mini_db_mcp/domain/domain_catalog.py index 9f126d5..f933ca7 100644 --- a/src/mini_db_mcp/domain/domain_catalog.py +++ b/src/mini_db_mcp/domain/domain_catalog.py @@ -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, + ), + ), + ) ) diff --git a/src/mini_db_mcp/domain/domain_handler_factory.py b/src/mini_db_mcp/domain/domain_handler_factory.py index 40f625b..42f9a6a 100644 --- a/src/mini_db_mcp/domain/domain_handler_factory.py +++ b/src/mini_db_mcp/domain/domain_handler_factory.py @@ -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 - ) -> DomainHandlerRegistry: +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) diff --git a/src/mini_db_mcp/domain/handler/query_handler.py b/src/mini_db_mcp/domain/handler/query_handler.py index e69de29..9e1b8cc 100644 --- a/src/mini_db_mcp/domain/handler/query_handler.py +++ b/src/mini_db_mcp/domain/handler/query_handler.py @@ -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) diff --git a/src/mini_db_mcp/domain/runtime/query_runtime.py b/src/mini_db_mcp/domain/runtime/query_runtime.py new file mode 100644 index 0000000..eed27aa --- /dev/null +++ b/src/mini_db_mcp/domain/runtime/query_runtime.py @@ -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)), + } + diff --git a/src/mini_db_mcp/server.py b/src/mini_db_mcp/server.py index 8a8b84d..17d50db 100644 --- a/src/mini_db_mcp/server.py +++ b/src/mini_db_mcp/server.py @@ -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, diff --git a/src/mini_db_mcp/tools.py b/src/mini_db_mcp/tools.py index 3bbe727..0b77d7a 100644 --- a/src/mini_db_mcp/tools.py +++ b/src/mini_db_mcp/tools.py @@ -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 {