Refine agent skill boundaries and secret handling

skills/database-operation/scripts/mp-db.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""
+MoviePilot 数据库操作脚本。
+
+脚本从项目配置读取数据库连接参数，不要求 Agent 在提示词中接触数据库密码。
+默认只允许查询语句；写操作必须显式传入 --write。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+from sqlalchemy import create_engine, inspect, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import SQLAlchemyError
+
+
+SCRIPT_PATH = Path(__file__).resolve()
+PROJECT_ROOT = SCRIPT_PATH.parents[3]
+WRITE_STATEMENT_RE = re.compile(
+    r"^\s*(insert|update|delete|drop|alter|truncate|create|replace)\b",
+    re.IGNORECASE,
+)
+WRITE_KEYWORD_RE = re.compile(
+    r"\b(insert|update|delete|drop|alter|truncate|create|replace)\b",
+    re.IGNORECASE,
+)
+SELECT_STATEMENT_RE = re.compile(r"^\s*(select|with|explain)\b", re.IGNORECASE)
+
+
+def _ensure_project_import() -> None:
+    """确保脚本可以从任意工作目录导入 MoviePilot 项目模块。"""
+    project_path = str(PROJECT_ROOT)
+    if project_path not in sys.path:
+        sys.path.insert(0, project_path)
+
+
+def _load_settings() -> Any:
+    """读取 MoviePilot 运行配置。"""
+    _ensure_project_import()
+    from app.core.config import settings  # pylint: disable=import-outside-toplevel
+
+    return settings
+
+
+def _build_engine() -> Engine:
+    """根据 MoviePilot 配置创建同步数据库引擎。"""
+    settings = _load_settings()
+    if str(settings.DB_TYPE).lower() == "postgresql":
+        return create_engine(
+            settings.DB_POSTGRESQL_URL(),
+            pool_pre_ping=settings.DB_POOL_PRE_PING,
+            pool_recycle=settings.DB_POOL_RECYCLE,
+        )
+    return create_engine(
+        f"sqlite:///{settings.CONFIG_PATH}/user.db",
+        pool_pre_ping=settings.DB_POOL_PRE_PING,
+        pool_recycle=settings.DB_POOL_RECYCLE,
+        connect_args={"timeout": settings.DB_TIMEOUT},
+    )
+
+
+def _normalize_sql(sql: str) -> str:
+    """去除 SQL 首尾空白和末尾分号。"""
+    return sql.strip().rstrip(";")
+
+
+def _contains_multiple_statements(sql: str) -> bool:
+    """判断 SQL 是否包含多语句分隔符。"""
+    return ";" in sql
+
+
+def _is_write_statement(sql: str) -> bool:
+    """判断 SQL 是否为写操作或结构变更操作。"""
+    return bool(WRITE_STATEMENT_RE.match(sql))
+
+
+def _is_supported_statement(sql: str, allow_write: bool) -> bool:
+    """判断 SQL 是否在当前权限模式下允许执行。"""
+    if _contains_multiple_statements(sql):
+        return False
+    if allow_write:
+        return True
+    return bool(SELECT_STATEMENT_RE.match(sql)) and not WRITE_KEYWORD_RE.search(sql)
+
+
+def _append_limit(sql: str, limit: int) -> str:
+    """为普通 SELECT 查询追加默认 LIMIT，避免输出过大。"""
+    if limit <= 0:
+        return sql
+    lowered = sql.lower()
+    if not lowered.lstrip().startswith("select"):
+        return sql
+    if re.search(r"\blimit\s+\d+\b", lowered):
+        return sql
+    return f"{sql} LIMIT {limit}"
+
+
+def _row_to_dict(row: Any) -> dict[str, Any]:
+    """将 SQLAlchemy 行对象转为普通字典。"""
+    return dict(row._mapping)
+
+
+def _print_json(payload: Any) -> None:
+    """输出 JSON 结果。"""
+    print(json.dumps(payload, ensure_ascii=False, indent=2, default=str))
+
+
+def list_tables() -> int:
+    """列出当前数据库中的数据表。"""
+    engine = _build_engine()
+    inspector = inspect(engine)
+    _print_json({"tables": sorted(inspector.get_table_names())})
+    return 0
+
+
+def show_schema(table_name: str) -> int:
+    """显示指定数据表的字段结构。"""
+    engine = _build_engine()
+    inspector = inspect(engine)
+    columns = [
+        {
+            "name": column.get("name"),
+            "type": str(column.get("type")),
+            "nullable": column.get("nullable"),
+            "default": column.get("default"),
+            "primary_key": bool(column.get("primary_key")),
+        }
+        for column in inspector.get_columns(table_name)
+    ]
+    _print_json({"table": table_name, "columns": columns})
+    return 0
+
+
+def run_query(sql: str, *, limit: int = 100, allow_write: bool = False) -> int:
+    """
+    执行 SQL 语句并输出 JSON 结果。
+
+    :param sql: 要执行的 SQL 语句
+    :param limit: 查询语句默认追加的最大行数
+    :param allow_write: 是否允许写操作或结构变更操作
+    :return: 进程退出码
+    """
+    normalized_sql = _normalize_sql(sql)
+    if not normalized_sql:
+        print("Error: SQL is empty", file=sys.stderr)
+        return 1
+    if not _is_supported_statement(normalized_sql, allow_write):
+        print("Error: write statements require --write", file=sys.stderr)
+        return 1
+
+    statement_sql = _append_limit(normalized_sql, limit)
+    engine = _build_engine()
+    try:
+        with engine.begin() as connection:
+            result = connection.execute(text(statement_sql))
+            if result.returns_rows:
+                rows = [_row_to_dict(row) for row in result.fetchall()]
+                _print_json({"rows": rows, "row_count": len(rows)})
+            else:
+                _print_json({"row_count": result.rowcount})
+    except SQLAlchemyError as err:
+        print(f"Error: {err}", file=sys.stderr)
+        return 1
+    return 0
+
+
+def _read_sql(sql: Optional[str], sql_file: Optional[str]) -> str:
+    """从参数或文件读取 SQL 文本。"""
+    if sql:
+        return sql
+    if sql_file:
+        return Path(sql_file).read_text(encoding="utf-8")
+    if not sys.stdin.isatty():
+        return sys.stdin.read()
+    return ""
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    """构建命令行参数解析器。"""
+    parser = argparse.ArgumentParser(description="MoviePilot database operation helper")
+    subparsers = parser.add_subparsers(dest="command")
+
+    query_parser = subparsers.add_parser("query", help="execute a SQL statement")
+    query_parser.add_argument("sql", nargs="?", help="SQL statement")
+    query_parser.add_argument("--file", dest="sql_file", help="read SQL from file")
+    query_parser.add_argument("--limit", type=int, default=100, help="default SELECT row limit")
+    query_parser.add_argument("--write", action="store_true", help="allow write statements")
+
+    write_parser = subparsers.add_parser("write", help="execute a write statement")
+    write_parser.add_argument("sql", nargs="?", help="SQL statement")
+    write_parser.add_argument("--file", dest="sql_file", help="read SQL from file")
+
+    subparsers.add_parser("tables", help="list tables")
+
+    schema_parser = subparsers.add_parser("schema", help="show table schema")
+    schema_parser.add_argument("table_name", help="table name")
+
+    return parser
+
+
+def main() -> int:
+    """执行命令行入口。"""
+    parser = _build_parser()
+    args = parser.parse_args()
+
+    if args.command == "tables":
+        return list_tables()
+    if args.command == "schema":
+        return show_schema(args.table_name)
+    if args.command == "query":
+        sql = _read_sql(args.sql, args.sql_file)
+        return run_query(sql, limit=args.limit, allow_write=args.write)
+    if args.command == "write":
+        sql = _read_sql(args.sql, args.sql_file)
+        return run_query(sql, limit=0, allow_write=True)
+
+    parser.print_help()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())