scripts/prove_b2_bootstrap_local.py

#!/usr/bin/env python3
"""Prove ``ensure_store_db`` can fill ``KINK_STORE_PATH`` from B2 S3 (same path as HF).

Uses parser ``secrets.toml`` ``[b2]`` and default object key
``kink/catalog/store_slim_bootstrap.db`` (override with ``--key``).

Example::

  python scripts/prove_b2_bootstrap_local.py
  python scripts/prove_b2_bootstrap_local.py --key kink/catalog/store_slim.db
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

from backend.b2_toml import load_b2_section, resolve_b2_config_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, help="Fallback TOML with [b2]")
    ap.add_argument("--b2-secrets-toml", type=Path, default=None, help="TOML with only [b2] (overrides env and --secrets-toml)")
    ap.add_argument("--key", default="kink/catalog/store_slim_bootstrap.db")
    ap.add_argument(
        "--dest",
        type=Path,
        default=Path("/tmp/kink_b2_bootstrap_proof.db"),
        help="Fresh path for downloaded DB",
    )
    args = ap.parse_args()

    b2_path = resolve_b2_config_path(cli_b2=args.b2_secrets_toml, cli_fallback=args.secrets_toml)
    if not b2_path.is_file():
        print(f"prove_b2_bootstrap_local: B2 secrets TOML not found: {b2_path}", file=sys.stderr)
        return 1
    try:
        b2 = load_b2_section(b2_path)
    except SystemExit as e:
        print(f"prove_b2_bootstrap_local: {e}", file=sys.stderr)
        return 1

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

    os.environ["B2_KEY_ID"] = str(b2["key_id"]).strip()
    os.environ["B2_APPLICATION_KEY"] = str(b2["application_key"]).strip()
    os.environ["B2_BUCKET"] = str(b2["bucket_name"]).strip()
    os.environ["B2_REGION"] = str(b2["region"]).strip()
    os.environ["KINK_B2_OBJECT_KEY"] = args.key.lstrip("/")
    os.environ.pop("KINK_CATALOG_URL", None)
    os.environ["KINK_STORE_PATH"] = str(dest)
    os.environ["KINK_HF_REQUIRE_FULL_CATALOG"] = "0"

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

    print(f"prove_b2_bootstrap_local: downloading B2 → {dest} …", flush=True)
    out = ensure_store_db(dest)
    assert out.resolve() == dest.resolve()

    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_b2_bootstrap_local: OK — {nk} kinks, {np} plays in {dest}", flush=True)
    return 0


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