diff --git a/docs/usage.md b/docs/usage.md index e3885f1..c1d7eb8 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -239,7 +239,59 @@ pridepy download-px-raw-files \ | --- | --- | --- | | `-a, --accession` | ProteomeXchange accession (e.g. `PXD039236`). `--px` is a deprecated alias | required | | `-o, --output-folder` | Destination directory | required | +| `-p, --protocol` | Transfer protocol: `ftp`, `aspera`, `globus`, `s3` (FTP-first with fallback) | `ftp` | +| `-w, --parallel-files` | Download 1–32 files concurrently (across-file concurrency) | `1` | +| `-t, --threads` | Parallel HTTP Range threads per file (1–32) for fast per-file downloads | `1` | | `--skip-if-downloaded-already` | Skip files already present locally | off | +| `--preserve-structure` | Recreate the dataset's subdirectory layout under the output folder | off | +| `--iprox-user` | iProX account username (with `--protocol aspera`; env fallback: `IPROX_USER`) | — | + +The iProX Aspera password is never accepted as a command-line flag. Set the +`IPROX_ASPERA_PASSWORD` environment variable, or omit it and `pridepy` will +prompt for it securely (hidden input) when `--protocol aspera` is used. + +### Fast downloads: parallel files and per-file segments + +Combine `-w` (files in parallel) and `-t` (Range segments per file) for fast bulk downloads. +The total concurrent connections is approximately `parallel_files × threads`. + +**Parallel across files (recommended for most users, no account required):** + +```bash +# Download up to 8 files concurrently from ProteomeXchange +pridepy download-px-raw-files \ + -a PXD077178 \ + -o ./PXD077178 \ + -w 8 +``` + +**Combine parallel files with per-file segments:** + +```bash +# Download 8 files in parallel, each split into 4 Range segments +pridepy download-px-raw-files \ + -a PXD077178 \ + -o ./out \ + -w 8 \ + -t 4 +``` + +### Fast downloads with iProX Aspera (account required) + +iProX offers Aspera for very large bulk transfers. Aspera is faster than HTTP +on high-bandwidth connections but requires an iProX account. Combine +`--protocol aspera` with `--iprox-user`; the password is read from +`IPROX_ASPERA_PASSWORD` or prompted for securely (never passed as a flag): + +```bash +# Download via iProX Aspera with 8-file parallelism +IPROX_ASPERA_PASSWORD=your_password pridepy download-px-raw-files \ + -a PXD077178 \ + -o ./out \ + --protocol aspera \ + --iprox-user your_username \ + -w 8 +``` ### Go directly to the hosting repository (native MassIVE / JPOST / iProX accessions) @@ -268,7 +320,7 @@ How each repository is enumerated: - **MassIVE** walks the FTPS tree at `massive-ftp.ucsd.edu` (the server requires TLS). MassIVE distributes datasets across versioned root directories (`/v01`–`/vNN`); `pridepy` discovers the correct root automatically. If FTP/FTPS is blocked by the network, `pridepy` falls back to HTTPS: it lists the dataset from the GNPS2 file index (`datasetcache.gnps2.org`) and downloads each file from the ProteoSAFe endpoint at `massive.ucsd.edu` (byte-identical to the FTPS copy). - **JPOST** lists files through the JSON PROXI endpoint at `https://repository.jpostdb.org/proxi/datasets/` and downloads from `ftp.jpostdb.org` over plain FTP. The PROXI listing avoids the source-IP connection limit JPOST enforces on FTP. -- **iProX** fetches the dataset's ProteomeXchange XML from `http://download.iprox.org//PX_.xml`, then downloads each referenced file from the same host over anonymous HTTP (with `Range` support for resume). iProX also exposes Aspera (`faspe://`) with username/password for very large bulk transfers; `pridepy` uses the public HTTP endpoint so no iProX credentials are required. +- **iProX** fetches the dataset's ProteomeXchange XML from `http://download.iprox.org//PX_.xml`, then downloads each referenced file from the same host over anonymous HTTP (with `Range` support for resume). iProX also exposes Aspera with username/password for very large bulk transfers; `pridepy` uses the public HTTP endpoint so no iProX credentials are required. `download-all-public-raw-files` retrieves the files stored under the dataset's `raw/` collection. These direct downloads support resume (REST for FTP, diff --git a/pridepy/download/client.py b/pridepy/download/client.py index 5dd10f7..6707810 100644 --- a/pridepy/download/client.py +++ b/pridepy/download/client.py @@ -318,8 +318,21 @@ def download_px_raw_files( output_folder: str, skip_if_downloaded_already: bool = True, flatten: bool = True, + parallel_files: int = 1, + download_threads: int = 1, + protocol: str = "ftp", + iprox_user: Optional[str] = None, + iprox_password: Optional[str] = None, ) -> None: """Delegate to :meth:`ProteomeXchangeProvider.download_from_accession_or_url`.""" return ProteomeXchangeProvider().download_from_accession_or_url( - px_id_or_url, output_folder, skip_if_downloaded_already, flatten=flatten + px_id_or_url, + output_folder, + skip_if_downloaded_already, + flatten=flatten, + parallel_files=parallel_files, + download_threads=download_threads, + protocol=protocol, + iprox_user=iprox_user, + iprox_password=iprox_password, ) diff --git a/pridepy/download/iprox.py b/pridepy/download/iprox.py index db5e7fc..608fd22 100644 --- a/pridepy/download/iprox.py +++ b/pridepy/download/iprox.py @@ -14,6 +14,8 @@ import logging import os import re +import subprocess +from concurrent.futures import ThreadPoolExecutor, as_completed import defusedxml.ElementTree as ET from typing import ClassVar, Dict, List, Optional from urllib.parse import urlparse @@ -23,6 +25,7 @@ from pridepy.download import registry from pridepy.download.base import Provider from pridepy.download.jpost import JpostProvider +from pridepy.download.transport import _safe_join @registry.register @@ -34,6 +37,9 @@ class IproxProvider(Provider): PX_XML_URL_TEMPLATE: ClassVar[str] = ( "http://download.iprox.org/{accession}/PX_{accession}.xml" ) + ASPERA_HOST: ClassVar[str] = "download.iprox.org" + ASPERA_PORT: ClassVar[str] = "33001" + ASPERA_ROOT: ClassVar[str] = "/data/iprox" # iProX PX XML uses the same PSI-MS cvParam "name" values as JPOST PROXI, # so we reuse JpostProvider's category map. PX_CATEGORY_MAP: ClassVar[Dict[str, str]] = JpostProvider.PROXI_CATEGORY_MAP @@ -45,6 +51,134 @@ def matches(accession: str) -> bool: return False return bool(re.fullmatch(r"IPX\d{7,10}", accession.upper())) + @staticmethod + def _ascp_binary() -> str: + # Reuse PRIDE's bundled ascp binary resolution. + from pridepy.download.pride import PrideProvider + return PrideProvider.get_ascp_binary() + + @classmethod + def _aspera_download_one( + cls, + ascp: str, + url: str, + relpath: Optional[str], + output_folder: str, + user: str, + password: str, + maximum_bandwidth: str, + skip_if_downloaded_already: bool, + env: Dict[str, str], + ) -> Optional[str]: + """Download a single URL via ascp. Returns ``url`` on failure, else None.""" + path = urlparse(url).path.lstrip("/") # e.g. IPX.../.../a.raw + source = f"{user}@{cls.ASPERA_HOST}:{cls.ASPERA_ROOT}/{path}" + if relpath: + dest = _safe_join(output_folder, relpath) + else: + dest = os.path.join(output_folder, os.path.basename(urlparse(url).path)) + dest_parent = os.path.dirname(dest) or output_folder + os.makedirs(dest_parent, exist_ok=True) + if ( + skip_if_downloaded_already + and os.path.isfile(dest) + and os.path.getsize(dest) > 0 + ): + logging.info(f"Skipping download as file already exists: {dest}") + return None + argv = [ + ascp, "-QT", "-P", cls.ASPERA_PORT, "-l", maximum_bandwidth, + "-k", "2", source, dest, + ] + logging.info( + "Aspera: %s -> %s", source.replace(password, "***"), dest + ) + try: + subprocess.run(argv, check=True, env=env) + return None + except subprocess.CalledProcessError as e: + logging.error(f"iProX Aspera failed for {url}: {e}") + return url + + @classmethod + def aspera_download( + cls, + urls: List[str], + output_folder: str, + relative_paths: List[Optional[str]], + user: Optional[str], + password: Optional[str], + maximum_bandwidth: str = "100M", + skip_if_downloaded_already: bool = False, + parallel_files: int = 1, + ) -> None: + """Download iProX-hosted URLs via ascp on port 33001. + + Requires iProX account credentials; the password is passed to the + subprocess through ASPERA_SCP_PASS (never argv). When + ``parallel_files`` > 1, transfers run concurrently: each ``ascp`` + invocation is its own subprocess writing its own destination file, so + this is safe. + """ + if not user or not password: + raise ValueError( + "iProX Aspera requires credentials: pass --iprox-user and set " + "IPROX_ASPERA_PASSWORD (or answer the password prompt), or use " + "the default parallel HTTP transport instead." + ) + ascp = cls._ascp_binary() + env = dict(os.environ) + env["ASPERA_SCP_PASS"] = password + os.makedirs(output_folder, exist_ok=True) + failed: List[str] = [] + workers = max(1, min(parallel_files, len(urls))) + if workers > 1: + with ThreadPoolExecutor(max_workers=workers) as executor: + future_to_url = { + executor.submit( + cls._aspera_download_one, + ascp, + url, + relative_paths[idx] if idx < len(relative_paths) else None, + output_folder, + user, + password, + maximum_bandwidth, + skip_if_downloaded_already, + env, + ): url + for idx, url in enumerate(urls) + } + for future in as_completed(future_to_url): + url = future_to_url[future] + try: + result = future.result() + if result is not None: + failed.append(result) + except Exception as e: + logging.error(f"iProX Aspera failed for {url}: {e}") + failed.append(url) + else: + for idx, url in enumerate(urls): + relpath = relative_paths[idx] if idx < len(relative_paths) else None + result = cls._aspera_download_one( + ascp, + url, + relpath, + output_folder, + user, + password, + maximum_bandwidth, + skip_if_downloaded_already, + env, + ) + if result is not None: + failed.append(result) + if failed: + raise RuntimeError( + f"iProX Aspera download failed for {len(failed)} file(s): {failed}" + ) + @staticmethod def _get_public_root(accession: str) -> str: return f"/{accession.upper()}" diff --git a/pridepy/download/proteomexchange.py b/pridepy/download/proteomexchange.py index 32ea72e..81bf508 100644 --- a/pridepy/download/proteomexchange.py +++ b/pridepy/download/proteomexchange.py @@ -26,10 +26,11 @@ import posixpath import re import defusedxml.ElementTree as ET -from typing import ClassVar, Dict, List +from typing import ClassVar, Dict, List, Optional from urllib.parse import urlparse from pridepy.download.base import Provider +from pridepy.download.util import flatten_relative_paths from pridepy.util.api_handling import Util @@ -170,23 +171,79 @@ def download_from_accession_or_url( output_folder: str, skip_if_downloaded_already: bool = True, flatten: bool = True, + parallel_files: int = 1, + download_threads: int = 1, + protocol: str = "ftp", + iprox_user: Optional[str] = None, + iprox_password: Optional[str] = None, ) -> None: """End-to-end: resolve XML, list files, partition by scheme, download. Convenience for the ``download-px-raw-files`` CLI command — combines :meth:`list_files` and :meth:`download_files` with the original ``download_px_raw_files`` defaults (skip-if-downloaded-already - defaults to ``True``, no parallel workers). + defaults to ``True``). ``parallel_files`` controls across-file + concurrency and ``download_threads`` controls per-file HTTP Range + segments; ``protocol`` flows into :meth:`download_files` (ftp/http(s) + are handled directly today). + + When ``protocol == "aspera"``, iProX-hosted files are routed through + :meth:`IproxProvider.aspera_download` instead of the HTTP/FTP path + (opt-in, requires ``iprox_user``/``iprox_password``). """ records = self.list_files(px_id_or_url) if not records: logging.info("No Associated raw file URIs found in PX XML") return + + if protocol.lower() == "aspera": + from pridepy.download.iprox import IproxProvider + iprox_urls, rels = [], [] + for r in records: + loc = self.get_download_url(r, protocol) + host = (urlparse(loc).hostname or "").lower() + if host == "download.iprox.org": + iprox_urls.append(loc) + rels.append(r.get("relativePath")) + if not iprox_urls: + raise ValueError( + "Aspera requested but no iProX-hosted files found in this dataset." + ) + if len(iprox_urls) < len(records): + logging.warning( + "%d of %d file(s) are NOT hosted on iProX and were NOT " + "downloaded: --protocol aspera only handles iProX-hosted " + "files. Use the default HTTP transport (omit --protocol, " + "or pass --protocol ftp) to download the full set.", + len(records) - len(iprox_urls), + len(records), + ) + if flatten: + sources = [ + rel if rel else urlparse(url).path + for url, rel in zip(iprox_urls, rels) + ] + dest_rels: List[Optional[str]] = flatten_relative_paths(sources) + else: + dest_rels = rels + IproxProvider.aspera_download( + urls=iprox_urls, + output_folder=output_folder, + relative_paths=dest_rels, + user=iprox_user, + password=iprox_password, + skip_if_downloaded_already=skip_if_downloaded_already, + parallel_files=parallel_files, + ) + return + self.download_files( accession=px_id_or_url, records=records, output_folder=output_folder, skip_if_downloaded_already=skip_if_downloaded_already, - protocol="ftp", + protocol=protocol, flatten=flatten, + parallel_files=parallel_files, + download_threads=download_threads, ) diff --git a/pridepy/download/transport.py b/pridepy/download/transport.py index feae6ab..5f6882e 100644 --- a/pridepy/download/transport.py +++ b/pridepy/download/transport.py @@ -19,6 +19,11 @@ from pridepy.util.api_handling import Util +# Combined cap on parallel_files (-w) x download_threads (-t): each factor is +# independently clamped to 32, but nested they can reach 1024 concurrent HTTP +# connections. Clamp the product to keep peak connections reasonable. +MAX_TOTAL_HTTP_CONNECTIONS = 64 + def _safe_join(output_folder: str, relative_path: str) -> str: """Join ``output_folder`` with a dataset-relative path. @@ -845,6 +850,19 @@ def _rel(idx: int) -> Optional[str]: failed: List[str] = [] workers = max(1, min(parallel_files, len(http_urls))) + if workers * download_threads > MAX_TOTAL_HTTP_CONNECTIONS: + clamped_threads = max(1, MAX_TOTAL_HTTP_CONNECTIONS // workers) + logging.warning( + "parallel_files (%d) x download_threads (%d) = %d exceeds the " + "combined connection cap of %d; reducing download_threads to %d " + "to keep peak HTTP connections bounded.", + workers, + download_threads, + workers * download_threads, + MAX_TOTAL_HTTP_CONNECTIONS, + clamped_threads, + ) + download_threads = clamped_threads if workers > 1: logging.info( f"Downloading {len(http_urls)} HTTP(S) file(s) with {workers} parallel workers" diff --git a/pridepy/pridepy.py b/pridepy/pridepy.py index 17723cf..590e089 100644 --- a/pridepy/pridepy.py +++ b/pridepy/pridepy.py @@ -1,6 +1,9 @@ #!/usr/bin/env python3 import asyncio import logging +import os +from typing import Optional + import click from pridepy.download.client import Client as Files from pridepy.pdc import download_pdc_files as run_pdc_download @@ -60,6 +63,15 @@ def main(): type=click.IntRange(1, 32), help="Number of threads for each file download. Default is 1.", ) +@click.option( + "-w", + "--parallel-files", + "parallel_files", + default=1, + type=click.IntRange(1, 32), + help="Number of files to download in parallel (across-file concurrency). " + "Combine with -t/--threads (per-file segments). Default is 1.", +) @click.option( "--preserve-structure", is_flag=True, @@ -75,6 +87,7 @@ def download_all_public_raw_files( aspera_maximum_bandwidth: str = "50M", checksum_check: bool = False, download_threads: int = 1, + parallel_files: int = 1, preserve_structure: bool = False, ): """ @@ -88,6 +101,7 @@ def download_all_public_raw_files( aspera_maximum_bandwidth (str): Maximum bandwidth for Aspera protocol. Default is 100M. checksum_check (bool): Flag to download checksum file for the project. Default is False. download_threads (int): Number of threads for each file download. Default is 1. + parallel_files (int): Number of files to download in parallel. Default is 1. """ raw_files = Files() @@ -105,6 +119,7 @@ def download_all_public_raw_files( aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, download_threads=download_threads, + parallel_files=parallel_files, flatten=not preserve_structure, ) @@ -161,6 +176,15 @@ def download_all_public_raw_files( type=click.IntRange(1, 32), help="Number of threads for each file download. Default is 1.", ) +@click.option( + "-w", + "--parallel-files", + "parallel_files", + default=1, + type=click.IntRange(1, 32), + help="Number of files to download in parallel (across-file concurrency). " + "Combine with -t/--threads (per-file segments). Default is 1.", +) @click.option( "--preserve-structure", is_flag=True, @@ -177,6 +201,7 @@ def download_all_public_category_files( checksum_check: bool = False, category: str = "RAW", download_threads: int = 1, + parallel_files: int = 1, preserve_structure: bool = False, ): """ @@ -191,6 +216,7 @@ def download_all_public_category_files( checksum_check (bool): If True, downloads the checksum file for the project. category (str): Comma-separated categories of files to download (e.g. RAW or RAW,SEARCH). download_threads (int): Number of threads for each file download. Default is 1. + parallel_files (int): Number of files to download in parallel. Default is 1. """ valid_categories = {"RAW", "PEAK", "SEARCH", "RESULT", "SPECTRUM_LIBRARY", "OTHER", "FASTA"} @@ -218,6 +244,7 @@ def download_all_public_category_files( checksum_check=checksum_check, categories=categories, download_threads=download_threads, + parallel_files=parallel_files, flatten=not preserve_structure, ) @@ -332,6 +359,30 @@ def download_file_by_name( default=False, help="Skip the download if the file has already been downloaded.", ) +@click.option( + "-p", + "--protocol", + default="ftp", + type=PROTOCOL_CHOICES, + help="Protocol to use for download: ftp, aspera, globus, s3. Default is ftp with fallback enabled.", +) +@click.option( + "-t", + "--threads", + "download_threads", + default=1, + type=click.IntRange(1, 32), + help="Number of threads for each file download. Default is 1.", +) +@click.option( + "-w", + "--parallel-files", + "parallel_files", + default=1, + type=click.IntRange(1, 32), + help="Number of files to download in parallel (across-file concurrency). " + "Combine with -t/--threads (per-file segments). Default is 1.", +) @click.option( "--preserve-structure", is_flag=True, @@ -339,20 +390,45 @@ def download_file_by_name( help="Recreate the dataset's subdirectory layout under the output folder. " "By default files are downloaded flat into the output folder.", ) +@click.option( + "--iprox-user", + "iprox_user", + envvar="IPROX_USER", + default=None, + type=str, + help="iProX account username. Only used with --protocol aspera. The " + "password is never accepted as a command-line flag: it is read from " + "the IPROX_ASPERA_PASSWORD environment variable, or prompted for " + "securely (hidden input) if not set.", +) def download_px_raw_files( accession: str, output_folder: str, skip_if_downloaded_already: bool, + protocol: str = "ftp", + download_threads: int = 1, + parallel_files: int = 1, preserve_structure: bool = False, + iprox_user: Optional[str] = None, ): """CLI wrapper to download raw files via ProteomeXchange XML.""" files = Files() logging.info(f"PX accession/URL: {accession}") + + password = os.environ.get("IPROX_ASPERA_PASSWORD") + if protocol.lower() == "aspera" and not password: + password = click.prompt("iProX Aspera password", hide_input=True) + files.download_px_raw_files( accession, output_folder, skip_if_downloaded_already, flatten=not preserve_structure, + protocol=protocol, + download_threads=download_threads, + parallel_files=parallel_files, + iprox_user=iprox_user, + iprox_password=password, ) @@ -602,6 +678,15 @@ def _read_url_arguments(url_list_path, urls_csv=None): type=click.IntRange(1, 32), help="Number of threads for each file download. Default is 1.", ) +@click.option( + "-w", + "--parallel-files", + "parallel_files", + default=1, + type=click.IntRange(1, 32), + help="Number of files to download in parallel (across-file concurrency). " + "Combine with -t/--threads (per-file segments). Default is 1.", +) @click.option( "--preserve-structure", is_flag=True, @@ -619,6 +704,7 @@ def download_files_by_list( aspera_maximum_bandwidth, checksum_check, download_threads, + parallel_files: int = 1, preserve_structure: bool = False, ): """Download a named subset of files from a PRIDE project.""" @@ -635,6 +721,7 @@ def download_files_by_list( aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, download_threads=download_threads, + parallel_files=parallel_files, flatten=not preserve_structure, ) @@ -693,6 +780,15 @@ def download_files_by_list( type=click.IntRange(1, 32), help="Number of threads for each file download. Default is 1.", ) +@click.option( + "-w", + "--parallel-files", + "parallel_files", + default=1, + type=click.IntRange(1, 32), + help="Number of files to download in parallel (across-file concurrency). " + "Combine with -t/--threads (per-file segments). Default is 1.", +) def download_files_by_url( url_list_path, urls_csv, @@ -701,6 +797,7 @@ def download_files_by_url( protocol, checksum_check, download_threads, + parallel_files: int = 1, ): """Download files from raw URLs (http/https/ftp), dispatched by scheme.""" urls = _read_url_arguments(url_list_path, urls_csv) @@ -711,6 +808,7 @@ def download_files_by_url( skip_if_downloaded_already=skip_if_downloaded_already, protocol=protocol, download_threads=download_threads, + parallel_files=parallel_files, checksum_check=checksum_check, ) diff --git a/pridepy/tests/test_cli_flatten.py b/pridepy/tests/test_cli_flatten.py index 92c00df..01f887c 100644 --- a/pridepy/tests/test_cli_flatten.py +++ b/pridepy/tests/test_cli_flatten.py @@ -23,7 +23,23 @@ def test_download_all_public_raw_files_flattens_by_default(self): kwargs = files_cls.return_value.download_all_raw_files.call_args.kwargs assert kwargs["flatten"] is True assert kwargs["download_threads"] == 1 - assert "parallel_files" not in kwargs + assert kwargs["parallel_files"] == 1 + + def test_download_all_public_raw_files_parallel_files(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-all-public-raw-files", + "-a", + "MSV000012345", + "-o", + "/tmp/x", + "-w", + "8", + ] + ) + kwargs = files_cls.return_value.download_all_raw_files.call_args.kwargs + assert kwargs["parallel_files"] == 8 def test_download_all_public_raw_files_threads(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -40,7 +56,7 @@ def test_download_all_public_raw_files_threads(self): ) kwargs = files_cls.return_value.download_all_raw_files.call_args.kwargs assert kwargs["download_threads"] == 4 - assert "parallel_files" not in kwargs + assert kwargs["parallel_files"] == 1 def test_download_all_public_raw_files_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -74,7 +90,25 @@ def test_download_all_public_category_files_preserve_structure(self): kwargs = files_cls.return_value.download_all_category_files.call_args.kwargs assert kwargs["flatten"] is False assert kwargs["download_threads"] == 1 - assert "parallel_files" not in kwargs + assert kwargs["parallel_files"] == 1 + + def test_download_all_public_category_files_parallel_files(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-all-public-category-files", + "-a", + "MSV000012345", + "-o", + "/tmp/x", + "-c", + "RAW", + "-w", + "8", + ] + ) + kwargs = files_cls.return_value.download_all_category_files.call_args.kwargs + assert kwargs["parallel_files"] == 8 def test_download_files_by_list_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -93,7 +127,25 @@ def test_download_files_by_list_preserve_structure(self): kwargs = files_cls.return_value.download_files_by_list.call_args.kwargs assert kwargs["flatten"] is False assert kwargs["download_threads"] == 1 - assert "parallel_files" not in kwargs + assert kwargs["parallel_files"] == 1 + + def test_download_files_by_list_parallel_files(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-files-by-list", + "-a", + "MSV000012345", + "-o", + "/tmp/x", + "-f", + "a.raw", + "-w", + "8", + ] + ) + kwargs = files_cls.return_value.download_files_by_list.call_args.kwargs + assert kwargs["parallel_files"] == 8 def test_download_files_by_url_threads(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -110,7 +162,23 @@ def test_download_files_by_url_threads(self): ) kwargs = files_cls.download_files_by_url.call_args.kwargs assert kwargs["download_threads"] == 4 - assert "parallel_files" not in kwargs + assert kwargs["parallel_files"] == 1 + + def test_download_files_by_url_parallel_files(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-files-by-url", + "-u", + "https://example.org/a.raw", + "-o", + "/tmp/x", + "-w", + "8", + ] + ) + kwargs = files_cls.download_files_by_url.call_args.kwargs + assert kwargs["parallel_files"] == 8 def test_download_px_raw_files_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -126,3 +194,25 @@ def test_download_px_raw_files_preserve_structure(self): ) kwargs = files_cls.return_value.download_px_raw_files.call_args.kwargs assert kwargs["flatten"] is False + + def test_download_px_raw_files_protocol_threads_parallel(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-px-raw-files", + "-a", + "PXD000001", + "-o", + "/tmp/x", + "-w", + "8", + "-t", + "4", + "-p", + "ftp", + ] + ) + kwargs = files_cls.return_value.download_px_raw_files.call_args.kwargs + assert kwargs["parallel_files"] == 8 + assert kwargs["download_threads"] == 4 + assert kwargs["protocol"] == "ftp" diff --git a/pridepy/tests/test_download_resilience.py b/pridepy/tests/test_download_resilience.py index 532f624..c2b8cc7 100644 --- a/pridepy/tests/test_download_resilience.py +++ b/pridepy/tests/test_download_resilience.py @@ -247,6 +247,39 @@ def test_download_http_urls_raises_when_a_file_fails(self): max_retries=1, ) + def test_download_http_urls_clamps_combined_connection_cap(self): + """parallel_files x download_threads must be clamped to + transport.MAX_TOTAL_HTTP_CONNECTIONS, with a warning explaining why, + so e.g. -w 32 -t 32 doesn't open 1024 connections.""" + seen_threads = [] + + def fake_http_download_one(url, output_folder, skip_if_downloaded_already, + max_retries=3, position=0, relative_path=None, + download_threads=1): + seen_threads.append(download_threads) + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch.object( + transport, "_http_download_one", side_effect=fake_http_download_one + ): + with self.assertLogs(level="WARNING") as log_ctx: + transport.download_http_urls( + http_urls=[ + "https://example.org/a.raw", + "https://example.org/b.raw", + "https://example.org/c.raw", + ], + output_folder=tmp_dir, + skip_if_downloaded_already=False, + parallel_files=32, + download_threads=32, + ) + assert any("connection cap" in m.lower() for m in log_ctx.output) + workers = min(32, 3) + expected_threads = max(1, transport.MAX_TOTAL_HTTP_CONNECTIONS // workers) + assert workers * expected_threads <= transport.MAX_TOTAL_HTTP_CONNECTIONS + assert all(t == expected_threads for t in seen_threads) + def test_download_ftp_urls_raises_when_a_file_fails(self): """A failed FTP transfer must surface as an exception.""" with tempfile.TemporaryDirectory() as tmp_dir: diff --git a/pridepy/tests/test_iprox_aspera.py b/pridepy/tests/test_iprox_aspera.py new file mode 100644 index 0000000..6fa6bab --- /dev/null +++ b/pridepy/tests/test_iprox_aspera.py @@ -0,0 +1,300 @@ +"""iProX Aspera command construction + credential handling.""" +import os +import re +import subprocess +import tempfile +from unittest import TestCase +from unittest.mock import patch, MagicMock + +import pytest + +from pridepy.download.iprox import IproxProvider + + +class TestIproxAspera(TestCase): + def test_builds_ascp_source_and_env(self): + url = "http://download.iprox.org/IPX0003578000/IPX0003578001/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + run.return_value = MagicMock(returncode=0) + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=["IPX0003578001/a.raw"], + user="bob", + password="secret", + maximum_bandwidth="100M", + ) + args, kwargs = run.call_args + argv = args[0] + assert argv[0] == "/bin/ascp" + assert "33001" in argv + assert "bob@download.iprox.org:/data/iprox/IPX0003578000/IPX0003578001/a.raw" in argv + # password only via env, never argv + assert "secret" not in argv + assert kwargs["env"]["ASPERA_SCP_PASS"] == "secret" + + def test_missing_credentials_raises(self): + with pytest.raises(ValueError, match="credentials"): + IproxProvider.aspera_download( + urls=["http://download.iprox.org/IPX1/a.raw"], + output_folder="/tmp/x", + relative_paths=["a.raw"], + user=None, + password=None, + ) + + def test_failed_transfer_raises_runtime_error(self): + url = "http://download.iprox.org/IPX0003578000/IPX0003578001/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + run.side_effect = subprocess.CalledProcessError(1, ["ascp"]) + with pytest.raises(RuntimeError, match=re.escape(url)): + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=["IPX0003578001/a.raw"], + user="bob", + password="secret", + maximum_bandwidth="100M", + ) + + def test_skip_if_downloaded_already_skips_existing_file(self): + url = "http://download.iprox.org/IPX0003578000/IPX0003578001/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + dest_dir = os.path.join(tmp, "IPX0003578001") + os.makedirs(dest_dir, exist_ok=True) + dest_file = os.path.join(dest_dir, "a.raw") + with open(dest_file, "w") as f: + f.write("already here") + + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=["IPX0003578001/a.raw"], + user="bob", + password="secret", + maximum_bandwidth="100M", + skip_if_downloaded_already=True, + ) + run.assert_not_called() + + def test_skip_if_downloaded_already_does_not_skip_when_only_dir_exists(self): + """Regression: when relpath is missing, dest used to resolve to the + output DIRECTORY, so os.path.exists(dest) was always True and every + such file was wrongly skipped.""" + url = "http://download.iprox.org/IPX0003578000/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + run.return_value = MagicMock(returncode=0) + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=[None], + user="bob", + password="secret", + maximum_bandwidth="100M", + skip_if_downloaded_already=True, + ) + run.assert_called_once() + + def test_skip_if_downloaded_already_does_not_skip_zero_byte_file(self): + url = "http://download.iprox.org/IPX0003578000/IPX0003578001/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + dest_dir = os.path.join(tmp, "IPX0003578001") + os.makedirs(dest_dir, exist_ok=True) + dest_file = os.path.join(dest_dir, "a.raw") + open(dest_file, "w").close() # 0-byte partial file + run.return_value = MagicMock(returncode=0) + + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=["IPX0003578001/a.raw"], + user="bob", + password="secret", + maximum_bandwidth="100M", + skip_if_downloaded_already=True, + ) + run.assert_called_once() + + def test_traversal_relpath_does_not_escape_output_folder(self): + """A relativePath like '../../etc/x' must not write outside output_folder.""" + url = "http://download.iprox.org/IPX0003578000/a.raw" + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run") as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + run.return_value = MagicMock(returncode=0) + IproxProvider.aspera_download( + urls=[url], + output_folder=tmp, + relative_paths=["../../etc/x"], + user="bob", + password="secret", + maximum_bandwidth="100M", + ) + args, kwargs = run.call_args + argv = args[0] + dest = argv[-1] + out_abs = os.path.abspath(tmp) + dest_abs = os.path.abspath(dest) + assert dest_abs == out_abs or dest_abs.startswith(out_abs + os.sep) + + def test_parallel_files_downloads_all_and_aggregates_failures(self): + urls = [ + "http://download.iprox.org/IPX0003578000/a.raw", + "http://download.iprox.org/IPX0003578000/b.raw", + "http://download.iprox.org/IPX0003578000/c.raw", + ] + rels = ["a.raw", "b.raw", "c.raw"] + + def fake_run(argv, check, env): + if argv[-1].endswith("b.raw"): + raise subprocess.CalledProcessError(1, argv) + return MagicMock(returncode=0) + + with tempfile.TemporaryDirectory() as tmp, \ + patch("pridepy.download.iprox.subprocess.run", side_effect=fake_run) as run, \ + patch.object(IproxProvider, "_ascp_binary", return_value="/bin/ascp"): + with pytest.raises(RuntimeError, match="b.raw"): + IproxProvider.aspera_download( + urls=urls, + output_folder=tmp, + relative_paths=rels, + user="bob", + password="secret", + maximum_bandwidth="100M", + parallel_files=2, + ) + assert run.call_count == 3 + + +class TestPxAsperaRouting(TestCase): + def test_px_aspera_routes_to_iprox(self): + from pridepy.download.proteomexchange import ProteomeXchangeProvider + prov = ProteomeXchangeProvider() + rec = { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org/IPX1/IPX2/a.raw"} + ], + "relativePath": "IPX2/a.raw", + } + with patch.object(prov, "list_files", return_value=[rec]), \ + patch("pridepy.download.iprox.IproxProvider.aspera_download") as asp: + prov.download_from_accession_or_url( + "PXD000001", "/tmp/x", protocol="aspera", + iprox_user="bob", iprox_password="secret", + ) + asp.assert_called_once() + assert asp.call_args.kwargs["user"] == "bob" + + def test_px_aspera_rejects_spoofed_host(self): + """A URL on a lookalike host (substring match, not exact) must NOT be + routed to iProX Aspera.""" + from pridepy.download.proteomexchange import ProteomeXchangeProvider + prov = ProteomeXchangeProvider() + rec = { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org.evil.example/IPX1/a.raw"} + ], + "relativePath": "a.raw", + } + with patch.object(prov, "list_files", return_value=[rec]), \ + patch("pridepy.download.iprox.IproxProvider.aspera_download") as asp: + with pytest.raises(ValueError, match="no iProX-hosted files"): + prov.download_from_accession_or_url( + "PXD000001", "/tmp/x", protocol="aspera", + iprox_user="bob", iprox_password="secret", + ) + asp.assert_not_called() + + def test_px_aspera_warns_on_mixed_dataset(self): + """Non-iProX files in a mixed dataset are dropped from the aspera + transfer; a warning should be logged naming how many were skipped.""" + from pridepy.download.proteomexchange import ProteomeXchangeProvider + prov = ProteomeXchangeProvider() + records = [ + { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org/IPX1/a.raw"} + ], + "relativePath": "a.raw", + }, + { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "ftp://massive-ftp.ucsd.edu/MSV1/b.raw"} + ], + "relativePath": "b.raw", + }, + ] + with patch.object(prov, "list_files", return_value=records), \ + patch("pridepy.download.iprox.IproxProvider.aspera_download") as asp, \ + self.assertLogs(level="WARNING") as log_ctx: + prov.download_from_accession_or_url( + "PXD000001", "/tmp/x", protocol="aspera", + iprox_user="bob", iprox_password="secret", + ) + asp.assert_called_once() + assert any("not" in m.lower() and "iprox" in m.lower() for m in log_ctx.output) + + def test_px_aspera_flattens_relative_paths_by_default(self): + """Aspera branch should honor flatten=True like the HTTP/FTP path: + dataset subtree paths collapse to deduplicated basenames.""" + from pridepy.download.proteomexchange import ProteomeXchangeProvider + prov = ProteomeXchangeProvider() + records = [ + { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org/IPX1/run1/a.raw"} + ], + "relativePath": "run1/a.raw", + }, + { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org/IPX1/run2/a.raw"} + ], + "relativePath": "run2/a.raw", + }, + ] + with patch.object(prov, "list_files", return_value=records), \ + patch("pridepy.download.iprox.IproxProvider.aspera_download") as asp: + prov.download_from_accession_or_url( + "PXD000001", "/tmp/x", protocol="aspera", flatten=True, + iprox_user="bob", iprox_password="secret", + ) + asp.assert_called_once() + rels = asp.call_args.kwargs["relative_paths"] + # Flattened + de-duped basenames, no subdirectories preserved. + assert set(rels) == {"a.raw", "a_1.raw"} + + def test_px_aspera_preserves_structure_when_not_flattened(self): + from pridepy.download.proteomexchange import ProteomeXchangeProvider + prov = ProteomeXchangeProvider() + rec = { + "publicFileLocations": [ + {"name": "FTP Protocol", + "value": "http://download.iprox.org/IPX1/IPX2/a.raw"} + ], + "relativePath": "IPX2/a.raw", + } + with patch.object(prov, "list_files", return_value=[rec]), \ + patch("pridepy.download.iprox.IproxProvider.aspera_download") as asp: + prov.download_from_accession_or_url( + "PXD000001", "/tmp/x", protocol="aspera", flatten=False, + iprox_user="bob", iprox_password="secret", + ) + asp.assert_called_once() + assert asp.call_args.kwargs["relative_paths"] == ["IPX2/a.raw"]