Harden GCS dataset loading

Author: anth-volk

policyengine_uk/data/dataset_sources.py

@@ -126,7 +126,13 @@ def _cached_dataset_path(
 
 def _download_blob(blob, local_path: Path) -> None:
     local_path.parent.mkdir(parents=True, exist_ok=True)
-    temporary_path = local_path.with_name(f"{local_path.name}.tmp")
+    fd, temporary_path_name = tempfile.mkstemp(
+        prefix=f".{local_path.name}.",
+        suffix=".tmp",
+        dir=local_path.parent,
+    )
+    os.close(fd)
+    temporary_path = Path(temporary_path_name)
     try:
         blob.download_to_filename(str(temporary_path))
         os.replace(temporary_path, local_path)

policyengine_uk/simulation.py

@@ -276,6 +276,11 @@ def build_from_dataset_source(
             self.build_from_url(dataset_source)
             return
         if dataset_source.startswith("gs://"):
+            if dataset_source in _url_dataset_cache:
+                multi_year_dataset = _url_dataset_cache[dataset_source]
+                self.build_from_multi_year_dataset(multi_year_dataset)
+                self.dataset = multi_year_dataset
+                return
             dataset_file = materialize_gcs_dataset_url(dataset_source)
             self.build_from_file(dataset_file, cache_key=dataset_source)
             return

policyengine_uk/tests/test_dataset_sources.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 import pytest
@@ -13,13 +14,15 @@ def __init__(self, name, generation, metadata=None, contents=b"dataset"):
         self.metadata = metadata
         self.contents = contents
         self.download_count = 0
+        self.download_filenames = []
         self.reload_count = 0
 
     def reload(self):
         self.reload_count += 1
 
     def download_to_filename(self, filename):
         self.download_count += 1
+        self.download_filenames.append(filename)
         Path(filename).write_bytes(self.contents)
 
 
@@ -187,3 +190,33 @@ def test_materialize_gcs_dataset_url_reuses_cached_file(monkeypatch, tmp_path):
     assert first_path == second_path
     assert Path(second_path).read_bytes() == b"current"
     assert current_blob.download_count == 1
+
+
+def test_download_blob_uses_unique_temp_path_for_each_download(monkeypatch, tmp_path):
+    local_path = tmp_path / "cache" / "file.h5"
+    created_temp_paths = []
+
+    def fake_mkstemp(*, prefix, suffix, dir):
+        temporary_path = Path(dir) / f"{prefix}{len(created_temp_paths)}{suffix}"
+        fd = os.open(temporary_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600)
+        created_temp_paths.append(temporary_path)
+        return fd, str(temporary_path)
+
+    monkeypatch.setattr(dataset_sources.tempfile, "mkstemp", fake_mkstemp)
+    blob = FakeBlob("data/file.h5", 444, contents=b"first")
+
+    dataset_sources._download_blob(blob, local_path)
+    local_path.unlink()
+    blob.contents = b"second"
+    dataset_sources._download_blob(blob, local_path)
+
+    assert [
+        Path(filename) for filename in blob.download_filenames
+    ] == created_temp_paths
+    assert len(set(created_temp_paths)) == 2
+    assert all(
+        temporary_path.parent == local_path.parent
+        for temporary_path in created_temp_paths
+    )
+    assert all(not temporary_path.exists() for temporary_path in created_temp_paths)
+    assert local_path.read_bytes() == b"second"

policyengine_uk/tests/test_simulation_dataset_sources.py

@@ -85,6 +85,44 @@ def fake_build_from_file(dataset_file, *, cache_key=None):
     }
 
 
+def test_dataset_source_reuses_cached_gcs_dataset_before_materializing(monkeypatch):
+    captured = {}
+    cached_dataset = object()
+    simulation = Simulation.__new__(Simulation)
+    url = "gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.55.10"
+
+    def fake_materialize_gcs_dataset_url(dataset_url):
+        raise AssertionError("Cached gs:// datasets should not be materialized.")
+
+    def fake_build_from_file(dataset_file, *, cache_key=None):
+        raise AssertionError("Cached gs:// datasets should not be read from disk.")
+
+    def fake_build_from_multi_year_dataset(dataset):
+        captured["dataset"] = dataset
+
+    simulation_module._url_dataset_cache.pop(url, None)
+    simulation_module._url_dataset_cache[url] = cached_dataset
+    monkeypatch.setattr(
+        simulation_module,
+        "materialize_gcs_dataset_url",
+        fake_materialize_gcs_dataset_url,
+    )
+    monkeypatch.setattr(simulation, "build_from_file", fake_build_from_file)
+    monkeypatch.setattr(
+        simulation,
+        "build_from_multi_year_dataset",
+        fake_build_from_multi_year_dataset,
+    )
+
+    try:
+        Simulation.build_from_dataset_source(simulation, url)
+
+        assert captured["dataset"] is cached_dataset
+        assert simulation.dataset is cached_dataset
+    finally:
+        simulation_module._url_dataset_cache.pop(url, None)
+
+
 def test_dataset_source_rejects_unsupported_remote_urls():
     simulation = Simulation.__new__(Simulation)