scripts/bootstrap_data.py

#!/usr/bin/env python3
"""Restore the checksummed seed without replacing any existing runtime file."""

from __future__ import annotations
import argparse
import hashlib
import json
import tarfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


def restore(root: Path, archive: Path, manifest: Path) -> tuple[int, int]:
    metadata = json.loads(manifest.read_text())
    digest = hashlib.sha256(archive.read_bytes()).hexdigest()
    if digest != metadata["sha256"]:
        raise ValueError(
            "Bootstrap checksum mismatch; download a clean copy of this branch."
        )
    created = skipped = 0
    root = root.resolve()
    with tarfile.open(archive, "r:gz") as bundle:
        members = bundle.getmembers()
        # Validate the whole archive before touching the destination.
        for member in members:
            target = (root / member.name).resolve()
            if not member.isfile() or not target.is_relative_to(root):
                raise ValueError(f"Unsafe bootstrap member: {member.name}")
        for member in members:
            target = root / member.name
            if target.exists():
                skipped += 1
                continue
            target.parent.mkdir(parents=True, exist_ok=True)
            with bundle.extractfile(member) as source, target.open("xb") as dest:
                dest.write(source.read())
            created += 1
    return created, skipped


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=ROOT)
    args = parser.parse_args()
    created, skipped = restore(
        args.root, ROOT / "bootstrap/seed.tar.gz", ROOT / "bootstrap/manifest.json"
    )
    print(
        f"Bootstrap verified: restored {created} files; preserved {skipped} existing files."
    )
    print(
        "Historical pending forecasts are not in this seed. New forecasts start from current observations."
    )


if __name__ == "__main__":
    main()