archery-sql-mcp 0.2.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.
@@ -0,0 +1,6 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .venv/
4
+ dist/
5
+ *.egg-info/
6
+ uv.lock
@@ -0,0 +1,130 @@
1
+ Metadata-Version: 2.5
2
+ Name: archery-sql-mcp
3
+ Version: 0.2.0
4
+ Summary: Archery SQL 查询 MCP 工具:自动登录 + 会话自愈 + 只读查询,供任意 MCP 客户端与 CLI 使用
5
+ License: MIT
6
+ Requires-Python: >=3.10
7
+ Requires-Dist: httpx>=0.27
8
+ Requires-Dist: mcp>=1.2
9
+ Description-Content-Type: text/markdown
10
+
11
+ # archery-sql-mcp
12
+
13
+ Archery SQL 平台(sql.tec-develop.com)查询工具:**自动登录 + 会话自愈 + 只读查询**。
14
+ 双模式运行——无参数启动为 MCP stdio server(任何 MCP 客户端可接入),带子命令则为 CLI。
15
+
16
+ ## 安装
17
+
18
+ 发布到 PyPI 后(最短路径):
19
+
20
+ ```bash
21
+ uvx archery-sql-mcp --help # 免安装直接跑
22
+ uv tool install archery-sql-mcp # 或安装为全局命令
23
+ ```
24
+
25
+ 从内部 Git 仓库:
26
+
27
+ ```bash
28
+ uv tool install git+https://git.tec-do.cn/<group>/archery-mcp.git
29
+ ```
30
+
31
+ 本地开发:
32
+
33
+ ```bash
34
+ git clone <仓库地址> && cd archery-mcp
35
+ uv run --project . archery-sql-mcp --help
36
+ ```
37
+
38
+ ## 初始化账号密码(一次性)
39
+
40
+ ```bash
41
+ archery-sql-mcp init
42
+ ```
43
+
44
+ 交互式录入 Archery 地址 / 用户名 / 密码(不回显)/ 默认实例 / 默认库,
45
+ 写入 `~/.archery-mcp/config.json` 并立即登录验证。
46
+
47
+ > 密码为本地明文存储(仅存于本机用户目录,不进任何 git 仓库)。
48
+ > 不想落盘的机器可改用环境变量:`ARCHERY_USERNAME` / `ARCHERY_PASSWORD` /
49
+ > `ARCHERY_BASE_URL` / `ARCHERY_INSTANCE` / `ARCHERY_DB`。
50
+
51
+ 会话缓存于 `~/.archery-mcp/session.json`,冷启动自动复用(最长约两周),
52
+ 过期自动重新登录,调用方无感。
53
+
54
+ ## 自检与使用
55
+
56
+ ```bash
57
+ archery-sql-mcp check # 登录 + SELECT 1 全链路验证
58
+ archery-sql-mcp instances # 列出有读权限的数据库实例
59
+ archery-sql-mcp query "SELECT id, status, error_message FROM cv_video_script_audit_tasks ORDER BY created_at DESC LIMIT 10"
60
+ archery-sql-mcp query "SELECT count(*) FROM open_api_request_logs WHERE created_at > now() - interval '1 day'" --limit 10
61
+ ```
62
+
63
+ 仅允许 `SELECT` / `WITH` 开头的语句(可在配置 `allowed_prefixes` 扩展)。
64
+
65
+ ## 接入 MCP 客户端
66
+
67
+ ZCode(`~/.zcode/mcp.json`)、Claude Desktop、Cursor(`~/.cursor/mcp.json`)同构:
68
+
69
+ ```jsonc
70
+ {
71
+ "mcpServers": {
72
+ "archery": {
73
+ "command": "uvx",
74
+ "args": ["--from", "git+https://git.tec-do.cn/<group>/archery-mcp.git", "archery-sql-mcp"]
75
+ }
76
+ }
77
+ }
78
+ ```
79
+
80
+ Claude Code:
81
+
82
+ ```bash
83
+ claude mcp add archery -- uvx --from git+https://git.tec-do.cn/<group>/archery-mcp.git archery-sql-mcp
84
+ ```
85
+
86
+ 接入后 agent 可调用 5 个工具:
87
+
88
+ | 工具 | 用途 |
89
+ |---|---|
90
+ | `query(sql, db_name?, instance_name?, limit=100)` | 只读查询(仅 SELECT / WITH),返回 `{status, columns, rows, row_count, elapsed_ms}` |
91
+ | `get_token()` | 登录并返回 token 概要(掩码显示,完整值在本机 session.json) |
92
+ | `refresh_token()` | 强制刷新 token(清缓存重新登录) |
93
+ | `list_instances()` | 列出有读权限(can_read)的数据库实例 |
94
+ | `instance_resources(resource_type, db_name?, instance_name?, schema_name?, tb_name?)` | 探查实例结构:schema / table / column |
95
+
96
+ ## 配置说明
97
+
98
+ 配置优先级:**环境变量 > `~/.archery-mcp/config.json` > 内置默认**。
99
+ 配置文件支持多 profile(`ARCHERY_PROFILE` 选择):
100
+
101
+ ```json
102
+ {
103
+ "default_profile": "prod",
104
+ "profiles": {
105
+ "prod": {
106
+ "base_url": "https://sql.tec-develop.com",
107
+ "username": "xxx",
108
+ "password": "yyy",
109
+ "instance_name": "alisg-haiyun-powerdata-pgsql-prod-01",
110
+ "db_name": "creativault_business",
111
+ "limit_num": 100,
112
+ "allowed_prefixes": ["select", "with"],
113
+ "poll_interval_s": 1,
114
+ "query_timeout_s": 60
115
+ }
116
+ }
117
+ }
118
+ ```
119
+
120
+ ## 错误行为
121
+
122
+ | 场景 | 行为 |
123
+ |---|---|
124
+ | session 过期 | 自动重登并重放请求,调用方无感 |
125
+ | 密码错 / 账号禁用 | 明确报错,不重试 |
126
+ | 非 SELECT 语句 | 客户端预检拒绝 |
127
+ | SQL 语法错 | 透传数据库错误原文 |
128
+ | 查询超时 | 报错并附 query_id |
129
+
130
+ 密码与 sessionid 永不出现在日志和工具返回值中。
@@ -0,0 +1,120 @@
1
+ # archery-sql-mcp
2
+
3
+ Archery SQL 平台(sql.tec-develop.com)查询工具:**自动登录 + 会话自愈 + 只读查询**。
4
+ 双模式运行——无参数启动为 MCP stdio server(任何 MCP 客户端可接入),带子命令则为 CLI。
5
+
6
+ ## 安装
7
+
8
+ 发布到 PyPI 后(最短路径):
9
+
10
+ ```bash
11
+ uvx archery-sql-mcp --help # 免安装直接跑
12
+ uv tool install archery-sql-mcp # 或安装为全局命令
13
+ ```
14
+
15
+ 从内部 Git 仓库:
16
+
17
+ ```bash
18
+ uv tool install git+https://git.tec-do.cn/<group>/archery-mcp.git
19
+ ```
20
+
21
+ 本地开发:
22
+
23
+ ```bash
24
+ git clone <仓库地址> && cd archery-mcp
25
+ uv run --project . archery-sql-mcp --help
26
+ ```
27
+
28
+ ## 初始化账号密码(一次性)
29
+
30
+ ```bash
31
+ archery-sql-mcp init
32
+ ```
33
+
34
+ 交互式录入 Archery 地址 / 用户名 / 密码(不回显)/ 默认实例 / 默认库,
35
+ 写入 `~/.archery-mcp/config.json` 并立即登录验证。
36
+
37
+ > 密码为本地明文存储(仅存于本机用户目录,不进任何 git 仓库)。
38
+ > 不想落盘的机器可改用环境变量:`ARCHERY_USERNAME` / `ARCHERY_PASSWORD` /
39
+ > `ARCHERY_BASE_URL` / `ARCHERY_INSTANCE` / `ARCHERY_DB`。
40
+
41
+ 会话缓存于 `~/.archery-mcp/session.json`,冷启动自动复用(最长约两周),
42
+ 过期自动重新登录,调用方无感。
43
+
44
+ ## 自检与使用
45
+
46
+ ```bash
47
+ archery-sql-mcp check # 登录 + SELECT 1 全链路验证
48
+ archery-sql-mcp instances # 列出有读权限的数据库实例
49
+ archery-sql-mcp query "SELECT id, status, error_message FROM cv_video_script_audit_tasks ORDER BY created_at DESC LIMIT 10"
50
+ archery-sql-mcp query "SELECT count(*) FROM open_api_request_logs WHERE created_at > now() - interval '1 day'" --limit 10
51
+ ```
52
+
53
+ 仅允许 `SELECT` / `WITH` 开头的语句(可在配置 `allowed_prefixes` 扩展)。
54
+
55
+ ## 接入 MCP 客户端
56
+
57
+ ZCode(`~/.zcode/mcp.json`)、Claude Desktop、Cursor(`~/.cursor/mcp.json`)同构:
58
+
59
+ ```jsonc
60
+ {
61
+ "mcpServers": {
62
+ "archery": {
63
+ "command": "uvx",
64
+ "args": ["--from", "git+https://git.tec-do.cn/<group>/archery-mcp.git", "archery-sql-mcp"]
65
+ }
66
+ }
67
+ }
68
+ ```
69
+
70
+ Claude Code:
71
+
72
+ ```bash
73
+ claude mcp add archery -- uvx --from git+https://git.tec-do.cn/<group>/archery-mcp.git archery-sql-mcp
74
+ ```
75
+
76
+ 接入后 agent 可调用 5 个工具:
77
+
78
+ | 工具 | 用途 |
79
+ |---|---|
80
+ | `query(sql, db_name?, instance_name?, limit=100)` | 只读查询(仅 SELECT / WITH),返回 `{status, columns, rows, row_count, elapsed_ms}` |
81
+ | `get_token()` | 登录并返回 token 概要(掩码显示,完整值在本机 session.json) |
82
+ | `refresh_token()` | 强制刷新 token(清缓存重新登录) |
83
+ | `list_instances()` | 列出有读权限(can_read)的数据库实例 |
84
+ | `instance_resources(resource_type, db_name?, instance_name?, schema_name?, tb_name?)` | 探查实例结构:schema / table / column |
85
+
86
+ ## 配置说明
87
+
88
+ 配置优先级:**环境变量 > `~/.archery-mcp/config.json` > 内置默认**。
89
+ 配置文件支持多 profile(`ARCHERY_PROFILE` 选择):
90
+
91
+ ```json
92
+ {
93
+ "default_profile": "prod",
94
+ "profiles": {
95
+ "prod": {
96
+ "base_url": "https://sql.tec-develop.com",
97
+ "username": "xxx",
98
+ "password": "yyy",
99
+ "instance_name": "alisg-haiyun-powerdata-pgsql-prod-01",
100
+ "db_name": "creativault_business",
101
+ "limit_num": 100,
102
+ "allowed_prefixes": ["select", "with"],
103
+ "poll_interval_s": 1,
104
+ "query_timeout_s": 60
105
+ }
106
+ }
107
+ }
108
+ ```
109
+
110
+ ## 错误行为
111
+
112
+ | 场景 | 行为 |
113
+ |---|---|
114
+ | session 过期 | 自动重登并重放请求,调用方无感 |
115
+ | 密码错 / 账号禁用 | 明确报错,不重试 |
116
+ | 非 SELECT 语句 | 客户端预检拒绝 |
117
+ | SQL 语法错 | 透传数据库错误原文 |
118
+ | 查询超时 | 报错并附 query_id |
119
+
120
+ 密码与 sessionid 永不出现在日志和工具返回值中。
@@ -0,0 +1,21 @@
1
+ [project]
2
+ name = "archery-sql-mcp"
3
+ version = "0.2.0"
4
+ description = "Archery SQL 查询 MCP 工具:自动登录 + 会话自愈 + 只读查询,供任意 MCP 客户端与 CLI 使用"
5
+ readme = "README.md"
6
+ requires-python = ">=3.10"
7
+ license = { text = "MIT" }
8
+ dependencies = [
9
+ "httpx>=0.27",
10
+ "mcp>=1.2",
11
+ ]
12
+
13
+ [project.scripts]
14
+ archery-sql-mcp = "archery_sql_mcp:main"
15
+
16
+ [build-system]
17
+ requires = ["hatchling"]
18
+ build-backend = "hatchling.build"
19
+
20
+ [tool.hatch.build.targets.wheel]
21
+ packages = ["src/archery_sql_mcp"]
@@ -0,0 +1,613 @@
1
+ """Archery SQL 查询 MCP 工具(单模块实现)。
2
+
3
+ 面向 Archery(Django)SQL 审核查询平台的只读查询封装:
4
+ - 自动登录(Django CSRF 协议:csrftoken cookie + csrfmiddlewaretoken 表单)
5
+ - 会话缓存(内存 + 磁盘),失效自动重登并重放请求
6
+ - 两步异步查询(POST /query/ 入队 → GET /queryresult/ 轮询至终态)
7
+
8
+ 双模式运行:
9
+ - 无参数启动 → MCP stdio server(任何 MCP host 可接入)
10
+ - 子命令启动 → CLI:init / check / query
11
+
12
+ 配置优先级:环境变量 > 配置文件(~/.archery-mcp/config.json,多 profile)> 内置默认。
13
+ 密码仅存本地配置文件(init 时明示风险),日志与工具返回值中永不出现密码 / sessionid。
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import argparse
19
+ import asyncio
20
+ import getpass
21
+ import json
22
+ import os
23
+ import re
24
+ import sys
25
+ import time
26
+ from pathlib import Path
27
+ from typing import Any, Optional
28
+
29
+ import httpx
30
+
31
+ VERSION = "0.2.0"
32
+
33
+ # ── 配置文件与会话缓存位置(可用 ARCHERY_CONFIG_DIR 整体改根目录) ──────────
34
+ _CONFIG_DIR = Path(os.environ.get("ARCHERY_CONFIG_DIR", str(Path.home() / ".archery-mcp")))
35
+ CONFIG_FILE = Path(os.environ.get("ARCHERY_CONFIG", str(_CONFIG_DIR / "config.json")))
36
+ SESSION_FILE = _CONFIG_DIR / "session.json"
37
+
38
+ # 内置默认值:公司线上库
39
+ _DEFAULTS: dict[str, Any] = {
40
+ "base_url": "https://sql.tec-develop.com",
41
+ "instance_name": "alisg-haiyun-powerdata-pgsql-prod-01",
42
+ "db_name": "creativault_business",
43
+ "limit_num": 100,
44
+ "allowed_prefixes": ["select", "with"],
45
+ "poll_interval_s": 1.0,
46
+ "query_timeout_s": 60.0,
47
+ "request_timeout_s": 30.0,
48
+ }
49
+
50
+ MAX_LIMIT = 1000
51
+
52
+ # 环境变量 → 配置键(优先级最高,不落盘,适合不想存密码的机器)
53
+ _ENV_MAP = {
54
+ "ARCHERY_BASE_URL": "base_url",
55
+ "ARCHERY_USERNAME": "username",
56
+ "ARCHERY_PASSWORD": "password",
57
+ "ARCHERY_INSTANCE": "instance_name",
58
+ "ARCHERY_DB": "db_name",
59
+ }
60
+
61
+ _CSRF_INPUT_RE = re.compile(r'name="csrfmiddlewaretoken"[^>]*value="([^"]+)"')
62
+
63
+
64
+ class ArcheryError(Exception):
65
+ """工具级错误,message 面向调用方展示。"""
66
+
67
+
68
+ # ── 配置 ──────────────────────────────────────────────────────────────
69
+
70
+ def load_profile() -> dict[str, Any]:
71
+ """合并三层配置:环境变量 > 配置文件 profile > 内置默认。"""
72
+ merged = dict(_DEFAULTS)
73
+ if CONFIG_FILE.exists():
74
+ try:
75
+ data = json.loads(CONFIG_FILE.read_text(encoding="utf-8"))
76
+ except json.JSONDecodeError as e:
77
+ raise ArcheryError(f"配置文件 JSON 解析失败({CONFIG_FILE}): {e}")
78
+ profile_name = os.environ.get("ARCHERY_PROFILE") or data.get("default_profile")
79
+ profiles = data.get("profiles") or {}
80
+ if profile_name and profile_name in profiles:
81
+ merged.update({k: v for k, v in profiles[profile_name].items() if v not in (None, "")})
82
+ elif profile_name:
83
+ raise ArcheryError(f"profile '{profile_name}' 不存在,可用: {list(profiles) or '无'}")
84
+ for env, key in _ENV_MAP.items():
85
+ val = os.environ.get(env)
86
+ if val:
87
+ merged[key] = val
88
+ if not merged.get("username") or not merged.get("password"):
89
+ raise ArcheryError(
90
+ f"未配置账号密码:请先运行 `archery-sql-mcp init`(写入 {CONFIG_FILE}),"
91
+ "或通过环境变量 ARCHERY_USERNAME / ARCHERY_PASSWORD 注入。"
92
+ )
93
+ return merged
94
+
95
+
96
+ # ── 会话管理:自动登录 + 缓存 + 失效重登 ────────────────────────────────
97
+
98
+ class SessionManager:
99
+ """Archery 登录态管理。Django session 过期表现为 302 → /login/。"""
100
+
101
+ def __init__(self, cfg: dict[str, Any]):
102
+ self.cfg = cfg
103
+ self.base = str(cfg["base_url"]).rstrip("/")
104
+ self.sessionid: Optional[str] = None
105
+ self.csrftoken: Optional[str] = None
106
+
107
+ # -- 磁盘缓存:冷启动优先复用(Django session 默认两周有效) --
108
+ def _load_cache(self) -> bool:
109
+ try:
110
+ if SESSION_FILE.exists():
111
+ data = json.loads(SESSION_FILE.read_text(encoding="utf-8"))
112
+ self.sessionid = data.get("sessionid")
113
+ self.csrftoken = data.get("csrftoken")
114
+ except Exception:
115
+ pass
116
+ return bool(self.sessionid)
117
+
118
+ def _save_cache(self) -> None:
119
+ try:
120
+ _CONFIG_DIR.mkdir(parents=True, exist_ok=True)
121
+ SESSION_FILE.write_text(
122
+ json.dumps(
123
+ {"sessionid": self.sessionid, "csrftoken": self.csrftoken,
124
+ "logged_in_at": time.strftime("%Y-%m-%d %H:%M:%S")},
125
+ ensure_ascii=False, indent=2,
126
+ ),
127
+ encoding="utf-8",
128
+ )
129
+ except Exception:
130
+ pass # 缓存写失败不影响本次会话
131
+
132
+ def invalidate(self) -> None:
133
+ self.sessionid = None
134
+ self.csrftoken = None
135
+ try:
136
+ SESSION_FILE.unlink(missing_ok=True)
137
+ except Exception:
138
+ pass
139
+
140
+ async def ensure(self) -> None:
141
+ if self.sessionid:
142
+ return
143
+ if self._load_cache():
144
+ return
145
+ await self.login()
146
+
147
+ async def login(self) -> None:
148
+ """Django CSRF 登录:GET /login/ 取 token → POST /authenticate/ → sessionid。"""
149
+ async with httpx.AsyncClient(
150
+ timeout=self.cfg["request_timeout_s"], follow_redirects=False
151
+ ) as client:
152
+ page = await client.get(f"{self.base}/login/")
153
+ page.raise_for_status()
154
+ m = _CSRF_INPUT_RE.search(page.text)
155
+ form_token = m.group(1) if m else ""
156
+ cookie_csrf = client.cookies.get("csrftoken") or ""
157
+ if not form_token or not cookie_csrf:
158
+ raise ArcheryError("登录页解析失败:未取到 csrfmiddlewaretoken / csrftoken")
159
+
160
+ resp = await client.post(
161
+ f"{self.base}/authenticate/",
162
+ data={
163
+ "username": self.cfg["username"],
164
+ "password": self.cfg["password"],
165
+ "csrfmiddlewaretoken": form_token,
166
+ },
167
+ headers={"Referer": f"{self.base}/login/"},
168
+ )
169
+ # 成功:302 跳往非 /login 路径;失败:302 回 /login 或 200 带错误页
170
+ location = resp.headers.get("location", "")
171
+ if resp.status_code not in (301, 302) or "/login" in location:
172
+ reason = _extract_login_error(resp.text) or (
173
+ f"HTTP {resp.status_code}"
174
+ + (f" → {location}" if location else "")
175
+ + "(用户名/密码错误,或账号被禁用)"
176
+ )
177
+ raise ArcheryError(f"Archery 登录失败: {reason}")
178
+
179
+ self.sessionid = client.cookies.get("sessionid")
180
+ self.csrftoken = client.cookies.get("csrftoken") or cookie_csrf
181
+ if not self.sessionid:
182
+ raise ArcheryError("登录成功但未取到 sessionid cookie")
183
+ self._save_cache()
184
+
185
+ def is_expired(self, resp: httpx.Response) -> bool:
186
+ """会话失效判定:302 重定向到登录页(Archery/Django 的'401'形态)。"""
187
+ if resp.status_code in (301, 302):
188
+ return "login" in (resp.headers.get("location") or "").lower()
189
+ return False
190
+
191
+ async def token_info(self, refresh: bool = False) -> dict[str, Any]:
192
+ """确保已登录并返回 token 概要。refresh=True 时强制重登(即'刷新 token')。
193
+
194
+ token(sessionid)以掩码返回:完整值仅存于本机 session.json,
195
+ 避免明文凭据进入 AI 对话历史 / 日志。
196
+ """
197
+ if refresh:
198
+ self.invalidate()
199
+ await self.login()
200
+ else:
201
+ await self.ensure()
202
+ sid = self.sessionid or ""
203
+ masked = (sid[:4] + "***" + sid[-4:]) if len(sid) > 8 else "***"
204
+ logged_in_at = ""
205
+ try:
206
+ if SESSION_FILE.exists():
207
+ logged_in_at = json.loads(
208
+ SESSION_FILE.read_text(encoding="utf-8")
209
+ ).get("logged_in_at", "")
210
+ except Exception:
211
+ pass
212
+ return {
213
+ "logged_in": True,
214
+ "token_masked": masked,
215
+ "logged_in_at": logged_in_at,
216
+ "base_url": self.base,
217
+ "note": f"完整 token 已缓存于本机 {SESSION_FILE},仅供本机调试使用",
218
+ }
219
+
220
+
221
+ def _extract_login_error(html: str) -> Optional[str]:
222
+ """尽力从登录失败页面提取可读错误(不同 Archery 版本模板不同)。"""
223
+ for pat in (r'class="[^"]*alert[^"]*"[^>]*>([^<]{4,120})<',
224
+ r'<div class="text-danger[^"]*">\s*([^<]{4,120})'):
225
+ m = re.search(pat, html)
226
+ if m:
227
+ return m.group(1).strip()
228
+ return None
229
+
230
+
231
+ # ── 查询客户端:两步异步 + 失效自动重登重放 ────────────────────────────
232
+
233
+ class ArcheryClient:
234
+ def __init__(self, cfg: dict[str, Any]):
235
+ self.cfg = cfg
236
+ self.sm = SessionManager(cfg)
237
+
238
+ async def _authed(
239
+ self,
240
+ method: str,
241
+ path: str,
242
+ *,
243
+ data: Optional[dict] = None,
244
+ params: Optional[dict] = None,
245
+ ) -> httpx.Response:
246
+ """带会话的请求;失效自动重登一次并重放(密码错等确定性失败不会循环)。"""
247
+ last_resp: Optional[httpx.Response] = None
248
+ for _attempt in range(2):
249
+ await self.sm.ensure()
250
+ async with httpx.AsyncClient(
251
+ timeout=self.cfg["request_timeout_s"], follow_redirects=False
252
+ ) as client:
253
+ resp = await client.request(
254
+ method, f"{self.sm.base}{path}",
255
+ data=data, params=params,
256
+ cookies={"csrftoken": self.sm.csrftoken or "",
257
+ "sessionid": self.sm.sessionid or ""},
258
+ headers={"X-CSRFToken": self.sm.csrftoken or "",
259
+ "Referer": self.sm.base + "/"},
260
+ )
261
+ if not self.sm.is_expired(resp):
262
+ return resp
263
+ last_resp = resp
264
+ self.sm.invalidate() # 触发下一轮 ensure() 重新登录
265
+ raise ArcheryError(
266
+ "会话失效且自动重新登录后仍未通过(账号可能被禁用或密码已变更),"
267
+ f"末次响应 HTTP {last_resp.status_code if last_resp else '?'}"
268
+ )
269
+
270
+ def _check_sql_allowed(self, sql: str) -> None:
271
+ body = re.sub(r"--[^\n]*", " ", sql)
272
+ body = re.sub(r"/\*.*?\*/", " ", body, flags=re.S)
273
+ first = body.strip().split(None, 1)[0].lower() if body.strip() else ""
274
+ allowed = [str(p).lower() for p in self.cfg.get("allowed_prefixes", [])]
275
+ if first not in allowed:
276
+ raise ArcheryError(
277
+ f"仅允许 {'/'.join(allowed)} 开头的只读语句,收到: '{first or '(空)'}'"
278
+ )
279
+
280
+ async def query(
281
+ self,
282
+ sql: str,
283
+ *,
284
+ db_name: Optional[str] = None,
285
+ instance_name: Optional[str] = None,
286
+ limit: Optional[int] = None,
287
+ ) -> dict[str, Any]:
288
+ self._check_sql_allowed(sql)
289
+ limit = min(max(int(limit or self.cfg["limit_num"]), 1), MAX_LIMIT)
290
+ started = time.monotonic()
291
+
292
+ resp = await self._authed(
293
+ "POST", "/query/",
294
+ data={
295
+ "instance_name": instance_name or self.cfg["instance_name"],
296
+ "db_name": db_name or self.cfg["db_name"],
297
+ "schema_name": "",
298
+ "tb_name": "",
299
+ "sql_content": sql,
300
+ "limit_num": limit,
301
+ },
302
+ )
303
+ try:
304
+ submitted = resp.json()
305
+ except Exception:
306
+ raise ArcheryError(f"/query/ 返回非 JSON (HTTP {resp.status_code}),请联系管理员确认平台版本兼容性")
307
+ if submitted.get("status") != 0:
308
+ raise ArcheryError(f"/query/ 提交失败: {submitted.get('msg') or submitted}")
309
+ query_id = (submitted.get("data") or {}).get("query_id")
310
+ if not query_id:
311
+ raise ArcheryError(f"/query/ 未返回 query_id: {submitted}")
312
+
313
+ # 轮询至终态
314
+ interval = float(self.cfg["poll_interval_s"])
315
+ timeout = float(self.cfg["query_timeout_s"])
316
+ while True:
317
+ resp = await self._authed("GET", "/queryresult/", params={"sql_query_id": query_id})
318
+ try:
319
+ payload = resp.json()
320
+ except Exception:
321
+ raise ArcheryError(f"/queryresult/ 返回非 JSON (HTTP {resp.status_code})")
322
+ data = payload.get("data") or {}
323
+ status = data.get("status", "")
324
+ if status == "finished":
325
+ err = data.get("err")
326
+ if err:
327
+ raise ArcheryError(f"SQL 执行错误: {err}")
328
+ columns = data.get("column_list") or []
329
+ rows = data.get("rows") or []
330
+ return {
331
+ "status": "finished",
332
+ "columns": columns,
333
+ "rows": rows,
334
+ "row_count": len(rows),
335
+ "elapsed_ms": int((time.monotonic() - started) * 1000),
336
+ }
337
+ if status in ("failed", "aborted"):
338
+ raise ArcheryError(f"查询失败: {data.get('err') or data}")
339
+ if time.monotonic() - started > timeout:
340
+ raise ArcheryError(
341
+ f"查询超时(>{timeout:.0f}s),query_id={query_id},可稍后重试或缩小结果集"
342
+ )
343
+ await asyncio.sleep(interval)
344
+
345
+ async def _get_json(self, path: str, params: Optional[dict] = None) -> Any:
346
+ """GET 平台接口并解包:status!=0 报错,否则返回 data(无 data 时返回原 payload)。"""
347
+ resp = await self._authed("GET", path, params=params)
348
+ try:
349
+ payload = resp.json()
350
+ except Exception:
351
+ raise ArcheryError(f"{path} 返回非 JSON (HTTP {resp.status_code})")
352
+ if payload.get("status") != 0:
353
+ raise ArcheryError(f"{path} 调用失败: {payload.get('msg') or payload}")
354
+ return payload.get("data", payload)
355
+
356
+ async def list_instances(self) -> Any:
357
+ """列出当前账号有读权限(can_read)的所有数据库实例。"""
358
+ # tag_codes[] 是数组参数:tag_codes%5B%5D=can_read
359
+ return await self._get_json(
360
+ "/group/user_all_instances/", params={"tag_codes[]": ["can_read"]}
361
+ )
362
+
363
+ async def instance_resources(
364
+ self,
365
+ *,
366
+ instance_name: Optional[str] = None,
367
+ db_name: Optional[str] = None,
368
+ resource_type: str = "schema",
369
+ schema_name: Optional[str] = None,
370
+ tb_name: Optional[str] = None,
371
+ ) -> Any:
372
+ """查询实例资源:resource_type=schema 列 schema;table 列表;column 需 schema_name+tb_name。"""
373
+ params: dict[str, Any] = {
374
+ "instance_name": instance_name or self.cfg["instance_name"],
375
+ "db_name": db_name or self.cfg["db_name"],
376
+ "resource_type": resource_type,
377
+ }
378
+ if schema_name:
379
+ params["schema_name"] = schema_name
380
+ if tb_name:
381
+ params["tb_name"] = tb_name
382
+ return await self._get_json("/instance/instance_resource/", params=params)
383
+
384
+
385
+ # ── CLI ───────────────────────────────────────────────────────────────
386
+
387
+ def _ensure_utf8_stdout() -> None:
388
+ try:
389
+ sys.stdout.reconfigure(encoding="utf-8") # type: ignore[union-attr]
390
+ sys.stderr.reconfigure(encoding="utf-8") # type: ignore[union-attr]
391
+ except Exception:
392
+ pass
393
+
394
+
395
+ def cmd_init(_args: argparse.Namespace) -> int:
396
+ _ensure_utf8_stdout()
397
+ print(f"将生成配置文件: {CONFIG_FILE}(密码明文存放于本机,不会进入任何 git 仓库)")
398
+ existing: dict[str, Any] = {}
399
+ if CONFIG_FILE.exists():
400
+ try:
401
+ existing = json.loads(CONFIG_FILE.read_text(encoding="utf-8"))
402
+ print("检测到已有配置,回车保留原值。")
403
+ except Exception:
404
+ existing = {}
405
+
406
+ old = existing.get("profiles", {}).get("prod", {})
407
+
408
+ def ask(prompt: str, default: str = "", secret: bool = False) -> str:
409
+ suffix = f" [{default}]" if default else ""
410
+ val = (getpass.getpass(prompt + suffix + ": ") if secret
411
+ else input(prompt + suffix + ": ")).strip()
412
+ return val or default
413
+
414
+ profile = {
415
+ "base_url": ask("Archery 地址", old.get("base_url", _DEFAULTS["base_url"])),
416
+ "username": ask("用户名", old.get("username", "")),
417
+ "password": ask("密码", old.get("password", ""), secret=True),
418
+ "instance_name": ask("默认实例", old.get("instance_name", _DEFAULTS["instance_name"])),
419
+ "db_name": ask("默认库", old.get("db_name", _DEFAULTS["db_name"])),
420
+ }
421
+ config = {"default_profile": "prod", "profiles": {"prod": profile}}
422
+ _CONFIG_DIR.mkdir(parents=True, exist_ok=True)
423
+ CONFIG_FILE.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8")
424
+ print("配置已写入,正在验证登录…")
425
+
426
+ ok = asyncio.run(_login_check(profile))
427
+ if ok:
428
+ print("初始化完成 ✓ 可运行 `archery-sql-mcp check` 做全链路自检。")
429
+ return 0
430
+ print("登录验证失败:请检查上面的错误,重新运行 init 修正。")
431
+ return 1
432
+
433
+
434
+ async def _login_check(profile: dict[str, Any]) -> bool:
435
+ cfg = dict(_DEFAULTS)
436
+ cfg.update({k: v for k, v in profile.items() if v})
437
+ try:
438
+ sm = SessionManager(cfg)
439
+ await sm.login()
440
+ print(f"登录成功(session 已缓存到 {SESSION_FILE})")
441
+ return True
442
+ except ArcheryError as e:
443
+ print(f"错误: {e}")
444
+ return False
445
+
446
+
447
+ def cmd_check(_args: argparse.Namespace) -> int:
448
+ _ensure_utf8_stdout()
449
+ try:
450
+ cfg = load_profile()
451
+ result = asyncio.run(ArcheryClient(cfg).query("SELECT 1 AS ok", limit=1))
452
+ print(f"登录 + 查询链路正常 ✓ SELECT 1 → {result['rows']} 耗时 {result['elapsed_ms']}ms")
453
+ return 0
454
+ except ArcheryError as e:
455
+ print(f"错误: {e}")
456
+ return 1
457
+
458
+
459
+ def cmd_query(args: argparse.Namespace) -> int:
460
+ _ensure_utf8_stdout()
461
+ try:
462
+ cfg = load_profile()
463
+ result = asyncio.run(ArcheryClient(cfg).query(
464
+ args.sql, db_name=args.db, instance_name=args.instance, limit=args.limit,
465
+ ))
466
+ print(json.dumps(result, ensure_ascii=False, indent=2))
467
+ return 0
468
+ except ArcheryError as e:
469
+ print(f"错误: {e}")
470
+ return 1
471
+
472
+
473
+ def _new_mcp_server(name: str):
474
+ """兼容 mcp SDK 1.x/2.x:2.x 中 FastMCP 改名为 MCPServer,用法一致。"""
475
+ try:
476
+ from mcp.server.mcpserver import MCPServer as Server # mcp >= 2
477
+ except ImportError:
478
+ from mcp.server.fastmcp import FastMCP as Server # mcp 1.x
479
+ return Server(name)
480
+
481
+
482
+ def cmd_instances(_args: argparse.Namespace) -> int:
483
+ _ensure_utf8_stdout()
484
+ try:
485
+ cfg = load_profile()
486
+ data = asyncio.run(ArcheryClient(cfg).list_instances())
487
+ print(json.dumps({"instances": data}, ensure_ascii=False, indent=2))
488
+ return 0
489
+ except ArcheryError as e:
490
+ print(f"错误: {e}")
491
+ return 1
492
+
493
+
494
+ def run_mcp() -> None:
495
+ """MCP stdio server 模式(惰性 import,CLI 场景无需安装 mcp 包)。"""
496
+ mcp = _new_mcp_server("archery")
497
+
498
+ @mcp.tool()
499
+ async def query(
500
+ sql: str,
501
+ db_name: Optional[str] = None,
502
+ instance_name: Optional[str] = None,
503
+ limit: int = 100,
504
+ ) -> str:
505
+ """查询 Archery SQL 平台上的数据库(只读,仅允许 SELECT / WITH)。
506
+
507
+ 默认连接线上生产库:实例 alisg-haiyun-powerdata-pgsql-prod-01,
508
+ 库 creativault_business。可通过 db_name / instance_name 参数切换。
509
+ 返回 JSON:{status, columns, rows, row_count, elapsed_ms}。
510
+
511
+ 典型用途:排查线上业务数据(任务状态、错误信息、日志表等)。
512
+ """
513
+ try:
514
+ cfg = load_profile()
515
+ result = await ArcheryClient(cfg).query(
516
+ sql, db_name=db_name, instance_name=instance_name, limit=limit,
517
+ )
518
+ return json.dumps(result, ensure_ascii=False)
519
+ except ArcheryError as e:
520
+ return json.dumps({"error": str(e)}, ensure_ascii=False)
521
+
522
+ @mcp.tool()
523
+ async def get_token() -> str:
524
+ """获取 Archery 登录 token(执行/复用登录,返回会话状态与掩码 token)。
525
+
526
+ token 仅以掩码显示,完整值缓存于本机 ~/.archery-mcp/session.json。
527
+ """
528
+ try:
529
+ cfg = load_profile()
530
+ info = await SessionManager(cfg).token_info(refresh=False)
531
+ return json.dumps(info, ensure_ascii=False)
532
+ except ArcheryError as e:
533
+ return json.dumps({"error": str(e)}, ensure_ascii=False)
534
+
535
+ @mcp.tool()
536
+ async def refresh_token() -> str:
537
+ """强制刷新 Archery 登录 token(清空缓存并重新登录,返回新会话状态)。"""
538
+ try:
539
+ cfg = load_profile()
540
+ info = await SessionManager(cfg).token_info(refresh=True)
541
+ return json.dumps(info, ensure_ascii=False)
542
+ except ArcheryError as e:
543
+ return json.dumps({"error": str(e)}, ensure_ascii=False)
544
+
545
+ @mcp.tool()
546
+ async def list_instances() -> str:
547
+ """列出你在 Archery 平台上有读权限(can_read)的所有数据库实例。
548
+
549
+ 用于 query 工具切换 instance_name 前确认实例名。
550
+ """
551
+ try:
552
+ cfg = load_profile()
553
+ data = await ArcheryClient(cfg).list_instances()
554
+ return json.dumps({"instances": data}, ensure_ascii=False)
555
+ except ArcheryError as e:
556
+ return json.dumps({"error": str(e)}, ensure_ascii=False)
557
+
558
+ @mcp.tool()
559
+ async def instance_resources(
560
+ resource_type: str = "schema",
561
+ db_name: Optional[str] = None,
562
+ instance_name: Optional[str] = None,
563
+ schema_name: Optional[str] = None,
564
+ tb_name: Optional[str] = None,
565
+ ) -> str:
566
+ """查询数据库实例的资源结构,便于编写 SQL 前探查。
567
+
568
+ resource_type=schema 列出库下 schema;table 列出表(可配 schema_name);
569
+ column 列出字段(需 schema_name + tb_name)。
570
+ 默认实例 alisg-haiyun-powerdata-pgsql-prod-01 / 库 creativault_business。
571
+ """
572
+ try:
573
+ cfg = load_profile()
574
+ data = await ArcheryClient(cfg).instance_resources(
575
+ instance_name=instance_name, db_name=db_name,
576
+ resource_type=resource_type, schema_name=schema_name, tb_name=tb_name,
577
+ )
578
+ return json.dumps({"resources": data}, ensure_ascii=False)
579
+ except ArcheryError as e:
580
+ return json.dumps({"error": str(e)}, ensure_ascii=False)
581
+
582
+ mcp.run()
583
+
584
+
585
+ def main() -> None:
586
+ parser = argparse.ArgumentParser(
587
+ prog="archery-sql-mcp",
588
+ description="Archery SQL 查询工具(无参数启动 = MCP server;子命令 = CLI)",
589
+ )
590
+ parser.add_argument("--version", action="version", version=f"%(prog)s {VERSION}")
591
+ sub = parser.add_subparsers(dest="command")
592
+
593
+ sub.add_parser("init", help="交互式初始化账号密码配置并验证登录")
594
+ sub.add_parser("check", help="全链路自检:登录 + SELECT 1")
595
+ sub.add_parser("instances", help="列出有读权限的数据库实例")
596
+
597
+ p_query = sub.add_parser("query", help="执行只读 SQL 查询")
598
+ p_query.add_argument("sql", help="SELECT / WITH 开头的 SQL 语句")
599
+ p_query.add_argument("--db", help="覆盖默认库名")
600
+ p_query.add_argument("--instance", help="覆盖默认实例名")
601
+ p_query.add_argument("--limit", type=int, help=f"返回行数上限(默认 100,最大 {MAX_LIMIT})")
602
+
603
+ args = parser.parse_args()
604
+ if args.command is None:
605
+ run_mcp() # 无参数 → MCP stdio 模式
606
+ return
607
+ handlers = {"init": cmd_init, "check": cmd_check, "query": cmd_query,
608
+ "instances": cmd_instances}
609
+ sys.exit(handlers[args.command](args))
610
+
611
+
612
+ if __name__ == "__main__":
613
+ main()