姚章浩
4455d02262
scripts/restore_drill.py：把留档 PG 行按原文 \copy FROM 回填进目标库，核对每张表行数与清单、 抽样页面正文/标题逐字段比对，再用临时跟踪表 _drill_restored 精确删除并断言残留为 0。 实测（本地测试库 weknora_full，留档 20261009T041913Z-0b217e8a）：13,539 行全部回填成功、 8 张表行数与清单逐一致、冲突 0、抽样页正文一致（388 字符）、残留 0。 同时记录一个格式事实：留档 JSONL 是 psql \copy text 格式（转义），按行 json.loads 会失败 （knowledges 3830/3830、wiki_pages 218/1982），回填/程序化读取须先按 text 格式反转义。
208 lines
8.8 KiB
Python
208 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
||
"""留档回填演练：把发布前留档（PG 行级 JSONL）真的插进一个目标库，核对行数与内容，再清理。
|
||
|
||
为什么需要它：哈希校验只能证明"文件没坏"，不能证明"能恢复"。这个演练完整走一遍
|
||
"psql 按原文回填 → 行数与清单一致 → 抽样逐字段比对 → 精确删除插入的行（残留必须为 0）"。
|
||
|
||
格式注意：留档 JSONL 是 psql `\\copy` **text 格式**输出（反斜杠与控制字符被转义），
|
||
所以回填必须用对称的 `\\copy FROM`（由 psql 反转义）；按行 `json.loads` 会失败，
|
||
程序化读取需先反转义（见 unescape_text_format）。
|
||
|
||
用法：
|
||
python3 restore_drill.py --snapshot-dir <快照目录> --target-url <目标库URL> [--keep] [--slug <页面slug>]
|
||
|
||
目标库选择：演练用本地测试库（如 WKNFull 的 weknora_full），不要用生产库。
|
||
演练会在目标库创建并用完即删一张跟踪表 `_drill_restored`（记录插入了哪些行，保证清理精确、无残留）。
|
||
"""
|
||
import argparse
|
||
import json
|
||
import pathlib
|
||
import subprocess
|
||
import sys
|
||
|
||
# 依赖顺序：主表 → 子表；resources 在 resource_bindings 之前
|
||
TABLES = [
|
||
"knowledge_bases",
|
||
"knowledges",
|
||
"knowledge_tags",
|
||
"wiki_folders",
|
||
"wiki_pages",
|
||
"resources",
|
||
"resource_bindings",
|
||
"knowledge_tag_relations",
|
||
]
|
||
TRACKING_TABLE = "_drill_restored"
|
||
|
||
|
||
def psql(image, snapshot_dir, url, sql):
|
||
command = [
|
||
"docker", "run", "--rm", "-i", "--network", "host",
|
||
"-e", f"PGURL={url}",
|
||
"-v", f"{snapshot_dir}:/snap:ro",
|
||
image, "sh", "-c", 'psql "$PGURL" -v ON_ERROR_STOP=1 -f -',
|
||
]
|
||
process = subprocess.run(command, input=sql, text=True, capture_output=True)
|
||
if process.returncode != 0:
|
||
sys.stderr.write(process.stdout[-4000:] + "\n" + process.stderr[-4000:])
|
||
raise SystemExit(f"回填失败（psql 退出码 {process.returncode}）")
|
||
return process.stdout
|
||
|
||
|
||
def build_insert(table):
|
||
"""回填一张表：staging(jsonb 原文) → 目标表（冲突跳过）→ 记录插入了哪些行。"""
|
||
return f"""WITH ins AS (
|
||
INSERT INTO {table}
|
||
SELECT (jsonb_populate_record(NULL::{table}, s.row)).*
|
||
FROM stage_{table} s
|
||
ON CONFLICT DO NOTHING
|
||
RETURNING to_jsonb({table}) AS row_json
|
||
), rec AS (
|
||
INSERT INTO {TRACKING_TABLE} (table_name, row)
|
||
SELECT '{table}', ins.row_json FROM ins
|
||
RETURNING 1
|
||
)
|
||
SELECT 'drill:{table}' AS marker,
|
||
(SELECT count(*) FROM stage_{table}) AS staged,
|
||
(SELECT count(*) FROM rec) AS inserted;"""
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--snapshot-dir", required=True)
|
||
parser.add_argument("--target-url", required=True)
|
||
parser.add_argument("--pg-image", default="postgres:17-alpine")
|
||
parser.add_argument("--slug", default="", help="抽样比对的 wiki 页面 slug（默认自动挑一个）")
|
||
parser.add_argument("--keep", action="store_true", help="保留插入的行（默认核对后精确删除）")
|
||
arguments = parser.parse_args()
|
||
|
||
snapshot = pathlib.Path(arguments.snapshot_dir).resolve()
|
||
manifest = json.loads((snapshot / "MANIFEST.json").read_text())
|
||
expected = {item["file"].removesuffix(".jsonl"): item.get("rows")
|
||
for item in manifest["files"] if item["file"].endswith(".jsonl")}
|
||
print(f"留档: {snapshot.name} | created_at={manifest['created_at']} "
|
||
f"| 库: {', '.join(manifest['knowledge_base_ids'])}")
|
||
|
||
slug = arguments.slug
|
||
if not slug:
|
||
with (snapshot / "wiki_pages.jsonl").open("rb") as handle:
|
||
for raw in handle:
|
||
try:
|
||
row = json.loads(unescape_text_format(raw))
|
||
except Exception:
|
||
continue
|
||
if row.get("content"):
|
||
slug = row["slug"]
|
||
break
|
||
print(f"抽样页面: {slug or '（留档里没有带正文的页面）'}")
|
||
|
||
statements = [
|
||
"\\set ON_ERROR_STOP on",
|
||
f"DROP TABLE IF EXISTS {TRACKING_TABLE};",
|
||
f"CREATE TABLE {TRACKING_TABLE} (table_name text NOT NULL, row jsonb NOT NULL);",
|
||
"BEGIN;",
|
||
]
|
||
for table in TABLES:
|
||
if not (snapshot / f"{table}.jsonl").exists():
|
||
continue
|
||
statements += [
|
||
f"CREATE TEMP TABLE stage_{table} (row jsonb);",
|
||
f"\\copy stage_{table} (row) FROM '/snap/{table}.jsonl'",
|
||
build_insert(table),
|
||
]
|
||
if slug:
|
||
statements.append(f"""SELECT 'drill:sample' AS marker,
|
||
((s.r).title = t.title) AS title_equal,
|
||
((s.r).content = t.content) AS content_equal,
|
||
length((s.r).content) AS chars
|
||
FROM (SELECT jsonb_populate_record(NULL::wiki_pages, row) AS r FROM stage_wiki_pages
|
||
WHERE (jsonb_populate_record(NULL::wiki_pages, row)).slug = '{slug}') s
|
||
JOIN wiki_pages t ON t.id = (s.r).id;""")
|
||
statements.append("COMMIT;")
|
||
|
||
if not arguments.keep:
|
||
statements.append("BEGIN;")
|
||
for table in reversed(TABLES):
|
||
if not (snapshot / f"{table}.jsonl").exists():
|
||
continue
|
||
statements.append(
|
||
f"DELETE FROM {table} t USING {TRACKING_TABLE} d "
|
||
f"WHERE d.table_name = '{table}' AND to_jsonb(t) = d.row;")
|
||
statements.append(
|
||
f"SELECT 'drill:cleanup:{table}' AS marker, count(*) AS left "
|
||
f"FROM {table} t JOIN {TRACKING_TABLE} d ON d.table_name = '{table}' AND to_jsonb(t) = d.row;")
|
||
statements.append("COMMIT;")
|
||
statements.append(f"DROP TABLE {TRACKING_TABLE};")
|
||
|
||
output = psql(arguments.pg_image, snapshot, arguments.target_url, "\n".join(statements) + "\n")
|
||
|
||
staged, inserted, left, sample = {}, {}, {}, {}
|
||
for line in output.splitlines():
|
||
parts = [part.strip() for part in line.split("|")]
|
||
if not parts or not parts[0].startswith("drill:"):
|
||
continue
|
||
marker = parts[0].split(":", 1)[1]
|
||
if marker.startswith("cleanup:"):
|
||
if len(parts) != 2:
|
||
raise SystemExit(f"清理结果列数异常（{len(parts)}）: {line}")
|
||
left[marker.split(":", 1)[1]] = int(parts[1])
|
||
elif marker == "sample":
|
||
if len(parts) != 4:
|
||
raise SystemExit(f"抽样结果列数异常（{len(parts)}）: {line}")
|
||
sample = {"title_equal": parts[1] == "t", "content_equal": parts[2] == "t", "chars": parts[3]}
|
||
else:
|
||
if len(parts) != 3:
|
||
raise SystemExit(f"回填结果列数异常（{len(parts)}）: {line}")
|
||
staged[marker] = int(parts[1])
|
||
inserted[marker] = int(parts[2])
|
||
|
||
print(f"\n{'表':<26}{'清单行数':>10}{'回填尝试':>10}{'已插入':>8}{'冲突跳过':>10}{'清理后残留':>12}")
|
||
bad = []
|
||
for table in TABLES:
|
||
if table not in staged:
|
||
continue
|
||
want, got, added = expected.get(table), staged[table], inserted.get(table, 0)
|
||
residue = left.get(table)
|
||
if want != got:
|
||
bad.append(f"{table} 行数与清单不符（清单 {want} / 回填 {got}）")
|
||
if arguments.keep:
|
||
residue_text = "（保留）"
|
||
elif residue is None:
|
||
residue_text = "未验证"
|
||
bad.append(f"{table} 清理结果未取到，无法证明无残留")
|
||
else:
|
||
residue_text = str(residue)
|
||
if residue:
|
||
bad.append(f"{table} 清理后仍有 {residue} 行")
|
||
print(f"{table:<26}{str(want):>10}{got:>10}{added:>8}{got - added:>10}{residue_text:>12}")
|
||
if not sample:
|
||
bad.append("抽样页面没有完成比对（结果未取到）")
|
||
else:
|
||
print(f"\n抽样比对: 标题{'一致' if sample['title_equal'] else '不一致'} / "
|
||
f"正文{'一致' if sample['content_equal'] else '不一致'}（{sample['chars']} 字符）")
|
||
if not sample["content_equal"] or not sample["title_equal"]:
|
||
bad.append("抽样页面的正文/标题与留档不一致")
|
||
if bad:
|
||
raise SystemExit("回填演练未通过: " + "; ".join(bad))
|
||
print(f"结论: {sum(staged.values())} 行按原文回填成功（与清单一致、抽样逐字段相同）；"
|
||
+ ("插入的行保留在目标库中。" if arguments.keep else "核对后已精确删除，残留为 0。"))
|
||
|
||
|
||
def unescape_text_format(raw: bytes) -> bytes:
|
||
"""psql text 格式反转义（\\\\→\\，\\n→换行，\\t→制表符…）。"""
|
||
out = bytearray()
|
||
index = 0
|
||
while index < len(raw):
|
||
char = raw[index]
|
||
if char == 0x5C and index + 1 < len(raw):
|
||
following = raw[index + 1]
|
||
out.append({0x62: 0x08, 0x66: 0x0C, 0x6E: 0x0A, 0x72: 0x0D,
|
||
0x74: 0x09, 0x76: 0x0B}.get(following, following))
|
||
index += 2
|
||
else:
|
||
out.append(char)
|
||
index += 1
|
||
return bytes(out)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |