fix: download Zenodo record files via the files API

brainles_preprocessing/utils/zenodo.py

@@ -1,15 +1,18 @@
 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
 
 import requests
 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,16 +113,16 @@ 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}")
 
         if not zenodo_response:
             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)

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",
     )