scripts/sync_hf_space_b2_catalog.py

#!/usr/bin/env python3
"""Upload the slim bootstrap SQLite to B2 and push B2 credentials to a Hugging Face Space.

Reads ``[b2]`` and ``[huggingface].token`` from parser ``secrets.toml`` (same defaults as
other scripts). Uploads ``data/store_slim_b2.db`` by default (output of
``shrink_store_for_remote_bootstrap.py`` — full kink rows, ~280 MiB).

Then calls ``HfApi.add_space_secret`` for ``B2_*`` and ``add_space_variable`` for
``KINK_B2_OBJECT_KEY`` so the Space bootstraps via S3 **before** Hub (see
``backend/hf_bootstrap``).

Requires:
  export HF_SPACE_REPO=owner/slug

Does not print secret values.

After success, rebuild the Space so the container picks up secrets::

  python scripts/publish_hf_space.py --verify
"""
from __future__ import annotations

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

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from backend.b2_toml import load_b2_section, resolve_b2_config_path
from backend.repo_dotenv import load_repo_dotenv
DEFAULT_SECRETS = Path(os.environ.get("KINK_PARSER_SECRETS_TOML", Path.home() / "PycharmProjects/parser/secrets.toml"))


def _load_hf_token(path: Path) -> str:
    t = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "").strip()
    if t:
        return t
    with path.open("rb") as f:
        data = tomllib.load(f)
    hf = (data.get("huggingface") or {}).get("token") or ""
    if not str(hf).strip():
        raise SystemExit(f"Missing [huggingface].token in {path}")
    return str(hf).strip()


def main() -> int:
    load_repo_dotenv()
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--secrets-toml",
        type=Path,
        default=DEFAULT_SECRETS,
        help="Parser secrets with [huggingface].token (and [b2] if --b2-secrets-toml not set)",
    )
    ap.add_argument(
        "--b2-secrets-toml",
        type=Path,
        default=None,
        help="TOML with only [b2] (overrides env KINK_B2_SECRETS_TOML and [b2] from --secrets-toml)",
    )
    ap.add_argument(
        "--local-path",
        type=Path,
        default=ROOT / "data" / "store_slim_b2.db",
        help="SQLite file to upload (default: shrunk bootstrap DB)",
    )
    ap.add_argument(
        "--key",
        default="kink/catalog/store_slim_bootstrap.db",
        help="S3 object key (Space default KINK_B2_OBJECT_KEY must match)",
    )
    ap.add_argument("--skip-upload", action="store_true", help="Only push HF secrets/variables")
    args = ap.parse_args()

    repo_id = (os.environ.get("HF_SPACE_REPO") or "").strip()
    if not repo_id:
        print("sync_hf_space_b2_catalog: set HF_SPACE_REPO=owner/slug", file=sys.stderr)
        return 1

    if not args.secrets_toml.is_file():
        print(f"sync_hf_space_b2_catalog: secrets not found: {args.secrets_toml}", file=sys.stderr)
        return 1
    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"sync_hf_space_b2_catalog: 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"sync_hf_space_b2_catalog: {e}", file=sys.stderr)
        return 1
    try:
        hf_token = _load_hf_token(args.secrets_toml)
    except SystemExit as e:
        print(f"sync_hf_space_b2_catalog: {e}", file=sys.stderr)
        return 1
    local = args.local_path.resolve()
    if not args.skip_upload:
        if not local.is_file():
            print(f"sync_hf_space_b2_catalog: missing {local} — run shrink_store_for_remote_bootstrap.py first", file=sys.stderr)
            return 1

    import boto3
    from huggingface_hub import HfApi

    region = b2["region"]
    endpoint = f"https://s3.{region}.backblazeb2.com"
    client = boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=b2["key_id"],
        aws_secret_access_key=b2["application_key"],
        region_name=region,
    )
    env_bucket = (os.environ.get("KINK_B2_CATALOG_BUCKET") or os.environ.get("B2_CATALOG_BUCKET") or "").strip()
    bucket = env_bucket or b2["bucket_name"]
    key = args.key.lstrip("/")
    if env_bucket:
        print(f"sync_hf_space_b2_catalog: using bucket from env: {bucket}", file=sys.stderr)

    if not args.skip_upload:
        print(
            f"sync_hf_space_b2_catalog: uploading {local} ({local.stat().st_size / 1e6:.1f} MB) → s3://{bucket}/{key}",
            file=sys.stderr,
        )
        try:
            client.upload_file(str(local), bucket, key, ExtraArgs={"ContentType": "application/x-sqlite3"})
        except Exception as e:
            err = str(e).lower()
            if "storage cap" in err or "cap exceeded" in err:
                print(
                    "sync_hf_space_b2_catalog: B2 **account** Caps & Alerts (storage $/day) hit — "
                    "not fixed by a new object key or new empty bucket on the same account. "
                    "Raise/remove caps in the B2 console, or use another account / KINK_B2_CATALOG_BUCKET on an account with room.",
                    file=sys.stderr,
                )
            raise

    api = HfApi(token=hf_token)

    def sec(name: str, value: str, desc: str) -> None:
        api.add_space_secret(repo_id, name, value, description=desc, token=hf_token)

    sec("B2_KEY_ID", b2["key_id"], "Backblaze B2 application key id (S3)")
    sec("B2_APPLICATION_KEY", b2["application_key"], "Backblaze B2 application key secret")
    sec("B2_BUCKET", bucket, "Backblaze B2 bucket for catalog object")
    sec("B2_REGION", region, "Backblaze B2 region (e.g. us-east-005)")

    api.add_space_variable(
        repo_id,
        "KINK_B2_OBJECT_KEY",
        key,
        description="S3 object key for store SQLite (bootstrap)",
        token=hf_token,
    )

    print("sync_hf_space_b2_catalog: Space secrets + KINK_B2_OBJECT_KEY updated.", file=sys.stderr)
    print("sync_hf_space_b2_catalog: Rebuild the Space:", file=sys.stderr)
    print(f"  HF_TOKEN=*** HF_SPACE_REPO={repo_id} python scripts/publish_hf_space.py --verify", file=sys.stderr)
    return 0


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