diff --git a/.agents/skills/xsql/SKILL.md b/.agents/skills/xsql/SKILL.md index 3501fc6e..edcf7cd9 100644 --- a/.agents/skills/xsql/SKILL.md +++ b/.agents/skills/xsql/SKILL.md @@ -89,6 +89,8 @@ For long-running queries, set a timeout (default: 30s): xsql query "" -p -f json --query-timeout 60 --attr source=codex-cli --attr agent=codex --attr env= --attr task=long-query ``` +Results are capped at **1000 rows** by default (`--max-rows`; MCP default **500**). Use explicit `LIMIT` in SQL for pagination. If `data.truncated` is true, follow `data.hint`. Pass `--max-rows 0` only when you truly need an unlimited result set. + ### Step 4: Validate the response Check the JSON envelope: diff --git a/README.md b/README.md index a2c848c1..10c3aba2 100644 --- a/README.md +++ b/README.md @@ -174,7 +174,7 @@ Note: Read-only mode by default. Writes require both profile `unsafe_allow_write | Command | Description | |---------|-------------| | `xsql ai [PROMPT]` | Start interactive AI assistant mode (TUI) | -| `xsql query ` | Execute SQL queries (read-only by default) | +| `xsql query ` | Execute SQL queries (read-only by default; 1000-row default cap, `--max-rows`) | | `xsql schema dump` | Export database schema (tables, columns, indexes, foreign keys) | | `xsql profile list` | List all profiles | | `xsql profile show ` | Show profile details (passwords are masked) | diff --git a/README_zh.md b/README_zh.md index fc6f1118..b85e1490 100644 --- a/README_zh.md +++ b/README_zh.md @@ -174,7 +174,7 @@ xsql query "" -p -f json --attr source=codex-cli --attr agent=cod | 命令 | 说明 | |------|------| | `xsql ai [PROMPT]` | 启动全屏交互式 AI 终端模式（TUI） | -| `xsql query ` | 执行 SQL 查询（默认只读） | +| `xsql query ` | 执行 SQL 查询（默认只读；默认最多 1000 行，`--max-rows` 可调） | | `xsql schema dump` | 导出数据库结构（表、列、索引、外键） | | `xsql profile list` | 列出所有 profile | | `xsql profile show ` | 查看 profile 详情（密码脱敏） | diff --git a/cmd/xsql/query.go b/cmd/xsql/query.go index a64c433d..8345907c 100644 --- a/cmd/xsql/query.go +++ b/cmd/xsql/query.go @@ -2,6 +2,7 @@ package main import ( "context" + "os" "time" "github.com/spf13/cobra" @@ -24,6 +25,8 @@ type QueryFlags struct { SSHSkipHostKey bool QueryTimeout int QueryTimeoutSet bool + MaxRows int + MaxRowsSet bool } // NewQueryCommand creates the query command @@ -36,6 +39,7 @@ func NewQueryCommand(w *output.Writer) *cobra.Command { Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { flags.QueryTimeoutSet = cmd.Flags().Changed("query-timeout") + flags.MaxRowsSet = cmd.Flags().Changed("max-rows") return runQuery(args, flags, w) }, } @@ -44,6 +48,7 @@ func NewQueryCommand(w *output.Writer) *cobra.Command { cmd.Flags().BoolVar(&flags.AllowPlaintext, "allow-plaintext", false, "Allow plaintext secrets in config") cmd.Flags().BoolVar(&flags.SSHSkipHostKey, "ssh-skip-known-hosts-check", false, "Skip SSH known_hosts check (dangerous)") cmd.Flags().IntVar(&flags.QueryTimeout, "query-timeout", 0, "Query timeout in seconds (default: 30)") + cmd.Flags().IntVar(&flags.MaxRows, "max-rows", 0, "Maximum rows to return (default: 1000; 0 = unlimited)") return cmd } @@ -57,6 +62,10 @@ func runQuery(args []string, flags *QueryFlags, w *output.Writer) error { } p := GlobalConfig.Resolved.Profile + maxRows, xe := app.ResolveQueryMaxRows(p, os.Getenv("XSQL_MAX_ROWS"), flags.MaxRows, flags.MaxRowsSet) + if xe != nil { + return xe + } timeout := app.QueryTimeout(p, flags.QueryTimeout, flags.QueryTimeoutSet, DefaultQueryTimeout) ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() @@ -68,6 +77,8 @@ func runQuery(args []string, flags *QueryFlags, w *output.Writer) error { AllowPlaintext: flags.AllowPlaintext, SkipHostKeyCheck: flags.SSHSkipHostKey, UnsafeAllowWrite: cliWriteAllowed(flags.UnsafeAllowWrite, p.UnsafeAllowWrite), + MaxRows: maxRows, + TruncationHint: app.QueryTruncationHint, }) // Record stats diff --git a/docs/ai.md b/docs/ai.md index 19d2d805..064c9265 100644 --- a/docs/ai.md +++ b/docs/ai.md @@ -77,7 +77,7 @@ xsql mcp server --transport streamable_http --http-addr 127.0.0.1:8787 --http-au ### MCP Tools MCP Server 提供以下 tools（均经 `internal/app`，SSH/连接行为与 CLI/Web 一致）： -- **query**: 执行 SQL 查询（默认只读；写入需 profile `unsafe_allow_write: true` 且 server 启动时 `--unsafe-allow-write`） +- **query**: 执行 SQL 查询（默认只读；写入需 profile `unsafe_allow_write: true` 且 server 启动时 `--unsafe-allow-write`）。可选 `max_rows`（`0` = 默认；受 `mcp.max_rows_limit` clamp；默认 500 或 profile/mcp 配置）。截断时 `data` 含 `truncated`/`truncated_reason`/`max_rows`/`max_response_bytes`/`hint` 等。 - **schema_list_tables**: 轻量表清单（可选 `table_pattern`、`include_system`），适合大库先做 discovery - **schema_dump**: 与 `xsql schema dump` 等价的结构化导出（建议配合 `table_pattern` 控制体积） - **schema_describe_table**: 单表结构（列、索引、外键）；PostgreSQL 需 `schema`，MySQL 可省略 diff --git a/docs/cli-spec.md b/docs/cli-spec.md index f7807870..2d7b2a7f 100644 --- a/docs/cli-spec.md +++ b/docs/cli-spec.md @@ -74,6 +74,8 @@ xsql query "INSERT INTO logs (msg) VALUES ('test')" --profile dev --unsafe-allow | `--unsafe-allow-write` | false | 本次命令申请写入；仅当 profile 同时设置 `unsafe_allow_write: true` 时生效 | | `--allow-plaintext` | false | 允许配置中使用明文密码（也可在配置文件中设置 `allow_plaintext: true`） | | `--ssh-skip-known-hosts-check` | false | 跳过 SSH 主机密钥验证（危险） | +| `--max-rows` | 1000 | 最多返回行数（`0` = 不限；ENV：`XSQL_MAX_ROWS`；可被 profile `max_rows` 覆盖默认值） | +| `--query-timeout` | 30 | 查询超时（秒）；profile 可用 `query_timeout` 覆盖 | **输出示例（JSON）：** ```json @@ -90,6 +92,23 @@ xsql query "INSERT INTO logs (msg) VALUES ('test')" --profile dev --unsafe-allow } ``` +**截断示例（JSON，字段只增）：** +```json +{ + "ok": true, + "schema_version": 1, + "data": { + "columns": ["n"], + "rows": [{"n": 1}, {"n": 2}], + "truncated": true, + "row_count": 2, + "max_rows": 2, + "truncated_reason": "max_rows", + "hint": "Result truncated by max_rows; add SQL LIMIT/OFFSET, use aggregation, or run MCP `schema_list_tables` or `xsql schema dump` for schema discovery." + } +} +``` + **输出示例（Table）：** ``` id name @@ -107,7 +126,7 @@ id,name 2,Bob ``` -> 注：Table 和 CSV 格式不包含 `ok` 和 `schema_version` 元数据，直接输出数据。 +> 注：Table 和 CSV 格式不包含 `ok` 和 `schema_version` 元数据，直接输出数据。CSV 若结果被截断，会在 **stderr** 打印一行截断说明与 hint（stdout 仍为纯 CSV，退出码 0）。 ### `xsql schema dump` @@ -273,6 +292,8 @@ Web UI 使用以下 HTTP API，统一返回 JSON Envelope： Web 查询强制只读，即使 profile 配置了 `unsafe_allow_write: true`，也不会在 Web 接口中生效。 +`POST /api/v1/query` JSON body：`{"profile":"...","sql":"...","max_rows":1000}`（`max_rows` 可选；省略时使用 profile / `XSQL_MAX_ROWS` / 默认 1000；`0` = 不限）。 + **配置 API 脱敏**：`GET /api/v1/config` 与 `GET /api/v1/profiles/{name}` 不会在响应中返回明文密码、DSN、SSH passphrase、AI `api_key` 或 HTTP `auth_token`。明文值以占位符 `***` 表示；`keyring:` 引用会原样返回（不解析 keyring 内容）。通过 Web UI 保存配置时，若客户端未修改某敏感字段（提交空值或 `***`），服务端会保留磁盘上已有值，而不会用占位符覆盖。 > **安全说明**：绑定 loopback 时 Web 默认不要求 Bearer token，本机其他进程仍可能调用上述只读配置接口获取脱敏后的配置与 keyring 引用。生产或共享环境请启用 `web.http.auth_token` 或非 loopback 绑定策略（见 `docs/config.md`）。 @@ -544,7 +565,8 @@ xsql mcp server --transport streamable_http --http-addr 127.0.0.1:8787 --http-au "type": "object", "properties": { "sql": {"type": "string", "description": "SQL query to execute"}, - "profile": {"type": "string", "description": "Profile name to use"} + "profile": {"type": "string", "description": "Profile name to use"}, + "max_rows": {"type": "integer", "description": "Maximum rows to return (0 = defaults; clamped by mcp.max_rows_limit; MCP default 500)"} }, "required": ["sql", "profile"] } diff --git a/docs/config.md b/docs/config.md index 105d8b97..543e95da 100644 --- a/docs/config.md +++ b/docs/config.md @@ -123,6 +123,11 @@ profiles: | `mcp.http.addr` | string | Streamable HTTP 监听地址 | | `mcp.http.auth_token` | string | Streamable HTTP 鉴权 token（支持 `keyring:` 引用） | | `mcp.http.allow_plaintext_token` | bool | 允许在配置中使用明文 token | +| `mcp.max_rows` | int | MCP `query` tool 默认行数上限（默认 500；profile `max_rows` 优先） | +| `mcp.max_rows_limit` | int | MCP `max_rows` 硬上限（默认 10000；更大请求会被 clamp） | +| `mcp.max_response_bytes` | int | MCP query 响应 JSON 体积上限（默认 524288） | + +Profile 还可设置 `max_rows`（query 默认行数上限，CLI/Web 默认 1000）、`query_timeout`、`schema_timeout`。 ## AI 配置项 diff --git a/docs/env.md b/docs/env.md index 9915f1d0..22e09d23 100644 --- a/docs/env.md +++ b/docs/env.md @@ -17,6 +17,7 @@ - `XSQL_AI_MODEL`：AI 模型名称 - `XSQL_AI_BASE_URL`：OpenAI 兼容服务 Base URL - `XSQL_AI_API_KEY`：AI 服务 API Key（明文仅保存在当前进程环境中） +- `XSQL_MAX_ROWS`：query 默认最大返回行数（非负整数；`0` = 不限；优先级低于 CLI `--max-rows`） ## 2. 连接参数（计划中，当前未实现） > 当前版本连接参数通过 config 文件的 profile 配置，ENV 支持计划在后续版本实现。 diff --git a/docs/error-contract.md b/docs/error-contract.md index be724ab1..9d8824b1 100644 --- a/docs/error-contract.md +++ b/docs/error-contract.md @@ -118,6 +118,15 @@ DB 类： } ``` +当结果被截断时，`data` 可附加（`schema_version` 仍为 `1`）： +- `truncated`（bool） +- `row_count`（int，已返回行数） +- `max_rows`（int，**生效的配置行数上限**，非当前行数） +- `truncated_reason`（string：`max_rows` 和/或 `max_response_bytes`） +- `max_response_bytes`（int，MCP 字节护栏生效时的阈值） +- `max_rows_requested` / `max_rows_limit`（MCP 请求被 clamp 时） +- `hint`（string，与 `truncated_reason` 一致；多原因时合并） + ### Table 格式（人类可读） ``` id name email @@ -146,5 +155,7 @@ id,name,email | 终端查看 | table (auto) | | 数据导出 | csv | -## 7. 大结果集（计划中） -- `--jsonl`（NDJSON）：每行一个 JSON 对象，用于大结果集/流式消费 +## 7. 大结果集 +- 默认对 query 结果施加 `max_rows`（CLI/Web 默认 1000；MCP 默认 500），扫描阶段提前停止。 +- CLI/Web 可用 `--max-rows 0` 或请求体 `max_rows: 0` 显式取消上限。 +- 未来可选：`--jsonl`（NDJSON）流式消费。 diff --git a/docs/rfcs/0017-query-result-size-limits.md b/docs/rfcs/0017-query-result-size-limits.md new file mode 100644 index 00000000..a6be3e19 --- /dev/null +++ b/docs/rfcs/0017-query-result-size-limits.md @@ -0,0 +1,69 @@ +# RFC 0017: Query result size limits + +Status: Accepted + +## 摘要 + +为 `xsql query`（CLI / Web / MCP / TUI）增加默认行数上限，在数据库驱动层扫描时提前停止，避免大结果集 OOM 与 MCP/agent 上下文膨胀。截断时在成功响应中附加 `truncated`、`row_count`、`max_rows` 与 agent 提示。CLI/Web 允许 `--max-rows 0` 显式取消上限；MCP 保持强制上限。顺带修复 schema 工具的空 `tables` 与 `include_system` 行为。 + +## 背景 / 动机 + +- `internal/db/query.go` 的 `scanRows` 无上限，整表 `SELECT *` 可拖垮进程。 +- MCP 将完整 JSON 放入 tool 文本，成本高于 TUI 的 `BoundToolFeedback`。 +- PR #82 后 MCP query 与 CLI 共用 `app.Query`，可在单一路径扩展。 + +## 方案（Proposed） + +### 默认与优先级 + +| 场景 | 默认上限 | 取消上限 | +|------|----------|----------| +| CLI / Web / TUI | 1000 行 | `--max-rows 0` 或 API `max_rows: 0`（须显式） | +| MCP | 500 行 | 不支持无上限；`max_rows: 0` 视为未设置，走默认链 | + +MCP 硬上限：**`mcp.max_rows_limit`（默认 10000）**，超过则 clamp 并输出 `max_rows_requested` / `max_rows_limit`。 + +优先级：**CLI/API 参数 > ENV (`XSQL_MAX_ROWS`) > profile `max_rows` > MCP `mcp.max_rows`（仅 MCP）> 上表默认值**。 + +选用 1000：足够日常探查，且单行 JSON 通常仍可控；MCP 500 更保守，减少 tool 响应 token。 + +### 输出（`schema_version` 仍为 1） + +成功时 `data` 在既有 `columns`/`rows` 上**只增**字段： + +```json +{ + "truncated": true, + "row_count": 1000, + "max_rows": 1000, + "truncated_reason": "max_rows", + "hint": "Result truncated by max_rows; add SQL LIMIT/OFFSET, use aggregation, or run MCP `schema_list_tables` or `xsql schema dump` for schema discovery." +} +``` + +字节护栏截断时追加 `truncated_reason: "max_response_bytes"`（可与 `max_rows` 并存）、`max_response_bytes`，且 **`max_rows` 始终为生效的配置行数上限**（不是当前行数）。`hint` 与 `truncated_reason` 一致（两种原因均出现时 hint 合并两句）。 + +未截断时不输出 `truncated`/`hint`（`row_count` 可选，与现有客户端兼容）。 + +### CSV + +stdout 仍为纯 CSV；若截断，在 **stderr** 输出一行 `xsql: query result truncated (...)` + hint，**退出码仍为 0**。 + +### MCP 响应字节护栏（第二层） + +配置 `mcp.max_response_bytes`（默认 524288）。序列化后若超限，用二分查找适配行数；metadata 见上。 + +### Schema 修复 + +- 空库 `schema_list_tables` 的 `tables` 统一为 `[]`。 +- `include_system: true` 时包含 MySQL 系统库表与 PostgreSQL `pg_catalog` / `information_schema` 等（与 RFC 0005 一致）。 + +## 兼容性与迁移 + +- 破坏性：默认新增 1000 行 cap；依赖全量无 LIMIT 查询的脚本需 `--max-rows 0` 或调大 profile。 +- `schema_version` 不变。 + +## 测试计划 + +- 单元：`ResolveQueryMaxRows`、`scanRows` 提前停止（sqlmock）、MCP 字节裁剪。 +- 集成 / E2E：MySQL 8.0、PostgreSQL 16 验证默认、自定义、无限、截断元数据。 diff --git a/docs/rfcs/README.md b/docs/rfcs/README.md index 60c210cf..3908fa3b 100644 --- a/docs/rfcs/README.md +++ b/docs/rfcs/README.md @@ -22,6 +22,7 @@ | [0014](0014-mcp-write-dual-authorization.md) | MCP 写入双重授权 | Implemented | | [0015](0015-mcp-app-layer-schema-tools.md) | MCP 复用 app 层与 Schema Tools | Implemented | | [0016](0016-error-guidance-for-agents.md) | AI Agent 错误指引（Error Guidance） | Accepted | +| [0017](0017-query-result-size-limits.md) | Query 结果行数/体积上限 | Accepted | ## 何时需要 RFC 当变更会影响以下任一项时，必须先提交/更新 RFC（并同步更新对应 `docs/*`）： diff --git a/internal/app/app.go b/internal/app/app.go index 48bdea18..55249f35 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -43,6 +43,8 @@ func (a App) BuildSpec() spec.Spec { spec.FlagSpec{Name: "unsafe-allow-write", Default: "false", Description: "Allow writes when profile unsafe_allow_write is true"}, spec.FlagSpec{Name: "allow-plaintext", Default: "false", Description: "Allow plaintext secrets in config"}, spec.FlagSpec{Name: "ssh-skip-known-hosts-check", Default: "false", Description: "Skip SSH known_hosts check (dangerous)"}, + spec.FlagSpec{Name: "max-rows", Env: "XSQL_MAX_ROWS", Default: "1000", Description: "Maximum rows to return (0 = unlimited)"}, + spec.FlagSpec{Name: "query-timeout", Default: "30", Description: "Query timeout in seconds"}, ), }, { diff --git a/internal/app/max_rows.go b/internal/app/max_rows.go new file mode 100644 index 00000000..3e245d2d --- /dev/null +++ b/internal/app/max_rows.go @@ -0,0 +1,107 @@ +package app + +import ( + "strconv" + "strings" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/errors" +) + +const ( + // DefaultQueryMaxRows is the default row cap for CLI, Web, and TUI queries. + DefaultQueryMaxRows = 1000 + // DefaultMCPQueryMaxRows is a lower default for MCP tool responses (smaller agent payloads). + DefaultMCPQueryMaxRows = 500 + // DefaultMCPMaxRowsLimit is the hard ceiling for MCP query max_rows (requests above are clamped). + DefaultMCPMaxRowsLimit = 10000 + // DefaultMCPMaxResponseBytes is the default serialized JSON size cap for MCP query results. + DefaultMCPMaxResponseBytes = 512 * 1024 +) + +// QueryTruncationHint is returned when results are truncated for agents/humans. +const QueryTruncationHint = errors.HintQueryResultTruncated + +// MCPResponseTruncationHint is appended when MCP byte guard trims rows. +const MCPResponseTruncationHint = errors.HintMCPResponseTruncated + +// MCPMaxRowsMeta describes MCP max_rows resolution/clamping for response metadata. +type MCPMaxRowsMeta struct { + Requested int + Clamped bool + Limit int +} + +// ResolveQueryMaxRows returns the effective row limit for CLI/Web/TUI. +// limit 0 means unlimited. Negative values are invalid. +func ResolveQueryMaxRows(profile config.Profile, envMaxRows string, cliValue int, cliSet bool) (limit int, xe *errors.XError) { + if cliSet { + if cliValue < 0 { + return 0, errors.New(errors.CodeCfgInvalid, "max_rows must be >= 0", map[string]any{"max_rows": cliValue}) + } + return cliValue, nil + } + if strings.TrimSpace(envMaxRows) != "" { + v, err := strconv.Atoi(strings.TrimSpace(envMaxRows)) + if err != nil || v < 0 { + return 0, errors.New(errors.CodeCfgInvalid, "XSQL_MAX_ROWS must be a non-negative integer", map[string]any{"value": envMaxRows}) + } + return v, nil + } + if profile.MaxRows > 0 { + return profile.MaxRows, nil + } + return DefaultQueryMaxRows, nil +} + +// ResolveMCPMaxRowsLimit returns the configured MCP max_rows ceiling. +func ResolveMCPMaxRowsLimit(mcpCfg config.MCPConfig) int { + if mcpCfg.MaxRowsLimit > 0 { + return mcpCfg.MaxRowsLimit + } + return DefaultMCPMaxRowsLimit +} + +// ResolveMCPQueryMaxRows returns the MCP row limit (never unlimited). +// MCP max_rows 0 means unset (use defaults). Negative values are rejected. +func ResolveMCPQueryMaxRows(profile config.Profile, mcpCfg config.MCPConfig, argValue int, argSet bool) (effective int, meta MCPMaxRowsMeta, xe *errors.XError) { + if argSet { + if argValue < 0 { + return 0, meta, errors.New(errors.CodeCfgInvalid, "max_rows must be >= 0 for MCP query", map[string]any{"max_rows": argValue}) + } + if argValue == 0 { + argSet = false + } + } + + resolved := 0 + if argSet && argValue > 0 { + resolved = argValue + meta.Requested = argValue + } else if profile.MaxRows > 0 { + resolved = profile.MaxRows + } else if mcpCfg.MaxRows > 0 { + resolved = mcpCfg.MaxRows + } else { + resolved = DefaultMCPQueryMaxRows + } + + ceiling := ResolveMCPMaxRowsLimit(mcpCfg) + if resolved > ceiling { + if meta.Requested == 0 { + meta.Requested = resolved + } + meta.Clamped = true + meta.Limit = ceiling + resolved = ceiling + } + return resolved, meta, nil +} + +// ResolveMCPMaxResponseBytes returns the MCP JSON response byte cap. +func ResolveMCPMaxResponseBytes(mcpCfg config.MCPConfig) int { + if mcpCfg.MaxResponseBytes > 0 { + return mcpCfg.MaxResponseBytes + } + return DefaultMCPMaxResponseBytes +} diff --git a/internal/app/max_rows_test.go b/internal/app/max_rows_test.go new file mode 100644 index 00000000..f21ac053 --- /dev/null +++ b/internal/app/max_rows_test.go @@ -0,0 +1,93 @@ +package app + +import ( + "testing" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/errors" +) + +func TestResolveQueryMaxRows_CLIOverride(t *testing.T) { + limit, xe := ResolveQueryMaxRows(config.Profile{MaxRows: 100}, "", 50, true) + if xe != nil || limit != 50 { + t.Fatalf("got limit=%d xe=%v", limit, xe) + } +} + +func TestResolveQueryMaxRows_CLIUnlimited(t *testing.T) { + limit, xe := ResolveQueryMaxRows(config.Profile{MaxRows: 100}, "", 0, true) + if xe != nil || limit != 0 { + t.Fatalf("got limit=%d xe=%v", limit, xe) + } +} + +func TestResolveQueryMaxRows_ProfileAndDefault(t *testing.T) { + limit, xe := ResolveQueryMaxRows(config.Profile{MaxRows: 200}, "", 0, false) + if xe != nil || limit != 200 { + t.Fatalf("got limit=%d xe=%v", limit, xe) + } + limit, xe = ResolveQueryMaxRows(config.Profile{}, "", 0, false) + if xe != nil || limit != DefaultQueryMaxRows { + t.Fatalf("got limit=%d xe=%v", limit, xe) + } +} + +func TestResolveQueryMaxRows_Invalid(t *testing.T) { + _, xe := ResolveQueryMaxRows(config.Profile{}, "", -1, true) + if xe == nil || xe.Code != errors.CodeCfgInvalid { + t.Fatalf("expected cfg invalid, got %v", xe) + } +} + +func TestResolveQueryMaxRows_Env(t *testing.T) { + limit, xe := ResolveQueryMaxRows(config.Profile{MaxRows: 100}, "250", 0, false) + if xe != nil || limit != 250 { + t.Fatalf("got limit=%d xe=%v", limit, xe) + } + _, xe = ResolveQueryMaxRows(config.Profile{}, "bad", 0, false) + if xe == nil { + t.Fatal("expected invalid env") + } +} + +func TestResolveMCPMaxResponseBytes(t *testing.T) { + if got := ResolveMCPMaxResponseBytes(config.MCPConfig{}); got != DefaultMCPMaxResponseBytes { + t.Fatalf("default bytes: %d", got) + } + if got := ResolveMCPMaxResponseBytes(config.MCPConfig{MaxResponseBytes: 1024}); got != 1024 { + t.Fatalf("custom bytes: %d", got) + } +} + +func TestResolveMCPMaxRowsLimit(t *testing.T) { + if got := ResolveMCPMaxRowsLimit(config.MCPConfig{}); got != DefaultMCPMaxRowsLimit { + t.Fatalf("default limit: %d", got) + } +} + +func TestResolveMCPQueryMaxRows(t *testing.T) { + v, meta, xe := ResolveMCPQueryMaxRows(config.Profile{}, config.MCPConfig{}, 0, false) + if xe != nil || v != DefaultMCPQueryMaxRows || meta.Clamped { + t.Fatalf("default mcp: %d meta=%+v xe=%v", v, meta, xe) + } + v, meta, xe = ResolveMCPQueryMaxRows(config.Profile{}, config.MCPConfig{}, 0, true) + if xe != nil || v != DefaultMCPQueryMaxRows { + t.Fatalf("mcp max_rows 0 unset: %d %v", v, xe) + } + _, _, xe = ResolveMCPQueryMaxRows(config.Profile{}, config.MCPConfig{}, -1, true) + if xe == nil { + t.Fatal("expected error for negative mcp max_rows") + } + v, meta, xe = ResolveMCPQueryMaxRows(config.Profile{MaxRows: 800}, config.MCPConfig{MaxRows: 300}, 0, false) + if xe != nil || v != 800 { + t.Fatalf("profile override: %d %v", v, xe) + } + v, meta, xe = ResolveMCPQueryMaxRows(config.Profile{}, config.MCPConfig{MaxRows: 300}, 0, false) + if xe != nil || v != 300 { + t.Fatalf("mcp config override: %d %v", v, xe) + } + v, meta, xe = ResolveMCPQueryMaxRows(config.Profile{}, config.MCPConfig{MaxRowsLimit: 100}, 500, true) + if xe != nil || v != 100 || !meta.Clamped || meta.Requested != 500 || meta.Limit != 100 { + t.Fatalf("clamp: v=%d meta=%+v", v, meta) + } +} diff --git a/internal/app/service.go b/internal/app/service.go index 08124ffb..3dd66bfb 100644 --- a/internal/app/service.go +++ b/internal/app/service.go @@ -45,6 +45,8 @@ type QueryRequest struct { AllowPlaintext bool SkipHostKeyCheck bool UnsafeAllowWrite bool + MaxRows int + TruncationHint string } // SchemaDumpRequest contains options for a schema dump operation. @@ -181,10 +183,18 @@ func Query(ctx context.Context, req QueryRequest) (*db.QueryResult, *errors.XErr } defer func() { _ = conn.Close() }() - return db.Query(ctx, conn.DB, req.SQL, db.QueryOptions{ + result, xe := db.Query(ctx, conn.DB, req.SQL, db.QueryOptions{ UnsafeAllowWrite: req.UnsafeAllowWrite, DBType: req.Profile.DB, + MaxRows: req.MaxRows, }) + if xe != nil { + return nil, xe + } + if result != nil && result.Truncated && result.Hint == "" { + db.RefreshTruncationHint(result) + } + return result, nil } // DumpSchema exports the schema using a resolved profile. diff --git a/internal/config/types.go b/internal/config/types.go index 08c0e606..09427e87 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -55,6 +55,9 @@ type Profile struct { QueryTimeout int `yaml:"query_timeout" json:"query_timeout"` // query timeout, default 30s SchemaTimeout int `yaml:"schema_timeout" json:"schema_timeout"` // schema export timeout, default 60s + // MaxRows caps rows returned by query (0 = unset; use CLI/API 0 for unlimited) + MaxRows int `yaml:"max_rows" json:"max_rows"` + // SSH proxy reference (refers to a name defined in ssh_proxies) SSHProxy string `yaml:"ssh_proxy" json:"ssh_proxy"` @@ -67,8 +70,11 @@ type Profile struct { // MCPConfig defines the MCP server configuration. type MCPConfig struct { - Transport string `yaml:"transport" json:"transport"` // stdio | streamable_http - HTTP MCPHTTPConfig `yaml:"http" json:"http"` + Transport string `yaml:"transport" json:"transport"` // stdio | streamable_http + HTTP MCPHTTPConfig `yaml:"http" json:"http"` + MaxRows int `yaml:"max_rows" json:"max_rows"` // default MCP query row cap + MaxRowsLimit int `yaml:"max_rows_limit" json:"max_rows_limit"` // hard MCP ceiling for max_rows + MaxResponseBytes int `yaml:"max_response_bytes" json:"max_response_bytes"` // MCP query JSON size cap } // MCPHTTPConfig defines the MCP Streamable HTTP transport configuration. diff --git a/internal/db/csv_notice.go b/internal/db/csv_notice.go new file mode 100644 index 00000000..c57568ab --- /dev/null +++ b/internal/db/csv_notice.go @@ -0,0 +1,19 @@ +package db + +import "fmt" + +// CSVTruncationNotice returns a stderr line when CSV output was truncated. +func (r *QueryResult) CSVTruncationNotice() (string, bool) { + if r == nil || !r.Truncated { + return "", false + } + reason := r.TruncatedReason + if reason == "" { + reason = TruncatedReasonMaxRows + } + msg := fmt.Sprintf("xsql: query result truncated (%s)", reason) + if r.Hint != "" { + msg += "; " + r.Hint + } + return msg, true +} diff --git a/internal/db/mysql/schema.go b/internal/db/mysql/schema.go index 5c97dbf5..041345ad 100644 --- a/internal/db/mysql/schema.go +++ b/internal/db/mysql/schema.go @@ -19,7 +19,7 @@ func (d *Driver) ListTables(ctx context.Context, conn *sql.DB, opts db.SchemaOpt return nil, xe } - tables, xe := d.listTables(ctx, conn, database, opts.TablePattern) + tables, xe := d.listTables(ctx, conn, database, opts.TablePattern, opts.IncludeSystem) if xe != nil { return nil, xe } @@ -96,18 +96,25 @@ func currentDatabase(ctx context.Context, conn *sql.DB) (string, *errors.XError) return database, nil } -func (d *Driver) listTables(ctx context.Context, conn *sql.DB, database, tablePattern string) ([]db.TableSummary, *errors.XError) { +func (d *Driver) listTables(ctx context.Context, conn *sql.DB, database, tablePattern string, includeSystem bool) ([]db.TableSummary, *errors.XError) { query := ` - SELECT table_name, table_comment + SELECT table_schema, table_name, table_comment FROM information_schema.tables - WHERE table_schema = ? AND table_type = 'BASE TABLE' + WHERE table_type = 'BASE TABLE' ` - args := []any{database} + args := []any{} + if includeSystem { + query += " AND (table_schema = ? OR table_schema IN ('information_schema','mysql','performance_schema','sys'))" + args = append(args, database) + } else { + query += " AND table_schema = ?" + args = append(args, database) + } if tablePattern != "" { query += " AND table_name LIKE ?" args = append(args, toLikePattern(tablePattern)) } - query += " ORDER BY table_name" + query += " ORDER BY table_schema, table_name" rows, err := conn.QueryContext(ctx, query, args...) if err != nil { @@ -115,14 +122,14 @@ func (d *Driver) listTables(ctx context.Context, conn *sql.DB, database, tablePa } defer rows.Close() - var tables []db.TableSummary + tables := make([]db.TableSummary, 0) for rows.Next() { - var name, comment string - if err := rows.Scan(&name, &comment); err != nil { + var schemaName, name, comment string + if err := rows.Scan(&schemaName, &name, &comment); err != nil { return nil, errors.Wrap(errors.CodeDBExecFailed, "failed to scan table row", nil, err) } tables = append(tables, db.TableSummary{ - Schema: database, + Schema: schemaName, Name: name, Comment: comment, }) diff --git a/internal/db/pg/schema.go b/internal/db/pg/schema.go index a2ede3b6..000814ad 100644 --- a/internal/db/pg/schema.go +++ b/internal/db/pg/schema.go @@ -96,10 +96,14 @@ func (d *Driver) listSchemas(ctx context.Context, conn *sql.DB, includeSystem bo query := ` SELECT schema_name FROM information_schema.schemata - WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast') ` - if !includeSystem { - query += " AND schema_name NOT LIKE 'pg_%'" + if includeSystem { + query += " WHERE schema_name NOT LIKE 'pg_toast%'" + } else { + query += ` + WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast') + AND schema_name NOT LIKE 'pg_%' + ` } query += " ORDER BY schema_name" @@ -145,7 +149,7 @@ func (d *Driver) listTables(ctx context.Context, conn *sql.DB, schemas []string, } defer rows.Close() - var tables []db.TableSummary + tables := make([]db.TableSummary, 0) for rows.Next() { var schema, name string var comment sql.NullString diff --git a/internal/db/query.go b/internal/db/query.go index bba4c009..abea0a8e 100644 --- a/internal/db/query.go +++ b/internal/db/query.go @@ -12,6 +12,17 @@ import ( type QueryResult struct { Columns []string `json:"columns" yaml:"columns"` Rows []map[string]any `json:"rows" yaml:"rows"` + + Truncated bool `json:"truncated,omitempty" yaml:"truncated,omitempty"` + RowCount int `json:"row_count,omitempty" yaml:"row_count,omitempty"` + MaxRows int `json:"max_rows,omitempty" yaml:"max_rows,omitempty"` + // MaxRowsRequested is set when MCP/user requested more rows than mcp.max_rows_limit. + MaxRowsRequested int `json:"max_rows_requested,omitempty" yaml:"max_rows_requested,omitempty"` + // MaxRowsLimitApplied is the configured MCP ceiling when clamping occurred. + MaxRowsLimitApplied int `json:"max_rows_limit,omitempty" yaml:"max_rows_limit,omitempty"` + TruncatedReason string `json:"truncated_reason,omitempty" yaml:"truncated_reason,omitempty"` + MaxResponseBytes int `json:"max_response_bytes,omitempty" yaml:"max_response_bytes,omitempty"` + Hint string `json:"hint,omitempty" yaml:"hint,omitempty"` } // ToTableData implements the output.TableFormatter interface for table output without JSON encoding/decoding. @@ -22,12 +33,22 @@ func (r *QueryResult) ToTableData() (columns []string, rows []map[string]any, ok return r.Columns, r.Rows, true } +// QueryMeta reports truncation metadata for table output. +func (r *QueryResult) QueryMeta() (truncated bool, rowCount, maxRows int, hint string, ok bool) { + if r == nil || !r.Truncated { + return false, 0, 0, "", false + } + return true, r.RowCount, r.MaxRows, r.Hint, true +} + // QueryOptions contains options for query execution. type QueryOptions struct { UnsafeAllowWrite bool // Allow write operations (bypass read-only protection) DBType string // Database type: mysql or pg // BypassStaticReadOnly skips SQL static analysis (integration tests for DB-layer enforcement only). BypassStaticReadOnly bool + // MaxRows limits rows scanned (0 = unlimited). + MaxRows int } // Query executes a SQL query and returns the result. @@ -38,7 +59,7 @@ type QueryOptions struct { func Query(ctx context.Context, db *sql.DB, query string, opts QueryOptions) (*QueryResult, *errors.XError) { // UnsafeAllowWrite bypasses all read-only protections if opts.UnsafeAllowWrite { - return executeQuery(ctx, db, query, opts.DBType) + return executeQuery(ctx, db, query, opts.DBType, opts.MaxRows) } // Enable dual read-only protection by default @@ -48,7 +69,7 @@ func Query(ctx context.Context, db *sql.DB, query string, opts QueryOptions) (*Q } } // Second layer: database transaction-level read-only - return queryWithReadOnlyTx(ctx, db, query, opts.DBType) + return queryWithReadOnlyTx(ctx, db, query, opts.DBType, opts.MaxRows) } func normalizeDBType(dbType string) string { @@ -88,7 +109,7 @@ func beginReadOnlyTx(ctx context.Context, db *sql.DB, dbType string) (*sql.Tx, * } // queryWithReadOnlyTx executes a query within a read-only transaction. -func queryWithReadOnlyTx(ctx context.Context, db *sql.DB, query string, dbType string) (*QueryResult, *errors.XError) { +func queryWithReadOnlyTx(ctx context.Context, db *sql.DB, query string, dbType string, maxRows int) (*QueryResult, *errors.XError) { tx, xe := beginReadOnlyTx(ctx, db, dbType) if xe != nil { return nil, xe @@ -104,28 +125,28 @@ func queryWithReadOnlyTx(ctx context.Context, db *sql.DB, query string, dbType s } defer rows.Close() - return scanRows(rows) + return scanRows(rows, maxRows) } // executeQuery executes a query directly (without a transaction). -func executeQuery(ctx context.Context, db *sql.DB, query string, dbType string) (*QueryResult, *errors.XError) { +func executeQuery(ctx context.Context, db *sql.DB, query string, dbType string, maxRows int) (*QueryResult, *errors.XError) { rows, err := db.QueryContext(ctx, query) if err != nil { return nil, wrapQueryFailed(dbType, err) } defer rows.Close() - return scanRows(rows) + return scanRows(rows, maxRows) } -// scanRows scans query result rows. -func scanRows(rows *sql.Rows) (*QueryResult, *errors.XError) { +// scanRows scans query result rows, stopping after maxRows when maxRows > 0. +func scanRows(rows *sql.Rows, maxRows int) (*QueryResult, *errors.XError) { cols, err := rows.Columns() if err != nil { return nil, errors.Wrap(errors.CodeDBExecFailed, "failed to get columns", nil, err) } - result := &QueryResult{Columns: cols, Rows: []map[string]any{}} + result := &QueryResult{Columns: cols, Rows: make([]map[string]any, 0)} for rows.Next() { vals := make([]any, len(cols)) ptrs := make([]any, len(cols)) @@ -140,10 +161,21 @@ func scanRows(rows *sql.Rows) (*QueryResult, *errors.XError) { row[c] = convertValue(vals[i]) } result.Rows = append(result.Rows, row) + + if maxRows > 0 && len(result.Rows) == maxRows { + if rows.Next() { + result.Truncated = true + result.MaxRows = maxRows + result.TruncatedReason = TruncatedReasonMaxRows + RefreshTruncationHint(result) + } + break + } } if err := rows.Err(); err != nil { return nil, errors.Wrap(errors.CodeDBExecFailed, "rows iteration error", nil, err) } + result.RowCount = len(result.Rows) return result, nil } diff --git a/internal/db/query_limit_test.go b/internal/db/query_limit_test.go new file mode 100644 index 00000000..db909613 --- /dev/null +++ b/internal/db/query_limit_test.go @@ -0,0 +1,88 @@ +package db + +import ( + "testing" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestScanRows_MaxRowsTruncated(t *testing.T) { + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer mockDB.Close() + + rows := sqlmock.NewRows([]string{"n"}). + AddRow(1). + AddRow(2). + AddRow(3) + mock.ExpectQuery("SELECT").WillReturnRows(rows) + + sqlRows, err := mockDB.Query("SELECT n FROM t") + if err != nil { + t.Fatal(err) + } + defer sqlRows.Close() + + result, xe := scanRows(sqlRows, 2) + if xe != nil { + t.Fatalf("scanRows: %v", xe) + } + if len(result.Rows) != 2 || !result.Truncated { + t.Fatalf("expected 2 rows truncated, got %d truncated=%v", len(result.Rows), result.Truncated) + } + if result.RowCount != 2 || result.MaxRows != 2 { + t.Fatalf("meta row_count=%d max_rows=%d", result.RowCount, result.MaxRows) + } +} + +func TestScanRows_MaxRowsExact(t *testing.T) { + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer mockDB.Close() + + rows := sqlmock.NewRows([]string{"n"}).AddRow(1).AddRow(2) + mock.ExpectQuery("SELECT").WillReturnRows(rows) + + sqlRows, err := mockDB.Query("SELECT n FROM t") + if err != nil { + t.Fatal(err) + } + defer sqlRows.Close() + + result, xe := scanRows(sqlRows, 2) + if xe != nil { + t.Fatalf("scanRows: %v", xe) + } + if len(result.Rows) != 2 || result.Truncated { + t.Fatalf("expected 2 rows not truncated, got truncated=%v", result.Truncated) + } +} + +func TestScanRows_Unlimited(t *testing.T) { + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer mockDB.Close() + + rows := sqlmock.NewRows([]string{"n"}).AddRow(1).AddRow(2).AddRow(3) + mock.ExpectQuery("SELECT").WillReturnRows(rows) + + sqlRows, err := mockDB.Query("SELECT n FROM t") + if err != nil { + t.Fatal(err) + } + defer sqlRows.Close() + + result, xe := scanRows(sqlRows, 0) + if xe != nil { + t.Fatalf("scanRows: %v", xe) + } + if len(result.Rows) != 3 || result.Truncated { + t.Fatalf("expected 3 rows, got %d truncated=%v", len(result.Rows), result.Truncated) + } +} diff --git a/internal/db/query_readonly_test.go b/internal/db/query_readonly_test.go index 60b2fd5d..e0417fc9 100644 --- a/internal/db/query_readonly_test.go +++ b/internal/db/query_readonly_test.go @@ -209,7 +209,7 @@ func TestQueryWithReadOnlyTx_QueryFailure(t *testing.T) { mock.ExpectQuery("SELECT 1").WillReturnError(errors.New("query failed")) mock.ExpectRollback() - _, xe := queryWithReadOnlyTx(context.Background(), db, "SELECT 1", "mysql") + _, xe := queryWithReadOnlyTx(context.Background(), db, "SELECT 1", "mysql", 0) if xe == nil || xe.Code != xsqlerrors.CodeDBExecFailed { t.Fatalf("expected query failure, got %v", xe) } @@ -225,7 +225,7 @@ func TestExecuteQuery_Success(t *testing.T) { rows := sqlmock.NewRows([]string{"msg"}).AddRow([]byte("hi")) mock.ExpectQuery("SELECT 'hi'").WillReturnRows(rows) - result, xe := executeQuery(context.Background(), db, "SELECT 'hi'", "mysql") + result, xe := executeQuery(context.Background(), db, "SELECT 'hi'", "mysql", 0) if xe != nil { t.Fatalf("executeQuery failed: %v", xe) } @@ -243,7 +243,7 @@ func TestExecuteQuery_Failure(t *testing.T) { mock.ExpectQuery("SELECT 1").WillReturnError(errors.New("boom")) - _, xe := executeQuery(context.Background(), db, "SELECT 1", "mysql") + _, xe := executeQuery(context.Background(), db, "SELECT 1", "mysql", 0) if xe == nil || xe.Code != xsqlerrors.CodeDBExecFailed { t.Fatalf("expected failure, got %v", xe) } @@ -265,7 +265,7 @@ func TestScanRows_MultipleRows(t *testing.T) { } defer rows.Close() - result, xe := scanRows(rows) + result, xe := scanRows(rows, 0) if xe != nil { t.Fatalf("scanRows failed: %v", xe) } diff --git a/internal/db/query_test.go b/internal/db/query_test.go index ac3aad42..39f20062 100644 --- a/internal/db/query_test.go +++ b/internal/db/query_test.go @@ -29,6 +29,17 @@ func TestConvertValue(t *testing.T) { // Query 函数的集成测试在 tests/integration/query_test.go 中 +func TestQueryResult_QueryMeta(t *testing.T) { + if _, _, _, _, ok := (&QueryResult{}).QueryMeta(); ok { + t.Fatal("non-truncated should not report meta") + } + tr := &QueryResult{Truncated: true, RowCount: 2, MaxRows: 2, Hint: "h"} + truncated, rowCount, maxRows, hint, ok := tr.QueryMeta() + if !ok || !truncated || rowCount != 2 || maxRows != 2 || hint != "h" { + t.Fatalf("meta=%v %d %d %q ok=%v", truncated, rowCount, maxRows, hint, ok) + } +} + func TestQueryResultToTableData(t *testing.T) { var result *QueryResult cols, rows, ok := result.ToTableData() diff --git a/internal/db/schema.go b/internal/db/schema.go index aa3054b2..d7867f5a 100644 --- a/internal/db/schema.go +++ b/internal/db/schema.go @@ -126,7 +126,14 @@ func ListTables(ctx context.Context, driverName string, db *sql.DB, opts SchemaO return nil, errors.New(errors.CodeDBDriverUnsupported, "driver does not support schema browsing: "+driverName, nil) } - return sd.ListTables(ctx, db, opts) + result, xe := sd.ListTables(ctx, db, opts) + if xe != nil { + return nil, xe + } + if result != nil && result.Tables == nil { + result.Tables = []TableSummary{} + } + return result, nil } // DescribeTable returns the schema details for a single table. diff --git a/internal/db/schema_test.go b/internal/db/schema_test.go index 004a85b3..95284b04 100644 --- a/internal/db/schema_test.go +++ b/internal/db/schema_test.go @@ -305,6 +305,20 @@ func TestDumpSchema_UnsupportedDriver(t *testing.T) { } } +func TestListTables_NormalizesNilTables(t *testing.T) { + driverName := fmt.Sprintf("schema-nil-tables-%s", t.Name()) + Register(driverName, &mockSchemaExplorerDriver{ + tableList: &TableList{Database: "app", Tables: nil}, + }) + list, xe := ListTables(context.Background(), driverName, &sql.DB{}, SchemaOptions{}) + if xe != nil { + t.Fatalf("ListTables: %v", xe) + } + if list.Tables == nil || len(list.Tables) != 0 { + t.Fatalf("expected non-nil empty slice, got %#v", list.Tables) + } +} + func TestListTables_UnsupportedDriver(t *testing.T) { _, xe := ListTables(nil, "nonexistent", nil, SchemaOptions{}) if xe == nil { diff --git a/internal/db/truncation.go b/internal/db/truncation.go new file mode 100644 index 00000000..c4ff514a --- /dev/null +++ b/internal/db/truncation.go @@ -0,0 +1,34 @@ +package db + +import ( + "strings" + + "github.com/zx06/xsql/internal/errors" +) + +const ( + TruncatedReasonMaxRows = "max_rows" + TruncatedReasonMaxResponseBytes = "max_response_bytes" +) + +// AppendTruncatedReason merges a truncation reason token into any existing value. +func AppendTruncatedReason(existing, reason string) string { + if reason == "" { + return existing + } + if existing == "" { + return reason + } + if strings.Contains(existing, reason) { + return existing + } + return existing + "," + reason +} + +// RefreshTruncationHint sets Hint from TruncatedReason when truncated. +func RefreshTruncationHint(r *QueryResult) { + if r == nil || !r.Truncated { + return + } + r.Hint = errors.BuildTruncationHint(r.TruncatedReason) +} diff --git a/internal/errors/hints.go b/internal/errors/hints.go index 35159cce..f986785d 100644 --- a/internal/errors/hints.go +++ b/internal/errors/hints.go @@ -1,22 +1,47 @@ package errors +import "strings" + // Shared agent hint strings (single source for Sonar duplicate-literal rules). const ( - hintUnknownTableMySQL = "Table or database not found; use MCP `schema_list_tables` or `xsql schema dump` to list valid tables, then fix the SQL." - hintUnknownTablePG = "Table or schema not found; use MCP `schema_list_tables` or `xsql schema dump` to list tables, then fix the SQL." - hintUnknownColumn = "Column not found; run `xsql schema dump --table ` to inspect columns, then fix the SQL." - hintSyntaxMySQL = "Fix SQL syntax; compare with the profile dialect (mysql) and re-run." - hintSyntaxPG = "Fix SQL syntax; compare with the profile dialect (pg) and re-run." - hintPermissionDenied = "Insufficient database privileges; use a read-only SELECT or ask for a profile with required grants." - hintAuthFailed = "Authentication failed; verify profile username/password or keyring secret reference for this profile." - hintReadOnlyWrite = "Write blocked in read-only mode; set profile `unsafe_allow_write: true` AND pass `--unsafe-allow-write` on query/ai/mcp server." - hintReadOnlyWritePG = "Write blocked in read-only transaction; set profile `unsafe_allow_write: true` AND pass `--unsafe-allow-write` on query/ai/mcp server." - hintDuplicateKey = "Duplicate key conflict; adjust INSERT/UPDATE values or use ON DUPLICATE handling if writes are explicitly allowed." - hintDuplicateKeyPG = "Duplicate key conflict; adjust values or constraints if writes are explicitly allowed." - hintMySQLConnection = "Cannot reach MySQL; verify profile host/port, network, and ssh_proxy settings." - hintPGConnection = "Cannot reach PostgreSQL; verify profile host/port, sslmode, and ssh_proxy settings." - hintQueryTimeoutMySQL = "Query timed out or lock wait exceeded; narrow the query, add LIMIT, or retry later." - hintQueryTimeoutPG = "Query canceled or lock timeout; narrow the query, add LIMIT, or retry later." - hintSerializationPG = "Serialization or deadlock failure; retry the read query or simplify concurrent access." - hintDBOther = "Database returned an error; inspect native_code and message, adjust SQL or connection settings." + hintUnknownTableMySQL = "Table or database not found; use MCP `schema_list_tables` or `xsql schema dump` to list valid tables, then fix the SQL." + hintUnknownTablePG = "Table or schema not found; use MCP `schema_list_tables` or `xsql schema dump` to list tables, then fix the SQL." + hintUnknownColumn = "Column not found; run `xsql schema dump --table ` to inspect columns, then fix the SQL." + hintSyntaxMySQL = "Fix SQL syntax; compare with the profile dialect (mysql) and re-run." + hintSyntaxPG = "Fix SQL syntax; compare with the profile dialect (pg) and re-run." + hintPermissionDenied = "Insufficient database privileges; use a read-only SELECT or ask for a profile with required grants." + hintAuthFailed = "Authentication failed; verify profile username/password or keyring secret reference for this profile." + hintReadOnlyWrite = "Write blocked in read-only mode; set profile `unsafe_allow_write: true` AND pass `--unsafe-allow-write` on query/ai/mcp server." + hintReadOnlyWritePG = "Write blocked in read-only transaction; set profile `unsafe_allow_write: true` AND pass `--unsafe-allow-write` on query/ai/mcp server." + hintDuplicateKey = "Duplicate key conflict; adjust INSERT/UPDATE values or use ON DUPLICATE handling if writes are explicitly allowed." + hintDuplicateKeyPG = "Duplicate key conflict; adjust values or constraints if writes are explicitly allowed." + hintMySQLConnection = "Cannot reach MySQL; verify profile host/port, network, and ssh_proxy settings." + hintPGConnection = "Cannot reach PostgreSQL; verify profile host/port, sslmode, and ssh_proxy settings." + hintQueryTimeoutMySQL = "Query timed out or lock wait exceeded; narrow the query, add LIMIT, or retry later." + hintQueryTimeoutPG = "Query canceled or lock timeout; narrow the query, add LIMIT, or retry later." + hintSerializationPG = "Serialization or deadlock failure; retry the read query or simplify concurrent access." + hintDBOther = "Database returned an error; inspect native_code and message, adjust SQL or connection settings." + hintQueryResultTruncated = "Result truncated by max_rows; add SQL LIMIT/OFFSET, use aggregation, or run MCP `schema_list_tables` or `xsql schema dump` for schema discovery." + hintMCPResponseTruncated = "Result truncated to fit MCP max_response_bytes; narrow the query with LIMIT or filters." ) + +// HintQueryResultTruncated is the agent hint when query rows are capped (success data.hint). +const HintQueryResultTruncated = hintQueryResultTruncated + +// HintMCPResponseTruncated is appended when MCP byte guard trims rows. +const HintMCPResponseTruncated = hintMCPResponseTruncated + +// BuildTruncationHint composes agent hints from a comma-separated truncated_reason value. +func BuildTruncationHint(truncatedReason string) string { + var parts []string + if truncatedReason == "" || strings.Contains(truncatedReason, "max_rows") { + parts = append(parts, HintQueryResultTruncated) + } + if strings.Contains(truncatedReason, "max_response_bytes") { + parts = append(parts, HintMCPResponseTruncated) + } + if len(parts) == 0 { + return HintQueryResultTruncated + } + return strings.Join(parts, " ") +} diff --git a/internal/errors/hints_test.go b/internal/errors/hints_test.go new file mode 100644 index 00000000..5987bc67 --- /dev/null +++ b/internal/errors/hints_test.go @@ -0,0 +1,19 @@ +package errors + +import ( + "strings" + "testing" +) + +func TestBuildTruncationHint(t *testing.T) { + if got := BuildTruncationHint("max_rows"); !strings.Contains(got, HintQueryResultTruncated) { + t.Fatalf("got %q", got) + } + if got := BuildTruncationHint("max_response_bytes"); !strings.Contains(got, HintMCPResponseTruncated) { + t.Fatalf("got %q", got) + } + got := BuildTruncationHint("max_rows,max_response_bytes") + if !strings.Contains(got, HintQueryResultTruncated) || !strings.Contains(got, HintMCPResponseTruncated) { + t.Fatalf("got %q", got) + } +} diff --git a/internal/mcp/query_result.go b/internal/mcp/query_result.go new file mode 100644 index 00000000..aae7f2f8 --- /dev/null +++ b/internal/mcp/query_result.go @@ -0,0 +1,73 @@ +package mcp + +import ( + "encoding/json" + + "github.com/zx06/xsql/internal/db" +) + +func applyMCPResponseByteLimit(result *db.QueryResult, maxBytes int, effectiveMaxRows int) *db.QueryResult { + if result == nil || maxBytes <= 0 { + return result + } + if effectiveMaxRows > 0 && result.MaxRows == 0 { + result.MaxRows = effectiveMaxRows + } + if len(result.Rows) == 0 { + return result + } + if mcpPayloadSize(result) <= maxBytes { + return result + } + + keep := binarySearchRowFit(result, maxBytes) + if keep >= len(result.Rows) { + return result + } + + result.Rows = result.Rows[:keep] + result.Truncated = true + result.RowCount = keep + result.MaxRows = effectiveMaxRows + result.MaxResponseBytes = maxBytes + result.TruncatedReason = db.AppendTruncatedReason(result.TruncatedReason, db.TruncatedReasonMaxResponseBytes) + db.RefreshTruncationHint(result) + return result +} + +func mcpPayloadSize(result *db.QueryResult) int { + payload, err := json.Marshal(map[string]any{ + "ok": true, + "schema_version": 1, + "data": result, + }) + if err != nil { + return 0 + } + return len(payload) +} + +func binarySearchRowFit(result *db.QueryResult, maxBytes int) int { + low, high := 0, len(result.Rows) + for low < high { + mid := (low + high + 1) / 2 + trial := &db.QueryResult{ + Columns: result.Columns, + Rows: result.Rows[:mid], + Truncated: result.Truncated, + RowCount: mid, + MaxRows: result.MaxRows, + MaxRowsRequested: result.MaxRowsRequested, + MaxRowsLimitApplied: result.MaxRowsLimitApplied, + TruncatedReason: result.TruncatedReason, + MaxResponseBytes: maxBytes, + Hint: result.Hint, + } + if mcpPayloadSize(trial) <= maxBytes { + low = mid + } else { + high = mid - 1 + } + } + return low +} diff --git a/internal/mcp/query_result_test.go b/internal/mcp/query_result_test.go new file mode 100644 index 00000000..1a7643b8 --- /dev/null +++ b/internal/mcp/query_result_test.go @@ -0,0 +1,74 @@ +package mcp + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" +) + +func TestApplyMCPResponseByteLimit_Noop(t *testing.T) { + result := &db.QueryResult{Columns: []string{"n"}, Rows: []map[string]any{{"n": 1}}} + out := applyMCPResponseByteLimit(result, 0, 500) + if out != result || out.Truncated { + t.Fatalf("expected noop, got %+v", out) + } +} + +func TestApplyMCPResponseByteLimit(t *testing.T) { + rows := make([]map[string]any, 0, 50) + for i := 0; i < 50; i++ { + rows = append(rows, map[string]any{"payload": string(make([]byte, 200))}) + } + result := &db.QueryResult{ + Columns: []string{"payload"}, + Rows: rows, + MaxRows: 500, + } + out := applyMCPResponseByteLimit(result, 2048, 500) + payload, _ := json.Marshal(map[string]any{"ok": true, "schema_version": 1, "data": out}) + if len(payload) > 2048 { + t.Fatalf("expected payload under 2048 bytes, got %d", len(payload)) + } + if !out.Truncated { + t.Fatal("expected truncated") + } + if out.MaxRows != 500 { + t.Fatalf("max_rows=%d want configured 500", out.MaxRows) + } + if !strings.Contains(out.TruncatedReason, db.TruncatedReasonMaxResponseBytes) { + t.Fatalf("reason=%q", out.TruncatedReason) + } + if out.MaxResponseBytes != 2048 { + t.Fatalf("max_response_bytes=%d", out.MaxResponseBytes) + } + if !strings.Contains(out.Hint, errors.HintMCPResponseTruncated) { + t.Fatalf("hint=%q", out.Hint) + } +} + +func TestApplyMCPResponseByteLimit_BothReasons(t *testing.T) { + rows := make([]map[string]any, 0, 20) + for i := 0; i < 20; i++ { + rows = append(rows, map[string]any{"payload": string(make([]byte, 100))}) + } + result := &db.QueryResult{ + Columns: []string{"payload"}, + Rows: rows, + Truncated: true, + TruncatedReason: db.TruncatedReasonMaxRows, + MaxRows: 20, + Hint: errors.HintQueryResultTruncated, + } + out := applyMCPResponseByteLimit(result, 1024, 20) + if !strings.Contains(out.TruncatedReason, db.TruncatedReasonMaxRows) || + !strings.Contains(out.TruncatedReason, db.TruncatedReasonMaxResponseBytes) { + t.Fatalf("reason=%q", out.TruncatedReason) + } + if !strings.Contains(out.Hint, errors.HintQueryResultTruncated) || + !strings.Contains(out.Hint, errors.HintMCPResponseTruncated) { + t.Fatalf("hint=%q", out.Hint) + } +} diff --git a/internal/mcp/schema_tools_test.go b/internal/mcp/schema_tools_test.go index 5cbd5ebd..96e74702 100644 --- a/internal/mcp/schema_tools_test.go +++ b/internal/mcp/schema_tools_test.go @@ -418,6 +418,31 @@ func assertToolError(t *testing.T, result *mcp.CallToolResult, code, messageCont } } +func TestSchemaListTables_EmptyDatabaseTablesArray(t *testing.T) { + prev := schemaApp.ListTables + t.Cleanup(func() { schemaApp.ListTables = prev }) + schemaApp.ListTables = func(context.Context, app.TableListRequest) (*db.TableList, *errors.XError) { + return &db.TableList{Database: "empty", Tables: []db.TableSummary{}}, nil + } + cfg := &config.File{ + Profiles: map[string]config.Profile{"dev": {DB: "mysql"}}, + } + handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) + result, _, err := handler.SchemaListTables(context.Background(), &mcp.CallToolRequest{}, SchemaListTablesInput{Profile: "dev"}) + if err != nil { + t.Fatal(err) + } + assertToolSuccess(t, result, func(data map[string]any) { + raw, ok := data["tables"].([]any) + if !ok { + t.Fatalf("tables should be array, got %T", data["tables"]) + } + if raw == nil || len(raw) != 0 { + t.Fatalf("expected empty tables array, got %v", data["tables"]) + } + }) +} + func assertToolSuccess(t *testing.T, result *mcp.CallToolResult, check func(map[string]any)) { t.Helper() if result == nil || result.IsError { diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go index 94426391..22c9fdf5 100644 --- a/internal/mcp/tools.go +++ b/internal/mcp/tools.go @@ -20,6 +20,7 @@ const defaultMCPQueryTimeout = 30 * time.Second type QueryInput struct { SQL string `json:"sql" jsonschema:"SQL query to execute"` Profile string `json:"profile" jsonschema:"Profile name to use"` + MaxRows *int `json:"max_rows,omitempty" jsonschema:"Maximum rows to return (0 = default; clamped by mcp.max_rows_limit; MCP default 500)"` } // ProfileShowInput represents the input for the profile_show tool @@ -89,6 +90,10 @@ func (h *ToolHandler) RegisterTools(server *mcp.Server) { Description: "Profile name to use", Enum: profileEnums, }, + "max_rows": { + Type: "integer", + Description: "Maximum rows to return (0 = use defaults; clamped by mcp.max_rows_limit; default 500)", + }, }, } server.AddTool(&mcp.Tool{ @@ -158,6 +163,16 @@ func (h *ToolHandler) Query(ctx context.Context, req *mcp.CallToolRequest, input return h.errorResult(xe), nil, nil } + argSet := input.MaxRows != nil + argVal := 0 + if argSet { + argVal = *input.MaxRows + } + maxRows, rowsMeta, xe := app.ResolveMCPQueryMaxRows(profile, h.config.MCP, argVal, argSet) + if xe != nil { + return h.errorResult(xe), nil, nil + } + timeout := app.QueryTimeout(profile, 0, false, defaultMCPQueryTimeout) ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() @@ -169,12 +184,22 @@ func (h *ToolHandler) Query(ctx context.Context, req *mcp.CallToolRequest, input AllowPlaintext: profile.AllowPlaintext, SkipHostKeyCheck: false, UnsafeAllowWrite: WriteAllowed(h.serverAllowsWrite, profile.UnsafeAllowWrite), + MaxRows: maxRows, + TruncationHint: app.QueryTruncationHint, }) h.recordMCPStats("query", input.Profile, xe == nil, time.Since(start), xe, input.SQL) if xe != nil { return h.errorResult(xe), nil, nil } + if rowsMeta.Clamped { + result.MaxRowsRequested = rowsMeta.Requested + result.MaxRowsLimitApplied = rowsMeta.Limit + } + if result.MaxRows == 0 && maxRows > 0 { + result.MaxRows = maxRows + } + result = applyMCPResponseByteLimit(result, app.ResolveMCPMaxResponseBytes(h.config.MCP), maxRows) return h.successResult(result) } diff --git a/internal/mcp/tools_test.go b/internal/mcp/tools_test.go index 59abe621..43c29fa0 100644 --- a/internal/mcp/tools_test.go +++ b/internal/mcp/tools_test.go @@ -308,6 +308,25 @@ func TestQuery_MissingSQL(t *testing.T) { } } +func TestQuery_MaxRowsNegativeRejected(t *testing.T) { + cfg := &config.File{ + Profiles: map[string]config.Profile{"dev": {DB: "mysql"}}, + } + handler := NewToolHandler(cfg, stats.StatsConfig{}, ServerOptions{}) + neg := -1 + result, _, err := handler.Query(context.Background(), &mcp.CallToolRequest{}, QueryInput{ + SQL: "SELECT 1", + Profile: "dev", + MaxRows: &neg, + }) + if err != nil { + t.Fatal(err) + } + if !result.IsError { + t.Fatal("expected error for negative max_rows on MCP") + } +} + func TestQuery_MissingProfile(t *testing.T) { cfg := &config.File{ Profiles: map[string]config.Profile{ diff --git a/internal/output/writer.go b/internal/output/writer.go index df15fc15..888dee55 100644 --- a/internal/output/writer.go +++ b/internal/output/writer.go @@ -24,8 +24,34 @@ func New(out, err io.Writer) Writer { return Writer{Out: out, Err: err} } +// CSVTruncationReporter is implemented by query results that may truncate row output. +type CSVTruncationReporter interface { + CSVTruncationNotice() (string, bool) +} + func (w Writer) WriteOK(format Format, data any) error { - return w.write(format, Envelope{OK: true, SchemaVersion: SchemaVersion, Data: data}) + if err := w.write(format, Envelope{OK: true, SchemaVersion: SchemaVersion, Data: data}); err != nil { + return err + } + if format == FormatCSV { + writeCSVTruncationNotice(w.Err, data) + } + return nil +} + +func writeCSVTruncationNotice(errOut io.Writer, data any) { + if errOut == nil { + return + } + rep, ok := data.(CSVTruncationReporter) + if !ok { + return + } + msg, ok := rep.CSVTruncationNotice() + if !ok || msg == "" { + return + } + _, _ = fmt.Fprintln(errOut, msg) } func (w Writer) WriteError(format Format, xe *errors.XError) error { @@ -66,6 +92,11 @@ type TableFormatter interface { ToTableData() (columns []string, rows []map[string]any, ok bool) } +// QueryMetaFormatter exposes truncation metadata for query table output. +type QueryMetaFormatter interface { + QueryMeta() (truncated bool, rowCount, maxRows int, hint string, ok bool) +} + // SchemaFormatter is the interface for data structures that support schema output. type SchemaFormatter interface { ToSchemaData() (database string, tables []SchemaTable, ok bool) @@ -96,7 +127,18 @@ func writeTable(out io.Writer, env Envelope) error { // First check if the data implements the TableFormatter interface (no JSON encode/decode) if formatter, ok := env.Data.(TableFormatter); ok { if cols, rows, ok := formatter.ToTableData(); ok { - return writeQueryResultTable(out, cols, rows) + if err := writeQueryResultTable(out, cols, rows); err != nil { + return err + } + if meta, ok := env.Data.(QueryMetaFormatter); ok { + if truncated, rowCount, maxRows, hint, ok := meta.QueryMeta(); ok && truncated { + _, _ = fmt.Fprintf(out, "(truncated at max_rows=%d; showing %d rows)\n", maxRows, rowCount) + if hint != "" { + _, _ = fmt.Fprintf(out, "hint: %s\n", hint) + } + } + } + return nil } } diff --git a/internal/output/writer_test.go b/internal/output/writer_test.go index 8b27d3c8..dcd26dbe 100644 --- a/internal/output/writer_test.go +++ b/internal/output/writer_test.go @@ -119,6 +119,72 @@ func TestWriteOK_TableFormat_QueryResult(t *testing.T) { } } +type truncatedQueryTable struct { + cols []string + rows []map[string]any +} + +func (truncatedQueryTable) QueryMeta() (truncated bool, rowCount, maxRows int, hint string, ok bool) { + return true, 1, 1, "use LIMIT", true +} + +func (t truncatedQueryTable) ToTableData() ([]string, []map[string]any, bool) { + return t.cols, t.rows, true +} + +type csvTruncStub struct { + cols []string + rows []map[string]any + msg string +} + +func (c csvTruncStub) ToTableData() ([]string, []map[string]any, bool) { + return c.cols, c.rows, true +} + +func (c csvTruncStub) CSVTruncationNotice() (string, bool) { + return c.msg, c.msg != "" +} + +func TestWriteOK_CSVTruncationNoticeOnStderr(t *testing.T) { + var out, errOut bytes.Buffer + w := New(&out, &errOut) + data := csvTruncStub{ + cols: []string{"n"}, + rows: []map[string]any{{"n": 1}}, + msg: "xsql: query result truncated (max_rows); hint text", + } + if err := w.WriteOK(FormatCSV, data); err != nil { + t.Fatal(err) + } + if strings.Contains(out.String(), "truncated") { + t.Fatal("stdout must stay pure CSV") + } + if !strings.Contains(errOut.String(), "truncated") { + t.Fatalf("stderr notice missing: %q", errOut.String()) + } +} + +func TestWriteOK_TableFormat_TruncatedQueryMeta(t *testing.T) { + var out bytes.Buffer + w := New(&out, &bytes.Buffer{}) + + result := truncatedQueryTable{ + cols: []string{"n"}, + rows: []map[string]any{{"n": 1}}, + } + if err := w.WriteOK(FormatTable, result); err != nil { + t.Fatal(err) + } + text := out.String() + if !strings.Contains(text, "(truncated at max_rows=1") { + t.Fatalf("missing truncation footer: %s", text) + } + if !strings.Contains(text, "hint: use LIMIT") { + t.Fatalf("missing hint line: %s", text) + } +} + func TestWriteOK_TableFormat_NonQueryResult(t *testing.T) { var out bytes.Buffer w := New(&out, &bytes.Buffer{}) diff --git a/internal/tui/model.go b/internal/tui/model.go index 9323033f..14bfe311 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -242,12 +242,18 @@ func (m Model) executeSQLCmd(sqlStr string) tea.Cmd { timeout := app.QueryTimeout(m.profile, 0, false, defaultTUIQueryTimeout) ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() + maxRows, mxe := app.ResolveQueryMaxRows(m.profile, "", 0, false) + if mxe != nil { + return queryExecutedMsg{err: mxe} + } res, xe := app.Query(ctx, app.QueryRequest{ Profile: m.profile, SQL: sqlStr, AllowPlaintext: m.profile.AllowPlaintext, SkipHostKeyCheck: m.profile.SSHConfig != nil && m.profile.SSHConfig.SkipHostKey, UnsafeAllowWrite: m.unsafeAllowWrite, + MaxRows: maxRows, + TruncationHint: app.QueryTruncationHint, }) elapsed := time.Since(start) return queryExecutedMsg{result: res, err: xe, duration: elapsed} diff --git a/internal/web/handler.go b/internal/web/handler.go index 7671f18c..cf9bb6ef 100644 --- a/internal/web/handler.go +++ b/internal/web/handler.go @@ -10,6 +10,7 @@ import ( "net" "net/http" "net/url" + "os" "path" "path/filepath" "strconv" @@ -50,6 +51,7 @@ type handler struct { type queryRequest struct { Profile string `json:"profile"` SQL string `json:"sql"` + MaxRows *int `json:"max_rows,omitempty"` } // NewHandler creates the web server handler. @@ -257,6 +259,22 @@ func (h *handler) handleQuery(w http.ResponseWriter, r *http.Request) { return } + var maxRows int + if req.MaxRows != nil { + if *req.MaxRows < 0 { + writeError(w, http.StatusBadRequest, errors.New(errors.CodeCfgInvalid, "max_rows must be >= 0", map[string]any{"max_rows": *req.MaxRows})) + return + } + maxRows = *req.MaxRows + } else { + var mxe *errors.XError + maxRows, mxe = app.ResolveQueryMaxRows(profile, os.Getenv("XSQL_MAX_ROWS"), 0, false) + if mxe != nil { + writeError(w, http.StatusBadRequest, mxe) + return + } + } + timeout := app.QueryTimeout(profile, 0, false, 30*time.Second) ctx, cancel := context.WithTimeout(r.Context(), timeout) defer cancel() @@ -268,6 +286,8 @@ func (h *handler) handleQuery(w http.ResponseWriter, r *http.Request) { AllowPlaintext: h.allowPlaintext, SkipHostKeyCheck: h.skipHostKeyCheck, UnsafeAllowWrite: false, + MaxRows: maxRows, + TruncationHint: app.QueryTruncationHint, }) // Record stats diff --git a/internal/web/handler_query_limits_test.go b/internal/web/handler_query_limits_test.go new file mode 100644 index 00000000..68b3d5ac --- /dev/null +++ b/internal/web/handler_query_limits_test.go @@ -0,0 +1,92 @@ +package web + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "testing/fstest" +) + +func TestHandler_HandleQuery_MaxRowsNegative(t *testing.T) { + configPath := createConfigFile(t, ` +profiles: + dev: + db: mysql + dsn: "root:root@tcp(127.0.0.1:3306)/testdb" +`) + handler := NewHandler(HandlerOptions{ + ConfigPath: configPath, + Assets: fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("ok")}}, + }) + + body := []byte(`{"profile":"dev","sql":"SELECT 1","max_rows":-1}`) + req := httptest.NewRequest(http.MethodPost, "/api/v1/query", strings.NewReader(string(body))) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandler_HandleQuery_MaxRowsExplicit(t *testing.T) { + dsn := os.Getenv("XSQL_TEST_MYSQL_DSN") + if dsn == "" { + t.Skip("XSQL_TEST_MYSQL_DSN not set") + } + configPath := createConfigFile(t, ` +profiles: + dev: + db: mysql + dsn: "`+dsn+`" +`) + handler := NewHandler(HandlerOptions{ + ConfigPath: configPath, + Assets: fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("ok")}}, + }) + + body := []byte(`{"profile":"dev","sql":"WITH RECURSIVE seq AS (SELECT 1 AS n UNION ALL SELECT n+1 FROM seq WHERE n < 10) SELECT n FROM seq","max_rows":2}`) + req := httptest.NewRequest(http.MethodPost, "/api/v1/query", strings.NewReader(string(body))) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp envelope + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if !resp.OK { + t.Fatalf("expected ok: %+v", resp) + } +} + +func TestHandler_HandleQuery_InvalidEnvMaxRows(t *testing.T) { + t.Setenv("XSQL_MAX_ROWS", "not-a-number") + configPath := createConfigFile(t, ` +profiles: + dev: + db: mysql + dsn: "root:root@tcp(127.0.0.1:3306)/testdb" +`) + handler := NewHandler(HandlerOptions{ + ConfigPath: configPath, + Assets: fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("ok")}}, + }) + + body := []byte(`{"profile":"dev","sql":"SELECT 1"}`) + req := httptest.NewRequest(http.MethodPost, "/api/v1/query", strings.NewReader(string(body))) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", rec.Code) + } +} diff --git a/tests/e2e/query_limits_test.go b/tests/e2e/query_limits_test.go new file mode 100644 index 00000000..3b002954 --- /dev/null +++ b/tests/e2e/query_limits_test.go @@ -0,0 +1,96 @@ +//go:build e2e + +package e2e + +import ( + "encoding/json" + "testing" +) + +type queryData struct { + Columns []string `json:"columns"` + Rows []map[string]any `json:"rows"` + Truncated bool `json:"truncated"` + RowCount int `json:"row_count"` + MaxRows int `json:"max_rows"` + Hint string `json:"hint"` +} + +func TestQuery_MaxRows_DefaultTruncation_MySQL(t *testing.T) { + dsn := mysqlDSN(t) + config := mysqlTestConfig(t, dsn) + sql := `WITH RECURSIVE seq AS (SELECT 1 AS n UNION ALL SELECT n+1 FROM seq WHERE n < 30) SELECT n FROM seq` + data := runQueryJSON(t, config, sql, "--max-rows", "5") + if len(data.Rows) != 5 || !data.Truncated || data.MaxRows != 5 { + t.Fatalf("unexpected data: %+v", data) + } +} + +func TestQuery_MaxRows_Unlimited_MySQL(t *testing.T) { + dsn := mysqlDSN(t) + config := mysqlTestConfig(t, dsn) + sql := `WITH RECURSIVE seq AS (SELECT 1 AS n UNION ALL SELECT n+1 FROM seq WHERE n < 4) SELECT n FROM seq` + data := runQueryJSON(t, config, sql, "--max-rows", "0") + if len(data.Rows) != 4 || data.Truncated { + t.Fatalf("unexpected data: %+v", data) + } +} + +func TestQuery_MaxRows_InvalidNegative(t *testing.T) { + dsn := mysqlDSN(t) + config := mysqlTestConfig(t, dsn) + _, _, exitCode := runXSQL(t, "query", "SELECT 1", "--config", config, "--profile", "test", "--format", "json", "--max-rows", "-1") + if exitCode != 2 { + t.Fatalf("expected exit 2, got %d", exitCode) + } +} + +func TestQuery_MaxRows_PG(t *testing.T) { + dsn := pgDSN(t) + config := createTempConfig(t, `profiles: + test: + db: pg + dsn: "`+dsn+`" +`) + data := runQueryJSON(t, config, "SELECT generate_series(1,20) AS n", "--max-rows", "3") + if len(data.Rows) != 3 || !data.Truncated { + t.Fatalf("unexpected data: %+v", data) + } +} + +func mysqlTestConfig(t *testing.T, dsn string) string { + t.Helper() + return createTempConfig(t, `profiles: + test: + db: mysql + dsn: "`+dsn+`" +`) +} + +func runQueryJSON(t *testing.T, config, sql string, extraArgs ...string) queryData { + t.Helper() + args := append([]string{"query", sql, "--config", config, "--profile", "test", "--format", "json"}, extraArgs...) + stdout, _, exitCode := runXSQL(t, args...) + if exitCode != 0 { + t.Fatalf("exit %d stdout=%s", exitCode, stdout) + } + var resp Response + if err := json.Unmarshal([]byte(stdout), &resp); err != nil { + t.Fatal(err) + } + data, err := decodeQueryData(resp.Data) + if err != nil { + t.Fatal(err) + } + return data +} + +func decodeQueryData(raw any) (queryData, error) { + b, err := json.Marshal(raw) + if err != nil { + return queryData{}, err + } + var data queryData + err = json.Unmarshal(b, &data) + return data, err +} diff --git a/tests/integration/app_query_test.go b/tests/integration/app_query_test.go new file mode 100644 index 00000000..39a1553e --- /dev/null +++ b/tests/integration/app_query_test.go @@ -0,0 +1,44 @@ +//go:build integration + +package integration + +import ( + "context" + "os" + "testing" + "time" + + "github.com/zx06/xsql/internal/app" + "github.com/zx06/xsql/internal/config" +) + +func TestAppQuery_TruncationHint_MySQL(t *testing.T) { + dsn := os.Getenv("XSQL_TEST_MYSQL_DSN") + if dsn == "" { + t.Skip("XSQL_TEST_MYSQL_DSN not set") + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + result, xe := app.Query(ctx, app.QueryRequest{ + Profile: config.Profile{ + DB: "mysql", + DSN: dsn, + }, + SQL: `WITH RECURSIVE seq AS (SELECT 1 AS n UNION ALL SELECT n+1 FROM seq WHERE n < 8) SELECT n FROM seq`, + MaxRows: 3, + TruncationHint: app.QueryTruncationHint, + }) + if xe != nil { + t.Fatalf("query failed: %v", xe) + } + if len(result.Rows) != 3 || !result.Truncated { + t.Fatalf("expected truncated 3 rows, got %d truncated=%v", len(result.Rows), result.Truncated) + } + if result.TruncatedReason != "max_rows" { + t.Fatalf("truncated_reason=%q", result.TruncatedReason) + } + if result.Hint != app.QueryTruncationHint { + t.Fatalf("hint=%q", result.Hint) + } +} diff --git a/tests/integration/query_test.go b/tests/integration/query_test.go index 04ae5b60..e20e00fa 100644 --- a/tests/integration/query_test.go +++ b/tests/integration/query_test.go @@ -581,3 +581,82 @@ func TestPg_Query_CTE(t *testing.T) { t.Errorf("expected 3 rows, got %d", len(result.Rows)) } } + +func TestMySQL_Query_MaxRowsDefaultTruncation(t *testing.T) { + dsn := os.Getenv("XSQL_TEST_MYSQL_DSN") + if dsn == "" { + t.Skip("XSQL_TEST_MYSQL_DSN not set") + } + drv, _ := db.Get("mysql") + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + conn, xe := drv.Open(ctx, db.ConnOptions{DSN: dsn}) + if xe != nil { + t.Fatalf("failed to open: %v", xe) + } + defer conn.Close() + + sql := ` + WITH RECURSIVE seq AS ( + SELECT 1 AS n + UNION ALL + SELECT n + 1 FROM seq WHERE n < 50 + ) + SELECT n FROM seq` + result, xe := db.Query(ctx, conn, sql, db.QueryOptions{DBType: "mysql", MaxRows: 10}) + if xe != nil { + t.Fatalf("query failed: %v", xe) + } + if len(result.Rows) != 10 || !result.Truncated { + t.Fatalf("expected 10 truncated rows, got len=%d truncated=%v", len(result.Rows), result.Truncated) + } +} + +func TestMySQL_Query_MaxRowsUnlimited(t *testing.T) { + dsn := os.Getenv("XSQL_TEST_MYSQL_DSN") + if dsn == "" { + t.Skip("XSQL_TEST_MYSQL_DSN not set") + } + drv, _ := db.Get("mysql") + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + conn, xe := drv.Open(ctx, db.ConnOptions{DSN: dsn}) + if xe != nil { + t.Fatalf("failed to open: %v", xe) + } + defer conn.Close() + + result, xe := db.Query(ctx, conn, ` + WITH RECURSIVE seq AS ( + SELECT 1 AS n UNION ALL SELECT n + 1 FROM seq WHERE n < 5 + ) SELECT n FROM seq`, db.QueryOptions{DBType: "mysql", MaxRows: 0}) + if xe != nil { + t.Fatalf("query failed: %v", xe) + } + if len(result.Rows) != 5 || result.Truncated { + t.Fatalf("expected 5 rows unlimited, got len=%d truncated=%v", len(result.Rows), result.Truncated) + } +} + +func TestPg_Query_MaxRowsTruncation(t *testing.T) { + dsn := os.Getenv("XSQL_TEST_PG_DSN") + if dsn == "" { + t.Skip("XSQL_TEST_PG_DSN not set") + } + drv, _ := db.Get("pg") + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + conn, xe := drv.Open(ctx, db.ConnOptions{DSN: dsn}) + if xe != nil { + t.Fatalf("failed to open: %v", xe) + } + defer conn.Close() + + result, xe := db.Query(ctx, conn, "SELECT generate_series(1, 40) AS n", db.QueryOptions{DBType: "pg", MaxRows: 5}) + if xe != nil { + t.Fatalf("query failed: %v", xe) + } + if len(result.Rows) != 5 || !result.Truncated || result.MaxRows != 5 { + t.Fatalf("got len=%d truncated=%v max_rows=%d", len(result.Rows), result.Truncated, result.MaxRows) + } +} diff --git a/tests/integration/schema_include_system_test.go b/tests/integration/schema_include_system_test.go new file mode 100644 index 00000000..9245a6ee --- /dev/null +++ b/tests/integration/schema_include_system_test.go @@ -0,0 +1,78 @@ +//go:build integration + +package integration + +import ( + "context" + "os" + "testing" + "time" + + "github.com/zx06/xsql/internal/db" + _ "github.com/zx06/xsql/internal/db/mysql" // register mysql driver + _ "github.com/zx06/xsql/internal/db/pg" // register pg driver +) + +func TestListTables_IncludeSystem(t *testing.T) { + cases := []struct { + name string + envDSN string + driver string + pattern string + wantSchema string + wantTable string + }{ + { + name: "mysql_db", + envDSN: "XSQL_TEST_MYSQL_DSN", + driver: "mysql", + pattern: "db", + wantSchema: "mysql", + wantTable: "db", + }, + { + name: "pg_pg_namespace", + envDSN: "XSQL_TEST_PG_DSN", + driver: "pg", + pattern: "pg_namespace", + wantSchema: "pg_catalog", + wantTable: "pg_namespace", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assertIncludeSystemTable(t, tc.envDSN, tc.driver, tc.pattern, tc.wantSchema, tc.wantTable) + }) + } +} + +func assertIncludeSystemTable(t *testing.T, envVar, driver, pattern, wantSchema, wantTable string) { + t.Helper() + dsn := os.Getenv(envVar) + if dsn == "" { + t.Skip(envVar + " not set") + } + drv, _ := db.Get(driver) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + conn, xe := drv.Open(ctx, db.ConnOptions{DSN: dsn}) + if xe != nil { + t.Fatalf("open: %v", xe) + } + defer conn.Close() + + list, xe := db.ListTables(ctx, driver, conn, db.SchemaOptions{ + TablePattern: pattern, + IncludeSystem: true, + }) + if xe != nil { + t.Fatalf("ListTables: %v", xe) + } + for _, tbl := range list.Tables { + if tbl.Schema == wantSchema && tbl.Name == wantTable { + return + } + } + t.Fatalf("expected %s.%s, got %d tables", wantSchema, wantTable, len(list.Tables)) +}