Add governed Statbel release promotion

scripts/provision_mol_population_history.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import argparse
 import csv
+from dataclasses import dataclass
 from hashlib import sha256
 import io
 import json
@@ -36,7 +37,11 @@ from urllib3.util.retry import Retry
 
 from geographic_scopes import GEOGRAPHIC_SCOPES, GeographicScope
 from statbel_population_preflight import (
+    MAX_GEOMETRY_ARCHIVE_BYTES,
+    MAX_POPULATION_ARCHIVE_BYTES,
     StatbelPreflightError,
+    validate_geometry_source_url,
+    validate_population_source_url,
     validate_statbel_release,
     write_manifest,
 )
@@ -67,6 +72,45 @@ POPULATION_URLS = {
 POPULATION_LAYOUTS = {year: ("new" if year == 2025 else "standard") for year in POPULATION_URLS}
 
 
+@dataclass(frozen=True)
+class PopulationReleaseConfig:
+    year: int
+    layout: str
+    population_url: str
+    geometry_url: str
+
+
+def resolve_release_config(
+    year: int,
+    *,
+    population_url: str | None = None,
+    geometry_url: str | None = None,
+    layout: str | None = None,
+) -> PopulationReleaseConfig:
+    overrides = (population_url, geometry_url, layout)
+    if any(value is not None for value in overrides):
+        if not all(value is not None for value in overrides):
+            raise ValueError("Population URL, geometry URL and population layout must be supplied together")
+        release = PopulationReleaseConfig(
+            year=year,
+            layout=str(layout),
+            population_url=str(population_url),
+            geometry_url=str(geometry_url),
+        )
+    elif year in POPULATION_URLS:
+        release = PopulationReleaseConfig(
+            year=year,
+            layout=POPULATION_LAYOUTS[year],
+            population_url=POPULATION_URLS[year],
+            geometry_url=SECTOR_URL.format(year=year),
+        )
+    else:
+        raise ValueError(f"Unsupported population year without an explicit governed release: {year}")
+    validate_population_source_url(release.population_url, release.year, release.layout)
+    validate_geometry_source_url(release.geometry_url, release.year)
+    return release
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Provision official annual Statbel population snapshots.")
     parser.add_argument("--scope", choices=sorted(GEOGRAPHIC_SCOPES), default=DEFAULT_SCOPE_KEY)
@@ -74,6 +118,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--project-name", default=None)
     parser.add_argument("--area-name", default=None, help="Case-insensitive fragment identifying the persisted Area.")
     parser.add_argument("--years", default="2021,2022,2023,2024,2025")
+    parser.add_argument(
+        "--population-url",
+        help="Exact official Statbel ZIP URL for one explicitly governed release; requires the other release flags",
+    )
+    parser.add_argument(
+        "--geometry-url",
+        help="Exact matching official statistical-sector GeoJSON ZIP URL for one explicitly governed release",
+    )
+    parser.add_argument("--population-layout", choices=("standard", "new"))
     parser.add_argument("--output-dir", type=Path, default=None)
     parser.add_argument("--boundary-path", type=Path, default=None)
     parser.add_argument(
@@ -107,6 +160,43 @@ def build_session() -> requests.Session:
     return session
 
 
+def download_archive(
+    session: requests.Session,
+    *,
+    url: str,
+    year: int,
+    layout: str | None,
+    max_bytes: int,
+    timeout: int,
+) -> bytes:
+    with session.get(url, timeout=timeout, stream=True) as response:
+        response.raise_for_status()
+        if layout is None:
+            validate_geometry_source_url(response.url, year)
+        else:
+            validate_population_source_url(response.url, year, layout)
+        content_length = response.headers.get("Content-Length")
+        if content_length:
+            try:
+                advertised_size = int(content_length)
+            except ValueError as exc:
+                raise RuntimeError(f"Official source returned an invalid Content-Length for {url}") from exc
+            if advertised_size <= 0 or advertised_size > max_bytes:
+                raise RuntimeError(f"Official source archive exceeds the {max_bytes}-byte download limit")
+        chunks: list[bytes] = []
+        size = 0
+        for chunk in response.iter_content(chunk_size=1024 * 1024):
+            if not chunk:
+                continue
+            size += len(chunk)
+            if size > max_bytes:
+                raise RuntimeError(f"Official source archive exceeds the {max_bytes}-byte download limit")
+            chunks.append(chunk)
+    if size <= 0:
+        raise RuntimeError(f"Official source archive is empty: {url}")
+    return b"".join(chunks)
+
+
 def response_data(response: requests.Response) -> Any:
     try:
         payload = response.json()
@@ -119,6 +209,29 @@ def response_data(response: requests.Response) -> Any:
     return payload["data"]
 
 
+def list_paginated_items(session: requests.Session, url: str, *, timeout: int) -> list[dict[str, Any]]:
+    items: list[dict[str, Any]] = []
+    offset = 0
+    total: int | None = None
+    while total is None or offset < total:
+        page = response_data(session.get(url, params={"limit": 200, "offset": offset}, timeout=timeout))
+        page_items = page.get("items") if isinstance(page, dict) else None
+        if not isinstance(page_items, list):
+            raise RuntimeError(f"GeoIntel list response for {url} has no items array")
+        page_total = int(page.get("total", len(page_items)))
+        if total is None:
+            total = page_total
+        elif page_total != total:
+            raise RuntimeError("GeoIntel pagination total changed while reading the population workspace")
+        items.extend(page_items)
+        if not page_items:
+            break
+        offset += len(page_items)
+    if total is not None and len(items) != total:
+        raise RuntimeError(f"GeoIntel list response for {url} returned {len(items)} of {total} items")
+    return items
+
+
 def series_key(scope: GeographicScope) -> str:
     return f"statbel:population-statistical-sector:{scope.key}"
 
@@ -340,9 +453,13 @@ def preflight_manifest_path(output_dir: Path, scope: GeographicScope, year: int)
 
 
 def previous_snapshot_path(output_dir: Path, scope: GeographicScope, year: int) -> Path | None:
-    candidates = [snapshot_path(output_dir, scope, candidate) for candidate in POPULATION_URLS if candidate < year]
-    available = [path for path in candidates if path.is_file()]
-    return max(available, key=lambda path: int(path.stem.rsplit("_", 1)[-1])) if available else None
+    prefix = f"{scope.key.replace('-', '_')}_statbel_population_"
+    available: list[tuple[int, Path]] = []
+    for path in output_dir.glob(f"{prefix}*.geojson"):
+        suffix = path.stem.removeprefix(prefix)
+        if suffix.isdigit() and int(suffix) < year and path.is_file():
+            available.append((int(suffix), path))
+    return max(available, key=lambda item: item[0])[1] if available else None
 
 
 def load_preflight_manifest(path: Path, snapshot: Path, year: int, scope: GeographicScope) -> dict[str, Any]:
@@ -383,23 +500,24 @@ def stage_release(
     output_dir: Path,
     boundary,
     scope: GeographicScope,
+    release: PopulationReleaseConfig | None = None,
 ) -> tuple[Path, Path, dict[str, Any]]:
-    layout = POPULATION_LAYOUTS[year]
-    population_url = POPULATION_URLS[year]
-    geometry_url = SECTOR_URL.format(year=year)
+    active_release = release or resolve_release_config(year)
+    if active_release.year != year:
+        raise ValueError("Population release year does not match the staged year")
     result = validate_statbel_release(
         year=year,
-        layout=layout,
+        layout=active_release.layout,
         population_content=population_content,
-        population_url=population_url,
+        population_url=active_release.population_url,
         geometry_content=geometry_content,
-        geometry_url=geometry_url,
+        geometry_url=active_release.geometry_url,
         scope=scope,
         baseline_snapshot=previous_snapshot_path(output_dir, scope, year),
     )
     raw_dir = output_dir / "raw" / str(year)
-    population_archive_path = raw_dir / Path(unquote(urlsplit(population_url).path)).name
-    geometry_archive_path = raw_dir / Path(unquote(urlsplit(geometry_url).path)).name
+    population_archive_path = raw_dir / Path(unquote(urlsplit(active_release.population_url).path)).name
+    geometry_archive_path = raw_dir / Path(unquote(urlsplit(active_release.geometry_url).path)).name
     write_bytes_atomic(population_archive_path, population_content)
     write_bytes_atomic(geometry_archive_path, geometry_content)
 
@@ -444,18 +562,18 @@ def locate_workspace(
     area_name: str,
     timeout: int,
 ):
-    projects = response_data(session.get(f"{base_url}/api/v1/projects", params={"limit": 200}, timeout=timeout))
-    project = next((item for item in projects.get("items") or [] if item.get("name") == project_name), None)
+    projects = list_paginated_items(session, f"{base_url}/api/v1/projects", timeout=timeout)
+    project = next((item for item in projects if item.get("name") == project_name), None)
     if not project:
         raise RuntimeError(f"Project {project_name!r} is missing")
     project_id = str(project["id"])
-    areas = response_data(session.get(f"{base_url}/api/v1/projects/{project_id}/areas", params={"limit": 200}, timeout=timeout))
+    areas = list_paginated_items(session, f"{base_url}/api/v1/projects/{project_id}/areas", timeout=timeout)
     area_fragment = area_name.strip().casefold()
-    matches = [item for item in areas.get("items") or [] if area_fragment in str(item.get("name") or "").casefold()]
+    matches = [item for item in areas if area_fragment in str(item.get("name") or "").casefold()]
     if len(matches) != 1:
         raise RuntimeError(f"Expected one official Area matching {area_name!r}, received {len(matches)}")
-    datasets = response_data(session.get(f"{base_url}/api/v1/projects/{project_id}/datasets", params={"limit": 200}, timeout=timeout))
-    return project_id, str(matches[0]["id"]), list(datasets.get("items") or [])
+    datasets = list_paginated_items(session, f"{base_url}/api/v1/projects/{project_id}/datasets", timeout=timeout)
+    return project_id, str(matches[0]["id"]), datasets
 
 
 def upload_snapshot(
@@ -468,7 +586,11 @@ def upload_snapshot(
     timeout: int,
     scope: GeographicScope,
     preflight_path: Path,
+    release: PopulationReleaseConfig | None = None,
 ) -> dict[str, Any]:
+    active_release = release or resolve_release_config(year)
+    if active_release.year != year:
+        raise ValueError("Population release year does not match the upload year")
     observed_at = f"{year}-01-01T00:00:00Z"
     preflight = load_preflight_manifest(preflight_path, path, year, scope)
     accounting = preflight["scope_accounting"]
@@ -506,8 +628,8 @@ def upload_snapshot(
         "operator_explicit_fetch": True,
         "scope_key": scope.key,
         "geometry_clipped_to_area": True,
-        "sector_geometry_url": SECTOR_URL.format(year=year),
-        "population_url": POPULATION_URLS[year],
+        "sector_geometry_url": active_release.geometry_url,
+        "population_url": active_release.population_url,
         "population_layout": preflight["release"]["population_layout"],
         "preflight_manifest_path": str(preflight_path),
         "preflight_manifest_sha256": sha256_path(preflight_path),
@@ -551,9 +673,23 @@ def main() -> int:
     except ValueError:
         print(json.dumps({"status": "error", "message": "Years must be comma-separated integers"}), file=sys.stderr)
         return 2
-    unsupported = [year for year in years if year not in POPULATION_URLS]
-    if unsupported or not years:
-        print(json.dumps({"status": "error", "message": f"Unsupported years: {unsupported}"}), file=sys.stderr)
+    custom_release_requested = any((args.population_url, args.geometry_url, args.population_layout))
+    if not years or (custom_release_requested and len(years) != 1):
+        message = "An explicit governed release requires exactly one population year" if custom_release_requested else "No years requested"
+        print(json.dumps({"status": "error", "message": message}), file=sys.stderr)
+        return 2
+    try:
+        releases = {
+            year: resolve_release_config(
+                year,
+                population_url=args.population_url if custom_release_requested else None,
+                geometry_url=args.geometry_url if custom_release_requested else None,
+                layout=args.population_layout if custom_release_requested else None,
+            )
+            for year in years
+        }
+    except (ValueError, StatbelPreflightError) as exc:
+        print(json.dumps({"status": "error", "message": str(exc)}, ensure_ascii=False), file=sys.stderr)
         return 2
 
     output_dir = resolve_output_dir(args, scope)
@@ -565,21 +701,35 @@ def main() -> int:
         prepared: list[dict[str, Any]] = []
         with build_session() as source_session:
             for year in years:
+                release = releases[year]
                 path = snapshot_path(output_dir, scope, year)
                 manifest_path = preflight_manifest_path(output_dir, scope, year)
                 preflight_status = "passed"
                 if args.force or not path.exists():
-                    sectors_response = source_session.get(SECTOR_URL.format(year=year), timeout=args.request_timeout)
-                    sectors_response.raise_for_status()
-                    population_response = source_session.get(POPULATION_URLS[year], timeout=args.request_timeout)
-                    population_response.raise_for_status()
+                    geometry_content = download_archive(
+                        source_session,
+                        url=release.geometry_url,
+                        year=year,
+                        layout=None,
+                        max_bytes=MAX_GEOMETRY_ARCHIVE_BYTES,
+                        timeout=args.request_timeout,
+                    )
+                    population_content = download_archive(
+                        source_session,
+                        url=release.population_url,
+                        year=year,
+                        layout=release.layout,
+                        max_bytes=MAX_POPULATION_ARCHIVE_BYTES,
+                        timeout=args.request_timeout,
+                    )
                     path, manifest_path, _manifest = stage_release(
                         year=year,
-                        population_content=population_response.content,
-                        geometry_content=sectors_response.content,
+                        population_content=population_content,
+                        geometry_content=geometry_content,
                         output_dir=output_dir,
                         boundary=boundary,
                         scope=scope,
+                        release=release,
                     )
                 elif manifest_path.is_file():
                     load_preflight_manifest(manifest_path, path, year, scope)
@@ -648,6 +798,7 @@ def main() -> int:
                         args.import_timeout,
                         scope,
                         item["manifest_path"],
+                        releases[year],
                     )
                     results.append({"year": year, "dataset_id": dataset["id"], "feature_count": dataset.get("feature_count"), "status": "imported"})
     except (OSError, RuntimeError, requests.RequestException, ValueError, KeyError, zipfile.BadZipFile) as exc: