backend/hf_bootstrap.py

"""Bootstrap the catalog SQLite before the app opens it.

Priority when the store file is missing:

1. ``KINK_CATALOG_URL`` — plain HTTPS GET (stream). Use for **public** URLs.
2. **Backblaze B2 (S3 API)** — if ``B2_KEY_ID``, ``B2_APPLICATION_KEY``, ``B2_BUCKET``, ``B2_REGION`` are set
   (same names as ``parser`` ``hf_deploy.py`` Space secrets), download ``KINK_B2_OBJECT_KEY`` (default
   ``kink/catalog/store_slim.db``). Works with **private** buckets; requires ``boto3``.
3. ``KINK_HF_DATASET_REPO`` — Hub **dataset** download (only when ``KINK_HF_REQUIRE_FULL_CATALOG`` is on, or when
   ``KINK_HF_USE_HUB_DATASET=1`` while full catalog is off — see below).
4. Bundled seed when full catalog is off and Hub dataset is not explicitly requested.

On Spaces: mount persistent storage and set ``KINK_STORE_PATH`` so the file survives restarts.
"""
from __future__ import annotations

import os
import shutil
import tarfile
from pathlib import Path

# Default dataset when KINK_HF_REQUIRE_FULL_CATALOG=1 but Space UI cleared KINK_HF_DATASET_REPO (empty overrides Dockerfile).
DEFAULT_KINK_HF_DATASET_REPO = "ronheichman/kink-catalog-slim"
DEFAULT_KINK_HF_REMOTE_DB_NAME = "store_slim.db"


def _env_truthy(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on")


def _b2_s3_env() -> tuple[str, str, str, str, str] | None:
    """Return (key_id, app_key, bucket, region, object_key) if B2 S3 download is configured."""
    key_id = (os.environ.get("B2_KEY_ID") or os.environ.get("KINK_B2_KEY_ID") or "").strip()
    app_key = (os.environ.get("B2_APPLICATION_KEY") or os.environ.get("KINK_B2_APPLICATION_KEY") or "").strip()
    bucket = (os.environ.get("B2_BUCKET") or os.environ.get("KINK_B2_BUCKET") or "").strip()
    region = (os.environ.get("B2_REGION") or os.environ.get("KINK_B2_REGION") or "").strip()
    object_key = (os.environ.get("KINK_B2_OBJECT_KEY") or "kink/catalog/store_slim.db").strip().lstrip("/")
    if key_id and app_key and bucket and region:
        return (key_id, app_key, bucket, region, object_key)
    return None


def _download_store_from_b2_s3(dest: Path) -> None:
    """Download catalog object from Backblaze via the S3-compatible API (private buckets OK)."""
    cfg = _b2_s3_env()
    if not cfg:
        raise RuntimeError("B2 S3 download requested but credentials are incomplete")
    key_id, app_key, bucket, region, object_key = cfg

    try:
        import boto3
    except ImportError as exc:
        raise ImportError("boto3 is required for B2 catalog download. Install boto3.") from exc

    endpoint = f"https://s3.{region}.backblazeb2.com"
    client = boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=key_id,
        aws_secret_access_key=app_key,
        region_name=region,
    )
    dest.parent.mkdir(parents=True, exist_ok=True)
    partial = dest.with_name(dest.name + ".partial")
    try:
        client.download_file(bucket, object_key, str(partial))
        partial.replace(dest)
    except BaseException:
        if partial.is_file():
            partial.unlink(missing_ok=True)
        raise


def _download_store_from_http(url: str, dest: Path) -> None:
    """Stream ``url`` to ``dest`` (atomic replace). Uses ``httpx`` (project dependency)."""
    import httpx

    dest.parent.mkdir(parents=True, exist_ok=True)
    partial = dest.with_name(dest.name + ".partial")
    headers: dict[str, str] = {}
    token = (os.environ.get("KINK_CATALOG_URL_BEARER") or "").strip()
    if token:
        headers["Authorization"] = f"Bearer {token}"

    timeout = httpx.Timeout(connect=30.0, read=None, write=30.0, pool=30.0)
    try:
        with httpx.Client(timeout=timeout, follow_redirects=True, headers=headers) as client:
            with client.stream("GET", url) as response:
                response.raise_for_status()
                with partial.open("wb") as out:
                    for chunk in response.iter_bytes(chunk_size=1024 * 1024):
                        if chunk:
                            out.write(chunk)
        partial.replace(dest)
    except BaseException:
        if partial.is_file():
            partial.unlink(missing_ok=True)
        raise


def ensure_media_cache(target_root: Path) -> Path:
    """Hydrate ``data/cached_assets`` from a Hub dataset archive when configured."""
    archive_name = os.environ.get("KINK_HF_MEDIA_ARCHIVE", "").strip()
    if not archive_name:
        return target_root

    target_root = target_root.resolve()
    in_tree_marker = target_root / f".{archive_name.replace('/', '_')}.ready"
    # /app/data is wiped + re-seeded by entrypoint on every container start, so the in-tree marker
    # never survives. Mirror it to /tmp (which survives restart_space) so we only pay the
    # 50-75s Hub download + extract once per /tmp lifetime instead of every restart.
    revision = os.environ.get("KINK_HF_MEDIA_DATASET_REVISION", "").strip() or "default"
    safe_archive = archive_name.replace('/', '_').replace(':', '_')
    persistent_marker = Path(f"/tmp/kink_media_cache.{safe_archive}.{revision}.ready")
    if in_tree_marker.is_file() or persistent_marker.is_file():
        print(f"[kink_cli media] hit in_tree={in_tree_marker.is_file()} persistent={persistent_marker.is_file()} archive={archive_name}", flush=True)
        return target_root
    print(f"[kink_cli media] miss in_tree={in_tree_marker} persistent={persistent_marker} archive={archive_name}", flush=True)

    repo = (
        os.environ.get("KINK_HF_MEDIA_DATASET_REPO", "").strip()
        or os.environ.get("KINK_HF_DATASET_REPO", "").strip()
        or DEFAULT_KINK_HF_DATASET_REPO
    )
    revision = os.environ.get("KINK_HF_MEDIA_DATASET_REVISION", "").strip() or None
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not repo:
        raise RuntimeError("KINK_HF_MEDIA_ARCHIVE is set but no Hub dataset repo is configured")

    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub is required when KINK_HF_MEDIA_ARCHIVE is set. "
            "Install with: pip install huggingface_hub"
        ) from exc

    target_root.mkdir(parents=True, exist_ok=True)
    downloaded = hf_hub_download(
        repo_id=repo,
        repo_type="dataset",
        filename=archive_name,
        revision=revision,
        token=token,
        local_dir=str(target_root.parent),
    )
    with tarfile.open(downloaded, "r:*") as tar:
        tar.extractall(target_root, filter="data")
    in_tree_marker.write_text("ok\n", encoding="utf-8")
    try:
        persistent_marker.parent.mkdir(parents=True, exist_ok=True)
        persistent_marker.write_text("ok\n", encoding="utf-8")
    except OSError:
        pass
    return target_root


def ensure_store_db(target: Path) -> Path:
    """Return ``target`` if it exists; otherwise download from Hub, or copy bundled seed when allowed."""
    target = target.resolve()
    require_full = _env_truthy("KINK_HF_REQUIRE_FULL_CATALOG")
    # Full slim DB is multi-GB; bundled seed is sub-megabyte. Ephemeral Spaces may keep an old seed file.
    stale_max = int(os.environ.get("KINK_HF_STALE_MAX_BYTES", str(100 * 1024 * 1024)))

    if target.is_file():
        if require_full and target.stat().st_size < stale_max:
            target.unlink()
        else:
            return target

    catalog_url = os.environ.get("KINK_CATALOG_URL", "").strip()
    if catalog_url:
        _download_store_from_http(catalog_url, target)
        if not target.is_file():
            raise RuntimeError(f"Expected database at {target} after HTTP download")
        return target

    if _b2_s3_env() is not None:
        _download_store_from_b2_s3(target)
        if not target.is_file():
            raise RuntimeError(f"Expected database at {target} after B2 S3 download")
        return target

    repo = os.environ.get("KINK_HF_DATASET_REPO", "").strip()
    if not require_full and repo and not _env_truthy("KINK_HF_USE_HUB_DATASET"):
        # HF Space "Variables" often keep KINK_HF_DATASET_REPO after switching the image to bundled-only.
        # That triggered anonymous hf_hub_download (slow / OOM) and ignored the bundled seed.
        repo = ""

    if require_full and not repo:
        repo = (
            os.environ.get("KINK_HF_DATASET_REPO_FALLBACK", DEFAULT_KINK_HF_DATASET_REPO).strip()
            or DEFAULT_KINK_HF_DATASET_REPO
        )

    if not repo:
        seed = Path(os.environ.get("KINK_SEED_PATH", "/app/deploy/hf/seed/hf_bundled_store.db")).resolve()
        if seed.is_file():
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(seed, target)
            return target
        raise FileNotFoundError(
            f"SQLite database not found: {target}. "
            "Locally: run python scripts/build_slim_store.py. "
            "On Hugging Face: set KINK_CATALOG_URL (public HTTPS), or B2_KEY_ID/B2_APPLICATION_KEY/B2_BUCKET/B2_REGION "
            "(and KINK_B2_OBJECT_KEY), or KINK_HF_REQUIRE_FULL_CATALOG=1 plus KINK_HF_DATASET_REPO (Hub dataset), "
            "or KINK_HF_USE_HUB_DATASET=1 with KINK_HF_DATASET_REPO when full catalog is off, "
            "KINK_STORE_PATH on a mounted bucket if you need persistence, "
            "or upload an existing DB to the bucket first. "
            "For the Docker image, ensure deploy/hf/seed/hf_bundled_store.db is included in the build context."
        )

    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub is required when KINK_HF_DATASET_REPO is set. "
            "Install with: pip install huggingface_hub"
        ) from exc

    target.parent.mkdir(parents=True, exist_ok=True)
    explicit_fn = os.environ.get("KINK_HF_DATASET_FILENAME", "").strip()
    if explicit_fn:
        remote_name = explicit_fn
    elif require_full:
        # KINK_STORE_PATH is often kink_store.db while the Hub file is store_slim.db; Space may clear FILENAME too.
        remote_name = DEFAULT_KINK_HF_REMOTE_DB_NAME
    else:
        remote_name = (target.name or "").strip() or target.name
    revision = os.environ.get("KINK_HF_DATASET_REVISION", "").strip() or None
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    downloaded = hf_hub_download(
        repo_id=repo,
        repo_type="dataset",
        filename=remote_name,
        revision=revision,
        token=token,
        local_dir=str(target.parent),
    )
    got = Path(downloaded).resolve()
    if got != target:
        shutil.move(str(got), str(target))
    if not target.is_file():
        raise RuntimeError(f"Expected database at {target} after Hub download")
    return target