scripts/e2e_deploy_kink_catalog.py

#!/usr/bin/env python3
"""One-shot end-to-end: shrink DB → (optional B2 bucket) → catalog to B2 or Hub → HF Space secrets/vars → publish + verify.

**Requires**
  ``export HF_SPACE_REPO=owner/slug`` (or pass ``--space-repo``).

Reads ``parser`` ``secrets.toml`` for ``[huggingface].token`` (and ``[b2]`` unless overridden).
Optional ``--b2-secrets-toml`` or ``KINK_B2_SECRETS_TOML`` supplies ``[b2]`` only (e.g. a new B2 account).

**Default bootstrap mode** (``--bootstrap auto``):
  1. Build ``data/store_slim_b2.db`` from ``data/store_slim.db`` (unless ``--skip-shrink``).
  2. Try **B2**: upload + push ``B2_*`` secrets + ``KINK_B2_OBJECT_KEY`` + ``verify_b2_catalog_http``.
  3. If B2 upload fails (e.g. storage cap), **Hub fallback**: remove B2 Space secrets so bootstrap does not stop on B2,
     upload to Hub dataset + set ``KINK_HF_DATASET_*`` variables, then continue.
  4. **Publish** the Space repo and **verify** live (same as ``publish_hf_space.py --verify``).

Does not print tokens or application keys.

Example::

  export HF_SPACE_REPO=ronheichman/kink-discovery
  python scripts/e2e_deploy_kink_catalog.py

  # Optional: create bucket first (needs writeBuckets on key or --admin master key env)
  python scripts/e2e_deploy_kink_catalog.py --create-bucket kink-catalog-myname

  # B2 only (no Hub fallback)
  python scripts/e2e_deploy_kink_catalog.py --bootstrap b2

  # Hub only (skip B2)
  python scripts/e2e_deploy_kink_catalog.py --bootstrap hub
"""
from __future__ import annotations

import argparse
import os
import subprocess
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 _run(py_args: list[str], env: dict | None = None) -> None:
    cmd = [sys.executable, str(ROOT / "scripts" / py_args[0])] + py_args[1:]
    e = {**os.environ, **(env or {})}
    print(f"e2e_deploy: {' '.join(cmd)}", file=sys.stderr, flush=True)
    r = subprocess.run(cmd, cwd=str(ROOT), env=e)
    if r.returncode != 0:
        raise SystemExit(r.returncode)


def _load_hf_token(secrets_path: Path) -> str:
    with secrets_path.open("rb") as f:
        data = tomllib.load(f)
    t = ((data.get("huggingface") or {}).get("token") or "").strip()
    if not t:
        raise SystemExit(f"e2e_deploy: missing [huggingface].token in {secrets_path}")
    return t


def _delete_b2_space_config(api, repo_id: str, hf_token: str) -> None:
    """Remove B2 secrets/vars so Space bootstrap can fall through to Hub."""
    from huggingface_hub.utils import HfHubHTTPError

    for key in ("B2_KEY_ID", "B2_APPLICATION_KEY", "B2_BUCKET", "B2_REGION"):
        try:
            api.delete_space_secret(repo_id, key, token=hf_token)
            print(f"e2e_deploy: removed Space secret {key}", file=sys.stderr)
        except HfHubHTTPError as exc:
            if getattr(exc, "response", None) is not None and exc.response.status_code in (400, 404):
                continue
            print(f"e2e_deploy: delete_space_secret {key}: {exc}", file=sys.stderr)
        except Exception as exc:
            print(f"e2e_deploy: delete_space_secret {key}: {exc}", file=sys.stderr)
    try:
        api.delete_space_variable(repo_id, "KINK_B2_OBJECT_KEY", token=hf_token)
        print("e2e_deploy: removed Space variable KINK_B2_OBJECT_KEY", file=sys.stderr)
    except Exception as exc:
        print(f"e2e_deploy: delete_space_variable KINK_B2_OBJECT_KEY: {exc}", file=sys.stderr)


def _hub_fallback(
    secrets_path: Path,
    space_repo: str,
    local_db: Path,
    hub_repo: str,
    remote_filename: str,
    hf_token: str,
) -> None:
    from huggingface_hub import HfApi

    api = HfApi(token=hf_token)
    _delete_b2_space_config(api, space_repo, hf_token)

    api.create_repo(hub_repo, repo_type="dataset", exist_ok=True)
    print(
        f"e2e_deploy: Hub fallback — uploading {local_db} → {hub_repo}/{remote_filename}",
        file=sys.stderr,
    )
    api.upload_file(
        path_or_fileobj=str(local_db),
        path_in_repo=remote_filename,
        repo_id=hub_repo,
        repo_type="dataset",
    )
    api.add_space_variable(
        space_repo,
        "KINK_HF_DATASET_FILENAME",
        remote_filename,
        description="Bootstrap SQLite (Hub) from e2e_deploy_kink_catalog",
        token=hf_token,
    )
    api.add_space_variable(
        space_repo,
        "KINK_HF_DATASET_REPO",
        hub_repo,
        description="Hub dataset for catalog",
        token=hf_token,
    )
    print("e2e_deploy: Hub fallback Space variables set.", file=sys.stderr)


def _try_b2_chain(
    secrets_path: Path,
    space_repo: str,
    local_db: Path,
    b2_object_key: str,
    b2_secrets_toml: Path | None,
) -> bool:
    """Return True if B2 upload + HF secrets + verify succeeded."""
    env = {
        **os.environ,
        "HF_SPACE_REPO": space_repo,
        "KINK_PARSER_SECRETS_TOML": str(secrets_path),
    }
    if b2_secrets_toml is not None:
        env["KINK_B2_SECRETS_TOML"] = str(b2_secrets_toml.resolve())
    cmd = [
        sys.executable,
        str(ROOT / "scripts" / "sync_hf_space_b2_catalog.py"),
        "--secrets-toml",
        str(secrets_path),
        "--local-path",
        str(local_db),
        "--key",
        b2_object_key,
    ]
    if b2_secrets_toml is not None:
        cmd.extend(["--b2-secrets-toml", str(b2_secrets_toml.resolve())])
    print(f"e2e_deploy: {' '.join(cmd)}", file=sys.stderr, flush=True)
    r = subprocess.run(cmd, cwd=str(ROOT), env=env)
    if r.returncode != 0:
        return False
    verify_args = [
        "verify_b2_catalog_http.py",
        "--secrets-toml",
        str(secrets_path),
        "--key",
        b2_object_key,
    ]
    if b2_secrets_toml is not None:
        verify_args.extend(["--b2-secrets-toml", str(b2_secrets_toml.resolve())])
    _run(verify_args)
    return True


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--secrets-toml", type=Path, default=DEFAULT_SECRETS)
    ap.add_argument(
        "--b2-secrets-toml",
        type=Path,
        default=None,
        help="TOML with only [b2] (also set KINK_B2_SECRETS_TOML); parser --secrets-toml still used for HF token",
    )
    ap.add_argument("--space-repo", default=os.environ.get("HF_SPACE_REPO", "").strip(), help="HF Space repo id")
    ap.add_argument("--source", type=Path, default=ROOT / "data" / "store_slim.db")
    ap.add_argument("--shrink-dest", type=Path, default=ROOT / "data" / "store_slim_b2.db")
    ap.add_argument("--skip-shrink", action="store_true")
    ap.add_argument(
        "--create-bucket",
        metavar="NAME",
        default="",
        help="If set, run b2_create_bucket.py --name NAME first (uses [b2] or B2_ADMIN_* env with --admin-bucket-create)",
    )
    ap.add_argument(
        "--admin-bucket-create",
        action="store_true",
        help="With --create-bucket, pass --admin to b2_create_bucket (needs B2_ADMIN_KEY_ID / B2_ADMIN_APPLICATION_KEY)",
    )
    ap.add_argument(
        "--bootstrap",
        choices=("auto", "b2", "hub"),
        default="auto",
        help="auto: try B2 then Hub on failure; b2: B2 only; hub: Hub only",
    )
    ap.add_argument("--b2-object-key", default="kink/catalog/store_slim_bootstrap.db")
    ap.add_argument("--hub-repo", default="ronheichman/kink-catalog-slim")
    ap.add_argument("--hub-filename", default="store_slim_bootstrap.db")
    ap.add_argument("--no-publish", action="store_true", help="Skip publish_hf_space.py (still shrink + catalog)")
    ap.add_argument(
        "--verify-min-kinks",
        type=int,
        default=90_000,
        help="HF_VERIFY_MIN_KINKS for verify_hf_stack (default 90000)",
    )
    args = ap.parse_args()

    space_repo = (args.space_repo or "").strip()
    if not space_repo:
        print("e2e_deploy: set --space-repo or HF_SPACE_REPO", file=sys.stderr)
        return 1

    secrets_path = args.secrets_toml.resolve()
    if not secrets_path.is_file():
        print(f"e2e_deploy: secrets not found: {secrets_path}", file=sys.stderr)
        return 1
    if args.b2_secrets_toml is not None and not args.b2_secrets_toml.is_file():
        print(f"e2e_deploy: --b2-secrets-toml not found: {args.b2_secrets_toml}", file=sys.stderr)
        return 1

    hf_token = _load_hf_token(secrets_path)
    os.environ["HF_TOKEN"] = hf_token
    os.environ["HF_SPACE_REPO"] = space_repo

    if args.create_bucket:
        bc = [
            "b2_create_bucket.py",
            "--secrets-toml",
            str(secrets_path),
            "--name",
            args.create_bucket,
        ]
        if args.admin_bucket_create:
            bc.append("--admin")
        if args.b2_secrets_toml is not None:
            bc.extend(["--b2-secrets-toml", str(args.b2_secrets_toml.resolve())])
        _run(bc)

    if not args.skip_shrink:
        src = args.source.resolve()
        dst = args.shrink_dest.resolve()
        if not src.is_file():
            print(f"e2e_deploy: missing source DB {src}", file=sys.stderr)
            return 1
        _run(
            [
                "shrink_store_for_remote_bootstrap.py",
                "--source",
                str(src),
                "--dest",
                str(dst),
                "--force",
            ]
        )

    local_db = args.shrink_dest.resolve()
    if not local_db.is_file():
        print(f"e2e_deploy: missing catalog file {local_db}", file=sys.stderr)
        return 1

    used_b2 = False
    if args.bootstrap in ("auto", "b2"):
        ok = _try_b2_chain(secrets_path, space_repo, local_db, args.b2_object_key, args.b2_secrets_toml)
        if ok:
            used_b2 = True
            print("e2e_deploy: B2 path OK.", file=sys.stderr)
        elif args.bootstrap == "b2":
            print("e2e_deploy: B2 path failed and --bootstrap b2 (no Hub fallback).", file=sys.stderr)
            return 1
        else:
            print("e2e_deploy: B2 failed — Hub fallback.", file=sys.stderr)
            _hub_fallback(
                secrets_path,
                space_repo,
                local_db,
                args.hub_repo,
                args.hub_filename,
                hf_token,
            )
    else:
        _hub_fallback(
            secrets_path,
            space_repo,
            local_db,
            args.hub_repo,
            args.hub_filename,
            hf_token,
        )

    if not used_b2:
        _run(
            [
                "prove_hub_bootstrap_local.py",
                "--secrets-toml",
                str(secrets_path),
                "--repo-id",
                args.hub_repo,
                "--filename",
                args.hub_filename,
            ]
        )

    if args.no_publish:
        print("e2e_deploy: --no-publish — done (no Space code push).", file=sys.stderr)
        return 0

    env_verify = {
        **os.environ,
        "HF_TOKEN": hf_token,
        "HF_SPACE_REPO": space_repo,
        "HF_VERIFY_MIN_KINKS": str(args.verify_min_kinks),
    }
    _run(["publish_hf_space.py", "--verify"], env=env_verify)

    print("e2e_deploy: complete.", file=sys.stderr)
    return 0


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