scanner/persistence.py

"""Optional persistence: round-trip the OHLCV cache to a private HF Dataset.

The HF Space container's ``/tmp`` is wiped on every restart.  If the user
sets ``HF_TOKEN`` (or ``HUGGING_FACE_HUB_TOKEN``) and ``FSCANNER_CACHE_REPO``,
we push the parquet cache to a private dataset repo and pull it back on
startup.  This makes cold-start scans near-instant for a previously-scanned
universe.
"""

from __future__ import annotations

import os
from typing import Optional

import pandas as pd

from . import paths

CACHE_REPO = os.environ.get("FSCANNER_CACHE_REPO", "").strip()
CACHE_FILE = "ohlcv_cache.parquet"


def _is_enabled() -> bool:
    return bool(CACHE_REPO) and bool(
        os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    )


def pull_remote_cache() -> int:
    """Try to download the remote cache and write to CACHE_PATH.  Returns the
    number of tickers loaded, 0 on any failure or if not enabled.

    Pulls whenever remote persistence is configured; the remote copy
    overrides whatever (possibly stale or missing) file exists locally.
    """
    if not _is_enabled():
        return 0
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(
            repo_id=CACHE_REPO,
            repo_type="dataset",
            filename=CACHE_FILE,
            force_download=False,
        )
        # Replace the local cache with the downloaded one
        df = pd.read_parquet(path)
        os.makedirs(os.path.dirname(paths.CACHE_PATH) or ".", exist_ok=True)
        df.to_parquet(paths.CACHE_PATH, index=False)
        tickers = df["Ticker"].nunique() if "Ticker" in df.columns else 0
        return int(tickers)
    except Exception:
        return 0


def push_remote_cache() -> bool:
    """Upload the local cache to the remote dataset repo.  Returns True on
    success.
    """
    if not _is_enabled() or not os.path.exists(paths.CACHE_PATH):
        return False
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(
            path_or_fileobj=paths.CACHE_PATH,
            path_in_repo=CACHE_FILE,
            repo_id=CACHE_REPO,
            repo_type="dataset",
            commit_message="Update OHLCV cache",
        )
        return True
    except Exception:
        return False