@hupan56/wlkj 3.0.0 → 3.1.2

package/templates/qoder/scripts/{mcp → protocol/mcp}/mysql_mcp_server.py

@@ -2,7 +2,12 @@
 # -*- coding: utf-8 -*-
 # v3.0 路径自举: 引导到 common/bootstrap, 统一 sys.path 逻辑
 import os as _o, sys as _s
-_cp = _o.path.join(_o.path.dirname(_o.path.dirname(_o.path.abspath(__file__))), 'common')
+_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()
 
@@ -38,44 +43,100 @@ try:
 except Exception:
     pass
 
-# ── 配置 (从环境变量, mcp.json env 注入) ──────────────────────────────
-MYSQL_HOST = os.environ.get('MYSQL_HOST', '10.54.6.9')
-MYSQL_PORT = int(os.environ.get('MYSQL_PORT', '3306'))
-MYSQL_USER = os.environ.get('MYSQL_USER', 'fywl_ics_test')
-MYSQL_PASSWORD = os.environ.get('MYSQL_PASSWORD', 'Fywl_ics_test56476')
-PING_HOST = MYSQL_HOST  # 连通性校验目标
+# ── 配置 (v3.x 多环境: 全部环境从 MYSQL_ENVIRONMENTS JSON 读, 运行时按 env 取) ──
+# launcher 把 config.yaml 的 mysql.environments 整体序列化成 JSON 注入这里。
+# server 启动后不再锁定单一连接, 而是按工具调用的 env 参数取对应环境建连。
 PING_TIMEOUT = 3  # 秒
 
-# ── 环境可信度声明 (随每次工具返回, AI 不可能"忘记"边界) ──────────────
-# QAS 测试环境: 表结构/字段/枚举取值范围 = 生产一致 (可信);
-#               数据行/计数分布 = 测试数据, 不代表线上 (不可信)。
-ENV_NAME = os.environ.get('MYSQL_ENV_NAME', 'QAS 测试环境')
-ENV_NOTICE_SCHEMA = (
-    '\n[环境: %s] 表结构/字段/枚举定义与生产一致，可作为字段规格的真源。' % ENV_NAME
-)
-ENV_NOTICE_DATA = (
-    '\n⚠️ [环境: %s] 以下为测试数据样本，不代表线上真实分布。'
-    '字段取值范围可信，但具体数值/计数仅测试环境参考。' % ENV_NAME
-)
+
+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():
+def check_intranet(env_key=None):
     """TCP 连通性校验 (比 ICMP ping 更可靠, 不依赖管理员权限)。
-    连 3306 端口, 通=内网可达, 不通=提示连内网。
-    返回 (ok: bool, msg: str)。
-    """
+    连目标环境 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((MYSQL_HOST, MYSQL_PORT), timeout=PING_TIMEOUT):
+        with socket.create_connection((host, port), timeout=PING_TIMEOUT):
             return True, ''
     except (socket.timeout, OSError) as e:
         return False, (
-            '无法连接内网 MySQL (%s:%s): %s\n'
-            '⚠️ 请确认你已连接公司内网 (VPN/办公网), 然后重试。'
-            'MySQL MCP 仅在内网可用。'
-        ) % (MYSQL_HOST, MYSQL_PORT, str(e)[:80])
+            '无法连接 MySQL [%s] (%s:%s): %s\n'
+            '⚠️ 请确认你已连接公司内网 (VPN/办公网) 且该环境地址可达, 然后重试。'
+        ) % (env_key, host, port, str(e)[:80])
 
 
 # ── 安全: SQL 只读白名单 ──────────────────────────────────────────────
@@ -119,33 +180,38 @@ def inject_limit(sql, max_rows=200):
     return s + ' LIMIT %d' % max_rows
 
 
-# ── MySQL 连接 (惰性) ─────────────────────────────────────────────────
-_conn = None
+# ── MySQL 连接 (惰性, 多环境连接池) ────────────────────────────────────
+# 每个环境一个连接, 按 env_key 缓存。运行时切库 = 取不同 env_key 的连接。
+_conns = {}  # env_key -> connection
 
 
-def _get_conn():
-    """惰性建连 (首次查询时才连)。"""
-    global _conn
-    if _conn is not None:
-        return _conn
+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
-    _conn = pymysql.connect(
-        host=MYSQL_HOST, port=MYSQL_PORT,
-        user=MYSQL_USER, password=MYSQL_PASSWORD,
+    _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 _conn
+    return _conns[env_key]
 
 
-def _execute_query(sql, db=None, max_rows=200):
-    """执行只读查询, 返回 (columns, rows)。失败抛异常。"""
+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)
-    conn = _get_conn()
+    env_key = env_key or _default_env()
+    conn = _get_conn(env_key)
     if db:
         conn.select_db(db)
     cur = conn.cursor()
@@ -159,22 +225,34 @@ def _execute_query(sql, db=None, max_rows=200):
 
 
 # ── 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') 或 query_schema(table='t_quality_case', db='fywl_ics')",
+        "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 行。[QAS测试环境] 数据行是测试数据, 不代表线上真实分布, 仅参考字段含义/格式。例: query_data(table='t_quality_case', columns='id,case_status', where='create_time > \"2026-06-01\"', limit=10)",
+        "description": "查数据: 按条件查表数据, 自动 LIMIT 200 行。数据行语义随环境不同 (测试=样本, 生产=真实数据)。例: query_data(table='t_quality_case', columns='id,case_status', env='test', limit=10)",
         "inputSchema": {
             "type": "object",
             "properties": {
@@ -183,13 +261,14 @@ TOOLS = [
                 "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": "查枚举值: 某字段的去重值。[QAS测试环境] 取值范围可信(与生产一致), 计数不可信(测试分布)。用于确认状态/类型字段有哪些取值。例: query_distinct(table='t_quality_case', column='case_status')",
+        "description": "查枚举值: 某字段的去重值。取值范围可信(各环境一致), 计数语义随环境不同。例: query_distinct(table='t_quality_case', column='case_status', env='test')",
         "inputSchema": {
             "type": "object",
             "properties": {
@@ -197,23 +276,25 @@ TOOLS = [
                 "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()",
-        "inputSchema": {"type": "object", "properties": {}},
+        "description": "列出某环境所有可见的数据库。用于发现有哪些库可查。例: list_databases(env='test')",
+        "inputSchema": {"type": "object", "properties": {"env": _ENV_PROP}},
     },
     {
         "name": "list_tables",
-        "description": "列出某库的表 (支持模糊匹配)。用于发现有哪些表。例: list_tables(db='fywl_ics', pattern='quality%')",
+        "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"],
         },
@@ -224,15 +305,37 @@ TOOLS = [
 # ── 工具执行 ──────────────────────────────────────────────────────────
 def _execute_tool(name, args):
     args = args or {}
-    # 内网校验 (每次工具调用都校验, 防 VPN 中途断开)
-    ok, msg = check_intranet()
+
+    # 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')
+        cols, rows = _execute_query('SHOW DATABASES', env_key=env_key)
         dbs = [list(r.values())[0] for r in rows]
-        return '数据库 (%d 个):\n  %s' % (len(dbs), '\n  '.join(dbs))
+        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', '')
@@ -240,9 +343,11 @@ def _execute_tool(name, args):
         sql = 'SHOW TABLES'
         if pattern:
             sql += ' LIKE "%s"' % pattern.replace('"', '')
-        cols, rows = _execute_query(sql, db=db)
+        cols, rows = _execute_query(sql, db=db, env_key=env_key)
         tables = [list(r.values())[0] for r in rows]
-        return '表 (%d 个, 库=%s):\n  %s' % (len(tables), db, '\n  '.join(tables[:100]))
+        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("'", '')
@@ -256,9 +361,9 @@ def _execute_tool(name, args):
         if db:
             sql += 'AND table_schema = "%s" ' % db.replace('"', '')
         sql += 'ORDER BY ordinal_position'
-        cols, rows = _execute_query(sql)
+        cols, rows = _execute_query(sql, env_key=env_key)
         if not rows:
-            return '未找到表: %s (检查表名或库名)' % table
+            return '未找到表: %s (检查表名或库名, 环境=%s)' % (table, env_key)
         lines = ['%-20s %-12s %-6s %-6s %s' % (
             '列名', '类型', '可空', '键', '注释')]
         for r in rows:
@@ -266,7 +371,7 @@ def _execute_tool(name, args):
                 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_SCHEMA
+        return '\n'.join(lines) + _env_notice(env_key, 'schema')
 
     if name == 'query_data':
         table = args.get('table', '').replace('"', '').replace("'", '')
@@ -278,14 +383,14 @@ def _execute_tool(name, args):
         if where:
             sql += ' WHERE %s' % where
         sql = inject_limit(sql, limit)
-        cols, rows = _execute_query(sql, db=db, max_rows=limit)
+        cols, rows = _execute_query(sql, db=db, max_rows=limit, env_key=env_key)
         if not rows:
-            return '无数据 (表=%s)' % table
+            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_DATA
+        return '\n'.join(lines) + _env_notice(env_key, 'data')
 
     if name == 'query_distinct':
         table = args.get('table', '').replace('"', '').replace("'", '')
@@ -296,14 +401,14 @@ def _execute_tool(name, args):
             '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)
+        cols, rows = _execute_query(sql, db=db, max_rows=limit, env_key=env_key)
         if not rows:
-            return '无数据 (表=%s 列=%s)' % (table, column)
-        # 取值范围可信(和生产一致), 但计数是测试环境分布(不可信)
+            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_DATA
+            lines.append('  %s: %s 次' % (r.get('val', ''), r.get('cnt', '')))
+        return '\n'.join(lines) + _env_notice(env_key, 'data')
 
     return '未知工具: %s' % name
 
@@ -315,7 +420,7 @@ def _execute_tool(name, args):
 def _bootstrap_base():
     """延迟 import BaseMCPServer (scripts/ 在 sys.path)。"""
     sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    from common.mcp_base import BaseMCPServer
+    from foundation.protocol.mcp_base import BaseMCPServer
     return BaseMCPServer
 
 
@@ -330,10 +435,19 @@ class MysqlMCPServer(_bootstrap_base()):
         return _execute_tool(name, args)
 
     def on_initialize_hook(self, info):
-        # 启动时内网校验, 不通塞 warning (与旧版行为一致)
-        ok, msg = check_intranet()
+        # 启动时校验默认环境连通性, 不通塞 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 只读校验拒绝, 单独友好提示