diff --git a/.agents/skills/xsql/SKILL.md b/.agents/skills/xsql/SKILL.md index be664d4..a1051d9 100644 --- a/.agents/skills/xsql/SKILL.md +++ b/.agents/skills/xsql/SKILL.md @@ -154,6 +154,7 @@ Table and CSV formats are for humans — they lack the `ok`/`schema_version` env ```bash # MCP server for AI assistant integration (read-only by default) +# Tools: query, schema_list_tables, schema_dump, schema_describe_table, profile_list, profile_show xsql mcp server --attr source=codex-cli --attr agent=codex # stdio transport (default) # Writes (only when user explicitly requests): add --unsafe-allow-write and ensure profile unsafe_allow_write: true xsql mcp server --unsafe-allow-write --attr source=codex-cli --attr agent=codex diff --git a/README.md b/README.md index 4c56671..a2c848c 100644 --- a/README.md +++ b/README.md @@ -148,7 +148,7 @@ Add the xsql MCP server to your Claude Desktop configuration: } ``` -Once started, Claude can query databases directly via the MCP protocol (read-only by default). To allow writes over MCP, set profile `unsafe_allow_write: true` **and** add `--unsafe-allow-write` to the MCP server `args`. +Once started, Claude can query databases directly via the MCP protocol (read-only by default). MCP also exposes schema discovery tools (`schema_list_tables`, `schema_dump`, `schema_describe_table`) with the same JSON contract as the CLI. To allow writes over MCP, set profile `unsafe_allow_write: true` **and** add `--unsafe-allow-write` to the MCP server `args`. ### Option 5: AGENTS.md / Rules (Cursor/Windsurf) diff --git a/README_zh.md b/README_zh.md index ef66f1b..fc6f111 100644 --- a/README_zh.md +++ b/README_zh.md @@ -148,7 +148,7 @@ xsql query "" -p -f json --attr source=codex-cli --attr agent=cod } ``` -启动后，Claude 可以直接通过 MCP 协议查询数据库（默认只读）。若需经 MCP 写入，除 profile 设置 `unsafe_allow_write: true` 外，还须在 `args` 中加入 `--unsafe-allow-write` 启动 MCP server。 +启动后，Claude 可以直接通过 MCP 协议查询数据库（默认只读）。MCP 还提供 `schema_list_tables`、`schema_dump`、`schema_describe_table` 等 schema discovery tools（与 CLI JSON 契约一致）。若需经 MCP 写入，除 profile 设置 `unsafe_allow_write: true` 外，还须在 `args` 中加入 `--unsafe-allow-write` 启动 MCP server。 ### 方式五：AGENTS.md / Rules（Cursor/Windsurf） diff --git a/docs/ai.md b/docs/ai.md index 9098c76..0a475c8 100644 --- a/docs/ai.md +++ b/docs/ai.md @@ -75,12 +75,15 @@ xsql mcp server --transport streamable_http --http-addr 127.0.0.1:8787 --http-au ``` ### MCP Tools -MCP Server 提供以下 tools： +MCP Server 提供以下 tools（均经 `internal/app`，SSH/连接行为与 CLI/Web 一致）： - **query**: 执行 SQL 查询（默认只读；写入需 profile `unsafe_allow_write: true` 且 server 启动时 `--unsafe-allow-write`） +- **schema_list_tables**: 轻量表清单（可选 `table_pattern`、`include_system`），适合大库先做 discovery +- **schema_dump**: 与 `xsql schema dump` 等价的结构化导出（建议配合 `table_pattern` 控制体积） +- **schema_describe_table**: 单表结构（列、索引、外键）；PostgreSQL 需 `schema`，MySQL 可省略 - **profile_list**: 列出所有配置的 profiles - **profile_show**: 查看 profile 详情 -MCP 写入与 CLI 一样采用双重授权：仅在 profile 与 MCP server 启动参数同时开启时，`query` tool 才会绕过只读保护。 +MCP 写入与 CLI 一样采用双重授权：仅在 profile 与 MCP server 启动参数同时开启时，`query` tool 才会绕过只读保护。Schema tools 始终只读。 ### 集成示例 在 Claude Desktop 配置中添加： diff --git a/docs/architecture.md b/docs/architecture.md index c674dae..0f99388 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -33,4 +33,5 @@ - DB driver：通过 registry 注册，实现新数据库无需改动核心逻辑。 - 输出 formatter：通过接口注册新格式。 - Frontend：CLI/MCP/TUI/Web 仅是调用 `internal/app` 的不同适配层。 +- MCP：`internal/mcp` 注册 MCP tools，数据库与 schema 操作经 `internal/app`（与 Web 相同），不重复实现连接/SSH。 - Web：本地 HTTP server 与嵌入式前端通过 `internal/web` 适配，复用 `internal/app` 业务服务。 diff --git a/docs/cli-spec.md b/docs/cli-spec.md index 5475448..f780787 100644 --- a/docs/cli-spec.md +++ b/docs/cli-spec.md @@ -575,7 +575,55 @@ xsql mcp server --transport streamable_http --http-addr 127.0.0.1:8787 --http-au } ``` -**输出格式：** MCP tool 调用返回 JSON 格式，遵循 xsql 的标准输出契约（`ok`、`schema_version`、`data`/`error`）。 +4. **schema_list_tables**: 轻量表清单（schema discovery 首选） + ```json + { + "name": "schema_list_tables", + "inputSchema": { + "type": "object", + "required": ["profile"], + "properties": { + "profile": {"type": "string"}, + "table_pattern": {"type": "string", "description": "Optional glob filter (*, ?)"}, + "include_system": {"type": "boolean", "description": "Include system tables"} + } + } + } + ``` + +5. **schema_dump**: 导出表结构（与 `xsql schema dump` JSON `data` 同形） + ```json + { + "name": "schema_dump", + "inputSchema": { + "type": "object", + "required": ["profile"], + "properties": { + "profile": {"type": "string"}, + "table_pattern": {"type": "string"}, + "include_system": {"type": "boolean"} + } + } + } + ``` + +6. **schema_describe_table**: 单表结构详情 + ```json + { + "name": "schema_describe_table", + "inputSchema": { + "type": "object", + "required": ["profile", "table"], + "properties": { + "profile": {"type": "string"}, + "schema": {"type": "string", "description": "PostgreSQL schema (optional on MySQL)"}, + "table": {"type": "string"} + } + } + } + ``` + +**输出格式：** MCP tool 调用返回 JSON 格式，遵循 xsql 的标准输出契约（`ok`、`schema_version`、`data`/`error`）。MCP 数据库路径复用 `internal/app`（SSH `skip_host_key` 等与 CLI 一致）。 **安全说明：** - query tool 默认只读模式（双重保护：SQL 静态分析 + DB 事务级只读） diff --git a/docs/rfcs/0003-mcp-server.md b/docs/rfcs/0003-mcp-server.md index cd3fb8d..6437d70 100644 --- a/docs/rfcs/0003-mcp-server.md +++ b/docs/rfcs/0003-mcp-server.md @@ -24,8 +24,9 @@ Status: Implemented auth_token: "keyring:mcp/http_token" allow_plaintext_token: false ``` -- MCP Server 提供的 tools： +- MCP Server 提供的 tools（详见 RFC 0015 增量）： - `query`: 执行 SQL 查询（支持只读模式） + - `schema_list_tables` / `schema_dump` / `schema_describe_table`: schema discovery - `profile_list`: 列出所有 profiles - `profile_show`: 查看 profile 详情 @@ -42,7 +43,7 @@ Status: Implemented - Streamable HTTP: 使用 MCP Streamable HTTP 传输，必须启用鉴权 - 兼容性策略： - 新增能力，不影响 CLI 现有功能 - - MCP server 使用独立的代码路径，复用 internal 层的核心能力 + - MCP server 经 `internal/app` 复用连接/SSH/schema 能力（RFC 0015） ### MCP Tools 设计 1. **query_tool**: diff --git a/docs/rfcs/0015-mcp-app-layer-schema-tools.md b/docs/rfcs/0015-mcp-app-layer-schema-tools.md new file mode 100644 index 0000000..cab7673 --- /dev/null +++ b/docs/rfcs/0015-mcp-app-layer-schema-tools.md @@ -0,0 +1,41 @@ +# RFC 0015: MCP 复用 app 层与 Schema Discovery Tools + +Status: Implemented + +## 摘要 +MCP Server 的 `query` 工具此前在 `internal/mcp` 内重复实现连接/SSH 逻辑，未复用 `internal/app`，导致 profile 的 `skip_host_key` 等行为与 CLI/Web 不一致。本 RFC 将 MCP 数据库操作统一经 `internal/app`（与 Web UI 相同），并新增轻量 schema discovery tools，使 AI 客户端无需猜测表名。 + +## 背景 / 动机 +- **痛点**：MCP 约 150 行重复连接代码；SSH `skip_host_key` 在 MCP 路径被忽略；无 schema 工具，agent 只能手写 information_schema 查询。 +- **目标**：MCP/Web/CLI 共享 `app.ResolveProfile` / `app.Query` / `app.ListTables` 等；暴露与 CLI JSON 契约一致的 schema tools；保留 RFC 0014 写入双重授权与 stats 记录。 +- **非目标**：MCP query 默认行数限制（后续 PR）；MCP 侧 `--ssh-skip-known-hosts-check` CLI 开关（仍仅 profile `skip_host_key` + 与 CLI 相同的 app 解析逻辑）。 + +## 方案（Proposed） +### MCP Tools（新增） +| Tool | 用途 | 输入 | +|------|------|------| +| `schema_list_tables` | 轻量表清单（LLM 友好） | `profile`, 可选 `table_pattern`, `include_system` | +| `schema_dump` | 与 `xsql schema dump` 等价的结构化导出 | `profile`, 可选 `table_pattern`, `include_system` | +| `schema_describe_table` | 单表列/索引/外键详情 | `profile`, `table`, 可选 `schema`（PG 必填） | + +输出均遵循标准信封：`{"ok":true,"schema_version":1,"data":{...}}`。 + +### 技术设计 +- `internal/mcp`：`query` / schema tools 调用 `internal/app`；`WriteAllowed` 不变；stats `cmd` 字段区分 `query`、`schema_list_tables`、`schema_dump`、`schema_describe_table`。 +- 导出 `app.SSHOptionsForProfile`（原 `resolveSSHOptions` 的公开包装）供一致性单测。 +- `query` 路径保留 `app.QueryRequest` 结构，便于后续扩展默认 `LIMIT` 等选项。 + +## 兼容性与迁移 +- **不破坏**：现有 tool 名称与 query/profile 工具 JSON 契约不变；新增 tools 为纯增量。 +- query 错误文案可能与 CLI 完全一致（如 `profile not found`），不再使用 MCP 专有 wording。 + +## 安全与隐私 +- schema tools 为只读；不暴露 secrets。 +- 写入仍仅 `query` + 双重授权（RFC 0014）。 + +## 测试计划 +- 单元：`SSHOptionsForProfile` profile `skip_host_key`；MCP handler 参数校验。 +- E2E：MCP schema tools 对 MySQL 8 / PostgreSQL 16；SSH skip_host_key 与 CLI 一致（dial 阶段错误码）。 + +## 未决问题 +- 无 diff --git a/docs/rfcs/README.md b/docs/rfcs/README.md index 3104f0c..f3700ca 100644 --- a/docs/rfcs/README.md +++ b/docs/rfcs/README.md @@ -20,6 +20,7 @@ | [0012](0012-goja-js-data-analysis.md) | goja Data Analysis & ReAct Loop | Implemented | | [0013](0013-ai-multi-tool-calls-sequential-execution.md) | AI Multi Tool Call Sequential Execution & Dialect Guidance | Implemented | | [0014](0014-mcp-write-dual-authorization.md) | MCP 写入双重授权 | Implemented | +| [0015](0015-mcp-app-layer-schema-tools.md) | MCP 复用 app 层与 Schema Tools | Implemented | ## 何时需要 RFC 当变更会影响以下任一项时，必须先提交/更新 RFC（并同步更新对应 `docs/*`）： diff --git a/internal/app/conn.go b/internal/app/conn.go index c76b480..7729032 100644 --- a/internal/app/conn.go +++ b/internal/app/conn.go @@ -177,6 +177,13 @@ func ResolveReconnectableSSH(ctx context.Context, profile config.Profile, allowP return rd, nil } +// SSHOptionsForProfile builds SSH options from a resolved profile, resolving secrets. +// skipHostKeyCheck is the runtime override (e.g. CLI --ssh-skip-known-hosts-check); +// profile ssh_proxy.skip_host_key is also honored. +func SSHOptionsForProfile(profile config.Profile, allowPlaintext, skipHostKeyCheck bool) (ssh.Options, *errors.XError) { + return resolveSSHOptions(profile, allowPlaintext, skipHostKeyCheck) +} + // resolveSSHOptions builds SSH options from a profile, resolving secrets. func resolveSSHOptions(profile config.Profile, allowPlaintext, skipHostKeyCheck bool) (ssh.Options, *errors.XError) { passphrase := profile.SSHConfig.Passphrase diff --git a/internal/app/conn_test.go b/internal/app/conn_test.go index f021cc2..07bf685 100644 --- a/internal/app/conn_test.go +++ b/internal/app/conn_test.go @@ -412,6 +412,23 @@ func TestResolveReconnectableSSH_NilCallback(t *testing.T) { _ = xe } +func TestSSHOptionsForProfile_ProfileSkipHostKey(t *testing.T) { + profile := config.Profile{ + SSHConfig: &config.SSHProxy{ + Host: "bastion.example.com", + SkipHostKey: true, + }, + } + + opts, xe := SSHOptionsForProfile(profile, false, false) + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + if !opts.SkipKnownHostsCheck { + t.Error("expected SkipKnownHostsCheck from profile skip_host_key") + } +} + func TestResolveSSHOptions_Success(t *testing.T) { profile := config.Profile{ SSHConfig: &config.SSHProxy{ @@ -422,7 +439,7 @@ func TestResolveSSHOptions_Success(t *testing.T) { }, } - opts, xe := resolveSSHOptions(profile, true, true) + opts, xe := SSHOptionsForProfile(profile, true, true) if xe != nil { t.Fatalf("unexpected error: %v", xe) } diff --git a/internal/mcp/app_consistency_test.go b/internal/mcp/app_consistency_test.go new file mode 100644 index 0000000..e93e20d --- /dev/null +++ b/internal/mcp/app_consistency_test.go @@ -0,0 +1,56 @@ +package mcp + +import ( + "testing" + + "github.com/zx06/xsql/internal/app" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/stats" +) + +// MCP database tools use internal/app, which applies the same SSH option rules as CLI/Web. +func TestMCPQueryPathMatchesCLISSHOptions(t *testing.T) { + cfg := config.File{ + Profiles: map[string]config.Profile{ + "remote": { + DB: "mysql", + SSHProxy: "bastion", + }, + }, + SSHProxies: map[string]config.SSHProxy{ + "bastion": { + Host: "jump.example.com", + Port: 22, + User: "jump", + SkipHostKey: true, + }, + }, + } + + handler := NewToolHandler(&cfg, stats.StatsConfig{}, ServerOptions{}) + profile, xe := handler.resolveProfile("remote") + if xe != nil { + t.Fatalf("resolveProfile: %v", xe) + } + + // MCP query uses SkipHostKeyCheck=false at runtime; profile skip_host_key still applies. + mcpSSH, xe := app.SSHOptionsForProfile(profile, false, false) + if xe != nil { + t.Fatalf("SSHOptionsForProfile: %v", xe) + } + + cliSSH, xe := app.SSHOptionsForProfile(profile, false, false) + if xe != nil { + t.Fatalf("SSHOptionsForProfile cli: %v", xe) + } + + if mcpSSH.SkipKnownHostsCheck != cliSSH.SkipKnownHostsCheck { + t.Fatalf("skip host key mismatch: mcp=%v cli=%v", mcpSSH.SkipKnownHostsCheck, cliSSH.SkipKnownHostsCheck) + } + if mcpSSH.Host != cliSSH.Host || mcpSSH.User != cliSSH.User { + t.Fatalf("ssh target mismatch: mcp=%+v cli=%+v", mcpSSH, cliSSH) + } + if !mcpSSH.SkipKnownHostsCheck { + t.Error("expected skip_host_key from profile to enable SkipKnownHostsCheck") + } +} diff --git a/internal/mcp/response.go b/internal/mcp/response.go new file mode 100644 index 0000000..0a7ddf9 --- /dev/null +++ b/internal/mcp/response.go @@ -0,0 +1,41 @@ +package mcp + +import ( + "encoding/json" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/zx06/xsql/internal/errors" +) + +func (h *ToolHandler) successResult(data any) (*mcp.CallToolResult, any, error) { + jsonData, err := json.MarshalIndent(map[string]any{ + "ok": true, + "schema_version": 1, + "data": data, + }, "", " ") + if err != nil { + return h.errorResult(errors.Wrap(errors.CodeInternal, "failed to marshal result", nil, err)), nil, nil + } + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: string(jsonData)}, + }, + }, nil, nil +} + +func (h *ToolHandler) errorResult(xe *errors.XError) *mcp.CallToolResult { + return &mcp.CallToolResult{ + IsError: true, + Content: []mcp.Content{ + &mcp.TextContent{Text: h.formatError(xe)}, + }, + } +} + +func (h *ToolHandler) errorFromErr(err error) *mcp.CallToolResult { + if err == nil { + return h.errorResult(errors.New(errors.CodeInternal, "unknown error", nil)) + } + return h.errorResult(errors.AsOrWrap(err)) +} diff --git a/internal/mcp/response_test.go b/internal/mcp/response_test.go new file mode 100644 index 0000000..de3ce77 --- /dev/null +++ b/internal/mcp/response_test.go @@ -0,0 +1,61 @@ +package mcp + +import ( + "strings" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/errors" + "github.com/zx06/xsql/internal/stats" +) + +func TestSuccessResult_Envelope(t *testing.T) { + handler := NewToolHandler(&config.File{Profiles: map[string]config.Profile{}}, stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.successResult(map[string]any{"n": 1}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.IsError { + t.Fatal("expected success") + } + text := result.Content[0].(*mcp.TextContent).Text + if !strings.Contains(text, `"ok": true`) || !strings.Contains(text, `"schema_version": 1`) { + t.Fatalf("unexpected payload: %s", text) + } +} + +func TestSuccessResult_MarshalFailure(t *testing.T) { + handler := NewToolHandler(&config.File{Profiles: map[string]config.Profile{}}, stats.StatsConfig{}, ServerOptions{}) + ch := make(chan int) + result, _, err := handler.successResult(ch) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil || !result.IsError { + t.Fatal("expected marshal failure to become error result") + } + if !strings.Contains(result.Content[0].(*mcp.TextContent).Text, "INTERNAL") { + t.Fatalf("expected INTERNAL, got %s", result.Content[0].(*mcp.TextContent).Text) + } +} + +func TestErrorFromErr_Nil(t *testing.T) { + handler := NewToolHandler(&config.File{Profiles: map[string]config.Profile{}}, stats.StatsConfig{}, ServerOptions{}) + result := handler.errorFromErr(nil) + if result == nil || !result.IsError { + t.Fatal("expected error result") + } +} + +func TestErrorResult_WrapsXError(t *testing.T) { + handler := NewToolHandler(&config.File{Profiles: map[string]config.Profile{}}, stats.StatsConfig{}, ServerOptions{}) + result := handler.errorResult(errors.New(errors.CodeCfgInvalid, "bad config", nil)) + if result == nil || !result.IsError { + t.Fatal("expected error result") + } + if !strings.Contains(result.Content[0].(*mcp.TextContent).Text, "CFG_INVALID") { + t.Fatalf("unexpected: %s", result.Content[0].(*mcp.TextContent).Text) + } +} diff --git a/internal/mcp/schema_tools.go b/internal/mcp/schema_tools.go new file mode 100644 index 0000000..4344209 --- /dev/null +++ b/internal/mcp/schema_tools.go @@ -0,0 +1,230 @@ +package mcp + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/zx06/xsql/internal/app" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" +) + +const ( + defaultMCPSchemaTimeout = 60 * time.Second + msgInvalidInput = "invalid input" + msgProfileRequired = "profile is required" +) + +// schemaApp delegates schema operations to internal/app (overridable in tests). +var schemaApp = struct { + ListTables func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) + DumpSchema func(context.Context, app.SchemaDumpRequest) (*db.SchemaInfo, *errors.XError) + DescribeTable func(context.Context, app.TableDescribeRequest) (*db.Table, *errors.XError) +}{ + ListTables: app.ListTables, + DumpSchema: app.DumpSchema, + DescribeTable: app.DescribeTable, +} + +// SchemaListTablesInput is input for schema_list_tables. +type SchemaListTablesInput struct { + Profile string `json:"profile" jsonschema:"Profile name to use"` + TablePattern string `json:"table_pattern,omitempty" jsonschema:"Table name filter (* and ? wildcards)"` + IncludeSystem bool `json:"include_system,omitempty" jsonschema:"Include system tables"` +} + +// SchemaDumpInput is input for schema_dump. +type SchemaDumpInput struct { + Profile string `json:"profile" jsonschema:"Profile name to use"` + TablePattern string `json:"table_pattern,omitempty" jsonschema:"Table name filter (* and ? wildcards)"` + IncludeSystem bool `json:"include_system,omitempty" jsonschema:"Include system tables"` +} + +// SchemaDescribeTableInput is input for schema_describe_table. +type SchemaDescribeTableInput struct { + Profile string `json:"profile" jsonschema:"Profile name to use"` + Schema string `json:"schema,omitempty" jsonschema:"Schema name (PostgreSQL; optional for MySQL)"` + Table string `json:"table" jsonschema:"Table name"` +} + +func (h *ToolHandler) registerSchemaTools(server *mcp.Server, profileEnums []any) { + profileProp := &jsonschema.Schema{ + Type: "string", + Description: "Profile name to use", + Enum: profileEnums, + } + optionalTablePattern := &jsonschema.Schema{ + Type: "string", + Description: "Table name filter (supports * and ? wildcards)", + } + optionalIncludeSystem := &jsonschema.Schema{ + Type: "boolean", + Description: "Include system tables", + } + + server.AddTool(&mcp.Tool{ + Name: "schema_list_tables", + Description: "List database tables (lightweight summary for schema discovery)", + InputSchema: &jsonschema.Schema{ + Type: "object", + Required: []string{"profile"}, + Properties: map[string]*jsonschema.Schema{ + "profile": profileProp, + "table_pattern": optionalTablePattern, + "include_system": optionalIncludeSystem, + }, + }, + }, h.schemaListTablesHandler) + + server.AddTool(&mcp.Tool{ + Name: "schema_dump", + Description: "Dump database schema (tables, columns, indexes, foreign keys); use table_pattern to limit size", + InputSchema: &jsonschema.Schema{ + Type: "object", + Required: []string{"profile"}, + Properties: map[string]*jsonschema.Schema{ + "profile": profileProp, + "table_pattern": optionalTablePattern, + "include_system": optionalIncludeSystem, + }, + }, + }, h.schemaDumpHandler) + + server.AddTool(&mcp.Tool{ + Name: "schema_describe_table", + Description: "Describe a single table (columns, indexes, foreign keys)", + InputSchema: &jsonschema.Schema{ + Type: "object", + Required: []string{"profile", "table"}, + Properties: map[string]*jsonschema.Schema{ + "profile": profileProp, + "schema": {Type: "string", Description: "Schema name (required for PostgreSQL when not public)"}, + "table": {Type: "string", Description: "Table name"}, + }, + }, + }, h.schemaDescribeTableHandler) +} + +func (h *ToolHandler) schemaListTablesHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + var input SchemaListTablesInput + if err := json.Unmarshal(req.Params.Arguments, &input); err != nil { + return h.errorFromErr(errors.Wrap(errors.CodeCfgInvalid, msgInvalidInput, nil, err)), nil + } + result, _, err := h.SchemaListTables(ctx, req, input) + return result, err +} + +func (h *ToolHandler) schemaDumpHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + var input SchemaDumpInput + if err := json.Unmarshal(req.Params.Arguments, &input); err != nil { + return h.errorFromErr(errors.Wrap(errors.CodeCfgInvalid, msgInvalidInput, nil, err)), nil + } + result, _, err := h.SchemaDump(ctx, req, input) + return result, err +} + +func (h *ToolHandler) schemaDescribeTableHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + var input SchemaDescribeTableInput + if err := json.Unmarshal(req.Params.Arguments, &input); err != nil { + return h.errorFromErr(errors.Wrap(errors.CodeCfgInvalid, msgInvalidInput, nil, err)), nil + } + result, _, err := h.SchemaDescribeTable(ctx, req, input) + return result, err +} + +// SchemaListTables lists tables via internal/app. +func (h *ToolHandler) SchemaListTables(ctx context.Context, req *mcp.CallToolRequest, input SchemaListTablesInput) (*mcp.CallToolResult, any, error) { + if input.Profile == "" { + return h.errorResult(errors.New(errors.CodeCfgInvalid, msgProfileRequired, nil)), nil, nil + } + profile, xe := h.resolveProfile(input.Profile) + if xe != nil { + return h.errorResult(xe), nil, nil + } + + timeout := app.SchemaTimeout(profile, 0, false, defaultMCPSchemaTimeout) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + start := time.Now() + result, xe := schemaApp.ListTables(ctx, app.TableListRequest{ + Profile: profile, + TablePattern: input.TablePattern, + IncludeSystem: input.IncludeSystem, + AllowPlaintext: profile.AllowPlaintext, + SkipHostKeyCheck: false, + }) + h.recordMCPStats("schema_list_tables", input.Profile, xe == nil, time.Since(start), xe, "") + + if xe != nil { + return h.errorResult(xe), nil, nil + } + return h.successResult(result) +} + +// SchemaDump dumps schema via internal/app. +func (h *ToolHandler) SchemaDump(ctx context.Context, req *mcp.CallToolRequest, input SchemaDumpInput) (*mcp.CallToolResult, any, error) { + if input.Profile == "" { + return h.errorResult(errors.New(errors.CodeCfgInvalid, msgProfileRequired, nil)), nil, nil + } + profile, xe := h.resolveProfile(input.Profile) + if xe != nil { + return h.errorResult(xe), nil, nil + } + + timeout := app.SchemaTimeout(profile, 0, false, defaultMCPSchemaTimeout) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + start := time.Now() + result, xe := schemaApp.DumpSchema(ctx, app.SchemaDumpRequest{ + Profile: profile, + TablePattern: input.TablePattern, + IncludeSystem: input.IncludeSystem, + AllowPlaintext: profile.AllowPlaintext, + SkipHostKeyCheck: false, + }) + h.recordMCPStats("schema_dump", input.Profile, xe == nil, time.Since(start), xe, "") + + if xe != nil { + return h.errorResult(xe), nil, nil + } + return h.successResult(result) +} + +// SchemaDescribeTable describes one table via internal/app. +func (h *ToolHandler) SchemaDescribeTable(ctx context.Context, req *mcp.CallToolRequest, input SchemaDescribeTableInput) (*mcp.CallToolResult, any, error) { + if input.Profile == "" { + return h.errorResult(errors.New(errors.CodeCfgInvalid, msgProfileRequired, nil)), nil, nil + } + if input.Table == "" { + return h.errorResult(errors.New(errors.CodeCfgInvalid, "table is required", nil)), nil, nil + } + profile, xe := h.resolveProfile(input.Profile) + if xe != nil { + return h.errorResult(xe), nil, nil + } + + timeout := app.SchemaTimeout(profile, 0, false, defaultMCPSchemaTimeout) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + start := time.Now() + result, xe := schemaApp.DescribeTable(ctx, app.TableDescribeRequest{ + Profile: profile, + Schema: input.Schema, + Name: input.Table, + AllowPlaintext: profile.AllowPlaintext, + SkipHostKeyCheck: false, + }) + h.recordMCPStats("schema_describe_table", input.Profile, xe == nil, time.Since(start), xe, "") + + if xe != nil { + return h.errorResult(xe), nil, nil + } + return h.successResult(result) +} diff --git a/internal/mcp/schema_tools_test.go b/internal/mcp/schema_tools_test.go new file mode 100644 index 0000000..e8322af --- /dev/null +++ b/internal/mcp/schema_tools_test.go @@ -0,0 +1,437 @@ +package mcp + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/zx06/xsql/internal/app" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" + "github.com/zx06/xsql/internal/stats" +) + +func TestSchemaListTables_MissingProfile(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaListTables(context.Background(), &mcp.CallToolRequest{}, SchemaListTablesInput{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgProfileRequired) +} + +func TestSchemaDump_MissingProfile(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaDump(context.Background(), &mcp.CallToolRequest{}, SchemaDumpInput{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgProfileRequired) +} + +func TestSchemaDescribeTable_MissingProfile(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaDescribeTable(context.Background(), &mcp.CallToolRequest{}, SchemaDescribeTableInput{ + Table: "users", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgProfileRequired) +} + +func TestSchemaDescribeTable_MissingTable(t *testing.T) { + cfg := &config.File{ + Profiles: map[string]config.Profile{ + "dev": {DB: "mysql"}, + }, + } + handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaDescribeTable(context.Background(), &mcp.CallToolRequest{}, SchemaDescribeTableInput{ + Profile: "dev", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", "table is required") +} + +func TestSchemaDescribeTable_ProfileNotFound(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaDescribeTable(context.Background(), &mcp.CallToolRequest{}, SchemaDescribeTableInput{ + Profile: "missing", + Table: "users", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", "") +} + +func TestRegisterTools_IncludesSchemaTools(t *testing.T) { + cfg := &config.File{ + Profiles: map[string]config.Profile{"dev": {DB: "mysql"}}, + } + handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) + server := mcp.NewServer(&mcp.Implementation{Name: "test", Version: "1"}, nil) + handler.RegisterTools(server) +} + +func TestSchemaListTablesHandler_InvalidJSON(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte("{")}, + } + result, err := handler.schemaListTablesHandler(context.Background(), req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgInvalidInput) +} + +func TestSchemaDumpHandler_InvalidJSON(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte("not-json")}, + } + result, err := handler.schemaDumpHandler(context.Background(), req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgInvalidInput) +} + +func TestSchemaDescribeTableHandler_InvalidJSON(t *testing.T) { + handler := NewToolHandler(emptyMCPConfig(), stats.StatsConfig{}, ServerOptions{}) + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte("{")}, + } + result, err := handler.schemaDescribeTableHandler(context.Background(), req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", msgInvalidInput) +} + +func TestSchemaListTables_SuccessAndRequestOptions(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + var captured app.TableListRequest + restore := swapSchemaHooks(t, schemaHooks{ + listTables: func(_ context.Context, req app.TableListRequest) (*db.TableList, *errors.XError) { + captured = req + return &db.TableList{ + Database: "appdb", + Tables: []db.TableSummary{ + {Schema: "appdb", Name: "users"}, + }, + }, nil + }, + }) + defer restore() + + result, _, err := handler.SchemaListTables(context.Background(), &mcp.CallToolRequest{}, SchemaListTablesInput{ + Profile: "dev", + TablePattern: "user*", + IncludeSystem: true, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolSuccess(t, result, func(data map[string]any) { + if data["database"] != "appdb" { + t.Fatalf("database: %v", data["database"]) + } + }) + if captured.TablePattern != "user*" || !captured.IncludeSystem { + t.Fatalf("unexpected list request: %+v", captured) + } + if captured.Profile.DB != "mysql" { + t.Fatalf("profile not passed: %+v", captured.Profile) + } +} + +func TestSchemaListTables_AppError(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + restore := swapSchemaHooks(t, schemaHooks{ + listTables: func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) { + return nil, errors.New(errors.CodeDBConnectFailed, "connect failed", nil) + }, + }) + defer restore() + + result, _, err := handler.SchemaListTables(context.Background(), &mcp.CallToolRequest{}, SchemaListTablesInput{ + Profile: "dev", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "DB_CONNECT_FAILED", "") +} + +func TestSchemaDump_SuccessWithPattern(t *testing.T) { + handler := testSchemaHandler(t, "pg") + var captured app.SchemaDumpRequest + restore := swapSchemaHooks(t, schemaHooks{ + dumpSchema: func(_ context.Context, req app.SchemaDumpRequest) (*db.SchemaInfo, *errors.XError) { + captured = req + return &db.SchemaInfo{ + Database: "proddb", + Tables: []db.Table{ + {Schema: "public", Name: "orders", Columns: []db.Column{{Name: "id", Type: "bigint"}}}, + }, + }, nil + }, + }) + defer restore() + + result, _, err := handler.SchemaDump(context.Background(), &mcp.CallToolRequest{}, SchemaDumpInput{ + Profile: "dev", + TablePattern: "order*", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolSuccess(t, result, func(data map[string]any) { + tables, ok := data["tables"].([]any) + if !ok || len(tables) != 1 { + t.Fatalf("tables: %v", data["tables"]) + } + }) + if captured.TablePattern != "order*" { + t.Fatalf("pattern not forwarded: %+v", captured) + } +} + +func TestSchemaDump_AppError(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + restore := swapSchemaHooks(t, schemaHooks{ + dumpSchema: func(context.Context, app.SchemaDumpRequest) (*db.SchemaInfo, *errors.XError) { + return nil, errors.New(errors.CodeInternal, "dump failed", nil) + }, + }) + defer restore() + + result, _, err := handler.SchemaDump(context.Background(), &mcp.CallToolRequest{}, SchemaDumpInput{Profile: "dev"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "INTERNAL", "") +} + +func TestSchemaDescribeTable_SuccessWithPGSchema(t *testing.T) { + handler := testSchemaHandler(t, "pg") + var captured app.TableDescribeRequest + restore := swapSchemaHooks(t, schemaHooks{ + describeTable: func(_ context.Context, req app.TableDescribeRequest) (*db.Table, *errors.XError) { + captured = req + return &db.Table{ + Schema: "public", + Name: "users", + Columns: []db.Column{{Name: "id", Type: "integer", PrimaryKey: true}}, + }, nil + }, + }) + defer restore() + + result, _, err := handler.SchemaDescribeTable(context.Background(), &mcp.CallToolRequest{}, SchemaDescribeTableInput{ + Profile: "dev", + Schema: "public", + Table: "users", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolSuccess(t, result, func(data map[string]any) { + if data["name"] != "users" { + t.Fatalf("name: %v", data["name"]) + } + }) + if captured.Schema != "public" || captured.Name != "users" { + t.Fatalf("describe request: %+v", captured) + } +} + +func TestSchemaDescribeTable_TableNotFound(t *testing.T) { + handler := testSchemaHandler(t, "pg") + restore := swapSchemaHooks(t, schemaHooks{ + describeTable: func(context.Context, app.TableDescribeRequest) (*db.Table, *errors.XError) { + return nil, errors.New(errors.CodeCfgInvalid, "table not found", map[string]any{"reason": "table_not_found"}) + }, + }) + defer restore() + + result, _, err := handler.SchemaDescribeTable(context.Background(), &mcp.CallToolRequest{}, SchemaDescribeTableInput{ + Profile: "dev", + Schema: "public", + Table: "missing_table", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + assertToolError(t, result, "CFG_INVALID", "table not found") +} + +func TestSchemaListTablesHandler_Success(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + restore := swapSchemaHooks(t, schemaHooks{ + listTables: func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) { + return &db.TableList{Database: "d"}, nil + }, + }) + defer restore() + + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte(`{"profile":"dev"}`)}, + } + result, err := handler.schemaListTablesHandler(context.Background(), req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil || result.IsError { + t.Fatal("expected success from handler") + } +} + +func TestSchemaDumpHandler_Success(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + restore := swapSchemaHooks(t, schemaHooks{ + dumpSchema: func(context.Context, app.SchemaDumpRequest) (*db.SchemaInfo, *errors.XError) { + return &db.SchemaInfo{Database: "d"}, nil + }, + }) + defer restore() + + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte(`{"profile":"dev"}`)}, + } + result, err := handler.schemaDumpHandler(context.Background(), req) + if err != nil || result == nil || result.IsError { + t.Fatalf("expected success, err=%v result=%+v", err, result) + } +} + +func TestSchemaDescribeTableHandler_Success(t *testing.T) { + handler := testSchemaHandler(t, "pg") + restore := swapSchemaHooks(t, schemaHooks{ + describeTable: func(context.Context, app.TableDescribeRequest) (*db.Table, *errors.XError) { + return &db.Table{Name: "t"}, nil + }, + }) + defer restore() + + req := &mcp.CallToolRequest{ + Params: &mcp.CallToolParamsRaw{Arguments: []byte(`{"profile":"dev","table":"t","schema":"public"}`)}, + } + result, err := handler.schemaDescribeTableHandler(context.Background(), req) + if err != nil || result == nil || result.IsError { + t.Fatalf("expected success, err=%v result=%+v", err, result) + } +} + +func TestSchemaListTables_RecordsStats(t *testing.T) { + handler := testSchemaHandler(t, "mysql") + handler.stats = stats.StatsConfig{Enabled: true, FilePath: t.TempDir() + "/stats.jsonl"} + restore := swapSchemaHooks(t, schemaHooks{ + listTables: func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) { + return &db.TableList{Database: "d", Tables: nil}, nil + }, + }) + defer restore() + + _, _, err := handler.SchemaListTables(context.Background(), &mcp.CallToolRequest{}, SchemaListTablesInput{Profile: "dev"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func emptyMCPConfig() *config.File { + return &config.File{Profiles: map[string]config.Profile{}} +} + +func testSchemaHandler(t *testing.T, dbType string) *ToolHandler { + t.Helper() + cfg := &config.File{ + Profiles: map[string]config.Profile{ + "dev": {DB: dbType, Database: "appdb"}, + }, + } + return NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) +} + +type schemaHooks struct { + listTables func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) + dumpSchema func(context.Context, app.SchemaDumpRequest) (*db.SchemaInfo, *errors.XError) + describeTable func(context.Context, app.TableDescribeRequest) (*db.Table, *errors.XError) +} + +func swapSchemaHooks(t *testing.T, h schemaHooks) func() { + t.Helper() + prevList := schemaApp.ListTables + prevDump := schemaApp.DumpSchema + prevDescribe := schemaApp.DescribeTable + if h.listTables != nil { + schemaApp.ListTables = h.listTables + } + if h.dumpSchema != nil { + schemaApp.DumpSchema = h.dumpSchema + } + if h.describeTable != nil { + schemaApp.DescribeTable = h.describeTable + } + return func() { + schemaApp.ListTables = prevList + schemaApp.DumpSchema = prevDump + schemaApp.DescribeTable = prevDescribe + } +} + +func assertToolError(t *testing.T, result *mcp.CallToolResult, code, messageContains string) { + t.Helper() + if result == nil || !result.IsError { + t.Fatal("expected error result") + } + text := result.Content[0].(*mcp.TextContent).Text + if !strings.Contains(text, code) { + t.Fatalf("expected code %s in %s", code, text) + } + if messageContains != "" && !strings.Contains(text, messageContains) { + t.Fatalf("expected message %q in %s", messageContains, text) + } + var envelope struct { + OK bool `json:"ok"` + SchemaVersion int `json:"schema_version"` + Error map[string]any `json:"error"` + } + if err := json.Unmarshal([]byte(text), &envelope); err != nil { + t.Fatalf("invalid error json: %v", err) + } + if envelope.OK || envelope.SchemaVersion != 1 || envelope.Error == nil { + t.Fatalf("unexpected envelope: %+v", envelope) + } +} + +func assertToolSuccess(t *testing.T, result *mcp.CallToolResult, check func(map[string]any)) { + t.Helper() + if result == nil || result.IsError { + t.Fatalf("expected success, got %+v", result) + } + text := result.Content[0].(*mcp.TextContent).Text + var envelope struct { + OK bool `json:"ok"` + SchemaVersion int `json:"schema_version"` + Data map[string]any `json:"data"` + } + if err := json.Unmarshal([]byte(text), &envelope); err != nil { + t.Fatalf("invalid success json: %v", err) + } + if !envelope.OK || envelope.SchemaVersion != 1 || envelope.Data == nil { + t.Fatalf("unexpected envelope: %+v", envelope) + } + if check != nil { + check(envelope.Data) + } +} diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go index 8cbc444..9442639 100644 --- a/internal/mcp/tools.go +++ b/internal/mcp/tools.go @@ -8,16 +8,14 @@ import ( "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/zx06/xsql/internal/app" "github.com/zx06/xsql/internal/config" - "github.com/zx06/xsql/internal/db" - _ "github.com/zx06/xsql/internal/db/mysql" - _ "github.com/zx06/xsql/internal/db/pg" "github.com/zx06/xsql/internal/errors" - "github.com/zx06/xsql/internal/secret" - "github.com/zx06/xsql/internal/ssh" "github.com/zx06/xsql/internal/stats" ) +const defaultMCPQueryTimeout = 30 * time.Second + // QueryInput represents the input for the query tool type QueryInput struct { SQL string `json:"sql" jsonschema:"SQL query to execute"` @@ -65,6 +63,10 @@ func (h *ToolHandler) getProfileNames() []string { return names } +func (h *ToolHandler) resolveProfile(name string) (config.Profile, *errors.XError) { + return app.ResolveProfile(*h.config, name) +} + // RegisterTools registers all tools with the MCP server func (h *ToolHandler) RegisterTools(server *mcp.Server) { profileNames := h.getProfileNames() @@ -118,18 +120,15 @@ func (h *ToolHandler) RegisterTools(server *mcp.Server) { Description: "Show profile details", InputSchema: profileShowSchema, }, h.profileShowHandler) + + h.registerSchemaTools(server, profileEnums) } // queryHandler is the raw handler for query tool func (h *ToolHandler) queryHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { var input QueryInput if err := json.Unmarshal(req.Params.Arguments, &input); err != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.Wrap(errors.CodeCfgInvalid, "invalid input", nil, err))}, - }, - }, nil + return h.errorFromErr(errors.Wrap(errors.CodeCfgInvalid, "invalid input", nil, err)), nil } result, _, err := h.Query(ctx, req, input) return result, err @@ -139,194 +138,44 @@ func (h *ToolHandler) queryHandler(ctx context.Context, req *mcp.CallToolRequest func (h *ToolHandler) profileShowHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { var input ProfileShowInput if err := json.Unmarshal(req.Params.Arguments, &input); err != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.Wrap(errors.CodeCfgInvalid, "invalid input", nil, err))}, - }, - }, nil + return h.errorFromErr(errors.Wrap(errors.CodeCfgInvalid, "invalid input", nil, err)), nil } result, _, err := h.ProfileShow(ctx, req, input) return result, err } -// Query executes a SQL query +// Query executes a SQL query via internal/app. func (h *ToolHandler) Query(ctx context.Context, req *mcp.CallToolRequest, input QueryInput) (*mcp.CallToolResult, any, error) { - // Validate required fields if input.SQL == "" { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "sql is required", nil))}, - }, - }, nil, nil + return h.errorResult(errors.New(errors.CodeCfgInvalid, "sql is required", nil)), nil, nil } - if input.Profile == "" { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "profile is required", nil))}, - }, - }, nil, nil - } - - // Get profile - profile := h.getProfile(input.Profile) - if profile == nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "profile does not exist", map[string]any{"name": input.Profile, "reason": "profile_not_found"}))}, - }, - }, nil, nil - } - - // Validate SSH proxy reference if present - if profile.SSHProxy != "" && profile.SSHConfig == nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "ssh_proxy not found", map[string]any{"profile": input.Profile, "ssh_proxy": profile.SSHProxy, "reason": "ssh_proxy_not_found"}))}, - }, - }, nil, nil + return h.errorResult(errors.New(errors.CodeCfgInvalid, "profile is required", nil)), nil, nil } - if profile.DB == "" { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "db type is required (mysql|pg)", nil))}, - }, - }, nil, nil - } - - // Parse password - password := profile.Password - if password != "" { - pw, xe := secret.Resolve(password, secret.Options{AllowPlaintext: profile.AllowPlaintext}) - if xe != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(xe)}, - }, - }, nil, nil - } - password = pw - } - - // SSH proxy - var sshClient *ssh.Client - if profile.SSHConfig != nil { - passphrase := profile.SSHConfig.Passphrase - if passphrase != "" { - pp, xe := secret.Resolve(passphrase, secret.Options{AllowPlaintext: profile.AllowPlaintext}) - if xe != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(xe)}, - }, - }, nil, nil - } - passphrase = pp - } - sshOpts := ssh.Options{ - Host: profile.SSHConfig.Host, - Port: profile.SSHConfig.Port, - User: profile.SSHConfig.User, - IdentityFile: profile.SSHConfig.IdentityFile, - Passphrase: passphrase, - KnownHostsFile: profile.SSHConfig.KnownHostsFile, - } - sc, xe := ssh.Connect(ctx, sshOpts) - if xe != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(xe)}, - }, - }, nil, nil - } - defer sc.Close() - sshClient = sc - } - - // Get driver - drv, ok := db.Get(profile.DB) - if !ok { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeDBDriverUnsupported, "unsupported db driver", map[string]any{"db": profile.DB}))}, - }, - }, nil, nil - } - - connOpts := db.ConnOptions{ - DSN: profile.DSN, - Host: profile.Host, - Port: profile.Port, - User: profile.User, - Password: password, - Database: profile.Database, - } - if sshClient != nil { - connOpts.Dialer = sshClient - } - - conn, xe := drv.Open(ctx, connOpts) + profile, xe := h.resolveProfile(input.Profile) if xe != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(xe)}, - }, - }, nil, nil + return h.errorResult(xe), nil, nil } - defer conn.Close() - // Query options - use read-only mode by default + timeout := app.QueryTimeout(profile, 0, false, defaultMCPQueryTimeout) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + start := time.Now() - result, xe := db.Query(ctx, conn, input.SQL, db.QueryOptions{ + result, xe := app.Query(ctx, app.QueryRequest{ + Profile: profile, + SQL: input.SQL, + AllowPlaintext: profile.AllowPlaintext, + SkipHostKeyCheck: false, UnsafeAllowWrite: WriteAllowed(h.serverAllowsWrite, profile.UnsafeAllowWrite), - DBType: profile.DB, }) - - // Record stats h.recordMCPStats("query", input.Profile, xe == nil, time.Since(start), xe, input.SQL) if xe != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(xe)}, - }, - }, nil, nil - } - - output := map[string]any{ - "ok": true, - "schema_version": 1, - "data": result, - } - jsonData, err := json.MarshalIndent(output, "", " ") - if err != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.Wrap(errors.CodeInternal, "failed to marshal result", nil, err))}, - }, - }, nil, nil + return h.errorResult(xe), nil, nil } - - // Return result directly in content per RFC - return &mcp.CallToolResult{ - Content: []mcp.Content{ - &mcp.TextContent{Text: string(jsonData)}, - }, - }, nil, nil + return h.successResult(result) } // ProfileList lists all profiles @@ -335,47 +184,20 @@ func (h *ToolHandler) ProfileList(ctx context.Context, req *mcp.CallToolRequest, for name, p := range h.config.Profiles { profiles = append(profiles, config.ProfileToInfo(name, p)) } - - output := map[string]any{ - "ok": true, - "schema_version": 1, - "data": map[string]any{ - "profiles": profiles, - }, - } - jsonData, err := json.MarshalIndent(output, "", " ") - if err != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.Wrap(errors.CodeInternal, "failed to marshal result", nil, err))}, - }, - }, nil, nil - } - - return &mcp.CallToolResult{ - Content: []mcp.Content{ - &mcp.TextContent{Text: string(jsonData)}, - }, - }, nil, nil + return h.successResult(map[string]any{"profiles": profiles}) } // ProfileShow shows profile details func (h *ToolHandler) ProfileShow(ctx context.Context, req *mcp.CallToolRequest, input ProfileShowInput) (*mcp.CallToolResult, any, error) { - profile, ok := h.config.Profiles[input.Name] - if !ok { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.New(errors.CodeCfgInvalid, "profile does not exist", map[string]any{"name": input.Name, "reason": "profile_not_found"}))}, - }, - }, nil, nil + profile, xe := h.resolveProfile(input.Name) + if xe != nil { + return h.errorResult(xe), nil, nil } + p := h.config.Profiles[input.Name] - // Redact sensitive information result := map[string]any{ "name": input.Name, - "description": profile.Description, + "description": p.Description, "db": profile.DB, "host": profile.Host, "port": profile.Port, @@ -384,71 +206,24 @@ func (h *ToolHandler) ProfileShow(ctx context.Context, req *mcp.CallToolRequest, "unsafe_allow_write": profile.UnsafeAllowWrite, "allow_plaintext": profile.AllowPlaintext, } - if profile.DSN != "" { + if p.DSN != "" { result["dsn"] = "***" } - if profile.Password != "" { + if p.Password != "" { result["password"] = "***" } if profile.SSHProxy != "" { result["ssh_proxy"] = profile.SSHProxy - if proxy, ok := h.config.SSHProxies[profile.SSHProxy]; ok { - result["ssh_host"] = proxy.Host - result["ssh_port"] = proxy.Port - result["ssh_user"] = proxy.User - if proxy.IdentityFile != "" { - result["ssh_identity_file"] = proxy.IdentityFile + if profile.SSHConfig != nil { + result["ssh_host"] = profile.SSHConfig.Host + result["ssh_port"] = profile.SSHConfig.Port + result["ssh_user"] = profile.SSHConfig.User + if profile.SSHConfig.IdentityFile != "" { + result["ssh_identity_file"] = profile.SSHConfig.IdentityFile } } } - - output := map[string]any{ - "ok": true, - "schema_version": 1, - "data": result, - } - jsonData, err := json.MarshalIndent(output, "", " ") - if err != nil { - return &mcp.CallToolResult{ - IsError: true, - Content: []mcp.Content{ - &mcp.TextContent{Text: h.formatError(errors.Wrap(errors.CodeInternal, "failed to marshal result", nil, err))}, - }, - }, nil, nil - } - - // Return result directly in content per RFC - return &mcp.CallToolResult{ - Content: []mcp.Content{ - &mcp.TextContent{Text: string(jsonData)}, - }, - }, nil, nil -} - -// getProfile gets a profile by name, or returns the default profile -// Also resolves SSH proxy reference if present -func (h *ToolHandler) getProfile(name string) *config.Profile { - if name == "" { - // Use default profile (first one) - for _, p := range h.config.Profiles { - return &p - } - return nil - } - - profile, ok := h.config.Profiles[name] - if !ok { - return nil - } - - // Resolve SSH proxy reference - if profile.SSHProxy != "" { - if proxy, ok := h.config.SSHProxies[profile.SSHProxy]; ok { - profile.SSHConfig = &proxy - } - } - - return &profile + return h.successResult(result) } // formatError formats an error as JSON diff --git a/internal/mcp/tools_test.go b/internal/mcp/tools_test.go index 5cbc1aa..59abe62 100644 --- a/internal/mcp/tools_test.go +++ b/internal/mcp/tools_test.go @@ -64,35 +64,7 @@ func TestNewToolHandler(t *testing.T) { } } -func TestGetProfile_Default(t *testing.T) { - cfg := &config.File{ - Profiles: map[string]config.Profile{ - "default": { - DB: "mysql", - Host: "localhost", - }, - "other": { - DB: "pg", - Host: "localhost", - }, - }, - } - - handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) - - // Get default profile (empty name) - profile := handler.getProfile("") - if profile == nil { - t.Fatal("expected non-nil profile") - } - - // Should return one of the profiles (order is not guaranteed) - if profile.DB != "mysql" && profile.DB != "pg" { - t.Errorf("expected db=mysql or db=pg, got %s", profile.DB) - } -} - -func TestGetProfile_ByName(t *testing.T) { +func TestResolveProfile_ByName(t *testing.T) { cfg := &config.File{ Profiles: map[string]config.Profile{ "dev": { @@ -108,10 +80,9 @@ func TestGetProfile_ByName(t *testing.T) { handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) - // Get dev profile - profile := handler.getProfile("dev") - if profile == nil { - t.Fatal("expected non-nil profile") + profile, xe := handler.resolveProfile("dev") + if xe != nil { + t.Fatalf("unexpected error: %v", xe) } if profile.DB != "mysql" { @@ -122,10 +93,9 @@ func TestGetProfile_ByName(t *testing.T) { t.Errorf("expected host=dev.example.com, got %s", profile.Host) } - // Get prod profile - profile = handler.getProfile("prod") - if profile == nil { - t.Fatal("expected non-nil profile") + profile, xe = handler.resolveProfile("prod") + if xe != nil { + t.Fatalf("unexpected error: %v", xe) } if profile.DB != "pg" { @@ -133,7 +103,7 @@ func TestGetProfile_ByName(t *testing.T) { } } -func TestGetProfile_NotFound(t *testing.T) { +func TestResolveProfile_NotFound(t *testing.T) { cfg := &config.File{ Profiles: map[string]config.Profile{ "dev": { @@ -144,24 +114,9 @@ func TestGetProfile_NotFound(t *testing.T) { handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) - // Get non-existent profile - profile := handler.getProfile("nonexistent") - if profile != nil { - t.Error("expected nil profile") - } -} - -func TestGetProfile_EmptyProfiles(t *testing.T) { - cfg := &config.File{ - Profiles: map[string]config.Profile{}, - } - - handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) - - // Get default profile with empty config - profile := handler.getProfile("") - if profile != nil { - t.Error("expected nil profile when no profiles exist") + _, xe := handler.resolveProfile("nonexistent") + if xe == nil { + t.Error("expected error for missing profile") } } @@ -774,12 +729,11 @@ func TestQuery_AllFields(t *testing.T) { // Test with all fields populated // This test requires actual database connection or mocking // For now, we'll just verify that the profile is correctly retrieved - profile := handler.getProfile("full") - if profile == nil { - t.Fatal("expected non-nil profile") + profile, xe := handler.resolveProfile("full") + if xe != nil { + t.Fatalf("resolveProfile: %v", xe) } - // Verify all fields are set if profile.DB != "mysql" { t.Errorf("expected DB=mysql, got %s", profile.DB) } @@ -863,7 +817,7 @@ func TestProfileShowHandler_InvalidJSON(t *testing.T) { } } -func TestGetProfile_ResolvesSSHProxy(t *testing.T) { +func TestResolveProfile_ResolvesSSHProxy(t *testing.T) { cfg := &config.File{ Profiles: map[string]config.Profile{ "dev": {DB: "mysql", SSHProxy: "bastion"}, @@ -874,9 +828,9 @@ func TestGetProfile_ResolvesSSHProxy(t *testing.T) { } handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) - profile := handler.getProfile("dev") - if profile == nil { - t.Fatal("expected non-nil profile") + profile, xe := handler.resolveProfile("dev") + if xe != nil { + t.Fatalf("resolveProfile: %v", xe) } if profile.SSHConfig == nil { t.Fatal("expected SSHConfig to be resolved") diff --git a/tests/e2e/MCP_E2E_TESTS.md b/tests/e2e/MCP_E2E_TESTS.md index 2bf0773..669bd55 100644 --- a/tests/e2e/MCP_E2E_TESTS.md +++ b/tests/e2e/MCP_E2E_TESTS.md @@ -49,7 +49,15 @@ - **TestMCPProfileShow_NotFound**: 测试不存在的配置文件的错误处理 -### 4. SSH 代理测试 +### 4. Schema Discovery 测试 + +- **TestMCPSchema_ListTables_MySQL** / **TestMCPSchema_ListTables_PostgreSQL**: 轻量表清单 +- **TestMCPSchema_Dump_MySQL**: 完整 schema dump JSON 契约 +- **TestMCPSchema_DescribeTable_PostgreSQL**: 单表结构详情 + +Tools: `schema_list_tables`, `schema_dump`, `schema_describe_table`（经 `internal/app`，与 CLI/Web 一致） + +### 5. SSH 代理测试 - **TestMCPQuery_SSHProxy_ConfigurationHandling**: 验证 SSH 代理配置被正确处理 - **TestMCPQuery_SSHProxy_InvalidConfiguration**: 测试无效的 SSH 代理配置错误 diff --git a/tests/e2e/mcp_test.go b/tests/e2e/mcp_test.go index 28d2f0f..de38142 100644 --- a/tests/e2e/mcp_test.go +++ b/tests/e2e/mcp_test.go @@ -4,7 +4,9 @@ package e2e import ( "bytes" + "database/sql" "encoding/json" + "fmt" "os" "os/exec" "path/filepath" @@ -14,7 +16,9 @@ import ( "testing" "time" + _ "github.com/go-sql-driver/mysql" "github.com/modelcontextprotocol/go-sdk/mcp" + _ "github.com/jackc/pgx/v5/stdlib" ) // TestMCPServer_Startup tests that MCP server starts correctly and returns tools list @@ -46,7 +50,14 @@ func TestMCPServer_Startup(t *testing.T) { toolNames[tool.Name] = true } - expectedTools := []string{"query", "profile_list", "profile_show"} + expectedTools := []string{ + "query", + "profile_list", + "profile_show", + "schema_list_tables", + "schema_dump", + "schema_describe_table", + } for _, name := range expectedTools { if !toolNames[name] { t.Errorf("expected tool %q not found", name) @@ -1587,3 +1598,238 @@ profiles: // Could fail with SSH_HOST_KEY_MISMATCH or other SSH errors t.Logf("Error code: %s", response.Error.Code) } + +// ============================================================================ +// MCP Schema Discovery E2E Tests +// ============================================================================ + +func createMySQLMCPE2ETable(t *testing.T) string { + t.Helper() + dsn := mysqlDSN(t) + conn, err := sql.Open("mysql", dsn) + if err != nil { + t.Fatalf("open mysql: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + table := fmt.Sprintf("mcp_e2e_%d", time.Now().UnixNano()) + _, err = conn.Exec(fmt.Sprintf("CREATE TABLE %s (id INT PRIMARY KEY, label VARCHAR(32))", table)) + if err != nil { + t.Fatalf("create table: %v", err) + } + t.Cleanup(func() { _, _ = conn.Exec("DROP TABLE IF EXISTS " + table) }) + return table +} + +func createPGMCPE2ETable(t *testing.T) (schema, table string) { + t.Helper() + dsn := pgDSN(t) + conn, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open pg: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + schema = "public" + table = fmt.Sprintf("mcp_e2e_%d", time.Now().UnixNano()) + _, err = conn.Exec(fmt.Sprintf(`CREATE TABLE %s.%s (id SERIAL PRIMARY KEY, label TEXT)`, schema, table)) + if err != nil { + t.Fatalf("create table: %v", err) + } + t.Cleanup(func() { _, _ = conn.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s.%s", schema, table)) }) + return schema, table +} + +func TestMCPSchema_ListTables_MySQL(t *testing.T) { + dsn := mysqlDSN(t) + table := createMySQLMCPE2ETable(t) + config := createTempConfig(t, `profiles: + test: + db: mysql + dsn: "`+dsn+`" +`) + + result, err := callMCPTool(t, config, "schema_list_tables", map[string]any{ + "profile": "test", + "table_pattern": table[:10] + "*", + }) + if err != nil { + t.Fatalf("schema_list_tables failed: %v", err) + } + if result.IsError { + t.Fatalf("expected success, got: %v", result.Content) + } + + textContent := result.Content[0].(*mcp.TextContent) + var response struct { + OK bool `json:"ok"` + Data struct { + Database string `json:"database"` + Tables []struct { + Name string `json:"name"` + } `json:"tables"` + } `json:"data"` + } + if err := json.Unmarshal([]byte(textContent.Text), &response); err != nil { + t.Fatalf("parse response: %v", err) + } + if !response.OK { + t.Fatal("expected ok=true") + } + if response.Data.Database == "" { + t.Error("expected non-empty database name") + } + found := false + for _, row := range response.Data.Tables { + if row.Name == table { + found = true + break + } + } + if !found { + t.Fatalf("expected created table %q in list, got %+v", table, response.Data.Tables) + } +} + +func TestMCPSchema_Dump_MySQL(t *testing.T) { + dsn := mysqlDSN(t) + table := createMySQLMCPE2ETable(t) + config := createTempConfig(t, `profiles: + test: + db: mysql + dsn: "`+dsn+`" +`) + + result, err := callMCPTool(t, config, "schema_dump", map[string]any{ + "profile": "test", + "table_pattern": table, + }) + if err != nil { + t.Fatalf("schema_dump failed: %v", err) + } + if result.IsError { + t.Fatalf("expected success, got: %v", result.Content) + } + + textContent := result.Content[0].(*mcp.TextContent) + var response struct { + OK bool `json:"ok"` + SchemaVersion int `json:"schema_version"` + Data struct { + Tables []any `json:"tables"` + } `json:"data"` + } + if err := json.Unmarshal([]byte(textContent.Text), &response); err != nil { + t.Fatalf("parse response: %v", err) + } + if response.SchemaVersion != 1 { + t.Errorf("schema_version want 1 got %d", response.SchemaVersion) + } + if len(response.Data.Tables) != 1 { + t.Fatalf("expected 1 table in dump, got %d", len(response.Data.Tables)) + } +} + +func TestMCPSchema_ListTables_PostgreSQL(t *testing.T) { + dsn := pgDSN(t) + schema, table := createPGMCPE2ETable(t) + config := createTempConfig(t, `profiles: + test: + db: pg + dsn: "`+dsn+`" +`) + + result, err := callMCPTool(t, config, "schema_list_tables", map[string]any{ + "profile": "test", + "table_pattern": table[:10] + "*", + }) + if err != nil { + t.Fatalf("schema_list_tables failed: %v", err) + } + if result.IsError { + t.Fatalf("expected success, got: %v", result.Content) + } + + textContent := result.Content[0].(*mcp.TextContent) + var listResp struct { + Data struct { + Tables []struct { + Schema string `json:"schema"` + Name string `json:"name"` + } `json:"tables"` + } `json:"data"` + } + if err := json.Unmarshal([]byte(textContent.Text), &listResp); err != nil { + t.Fatalf("parse list: %v", err) + } + found := false + for _, row := range listResp.Data.Tables { + if row.Schema == schema && row.Name == table { + found = true + break + } + } + if !found { + t.Fatalf("expected table %s.%s in list", schema, table) + } +} + +func TestMCPSchema_DescribeTable_PostgreSQL(t *testing.T) { + dsn := pgDSN(t) + schema, table := createPGMCPE2ETable(t) + config := createTempConfig(t, `profiles: + test: + db: pg + dsn: "`+dsn+`" +`) + + descResult, err := callMCPTool(t, config, "schema_describe_table", map[string]any{ + "profile": "test", + "schema": schema, + "table": table, + }) + if err != nil { + t.Fatalf("schema_describe_table failed: %v", err) + } + if descResult.IsError { + t.Fatalf("expected success, got: %v", descResult.Content) + } + + text := descResult.Content[0].(*mcp.TextContent).Text + var descResp struct { + OK bool `json:"ok"` + Data struct { + Name string `json:"name"` + Columns []struct { + Name string `json:"name"` + } `json:"columns"` + } `json:"data"` + } + if err := json.Unmarshal([]byte(text), &descResp); err != nil { + t.Fatalf("parse describe: %v", err) + } + if descResp.Data.Name != table || len(descResp.Data.Columns) == 0 { + t.Fatalf("unexpected describe payload: %+v", descResp.Data) + } +} + +func TestMCPSchema_DescribeTable_MySQL(t *testing.T) { + dsn := mysqlDSN(t) + table := createMySQLMCPE2ETable(t) + config := createTempConfig(t, `profiles: + test: + db: mysql + dsn: "`+dsn+`" +`) + + descResult, err := callMCPTool(t, config, "schema_describe_table", map[string]any{ + "profile": "test", + "table": table, + }) + if err != nil { + t.Fatalf("schema_describe_table failed: %v", err) + } + if descResult.IsError { + t.Fatalf("expected success, got: %v", descResult.Content) + } +}