Handle reference genome download failures

SigProfilerMatrixGenerator/controllers/cli_controller.py

@@ -60,7 +60,8 @@ def parse_arguments_install(args: List[str]) -> argparse.Namespace:
         "--local_genome",
         help="""
             Install an offline reference genome downloaded from the Alexandrov Lab's FTP server.
-            Provide the absolute path to the locally-stored genome file.
+            Provide the absolute path to the directory containing the locally-stored
+            <genome>.tar.gz file.
             For downloads, visit AlexandrovLab's ftp server:
             ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/SigProfilerMatrixGenerator/
             """,

SigProfilerMatrixGenerator/scripts/reference_genome_manager.py

@@ -348,6 +348,10 @@
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
 
 
+class GenomeDownloadError(RuntimeError):
+    """Raised when a reference genome archive cannot be downloaded or installed."""
+
+
 class ReferenceGenomeManager:
     """
     A class for downloading and managing reference genomes.
@@ -370,38 +374,77 @@ def download_genome(self, genome_name):
 
         file_name = f"{genome_name}.tar.gz"
         local_filepath = self.reference_dir.get_tsb_dir() / file_name
+        local_filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        downloaded = False
+        failures = []
 
         # First try to download using FTP
         for server_key in FTP_SERVERS:
             server_info = FTP_SERVERS[server_key]
             try:
+                self._remove_partial_archive(local_filepath)
                 self._download_via_ftplib(
                     server_info["address"],
                     server_info["path"],
                     file_name,
                     local_filepath,
                 )
-                logging.info(f"Downloaded {genome_name} from {server_key} using FTP.")
-                break  # Exit the loop if download is successful
+                if self._downloaded_archive_exists(local_filepath):
+                    logging.info(
+                        f"Downloaded {genome_name} from {server_key} using FTP."
+                    )
+                    downloaded = True
+                    break  # Exit the loop if download is successful
+                failures.append(
+                    f"{server_key} FTP: download completed but did not create "
+                    f"{local_filepath}."
+                )
             except ftplib.all_errors as e:
+                failures.append(f"{server_key} FTP: {e}")
                 logging.info(f"Attempt to download from {server_key} failed: {e}")
+            if downloaded:
+                break
             if shutil.which("curl"):
                 try:
+                    self._remove_partial_archive(local_filepath)
                     self._download_via_curl(
                         server_info["address"],
                         server_info["path"],
                         file_name,
                         local_filepath,
                     )
-                    logging.info(f"Downloaded {genome_name} using curl.")
-                    break  # Exit the loop if download is successful
-                except (subprocess.CalledProcessError, KeyboardInterrupt) as e:
+                    if self._downloaded_archive_exists(local_filepath):
+                        logging.info(f"Downloaded {genome_name} using curl.")
+                        downloaded = True
+                        break  # Exit the loop if download is successful
+                    failures.append(
+                        f"{server_key} curl: download completed but did not create "
+                        f"{local_filepath}."
+                    )
+                except (subprocess.CalledProcessError, OSError) as e:
+                    failures.append(f"{server_key} curl: {e}")
                     logging.error(f"Curl download failed with error: {e}")
             else:
+                failures.append(f"{server_key} curl: curl is not available.")
                 logging.error("Curl is not available. Unable to download the file.")
-                return
 
-        self._unzip_file(local_filepath)
+        if not downloaded:
+            self._remove_partial_archive(local_filepath)
+            raise GenomeDownloadError(
+                self._format_download_failure_message(file_name, failures)
+            )
+
+        try:
+            self._unzip_file(local_filepath)
+        except (tarfile.TarError, OSError) as e:
+            self._remove_partial_archive(local_filepath)
+            raise GenomeDownloadError(
+                f"Downloaded archive {local_filepath} could not be extracted. "
+                "Please retry the installation; if the problem persists, manually "
+                "download a fresh archive and install it with --local_genome in the "
+                "CLI or offline_files_path in the Python API."
+            ) from e
         local_filepath.unlink()
         logging.info(f"{genome_name} has been successfully installed.")
 
@@ -568,16 +611,57 @@ def _download_via_curl(self, ftp_server, ftp_path, filename, local_filepath):
         Helper function for download_genome.
         """
         try:
-            url = f"ftp://{ftp_server}/{ftp_path}/{filename}"
-            command = ["curl", "-o", str(local_filepath), url]
+            local_filepath.parent.mkdir(parents=True, exist_ok=True)
+            url = f"ftp://{ftp_server}/{ftp_path.rstrip('/')}/{filename}"
+            command = [
+                "curl",
+                "--fail",
+                "--location",
+                "--retry",
+                "3",
+                "--connect-timeout",
+                "30",
+                "-o",
+                str(local_filepath),
+                url,
+            ]
             subprocess.run(command, check=True)
         except subprocess.CalledProcessError as e:
             logging.error(f"Curl download failed with error: {e}.")
             raise
 
+    def _downloaded_archive_exists(self, local_filepath):
+        return local_filepath.exists() and local_filepath.stat().st_size > 0
+
+    def _remove_partial_archive(self, local_filepath):
+        try:
+            local_filepath.unlink()
+        except FileNotFoundError:
+            pass
+
+    def _format_download_failure_message(self, file_name, failures):
+        message = [
+            f"Unable to download {file_name} from the configured reference genome mirrors.",
+        ]
+        if failures:
+            message.append("Download attempts failed with:")
+            message.extend(f"  - {failure}" for failure in failures)
+        message.append(
+            "Please check FTP/network access or manually download the archive from "
+            "ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/SigProfilerMatrixGenerator/ "
+            "and install it with --local_genome in the CLI or offline_files_path in "
+            "the Python API. Both options expect the directory containing the archive."
+        )
+        return "\n".join(message)
+
     def _unzip_file(self, file_path):
         with tarfile.open(file_path, "r:gz") as tar:
-            tar.extractall(path=self.reference_dir.get_tsb_dir())
+            extraction_options = {}
+            if hasattr(tarfile, "data_filter"):
+                extraction_options["filter"] = "data"
+            tar.extractall(
+                path=self.reference_dir.get_tsb_dir(), **extraction_options
+            )
 
     def _verify_checksum(self, file_path, checksum):
         """

tests/scripts/test_reference_genome_manager.py

@@ -0,0 +1,156 @@
+import ftplib
+import io
+import subprocess
+import tarfile
+
+import pytest
+
+from SigProfilerMatrixGenerator.scripts import reference_genome_manager
+
+
+def write_test_archive(archive_path, genome_name="test_genome"):
+    archive_path.parent.mkdir(parents=True, exist_ok=True)
+    contents = b"test chromosome contents"
+    member = tarfile.TarInfo(f"{genome_name}/1.txt")
+    member.size = len(contents)
+    with tarfile.open(archive_path, "w:gz") as archive:
+        archive.addfile(member, io.BytesIO(contents))
+
+
+def test_download_genome_installs_archive_from_ftp(monkeypatch, tmp_path):
+    manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path)
+    archive_path = tmp_path.resolve() / "tsb" / "test_genome.tar.gz"
+    installed_file = tmp_path.resolve() / "tsb" / "test_genome" / "1.txt"
+
+    monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False)
+    monkeypatch.setattr(
+        manager,
+        "_download_via_ftplib",
+        lambda *args: write_test_archive(args[-1]),
+    )
+
+    def unexpected_curl(*args):
+        raise AssertionError("curl should not run after a successful FTP download")
+
+    monkeypatch.setattr(manager, "_download_via_curl", unexpected_curl)
+
+    manager.download_genome("test_genome")
+
+    assert installed_file.read_bytes() == b"test chromosome contents"
+    assert not archive_path.exists()
+
+
+def test_download_genome_falls_back_to_curl_on_same_mirror(monkeypatch, tmp_path):
+    manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path)
+    calls = []
+
+    monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False)
+    monkeypatch.setattr(reference_genome_manager.shutil, "which", lambda name: "curl")
+
+    def fail_ftp(server, path, filename, local_filepath):
+        calls.append(("ftp", server))
+        raise ftplib.error_temp("Connection reset by peer")
+
+    def successful_curl(server, path, filename, local_filepath):
+        calls.append(("curl", server))
+        write_test_archive(local_filepath)
+
+    monkeypatch.setattr(manager, "_download_via_ftplib", fail_ftp)
+    monkeypatch.setattr(manager, "_download_via_curl", successful_curl)
+
+    manager.download_genome("test_genome")
+
+    assert calls == [
+        ("ftp", "alexandrovlab-ftp.ucsd.edu"),
+        ("curl", "alexandrovlab-ftp.ucsd.edu"),
+    ]
+    assert (tmp_path.resolve() / "tsb" / "test_genome" / "1.txt").exists()
+
+
+def test_download_via_curl_uses_retrying_ftp_command(monkeypatch, tmp_path):
+    manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path)
+    archive_path = tmp_path / "downloads" / "ebv.tar.gz"
+    observed = {}
+
+    def capture_run(command, check):
+        observed["command"] = command
+        observed["check"] = check
+
+    monkeypatch.setattr(reference_genome_manager.subprocess, "run", capture_run)
+
+    manager._download_via_curl(
+        "alexandrovlab-ftp.ucsd.edu",
+        "pub/tools/SigProfilerMatrixGenerator/",
+        "ebv.tar.gz",
+        archive_path,
+    )
+
+    assert observed == {
+        "command": [
+            "curl",
+            "--fail",
+            "--location",
+            "--retry",
+            "3",
+            "--connect-timeout",
+            "30",
+            "-o",
+            str(archive_path),
+            "ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/"
+            "SigProfilerMatrixGenerator/ebv.tar.gz",
+        ],
+        "check": True,
+    }
+
+
+def test_download_genome_raises_clear_error_when_all_downloads_fail(
+    monkeypatch, tmp_path
+):
+    manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path)
+    archive_path = tmp_path.resolve() / "tsb" / "mm10.tar.gz"
+
+    monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False)
+    monkeypatch.setattr(reference_genome_manager.shutil, "which", lambda name: "curl")
+
+    def fail_ftp(*args):
+        raise ftplib.error_temp("Connection reset by peer")
+
+    def fail_curl(*args):
+        local_filepath = args[-1]
+        local_filepath.write_bytes(b"partial archive")
+        raise subprocess.CalledProcessError(56, ["curl"])
+
+    monkeypatch.setattr(manager, "_download_via_ftplib", fail_ftp)
+    monkeypatch.setattr(manager, "_download_via_curl", fail_curl)
+
+    with pytest.raises(reference_genome_manager.GenomeDownloadError) as error:
+        manager.download_genome("mm10")
+
+    message = str(error.value)
+    assert "Unable to download mm10.tar.gz" in message
+    assert "alexandrovlab FTP" in message
+    assert "sanger curl" in message
+    assert "--local_genome" in message
+    assert "offline_files_path" in message
+    assert "directory containing the archive" in message
+    assert not archive_path.exists()
+
+
+def test_download_genome_removes_bad_archive_when_extract_fails(monkeypatch, tmp_path):
+    manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path)
+    archive_path = tmp_path.resolve() / "tsb" / "mm10.tar.gz"
+
+    monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False)
+
+    def write_bad_archive(*args):
+        local_filepath = args[-1]
+        local_filepath.write_bytes(b"not a tar.gz archive")
+
+    monkeypatch.setattr(manager, "_download_via_ftplib", write_bad_archive)
+
+    with pytest.raises(reference_genome_manager.GenomeDownloadError) as error:
+        manager.download_genome("mm10")
+
+    assert "could not be extracted" in str(error.value)
+    assert "offline_files_path" in str(error.value)
+    assert not archive_path.exists()