duckquery-mcp 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- duckquery_mcp-0.1.0/.gitignore +206 -0
- duckquery_mcp-0.1.0/PKG-INFO +11 -0
- duckquery_mcp-0.1.0/README.md +44 -0
- duckquery_mcp-0.1.0/duckquery_mcp/__init__.py +1 -0
- duckquery_mcp-0.1.0/duckquery_mcp/__main__.py +7 -0
- duckquery_mcp-0.1.0/duckquery_mcp/client.py +94 -0
- duckquery_mcp-0.1.0/duckquery_mcp/config.py +26 -0
- duckquery_mcp-0.1.0/duckquery_mcp/safety.py +17 -0
- duckquery_mcp-0.1.0/duckquery_mcp/server.py +17 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/__init__.py +57 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/ai_settings.py +18 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/discover.py +49 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/export.py +11 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/passthrough.py +11 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/query.py +95 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/sources.py +142 -0
- duckquery_mcp-0.1.0/duckquery_mcp/tools/transform.py +20 -0
- duckquery_mcp-0.1.0/duckquery_mcp/util.py +6 -0
- duckquery_mcp-0.1.0/pyproject.toml +19 -0
- duckquery_mcp-0.1.0/tests/conftest.py +8 -0
- duckquery_mcp-0.1.0/tests/test_client.py +77 -0
- duckquery_mcp-0.1.0/tests/test_config.py +26 -0
- duckquery_mcp-0.1.0/tests/test_integration.py +15 -0
- duckquery_mcp-0.1.0/tests/test_normalize.py +48 -0
- duckquery_mcp-0.1.0/tests/test_passthrough.py +19 -0
- duckquery_mcp-0.1.0/tests/test_safety.py +19 -0
- duckquery_mcp-0.1.0/tests/test_tools_ai_export.py +14 -0
- duckquery_mcp-0.1.0/tests/test_tools_discover.py +14 -0
- duckquery_mcp-0.1.0/tests/test_tools_query.py +66 -0
- duckquery_mcp-0.1.0/tests/test_tools_sources.py +17 -0
- duckquery_mcp-0.1.0/tests/test_tools_transform.py +13 -0
|
@@ -0,0 +1,206 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# C extensions
|
|
7
|
+
*.so
|
|
8
|
+
|
|
9
|
+
# Distribution / packaging
|
|
10
|
+
.Python
|
|
11
|
+
build/
|
|
12
|
+
develop-eggs/
|
|
13
|
+
dist/
|
|
14
|
+
downloads/
|
|
15
|
+
eggs/
|
|
16
|
+
.eggs/
|
|
17
|
+
lib/
|
|
18
|
+
!frontend/src/lib/
|
|
19
|
+
lib64/
|
|
20
|
+
parts/
|
|
21
|
+
sdist/
|
|
22
|
+
var/
|
|
23
|
+
wheels/
|
|
24
|
+
*.egg-info/
|
|
25
|
+
.installed.cfg
|
|
26
|
+
*.egg
|
|
27
|
+
|
|
28
|
+
# PyInstaller
|
|
29
|
+
# Usually these files are written by a python script from a template
|
|
30
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
31
|
+
*.manifest
|
|
32
|
+
*.spec
|
|
33
|
+
!api/duckquery.spec
|
|
34
|
+
|
|
35
|
+
# PyInstaller build artifacts
|
|
36
|
+
api/dist/
|
|
37
|
+
api/build/
|
|
38
|
+
|
|
39
|
+
# Installer logs
|
|
40
|
+
pip-log.txt
|
|
41
|
+
pip-delete-this-directory.txt
|
|
42
|
+
|
|
43
|
+
# Unit test / coverage reports
|
|
44
|
+
htmlcov/
|
|
45
|
+
.tox/
|
|
46
|
+
.nox/
|
|
47
|
+
.coverage
|
|
48
|
+
.coverage.*
|
|
49
|
+
.cache
|
|
50
|
+
.pytest_cache/
|
|
51
|
+
.hypothesis/
|
|
52
|
+
|
|
53
|
+
# Generated lint reports (do not commit)
|
|
54
|
+
pylint-report.json
|
|
55
|
+
api/pylint-report.json
|
|
56
|
+
|
|
57
|
+
# Translations
|
|
58
|
+
*.mo
|
|
59
|
+
*.pot
|
|
60
|
+
|
|
61
|
+
# Django stuff:
|
|
62
|
+
*.log
|
|
63
|
+
local_settings.py
|
|
64
|
+
db.sqlite3
|
|
65
|
+
|
|
66
|
+
# Flask stuff:
|
|
67
|
+
instance/
|
|
68
|
+
.webassets-cache
|
|
69
|
+
|
|
70
|
+
# Scrapy stuff:
|
|
71
|
+
.scrapy
|
|
72
|
+
|
|
73
|
+
# Sphinx documentation
|
|
74
|
+
docs/_build/
|
|
75
|
+
|
|
76
|
+
# PyBuilder
|
|
77
|
+
target/
|
|
78
|
+
|
|
79
|
+
# Jupyter Notebook
|
|
80
|
+
.ipynb_checkpoints
|
|
81
|
+
|
|
82
|
+
# Environments
|
|
83
|
+
.env
|
|
84
|
+
.venv
|
|
85
|
+
env/
|
|
86
|
+
venv/
|
|
87
|
+
ENV/
|
|
88
|
+
env.bak/
|
|
89
|
+
venv.bak/
|
|
90
|
+
|
|
91
|
+
# IDEs
|
|
92
|
+
.idea/
|
|
93
|
+
.vscode/
|
|
94
|
+
|
|
95
|
+
# Node.js
|
|
96
|
+
node_modules/
|
|
97
|
+
dist/
|
|
98
|
+
.DS_Store
|
|
99
|
+
|
|
100
|
+
# Temp files
|
|
101
|
+
*.tmp
|
|
102
|
+
*.swp
|
|
103
|
+
|
|
104
|
+
# Project specific files - 敏感信息和数据文件
|
|
105
|
+
# Sensitive information and data files
|
|
106
|
+
api/mysql_configs.json
|
|
107
|
+
api/postgresql_configs.json
|
|
108
|
+
api/mysql_datasources.json
|
|
109
|
+
api/temp_files/*
|
|
110
|
+
!api/temp_files/.gitkeep
|
|
111
|
+
api/core/temp_files/*
|
|
112
|
+
!api/core/temp_files/.gitkeep
|
|
113
|
+
api/exports/*
|
|
114
|
+
!api/exports/.gitkeep
|
|
115
|
+
api/host_documents/
|
|
116
|
+
api/host_downloads/
|
|
117
|
+
api/server_mounts/
|
|
118
|
+
exports/*
|
|
119
|
+
!exports/.gitkeep
|
|
120
|
+
temp_files/*
|
|
121
|
+
!temp_files/.gitkeep
|
|
122
|
+
config/mysql-configs.json
|
|
123
|
+
config/postgresql-configs.json
|
|
124
|
+
config/mysql-production.json
|
|
125
|
+
config/datasources.json
|
|
126
|
+
/data/
|
|
127
|
+
!/data/uploads/.gitkeep
|
|
128
|
+
!/data/duckdb/.gitkeep
|
|
129
|
+
api/data/
|
|
130
|
+
|
|
131
|
+
# DuckDB相关文件
|
|
132
|
+
.duckdb/
|
|
133
|
+
*.db
|
|
134
|
+
*.duckdb
|
|
135
|
+
|
|
136
|
+
# 清理报告文件
|
|
137
|
+
cleanup-report.txt
|
|
138
|
+
|
|
139
|
+
# 真实配置文件(只保留.example文件)
|
|
140
|
+
# Real configuration files (keep only .example files)
|
|
141
|
+
**/mysql-configs.json
|
|
142
|
+
**/datasources.json
|
|
143
|
+
**/app-config.json
|
|
144
|
+
**/file-datasources.json
|
|
145
|
+
**/sql-favorites.json
|
|
146
|
+
**/secret.key
|
|
147
|
+
!**/*.example
|
|
148
|
+
|
|
149
|
+
# Frontend build outputs
|
|
150
|
+
frontend/dist/
|
|
151
|
+
frontend/build/
|
|
152
|
+
|
|
153
|
+
# Additional Node.js ignores
|
|
154
|
+
npm-debug.log*
|
|
155
|
+
yarn-debug.log*
|
|
156
|
+
yarn-error.log*
|
|
157
|
+
lerna-debug.log*
|
|
158
|
+
.npm
|
|
159
|
+
.eslintcache
|
|
160
|
+
|
|
161
|
+
# OS specific
|
|
162
|
+
Thumbs.db
|
|
163
|
+
ehthumbs.db
|
|
164
|
+
.Spotlight-V100
|
|
165
|
+
.Trashes
|
|
166
|
+
|
|
167
|
+
# Local clipboard and temporary files
|
|
168
|
+
.gemini-clipboard/
|
|
169
|
+
|
|
170
|
+
# Backend writes a per-deployment runtime descriptor here
|
|
171
|
+
# (Docker sets APP_ROOT=/app -> ./api/runtime.json). Generated, not source.
|
|
172
|
+
api/runtime.json
|
|
173
|
+
|
|
174
|
+
# Logs
|
|
175
|
+
logs/
|
|
176
|
+
*.log
|
|
177
|
+
|
|
178
|
+
# Database files
|
|
179
|
+
*.db
|
|
180
|
+
*.sqlite
|
|
181
|
+
*.sqlite3
|
|
182
|
+
|
|
183
|
+
# Backup files
|
|
184
|
+
*.bak
|
|
185
|
+
*.backup
|
|
186
|
+
|
|
187
|
+
# Cursor IDE files
|
|
188
|
+
.cursorrules
|
|
189
|
+
.cursor/
|
|
190
|
+
.claude/
|
|
191
|
+
# Git hook scripts at repo root only (do not ignore frontend/src/hooks)
|
|
192
|
+
/hooks/
|
|
193
|
+
CLAUDE.md
|
|
194
|
+
# runtime file data sources cache
|
|
195
|
+
data/file_datasources.json
|
|
196
|
+
|
|
197
|
+
# 本地运行时数据库/扩展(联邦扩展二进制、DuckDB 数据与 WAL),不入库
|
|
198
|
+
data-test/
|
|
199
|
+
*.duckdb_extension
|
|
200
|
+
*.db.wal
|
|
201
|
+
*.db.wal.broken.*
|
|
202
|
+
|
|
203
|
+
# DuckDB 离线扩展预下载产物(由 fetch_duckdb_extensions.py 生成)
|
|
204
|
+
api/extensions/
|
|
205
|
+
# Tauri sidecar(打包时由 CI 暂存的 PyInstaller onedir,勿入库)
|
|
206
|
+
frontend/src-tauri/binaries/
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: duckquery-mcp
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MCP server for a locally-running DuckQuery backend
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Requires-Dist: httpx>=0.27
|
|
7
|
+
Requires-Dist: mcp>=1.2.0
|
|
8
|
+
Provides-Extra: dev
|
|
9
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
10
|
+
Requires-Dist: pytest>=8; extra == 'dev'
|
|
11
|
+
Requires-Dist: respx>=0.21; extra == 'dev'
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
# duckquery-mcp
|
|
2
|
+
|
|
3
|
+
An MCP server that exposes a locally-running DuckQuery backend to MCP clients
|
|
4
|
+
(Claude Code, Cursor, Codex). Start DuckQuery first (desktop app, Docker, or
|
|
5
|
+
`uvicorn main:app --port 48001`); this server auto-discovers it.
|
|
6
|
+
|
|
7
|
+
## Run
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
uvx duckquery-mcp # zero-install
|
|
11
|
+
# or: pipx run duckquery-mcp
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
Env:
|
|
15
|
+
- `DUCKQUERY_API_BASE` — explicit backend URL (e.g. `http://127.0.0.1:48001`). Optional; auto-discovered otherwise (runtime.json, then probes 48001/8000/8001).
|
|
16
|
+
- `DUCKQUERY_MCP_MODE` — `read-only` | `normal` (default) | `full`.
|
|
17
|
+
|
|
18
|
+
## Add to a CLI
|
|
19
|
+
|
|
20
|
+
Claude Code:
|
|
21
|
+
```bash
|
|
22
|
+
claude mcp add duckquery -- uvx duckquery-mcp
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
Cursor / Codex (`mcp.json`):
|
|
26
|
+
```json
|
|
27
|
+
{
|
|
28
|
+
"mcpServers": {
|
|
29
|
+
"duckquery": {
|
|
30
|
+
"command": "uvx",
|
|
31
|
+
"args": ["duckquery-mcp"],
|
|
32
|
+
"env": { "DUCKQUERY_MCP_MODE": "normal" }
|
|
33
|
+
}
|
|
34
|
+
}
|
|
35
|
+
}
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
## Tools
|
|
39
|
+
|
|
40
|
+
High-level tools (query, ask, discover, add sources, configure LLM, transform,
|
|
41
|
+
export) plus a generic `duckquery_request` passthrough. Safety mode gates which
|
|
42
|
+
tools are exposed: `read-only` hides all mutating tools; `normal` exposes them
|
|
43
|
+
but destructive raw SQL and non-GET passthrough require `confirm=true`; `full`
|
|
44
|
+
removes the gate.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import os
|
|
3
|
+
import sys
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import httpx
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class BackendNotFound(Exception):
|
|
11
|
+
pass
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class BackendError(Exception):
|
|
15
|
+
pass
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def runtime_file() -> Path:
|
|
19
|
+
"""Mirror api/core/common/paths.get_user_data_dir() / 'runtime.json'."""
|
|
20
|
+
override = os.getenv("APP_ROOT")
|
|
21
|
+
if override:
|
|
22
|
+
return Path(override) / "runtime.json"
|
|
23
|
+
home = Path(os.path.expanduser("~"))
|
|
24
|
+
if sys.platform == "darwin":
|
|
25
|
+
base = home / "Library" / "Application Support" / "DuckQuery"
|
|
26
|
+
elif sys.platform.startswith("win"):
|
|
27
|
+
base = Path(os.getenv("APPDATA") or home) / "DuckQuery"
|
|
28
|
+
else:
|
|
29
|
+
base = home / ".local" / "share" / "DuckQuery"
|
|
30
|
+
return base / "runtime.json"
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class DuckQueryClient:
|
|
34
|
+
def __init__(self, cfg) -> None:
|
|
35
|
+
self.cfg = cfg
|
|
36
|
+
self._base: str | None = None
|
|
37
|
+
self._http = httpx.AsyncClient(timeout=cfg.timeout)
|
|
38
|
+
|
|
39
|
+
async def _healthy(self, base: str) -> bool:
|
|
40
|
+
try:
|
|
41
|
+
r = await self._http.get(f"{base}/health")
|
|
42
|
+
return r.status_code == 200 and r.json().get("status") == "healthy"
|
|
43
|
+
except Exception:
|
|
44
|
+
return False
|
|
45
|
+
|
|
46
|
+
async def base(self) -> str:
|
|
47
|
+
if self._base and await self._healthy(self._base):
|
|
48
|
+
return self._base
|
|
49
|
+
# 1. explicit env
|
|
50
|
+
if self.cfg.api_base and await self._healthy(self.cfg.api_base):
|
|
51
|
+
self._base = self.cfg.api_base
|
|
52
|
+
return self._base
|
|
53
|
+
# 2. runtime.json
|
|
54
|
+
rf = runtime_file()
|
|
55
|
+
if rf.exists():
|
|
56
|
+
try:
|
|
57
|
+
b = json.loads(rf.read_text()).get("base")
|
|
58
|
+
if b and await self._healthy(b):
|
|
59
|
+
self._base = b
|
|
60
|
+
return self._base
|
|
61
|
+
except Exception:
|
|
62
|
+
pass
|
|
63
|
+
# 3. probe known ports
|
|
64
|
+
for port in self.cfg.probe_ports:
|
|
65
|
+
b = f"http://127.0.0.1:{port}"
|
|
66
|
+
if await self._healthy(b):
|
|
67
|
+
self._base = b
|
|
68
|
+
return self._base
|
|
69
|
+
raise BackendNotFound(
|
|
70
|
+
"DuckQuery backend not found — start the DuckQuery app "
|
|
71
|
+
"or set DUCKQUERY_API_BASE."
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
async def call(self, method: str, path: str, *, json_body: Any = None,
|
|
75
|
+
params: dict | None = None) -> Any:
|
|
76
|
+
base = await self.base()
|
|
77
|
+
r = await self._http.request(method, f"{base}{path}", json=json_body, params=params)
|
|
78
|
+
try:
|
|
79
|
+
payload = r.json()
|
|
80
|
+
except Exception:
|
|
81
|
+
r.raise_for_status()
|
|
82
|
+
return {"raw": r.text}
|
|
83
|
+
if r.status_code >= 400:
|
|
84
|
+
msg = None
|
|
85
|
+
if isinstance(payload, dict):
|
|
86
|
+
# FastAPI uses "detail"; DuckQuery's envelope uses "message"/"messageCode"
|
|
87
|
+
msg = payload.get("detail") or payload.get("message") or payload.get("messageCode")
|
|
88
|
+
raise BackendError(str(msg) if msg else f"HTTP {r.status_code}")
|
|
89
|
+
if isinstance(payload, dict):
|
|
90
|
+
if payload.get("success") is False:
|
|
91
|
+
raise BackendError(payload.get("message") or payload.get("messageCode") or "request failed")
|
|
92
|
+
if "data" in payload:
|
|
93
|
+
return payload["data"]
|
|
94
|
+
return payload
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
|
|
4
|
+
MODES = ("read-only", "normal", "full")
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@dataclass(frozen=True)
|
|
8
|
+
class Config:
|
|
9
|
+
api_base: str | None
|
|
10
|
+
mode: str
|
|
11
|
+
timeout: float
|
|
12
|
+
row_cap: int
|
|
13
|
+
probe_ports: tuple[int, ...]
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def load_config() -> Config:
|
|
17
|
+
mode = os.getenv("DUCKQUERY_MCP_MODE", "normal")
|
|
18
|
+
if mode not in MODES:
|
|
19
|
+
raise SystemExit(f"DUCKQUERY_MCP_MODE must be one of {MODES}, got {mode!r}")
|
|
20
|
+
return Config(
|
|
21
|
+
api_base=os.getenv("DUCKQUERY_API_BASE") or None,
|
|
22
|
+
mode=mode,
|
|
23
|
+
timeout=float(os.getenv("DUCKQUERY_MCP_TIMEOUT", "120")),
|
|
24
|
+
row_cap=int(os.getenv("DUCKQUERY_MCP_ROW_CAP", "200")),
|
|
25
|
+
probe_ports=(48001, 8000, 8001),
|
|
26
|
+
)
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
import re
|
|
2
|
+
|
|
3
|
+
_READ = re.compile(r"^\s*(SELECT|WITH|EXPLAIN|PRAGMA|DESCRIBE|SHOW)\b", re.I)
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def is_write_sql(sql: str) -> bool:
|
|
7
|
+
"""True unless the statement is clearly read-only."""
|
|
8
|
+
return _READ.match(sql or "") is None
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def tool_allowed(tier: str, mode: str) -> bool:
|
|
12
|
+
"""tier: 'read' | 'write'. mode: 'read-only' | 'normal' | 'full'."""
|
|
13
|
+
if mode == "full":
|
|
14
|
+
return True
|
|
15
|
+
if mode == "normal":
|
|
16
|
+
return True
|
|
17
|
+
return tier == "read" # read-only
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from mcp.server.fastmcp import FastMCP
|
|
2
|
+
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.config import load_config
|
|
5
|
+
from duckquery_mcp.tools import register_all
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def build() -> FastMCP:
|
|
9
|
+
cfg = load_config()
|
|
10
|
+
client = DuckQueryClient(cfg)
|
|
11
|
+
mcp = FastMCP("duckquery")
|
|
12
|
+
register_all(mcp, client, cfg)
|
|
13
|
+
return mcp
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def run() -> None:
|
|
17
|
+
build().run() # stdio transport by default
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
import inspect
|
|
2
|
+
|
|
3
|
+
from mcp.server.fastmcp import FastMCP
|
|
4
|
+
|
|
5
|
+
from duckquery_mcp.client import BackendError, BackendNotFound, DuckQueryClient
|
|
6
|
+
from duckquery_mcp.config import Config
|
|
7
|
+
from duckquery_mcp.safety import tool_allowed
|
|
8
|
+
from duckquery_mcp.tools import ai_settings, discover, export, passthrough, query, sources, transform
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def register_all(mcp: FastMCP, client: DuckQueryClient, cfg: Config) -> None:
|
|
12
|
+
def add(tier: str):
|
|
13
|
+
"""Register tool (only if its tier is allowed by mode), binding client/cfg and
|
|
14
|
+
preserving the user-facing parameter signature for the MCP schema."""
|
|
15
|
+
def deco(fn):
|
|
16
|
+
if not tool_allowed(tier, cfg.mode):
|
|
17
|
+
return fn
|
|
18
|
+
sig = inspect.signature(fn)
|
|
19
|
+
user_params = [p for n, p in sig.parameters.items() if n not in ("client", "cfg")]
|
|
20
|
+
|
|
21
|
+
async def wrapped(**kwargs):
|
|
22
|
+
try:
|
|
23
|
+
return await fn(client, cfg, **kwargs)
|
|
24
|
+
except (BackendNotFound, BackendError) as exc:
|
|
25
|
+
return {"error": str(exc)}
|
|
26
|
+
|
|
27
|
+
wrapped.__name__ = fn.__name__
|
|
28
|
+
wrapped.__doc__ = fn.__doc__
|
|
29
|
+
wrapped.__signature__ = sig.replace(parameters=user_params)
|
|
30
|
+
mcp.tool()(wrapped)
|
|
31
|
+
return fn
|
|
32
|
+
return deco
|
|
33
|
+
|
|
34
|
+
add("read")(discover.list_tables)
|
|
35
|
+
add("read")(discover.describe_table)
|
|
36
|
+
add("read")(discover.list_connections)
|
|
37
|
+
add("read")(discover.list_db_objects)
|
|
38
|
+
add("read")(query.run_sql)
|
|
39
|
+
add("read")(query.federated_query)
|
|
40
|
+
add("read")(query.ask)
|
|
41
|
+
add("read")(query.explain_sql)
|
|
42
|
+
add("read")(query.suggest_chart)
|
|
43
|
+
add("read")(query.chat)
|
|
44
|
+
add("read")(query.error_fix)
|
|
45
|
+
add("write")(sources.add_connection)
|
|
46
|
+
add("write")(sources.add_local_file_source)
|
|
47
|
+
add("write")(sources.import_excel)
|
|
48
|
+
add("write")(sources.paste_data)
|
|
49
|
+
add("write")(sources.read_url)
|
|
50
|
+
add("write")(transform.save_as_table)
|
|
51
|
+
add("write")(transform.pivot)
|
|
52
|
+
add("write")(transform.set_operations)
|
|
53
|
+
add("read")(ai_settings.get_ai_settings)
|
|
54
|
+
add("write")(ai_settings.configure_llm)
|
|
55
|
+
add("write")(ai_settings.test_llm_provider)
|
|
56
|
+
add("read")(export.export_results)
|
|
57
|
+
add("read")(passthrough.duckquery_request)
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
async def get_ai_settings(client, cfg) -> Any:
|
|
5
|
+
"""Current AI/LLM settings (api keys are masked by the backend)."""
|
|
6
|
+
return await client.call("GET", "/api/settings/ai")
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
async def configure_llm(client, cfg, *, settings: dict) -> Any:
|
|
10
|
+
"""Update AI/LLM settings. `settings` matches AISettingsPayload fields:
|
|
11
|
+
enabled (bool), default_provider (str|None), providers (list[dict]),
|
|
12
|
+
features (dict), timeout_seconds (int), num_retries (int)."""
|
|
13
|
+
return await client.call("PUT", "/api/settings/ai", json_body=settings)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
async def test_llm_provider(client, cfg, *, provider_id: str) -> Any:
|
|
17
|
+
"""Test connectivity/credentials for one configured provider."""
|
|
18
|
+
return await client.call("POST", f"/api/ai/providers/{provider_id}/test")
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
from duckquery_mcp.util import normalize_connection_id
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
async def list_tables(client, cfg) -> Any:
|
|
7
|
+
"""List DuckDB tables currently loaded in the local engine."""
|
|
8
|
+
return await client.call("GET", "/api/duckdb/tables")
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
async def describe_table(client, cfg, *, name: str) -> Any:
|
|
12
|
+
"""Columns/types/sample for one DuckDB table."""
|
|
13
|
+
return await client.call("GET", f"/api/duckdb/tables/detail/{name}")
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
async def list_connections(client, cfg) -> Any:
|
|
17
|
+
"""List saved external database connections (MySQL/Postgres)."""
|
|
18
|
+
return await client.call("GET", "/api/datasources/databases/list")
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
async def list_db_objects(client, cfg, *, connection_id: str, kind: str = "tables") -> Any:
|
|
22
|
+
"""List schemas or tables in an external connection. kind: 'schemas' | 'tables'.
|
|
23
|
+
|
|
24
|
+
The connection_id from list_connections (e.g. 'db_SORDER') is accepted; its 'db_'
|
|
25
|
+
prefix is normalized. For 'tables', returns COMPACT entries (table name + comment +
|
|
26
|
+
column_count) capped at 200 — a big schema's full column lists would be huge. To get
|
|
27
|
+
one table's columns, run `federated_query("SELECT * FROM alias.<table> LIMIT 0", ...)`.
|
|
28
|
+
"""
|
|
29
|
+
cid = normalize_connection_id(connection_id)
|
|
30
|
+
data = await client.call("GET", f"/api/datasources/databases/{cid}/{kind}")
|
|
31
|
+
tables = data.get("tables") if isinstance(data, dict) else None
|
|
32
|
+
if isinstance(tables, list):
|
|
33
|
+
cap = 200
|
|
34
|
+
compact = [
|
|
35
|
+
{
|
|
36
|
+
"table_name": t.get("table_name"),
|
|
37
|
+
"comment": t.get("table_comment"),
|
|
38
|
+
"column_count": len(t.get("columns") or []),
|
|
39
|
+
}
|
|
40
|
+
for t in tables[:cap]
|
|
41
|
+
]
|
|
42
|
+
return {
|
|
43
|
+
"connection_id": data.get("connection_id"),
|
|
44
|
+
"database": data.get("database"),
|
|
45
|
+
"table_count": len(tables),
|
|
46
|
+
"truncated": len(tables) > cap,
|
|
47
|
+
"tables": compact,
|
|
48
|
+
}
|
|
49
|
+
return data
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
async def export_results(client, cfg, *, sql: str, format: str = "parquet",
|
|
5
|
+
attach_databases: list | None = None) -> Any:
|
|
6
|
+
"""Export a query result to a file; returns file_id, download_url, format, row_count_estimate.
|
|
7
|
+
format: 'parquet' (default) or 'csv'. Only SELECT queries are allowed."""
|
|
8
|
+
body: dict = {"sql": sql, "format": format}
|
|
9
|
+
if attach_databases is not None:
|
|
10
|
+
body["attach_databases"] = attach_databases
|
|
11
|
+
return await client.call("POST", "/api/query-results/export", json_body=body)
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
async def duckquery_request(client, cfg, *, method: str, path: str,
|
|
5
|
+
json: dict | None = None, confirm: bool = False) -> Any:
|
|
6
|
+
"""Call any DuckQuery API endpoint directly (escape hatch for features without a dedicated tool).
|
|
7
|
+
Non-GET methods require confirm=True unless mode is 'full'."""
|
|
8
|
+
method = method.upper()
|
|
9
|
+
if method != "GET" and cfg.mode != "full" and not confirm:
|
|
10
|
+
return {"error": "This is a mutating request; pass confirm=true to proceed."}
|
|
11
|
+
return await client.call(method, path, json_body=json)
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.config import Config
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def _truncate(data: dict, cfg: Config) -> dict:
|
|
8
|
+
"""Compact a /execute-style result; cap rows at cfg.row_cap."""
|
|
9
|
+
rows = data.get("data") or []
|
|
10
|
+
capped = rows[: cfg.row_cap]
|
|
11
|
+
return {
|
|
12
|
+
"columns": data.get("columns"),
|
|
13
|
+
"rows": capped,
|
|
14
|
+
"row_count": data.get("row_count", len(rows)),
|
|
15
|
+
"truncated": len(rows) > len(capped),
|
|
16
|
+
}
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
async def run_sql(client: DuckQueryClient, cfg: Config, *, sql: str, preview: bool = True) -> Any:
|
|
20
|
+
"""Run DuckDB SQL against local tables. Returns columns + (capped) rows."""
|
|
21
|
+
from duckquery_mcp.safety import is_write_sql
|
|
22
|
+
if cfg.mode == "read-only" and is_write_sql(sql):
|
|
23
|
+
return {"error": "read-only mode: only SELECT / WITH / EXPLAIN are allowed."}
|
|
24
|
+
data = await client.call("POST", "/api/duckdb/execute",
|
|
25
|
+
json_body={"sql": sql, "is_preview": preview})
|
|
26
|
+
return _truncate(data, cfg)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
async def federated_query(client: DuckQueryClient, cfg: Config, *, sql: str, attach_databases: list) -> Any:
|
|
30
|
+
"""Run SQL across attached external DBs (MySQL/Postgres) + local tables.
|
|
31
|
+
|
|
32
|
+
attach_databases: [{"alias": "m", "connection_id": "SORDER"}]. Connection ids from
|
|
33
|
+
list_connections (e.g. "db_SORDER") are accepted — the "db_" prefix is normalized.
|
|
34
|
+
Reference an attached table as alias.table, e.g. SELECT * FROM m.orders LIMIT 100.
|
|
35
|
+
"""
|
|
36
|
+
from duckquery_mcp.safety import is_write_sql
|
|
37
|
+
from duckquery_mcp.util import normalize_connection_id
|
|
38
|
+
if cfg.mode == "read-only" and is_write_sql(sql):
|
|
39
|
+
return {"error": "read-only mode: only SELECT / WITH / EXPLAIN are allowed."}
|
|
40
|
+
attach = [
|
|
41
|
+
{**db, "connection_id": normalize_connection_id(db["connection_id"])}
|
|
42
|
+
if isinstance(db, dict) and db.get("connection_id") else db
|
|
43
|
+
for db in (attach_databases or [])
|
|
44
|
+
]
|
|
45
|
+
data = await client.call("POST", "/api/duckdb/federated-query",
|
|
46
|
+
json_body={"sql": sql, "attach_databases": attach, "is_preview": True})
|
|
47
|
+
return _truncate(data, cfg)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
async def ask(client: DuckQueryClient, cfg: Config, *, question: str, tables: list | None = None, locale: str = "zh") -> Any:
|
|
51
|
+
"""Natural-language question -> generated DuckDB SQL -> executed result."""
|
|
52
|
+
gen = await client.call("POST", "/api/ai/nl-to-sql",
|
|
53
|
+
json_body={"question": question, "tables": tables or [], "locale": locale})
|
|
54
|
+
sql = gen.get("sql") if isinstance(gen, dict) else None
|
|
55
|
+
if not sql:
|
|
56
|
+
return {"error": "no SQL generated", "raw": gen}
|
|
57
|
+
result = await run_sql(client, cfg, sql=sql)
|
|
58
|
+
return {"generated_sql": sql, **result}
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
async def explain_sql(client: DuckQueryClient, cfg: Config, *, sql: str, locale: str = "zh") -> Any:
|
|
62
|
+
"""Plain-language explanation of a SQL statement."""
|
|
63
|
+
return await client.call("POST", "/api/ai/explain-sql", json_body={"sql": sql, "locale": locale})
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
async def suggest_chart(
|
|
67
|
+
client: DuckQueryClient,
|
|
68
|
+
cfg: Config,
|
|
69
|
+
*,
|
|
70
|
+
columns: list,
|
|
71
|
+
sample: list | None = None,
|
|
72
|
+
locale: str = "zh",
|
|
73
|
+
) -> Any:
|
|
74
|
+
"""Suggest a chart type for a query result. Pass columns (list of {name, type}) and optional sample rows."""
|
|
75
|
+
return await client.call("POST", "/api/ai/suggest-chart",
|
|
76
|
+
json_body={"columns": columns, "sample": sample or [], "locale": locale})
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
async def chat(
|
|
80
|
+
client: DuckQueryClient,
|
|
81
|
+
cfg: Config,
|
|
82
|
+
*,
|
|
83
|
+
messages: list,
|
|
84
|
+
tables: list | None = None,
|
|
85
|
+
locale: str = "zh",
|
|
86
|
+
) -> Any:
|
|
87
|
+
"""Free-form data conversation with the configured LLM. messages is [{role, content}, ...]."""
|
|
88
|
+
return await client.call("POST", "/api/ai/chat",
|
|
89
|
+
json_body={"messages": messages, "tables": tables or [], "locale": locale})
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
async def error_fix(client: DuckQueryClient, cfg: Config, *, sql: str, error_message: str) -> Any:
|
|
93
|
+
"""Error doctor: suggest a fix for a failing query, given the error."""
|
|
94
|
+
return await client.call("POST", "/api/ai/error-fix",
|
|
95
|
+
json_body={"sql": sql, "error": error_message})
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
async def add_connection(client, cfg, *, connection: dict, test: bool = True) -> Any:
|
|
5
|
+
"""Save (and optionally test) an external DB connection.
|
|
6
|
+
|
|
7
|
+
`connection` must follow the DatabaseConnection shape:
|
|
8
|
+
{name, type, params: {host, port, database, username, password, ...}}
|
|
9
|
+
where `type` is one of: mysql, postgresql, sqlite, etc.
|
|
10
|
+
`test=True` (default) verifies the connection before saving; set to False to save without testing.
|
|
11
|
+
"""
|
|
12
|
+
return await client.call(
|
|
13
|
+
"POST",
|
|
14
|
+
"/api/datasources/databases",
|
|
15
|
+
params={"test_connection": str(test).lower()},
|
|
16
|
+
json_body=connection,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
async def add_local_file_source(
|
|
21
|
+
client,
|
|
22
|
+
cfg,
|
|
23
|
+
*,
|
|
24
|
+
path: str,
|
|
25
|
+
table_alias: str | None = None,
|
|
26
|
+
import_mode: str = "auto",
|
|
27
|
+
csv_delimiter: str | None = None,
|
|
28
|
+
csv_has_header: bool | None = None,
|
|
29
|
+
csv_encoding: str | None = None,
|
|
30
|
+
) -> Any:
|
|
31
|
+
"""Register a local CSV/Parquet/JSON/Excel file as a DuckDB table.
|
|
32
|
+
|
|
33
|
+
Desktop mode allows any local path (ALLOW_ARBITRARY_LOCAL_PATHS=1).
|
|
34
|
+
`import_mode`: auto | smart | raw (controls type inference on load).
|
|
35
|
+
CSV-specific options (`csv_delimiter`, `csv_has_header`, `csv_encoding`) are ignored for non-CSV files.
|
|
36
|
+
"""
|
|
37
|
+
body: dict = {"path": path, "import_mode": import_mode}
|
|
38
|
+
if table_alias is not None:
|
|
39
|
+
body["table_alias"] = table_alias
|
|
40
|
+
if csv_delimiter is not None:
|
|
41
|
+
body["csv_delimiter"] = csv_delimiter
|
|
42
|
+
if csv_has_header is not None:
|
|
43
|
+
body["csv_has_header"] = csv_has_header
|
|
44
|
+
if csv_encoding is not None:
|
|
45
|
+
body["csv_encoding"] = csv_encoding
|
|
46
|
+
return await client.call("POST", "/api/server-files/import", json_body=body)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
async def import_excel(
|
|
50
|
+
client,
|
|
51
|
+
cfg,
|
|
52
|
+
*,
|
|
53
|
+
path: str,
|
|
54
|
+
sheets: list,
|
|
55
|
+
import_mode: str = "auto",
|
|
56
|
+
) -> Any:
|
|
57
|
+
"""Import selected Excel sheets as DuckDB tables.
|
|
58
|
+
|
|
59
|
+
Each item in `sheets` is a dict matching ExcelSheetImportConfig:
|
|
60
|
+
{name: str, target_table: str, header_rows?: int, header_row_index?: int,
|
|
61
|
+
fill_merged?: bool, mode?: "create"|"append"|"replace"}
|
|
62
|
+
`import_mode`: auto | smart | raw.
|
|
63
|
+
"""
|
|
64
|
+
return await client.call(
|
|
65
|
+
"POST",
|
|
66
|
+
"/api/server-files/excel/import",
|
|
67
|
+
json_body={"path": path, "sheets": sheets, "import_mode": import_mode},
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
async def paste_data(
|
|
72
|
+
client,
|
|
73
|
+
cfg,
|
|
74
|
+
*,
|
|
75
|
+
table_name: str,
|
|
76
|
+
column_names: list,
|
|
77
|
+
column_types: list,
|
|
78
|
+
data_rows: list,
|
|
79
|
+
delimiter: str = ",",
|
|
80
|
+
has_header: bool = False,
|
|
81
|
+
) -> Any:
|
|
82
|
+
"""Create a DuckDB table from pasted tabular data.
|
|
83
|
+
|
|
84
|
+
`column_names`: list of column name strings, e.g. ["id", "name"].
|
|
85
|
+
`column_types`: list of type strings matching column_names, e.g. ["INTEGER", "VARCHAR"].
|
|
86
|
+
Supported types: VARCHAR, INTEGER, DOUBLE, DATE, BOOLEAN.
|
|
87
|
+
`data_rows`: list of rows, each a list of string cell values,
|
|
88
|
+
e.g. [["1", "Alice"], ["2", "Bob"]].
|
|
89
|
+
`delimiter`: used only for context (data is already parsed into rows).
|
|
90
|
+
`has_header`: whether the first row of data_rows is a header (usually False since
|
|
91
|
+
column_names is explicit).
|
|
92
|
+
"""
|
|
93
|
+
return await client.call(
|
|
94
|
+
"POST",
|
|
95
|
+
"/api/paste-data",
|
|
96
|
+
json_body={
|
|
97
|
+
"table_name": table_name,
|
|
98
|
+
"column_names": column_names,
|
|
99
|
+
"column_types": column_types,
|
|
100
|
+
"data_rows": data_rows,
|
|
101
|
+
"delimiter": delimiter,
|
|
102
|
+
"has_header": has_header,
|
|
103
|
+
},
|
|
104
|
+
)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
async def read_url(
|
|
108
|
+
client,
|
|
109
|
+
cfg,
|
|
110
|
+
*,
|
|
111
|
+
url: str,
|
|
112
|
+
table_alias: str,
|
|
113
|
+
file_type: str | None = None,
|
|
114
|
+
import_mode: str = "auto",
|
|
115
|
+
encoding: str = "utf-8",
|
|
116
|
+
delimiter: str = ",",
|
|
117
|
+
header: bool = True,
|
|
118
|
+
prefer_native: bool = True,
|
|
119
|
+
) -> Any:
|
|
120
|
+
"""Download a remote file URL and load it into a DuckDB table.
|
|
121
|
+
|
|
122
|
+
`url`: public HTTP/HTTPS URL (GitHub blob URLs are auto-converted to raw).
|
|
123
|
+
`table_alias`: desired table name (de-duplicated if it already exists).
|
|
124
|
+
`file_type`: csv | json | parquet | excel (auto-detected from URL/Content-Type if omitted).
|
|
125
|
+
`import_mode`: auto | smart | raw.
|
|
126
|
+
`prefer_native`: try DuckDB/httpfs direct read first (faster), fall back to HTTP download.
|
|
127
|
+
Internal/loopback addresses and S3 URLs are blocked by the backend.
|
|
128
|
+
"""
|
|
129
|
+
return await client.call(
|
|
130
|
+
"POST",
|
|
131
|
+
"/api/read_from_url",
|
|
132
|
+
json_body={
|
|
133
|
+
"url": url,
|
|
134
|
+
"table_alias": table_alias,
|
|
135
|
+
"file_type": file_type,
|
|
136
|
+
"import_mode": import_mode,
|
|
137
|
+
"encoding": encoding,
|
|
138
|
+
"delimiter": delimiter,
|
|
139
|
+
"header": header,
|
|
140
|
+
"prefer_native": prefer_native,
|
|
141
|
+
},
|
|
142
|
+
)
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
async def save_as_table(client, cfg, *, sql: str, table_name: str) -> Any:
|
|
5
|
+
"""Materialize a query's result as a new DuckDB table."""
|
|
6
|
+
# Router reads "table_alias" (or "tableAlias"); public param kept as table_name for clarity.
|
|
7
|
+
return await client.call("POST", "/api/save_query_to_duckdb",
|
|
8
|
+
json_body={"sql": sql, "table_alias": table_name})
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
async def pivot(client, cfg, *, config: dict, pivot_config: dict, execute: bool = False) -> Any:
|
|
12
|
+
"""Pivot a table. execute=False previews; execute=True writes the pivoted result."""
|
|
13
|
+
path = "/api/pivot-query/generate" if execute else "/api/pivot-query/preview"
|
|
14
|
+
return await client.call("POST", path, json_body={"config": config, "pivot_config": pivot_config})
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
async def set_operations(client, cfg, *, config: dict, execute: bool = False) -> Any:
|
|
18
|
+
"""UNION/INTERSECT/EXCEPT across DuckDB tables. Pass config (SetOperationConfig dict). execute=False previews."""
|
|
19
|
+
path = "/api/set-operations/execute" if execute else "/api/set-operations/preview"
|
|
20
|
+
return await client.call("POST", path, json_body={"config": config})
|
|
@@ -0,0 +1,6 @@
|
|
|
1
|
+
def normalize_connection_id(connection_id: str) -> str:
|
|
2
|
+
"""Strip the ``db_`` prefix that ``/databases/list`` adds to connection ids but the
|
|
3
|
+
introspect / federated-query endpoints don't accept. So an id like ``db_SORDER``
|
|
4
|
+
(from list_connections) becomes ``SORDER`` (what those endpoints look up)."""
|
|
5
|
+
cid = str(connection_id)
|
|
6
|
+
return cid[3:] if cid.startswith("db_") else cid
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "duckquery-mcp"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "MCP server for a locally-running DuckQuery backend"
|
|
5
|
+
requires-python = ">=3.10"
|
|
6
|
+
dependencies = ["mcp>=1.2.0", "httpx>=0.27"]
|
|
7
|
+
|
|
8
|
+
[project.scripts]
|
|
9
|
+
duckquery-mcp = "duckquery_mcp.__main__:main"
|
|
10
|
+
|
|
11
|
+
[project.optional-dependencies]
|
|
12
|
+
dev = ["pytest>=8", "pytest-asyncio>=0.23", "respx>=0.21"]
|
|
13
|
+
|
|
14
|
+
[build-system]
|
|
15
|
+
requires = ["hatchling"]
|
|
16
|
+
build-backend = "hatchling.build"
|
|
17
|
+
|
|
18
|
+
[tool.pytest.ini_options]
|
|
19
|
+
asyncio_mode = "auto"
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
import respx
|
|
2
|
+
import httpx
|
|
3
|
+
import pytest
|
|
4
|
+
from duckquery_mcp.client import DuckQueryClient, BackendNotFound, BackendError
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@respx.mock
|
|
8
|
+
async def test_probe_finds_healthy_backend(cfg):
|
|
9
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
10
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
11
|
+
client = DuckQueryClient(cfg)
|
|
12
|
+
assert await client.base() == "http://127.0.0.1:48001"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@respx.mock
|
|
16
|
+
async def test_env_base_wins(cfg):
|
|
17
|
+
cfg = cfg.__class__(**{**cfg.__dict__, "api_base": "http://127.0.0.1:7000"})
|
|
18
|
+
respx.get("http://127.0.0.1:7000/health").mock(
|
|
19
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
20
|
+
client = DuckQueryClient(cfg)
|
|
21
|
+
assert await client.base() == "http://127.0.0.1:7000"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@respx.mock
|
|
25
|
+
async def test_none_found_raises(cfg):
|
|
26
|
+
respx.get(url__regex=r".*/health").mock(return_value=httpx.Response(503))
|
|
27
|
+
client = DuckQueryClient(cfg)
|
|
28
|
+
with pytest.raises(BackendNotFound):
|
|
29
|
+
await client.base()
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@respx.mock
|
|
33
|
+
async def test_call_unwraps_success_envelope(cfg):
|
|
34
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
35
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
36
|
+
respx.post("http://127.0.0.1:48001/api/duckdb/execute").mock(
|
|
37
|
+
return_value=httpx.Response(200, json={
|
|
38
|
+
"success": True, "data": {"row_count": 1, "data": [{"n": 16}]},
|
|
39
|
+
"messageCode": "QUERY_EXECUTED"}))
|
|
40
|
+
client = DuckQueryClient(cfg)
|
|
41
|
+
out = await client.call("POST", "/api/duckdb/execute", json_body={"sql": "SELECT 8+8 AS n"})
|
|
42
|
+
assert out["row_count"] == 1
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@respx.mock
|
|
46
|
+
async def test_call_raises_on_failure_envelope(cfg):
|
|
47
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
48
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
49
|
+
respx.post("http://127.0.0.1:48001/api/duckdb/execute").mock(
|
|
50
|
+
return_value=httpx.Response(200, json={
|
|
51
|
+
"success": False, "message": "syntax error", "messageCode": "QUERY_FAILED"}))
|
|
52
|
+
client = DuckQueryClient(cfg)
|
|
53
|
+
with pytest.raises(BackendError, match="syntax error"):
|
|
54
|
+
await client.call("POST", "/api/duckdb/execute", json_body={"sql": "SELEC 1"})
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
@respx.mock
|
|
58
|
+
async def test_call_raises_on_http_4xx(cfg):
|
|
59
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
60
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
61
|
+
respx.post("http://127.0.0.1:48001/api/duckdb/execute").mock(
|
|
62
|
+
return_value=httpx.Response(404, json={"detail": "Not Found"}))
|
|
63
|
+
client = DuckQueryClient(cfg)
|
|
64
|
+
with pytest.raises(BackendError, match="Not Found"):
|
|
65
|
+
await client.call("POST", "/api/duckdb/execute", json_body={"sql": "x"})
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
@respx.mock
|
|
69
|
+
async def test_call_4xx_surfaces_envelope_message(cfg):
|
|
70
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
71
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
72
|
+
respx.post("http://127.0.0.1:48001/api/ai/explain-sql").mock(
|
|
73
|
+
return_value=httpx.Response(400, json={
|
|
74
|
+
"success": False, "messageCode": "ai_not_configured", "message": "AI not configured"}))
|
|
75
|
+
client = DuckQueryClient(cfg)
|
|
76
|
+
with pytest.raises(BackendError, match="not configured"):
|
|
77
|
+
await client.call("POST", "/api/ai/explain-sql", json_body={"sql": "x"})
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
from duckquery_mcp.config import load_config, MODES
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def test_defaults(monkeypatch):
|
|
6
|
+
for k in ("DUCKQUERY_API_BASE", "DUCKQUERY_MCP_MODE", "DUCKQUERY_MCP_ROW_CAP"):
|
|
7
|
+
monkeypatch.delenv(k, raising=False)
|
|
8
|
+
cfg = load_config()
|
|
9
|
+
assert cfg.api_base is None
|
|
10
|
+
assert cfg.mode == "normal"
|
|
11
|
+
assert cfg.row_cap == 200
|
|
12
|
+
assert cfg.probe_ports == (48001, 8000, 8001)
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def test_env_override(monkeypatch):
|
|
16
|
+
monkeypatch.setenv("DUCKQUERY_API_BASE", "http://127.0.0.1:9999")
|
|
17
|
+
monkeypatch.setenv("DUCKQUERY_MCP_MODE", "read-only")
|
|
18
|
+
cfg = load_config()
|
|
19
|
+
assert cfg.api_base == "http://127.0.0.1:9999"
|
|
20
|
+
assert cfg.mode == "read-only"
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def test_bad_mode(monkeypatch):
|
|
24
|
+
monkeypatch.setenv("DUCKQUERY_MCP_MODE", "bogus")
|
|
25
|
+
with pytest.raises(SystemExit):
|
|
26
|
+
load_config()
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import pytest
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.config import load_config
|
|
5
|
+
from duckquery_mcp.tools.query import run_sql
|
|
6
|
+
|
|
7
|
+
pytestmark = pytest.mark.skipif(
|
|
8
|
+
not os.getenv("DUCKQUERY_INTEGRATION"),
|
|
9
|
+
reason="set DUCKQUERY_INTEGRATION=1 with a running backend")
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
async def test_run_sql_live():
|
|
13
|
+
cfg = load_config()
|
|
14
|
+
out = await run_sql(DuckQueryClient(cfg), cfg, sql="SELECT 6*7 AS n")
|
|
15
|
+
assert out["rows"] == [{"n": 42}]
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
import json
|
|
2
|
+
|
|
3
|
+
import httpx
|
|
4
|
+
import respx
|
|
5
|
+
|
|
6
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
7
|
+
from duckquery_mcp.tools.discover import list_db_objects
|
|
8
|
+
from duckquery_mcp.tools.query import federated_query
|
|
9
|
+
from duckquery_mcp.util import normalize_connection_id
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_normalize_strips_db_prefix():
|
|
13
|
+
assert normalize_connection_id("db_SORDER") == "SORDER"
|
|
14
|
+
assert normalize_connection_id("SORDER") == "SORDER"
|
|
15
|
+
assert normalize_connection_id("db_db_x") == "db_x" # only one prefix stripped
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@respx.mock
|
|
19
|
+
async def test_list_db_objects_strips_id_and_compacts(cfg):
|
|
20
|
+
base = "http://127.0.0.1:48001"
|
|
21
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
22
|
+
# must hit the STRIPPED id (SORDER), not db_SORDER
|
|
23
|
+
route = respx.get(f"{base}/api/datasources/databases/SORDER/tables").mock(
|
|
24
|
+
return_value=httpx.Response(200, json={"success": True, "data": {
|
|
25
|
+
"connection_id": "SORDER", "database": "store_order",
|
|
26
|
+
"tables": [
|
|
27
|
+
{"table_name": "t1", "table_comment": "c1", "columns": [{"name": "a"}, {"name": "b"}]},
|
|
28
|
+
{"table_name": "t2", "table_comment": "c2", "columns": [{"name": "x"}]},
|
|
29
|
+
]}}))
|
|
30
|
+
out = await list_db_objects(DuckQueryClient(cfg), cfg, connection_id="db_SORDER")
|
|
31
|
+
assert route.called
|
|
32
|
+
assert out["table_count"] == 2
|
|
33
|
+
assert out["truncated"] is False
|
|
34
|
+
assert out["tables"][0] == {"table_name": "t1", "comment": "c1", "column_count": 2}
|
|
35
|
+
assert "columns" not in out["tables"][0] # heavy per-table columns dropped
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
@respx.mock
|
|
39
|
+
async def test_federated_query_strips_connection_id(cfg):
|
|
40
|
+
base = "http://127.0.0.1:48001"
|
|
41
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
42
|
+
route = respx.post(f"{base}/api/duckdb/federated-query").mock(
|
|
43
|
+
return_value=httpx.Response(200, json={"success": True, "data": {
|
|
44
|
+
"columns": ["n"], "data": [{"n": 1}], "row_count": 1}}))
|
|
45
|
+
await federated_query(DuckQueryClient(cfg), cfg, sql="SELECT 1 AS n FROM m.t",
|
|
46
|
+
attach_databases=[{"alias": "m", "connection_id": "db_SORDER"}])
|
|
47
|
+
sent = json.loads(route.calls.last.request.content)
|
|
48
|
+
assert sent["attach_databases"][0]["connection_id"] == "SORDER"
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
import respx, httpx, pytest
|
|
2
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
3
|
+
from duckquery_mcp.tools.passthrough import duckquery_request
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@respx.mock
|
|
7
|
+
async def test_get_passthrough(cfg):
|
|
8
|
+
base = "http://127.0.0.1:48001"
|
|
9
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
10
|
+
respx.get(f"{base}/api/async-tasks").mock(
|
|
11
|
+
return_value=httpx.Response(200, json={"success": True, "data": {"tasks": []}}))
|
|
12
|
+
out = await duckquery_request(DuckQueryClient(cfg), cfg, method="GET", path="/api/async-tasks")
|
|
13
|
+
assert out == {"tasks": []}
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
async def test_non_get_needs_confirm_in_normal(cfg):
|
|
17
|
+
out = await duckquery_request(DuckQueryClient(cfg), cfg, method="DELETE",
|
|
18
|
+
path="/api/duckdb/tables/t", confirm=False)
|
|
19
|
+
assert "confirm" in out["error"].lower()
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
from duckquery_mcp.safety import is_write_sql, tool_allowed
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_read_sql():
|
|
5
|
+
assert is_write_sql("SELECT * FROM t") is False
|
|
6
|
+
assert is_write_sql(" with x as (select 1) select * from x") is False
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def test_write_sql():
|
|
10
|
+
assert is_write_sql("DROP TABLE t") is True
|
|
11
|
+
assert is_write_sql("delete from t") is True
|
|
12
|
+
assert is_write_sql("garbage") is True # unknown -> treat as write
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def test_tool_allowed_by_mode():
|
|
16
|
+
assert tool_allowed("read", "read-only") is True
|
|
17
|
+
assert tool_allowed("write", "read-only") is False
|
|
18
|
+
assert tool_allowed("write", "normal") is True
|
|
19
|
+
assert tool_allowed("write", "full") is True
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
import respx, httpx
|
|
2
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
3
|
+
from duckquery_mcp.tools.ai_settings import get_ai_settings
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@respx.mock
|
|
7
|
+
async def test_get_ai_settings_masked(cfg):
|
|
8
|
+
base = "http://127.0.0.1:48001"
|
|
9
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
10
|
+
respx.get(f"{base}/api/settings/ai").mock(
|
|
11
|
+
return_value=httpx.Response(200, json={"success": True,
|
|
12
|
+
"data": {"default_provider": "openai", "providers": [{"id": "openai", "api_key": "****"}]}}))
|
|
13
|
+
out = await get_ai_settings(DuckQueryClient(cfg), cfg)
|
|
14
|
+
assert out["providers"][0]["api_key"] == "****"
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
import respx
|
|
2
|
+
import httpx
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.tools.discover import list_tables
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@respx.mock
|
|
8
|
+
async def test_list_tables(cfg):
|
|
9
|
+
base = "http://127.0.0.1:48001"
|
|
10
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
11
|
+
respx.get(f"{base}/api/duckdb/tables").mock(
|
|
12
|
+
return_value=httpx.Response(200, json={"success": True, "data": {"tables": ["a", "b"]}}))
|
|
13
|
+
out = await list_tables(DuckQueryClient(cfg), cfg)
|
|
14
|
+
assert out == {"tables": ["a", "b"]}
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
import respx
|
|
2
|
+
import httpx
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.tools.query import run_sql
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@respx.mock
|
|
8
|
+
async def test_run_sql_returns_rows(cfg):
|
|
9
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
10
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
11
|
+
respx.post("http://127.0.0.1:48001/api/duckdb/execute").mock(
|
|
12
|
+
return_value=httpx.Response(200, json={
|
|
13
|
+
"success": True,
|
|
14
|
+
"data": {"columns": ["n"], "data": [{"n": 16}], "row_count": 1},
|
|
15
|
+
"messageCode": "QUERY_EXECUTED"}))
|
|
16
|
+
client = DuckQueryClient(cfg)
|
|
17
|
+
out = await run_sql(client, cfg, sql="SELECT 8+8 AS n")
|
|
18
|
+
assert out["row_count"] == 1
|
|
19
|
+
assert out["rows"] == [{"n": 16}]
|
|
20
|
+
assert out["truncated"] is False
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@respx.mock
|
|
24
|
+
async def test_run_sql_truncates(cfg):
|
|
25
|
+
cfg = cfg.__class__(**{**cfg.__dict__, "row_cap": 2})
|
|
26
|
+
respx.get("http://127.0.0.1:48001/health").mock(
|
|
27
|
+
return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
28
|
+
respx.post("http://127.0.0.1:48001/api/duckdb/execute").mock(
|
|
29
|
+
return_value=httpx.Response(200, json={
|
|
30
|
+
"success": True,
|
|
31
|
+
"data": {"columns": ["n"], "data": [{"n": i} for i in range(5)], "row_count": 5},
|
|
32
|
+
"messageCode": "QUERY_EXECUTED"}))
|
|
33
|
+
client = DuckQueryClient(cfg)
|
|
34
|
+
out = await run_sql(client, cfg, sql="SELECT * FROM big")
|
|
35
|
+
assert len(out["rows"]) == 2
|
|
36
|
+
assert out["truncated"] is True
|
|
37
|
+
assert out["row_count"] == 5
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
async def test_run_sql_blocks_write_in_readonly(cfg):
|
|
41
|
+
ro = cfg.__class__(**{**cfg.__dict__, "mode": "read-only"})
|
|
42
|
+
out = await run_sql(None, ro, sql="DROP TABLE t") # short-circuits before any HTTP call
|
|
43
|
+
assert "read-only" in out["error"].lower()
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
async def test_federated_query_blocks_write_in_readonly(cfg):
|
|
47
|
+
ro = cfg.__class__(**{**cfg.__dict__, "mode": "read-only"})
|
|
48
|
+
from duckquery_mcp.tools.query import federated_query
|
|
49
|
+
out = await federated_query(None, ro, sql="DROP TABLE t", attach_databases=[])
|
|
50
|
+
assert "read-only" in out["error"].lower()
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
@respx.mock
|
|
54
|
+
async def test_ask_generates_then_runs(cfg):
|
|
55
|
+
base = "http://127.0.0.1:48001"
|
|
56
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
57
|
+
respx.post(f"{base}/api/ai/nl-to-sql").mock(
|
|
58
|
+
return_value=httpx.Response(200, json={"success": True, "data": {"sql": "SELECT 1 AS n"}}))
|
|
59
|
+
respx.post(f"{base}/api/duckdb/execute").mock(
|
|
60
|
+
return_value=httpx.Response(200, json={"success": True,
|
|
61
|
+
"data": {"columns": ["n"], "data": [{"n": 1}], "row_count": 1}}))
|
|
62
|
+
from duckquery_mcp.tools.query import ask
|
|
63
|
+
client = DuckQueryClient(cfg)
|
|
64
|
+
out = await ask(client, cfg, question="how many?")
|
|
65
|
+
assert out["generated_sql"] == "SELECT 1 AS n"
|
|
66
|
+
assert out["rows"] == [{"n": 1}]
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
import respx
|
|
2
|
+
import httpx
|
|
3
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
4
|
+
from duckquery_mcp.tools.sources import add_local_file_source
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@respx.mock
|
|
8
|
+
async def test_add_local_file_source(cfg):
|
|
9
|
+
base = "http://127.0.0.1:48001"
|
|
10
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
11
|
+
route = respx.post(f"{base}/api/server-files/import").mock(
|
|
12
|
+
return_value=httpx.Response(200, json={"success": True,
|
|
13
|
+
"data": {"table_name": "sales", "row_count": 42}}))
|
|
14
|
+
out = await add_local_file_source(DuckQueryClient(cfg), cfg, path="/data/sales.csv")
|
|
15
|
+
assert out["table_name"] == "sales"
|
|
16
|
+
sent = route.calls.last.request
|
|
17
|
+
assert b"/data/sales.csv" in sent.content
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
import respx, httpx
|
|
2
|
+
from duckquery_mcp.client import DuckQueryClient
|
|
3
|
+
from duckquery_mcp.tools.transform import save_as_table
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@respx.mock
|
|
7
|
+
async def test_save_as_table(cfg):
|
|
8
|
+
base = "http://127.0.0.1:48001"
|
|
9
|
+
respx.get(f"{base}/health").mock(return_value=httpx.Response(200, json={"status": "healthy"}))
|
|
10
|
+
respx.post(f"{base}/api/save_query_to_duckdb").mock(
|
|
11
|
+
return_value=httpx.Response(200, json={"success": True, "data": {"table_name": "t2"}}))
|
|
12
|
+
out = await save_as_table(DuckQueryClient(cfg), cfg, sql="SELECT 1", table_name="t2")
|
|
13
|
+
assert out["table_name"] == "t2"
|