muse-agent-example/.claude/skills/db/scripts/db.py

#!/usr/bin/env python3
"""muse-example 唯一数据库通道（db skill 脚本层）。

- DSN 锁死 muse-example：严禁触碰共享 PG 上其他库（muse_local / muse_slice_live / *_test）。
- query 卡片式打印=审查面；exec 报影响行数；apply 整文件一个事务失败全回滚。
- 失败原样抛错不静默（公约）。
"""
import json
import sys

import click
import psycopg
from psycopg import sql

# 连接事实与凭据来源：db/连接信息.md（内网 Tailscale 段，凭据明文入仓为既定政策）
DSN = "postgresql://root:f6710e2d0294eb1c10e26a805a64bc54@100.64.0.8:5433/muse-example"
TRUNC = 160  # 卡片模式长值截断阈值（字符）


def connect(readonly: bool = False):
    """统一连接入口：复用锁死的 muse-example DSN（即开即关，Tailscale 不持长事务）。

    readonly=True 时会话级锁死只读（写语句被 PG 直接拒）——query 命令与看板用。
    其它 skill 的写路径需要参数化短连接时，`from db import connect` 复用同一 DSN，
    不要各自硬编码连接串（仿 llm skill 的 _bump_window）。
    """
    if readonly:
        return psycopg.connect(DSN, options="-c default_transaction_read_only=on")
    return psycopg.connect(DSN)


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


@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/DDL），报影响行数。"""
    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 = ("INSERT", "UPDATE", "DELETE", "WITH")


def _guard_execparams(sql_text: str) -> None:
    stripped = sql_text.strip()
    if ";" in stripped.rstrip(";"):
        raise click.ClickException("execparams 只接受单条语句（禁止多语句）")
    first = stripped.split(None, 1)[0].upper() if stripped else ""
    if first not in _EXECPARAMS_ALLOWED:
        raise click.ClickException(
            f"execparams 只做参数化写（INSERT/UPDATE/DELETE/WITH），收到 {first or '空'} 语句")


@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)