src/densefeed/storage/sync.py

"""HF Dataset sync — download/upload DuckDB file for persistence across Space restarts."""

from __future__ import annotations

import logging
import shutil
from pathlib import Path

logger = logging.getLogger(__name__)

DUCKDB_HF_PATH = "densefeed.duckdb"


def download_db(repo_id: str, local_path: str | Path) -> bool:
    """Pull latest DuckDB file from HF Dataset. Returns True if downloaded, False if fresh.

    On corruption: deletes local file and re-creates fresh.
    """
    from densefeed.auth import hf_token

    token = hf_token.get()
    if not token or not repo_id:
        logger.info("[sync] no token or repo, starting with fresh DB")
        return False

    local_path = Path(local_path)

    try:
        from huggingface_hub import hf_hub_download

        downloaded = hf_hub_download(
            repo_id=repo_id,
            filename=DUCKDB_HF_PATH,
            repo_type="dataset",
            token=token,
            local_dir=str(local_path.parent),
            force_download=True,
        )
        # hf_hub_download preserves filename, rename if needed
        downloaded_path = Path(downloaded)
        if downloaded_path != local_path:
            shutil.move(str(downloaded_path), str(local_path))

        logger.info("[sync] downloaded DB from %s", repo_id)
        return True

    except Exception as e:
        logger.warning("[sync] download failed (starting fresh): %s", e)
        # Do not delete an existing local DuckDB here. A missing remote file is
        # expected for first-run user datasets, and deleting the local DB can
        # unlink an already-open DuckDB connection before its initial upload.
        return False


def upload_db(repo_id: str, local_path: str | Path) -> bool:
    """Push DuckDB file back to HF Dataset after a run completes."""
    from densefeed.auth import hf_token

    token = hf_token.get()
    if not token or not repo_id:
        logger.warning("[sync] no token or repo, skipping upload")
        return False

    local_path = Path(local_path)
    if local_path.exists() and not local_path.is_file():
        logger.warning("[sync] DB path exists but is not a file, replacing: %s", local_path)
        shutil.rmtree(local_path, ignore_errors=True)
    if not local_path.exists():
        logger.info("[sync] DB file missing at %s; creating empty initialized DB", local_path)
        try:
            import duckdb
            from densefeed.schema import init_schema

            local_path.parent.mkdir(parents=True, exist_ok=True)
            conn = duckdb.connect(str(local_path))
            init_schema(conn)
            conn.execute("CHECKPOINT")
            conn.close()
            if not local_path.is_file():
                raise FileNotFoundError(f"DuckDB did not create file: {local_path}")
        except Exception as e:
            logger.warning("[sync] failed to create initial DB at %s: %s", local_path, e)
            return False

    try:
        from huggingface_hub import HfApi

        api = HfApi(token=token)
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=DUCKDB_HF_PATH,
            repo_id=repo_id,
            repo_type="dataset",
        )
        logger.info("[sync] uploaded DB to %s", repo_id)
        return True

    except Exception as e:
        logger.warning("[sync] upload failed: %s", e)
        return False


def sync_after_stage(
    repo_id: str, local_path: str | Path, run_id: str, stage: int
) -> bool:
    """Best-effort DuckDB upload after a successful stage."""
    logger.info("[sync] syncing run %s after stage %d", run_id, stage)
    try:
        return upload_db(repo_id, local_path)
    except Exception as e:
        logger.warning(
            "[sync] stage sync failed for run %s stage %d: %s", run_id, stage, e
        )
        return False