add query domain handler
This commit is contained in:
+1
-1
@@ -16,7 +16,7 @@ DB_POOL_PRE_PING=true
|
|||||||
# --- Query limits ---
|
# --- Query limits ---
|
||||||
QUERY_MAX_ROWS=1000
|
QUERY_MAX_ROWS=1000
|
||||||
QUERY_MAX_BYTES=262144
|
QUERY_MAX_BYTES=262144
|
||||||
QUERY_MAX_CHARS=4096
|
QUERY_MAX_CELL_CHARS=4096
|
||||||
QUERY_TIMEOUT_SECONDS=30
|
QUERY_TIMEOUT_SECONDS=30
|
||||||
QUERY_ALLOW_WRITE=False
|
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 和网络访问控制;当前服务默认仅监听本机地址。
|
||||||
+86
-3
@@ -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 datetime import datetime, date, time
|
||||||
from decimal import Decimal
|
from decimal import Decimal
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
import base64
|
|
||||||
from sqlalchemy import URL, inspect, Table, MetaData, select
|
from sqlalchemy import URL, inspect, Table, MetaData, select
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||||
from sqlalchemy.exc import SQLAlchemyError
|
from sqlalchemy.exc import SQLAlchemyError
|
||||||
@@ -12,7 +15,17 @@ from mini_db_mcp.utils.pagination_util import paginate_list
|
|||||||
|
|
||||||
|
|
||||||
class QueryExecutionError(RuntimeError):
|
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:
|
class Database:
|
||||||
def __init__(self, settings: Config):
|
def __init__(self, settings: Config):
|
||||||
@@ -218,6 +231,76 @@ class Database:
|
|||||||
for value in values[:limit]
|
for value in values[:limit]
|
||||||
], truncated
|
], 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:
|
def _serialize_scalar(value: Any, *, max_chars: int) -> Any:
|
||||||
if value is None or isinstance(value, (bool, int, float)):
|
if value is None or isinstance(value, (bool, int, float)):
|
||||||
|
|||||||
@@ -117,6 +117,29 @@ _SAMPLE_VALUES_OUTPUT = {
|
|||||||
"additionalProperties": False,
|
"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(
|
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.handler.metadata_handler import MetadataHandler
|
||||||
from mini_db_mcp.domain.runtime.metadata_runtime import MetadataRuntime, MetadataStore
|
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(*,
|
def create_domain_registry(
|
||||||
metadata_store: MetadataStore
|
*,
|
||||||
|
metadata_store: MetadataStore,
|
||||||
|
query_store: QueryStore,
|
||||||
|
query_allow_write: bool = False,
|
||||||
|
query_max_rows: int = 1000,
|
||||||
) -> DomainHandlerRegistry:
|
) -> DomainHandlerRegistry:
|
||||||
"""Build the domain runtimes, handlers, and registry."""
|
"""Build the domain runtimes, handlers, and registry."""
|
||||||
metadata_runtime = MetadataRuntime(metadata_store)
|
metadata_runtime = MetadataRuntime(metadata_store)
|
||||||
return DomainHandlerRegistry(
|
query_runtime = QueryRuntime(
|
||||||
[
|
query_store,
|
||||||
MetadataHandler(metadata_runtime)
|
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.database = Database(config)
|
||||||
self.handler_registry = create_domain_registry(
|
self.handler_registry = create_domain_registry(
|
||||||
metadata_store=self.database,
|
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(
|
self.tools_manager = ToolsManager(
|
||||||
catalog=DOMAIN_CATALOG,
|
catalog=DOMAIN_CATALOG,
|
||||||
|
|||||||
@@ -96,12 +96,17 @@ class ToolsManager:
|
|||||||
str(exc),
|
str(exc),
|
||||||
)
|
)
|
||||||
except Exception as 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(
|
return self._error_result(
|
||||||
name,
|
name,
|
||||||
child_tool,
|
child_tool,
|
||||||
DomainErrorCode.CHILD_EXECUTION_FAILED,
|
error_code,
|
||||||
str(exc),
|
str(exc),
|
||||||
retryable=True,
|
retryable=getattr(exc, "retryable", True),
|
||||||
)
|
)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
|
|||||||
Reference in New Issue
Block a user