#!/usr/bin/env python3 """muse-example 唯一数据库通道（access-database Skill 脚本层）。 - query 卡片式打印=审查面；exec 报影响行数；apply 整文件一个事务失败全回滚。 - 失败原样抛错不静默（公约）。 连接本身归共享模块 `muse_db`：本文件只提供 CLI，不再作为其它 Skill 的 import 目标。 """ import json import sys import click import psycopg from muse_db import connect from psycopg import sql TRUNC = 160 # 卡片模式长值截断阈值（字符） # exec / execparams 只允许这些写语句；DDL 与管理类命令一律机械拒绝。 _WRITE_ALLOWED = ("INSERT", "UPDATE", "DELETE", "WITH") _DDL_OR_ADMIN = frozenset({ "CREATE", "ALTER", "DROP", "TRUNCATE", "GRANT", "REVOKE", "COMMENT", "VACUUM", "ANALYZE", "REINDEX", "CLUSTER", "REFRESH", "SECURITY", "CALL", "DO", "COPY", "LISTEN", "NOTIFY", "LOAD", "DISCARD", "SET", }) def _fmt(value, full: bool) -> str: """卡片值格式化：NULL 显示 ∅；长值默认截断并标注总长。""" if value is None: return "∅" s = str(value) if not full and len(s) > TRUNC: return s[:TRUNC] + f"…(共{len(s)}字)" return s def _first_keyword(sql_text: str) -> str: stripped = sql_text.strip() return stripped.split(None, 1)[0].upper() if stripped else "" def _require_single_statement(sql_text: str, cmd: str) -> str: stripped = sql_text.strip() if not stripped: raise click.ClickException(f"{cmd} 拒绝空语句") if ";" in stripped.rstrip(";"): raise click.ClickException(f"{cmd} 只接受单条语句（禁止多语句）") return stripped def _require_delete_has_where(sql_text: str, cmd: str) -> None: """裸 DELETE（无 WHERE）会整表清空，红线禁止。""" if _first_keyword(sql_text) != "DELETE": return if " WHERE " not in sql_text.upper(): raise click.ClickException(f"{cmd} 禁止无 WHERE 的 DELETE（裸删）") def _guard_exec(sql_text: str) -> None: """exec：拒绝 DDL/管理语句与裸 DELETE；只放行 INSERT/UPDATE/DELETE/WITH。""" stripped = _require_single_statement(sql_text, "exec") first = _first_keyword(stripped) if first in _DDL_OR_ADMIN or first not in _WRITE_ALLOWED: raise click.ClickException( f"exec 禁止 DDL/非写白名单语句（收到 {first or '空'}）；" "只允许 INSERT/UPDATE/DELETE/WITH" ) _require_delete_has_where(stripped, "exec") @click.group() def cli(): """muse-example 数据库通道""" @cli.command() @click.argument("sql_text") @click.option("--json", "as_json", is_flag=True, help="按 JSON 数组输出（给脚本消费）") @click.option("--full", is_flag=True, help="长值不截断") @click.option("--max", "max_rows", default=50, show_default=True, help="最多打印行数") def query(sql_text: str, as_json: bool, full: bool, max_rows: int): """执行 SELECT，卡片式打印（审查面）。连接只读，写语句（含 INSERT...RETURNING）会被 PG 拒。""" with connect(readonly=True) as conn: cur = conn.execute(sql_text) if cur.description is None: raise click.ClickException("该语句无结果集——写操作请用 exec") cols = [d.name for d in cur.description] rows = cur.fetchall() if as_json: out = [dict(zip(cols, r)) for r in rows[:max_rows]] click.echo(json.dumps(out, ensure_ascii=False, default=str, indent=1)) else: width = max((len(c) for c in cols), default=0) for i, r in enumerate(rows[:max_rows], 1): click.echo(f"── 行 {i} " + "─" * 30) for c, v in zip(cols, r): click.echo(f"{c:<{width}} │ {_fmt(v, full)}") if len(rows) > max_rows: click.echo(f"…（共 {len(rows)} 行，仅示 {max_rows} 行，--max 放宽）") else: click.echo(f"（共 {len(rows)} 行）") @cli.command("exec") @click.argument("sql_text") def exec_(sql_text: str): """执行单条写操作（INSERT/UPDATE/DELETE，须带 WHERE），报影响行数。禁止 DDL。""" _guard_exec(sql_text) with connect() as conn: cur = conn.execute(sql_text) conn.commit() n = cur.rowcount click.echo(f"OK（影响 {n} 行）" if n >= 0 else "OK") def _read_params(stdin_json: bool, param_tuple: tuple, stream=None) -> list: """收集参数：--stdin 时从 stream（默认 stdin）读 JSON 数组；否则用 --param 按序。两者互斥。 大对象（raw 全文）走 --stdin：内容经 JSON 数组从标准输入流入，避开 shell 转义与 ARG_MAX。 """ if stdin_json: if param_tuple: raise click.ClickException("--stdin 与 --param 互斥，二选一") raw = (stream if stream is not None else sys.stdin).read() try: data = json.loads(raw) except json.JSONDecodeError as e: raise click.ClickException(f"--stdin 参数不是合法 JSON：{e}") if not isinstance(data, list): raise click.ClickException("--stdin 参数必须是 JSON 数组，元素按 %s 顺序排列") return data return list(param_tuple) # execparams 只允许这些种类的单条语句（参数化写）；拒绝 DO/CALL/DDL/TRUNCATE/多语句 _EXECPARAMS_ALLOWED = _WRITE_ALLOWED def _guard_execparams(sql_text: str) -> None: stripped = _require_single_statement(sql_text, "execparams") first = _first_keyword(stripped) if first not in _EXECPARAMS_ALLOWED: raise click.ClickException( f"execparams 只做参数化写（INSERT/UPDATE/DELETE/WITH），收到 {first or '空'} 语句") _require_delete_has_where(stripped, "execparams") @cli.command("execparams") @click.argument("sql_text") @click.option("--stdin", "stdin_json", is_flag=True, help="参数从 stdin 读入（JSON 数组，按 %s 顺序）；大对象用，避开 shell 转义与 ARG_MAX") @click.option("--param", "param_tuple", multiple=True, help="按序传参（小值，可重复）；与 --stdin 互斥") def execparams(sql_text: str, stdin_json: bool, param_tuple: tuple): """参数化写操作：SQL 用 %s 占位，参数走服务端绑定（防注入；大内容不拼命令行）。""" _guard_execparams(sql_text) params = _read_params(stdin_json, param_tuple) with connect() as conn: cur = conn.execute(sql_text, params) if cur.description is not None: raise click.ClickException("该语句有结果集——查询请用 query（execparams 只做写）") conn.commit() n = cur.rowcount click.echo(f"OK（影响 {n} 行）" if n >= 0 else "OK") @cli.command() @click.argument("file", type=click.Path(exists=True)) def apply(file: str): """应用 SQL 文件：整文件一个事务（无参数走 simple protocol 支持多语句），失败全回滚。""" with open(file, encoding="utf-8") as f: sql = f.read() with connect() as conn: conn.execute(sql) conn.commit() click.echo(f"已应用: {file}") @cli.command() def tables(): """表清单 + 活行数（有 deleted 列的表计 deleted=FALSE，其余计全行）。""" with connect() as conn: names = [r[0] for r in conn.execute( "SELECT tablename FROM pg_tables WHERE schemaname='public' ORDER BY tablename")] for t in names: has_deleted = conn.execute( "SELECT 1 FROM information_schema.columns WHERE table_name=%s AND column_name='deleted'", (t,)).fetchone() # 表名用 sql.Identifier 安全拼接（不手工拼双引号），避免恶意表名破坏语句 if has_deleted: stmt = sql.SQL("SELECT count(*) FROM {} WHERE deleted = FALSE").format(sql.Identifier(t)) else: stmt = sql.SQL("SELECT count(*) FROM {}").format(sql.Identifier(t)) n = conn.execute(stmt).fetchone()[0] click.echo(f"{t:<45} {n}") if __name__ == "__main__": try: cli() except psycopg.Error as e: # 数据库错误原样透出（可追踪、不静默） click.echo(f"[db错误] {type(e).__name__}: {e}", err=True) sys.exit(1)