scripts/prove_hub_bootstrap_local.py

#!/usr/bin/env python3
"""Download catalog from Hub into a temp DB (same path as HF bootstrap after first boot).

Uses ``[huggingface].token`` from parser ``secrets.toml`` for private datasets; public
datasets work without it.

Example::

  python scripts/prove_hub_bootstrap_local.py --filename store_slim_bootstrap.db
"""
from __future__ import annotations

import argparse
import os
import sys
import tomllib
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_SECRETS = Path(os.environ.get("KINK_PARSER_SECRETS_TOML", Path.home() / "PycharmProjects/parser/secrets.toml"))


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--secrets-toml", type=Path, default=DEFAULT_SECRETS)
    ap.add_argument("--repo-id", default="ronheichman/kink-catalog-slim")
    ap.add_argument("--filename", default="store_slim_bootstrap.db")
    ap.add_argument("--dest", type=Path, default=Path("/tmp/kink_hub_bootstrap_proof.db"))
    args = ap.parse_args()

    dest = args.dest.resolve()
    if dest.exists():
        dest.unlink()

    token = ""
    if args.secrets_toml.is_file():
        with args.secrets_toml.open("rb") as f:
            data = tomllib.load(f)
        token = ((data.get("huggingface") or {}).get("token") or "").strip()
    if token:
        os.environ["HF_TOKEN"] = token

    os.environ["KINK_HF_DATASET_REPO"] = args.repo_id
    os.environ["KINK_HF_DATASET_FILENAME"] = args.filename
    os.environ["KINK_STORE_PATH"] = str(dest)
    os.environ["KINK_HF_REQUIRE_FULL_CATALOG"] = "1"
    os.environ.pop("KINK_CATALOG_URL", None)
    for k in ("B2_KEY_ID", "B2_APPLICATION_KEY", "B2_BUCKET", "B2_REGION"):
        os.environ.pop(k, None)

    sys.path.insert(0, str(ROOT))
    from backend.hf_bootstrap import ensure_store_db

    print(f"prove_hub_bootstrap_local: hf_hub_download → {dest} …", flush=True)
    ensure_store_db(dest)

    import sqlite3

    conn = sqlite3.connect(str(dest))
    try:
        nk = conn.execute("SELECT COUNT(*) FROM kink").fetchone()[0]
        np = conn.execute(
            "SELECT COUNT(*) FROM kinkcontenttype WHERE content_kind='play'"
        ).fetchone()[0]
    finally:
        conn.close()

    print(f"prove_hub_bootstrap_local: OK — {nk} kinks, {np} plays", flush=True)
    return 0


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