Create-knowledge-find-decode/scripts/validate_db.py

"""M1 校验：连库、确认 schema/表/列齐全、做一次可逆写探测。

用法（在能连到 RDS 的机器上，如海外开发机）：
    python scripts/validate_db.py [path/to/.env]
"""
from __future__ import annotations

import sys

from creation_knowledge.config import PgConfig
from creation_knowledge.integrations.db import CkStore, _connect
from creation_knowledge.models import Post

EXPECTED = {
    "ck_post": {"id", "platform", "url", "raw", "extracted", "screening",
                "stage", "created_at", "updated_at"},
    "ck_knowledge_item": {"id", "post_id", "item", "deconstruction",
                          "ingest_payload", "ingest_status", "knowledge_id",
                          "created_at", "updated_at"},
}
PROBE_ID = "_validate_probe"


def main() -> int:
    env_file = sys.argv[1] if len(sys.argv) > 1 else ".env"
    cfg = PgConfig.from_env(env_file)
    print(f"[cfg] host={cfg.host} db={cfg.database} schema={cfg.schema} user={cfg.user}")

    # 1) 连通 + 版本
    with _connect(cfg) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT version()")
            print("[conn] OK:", cur.fetchone()[0][:70])

    store = CkStore(cfg)

    # 2) 表/列齐全
    ok = True
    for table, expected_cols in EXPECTED.items():
        cols = set(store.table_columns(table))
        if not cols:
            print(f"[schema] MISSING table: {cfg.schema}.{table}")
            ok = False
            continue
        missing = expected_cols - cols
        if missing:
            print(f"[schema] {table} 缺列: {sorted(missing)}")
            ok = False
        else:
            print(f"[schema] {table} OK ({len(cols)} cols)")
    if not ok:
        print("RESULT: FAIL（表结构不符）")
        return 1

    # 3) 可逆写探测：插入一帖+一片段，读回，删除
    store.delete_post(PROBE_ID)  # 清理上次残留
    try:
        store.upsert_post(Post(id=PROBE_ID, url="probe://x", content_id="x",
                               raw={"probe": True}))
        store.set_extracted(PROBE_ID, {"text": "probe", "is_empty": False})
        item_id = store.save_item(PROBE_ID, {"title": "probe item",
                                             "knowledge_types": ["how"]},
                                  {"stages": ["脚本"]}, {"title": "probe"})
        post = store.read_post(PROBE_ID)
        items = store.read_items(PROBE_ID)
        assert post and post["stage"] == "extracted", post
        assert items and items[0]["id"] == item_id, items
        print(f"[write] OK: post.stage={post['stage']}, item_id={item_id}, "
              f"ingest_status={items[0]['ingest_status']}")
    finally:
        store.delete_post(PROBE_ID)
        print("[write] 探测数据已清理")

    print("RESULT: PASS")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())