SPlit up

pipeline/utils/download.py

@@ -1,11 +1,19 @@
 """Shared download and extraction helpers for pipeline scripts."""
 
+import time
 import zipfile
+from io import BytesIO
 from pathlib import Path
 
 import httpx
+import polars as pl
 from tqdm import tqdm
 
+# Census 2021 LSOAs (TYPE151) with an E prefix. The Census 2021 geography is
+# frozen, so NOMIS England-level downloads must yield exactly this many LSOAs;
+# fewer means the download was truncated.
+ENGLAND_LSOA_COUNT_2021 = 33_755
+
 
 def download(url: str, output_path: Path, *, timeout: float = 120) -> None:
     """Stream-download a URL to a local file with a tqdm progress bar."""
@@ -38,3 +46,86 @@ def extract_zip(zip_path: Path, extract_dir: Path) -> None:
     extract_dir.mkdir(parents=True, exist_ok=True)
     with zipfile.ZipFile(zip_path, "r") as zf:
         zf.extractall(extract_dir)
+
+
+def download_nomis_csv(base_url: str, *, page_size: int = 25_000) -> pl.DataFrame:
+    """Download a NOMIS CSV dataset, paging with recordoffset/RecordLimit.
+
+    The page size is sent explicitly as ``RecordLimit``: last-page detection is
+    ``rows < page_size``, so relying on NOMIS's implicit default would silently
+    truncate the dataset to one page if that default ever differed.
+    """
+    frames = []
+    offset = 0
+    while True:
+        url = f"{base_url}&RecordLimit={page_size}&recordoffset={offset}"
+        response = httpx.get(url, follow_redirects=True, timeout=120)
+        response.raise_for_status()  # pyright: ignore[reportUnusedCallResult]
+        if len(response.content) == 0:
+            break
+        chunk = pl.read_csv(BytesIO(response.content))
+        if chunk.height == 0:
+            break
+        frames.append(chunk)
+        print(f"  Fetched {chunk.height} rows (offset={offset})")
+        if chunk.height < page_size:
+            break
+        offset += page_size
+
+    if not frames:
+        raise RuntimeError(f"NOMIS returned no rows for {base_url}")
+    return pl.concat(frames)
+
+
+def download_arcgis_hub_export(
+    url: str,
+    output_path: Path,
+    *,
+    expected_geometry: str | None = None,
+    poll_interval_s: float = 5,
+    poll_timeout_s: float = 600,
+) -> int:
+    """Download an ArcGIS Hub `api/download/v1` export, handling deferred jobs.
+
+    The endpoint returns HTTP 202 with a JSON status body while the export is
+    still being prepared; a plain download would save that placeholder as the
+    output file with a success exit code. Poll until the file is ready, then
+    validate the result with pyogrio (feature count > 0 and, optionally, a
+    geometry-type substring) before moving it into place. Returns the feature
+    count.
+    """
+    import pyogrio
+
+    tmp_path = output_path.with_name(f"{output_path.stem}.tmp{output_path.suffix}")
+    deadline = time.monotonic() + poll_timeout_s
+    with httpx.Client(follow_redirects=True, timeout=300) as client:
+        while True:
+            with client.stream("GET", url) as response:
+                if response.status_code == 202:
+                    response.read()
+                    if time.monotonic() > deadline:
+                        raise TimeoutError(
+                            f"Export did not finish within {poll_timeout_s}s: "
+                            f"{response.text}"
+                        )
+                    time.sleep(poll_interval_s)
+                    continue
+                response.raise_for_status()  # pyright: ignore[reportUnusedCallResult]
+                with tmp_path.open("wb") as fh:
+                    for chunk in response.iter_bytes():
+                        fh.write(chunk)
+                break
+
+    info = pyogrio.read_info(tmp_path)
+    features = int(info.get("features", 0))
+    geometry_type = str(info.get("geometry_type") or "")
+    if features <= 0:
+        raise ValueError(f"Downloaded file {output_path.name} contains no features")
+    if expected_geometry is not None and expected_geometry not in geometry_type:
+        raise ValueError(
+            f"Expected {expected_geometry!r} geometry in {output_path.name}, "
+            f"got {geometry_type!r}"
+        )
+
+    tmp_path.replace(output_path)
+    return features