diff --git a/brainles_preprocessing/utils/zenodo.py b/brainles_preprocessing/utils/zenodo.py index fe7183f..640149e 100644 --- a/brainles_preprocessing/utils/zenodo.py +++ b/brainles_preprocessing/utils/zenodo.py @@ -1,8 +1,8 @@ from __future__ import annotations import shutil +import tempfile import zipfile -from io import BytesIO from pathlib import Path from typing import ClassVar, Dict, List, Tuple @@ -10,6 +10,9 @@ from loguru import logger from rich.progress import Progress, SpinnerColumn, TextColumn +# (connect, read) timeout in seconds for every request to Zenodo. +REQUEST_TIMEOUT = (10, 60) + ATLASES_FOLDER = Path(__file__).parent.parent / "registration" / "atlases" ATLASES_RECORD_ID = "15236131" @@ -97,7 +100,7 @@ def fetch(self) -> Path: def _fetch_uncached(self) -> Path: """Perform the actual Zenodo check / download without consulting the cache.""" - zenodo_response = self._get_metadata_and_archive_url() + zenodo_response = self._get_metadata_and_files() pattern = self._glob_pattern() matching_folders = list(self.target_dir.glob(pattern)) @@ -110,8 +113,8 @@ def _fetch_uncached(self) -> Path: raise ZenodoException(msg) logger.info(f"{self.label.title()} not found locally.") - metadata, archive_url = zenodo_response - return self._download(metadata, archive_url) + metadata, files = zenodo_response + return self._download(metadata, files) logger.info(f"Found local {self.label}: {latest_local}") @@ -119,7 +122,7 @@ def _fetch_uncached(self) -> Path: logger.warning(f"Zenodo unreachable. Using latest downloaded {self.label}.") return self.target_dir / latest_local - metadata, archive_url = zenodo_response + metadata, files = zenodo_response remote_version = metadata["version"] local_version = latest_local.split("_v")[1] @@ -130,13 +133,14 @@ def _fetch_uncached(self) -> Path: logger.info( f"New version of {self.label} available on Zenodo ({remote_version}). Replacing local copy..." ) - shutil.rmtree( - self.target_dir / latest_local, - onerror=lambda func, path, excinfo: logger.warning( - f"Failed to delete {path}: {excinfo}" - ), + # The old version is only removed once the new one has been fully + # downloaded and staged, so a failed download never destroys a + # known-good local copy. + return self._download( + metadata, + files, + stale_folder=self.target_dir / latest_local, ) - return self._download(metadata, archive_url) def _glob_pattern(self) -> str: return f"{self.record_id}_v*.*.*" @@ -162,9 +166,22 @@ def _get_latest_version_folder_name( return None return latest.name - def _get_metadata_and_archive_url(self) -> Tuple[Dict, str] | None: + def _get_metadata_and_files(self) -> Tuple[Dict, List[Dict]] | None: + """Retrieve the record metadata and its file list from Zenodo. + + Zenodo's API used to expose a single ``links["archive"]`` URL pointing to + a zip of all files. That link has been removed, so the record's ``files`` + list is used instead: each entry provides a ``key`` and a ``links["self"]`` + download URL. + + Returns: + A ``(metadata, files)`` tuple, or ``None`` if Zenodo is unreachable or + returns an unexpected response. + """ try: - response = requests.get(f"{self.BASE_URL}/{self.record_id}") + response = requests.get( + f"{self.BASE_URL}/{self.record_id}", timeout=REQUEST_TIMEOUT + ) if response.status_code != 200: error_msg = ( f"Cannot find record '{self.record_id}' on Zenodo " @@ -174,65 +191,162 @@ def _get_metadata_and_archive_url(self) -> Tuple[Dict, str] | None: raise ZenodoException(error_msg) data = response.json() - return data["metadata"], data["links"]["archive"] + metadata = data["metadata"] + if not isinstance(metadata, dict) or not isinstance( + metadata.get("version"), str + ): + raise ValueError("record metadata is missing a string 'version'") + + files = data.get("files") + if files: + for entry in files: + if ( + not isinstance(entry, dict) + or not isinstance(entry.get("key"), str) + or not isinstance(entry.get("links", {}).get("self"), str) + ): + raise ValueError( + f"invalid file descriptor in record: {entry!r}" + ) + return metadata, files + + # Fallback for the deprecated archive link (kept for compatibility). + archive_url = data.get("links", {}).get("archive") + if isinstance(archive_url, str): + return metadata, [ + {"key": f"{self.record_id}.zip", "links": {"self": archive_url}} + ] + + logger.warning( + f"Zenodo record '{self.record_id}' does not expose any downloadable files." + ) + return None except requests.exceptions.RequestException as e: logger.warning(f"Failed to fetch metadata from Zenodo: {e}") return None + except (KeyError, ValueError) as e: + logger.warning( + f"Unexpected Zenodo response for record '{self.record_id}': {e}" + ) + return None def _download( self, metadata: Dict, - archive_url: str, + files: List[Dict], + stale_folder: Path | None = None, ) -> Path: folder = self._build_folder_path(metadata["version"]) - folder.mkdir(parents=True, exist_ok=True) + self.target_dir.mkdir(parents=True, exist_ok=True) + + # Stage everything in a temporary directory first, so a failed download + # never leaves a partial install or destroys an existing good version. + temp_folder = Path( + tempfile.mkdtemp(dir=self.target_dir, prefix=f".{folder.name}.tmp.") + ) logger.info(f"Downloading {self.label} from Zenodo. This may take a while...") - response = requests.get(archive_url, stream=True) - if response.status_code != 200: - msg = ( - f"Failed to download {self.label}. Status code: {response.status_code}" - ) - logger.error(msg) - raise ZenodoException(msg) + try: + for file in files: + key = file["key"] + file_url = file["links"]["self"] + destination = self._safe_destination(temp_folder, key) + self._download_file(file_url, destination) + + if destination.suffix == ".zip": + self._extract_zip(destination, temp_folder) + + if folder.exists(): + shutil.rmtree(folder, ignore_errors=True) + + # Install the fully staged version before removing the previous one, + # so an interrupted/failed swap can never leave us with neither. + temp_folder.rename(folder) + + if ( + stale_folder is not None + and stale_folder != folder + and stale_folder.exists() + ): + shutil.rmtree( + stale_folder, + onerror=lambda func, path, excinfo: logger.warning( + f"Failed to delete {path}: {excinfo}" + ), + ) + except BaseException: + shutil.rmtree(temp_folder, ignore_errors=True) + raise - self._extract_archive(response, folder) - logger.info(f"{self.label.title()} extracted to {folder}") + logger.info(f"{self.label.title()} downloaded to {folder}") return folder - def _extract_archive( + @staticmethod + def _safe_destination(base: Path, key: str) -> Path: + """Join ``key`` onto ``base`` and reject paths that escape it. + + ``key`` comes from the (remote) record metadata, so it must not be able + to write outside the staging directory via ``..`` or an absolute path. + """ + base_resolved = base.resolve() + destination = (base_resolved / key).resolve() + if destination != base_resolved and base_resolved not in destination.parents: + raise ZenodoException(f"Unsafe file key in record: {key!r}") + return destination + + def _download_file(self, url: str, destination: Path) -> None: + chunk_size = 1024 * 1024 # 1 MiB + destination.parent.mkdir(parents=True, exist_ok=True) + + try: + response = requests.get(url, stream=True, timeout=REQUEST_TIMEOUT) + if response.status_code != 200: + msg = ( + f"Failed to download {self.label} file '{destination.name}'. " + f"Status code: {response.status_code}" + ) + logger.error(msg) + raise ZenodoException(msg) + + with Progress( + SpinnerColumn(), + TextColumn(f"[cyan]Downloading {destination.name}..."), + TextColumn("[cyan]{task.completed:.2f} MB"), + transient=True, + ) as progress: + task = progress.add_task("", total=None) + with open(destination, "wb") as f: + for chunk in response.iter_content(chunk_size=chunk_size): + if not chunk: + continue + f.write(chunk) + progress.update(task, advance=len(chunk) / (1024**2)) + except requests.exceptions.RequestException as e: + msg = f"Failed to download {self.label} file '{destination.name}': {e}" + logger.error(msg) + raise ZenodoException(msg) from e + + def _extract_zip( self, - response: requests.Response, + zip_path: Path, folder: Path, - ): - chunk_size = 1024 # 1KB - buffer = BytesIO() - - with Progress( - SpinnerColumn(), - TextColumn(f"[cyan]Downloading {self.label}..."), - TextColumn("[cyan]{task.completed:.2f} MB"), - transient=True, - ) as progress: - task = progress.add_task("", total=None) - for chunk in response.iter_content(chunk_size=chunk_size): - buffer.write(chunk) - progress.update(task, advance=len(chunk) / (chunk_size**2)) - - with zipfile.ZipFile(buffer) as z: - z.extractall(folder) - - for file in folder.iterdir(): - if file.is_file() and file.suffix == ".zip": - with zipfile.ZipFile(file) as inner_zip: - files = inner_zip.namelist() - with Progress(transient=True) as progress: - task = progress.add_task( - f"[cyan]Extracting inner zip...", total=len(files) - ) - for i, f in enumerate(files): - inner_zip.extract(f, folder) - progress.update(task, completed=i + 1) - file.unlink() # Remove inner zip after extraction + ) -> None: + with zipfile.ZipFile(zip_path) as z: + names = z.namelist() + with Progress(transient=True) as progress: + task = progress.add_task( + f"[cyan]Extracting {zip_path.name}...", total=len(names) + ) + for i, name in enumerate(names): + self._safe_destination(folder, name) + z.extract(name, folder) + progress.update(task, completed=i + 1) + zip_path.unlink() # Remove the archive after extraction + + # Extract any archives nested inside this one (e.g. a record whose + # archive contains the SynthStrip weights zip). + for nested in folder.rglob("*.zip"): + if nested.is_file(): + self._extract_zip(nested, folder) diff --git a/tests/test_preprocessor.py b/tests/test_preprocessor.py index 07ee5bc..b70d0df 100644 --- a/tests/test_preprocessor.py +++ b/tests/test_preprocessor.py @@ -47,6 +47,7 @@ def test_no_name_conflicts(dummy_registrator, dummy_brain_extractor, dummy_defac registrator=dummy_registrator, brain_extractor=dummy_brain_extractor, defacer=dummy_defacer, + atlas_image_path="tmp/atlas.nii.gz", ) diff --git a/tests/test_zenodo.py b/tests/test_zenodo.py index 724cfe6..e65cf2a 100644 --- a/tests/test_zenodo.py +++ b/tests/test_zenodo.py @@ -1,5 +1,7 @@ +from io import BytesIO from pathlib import Path from unittest.mock import MagicMock, patch +import zipfile import pytest import requests @@ -31,38 +33,83 @@ def dummy_metadata(): @pytest.fixture -def dummy_archive_url(): - return "https://zenodo.org/record/dummy/archive.zip" +def dummy_files(): + return [ + { + "key": "dummy.nii", + "links": { + "self": "https://zenodo.org/record/dummy/files/dummy.nii/content" + }, + } + ] @pytest.fixture -def dummy_zenodo_response(dummy_metadata, dummy_archive_url): - return dummy_metadata, dummy_archive_url +def dummy_zenodo_response(dummy_metadata, dummy_files): + return dummy_metadata, dummy_files -# ---- Tests for _get_metadata_and_archive_url ---- +# ---- Tests for _get_metadata_and_files ---- -def test_get_metadata_and_archive_url_success( - monkeypatch, dummy_metadata, dummy_archive_url -): +def test_get_metadata_and_files_success(monkeypatch, dummy_metadata, dummy_files): + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.json.return_value = { + "metadata": dummy_metadata, + "files": dummy_files, + } + + monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) + record = ZenodoRecord("123", Path("/tmp"), "test") + + metadata, files = record._get_metadata_and_files() + + assert metadata == dummy_metadata + assert files == dummy_files + + +def test_get_metadata_and_files_falls_back_to_archive(monkeypatch, dummy_metadata): + archive_url = "https://zenodo.org/record/dummy/archive.zip" response_mock = MagicMock() response_mock.status_code = 200 response_mock.json.return_value = { "metadata": dummy_metadata, - "links": {"archive": dummy_archive_url}, + "links": {"archive": archive_url}, } monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) record = ZenodoRecord("123", Path("/tmp"), "test") - metadata, url = record._get_metadata_and_archive_url() + metadata, files = record._get_metadata_and_files() assert metadata == dummy_metadata - assert url == dummy_archive_url + assert files[0]["links"]["self"] == archive_url + + +def test_get_metadata_and_files_no_files_returns_none(monkeypatch, dummy_metadata): + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.json.return_value = {"metadata": dummy_metadata} + + monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) + record = ZenodoRecord("123", Path("/tmp"), "test") + + assert record._get_metadata_and_files() is None + + +def test_get_metadata_and_files_malformed_response_returns_none(monkeypatch): + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.json.return_value = {"unexpected": "shape"} + + monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) + record = ZenodoRecord("123", Path("/tmp"), "test") + assert record._get_metadata_and_files() is None -def test_get_metadata_and_archive_url_failure(monkeypatch): + +def test_get_metadata_and_files_failure(monkeypatch): response_mock = MagicMock() response_mock.status_code = 404 @@ -70,16 +117,16 @@ def test_get_metadata_and_archive_url_failure(monkeypatch): record = ZenodoRecord("invalid", Path("/tmp"), "test") with pytest.raises(ZenodoException): - record._get_metadata_and_archive_url() + record._get_metadata_and_files() @patch( "requests.get", side_effect=requests.exceptions.RequestException("Connection error") ) -def test_get_metadata_and_archive_url_connection_error(mock_get): +def test_get_metadata_and_files_connection_error(mock_get): record = ZenodoRecord("123", Path("/tmp"), "test") - assert record._get_metadata_and_archive_url() is None + assert record._get_metadata_and_files() is None # ---- Tests for _get_latest_version_folder_name ---- @@ -111,7 +158,7 @@ def test_get_latest_version_folder_name_ignores_empty_folder(tmp_path): # ---- Tests for fetch() method ---- -@patch.object(ZenodoRecord, "_get_metadata_and_archive_url") +@patch.object(ZenodoRecord, "_get_metadata_and_files") @patch.object(ZenodoRecord, "_download") def test_fetch_downloads_new_if_no_local( mock_download, mock_metadata, tmp_path, dummy_zenodo_response @@ -126,7 +173,7 @@ def test_fetch_downloads_new_if_no_local( mock_download.assert_called_once() -@patch.object(ZenodoRecord, "_get_metadata_and_archive_url", return_value=None) +@patch.object(ZenodoRecord, "_get_metadata_and_files", return_value=None) def test_fetch_zenodo_unreachable_raises(mock_metadata, tmp_path): record = ZenodoRecord("123", tmp_path, "test") @@ -134,7 +181,7 @@ def test_fetch_zenodo_unreachable_raises(mock_metadata, tmp_path): record.fetch() -@patch.object(ZenodoRecord, "_get_metadata_and_archive_url") +@patch.object(ZenodoRecord, "_get_metadata_and_files") @patch.object(ZenodoRecord, "_download") def test_fetch_skips_if_latest_present( mock_download, mock_metadata, tmp_path, dummy_zenodo_response @@ -151,7 +198,7 @@ def test_fetch_skips_if_latest_present( mock_download.assert_not_called() -@patch.object(ZenodoRecord, "_get_metadata_and_archive_url") +@patch.object(ZenodoRecord, "_get_metadata_and_files") @patch.object(ZenodoRecord, "_download") def test_fetch_replaces_old_version( mock_download, mock_metadata, tmp_path, dummy_zenodo_response @@ -167,8 +214,195 @@ def test_fetch_replaces_old_version( result = record.fetch() assert result.name == "123_v1.2.3" - assert not old_folder.exists() mock_download.assert_called_once() + # Removal of the old version is delegated to _download (atomic replace). + assert mock_download.call_args.kwargs["stale_folder"] == old_folder + + +# ---- Tests for downloading record files ---- + + +def _fake_response(content: bytes, status_code: int = 200): + response = MagicMock() + response.status_code = status_code + response.iter_content.return_value = [content] + return response + + +def test_download_writes_individual_files(tmp_path): + files = [ + {"key": "a.nii", "links": {"self": "https://example.org/a.nii"}}, + {"key": "sub/b.nii", "links": {"self": "https://example.org/b.nii"}}, + ] + responses = [_fake_response(b"AAAA"), _fake_response(b"BBBB")] + with patch("requests.get", side_effect=responses): + record = ZenodoRecord("123", tmp_path, "test") + folder = record._download({"version": "1.2.3"}, files) + + assert folder == tmp_path / "123_v1.2.3" + assert (folder / "a.nii").read_bytes() == b"AAAA" + assert (folder / "sub" / "b.nii").read_bytes() == b"BBBB" + + +def test_download_extracts_zip_file(tmp_path): + buffer = BytesIO() + with zipfile.ZipFile(buffer, "w") as z: + z.writestr("inner.txt", "hello") + zip_bytes = buffer.getvalue() + + files = [{"key": "weights.zip", "links": {"self": "https://example.org/w.zip"}}] + with patch("requests.get", return_value=_fake_response(zip_bytes)): + record = ZenodoRecord("123", tmp_path, "test") + folder = record._download({"version": "1.0.0"}, files) + + assert (folder / "inner.txt").read_text() == "hello" + assert not (folder / "weights.zip").exists() + + +def test_download_file_raises_on_error(tmp_path): + response = _fake_response(b"", status_code=500) + with patch("requests.get", return_value=response): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(ZenodoException): + record._download_file("https://example.org/a.nii", tmp_path / "a.nii") + + +def test_download_file_request_exception_raises_zenodo_exception(tmp_path): + with patch( + "requests.get", + side_effect=requests.exceptions.RequestException("connection reset"), + ): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(ZenodoException): + record._download_file("https://example.org/a.nii", tmp_path / "a.nii") + + +def test_download_extracts_nested_zip(tmp_path): + inner = BytesIO() + with zipfile.ZipFile(inner, "w") as z: + z.writestr("deep.txt", "deep") + outer = BytesIO() + with zipfile.ZipFile(outer, "w") as z: + z.writestr("nested.zip", inner.getvalue()) + + files = [{"key": "outer.zip", "links": {"self": "https://example.org/o.zip"}}] + with patch("requests.get", return_value=_fake_response(outer.getvalue())): + record = ZenodoRecord("123", tmp_path, "test") + folder = record._download({"version": "1.0.0"}, files) + + assert (folder / "deep.txt").read_text() == "deep" + assert not list(folder.rglob("*.zip")) + + +def test_download_failure_keeps_previous_version_and_cleans_temp(tmp_path): + old_folder = tmp_path / "123_v1.0.0" + old_folder.mkdir() + (old_folder / "keep.nii").write_text("old") + + files = [{"key": "a.nii", "links": {"self": "https://example.org/a.nii"}}] + with patch("requests.get", return_value=_fake_response(b"", status_code=500)): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(ZenodoException): + record._download({"version": "1.2.3"}, files, stale_folder=old_folder) + + # Previous version is untouched and no partial/temp installs remain. + assert (old_folder / "keep.nii").read_text() == "old" + assert not (tmp_path / "123_v1.2.3").exists() + assert not list(tmp_path.glob(".123_v1.2.3.tmp.*")) + + +def test_download_replaces_stale_version_on_success(tmp_path): + old_folder = tmp_path / "123_v1.0.0" + old_folder.mkdir() + (old_folder / "keep.nii").write_text("old") + + files = [{"key": "a.nii", "links": {"self": "https://example.org/a.nii"}}] + with patch("requests.get", return_value=_fake_response(b"AAAA")): + record = ZenodoRecord("123", tmp_path, "test") + folder = record._download({"version": "1.2.3"}, files, stale_folder=old_folder) + + assert not old_folder.exists() + assert (folder / "a.nii").read_bytes() == b"AAAA" + assert not list(tmp_path.glob(".123_v1.2.3.tmp.*")) + + +@pytest.mark.parametrize("unsafe_key", ["../escape.nii", "/tmp/escape.nii"]) +def test_download_rejects_unsafe_key(tmp_path, unsafe_key): + files = [{"key": unsafe_key, "links": {"self": "https://example.org/a.nii"}}] + with patch("requests.get", return_value=_fake_response(b"AAAA")): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(ZenodoException): + record._download({"version": "1.2.3"}, files) + + assert not (tmp_path / "escape.nii").exists() + assert not (tmp_path.parent / "escape.nii").exists() + assert not (tmp_path / "123_v1.2.3").exists() + assert not list(tmp_path.glob(".123_v1.2.3.tmp.*")) + + +def test_download_extract_rejects_unsafe_zip_member(tmp_path): + buffer = BytesIO() + with zipfile.ZipFile(buffer, "w") as z: + z.writestr("../escape.txt", "evil") + + files = [{"key": "evil.zip", "links": {"self": "https://example.org/e.zip"}}] + with patch("requests.get", return_value=_fake_response(buffer.getvalue())): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(ZenodoException): + record._download({"version": "1.2.3"}, files) + + assert not (tmp_path / "escape.txt").exists() + assert not (tmp_path.parent / "escape.txt").exists() + assert not list(tmp_path.glob(".123_v1.2.3.tmp.*")) + + +def test_download_keeps_stale_version_if_install_fails(tmp_path): + old_folder = tmp_path / "123_v1.0.0" + old_folder.mkdir() + (old_folder / "keep.nii").write_text("old") + + files = [{"key": "a.nii", "links": {"self": "https://example.org/a.nii"}}] + with patch("requests.get", return_value=_fake_response(b"AAAA")), patch( + "pathlib.Path.rename", side_effect=OSError("boom") + ): + record = ZenodoRecord("123", tmp_path, "test") + with pytest.raises(OSError): + record._download({"version": "1.2.3"}, files, stale_folder=old_folder) + + # The known-good version must survive a failed install. + assert (old_folder / "keep.nii").read_text() == "old" + assert not (tmp_path / "123_v1.2.3").exists() + assert not list(tmp_path.glob(".123_v1.2.3.tmp.*")) + + +def test_get_metadata_and_files_invalid_file_descriptor_returns_none( + monkeypatch, dummy_metadata +): + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.json.return_value = { + "metadata": dummy_metadata, + "files": [{"key": "a.nii"}], # missing links.self + } + + monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) + record = ZenodoRecord("123", Path("/tmp"), "test") + + assert record._get_metadata_and_files() is None + + +def test_get_metadata_and_files_missing_version_returns_none(monkeypatch): + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.json.return_value = { + "metadata": {"title": "no version"}, + "files": [{"key": "a.nii", "links": {"self": "https://example.org/a.nii"}}], + } + + monkeypatch.setattr("requests.get", lambda *args, **kwargs: response_mock) + record = ZenodoRecord("123", Path("/tmp"), "test") + + assert record._get_metadata_and_files() is None # ---- Tests for process-level caching ----