core/hubio.py

"""Low-level JSONL read/write against the HF Dataset repo.

Shared by core.auth (users registry) and core.storage (per-user transactions).
The Space disk is ephemeral, so everything persists to the dataset. Auth uses
the HF_TOKEN environment variable.

Note: Hugging Face has no real "append" API — a write uploads the whole file.
We isolate writers by giving each user their own file (data/<user>.jsonl), so
one user's save never clobbers another's. Within a single user, saves are
sequential. The shared users.jsonl can in theory race on simultaneous signups;
acceptable for this app, and we always re-read the latest before writing.
"""

from __future__ import annotations

import json
import os
from typing import Any

DATASET_REPO = "build-small-hackathon/budget-tracker-data"


def token() -> str | None:
    return os.environ.get("HF_TOKEN")


def ensure_dataset() -> None:
    """Create the (private) dataset on first use; no-op if it exists."""
    from huggingface_hub import HfApi

    HfApi(token=token()).create_repo(
        repo_id=DATASET_REPO, repo_type="dataset", private=True, exist_ok=True
    )


def read_jsonl(path_in_repo: str, force: bool = False) -> list[dict[str, Any]]:
    """Return parsed JSONL rows from a dataset file ([] if missing/new)."""
    from huggingface_hub import hf_hub_download
    from huggingface_hub.errors import EntryNotFoundError, RepositoryNotFoundError

    try:
        path = hf_hub_download(
            repo_id=DATASET_REPO,
            filename=path_in_repo,
            repo_type="dataset",
            token=token(),
            force_download=force,
        )
    except (EntryNotFoundError, RepositoryNotFoundError):
        return []
    except Exception as e:  # pragma: no cover - network/permission edge cases
        print(f"[hubio] read_jsonl({path_in_repo}) failed: {e}")
        return []

    rows: list[dict[str, Any]] = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return rows


def write_jsonl(path_in_repo: str, rows: list[dict[str, Any]], message: str) -> None:
    """Overwrite a dataset file with the given rows (creates repo if needed)."""
    from huggingface_hub import HfApi

    ensure_dataset()
    content = "\n".join(json.dumps(r, ensure_ascii=False) for r in rows)
    if content:
        content += "\n"
    HfApi(token=token()).upload_file(
        path_or_fileobj=content.encode("utf-8"),
        path_in_repo=path_in_repo,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=token(),
        commit_message=message,
    )


def append_jsonl(path_in_repo: str, obj: dict[str, Any], message: str) -> list[dict[str, Any]]:
    """Append one row to a dataset file (read latest → add → upload). Returns all rows."""
    rows = read_jsonl(path_in_repo, force=True)
    rows.append(obj)
    write_jsonl(path_in_repo, rows, message)
    return rows