413 lines
15 KiB
Python
413 lines
15 KiB
Python
import os
|
||
import re
|
||
import contextlib
|
||
from collections import defaultdict
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
import sqlparse
|
||
from dotenv import load_dotenv
|
||
from psycopg_pool import ConnectionPool
|
||
|
||
from mcp.server.fastmcp import FastMCP
|
||
from mcp.server.transport_security import TransportSecuritySettings
|
||
|
||
from starlette.applications import Starlette
|
||
from starlette.middleware.base import BaseHTTPMiddleware
|
||
from starlette.requests import Request
|
||
from starlette.responses import JSONResponse
|
||
from starlette.routing import Mount
|
||
|
||
# 加载配置：.env.local > 同级 .env > 项目根 .env
|
||
_here = Path(__file__).resolve().parent
|
||
_root = _here.parent.parent # apps/mcp-server -> apps -> NeoZQYY
|
||
load_dotenv(_here / ".env.local", override=True)
|
||
load_dotenv(_here / ".env", override=False)
|
||
load_dotenv(_root / ".env", override=False)
|
||
|
||
|
||
# ----------------------------
|
||
# 工具：环境变量解析（避免 int("") 报错）
|
||
# ----------------------------
|
||
def env_str(name: str, default: str = "", required: bool = False) -> str:
|
||
v = os.getenv(name, default)
|
||
v = v if v is not None else default
|
||
v = v.strip() if isinstance(v, str) else v
|
||
if required and (v is None or v == ""):
|
||
raise RuntimeError(f"Missing required env var: {name}")
|
||
return v
|
||
|
||
|
||
def env_int(name: str, default: Optional[int] = None, required: bool = False) -> int:
|
||
raw = os.getenv(name, "")
|
||
if raw is None or raw.strip() == "":
|
||
if required and default is None:
|
||
raise RuntimeError(f"Missing required env var: {name}")
|
||
if default is None:
|
||
raise RuntimeError(f"Missing env var: {name}")
|
||
return default
|
||
try:
|
||
return int(raw.strip())
|
||
except ValueError as e:
|
||
raise RuntimeError(f"Invalid int env var {name}={raw!r}") from e
|
||
|
||
|
||
# ----------------------------
|
||
# 配置（用环境变量注入）
|
||
# MCP_PG_* 优先（独立部署），回退到项目公共 DB_* / ETL_DB_NAME / PG_NAME（分离式配置）
|
||
# ----------------------------
|
||
PGHOST = env_str("MCP_PG_HOST", default="") or env_str("DB_HOST", required=True)
|
||
PGPORT = env_int("MCP_PG_PORT", default=0) or env_int("DB_PORT", default=5432)
|
||
PGDATABASE = env_str("MCP_PG_DATABASE", default="") or env_str("ETL_DB_NAME", default="") or env_str("PG_NAME", required=True)
|
||
PGUSER = env_str("MCP_PG_USER", default="") or env_str("DB_USER", required=True)
|
||
PGPASSWORD = env_str("MCP_PG_PASSWORD", default="") or env_str("DB_PASSWORD", required=True)
|
||
|
||
MCP_TOKEN = env_str("MCP_TOKEN", default="") # 鉴权 token（可空：不启用鉴权）
|
||
MAX_ROWS = env_int("MCP_MAX_ROWS", default=500) # query_sql 默认最大行数
|
||
PORT = env_int("PORT", default=9000) # uvicorn 端口
|
||
|
||
# etl_feiqiu 库的六层 schema 架构
|
||
ALLOWED_SCHEMAS = ("ods", "dwd", "dws", "core", "meta", "app")
|
||
ALLOWED_SCHEMA_SET = set(ALLOWED_SCHEMAS)
|
||
|
||
# psycopg DSN（如果密码包含空格等特殊字符，建议改用 URL 形式并做编码）
|
||
DSN = (
|
||
f"host={PGHOST} port={PGPORT} dbname={PGDATABASE} "
|
||
f"user={PGUSER} password={PGPASSWORD}"
|
||
)
|
||
|
||
# 连接池：不要 open=True（避免解释器退出时 __del__ 清理触发异常）
|
||
pool = ConnectionPool(conninfo=DSN, min_size=1, max_size=10, timeout=60, open=False)
|
||
|
||
|
||
# ----------------------------
|
||
# SQL 只读门禁（最终底线仍是 DB 只读账号）
|
||
# ----------------------------
|
||
FORBIDDEN = re.compile(
|
||
r"\b(insert|update|delete|drop|alter|truncate|create|grant|revoke|copy|call|execute|do)\b",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
# 额外禁止显式跨 schema 访问（避免越权）
|
||
# 匹配 schema.table 模式，但排除单字母别名（如 t.id、o.amount）
|
||
SCHEMA_QUAL = re.compile(r"\b([a-zA-Z_][a-zA-Z0-9_]{1,})\s*\.\s*[a-zA-Z_]", re.IGNORECASE)
|
||
|
||
|
||
def _is_probably_readonly(sql: str) -> bool:
|
||
if FORBIDDEN.search(sql):
|
||
return False
|
||
parsed = sqlparse.parse(sql)
|
||
if not parsed:
|
||
return False
|
||
stmt = parsed[0]
|
||
for tok in stmt.tokens:
|
||
if tok.is_whitespace:
|
||
continue
|
||
first = str(tok).strip().lower()
|
||
return first in ("select", "with", "show", "explain")
|
||
return False
|
||
|
||
|
||
def _validate_schema(schema: str) -> Optional[Dict[str, Any]]:
|
||
if schema not in ALLOWED_SCHEMA_SET:
|
||
return {"error": f"schema 不允许：{schema}。仅允许：{sorted(ALLOWED_SCHEMA_SET)}"}
|
||
return None
|
||
|
||
|
||
def _reject_cross_schema(sql: str, allowed_schema: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
简单防护：如果出现显式 schema 前缀（xxx.），要求必须是白名单内的 schema 或系统 schema。
|
||
注：这不是 SQL parser 级别的严格策略，但能挡住绝大多数越权写法。
|
||
"""
|
||
matches = set(m.group(1) for m in SCHEMA_QUAL.finditer(sql or ""))
|
||
# 允许所有业务 schema + 系统 schema
|
||
safe = ALLOWED_SCHEMA_SET | {"pg_catalog", "information_schema"}
|
||
bad = sorted([s for s in matches if s.lower() not in {a.lower() for a in safe}])
|
||
if bad:
|
||
return {"error": f"SQL 被拒绝：检测到不允许的 schema 引用 {bad}，仅允许 {sorted(ALLOWED_SCHEMA_SET)} / 系统 schema。"}
|
||
return None
|
||
|
||
|
||
# ----------------------------
|
||
# FastMCP：Streamable HTTP + JSON 响应
|
||
# ----------------------------
|
||
mcp = FastMCP(
|
||
"postgres-mcp",
|
||
stateless_http=True,
|
||
json_response=True,
|
||
transport_security=TransportSecuritySettings(
|
||
enable_dns_rebinding_protection=True,
|
||
allowed_hosts=[
|
||
# 关键：既允许不带端口，也允许带端口
|
||
"mcp.langlangzhuoqiu.cn",
|
||
"mcp.langlangzhuoqiu.cn:*",
|
||
"localhost",
|
||
"localhost:*",
|
||
"127.0.0.1",
|
||
"127.0.0.1:*",
|
||
"100.64.0.4",
|
||
"100.64.0.4:*",
|
||
"100.64.0.1",
|
||
"100.64.0.1:*",
|
||
"106.52.16.235",
|
||
"106.52.16.235:*",
|
||
],
|
||
allowed_origins=[
|
||
"https://mcp.langlangzhuoqiu.cn",
|
||
"https://mcp.langlangzhuoqiu.cn:*",
|
||
"http://localhost",
|
||
"http://localhost:*",
|
||
"http://127.0.0.1",
|
||
"http://127.0.0.1:*",
|
||
],
|
||
),
|
||
)
|
||
|
||
|
||
# ----------------------------
|
||
# Tools：面向 etl_feiqiu 六层 schema
|
||
# ----------------------------
|
||
@mcp.tool()
|
||
def list_tables(schema: str = "dwd", include_views: bool = False) -> Dict[str, Any]:
|
||
"""列出指定 schema（ods/dwd/dws/core/meta/app）下的表（可选包含视图）"""
|
||
err = _validate_schema(schema)
|
||
if err:
|
||
return err
|
||
|
||
table_types = ("BASE TABLE", "VIEW") if include_views else ("BASE TABLE",)
|
||
sql = """
|
||
SELECT table_name, table_type
|
||
FROM information_schema.tables
|
||
WHERE table_schema = %s AND table_type = ANY(%s)
|
||
ORDER BY table_name;
|
||
"""
|
||
with pool.connection() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, (schema, list(table_types)))
|
||
rows = cur.fetchall()
|
||
|
||
return {
|
||
"schema": schema,
|
||
"include_views": include_views,
|
||
"tables": [{"name": r[0], "type": r[1]} for r in rows],
|
||
"table_count": len(rows),
|
||
}
|
||
|
||
|
||
@mcp.tool()
|
||
def describe_table(table: str, schema: str = "dwd") -> Dict[str, Any]:
|
||
"""查看表结构（字段、类型、是否可空、默认值）"""
|
||
err = _validate_schema(schema)
|
||
if err:
|
||
return err
|
||
|
||
sql = """
|
||
SELECT column_name, data_type, is_nullable, column_default, ordinal_position
|
||
FROM information_schema.columns
|
||
WHERE table_schema=%s AND table_name=%s
|
||
ORDER BY ordinal_position;
|
||
"""
|
||
with pool.connection() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, (schema, table))
|
||
rows = cur.fetchall()
|
||
|
||
return {
|
||
"schema": schema,
|
||
"table": table,
|
||
"columns": [
|
||
{"name": r[0], "type": r[1], "nullable": r[2], "default": r[3], "position": r[4]}
|
||
for r in rows
|
||
],
|
||
"column_count": len(rows),
|
||
}
|
||
|
||
|
||
@mcp.tool()
|
||
def describe_schemas(
|
||
schemas: Optional[List[str]] = None,
|
||
include_views: bool = False,
|
||
max_tables_per_schema: int = 500,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
返回 ods/dwd/dws/core/meta/app schema 下的表结构（含主键）。
|
||
不传 schemas 则返回全部六个 schema。
|
||
"""
|
||
schemas = schemas or list(ALLOWED_SCHEMAS)
|
||
|
||
invalid = [s for s in schemas if s not in ALLOWED_SCHEMA_SET]
|
||
if invalid:
|
||
return {"error": f"存在不允许的 schema：{invalid}。仅允许：{sorted(ALLOWED_SCHEMA_SET)}"}
|
||
|
||
table_types = ("BASE TABLE", "VIEW") if include_views else ("BASE TABLE",)
|
||
|
||
with pool.connection() as conn:
|
||
with conn.cursor() as cur:
|
||
# 1) 表清单
|
||
cur.execute(
|
||
"""
|
||
SELECT table_schema, table_name, table_type
|
||
FROM information_schema.tables
|
||
WHERE table_schema = ANY(%s)
|
||
AND table_type = ANY(%s)
|
||
ORDER BY table_schema, table_name;
|
||
""",
|
||
(schemas, list(table_types)),
|
||
)
|
||
table_rows = cur.fetchall()
|
||
|
||
tables_by_schema: Dict[str, List[Tuple[str, str]]] = defaultdict(list)
|
||
for s, t, tt in table_rows:
|
||
if len(tables_by_schema[s]) < max_tables_per_schema:
|
||
tables_by_schema[s].append((t, tt))
|
||
|
||
# 2) 所有列（一次性取；如表非常多，可考虑拆分/分页）
|
||
cur.execute(
|
||
"""
|
||
SELECT table_schema, table_name, column_name, data_type, is_nullable, column_default, ordinal_position
|
||
FROM information_schema.columns
|
||
WHERE table_schema = ANY(%s)
|
||
ORDER BY table_schema, table_name, ordinal_position;
|
||
""",
|
||
(schemas,),
|
||
)
|
||
col_rows = cur.fetchall()
|
||
|
||
cols_map: Dict[Tuple[str, str], List[Dict[str, Any]]] = defaultdict(list)
|
||
for s, t, c, dt, nul, dft, pos in col_rows:
|
||
cols_map[(s, t)].append(
|
||
{"name": c, "type": dt, "nullable": nul, "default": dft, "position": pos}
|
||
)
|
||
|
||
# 3) 主键
|
||
cur.execute(
|
||
"""
|
||
SELECT kcu.table_schema, kcu.table_name, kcu.column_name, kcu.ordinal_position
|
||
FROM information_schema.table_constraints tc
|
||
JOIN information_schema.key_column_usage kcu
|
||
ON tc.constraint_name = kcu.constraint_name
|
||
AND tc.table_schema = kcu.table_schema
|
||
AND tc.table_name = kcu.table_name
|
||
WHERE tc.constraint_type = 'PRIMARY KEY'
|
||
AND tc.table_schema = ANY(%s)
|
||
ORDER BY kcu.table_schema, kcu.table_name, kcu.ordinal_position;
|
||
""",
|
||
(schemas,),
|
||
)
|
||
pk_rows = cur.fetchall()
|
||
|
||
pk_map: Dict[Tuple[str, str], List[str]] = defaultdict(list)
|
||
for s, t, col, _pos in pk_rows:
|
||
pk_map[(s, t)].append(col)
|
||
|
||
# 4) 组装
|
||
result: Dict[str, Any] = {
|
||
"schemas": {},
|
||
"include_views": include_views,
|
||
"limits": {"max_tables_per_schema": max_tables_per_schema},
|
||
}
|
||
|
||
for s in schemas:
|
||
schema_tables = tables_by_schema.get(s, [])
|
||
result["schemas"][s] = {"table_count": len(schema_tables), "tables": {}}
|
||
for t, tt in schema_tables:
|
||
key = (s, t)
|
||
result["schemas"][s]["tables"][t] = {
|
||
"type": tt,
|
||
"primary_key": pk_map.get(key, []),
|
||
"columns": cols_map.get(key, []),
|
||
"column_count": len(cols_map.get(key, [])),
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
@mcp.tool()
|
||
def query_sql(schema: str, sql: str, max_rows: int = MAX_ROWS) -> Dict[str, Any]:
|
||
"""
|
||
在指定 schema 内执行只读 SQL（会 SET LOCAL search_path），并限制显式跨 schema 引用。
|
||
"""
|
||
err = _validate_schema(schema)
|
||
if err:
|
||
return err
|
||
|
||
sql = (sql or "").strip().rstrip(";")
|
||
if not _is_probably_readonly(sql):
|
||
return {"error": "SQL 被拒绝：仅允许只读（select/with/show/explain）并禁止危险关键字。"}
|
||
|
||
cross = _reject_cross_schema(sql, allowed_schema=schema)
|
||
if cross:
|
||
return cross
|
||
|
||
with pool.connection() as conn:
|
||
with conn.cursor() as cur:
|
||
# schema 已白名单校验，可安全拼接
|
||
cur.execute(f"SET LOCAL search_path TO {schema}")
|
||
cur.execute(sql)
|
||
|
||
cols = [d.name for d in (cur.description or [])]
|
||
rows = cur.fetchmany(max_rows + 1)
|
||
|
||
truncated = len(rows) > max_rows
|
||
rows = rows[:max_rows]
|
||
|
||
safe_rows: List[List[Any]] = []
|
||
for r in rows:
|
||
safe_rows.append([v if isinstance(v, (int, float, str, bool)) or v is None else str(v) for v in r])
|
||
|
||
return {
|
||
"schema": schema,
|
||
"columns": cols,
|
||
"rows": safe_rows,
|
||
"row_count": len(safe_rows),
|
||
"truncated": truncated,
|
||
"max_rows": max_rows,
|
||
}
|
||
|
||
|
||
# ----------------------------
|
||
# 鉴权 Middleware（支持 Bearer 或 query token）
|
||
# ----------------------------
|
||
class AuthMiddleware(BaseHTTPMiddleware):
|
||
async def dispatch(self, request: Request, call_next):
|
||
if MCP_TOKEN and request.url.path.startswith("/mcp"):
|
||
auth = request.headers.get("authorization", "")
|
||
token_q = request.query_params.get("token", "")
|
||
if auth != f"Bearer {MCP_TOKEN}" and token_q != MCP_TOKEN:
|
||
return JSONResponse({"error": "unauthorized"}, status_code=401)
|
||
return await call_next(request)
|
||
|
||
|
||
# ----------------------------
|
||
# lifespan：显式 open/close pool，并运行 session_manager
|
||
# ----------------------------
|
||
@contextlib.asynccontextmanager
|
||
async def lifespan(app: Starlette):
|
||
pool.open(wait=True, timeout=30)
|
||
try:
|
||
async with mcp.session_manager.run():
|
||
yield
|
||
finally:
|
||
# 避免解释器退出阶段 __del__ 清理导致异常
|
||
pool.close(timeout=5)
|
||
|
||
|
||
# MCP endpoint：/mcp（默认 streamable_http_path="/mcp"）
|
||
app = Starlette(
|
||
routes=[Mount("/", app=mcp.streamable_http_app())],
|
||
lifespan=lifespan,
|
||
)
|
||
app.add_middleware(AuthMiddleware)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(
|
||
app,
|
||
host="0.0.0.0",
|
||
port=PORT,
|
||
proxy_headers=True,
|
||
forwarded_allow_ips="*",
|
||
)
|