scripts/bootstrap_hf_data.py

"""Build the Space's read-only DuckDB from the published daily Parquet files."""
from __future__ import annotations

import argparse
import os
from pathlib import Path

import duckdb
from huggingface_hub import hf_hub_download


DEFAULT_REPO = "egoh33/ibkr-daily-stock-data"
FILES = ("polygon_bars.parquet", "silver_stock_features.parquet")


def materialize_database(files: dict[str, Path], db_path: Path) -> None:
    """Atomically create a DuckDB containing the two public daily-stock tables."""
    db_path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = db_path.with_suffix(db_path.suffix + ".tmp")
    temp_path.unlink(missing_ok=True)
    with duckdb.connect(str(temp_path)) as conn:
        for filename, source in files.items():
            table = Path(filename).stem
            safe_source = source.as_posix().replace("'", "''")
            conn.execute(f"CREATE TABLE {table} AS SELECT * FROM read_parquet('{safe_source}')")
        conn.execute("CHECKPOINT")
    db_path.unlink(missing_ok=True)
    temp_path.replace(db_path)


def download_files(repo_id: str, token: str | None = None) -> dict[str, Path]:
    return {
        filename: Path(
            hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)
        )
        for filename in FILES
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", default=os.getenv("HF_DATASET_REPO", DEFAULT_REPO))
    parser.add_argument("--db-path", default=os.getenv("DB_PATH", "./data/equity.duckdb"))
    parser.add_argument("--source-dir", help="Use local Parquet files instead of downloading")
    args = parser.parse_args()

    if args.source_dir:
        source_dir = Path(args.source_dir)
        files = {filename: source_dir / filename for filename in FILES}
    else:
        files = download_files(args.repo, token=os.getenv("HF_TOKEN"))
    missing = [str(path) for path in files.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"Missing required daily-stock files: {', '.join(missing)}")
    materialize_database(files, Path(args.db_path))
    print(f"Materialized {args.db_path} from {args.repo}")


if __name__ == "__main__":
    main()