280 lines
12 KiB
Python
280 lines
12 KiB
Python
"""对**真实 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())
|