@hupan56/wlkj 2.7.12 → 3.1.2

package/templates/qoder/scripts/protocol/mcp/mysql_mcp_server.py

@@ -0,0 +1,461 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# v3.0 路径自举: 引导到 common/bootstrap, 统一 sys.path 逻辑
+import os as _o, sys as _s
+_f = _o.path.abspath(__file__)
+for _ in range(10):
+    _f = _o.path.dirname(_f)
+    _cp = _o.path.join(_f, 'foundation')
+    if _o.path.isfile(_o.path.join(_cp, 'bootstrap.py')):
+        break
+if _cp not in _s.path: _s.path.insert(0, _cp)
+from bootstrap import setup; setup()
+
+"""MySQL Read-Only MCP Server (pure Python stdio, no external SDK).
+
+提供 3 个语义化只读工具查询 MySQL 数据库 (真实数据层真相源):
+  - query_schema:  查表结构 (列名/类型/注释/约束)
+  - query_data:    查数据 (自动 LIMIT, 防 SELECT * 拉爆)
+  - query_distinct:查枚举值 (DISTINCT 某字段的去重值 + 计数)
+
+安全设计 (硬约束):
+  1. 启动前 ping 内网, 不通 → 输出"请连接内网"并拒绝启动
+  2. SQL 只读白名单: 只允许 SELECT/SHOW/DESCRIBE/EXPLAIN/WITH
+     任何写操作 (INSERT/UPDATE/DELETE/DROP/ALTER/...) 直接拒绝
+  3. query_data 自动注入 LIMIT, 上限 200 行
+  4. 账号密码从环境变量读 (mcp.json env 配置, 不落代码)
+
+注册: ~/.qoderwork/mcp.json 的 "qoder-mysql"。
+工具名: mcp__qoder-mysql__query_schema / query_data / query_distinct
+"""
+import io
+import os
+import re
+import sys
+import json
+import socket
+import asyncio
+
+# UTF-8
+try:
+    sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+    sys.stderr.reconfigure(encoding='utf-8', errors='replace')
+except Exception:
+    pass
+
+# ── 配置 (v3.x 多环境: 全部环境从 MYSQL_ENVIRONMENTS JSON 读, 运行时按 env 取) ──
+# launcher 把 config.yaml 的 mysql.environments 整体序列化成 JSON 注入这里。
+# server 启动后不再锁定单一连接, 而是按工具调用的 env 参数取对应环境建连。
+PING_TIMEOUT = 3  # 秒
+
+
+def _load_environments():
+    """解析全部环境配置。返回 dict[env_key -> {host,port,user,password,env_name}]。
+    数据源优先级: MYSQL_ENVIRONMENTS (JSON) > 旧扁平 MYSQL_HOST/... (包成单环境)。
+    解析结果缓存到模块级 _envs。"""
+    global _envs
+    if _envs is not None:
+        return _envs
+    _envs = {}
+    raw = os.environ.get('MYSQL_ENVIRONMENTS', '')
+    if raw:
+        try:
+            data = json.loads(raw)
+            if isinstance(data, dict):
+                for k, v in data.items():
+                    if isinstance(v, dict) and v.get('host'):
+                        _envs[k] = v
+        except (json.JSONDecodeError, TypeError):
+            pass
+    # 兼容旧扁平 env (launcher 没注入 JSON 时, 直接从 MYSQL_HOST 等读)
+    if not _envs:
+        h = os.environ.get('MYSQL_HOST', '')
+        if h:
+            _envs = {'test': {
+                'host': h,
+                'port': int(os.environ.get('MYSQL_PORT', '3306')),
+                'user': os.environ.get('MYSQL_USER', ''),
+                'password': os.environ.get('MYSQL_PASSWORD', ''),
+                'env_name': os.environ.get('MYSQL_ENV_NAME', h),
+            }}
+    return _envs
+
+
+def _default_env():
+    """默认环境 key。优先 MYSQL_DEFAULT_ENV, 否则取第一个环境。"""
+    d = os.environ.get('MYSQL_DEFAULT_ENV', '')
+    if d and d in _load_environments():
+        return d
+    envs = _load_environments()
+    return next(iter(envs), 'test') if envs else 'test'
+
+
+def _env_config(env_key):
+    """取某环境的配置 dict。不存在返回 None。"""
+    return _load_environments().get(env_key)
+
+
+def _env_notice(env_key, kind='data'):
+    """生成当前环境可信度声明 (随查询结果返回, AI 不可能"忘记"边界)。
+    kind='schema' = 表结构 (可信); kind='data' = 数据行 (测试/生产语义不同)。"""
+    cfg = _env_config(env_key) or {}
+    name = cfg.get('env_name', env_key)
+    is_prod = 'prod' in str(env_key).lower() or '生产' in name
+    if kind == 'schema':
+        if is_prod:
+            return '\n🔴 [环境: %s] 表结构/字段/枚举定义 (生产环境, 仅核对用)。' % name
+        return '\n[环境: %s] 表结构/字段/枚举定义与生产一致，可作为字段规格的真源。' % name
+    # data
+    if is_prod:
+        return ('\n🔴⚠️ [环境: %s] 以下为生产真实数据! 仅核对字段含义/格式, '
+                '禁止作为业务结论外传。' % name)
+    return ('\n⚠️ [环境: %s] 以下为测试数据样本，不代表线上真实分布。'
+            '字段取值范围可信，但具体数值/计数仅测试环境参考。' % name)
+
+
+_envs = None  # 缓存: 解析一次
+
+
+# stdout 重定向 + 协议层由 BaseMCPServer 接管 (v3.0 重构)
+
+
+# ── 安全: 内网连通性校验 ──────────────────────────────────────────────
+def check_intranet(env_key=None):
+    """TCP 连通性校验 (比 ICMP ping 更可靠, 不依赖管理员权限)。
+    连目标环境 3306 端口, 通=内网可达, 不通=提示连内网。
+    env_key=None 时用默认环境。返回 (ok: bool, msg: str)。"""
+    env_key = env_key or _default_env()
+    cfg = _env_config(env_key)
+    if not cfg:
+        return False, '未知环境 "%s"，可选: %s' % (env_key, list(_load_environments()))
+    host, port = cfg.get('host', ''), int(cfg.get('port', 3306))
+    try:
+        with socket.create_connection((host, port), timeout=PING_TIMEOUT):
+            return True, ''
+    except (socket.timeout, OSError) as e:
+        return False, (
+            '无法连接 MySQL [%s] (%s:%s): %s\n'
+            '⚠️ 请确认你已连接公司内网 (VPN/办公网) 且该环境地址可达, 然后重试。'
+        ) % (env_key, host, port, str(e)[:80])
+
+
+# ── 安全: SQL 只读白名单 ──────────────────────────────────────────────
+# 允许的 SQL 首关键字 (只读操作)
+READONLY_PREFIXES = {
+    'select', 'show', 'describe', 'desc', 'explain', 'with',
+}
+# 明确禁止的关键字 (即使出现在子查询也拦截, 防注入写操作)
+FORBIDDEN_KEYWORDS = {
+    'insert', 'update', 'delete', 'drop', 'alter', 'create', 'truncate',
+    'replace', 'rename', 'grant', 'revoke', 'lock', 'unlock', 'call',
+    'handler', 'load', 'outfile', 'dumpfile',
+}
+
+
+def validate_readonly(sql):
+    """校验 SQL 是否只读。返回 (ok: bool, reason: str)。"""
+    if not sql or not sql.strip():
+        return False, '空 SQL'
+    s = sql.strip().lower().lstrip('(').strip()
+    # 去掉开头的 WITH ... AS ( ... ) 后看主语句
+    # 简化: 取第一个词
+    first_word = re.split(r'[\s(]', s, 1)[0]
+    if first_word not in READONLY_PREFIXES:
+        return False, '非只读操作: 首关键字 "%s" 不在白名单 %s' % (
+            first_word, sorted(READONLY_PREFIXES))
+    # 扫描是否含危险关键字 (作为独立词, 防 delete 嵌在字段名里)
+    tokens = set(re.findall(r'\b\w+\b', s))
+    hit = tokens & FORBIDDEN_KEYWORDS
+    if hit:
+        return False, 'SQL 含禁止的关键字: %s (只读模式拒绝写操作)' % sorted(hit)
+    return True, ''
+
+
+def inject_limit(sql, max_rows=200):
+    """给 SELECT 自动加 LIMIT (若用户没写)。防拉爆。"""
+    s = sql.strip().rstrip(';')
+    # 已有 limit → 不动 (但截断到 max_rows 由 fetch 控制)
+    if re.search(r'\blimit\s+\d+', s, re.IGNORECASE):
+        return s
+    return s + ' LIMIT %d' % max_rows
+
+
+# ── MySQL 连接 (惰性, 多环境连接池) ────────────────────────────────────
+# 每个环境一个连接, 按 env_key 缓存。运行时切库 = 取不同 env_key 的连接。
+_conns = {}  # env_key -> connection
+
+
+def _get_conn(env_key):
+    """惰性建连 (首次查某环境时才连)。连接按 env_key 缓存复用。"""
+    if env_key in _conns:
+        return _conns[env_key]
+    cfg = _env_config(env_key)
+    if not cfg:
+        raise ValueError('未知环境 "%s"，可选: %s' % (env_key, list(_load_environments())))
+    # 用最普及的 pymysql, 纯 Python, 无编译依赖
+    import pymysql
+    _conns[env_key] = pymysql.connect(
+        host=cfg.get('host', ''), port=int(cfg.get('port', 3306)),
+        user=cfg.get('user', ''), password=cfg.get('password', ''),
+        charset='utf8mb4',
+        cursorclass=pymysql.cursors.DictCursor,
+        connect_timeout=5, read_timeout=10,
+    )
+    return _conns[env_key]
+
+
+def _execute_query(sql, db=None, max_rows=200, env_key=None):
+    """执行只读查询, 返回 (columns, rows)。失败抛异常。
+    env_key 指定走哪个环境 (默认 default_env)。"""
+    ok, reason = validate_readonly(sql)
+    if not ok:
+        raise ValueError('SQL 被拒: ' + reason)
+    env_key = env_key or _default_env()
+    conn = _get_conn(env_key)
+    if db:
+        conn.select_db(db)
+    cur = conn.cursor()
+    try:
+        cur.execute(sql)
+        rows = cur.fetchmany(max_rows)  # 硬上限, 不全拉
+        cols = [d[0] for d in cur.description] if cur.description else []
+        return cols, rows
+    finally:
+        cur.close()
+
+
+# ── MCP 工具定义 ──────────────────────────────────────────────────────
+# 所有查询工具都带可选 env 参数, 取 config.yaml mysql.environments 的 key
+# (如 test/pre/prod)。不填用 default_env。调用方 (AI) 应先问用户查哪个环境。
+_ENV_PROP = {"type": "string", "description":
+    "环境标签 (config.yaml mysql.environments 的 key, 如 test/pre/prod)。"
+    "不填用默认环境。先问用户查哪个环境再传。"}
+
+TOOLS = [
+    {
+        "name": "list_envs",
+        "description": "列出所有可用的数据库环境 (config.yaml mysql.environments)。AI 查库前先调用此工具, 把可选环境列给用户选。例: list_envs()",
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+    {
+        "name": "query_schema",
+        "description": "查表结构: 列名/类型/注释/是否可空/主键。用于写 PRD/spec 时核对真实字段。例: query_schema(table='t_quality_case', env='test')",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "table": {"type": "string", "description": "表名 (如 t_quality_case)"},
+                "db": {"type": "string", "description": "库名 (可选, 不填用默认库)"},
+                "env": _ENV_PROP,
+            },
+            "required": ["table"],
+        },
+    },
+    {
+        "name": "query_data",
+        "description": "查数据: 按条件查表数据, 自动 LIMIT 200 行。数据行语义随环境不同 (测试=样本, 生产=真实数据)。例: query_data(table='t_quality_case', columns='id,case_status', env='test', limit=10)",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "table": {"type": "string", "description": "表名"},
+                "columns": {"type": "string", "description": "列名 (逗号分隔, 默认 *)"},
+                "where": {"type": "string", "description": "WHERE 条件 (不带 WHERE 关键字)"},
+                "db": {"type": "string", "description": "库名 (可选)"},
+                "limit": {"type": "integer", "description": "行数上限 (默认200, 最大200)"},
+                "env": _ENV_PROP,
+            },
+            "required": ["table"],
+        },
+    },
+    {
+        "name": "query_distinct",
+        "description": "查枚举值: 某字段的去重值。取值范围可信(各环境一致), 计数语义随环境不同。例: query_distinct(table='t_quality_case', column='case_status', env='test')",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "table": {"type": "string", "description": "表名"},
+                "column": {"type": "string", "description": "要查枚举的列名"},
+                "db": {"type": "string", "description": "库名 (可选)"},
+                "limit": {"type": "integer", "description": "返回枚举值上限 (默认50)"},
+                "env": _ENV_PROP,
+            },
+            "required": ["table", "column"],
+        },
+    },
+    {
+        "name": "list_databases",
+        "description": "列出某环境所有可见的数据库。用于发现有哪些库可查。例: list_databases(env='test')",
+        "inputSchema": {"type": "object", "properties": {"env": _ENV_PROP}},
+    },
+    {
+        "name": "list_tables",
+        "description": "列出某环境某库的表 (支持模糊匹配)。用于发现有哪些表。例: list_tables(db='fywl_ics', pattern='quality%', env='test')",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "db": {"type": "string", "description": "库名"},
+                "pattern": {"type": "string", "description": "表名 LIKE 模式 (可选, 如 quality%)"},
+                "env": _ENV_PROP,
+            },
+            "required": ["db"],
+        },
+    },
+]
+
+
+# ── 工具执行 ──────────────────────────────────────────────────────────
+def _execute_tool(name, args):
+    args = args or {}
+
+    # list_envs: 不需要连库, 直接返回可选环境 (供 AI 列给用户选)
+    if name == 'list_envs':
+        envs = _load_environments()
+        if not envs:
+            return ('未配置任何数据库环境。请在 .qoder/config.yaml 的 mysql.environments '
+                    '里配置 (参考 mysql.default_env + environments 结构)。')
+        default = _default_env()
+        lines = ['可选数据库环境 (%d 个, 标 * 为默认):' % len(envs)]
+        for k, v in envs.items():
+            mark = ' *' if k == default else ''
+            lines.append('  %s%s - %s (%s)' % (
+                k, mark, v.get('env_name', k), v.get('host', '?')))
+        return '\n'.join(lines)
+
+    # 其它工具: 解析 env 参数, 校验环境存在
+    env_key = args.get('env') or _default_env()
+    if not _env_config(env_key):
+        return '❌ 未知环境 "%s"，可选: %s\n(在 config.yaml mysql.environments 里配置)' % (
+            env_key, list(_load_environments()))
+
+    # 内网/连通性校验 (按目标环境, 每次调用都校验, 防 VPN 中途断开)
+    ok, msg = check_intranet(env_key)
+    if not ok:
+        return msg
+
+    if name == 'list_databases':
+        cols, rows = _execute_query('SHOW DATABASES', env_key=env_key)
+        dbs = [list(r.values())[0] for r in rows]
+        return '数据库 (%d 个, 环境=%s):\n  %s%s' % (
+            len(dbs), env_key, '\n  '.join(dbs), _env_notice(env_key, 'schema'))
+
+    if name == 'list_tables':
+        db = args.get('db', '')
+        pattern = args.get('pattern', '')
+        sql = 'SHOW TABLES'
+        if pattern:
+            sql += ' LIKE "%s"' % pattern.replace('"', '')
+        cols, rows = _execute_query(sql, db=db, env_key=env_key)
+        tables = [list(r.values())[0] for r in rows]
+        return '表 (%d 个, 库=%s, 环境=%s):\n  %s%s' % (
+            len(tables), db, env_key, '\n  '.join(tables[:100]),
+            _env_notice(env_key, 'schema'))
+
+    if name == 'query_schema':
+        table = args.get('table', '').replace('"', '').replace("'", '')
+        db = args.get('db')
+        sql = (
+            'SELECT column_name, data_type, column_type, is_nullable, '
+            'column_key, column_default, column_comment '
+            'FROM information_schema.columns '
+            'WHERE table_name = "%s" ' % table
+        )
+        if db:
+            sql += 'AND table_schema = "%s" ' % db.replace('"', '')
+        sql += 'ORDER BY ordinal_position'
+        cols, rows = _execute_query(sql, env_key=env_key)
+        if not rows:
+            return '未找到表: %s (检查表名或库名, 环境=%s)' % (table, env_key)
+        lines = ['%-20s %-12s %-6s %-6s %s' % (
+            '列名', '类型', '可空', '键', '注释')]
+        for r in rows:
+            lines.append('%-20s %-12s %-6s %-6s %s' % (
+                r.get('column_name', ''), r.get('column_type', '')[:12],
+                r.get('is_nullable', '')[:3], r.get('column_key', ''),
+                (r.get('column_comment', '') or '')[:30]))
+        return '\n'.join(lines) + _env_notice(env_key, 'schema')
+
+    if name == 'query_data':
+        table = args.get('table', '').replace('"', '').replace("'", '')
+        columns = args.get('columns') or '*'
+        where = args.get('where', '').strip()
+        db = args.get('db')
+        limit = min(int(args.get('limit', 200)), 200)
+        sql = 'SELECT %s FROM `%s`' % (columns, table)
+        if where:
+            sql += ' WHERE %s' % where
+        sql = inject_limit(sql, limit)
+        cols, rows = _execute_query(sql, db=db, max_rows=limit, env_key=env_key)
+        if not rows:
+            return '无数据 (表=%s, 环境=%s)' % (table, env_key)
+        lines = [' | '.join(cols)]
+        lines.append('-' * 40)
+        for r in rows[:limit]:
+            lines.append(' | '.join(str(r.get(c, ''))[:30] for c in cols))
+        return '\n'.join(lines) + _env_notice(env_key, 'data')
+
+    if name == 'query_distinct':
+        table = args.get('table', '').replace('"', '').replace("'", '')
+        column = args.get('column', '').replace('"', '').replace("'", '')
+        db = args.get('db')
+        limit = min(int(args.get('limit', 50)), 50)
+        sql = (
+            'SELECT `%s` AS val, COUNT(*) AS cnt FROM `%s` '
+            'GROUP BY `%s` ORDER BY cnt DESC LIMIT %d'
+        ) % (column, table, column, limit)
+        cols, rows = _execute_query(sql, db=db, max_rows=limit, env_key=env_key)
+        if not rows:
+            return '无数据 (表=%s 列=%s, 环境=%s)' % (table, column, env_key)
+        # 取值范围可信(各环境一致), 计数语义随环境 (测试=样本/生产=真实)
+        lines = ['枚举值 (表=%s 列=%s, 取值范围可信):' % (table, column)]
+        for r in rows:
+            lines.append('  %s: %s 次' % (r.get('val', ''), r.get('cnt', '')))
+        return '\n'.join(lines) + _env_notice(env_key, 'data')
+
+    return '未知工具: %s' % name
+
+
+# ── MCP Server (继承 BaseMCPServer, 协议层由基类接管) ──────────────────
+# v3.0 重构: 删除 _send/_handle/main 三段逐字复制代码 (~90行),
+# 改为继承 BaseMCPServer, 覆写 on_tool_error 处理 ValueError(只读拒绝)。
+
+def _bootstrap_base():
+    """延迟 import BaseMCPServer (scripts/ 在 sys.path)。"""
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    from foundation.protocol.mcp_base import BaseMCPServer
+    return BaseMCPServer
+
+
+class MysqlMCPServer(_bootstrap_base()):
+    """MySQL 只读 MCP Server。继承基类协议层, 只填工具逻辑 + 异常分类。"""
+
+    SERVER_NAME = "qoder-mysql"
+    SERVER_VERSION = "1.0.0"
+    TOOLS = TOOLS  # 引用上方模块级 TOOLS
+
+    def handle_tool(self, name, args):
+        return _execute_tool(name, args)
+
+    def on_initialize_hook(self, info):
+        # 启动时校验默认环境连通性, 不通塞 warning。
+        # 附带把可用环境列出来, 方便调用方知道有哪些环境可查。
+        envs = _load_environments()
+        if not envs:
+            info['warning'] = ('未配置任何 mysql 环境。请在 config.yaml 的 '
+                               'mysql.environments 里配置。')
+            return
+        default = _default_env()
+        ok, msg = check_intranet(default)
+        if not ok:
+            info['warning'] = msg
+        info['mysqlEnvs'] = list(envs.keys())
+        info['mysqlDefaultEnv'] = default
+
+    def on_tool_error(self, e):
+        # mysql 特有: ValueError = SQL 只读校验拒绝, 单独友好提示
+        if isinstance(e, ValueError):
+            return ("❌ " + str(e), True)
+        # 其它异常走基类默认
+        return super().on_tool_error(e)
+
+
+if __name__ == '__main__':
+    MysqlMCPServer().run()