tools/pg_check.py

"""对**真实 PostgreSQL 服务器**做台账契约一致性校验。

为什么需要它：本工具开发机上没有任何 PostgreSQL（无服务、无 Docker），因此
PG 后端的真实读写语义无法在开发环境里验证。这个脚本让你**在自己的库上**
一次性验完，而且：

* **全程在一个事务里，最后 ``ROLLBACK``** —— 连建表都在事务内，
  对你的真实库**零残留**（表本来存在就什么都不变）；
* 用**同一套契约**分别跑本地 SQLite 与 PostgreSQL，然后**逐项对比**，
  这样"两个后端行为不一致"会直接暴露成 FAIL，而不是等同事用出问题才发现；
* 写入的记录 ``lic_id`` 一律带 ``PGCHECK-`` 前缀，且**所有观测值都只统计这些
  记录** —— 库里已有的真实台账既不参与对比，也不会被打印出来；
* 不会改动你的任何真实数据（验证结束后事务回滚）。

用法::

    # 1) 显式参数（密码走环境变量或交互输入，避免留在命令历史里）
    uv run python tools/pg_check.py --host 127.0.0.1 --port 5432 \\
        --dbname chongming_licenses --user license_ro

    # 2) 或用 DSN
    uv run python tools/pg_check.py --dsn "postgresql://license_ro@127.0.0.1:5432/chongming_licenses"

    # 密码来源：环境变量 PGPASSWORD 或交互式提问（推荐）；
    # 也支持 --password（会留在命令历史里，仅用于自动化）

退出码 0 = 两个后端行为一致。
"""
from __future__ import annotations

import argparse
import getpass
import os
import sys
import tempfile
from datetime import timedelta
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

from chongming_license import (  # noqa: E402
    Registry,
    build_license,
    build_payload,
    generate_keypair,
    machine_code,
    utc_today,
    write_license,
)
from chongming_license_gui import pg_store  # noqa: E402
from chongming_license_gui.pg_store import (  # noqa: E402
    DEFAULT_DBNAME,
    DEFAULT_HOST,
    DEFAULT_PORT,
    DEFAULT_SCHEMA,
    DEFAULT_SSLMODE,
    SSL_MODES,
    PgConfig,
    PostgresRegistry,
)

PREFIX = "PGCHECK-"
STATUSES = ("expiring", "forever", "active", "expired", "not_yet_valid")
failures: list[str] = []


def check(name: str, ok: bool, detail: str = "") -> None:
    print(f"[{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))
    if not ok:
        failures.append(name)


def _mine(records) -> list[str]:
    """只保留本次校验写入的记录（真实台账不参与对比、也不被打印）。"""
    return [r.lic_id for r in records if r.lic_id.startswith(PREFIX)]


# ---------------------------------------------------------------------------
# 契约：在任意后端上跑同一套操作，返回"可比较的观测值"
# ---------------------------------------------------------------------------
def _demo_docs(sandbox: Path) -> list[dict]:
    """两条用于校验的授权：一条即将到期、一条永久（覆盖排序哨兵与状态分桶）。

    客户名刻意混入 ASCII（``PGCheck``），用来验证"大小写不敏感"这个方言差异点。
    """
    private_key, _ = generate_keypair()
    today = utc_today()
    sources = {
        "board_serial": "pgcheckboard01",
        "product_uuid": "pgcheckproductuuid0123456789ab",
        "mac": "001b44pgcheck",
        "machine_id": "pgcheckmachineid0123456789abcdef",
    }
    specs = [
        (f"{PREFIX}EXPIRING", "PGCheck 契约校验（示例）", ["aibox_sophon_adapter"],
         (today + timedelta(days=20)).isoformat()),
        (f"{PREFIX}FOREVER", "PGCheck 永久（示例）",
         ["aibox_sophon_adapter", "ip_sound_column"], None),
    ]
    docs = []
    for lic_id, customer, plugins, expires in specs:
        payload = build_payload(
            machine_code=machine_code(sources), plugins=plugins, sources=sources,
            customer=customer, lic_id=lic_id, expires_at=expires, note="载荷备注",
        )
        doc = build_license(payload, private_key, kid="pg-check")
        write_license(doc, sandbox / f"{lic_id}.lic")
        docs.append(doc)
    return docs


def contract(reg, docs: list[dict], sandbox: Path) -> dict:
    """同一套操作 → 观测值。两边结果必须逐项相同。"""
    ids = [d["payload"]["lic_id"] for d in docs]
    mine = lambda records: _mine(records)  # noqa: E731 - 短小，读起来更直观
    obs: dict[str, object] = {}

    for doc in docs:
        reg.save(doc, public_key_fp="PGCHECK-FP", source="pg-check")

    # 排序：到期日升序，永久（NULL）用哨兵排最后
    obs["顺序(到期升序/永久最后)"] = mine(reg.query())
    obs["客户模糊检索"] = mine(reg.query(customer="永久"))
    obs["客户检索大小写不敏感"] = mine(reg.query(customer="pgcheck"))
    obs["插件精确检索"] = sorted(mine(reg.query(plugin="ip_sound_column")))
    obs["关键字检索"] = sorted(mine(reg.query(keyword="契约校验")))
    obs["状态筛选(永久)"] = mine(reg.query(status="forever"))
    obs["状态分桶"] = {s: len(mine(reg.query(status=s))) for s in STATUSES}

    rec = reg.get(ids[0])
    obs["get 命中"] = rec.lic_id if rec else None
    obs["机器码已大写"] = bool(rec) and rec.machine_code == rec.machine_code.upper()
    obs["来源指纹回读一致"] = bool(rec) and rec.sources == docs[0]["payload"]["sources"]

    # 人工批注：手工改过之后，重新入库不得被载荷备注冲掉
    obs["update_note 返回值"] = reg.update_note(ids[0], "人工批注")
    reg.save(docs[0], public_key_fp="PGCHECK-FP", source="renew")
    obs["重新入库后仍是人工批注"] = reg.get(ids[0]).note

    # 首次入库时间要保留（库的 SQLite 版把 created_at 排除在 DO UPDATE 之外）
    created_before = reg.get(ids[0]).created_at
    reg.save(docs[0], public_key_fp="PGCHECK-FP", source="renew")
    obs["重新入库后 created_at 不变"] = reg.get(ids[0]).created_at == created_before

    # 导出：文件名规则、CSV 表头与编码（这些方法是从库继承的，两边应完全一致）
    my_records = [r for r in reg.query() if r.lic_id.startswith(PREFIX)]
    out = sandbox / "export"
    obs["导出 .lic 文件名"] = sorted(p.name for p in reg.export_licenses(out, my_records))
    csv_path = reg.export_csv(sandbox / "台账.csv", my_records)
    raw = csv_path.read_bytes()
    obs["CSV 带 BOM"] = raw[:3] == b"\xef\xbb\xbf"
    lines = csv_path.read_text(encoding="utf-8-sig").splitlines()
    obs["CSV 行数(含表头)"] = len(lines)
    obs["CSV 表头列数"] = len(lines[0].split(",")) if lines else 0

    obs["delete 返回值"] = reg.delete(ids[1])
    obs["删除后剩余"] = mine(reg.query())
    return obs


# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="在真实 PostgreSQL 上校验台账后端与本地 SQLite 行为一致",
    )
    parser.add_argument("--dsn",
                        help="连接串（postgresql://用户@主机:端口/库名）；给了它就不用下面各项")
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"默认 {DEFAULT_PORT}")
    parser.add_argument("--dbname", default=DEFAULT_DBNAME)
    parser.add_argument("--user", default="")
    parser.add_argument("--password", default=None,
                        help="不建议：会留在命令历史里。优先用 PGPASSWORD 或交互输入")
    parser.add_argument("--sslmode", default=DEFAULT_SSLMODE, choices=list(SSL_MODES))
    parser.add_argument("--schema", default=DEFAULT_SCHEMA)
    return parser.parse_args()


def resolve_config(args: argparse.Namespace) -> tuple[PgConfig, str]:
    """返回 (配置, 密码)。密码优先级：--password > DSN 内 > PGPASSWORD。"""
    if args.dsn:
        # 连接串解析交给 pg_store（进而交给 psycopg），避免自己写一套连接串文法
        config = PgConfig.from_dsn(args.dsn, schema=args.schema)
        password = config.password
    else:
        config = PgConfig(host=args.host, port=args.port, dbname=args.dbname,
                          user=args.user, sslmode=args.sslmode, schema=args.schema)
        password = ""
    if args.password:
        password = args.password
    return config, password or os.environ.get("PGPASSWORD", "")


def main() -> int:
    args = parse_args()
    try:
        config, password = resolve_config(args)
    except ImportError:
        print("未安装 PostgreSQL 驱动：uv sync --extra postgres")
        return 2
    if not config.dbname:
        print("请指定数据库名（--dbname 或 --dsn）")
        return 2
    if not password:
        password = getpass.getpass(f"{config.user or 'postgres'}@{config.host} 的密码（不显示）：")

    print("=" * 70)
    print("台账契约一致性校验：本地 SQLite vs PostgreSQL")
    print(f"目标：{config.display()}")          # 展示串不含密码
    print("说明：所有写入都在一个事务里，结束时 ROLLBACK；")
    print("      观测值只统计本次的 PGCHECK-* 记录，你的真实台账不参与、也不会被打印。")
    print("=" * 70)

    sandbox = Path(tempfile.mkdtemp(prefix="clg-pgcheck-"))
    docs = _demo_docs(sandbox)

    # ---- 1) 本地 SQLite 基准 ----
    print("\n--- 本地 SQLite（基准）---")
    sqlite_sandbox = sandbox / "sqlite"
    sqlite_sandbox.mkdir(parents=True, exist_ok=True)
    with Registry(sqlite_sandbox / "ledger.db") as sqlite_reg:
        sqlite_obs = contract(sqlite_reg, docs, sqlite_sandbox)
    for key, value in sqlite_obs.items():
        print(f"    {key} = {value}")

    # ---- 2) PostgreSQL（一个事务，最后回滚）----
    print("\n--- PostgreSQL（事务内执行，结束回滚）---")
    pg_sandbox = sandbox / "pg"
    pg_sandbox.mkdir(parents=True, exist_ok=True)
    try:
        reg = PostgresRegistry(config, password=password, autocommit=False)
    except Exception as exc:  # noqa: BLE001
        print(f"[FAIL] 连接失败：{exc}")
        return 1

    try:
        pg_obs = contract(reg, docs, pg_sandbox)
    except Exception as exc:  # noqa: BLE001 - 把差异当成校验失败而不是崩溃
        print(f"[FAIL] 在 PostgreSQL 上执行契约时出错：{type(exc).__name__}: {exc}")
        reg._conn.rollback()
        reg.close()
        return 1

    # ---- 3) 逐项对比 ----
    print("\n--- 逐项对比（SQLite 基准 vs PostgreSQL）---")
    for key, expected in sqlite_obs.items():
        actual = pg_obs.get(key)
        check(key, actual == expected, "" if actual == expected
              else f"SQLite={expected!r}  PG={actual!r}")

    # ---- 4) 回滚并确认没有残留 ----
    reg._conn.rollback()
    table = f'"{config.schema}"."{pg_store.TABLE}"'
    try:
        row = reg._conn.execute(
            f"SELECT count(*) AS n FROM {table} WHERE lic_id LIKE '{PREFIX}%'"
        ).fetchone()
        residue = row["n"] if row else 0
        print(f"\n回滚完成：库里残留的校验记录数 = {residue}")
        if residue:
            print("[FAIL] 残留了校验记录 —— 这不该发生")
            failures.append("回滚后有残留")
    except Exception as exc:  # noqa: BLE001
        # 表本来不存在 → 连建表一起被回滚了，这正是期望结果
        print(f"\n回滚完成：表不存在（建表也随事务一起回滚）—— {type(exc).__name__}")
    reg.close()

    print("\n" + "=" * 70)
    if failures:
        print(f"❌ {len(failures)} 项在两个后端之间不一致：{failures}")
        print("   请把上面的差异贴给研发 —— 这正是本脚本存在的意义。")
        return 1
    print("✅ PostgreSQL 后端与本地 SQLite 行为完全一致（且已在事务中回滚，无残留）")
    return 0


if __name__ == "__main__":
    sys.exit(main())