diff --git a/README.md b/README.md index 423db45..f6e7933 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ pridepy --help | `download-file-by-name` | Download a single file (public or private) | | `download-files-by-list` | Download a named subset of files from a manifest/CSV | | `download-files-by-url` | Download files from raw `http`/`https`/`ftp` URLs | +| `download-pdc-files` | Download PDC/CPTAC files via PDC signed HTTPS URLs | | `download-px-raw-files` | Download RAW files resolved from a ProteomeXchange accession | | `list-private-files` | List files of a private project (needs credentials) | | `stream-files-metadata` | Stream file metadata (one project or all) to JSON | @@ -118,6 +119,9 @@ pridepy download-px-raw-files -a PXD039236 -o ./downloads/PXD039236 # Download a native MassIVE / JPOST / iProX dataset pridepy download-all-public-raw-files -a MSV000082297 -o ./downloads/MSV000082297 + +# Download PDC/CPTAC files for a study (single type, or per-row types via CSV) +pridepy download-pdc-files -a PDC000109 --file-type psm -o ./downloads/pdc ``` Full option tables and more examples are in [docs/usage.md](docs/usage.md). diff --git a/pridepy/download/base.py b/pridepy/download/base.py index 3382f08..b75005e 100644 --- a/pridepy/download/base.py +++ b/pridepy/download/base.py @@ -125,6 +125,7 @@ def download_all_raw( checksum_check: bool = False, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ) -> None: """Download all RAW files for the dataset.""" self.download_files( @@ -137,6 +138,7 @@ def download_all_raw( checksum_check=checksum_check, aspera_maximum_bandwidth=aspera_maximum_bandwidth, flatten=flatten, + download_threads=download_threads, ) def download_category( @@ -150,6 +152,7 @@ def download_category( checksum_check: bool = False, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ) -> None: """Download all files of the given categories for the dataset.""" self.download_files( @@ -162,6 +165,7 @@ def download_category( checksum_check=checksum_check, aspera_maximum_bandwidth=aspera_maximum_bandwidth, flatten=flatten, + download_threads=download_threads, ) def download_by_name( @@ -203,6 +207,7 @@ def download_by_filenames( checksum_check: bool = False, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ) -> None: """Download a subset of project files identified by a filename list. @@ -232,6 +237,7 @@ def download_by_filenames( checksum_check=checksum_check, aspera_maximum_bandwidth=aspera_maximum_bandwidth, flatten=flatten, + download_threads=download_threads, ) # ------------------------------------------------------------------ @@ -251,6 +257,7 @@ def download_files( username: Optional[str] = None, password: Optional[str] = None, flatten: bool = True, + download_threads: int = 1, ) -> None: """Partition record URLs by scheme and route to the matching transport. @@ -322,4 +329,5 @@ def download_files( skip_if_downloaded_already=skip_if_downloaded_already, parallel_files=parallel_files, relative_paths=http_relpaths, + download_threads=download_threads, ) diff --git a/pridepy/download/by_url.py b/pridepy/download/by_url.py index 0b4d03a..1a783d6 100644 --- a/pridepy/download/by_url.py +++ b/pridepy/download/by_url.py @@ -87,7 +87,13 @@ def _callback(data: bytes) -> None: ) -def _dispatch_url_scheme(parsed, target: str, protocol: str = "ftp", position: int = 0) -> None: +def _dispatch_url_scheme( + parsed, + target: str, + protocol: str = "ftp", + position: int = 0, + download_threads: int = 1, +) -> None: """Route a parsed URL to its protocol-specific downloader. ``protocol='globus'`` swaps the http/https single-connection streamer @@ -96,7 +102,11 @@ def _dispatch_url_scheme(parsed, target: str, protocol: str = "ftp", position: i """ scheme = (parsed.scheme or "").lower() if scheme in ("http", "https"): - if protocol == "globus": + if download_threads and download_threads > 1: + transport._multipart_download( + parsed.geturl(), target, threads=download_threads, position=position + ) + elif protocol == "globus": transport._parallel_download(parsed.geturl(), target, position=position) else: _http_download_url(parsed.geturl(), target) @@ -112,6 +122,7 @@ def _download_single_url( skip_if_exists: bool = False, protocol: str = "ftp", position: int = 0, + download_threads: int = 1, ) -> str: """Download one URL, dispatched by scheme; return the local file path.""" parsed = urlparse(url) @@ -128,7 +139,13 @@ def _download_single_url( return target try: - _dispatch_url_scheme(parsed, target, protocol, position=position) + _dispatch_url_scheme( + parsed, + target, + protocol, + position=position, + download_threads=download_threads, + ) except Exception: # Don't leave a truncated/partial file behind — a non-empty partial # would otherwise be wrongly skipped on the next run. @@ -149,6 +166,7 @@ def download_files_by_url( protocol: str = "ftp", parallel_files: int = 1, checksum_check: bool = False, + download_threads: int = 1, ) -> None: """Download files from a list of raw URLs, dispatched by URL scheme. @@ -173,39 +191,34 @@ def download_files_by_url( os.makedirs(output_folder, exist_ok=True) - parallel_files = min(parallel_files, 3, len(urls)) + workers = min(parallel_files, 3, len(urls)) failures: List[Tuple[str, str]] = [] - if parallel_files < 2: - for url in urls: + if workers > 1: + logging.info( + "Downloading %d URL(s) with %d parallel workers", + len(urls), workers, + ) + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = { + executor.submit( + _download_single_url, + url, + output_folder, + skip_if_downloaded_already, + protocol, + position=idx, + download_threads=download_threads, + ): url + for idx, url in enumerate(urls) + } + for future in as_completed(futures): + url = futures[future] try: - _download_single_url( - url, output_folder, skip_if_downloaded_already, protocol, - ) + future.result() except Exception as exc: # pylint: disable=broad-except logging.error("Failed to download %s: %s", url, exc) failures.append((url, str(exc))) - else: - logging.info( - "Downloading %d URL(s) with %d parallel workers", - len(urls), parallel_files, - ) - with ThreadPoolExecutor(max_workers=parallel_files) as executor: - futures = { - executor.submit( - _download_single_url, - url, output_folder, skip_if_downloaded_already, protocol, - position=idx, - ): url - for idx, url in enumerate(urls) - } - for future in as_completed(futures): - url = futures[future] - try: - future.result() - except Exception as exc: # pylint: disable=broad-except - logging.error("Failed to download %s: %s", url, exc) - failures.append((url, str(exc))) if failures: summary = ", ".join(f"{u} ({e})" for u, e in failures) diff --git a/pridepy/download/client.py b/pridepy/download/client.py index 979c650..5dd10f7 100644 --- a/pridepy/download/client.py +++ b/pridepy/download/client.py @@ -106,6 +106,7 @@ def download_http_urls( skip_if_downloaded_already: bool, parallel_files: int = 1, max_retries: int = 3, + download_threads: int = 1, ) -> None: """Shim — see :func:`pridepy.download.transport.download_http_urls`.""" return transport.download_http_urls( @@ -114,6 +115,7 @@ def download_http_urls( skip_if_downloaded_already=skip_if_downloaded_already, parallel_files=parallel_files, max_retries=max_retries, + download_threads=download_threads, ) # Accession-matcher convenience helpers (useful public API). @@ -187,6 +189,7 @@ def download_all_raw_files( checksum_check: bool = False, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ): """Download all RAW files for any registered provider.""" return registry.resolve(accession).download_all_raw( @@ -198,6 +201,7 @@ def download_all_raw_files( checksum_check=checksum_check, parallel_files=parallel_files, flatten=flatten, + download_threads=download_threads, ) def download_all_category_files( @@ -212,6 +216,7 @@ def download_all_category_files( category: str = None, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ): """Download all files of the given categories from a project.""" if categories is None: @@ -226,6 +231,7 @@ def download_all_category_files( checksum_check=checksum_check, parallel_files=parallel_files, flatten=flatten, + download_threads=download_threads, ) def download_file_by_name( @@ -269,6 +275,7 @@ def download_files_by_list( checksum_check: bool = False, parallel_files: int = 1, flatten: bool = True, + download_threads: int = 1, ) -> None: """Download a subset of project files identified by a filename list.""" return registry.resolve(accession).download_by_filenames( @@ -281,6 +288,7 @@ def download_files_by_list( checksum_check=checksum_check, parallel_files=parallel_files, flatten=flatten, + download_threads=download_threads, ) @staticmethod @@ -291,6 +299,7 @@ def download_files_by_url( protocol: str = "ftp", parallel_files: int = 1, checksum_check: bool = False, + download_threads: int = 1, ) -> None: """Delegate to :func:`pridepy.download.by_url.download_files_by_url`.""" return by_url.download_files_by_url( @@ -300,6 +309,7 @@ def download_files_by_url( protocol=protocol, parallel_files=parallel_files, checksum_check=checksum_check, + download_threads=download_threads, ) def download_px_raw_files( diff --git a/pridepy/download/pride.py b/pridepy/download/pride.py index 8ab017f..80f9348 100644 --- a/pridepy/download/pride.py +++ b/pridepy/download/pride.py @@ -317,8 +317,20 @@ def save_checksum_file(accession, output_folder): # ------------------------------------------------------------------ @staticmethod - def _globus_download_one(file, output_folder, skip_if_downloaded_already, max_retries=6, position=0): - """Download a single file via globus; used as a worker target.""" + def _globus_download_one( + file, + output_folder, + skip_if_downloaded_already, + max_retries=6, + position=0, + download_threads: int = 1, + ): + """Download a single file via globus; used as a worker target. + + When ``download_threads`` > 1 the file is fetched via parallel HTTP + Range segments (:func:`transport._multipart_download`); otherwise a + single-connection stream is used. + """ download_url = PrideProvider._get_download_url(file, "globus") new_file_path = PrideProvider.get_output_file_name(download_url, file, output_folder) @@ -328,7 +340,12 @@ def _globus_download_one(file, output_folder, skip_if_downloaded_already, max_re for attempt in range(1, max_retries + 1): try: - transport._parallel_download(download_url, new_file_path, position=position) + if download_threads and download_threads > 1: + transport._multipart_download( + download_url, new_file_path, threads=download_threads, position=position, + ) + else: + transport._parallel_download(download_url, new_file_path, position=position) return except Exception as e: logging.warning(f"Attempt {attempt}/{max_retries} failed for {file.get('fileName', '?')}: {e}") @@ -404,6 +421,7 @@ def download_files_from_globus( file_list_json: List[Dict], output_folder, skip_if_downloaded_already, parallel_files: int = 1, checksum_map: Optional[Dict[str, str]] = None, + download_threads: int = 1, ): """ Download files using globus transfer url with progress bar for each file. @@ -457,7 +475,7 @@ def download_files_from_globus( for file in files_to_download: try: PrideProvider._globus_download_one( - file, output_folder, False + file, output_folder, False, download_threads=download_threads, ) new_file_path = PrideProvider.get_output_file_name( PrideProvider._get_download_url(file, "globus"), file, output_folder @@ -472,8 +490,11 @@ def download_files_from_globus( futures = { executor.submit( PrideProvider._globus_download_one, - file, output_folder, False, + file, + output_folder, + False, position=idx, + download_threads=download_threads, ): file for idx, file in enumerate(files_to_download) } @@ -666,6 +687,7 @@ def _batch_download_by_protocol( aspera_maximum_bandwidth: str, parallel_files: int = 1, checksum_map: Optional[Dict[str, str]] = None, + download_threads: int = 1, ) -> None: """ Transfer a batch of files with one protocol, reusing a single @@ -706,6 +728,7 @@ def _batch_download_by_protocol( skip_if_downloaded_already=skip_if_downloaded_already, parallel_files=parallel_files, checksum_map=checksum_map or {}, + download_threads=download_threads, ) return if protocol == "s3": @@ -726,6 +749,7 @@ def _download_with_fallback( aspera_maximum_bandwidth: str, max_protocol_retries: int = 2, parallel_files: int = 1, + download_threads: int = 1, ) -> bool: """ Download one file by trying each protocol in sequence, validating @@ -749,6 +773,7 @@ def _download_with_fallback( skip_if_downloaded_already=False, aspera_maximum_bandwidth=aspera_maximum_bandwidth, parallel_files=parallel_files, + download_threads=download_threads, ) except Exception as error: logging.error( @@ -787,6 +812,7 @@ def download_files( username: Optional[str] = None, password: Optional[str] = None, flatten: bool = True, + download_threads: int = 1, ): """Override Provider.download_files with the multi-protocol orchestrator. @@ -808,6 +834,7 @@ def download_files( aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, parallel_files=parallel_files, + download_threads=download_threads, ) def download_by_name( @@ -887,6 +914,7 @@ def _download_files_batch( aspera_maximum_bandwidth: str = "100M", # Aspera maximum bandwidth checksum_check=False, parallel_files: int = 1, + download_threads: int = 1, ): """ Download files using the ftp, aspera, globus, or s3 transfer protocol. @@ -933,6 +961,7 @@ def _download_files_batch( aspera_maximum_bandwidth=aspera_maximum_bandwidth, parallel_files=parallel_files, checksum_map=checksum_map, + download_threads=download_threads, ) except Exception as exc: logging.warning( @@ -968,6 +997,7 @@ def _download_files_batch( expected_checksum=expected_checksum, aspera_maximum_bandwidth=aspera_maximum_bandwidth, parallel_files=parallel_files, + download_threads=download_threads, ) if not success: failed_files.append(file_record.get("fileName", "")) diff --git a/pridepy/download/transport.py b/pridepy/download/transport.py index 122e7de..feae6ab 100644 --- a/pridepy/download/transport.py +++ b/pridepy/download/transport.py @@ -83,22 +83,59 @@ def _open_ftp_connection(host: str, use_tls: bool, timeout: int = 30) -> FTP: return ftp -def _walk_ftp_tree(ftp: FTP, remote_dir: str) -> List[str]: +# Emit a progress line every this many directories while walking a remote +# tree. Large deposits (e.g. a MassIVE timsTOF dataset with thousands of .d +# directories, each needing its own TLS data connection to list) can take +# many minutes to enumerate; without progress the caller looks hung. +_WALK_PROGRESS_EVERY_DIRS = 100 + + +def _walk_ftp_tree( + ftp: FTP, remote_dir: str, _progress: Optional[dict] = None +) -> List[str]: """ Recursively list files under a remote FTP directory. + + Emits an INFO progress heartbeat every ``_WALK_PROGRESS_EVERY_DIRS`` + directories, plus a final summary, so enumerating a large deposit does + not look like a hang. ``_progress`` is internal recursion state; callers + invoke this with ``(ftp, remote_dir)`` only. """ import posixpath + + top_level = _progress is None + if top_level: + _progress = {"dirs": 0, "files": 0} + + def _note_dir_listed() -> None: + _progress["dirs"] += 1 + if _progress["dirs"] % _WALK_PROGRESS_EVERY_DIRS == 0: + logging.info( + "Listing remote tree: %d directories scanned, " + "%d files found so far...", + _progress["dirs"], + _progress["files"], + ) + file_paths: List[str] = [] try: entries = list(ftp.mlsd(remote_dir)) + _note_dir_listed() for name, facts in entries: if name in {".", ".."}: continue child_path = posixpath.join(remote_dir.rstrip("/"), name) if facts.get("type") == "dir": - file_paths.extend(_walk_ftp_tree(ftp, child_path)) + file_paths.extend(_walk_ftp_tree(ftp, child_path, _progress)) elif facts.get("type") == "file": file_paths.append(child_path) + _progress["files"] += 1 + if top_level: + logging.info( + "Listing remote tree complete: %d directories, %d files.", + _progress["dirs"], + _progress["files"], + ) return file_paths except (AttributeError, ftplib.error_perm): pass @@ -108,6 +145,7 @@ def _walk_ftp_tree(ftp: FTP, remote_dir: str) -> List[str]: try: ftp.cwd(remote_dir) ftp.retrlines("LIST", listing.append) + _note_dir_listed() for entry in listing: parts = entry.split(maxsplit=8) if len(parts) < 9: @@ -117,11 +155,18 @@ def _walk_ftp_tree(ftp: FTP, remote_dir: str) -> List[str]: continue child_path = posixpath.join(remote_dir.rstrip("/"), name) if entry.startswith("d"): - file_paths.extend(_walk_ftp_tree(ftp, child_path)) + file_paths.extend(_walk_ftp_tree(ftp, child_path, _progress)) else: file_paths.append(child_path) + _progress["files"] += 1 finally: ftp.cwd(current_dir) + if top_level: + logging.info( + "Listing remote tree complete: %d directories, %d files.", + _progress["dirs"], + _progress["files"], + ) return file_paths @@ -580,6 +625,142 @@ def _parallel_download(url, file_path, position=0): ) +def _download_range(url, file_path, start, end, pbar, max_retries=3): + """Download a byte range directly into the target file using seek. + + On transient failures (e.g. ``IncompleteRead``) the segment resumes + from the last successfully written byte using a fresh Range request + for the remaining bytes, instead of restarting the whole segment. + ``pbar`` only accumulates bytes that were actually written, so the + progress bar stays in sync with on-disk state across retries. + """ + cursor = start + for attempt in range(1, max_retries + 1): + try: + session = Util.create_session_with_retries() + headers = {"Range": f"bytes={cursor}-{end}"} + with session.get(url, headers=headers, stream=True, timeout=(15, 15)) as r: + r.raise_for_status() + if r.status_code != 206: + raise RuntimeError(f"Server did not honor Range request: {r.status_code}") + content_range = r.headers.get("Content-Range", "") + if not content_range.lower().startswith(f"bytes {cursor}-{end}/"): + raise RuntimeError(f"Unexpected Content-Range header: {content_range}") + with open(file_path, "r+b") as f: + f.seek(cursor) + for chunk in r.iter_content(chunk_size=8 * 1024 * 1024): + if chunk: + f.write(chunk) + cursor += len(chunk) + pbar.update(len(chunk)) + return + except (requests.RequestException, RuntimeError, OSError) as exc: + logging.warning( + f"Range {start}-{end} attempt {attempt}/{max_retries} failed " + f"(resumed at {cursor}/{end + 1}): {exc}" + ) + if attempt >= max_retries: + raise + time.sleep(2 * attempt) + + +def _can_use_multipart(threads, accept_ranges, total_size, min_size_bytes): + return ( + threads >= 2 + and accept_ranges == "bytes" + and total_size > 0 + and total_size >= min_size_bytes + ) + + +def _read_http_download_metadata(url): + session = Util.create_session_with_retries() + head = session.head(url, timeout=(30, 30), allow_redirects=True) + head.raise_for_status() + total_size = int(head.headers.get("content-length", 0)) + accept_ranges = head.headers.get("accept-ranges", "none").strip().lower() + return total_size, accept_ranges + + +def _prepare_multipart_target(file_path, total_size): + if os.path.exists(file_path) and os.path.getsize(file_path) == total_size: + logging.info("File already complete: %s", file_path) + return False + + with open(file_path, "wb") as pre: + pre.truncate(total_size) + return True + + +def _build_download_ranges(total_size, threads): + part_size = total_size // threads + ranges = [] + for index in range(threads): + start = index * part_size + end = total_size - 1 if index == threads - 1 else (start + part_size - 1) + ranges.append((start, end)) + return ranges + + +def _download_multipart_ranges(url, file_path, ranges, total_size, position): + try: + with tqdm( + total=total_size, + unit="B", + unit_scale=True, + desc=file_path, + position=position, + leave=True, + ) as pbar: + with ThreadPoolExecutor(max_workers=len(ranges)) as executor: + futures = [ + executor.submit(_download_range, url, file_path, start, end, pbar) + for start, end in ranges + ] + for future in as_completed(futures): + future.result() + except (requests.RequestException, RuntimeError, OSError): + if os.path.exists(file_path): + os.remove(file_path) + raise + + +def _multipart_download(url, file_path, threads=8, position=0, min_size_bytes=10 * 1024 * 1024): + """Download a single file via parallel HTTP Range requests. + + Falls back to :func:`_parallel_download` if the server does not advertise + ``Accept-Ranges: bytes``, the total size is unknown, or the file is + smaller than ``min_size_bytes`` (default 10 MB). + + Threads are clamped to ``[1, 32]``. Resume is best-effort: if an existing + file matches the expected total size, it is treated as complete. + """ + parent = os.path.dirname(file_path) + if parent: + os.makedirs(parent, exist_ok=True) + + threads = max(1, min(32, int(threads or 1))) + try: + total_size, accept_ranges = _read_http_download_metadata(url) + except (requests.RequestException, ValueError) as exc: + logging.info( + "HEAD failed for multipart, falling back to single stream: %s", + exc, + ) + _parallel_download(url, file_path, position=position) + return + + if not _can_use_multipart(threads, accept_ranges, total_size, min_size_bytes): + _parallel_download(url, file_path, position=position) + return + + if not _prepare_multipart_target(file_path, total_size): + return + + ranges = _build_download_ranges(total_size, threads) + _download_multipart_ranges(url, file_path, ranges, total_size, position) + + def _http_download_one( url: str, output_folder: str, @@ -587,6 +768,7 @@ def _http_download_one( max_retries: int = 3, position: int = 0, relative_path: Optional[str] = None, + download_threads: int = 1, ) -> None: """ Download a single HTTP(S) URL with HEAD-then-Range resume and retry. @@ -597,6 +779,10 @@ def _http_download_one( ``relative_path`` (when given) is the dataset-relative destination, so files keep their collection layout instead of being flattened to the URL basename. + + When ``download_threads`` > 1 a single file is split into parallel HTTP + Range segments via :func:`_multipart_download`; otherwise a single + connection stream is used. """ local_path = _dest_path(output_folder, urlparse(url).path, relative_path) if skip_if_downloaded_already and os.path.exists(local_path): @@ -605,7 +791,10 @@ def _http_download_one( last_error: Optional[Exception] = None for attempt in range(1, max_retries + 1): try: - _parallel_download(url, local_path, position=position) + if download_threads and download_threads > 1: + _multipart_download(url, local_path, threads=download_threads, position=position) + else: + _parallel_download(url, local_path, position=position) logging.info(f"Successfully downloaded {local_path}") return except Exception as e: @@ -625,6 +814,7 @@ def download_http_urls( parallel_files: int = 1, max_retries: int = 3, relative_paths: Optional[List[str]] = None, + download_threads: int = 1, ) -> None: """ Download a list of HTTP(S) URLs with HEAD-then-Range resume, per-file @@ -635,6 +825,9 @@ def download_http_urls( ``requests`` session is opened inside ``_parallel_download``) so the only shared resource is the output directory. + When ``download_threads`` > 1, each individual file is split into that + many parallel HTTP Range segments for higher single-file throughput. + :param relative_paths: Optional per-URL dataset-relative destination paths (parallel to ``http_urls``); see :func:`download_ftp_urls`. :raises RuntimeError: after attempting every URL, if one or more failed. @@ -666,6 +859,7 @@ def _rel(idx: int) -> Optional[str]: max_retries, idx, _rel(idx), + download_threads, ): url for idx, url in enumerate(http_urls) } @@ -685,6 +879,7 @@ def _rel(idx: int) -> Optional[str]: skip_if_downloaded_already, max_retries, relative_path=_rel(idx), + download_threads=download_threads, ) except Exception as e: logging.error(f"HTTP download failed for {url}: {e}") diff --git a/pridepy/pdc/__init__.py b/pridepy/pdc/__init__.py new file mode 100644 index 0000000..05ae0cd --- /dev/null +++ b/pridepy/pdc/__init__.py @@ -0,0 +1,20 @@ +from pridepy.pdc.client import ( + PDCDownloadRequest, + PDCFile, + fetch_study_files, + parse_accessions, + parse_download_requests, + refresh_signed_url, +) +from pridepy.pdc.downloader import PDCDownloadStats, download_pdc_files + +__all__ = [ + "PDCDownloadRequest", + "PDCDownloadStats", + "PDCFile", + "download_pdc_files", + "fetch_study_files", + "parse_accessions", + "parse_download_requests", + "refresh_signed_url", +] diff --git a/pridepy/pdc/client.py b/pridepy/pdc/client.py new file mode 100644 index 0000000..73a7edb --- /dev/null +++ b/pridepy/pdc/client.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import csv +import logging +import os +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Tuple + +from pridepy.util.api_handling import Util + +PDC_API = "https://pdc.cancer.gov/graphql" +USER_AGENT = ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" +) +FETCH_TIMEOUT_SECONDS = 600 + +LOGGER = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class PDCFileTypeFilter: + data_category: Optional[str] + file_format: Optional[str] + filename_suffix: Optional[str] = None + + +@dataclass(frozen=True) +class PDCFile: + study_id: str + file_id: str + file_name: str + file_format: Optional[str] + file_size: int + data_category: Optional[str] + file_type: Optional[str] + file_location: Optional[str] + md5sum: Optional[str] + url: str + + +@dataclass(frozen=True) +class PDCDownloadRequest: + study_id: str + file_type: str + + +PDC_FILE_TYPE_FILTERS: Dict[str, PDCFileTypeFilter] = { + "mzid": PDCFileTypeFilter("Peptide Spectral Matches", "mzIdentML"), + "psm": PDCFileTypeFilter("Peptide Spectral Matches", "tsv", ".psm"), + "raw": PDCFileTypeFilter("Raw Mass Spectra", "vendor-specific"), + "mzml": PDCFileTypeFilter("Processed Mass Spectra", "mzML"), +} + +PDC_FILE_TYPE_COLUMNS = ("file-type", "file_type", "filetype") + + +FILES_PER_STUDY_QUERY = """ +query FilesPerStudy($studyId: String!) { + filesPerStudy(pdc_study_id: $studyId, acceptDUA: true) { + file_id + pdc_study_id + file_name + file_format + file_size + data_category + file_type + file_location + md5sum + signedUrl { + url + } + } +} +""" + +_refresh_lock = threading.Lock() + + +def split_accession_text(text: str) -> List[str]: + items: List[str] = [] + for raw_item in text.replace(",", "\n").splitlines(): + item = raw_item.strip() + if item and not item.startswith("#"): + items.append(item) + return items + + +def dedupe_keep_order(items: Iterable[str]) -> List[str]: + return list(dict.fromkeys(item for item in items if item)) + + +def _study_id_column(columns: List[str], path: Path) -> str: + if "pdc_id" in columns: + return "pdc_id" + if "pdc_study_id" in columns: + return "pdc_study_id" + raise ValueError(f"CSV must contain pdc_id or pdc_study_id column: {path}") + + +def _file_type_column(columns: List[str]) -> Optional[str]: + for column in PDC_FILE_TYPE_COLUMNS: + if column in columns: + return column + return None + + +def _read_accession_rows_from_csv(path: Path) -> Tuple[List[Tuple[str, Optional[str]]], bool]: + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + columns = reader.fieldnames or [] + study_column = _study_id_column(columns, path) + file_type_column = _file_type_column(columns) + + rows = [] + for row in reader: + study_id = str(row.get(study_column) or "").strip() + if not study_id: + continue + row_file_type = None + if file_type_column: + row_file_type = str(row.get(file_type_column) or "").strip() or None + rows.append((study_id, row_file_type)) + + if not rows: + raise ValueError(f"Download accession list is empty: {path}") + return rows, file_type_column is not None + + +def _read_accessions_from_csv(path: Path) -> List[str]: + rows, _has_file_type_column = _read_accession_rows_from_csv(path) + return dedupe_keep_order(study_id for study_id, _file_type in rows) + + +def parse_accessions(accession: str) -> List[str]: + if not accession or not accession.strip(): + raise ValueError("--accession must not be empty") + + source_path = Path(accession).expanduser() + if source_path.exists(): + if source_path.suffix.lower() != ".csv": + raise ValueError(f"Only CSV accession files are supported: {source_path}") + return _read_accessions_from_csv(source_path) + + if source_path.suffix.lower() == ".csv" or os.path.sep in accession: + raise ValueError(f"Accession CSV not found: {source_path}") + + accessions = dedupe_keep_order(split_accession_text(accession)) + if not accessions: + raise ValueError("--accession did not contain any PDC study ID") + return accessions + + +def normalize_file_type(file_type: str) -> str: + normalized = str(file_type or "").strip().lower() + if normalized not in PDC_FILE_TYPE_FILTERS: + valid = ", ".join(sorted(PDC_FILE_TYPE_FILTERS)) + raise ValueError(f"Unsupported PDC file type: {file_type}. Valid values: {valid}") + return normalized + + +def parse_download_requests(accession: str, file_type: Optional[str] = None) -> List[PDCDownloadRequest]: + if not accession or not accession.strip(): + raise ValueError("--accession must not be empty") + + command_file_type = normalize_file_type(file_type) if file_type else None + source_path = Path(accession).expanduser() + + if source_path.exists(): + if source_path.suffix.lower() != ".csv": + raise ValueError(f"Only CSV accession files are supported: {source_path}") + + rows, has_csv_file_type = _read_accession_rows_from_csv(source_path) + if has_csv_file_type and command_file_type: + LOGGER.warning( + "CSV contains file-type column; --file-type=%s overrides CSV file-type values", + command_file_type, + ) + + requests = [] + for study_id, csv_file_type in rows: + request_file_type = command_file_type + if request_file_type is None: + if csv_file_type is None: + raise ValueError( + "CSV rows must contain file-type/filetype when --file-type is not set: " + f"{study_id}" + ) + request_file_type = normalize_file_type(csv_file_type) + requests.append(PDCDownloadRequest(study_id, request_file_type)) + return list(dict.fromkeys(requests)) + + study_ids = parse_accessions(accession) + if command_file_type is None: + raise ValueError("--file-type is required unless --accession is a CSV with file-type/filetype column") + return [PDCDownloadRequest(study_id, command_file_type) for study_id in study_ids] + + +def get_file_type_filter(file_type: str) -> PDCFileTypeFilter: + return PDC_FILE_TYPE_FILTERS[normalize_file_type(file_type)] + + +def entry_matches_file_type(entry: Dict, file_type: str) -> bool: + file_filter = get_file_type_filter(file_type) + if file_filter.data_category is not None and entry.get("data_category") != file_filter.data_category: + return False + if file_filter.file_format is not None and entry.get("file_format") != file_filter.file_format: + return False + if file_filter.filename_suffix is not None: + file_name = str(entry.get("file_name") or "") + if not file_name.endswith(file_filter.filename_suffix): + return False + return True + + +def _safe_int(value) -> int: + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + +def normalize_pdc_file(entry: Dict, fallback_study_id: str) -> Optional[PDCFile]: + file_name = str(entry.get("file_name") or "").strip() + signed_url = (entry.get("signedUrl") or {}).get("url") + if not file_name or not signed_url: + return None + + md5sum = entry.get("md5sum") + return PDCFile( + study_id=str(entry.get("pdc_study_id") or fallback_study_id), + file_id=str(entry.get("file_id") or ""), + file_name=file_name, + file_format=entry.get("file_format"), + file_size=_safe_int(entry.get("file_size")), + data_category=entry.get("data_category"), + file_type=entry.get("file_type"), + file_location=entry.get("file_location"), + md5sum=str(md5sum).lower() if md5sum else None, + url=str(signed_url), + ) + + +def post_graphql(query: str, variables: Dict, session=None) -> Dict: + active_session = session or Util.create_session_with_retries() + response = active_session.post( + PDC_API, + json={"query": query, "variables": variables}, + headers={"Content-Type": "application/json", "User-Agent": USER_AGENT}, + timeout=FETCH_TIMEOUT_SECONDS, + ) + response.raise_for_status() + payload = response.json() + if payload.get("errors"): + raise RuntimeError(f"PDC GraphQL returned errors: {payload['errors']}") + return payload + + +def fetch_study_files(study_id: str, file_type: str, session=None) -> List[PDCFile]: + payload = post_graphql(FILES_PER_STUDY_QUERY, {"studyId": study_id}, session=session) + raw_files = payload.get("data", {}).get("filesPerStudy", []) or [] + files: List[PDCFile] = [] + for entry in raw_files: + if not entry_matches_file_type(entry, file_type): + continue + pdc_file = normalize_pdc_file(entry, study_id) + if pdc_file is None: + LOGGER.warning("Skipping PDC file without name or signed URL in study %s", study_id) + continue + files.append(pdc_file) + return files + + +def refresh_signed_url(study_id: str, file_name: str, file_type: str, session=None) -> Optional[str]: + with _refresh_lock: + for pdc_file in fetch_study_files(study_id, file_type, session=session): + if pdc_file.file_name == file_name: + return pdc_file.url + return None diff --git a/pridepy/pdc/downloader.py b/pridepy/pdc/downloader.py new file mode 100644 index 0000000..ecb021c --- /dev/null +++ b/pridepy/pdc/downloader.py @@ -0,0 +1,386 @@ +from __future__ import annotations + +import logging +import os +import time +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Callable, List, Optional, Tuple + +from pridepy.download import transport +from pridepy.download.util import compute_md5 +from pridepy.pdc.client import PDCFile, fetch_study_files, parse_download_requests, refresh_signed_url + +LOGGER = logging.getLogger(__name__) + +FetchFiles = Callable[[str, str], List[PDCFile]] +RefreshUrl = Callable[[str, str, str], Optional[str]] + + +@dataclass +class PDCDownloadStats: + studies: int = 0 + total_files: int = 0 + downloaded: int = 0 + skipped: int = 0 + failed: int = 0 + + +@dataclass(frozen=True) +class PDCTransferResult: + success: bool + message: str + http_status: Optional[int] = None + + +@dataclass(frozen=True) +class PDCDownloadFailure: + study_id: str + file_name: str + message: str + file_type: Optional[str] = None + pdc_file: Optional[PDCFile] = None + http_status: Optional[int] = None + + +def validate_pdc_file(path: Path, pdc_file: PDCFile, checksum_check: bool) -> Tuple[bool, str]: + if not path.exists(): + return False, "file does not exist" + + actual_size = path.stat().st_size + if pdc_file.file_size > 0 and actual_size != pdc_file.file_size: + return False, f"size mismatch (expected={pdc_file.file_size}, actual={actual_size})" + if pdc_file.file_size <= 0 and actual_size == 0: + return False, "file is empty and PDC did not provide a positive file_size" + + if checksum_check: + if pdc_file.md5sum: + actual_md5 = compute_md5(str(path)) + if actual_md5.lower() != pdc_file.md5sum.lower(): + return False, ( + f"checksum mismatch (expected={pdc_file.md5sum.lower()}, " + f"actual={actual_md5.lower()})" + ) + else: + LOGGER.warning( + "PDC did not provide md5sum for %s/%s; falling back to size validation", + pdc_file.study_id, + pdc_file.file_name, + ) + + return True, "ok" + + +def _is_download_complete(path: Path, pdc_file: PDCFile, checksum_check: bool) -> bool: + valid, reason = validate_pdc_file(path, pdc_file, checksum_check) + if not valid: + LOGGER.info("Existing file is incomplete: %s (%s)", path, reason) + return valid + + +def _http_status_from_exception(exc: Exception) -> Optional[int]: + response = getattr(exc, "response", None) + status_code = getattr(response, "status_code", None) + if status_code is not None: + try: + return int(status_code) + except (TypeError, ValueError): + return None + if "403" in str(exc): + return 403 + return None + + +def _part_path(target: Path) -> Path: + return target.with_name(target.name + ".part") + + +def _remove_file(path: Path) -> None: + try: + path.unlink(missing_ok=True) + except FileNotFoundError: + return + + +def _download_to_part(pdc_file: PDCFile, target: Path, checksum_check: bool, download_threads: int) -> PDCTransferResult: + part = _part_path(target) + _remove_file(part) + target.parent.mkdir(parents=True, exist_ok=True) + + try: + LOGGER.info("Downloading %s/%s", pdc_file.study_id, pdc_file.file_name) + if download_threads > 1: + transport._multipart_download(pdc_file.url, str(part), threads=download_threads) + else: + transport._parallel_download(pdc_file.url, str(part)) + + valid, reason = validate_pdc_file(part, pdc_file, checksum_check) + if not valid: + _remove_file(part) + return PDCTransferResult(False, reason) + + os.replace(part, target) + return PDCTransferResult(True, "ok") + except Exception as exc: # pylint: disable=broad-except + _remove_file(part) + return PDCTransferResult(False, str(exc), _http_status_from_exception(exc)) + + +def _download_with_retries( + pdc_file: PDCFile, + target: Path, + checksum_check: bool, + download_threads: int, + file_type: str, + refresh_url: RefreshUrl, + refresh_on_403: bool, + max_attempts: int, +) -> PDCTransferResult: + current_file = pdc_file + result = PDCTransferResult(False, "not attempted") + for attempt in range(1, max_attempts + 1): + if attempt > 1: + LOGGER.info( + "Retrying %s/%s (%d/%d)", + current_file.study_id, + current_file.file_name, + attempt, + max_attempts, + ) + + result = _download_to_part(current_file, target, checksum_check, download_threads) + if result.success: + return result + + if result.http_status == 403 and refresh_on_403: + fresh_url = refresh_url(current_file.study_id, current_file.file_name, file_type) + if fresh_url: + current_file = replace(current_file, url=fresh_url) + LOGGER.info("Refreshed signed URL for %s/%s", current_file.study_id, current_file.file_name) + else: + LOGGER.warning( + "Could not refresh signed URL for %s/%s", + current_file.study_id, + current_file.file_name, + ) + + if attempt < max_attempts: + time.sleep(min(60, 2 ** attempt)) + + return result + + +def _target_path(output_folder: Path, pdc_file: PDCFile) -> Path: + return output_folder / pdc_file.study_id / pdc_file.file_name + + +def _write_failed_files(output_folder: Path, failures: List[PDCDownloadFailure]) -> Path: + failed_log = output_folder / "failed_files.txt" + with failed_log.open("w", encoding="utf-8") as handle: + handle.write("# PDC download failures\n") + handle.write("# study_id\tfile_name\tmessage\n") + for failure in failures: + handle.write(f"{failure.study_id}\t{failure.file_name}\t{failure.message}\n") + return failed_log + + +def _failure_from_result(pdc_file: PDCFile, result: PDCTransferResult, file_type: str) -> PDCDownloadFailure: + return PDCDownloadFailure( + study_id=pdc_file.study_id, + file_name=pdc_file.file_name, + message=result.message, + file_type=file_type, + pdc_file=pdc_file, + http_status=result.http_status, + ) + + +def _fetch_study_file_list( + study_id: str, + file_type: str, + fetch_files: FetchFiles, +) -> Tuple[List[PDCFile], Optional[PDCDownloadFailure]]: + try: + return fetch_files(study_id, file_type), None + except Exception as exc: # pylint: disable=broad-except + return [], PDCDownloadFailure(study_id, "", str(exc), file_type=file_type) + + +def _download_study_files( + study_id: str, + study_files: List[PDCFile], + output_path: Path, + options: "PDCDownloadOptions", + stats: PDCDownloadStats, +) -> List[PDCDownloadFailure]: + failures: List[PDCDownloadFailure] = [] + for file_index, pdc_file in enumerate(study_files, 1): + target = _target_path(output_path, pdc_file) + LOGGER.info("[%s %d/%d] %s", study_id, file_index, len(study_files), pdc_file.file_name) + + if options.skip_existing and _is_download_complete(target, pdc_file, options.checksum_check): + stats.skipped += 1 + LOGGER.info("Skipped existing file: %s", target) + continue + + result = _download_with_retries( + pdc_file, + target, + options.checksum_check, + options.download_threads, + options.file_type, + options.refresh_url, + refresh_on_403=False, + max_attempts=1, + ) + if result.success: + stats.downloaded += 1 + else: + failures.append(_failure_from_result(pdc_file, result, options.file_type)) + LOGGER.error("Failed %s/%s: %s", pdc_file.study_id, pdc_file.file_name, result.message) + return failures + + +@dataclass(frozen=True) +class PDCDownloadOptions: + file_type: str + skip_existing: bool + checksum_check: bool + download_threads: int + refresh_url: RefreshUrl + + +def _refresh_before_retry( + failure: PDCDownloadFailure, + file_type: str, + refresh_url: RefreshUrl, +) -> PDCFile: + retry_file = failure.pdc_file + if retry_file is None: + raise ValueError("Cannot retry metadata failure") + if failure.http_status != 403: + return retry_file + + fresh_url = refresh_url(retry_file.study_id, retry_file.file_name, file_type) + if fresh_url: + retry_file = replace(retry_file, url=fresh_url) + LOGGER.info("Refreshed signed URL for %s/%s", retry_file.study_id, retry_file.file_name) + return retry_file + + +def _retry_failed_downloads( + failures: List[PDCDownloadFailure], + output_path: Path, + options: PDCDownloadOptions, + stats: PDCDownloadStats, +) -> List[PDCDownloadFailure]: + retry_failures: List[PDCDownloadFailure] = [] + LOGGER.info("Retrying %d failed PDC file(s)", len(failures)) + for failure in failures: + if failure.pdc_file is None: + retry_failures.append(failure) + continue + + retry_file_type = failure.file_type or options.file_type + retry_file = _refresh_before_retry(failure, retry_file_type, options.refresh_url) + result = _download_with_retries( + retry_file, + _target_path(output_path, retry_file), + options.checksum_check, + options.download_threads, + retry_file_type, + options.refresh_url, + refresh_on_403=True, + max_attempts=3, + ) + if result.success: + stats.downloaded += 1 + LOGGER.info("Recovered %s/%s", failure.study_id, failure.file_name) + else: + retry_failures.append(_failure_from_result(retry_file, result, retry_file_type)) + return retry_failures + + +def download_pdc_files( + accession: str, + file_type: Optional[str], + output_folder: str, + skip_if_downloaded_already: bool = False, + checksum_check: bool = True, + download_threads: int = 1, + retry: bool = False, + fetch_files: FetchFiles = fetch_study_files, + refresh_url: RefreshUrl = refresh_signed_url, +) -> PDCDownloadStats: + requests = parse_download_requests(accession, file_type) + output_path = Path(output_folder).expanduser() + output_path.mkdir(parents=True, exist_ok=True) + active_download_threads = max(1, min(32, int(download_threads or 1))) + retry_options = PDCDownloadOptions( + file_type=requests[0].file_type, + skip_existing=skip_if_downloaded_already, + checksum_check=checksum_check, + download_threads=active_download_threads, + refresh_url=refresh_url, + ) + + stats = PDCDownloadStats(studies=len(dict.fromkeys(request.study_id for request in requests))) + failures: List[PDCDownloadFailure] = [] + empty_match_failures: List[PDCDownloadFailure] = [] + + for request_index, request in enumerate(requests, 1): + LOGGER.info( + "[%d/%d] Fetching PDC files for %s file_type=%s", + request_index, + len(requests), + request.study_id, + request.file_type, + ) + study_files, metadata_failure = _fetch_study_file_list(request.study_id, request.file_type, fetch_files) + if metadata_failure: + failures.append(metadata_failure) + continue + + stats.total_files += len(study_files) + if not study_files: + message = f"No PDC files matched file_type={request.file_type}" + LOGGER.warning("%s for %s", message, request.study_id) + empty_match_failures.append( + PDCDownloadFailure( + request.study_id, + f"<{request.file_type}>", + message, + file_type=request.file_type, + ) + ) + continue + + options = PDCDownloadOptions( + file_type=request.file_type, + skip_existing=skip_if_downloaded_already, + checksum_check=checksum_check, + download_threads=active_download_threads, + refresh_url=refresh_url, + ) + failures.extend(_download_study_files(request.study_id, study_files, output_path, options, stats)) + + if failures and retry: + failures = _retry_failed_downloads(failures, output_path, retry_options, stats) + + if stats.downloaded == 0 and stats.skipped == 0 and empty_match_failures: + failures.extend(empty_match_failures) + + stats.failed = len(failures) + if failures: + failed_log = _write_failed_files(output_path, failures) + raise RuntimeError(f"Failed to download {len(failures)} PDC file(s). See {failed_log}") + + LOGGER.info( + "PDC download finished: studies=%d total=%d downloaded=%d skipped=%d failed=%d", + stats.studies, + stats.total_files, + stats.downloaded, + stats.skipped, + stats.failed, + ) + return stats diff --git a/pridepy/pridepy.py b/pridepy/pridepy.py index e307a61..bdea8af 100644 --- a/pridepy/pridepy.py +++ b/pridepy/pridepy.py @@ -3,6 +3,7 @@ import logging import click from pridepy.download.client import Client as Files +from pridepy.pdc import download_pdc_files as run_pdc_download from pridepy.project.project import Project PROTOCOL_CHOICES = click.Choice(["ftp", "aspera", "globus", "s3"], case_sensitive=False) @@ -52,11 +53,12 @@ def main(): default=False, ) @click.option( - "-w", - "--parallel-files", + "-t", + "--threads", + "download_threads", default=1, - type=click.IntRange(1, 3), - help="Number of files to download simultaneously (1-3). Primarily used by globus protocol. Default is 1.", + type=click.IntRange(1, 32), + help="Number of threads for each file download. Default is 1.", ) @click.option( "--preserve-structure", @@ -72,7 +74,7 @@ def download_all_public_raw_files( skip_if_downloaded_already, aspera_maximum_bandwidth: str = "50M", checksum_check: bool = False, - parallel_files: int = 1, + download_threads: int = 1, preserve_structure: bool = False, ): """ @@ -85,7 +87,7 @@ def download_all_public_raw_files( skip_if_downloaded_already (bool): Skip download if files already exist. Default is False. 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. - parallel_files (int): Number of files to download simultaneously. Default is 1. + download_threads (int): Number of threads for each file download. Default is 1. """ raw_files = Files() @@ -102,7 +104,7 @@ def download_all_public_raw_files( protocol, aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, - parallel_files=parallel_files, + download_threads=download_threads, flatten=not preserve_structure, ) @@ -152,11 +154,12 @@ def download_all_public_raw_files( "Valid values: RAW, PEAK, SEARCH, RESULT, SPECTRUM_LIBRARY, OTHER, FASTA", ) @click.option( - "-w", - "--parallel-files", + "-t", + "--threads", + "download_threads", default=1, - type=click.IntRange(1, 3), - help="Number of files to download simultaneously (1-3). Primarily used by globus protocol. Default is 1.", + type=click.IntRange(1, 32), + help="Number of threads for each file download. Default is 1.", ) @click.option( "--preserve-structure", @@ -173,7 +176,7 @@ def download_all_public_category_files( aspera_maximum_bandwidth: str = "50M", checksum_check: bool = False, category: str = "RAW", - parallel_files: int = 1, + download_threads: int = 1, preserve_structure: bool = False, ): """ @@ -187,7 +190,7 @@ def download_all_public_category_files( aspera_maximum_bandwidth (str): Maximum bandwidth for Aspera transfers. 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). - parallel_files (int): Number of files to download simultaneously. Default is 1. + download_threads (int): Number of threads for each file download. Default is 1. """ valid_categories = {"RAW", "PEAK", "SEARCH", "RESULT", "SPECTRUM_LIBRARY", "OTHER", "FASTA"} @@ -214,7 +217,7 @@ def download_all_public_category_files( aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, categories=categories, - parallel_files=parallel_files, + download_threads=download_threads, flatten=not preserve_structure, ) @@ -592,11 +595,12 @@ def _read_url_arguments(url_list_path, urls_csv=None): help="Download project checksums and validate downloaded files.", ) @click.option( - "-w", - "--parallel-files", + "-t", + "--threads", + "download_threads", default=1, - type=click.IntRange(1, 3), - help="Number of files to download simultaneously (1-3). Primarily used by globus protocol. Default is 1.", + type=click.IntRange(1, 32), + help="Number of threads for each file download. Default is 1.", ) @click.option( "--preserve-structure", @@ -614,7 +618,7 @@ def download_files_by_list( skip_if_downloaded_already, aspera_maximum_bandwidth, checksum_check, - parallel_files, + download_threads, preserve_structure: bool = False, ): """Download a named subset of files from a PRIDE project.""" @@ -630,7 +634,7 @@ def download_files_by_list( protocol=protocol, aspera_maximum_bandwidth=aspera_maximum_bandwidth, checksum_check=checksum_check, - parallel_files=parallel_files, + download_threads=download_threads, flatten=not preserve_structure, ) @@ -682,11 +686,12 @@ def download_files_by_list( "Accessions are inferred from PRIDE URL paths (only PRIDE URLs supported).", ) @click.option( - "-w", - "--parallel-files", + "-t", + "--threads", + "download_threads", default=1, - type=click.IntRange(1, 3), - help="Number of files to download simultaneously (1-3), for any URL scheme. Default is 1.", + type=click.IntRange(1, 32), + help="Number of threads for each file download. Default is 1.", ) def download_files_by_url( url_list_path, @@ -695,7 +700,7 @@ def download_files_by_url( skip_if_downloaded_already, protocol, checksum_check, - parallel_files, + download_threads, ): """Download files from raw URLs (http/https/ftp), dispatched by scheme.""" urls = _read_url_arguments(url_list_path, urls_csv) @@ -705,10 +710,93 @@ def download_files_by_url( output_folder=output_folder, skip_if_downloaded_already=skip_if_downloaded_already, protocol=protocol, - parallel_files=parallel_files, + download_threads=download_threads, checksum_check=checksum_check, ) +@main.command( + "download-pdc-files", + help="Download files from PDC/CPTAC studies via PDC signed HTTPS URLs", +) +@click.option( + "-a", + "--accession", + required=True, + help="PDC study ID, comma-separated PDC study IDs, or a CSV with pdc_id/pdc_study_id and optional file-type/filetype.", +) +@click.option( + "--file-type", + required=False, + type=click.Choice(["mzid", "psm", "raw", "mzml"], case_sensitive=False), + help="PDC file type to download: mzid, psm, raw, or mzml. Overrides CSV file-type/filetype values.", +) +@click.option( + "-o", + "--output-folder", + required=True, + help="Output folder. Files are written as //.", +) +@click.option( + "--skip-if-downloaded-already", + is_flag=True, + default=False, + help="Skip files that already exist locally and match PDC size/checksum metadata.", +) +@click.option( + "--checksum-check/--no-checksum-check", + "checksum_check", + default=True, + help="Validate downloads against PDC md5sum values. Enabled by default.", +) +@click.option( + "-t", + "--threads", + "download_threads", + default=1, + type=click.IntRange(1, 32), + help="Number of parallel HTTP Range threads per file (1-32). Default is 1.", +) +@click.option( + "--retry", + is_flag=True, + default=False, + help="Retry failed files; HTTP 403 retries refresh the PDC signed URL first.", +) +def download_pdc_files( + accession, + file_type, + output_folder, + skip_if_downloaded_already, + checksum_check, + download_threads, + retry, +): + """Download PDC/CPTAC files with PDC GraphQL metadata and signed HTTPS URLs.""" + try: + stats = run_pdc_download( + accession=accession, + file_type=file_type, + output_folder=output_folder, + skip_if_downloaded_already=skip_if_downloaded_already, + checksum_check=checksum_check, + download_threads=download_threads, + retry=retry, + ) + except ValueError as exc: + raise click.BadParameter(str(exc)) from exc + except RuntimeError as exc: + raise click.ClickException(str(exc)) from exc + + logging.info( + "PDC download completed: studies=%d total=%d downloaded=%d skipped=%d failed=%d", + stats.studies, + stats.total_files, + stats.downloaded, + stats.skipped, + stats.failed, + ) + + if __name__ == "__main__": main() diff --git a/pridepy/tests/test_cli_flatten.py b/pridepy/tests/test_cli_flatten.py index b068a40..92c00df 100644 --- a/pridepy/tests/test_cli_flatten.py +++ b/pridepy/tests/test_cli_flatten.py @@ -22,6 +22,25 @@ 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 + + def test_download_all_public_raw_files_threads(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-all-public-raw-files", + "-a", + "MSV000012345", + "-o", + "/tmp/x", + "--threads", + "4", + ] + ) + kwargs = files_cls.return_value.download_all_raw_files.call_args.kwargs + assert kwargs["download_threads"] == 4 + assert "parallel_files" not in kwargs def test_download_all_public_raw_files_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -54,6 +73,8 @@ 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 def test_download_files_by_list_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: @@ -71,6 +92,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 + + def test_download_files_by_url_threads(self): + with patch("pridepy.pridepy.Files") as files_cls: + self._invoke( + [ + "download-files-by-url", + "-u", + "https://example.org/a.raw", + "-o", + "/tmp/x", + "-t", + "4", + ] + ) + kwargs = files_cls.download_files_by_url.call_args.kwargs + assert kwargs["download_threads"] == 4 + assert "parallel_files" not in kwargs def test_download_px_raw_files_preserve_structure(self): with patch("pridepy.pridepy.Files") as files_cls: diff --git a/pridepy/tests/test_download_by_url.py b/pridepy/tests/test_download_by_url.py index b8c66da..fa5644a 100644 --- a/pridepy/tests/test_download_by_url.py +++ b/pridepy/tests/test_download_by_url.py @@ -12,7 +12,7 @@ import click import pytest -from pridepy.download import by_url +from pridepy.download import by_url, transport from pridepy.download.client import Client as Files from pridepy.pridepy import _read_url_arguments @@ -49,6 +49,55 @@ def fake_http(_url, target_path): mock_http.assert_called_once() assert os.path.exists(target) + def test_threads_use_multipart_http(self): + with tempfile.TemporaryDirectory() as tmp_dir: + target = os.path.join(tmp_dir, "sample.raw") + + def fake_multipart(_url, target_path, threads=1, position=0): + assert threads == 4 + assert position == 0 + _touch_valid(target_path) + + with patch.object( + by_url.transport, "_multipart_download", side_effect=fake_multipart + ) as mock_multipart: + Files.download_files_by_url( + urls=["https://example.org/sample.raw"], + output_folder=tmp_dir, + download_threads=4, + ) + + mock_multipart.assert_called_once() + assert os.path.exists(target) + + def test_multipart_failure_removes_preallocated_file(self): + class FakeHeadResponse: + headers = {"content-length": str(20 * 1024 * 1024), "accept-ranges": "bytes"} + + def raise_for_status(self): + return None + + class FakeSession: + def head(self, *_args, **_kwargs): + return FakeHeadResponse() + + with tempfile.TemporaryDirectory() as tmp_dir: + target = os.path.join(tmp_dir, "large.raw") + with patch.object( + transport.Util, "create_session_with_retries", return_value=FakeSession() + ), patch.object( + transport, "_download_range", side_effect=RuntimeError("range failed") + ): + with pytest.raises(RuntimeError, match="range failed"): + transport._multipart_download( + "https://example.org/large.raw", + target, + threads=2, + min_size_bytes=1, + ) + + assert not os.path.exists(target) + def test_dispatches_ftp(self): with tempfile.TemporaryDirectory() as tmp_dir: target = os.path.join(tmp_dir, "sample.raw") diff --git a/pridepy/tests/test_pdc_cli.py b/pridepy/tests/test_pdc_cli.py new file mode 100644 index 0000000..57f4582 --- /dev/null +++ b/pridepy/tests/test_pdc_cli.py @@ -0,0 +1,90 @@ +from unittest import TestCase +from unittest.mock import Mock, patch + +from click.testing import CliRunner + +from pridepy.pdc.downloader import PDCDownloadStats +from pridepy.pridepy import main + + +class TestPDCCli(TestCase): + def test_help_exposes_pdc_options_without_pride_protocol_options(self): + result = CliRunner().invoke(main, ["download-pdc-files", "--help"]) + + assert result.exit_code == 0 + assert "--accession" in result.output + assert "--file-type" in result.output + assert "--output-folder" in result.output + assert "--skip-if-downloaded-already" in result.output + assert "--checksum-check / --no-checksum-check" in result.output + assert "--threads" in result.output + assert "--retry" in result.output + assert "--protocol" not in result.output + assert "--aspera-maximum-bandwidth" not in result.output + assert "--workers" not in result.output + assert "--proxy-port" not in result.output + + def test_command_delegates_to_pdc_downloader(self): + stats = PDCDownloadStats(studies=1, total_files=1, downloaded=1) + runner = CliRunner() + with patch("pridepy.pridepy.run_pdc_download", Mock(return_value=stats)) as mock_run: + result = runner.invoke( + main, + [ + "download-pdc-files", + "--accession", + "PDC000109", + "--file-type", + "psm", + "--output-folder", + "downloads", + "--skip-if-downloaded-already", + "--no-checksum-check", + "--threads", + "4", + "--retry", + ], + ) + + assert result.exit_code == 0 + mock_run.assert_called_once_with( + accession="PDC000109", + file_type="psm", + output_folder="downloads", + skip_if_downloaded_already=True, + checksum_check=False, + download_threads=4, + retry=True, + ) + + def test_command_allows_file_type_to_be_omitted_for_csv(self): + stats = PDCDownloadStats(studies=2, total_files=2, downloaded=2) + runner = CliRunner() + with patch("pridepy.pridepy.run_pdc_download", Mock(return_value=stats)) as mock_run: + result = runner.invoke( + main, + [ + "download-pdc-files", + "--accession", + "studies.csv", + "--output-folder", + "downloads", + ], + ) + + assert result.exit_code == 0 + mock_run.assert_called_once_with( + accession="studies.csv", + file_type=None, + output_folder="downloads", + skip_if_downloaded_already=False, + checksum_check=True, + download_threads=1, + retry=False, + ) + + def test_root_help_lists_pdc_command(self): + result = CliRunner().invoke(main, ["--help"]) + + assert result.exit_code == 0 + assert "download-pdc-files" in result.output diff --git a/pridepy/tests/test_pdc_client.py b/pridepy/tests/test_pdc_client.py new file mode 100644 index 0000000..13fe0fe --- /dev/null +++ b/pridepy/tests/test_pdc_client.py @@ -0,0 +1,177 @@ +import os +import tempfile +from unittest import TestCase +from unittest.mock import patch + +import pytest + +from pridepy.pdc.client import ( + PDCDownloadRequest, + PDCFile, + entry_matches_file_type, + fetch_study_files, + parse_accessions, + parse_download_requests, + refresh_signed_url, +) + + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeSession: + def __init__(self, payload): + self.payload = payload + self.calls = [] + + def post(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return FakeResponse(self.payload) + + +def _payload(entries): + return {"data": {"filesPerStudy": entries}} + + +def _entry(name, file_format="tsv", data_category="Peptide Spectral Matches", md5sum="900150983cd24fb0d6963f7d28e17f72"): + return { + "file_id": "file-1", + "pdc_study_id": "PDC000109", + "file_name": name, + "file_format": file_format, + "file_size": "3", + "data_category": data_category, + "file_type": "Text", + "file_location": "s3://bucket/key", + "md5sum": md5sum, + "signedUrl": {"url": f"https://example.org/{name}"}, + } + + +class TestPDCAccessions(TestCase): + def test_single_accession(self): + assert parse_accessions("PDC000109") == ["PDC000109"] + + def test_comma_accessions_are_deduped(self): + assert parse_accessions("PDC000109,PDC000110,PDC000109") == [ + "PDC000109", + "PDC000110", + ] + + def test_csv_pdc_id_column(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_id,name\nPDC000109,a\nPDC000110,b\nPDC000109,c\n") + assert parse_accessions(path) == ["PDC000109", "PDC000110"] + + def test_csv_pdc_study_id_column(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_study_id\nPDC000111\n") + assert parse_accessions(path) == ["PDC000111"] + + def test_csv_file_type_column_creates_mixed_download_requests(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_id,file-type\nPDC000109,raw\nPDC000110,mzml\n") + assert parse_download_requests(path) == [ + PDCDownloadRequest("PDC000109", "raw"), + PDCDownloadRequest("PDC000110", "mzml"), + ] + + def test_csv_filetype_column_is_supported(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_id,filetype\nPDC000109,psm\n") + assert parse_download_requests(path) == [PDCDownloadRequest("PDC000109", "psm")] + + def test_command_line_file_type_overrides_csv_file_type(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_id,file-type\nPDC000109,not-a-type\n") + with self.assertLogs("pridepy.pdc.client", level="WARNING") as logs: + requests = parse_download_requests(path, file_type="raw") + assert requests == [PDCDownloadRequest("PDC000109", "raw")] + assert "overrides CSV file-type values" in "\n".join(logs.output) + + def test_file_type_is_required_without_csv_file_type(self): + with pytest.raises(ValueError, match="--file-type is required"): + parse_download_requests("PDC000109") + + def test_csv_file_type_is_required_when_command_line_file_type_is_missing(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_id\nPDC000109\n") + with pytest.raises(ValueError, match="file-type/filetype"): + parse_download_requests(path) + + def test_csv_missing_column_raises(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("study\nPDC000111\n") + with pytest.raises(ValueError, match="pdc_id or pdc_study_id"): + parse_accessions(path) + + def test_empty_csv_raises(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "studies.csv") + with open(path, "w", encoding="utf-8") as handle: + handle.write("pdc_study_id\n\n") + with pytest.raises(ValueError, match="empty"): + parse_accessions(path) + + +class TestPDCClient(TestCase): + def test_psm_filter_keeps_only_psm_suffix(self): + assert entry_matches_file_type(_entry("sample.psm"), "psm") + assert not entry_matches_file_type(_entry("sample.tmt11.tsv"), "psm") + + def test_fetch_study_files_maps_graphql_entries(self): + session = FakeSession(_payload([_entry("sample.psm")])) + + files = fetch_study_files("PDC000109", "psm", session=session) + + assert len(files) == 1 + assert files[0] == PDCFile( + study_id="PDC000109", + file_id="file-1", + file_name="sample.psm", + file_format="tsv", + file_size=3, + data_category="Peptide Spectral Matches", + file_type="Text", + file_location="s3://bucket/key", + md5sum="900150983cd24fb0d6963f7d28e17f72", + url="https://example.org/sample.psm", + ) + assert session.calls[0][1]["json"]["variables"] == {"studyId": "PDC000109"} + + def test_fetch_study_files_skips_missing_signed_url(self): + entry = _entry("sample.psm") + entry["signedUrl"] = None + session = FakeSession(_payload([entry])) + + assert fetch_study_files("PDC000109", "psm", session=session) == [] + + def test_refresh_signed_url_returns_matching_file(self): + files = [ + PDCFile("PDC000109", "1", "a.psm", "tsv", 3, "Peptide Spectral Matches", None, None, None, "old"), + PDCFile("PDC000109", "2", "b.psm", "tsv", 3, "Peptide Spectral Matches", None, None, None, "new"), + ] + with patch("pridepy.pdc.client.fetch_study_files", return_value=files): + assert refresh_signed_url("PDC000109", "b.psm", "psm") == "new" diff --git a/pridepy/tests/test_pdc_downloader.py b/pridepy/tests/test_pdc_downloader.py new file mode 100644 index 0000000..164d611 --- /dev/null +++ b/pridepy/tests/test_pdc_downloader.py @@ -0,0 +1,260 @@ +import hashlib +import os +import tempfile +from unittest import TestCase +from unittest.mock import Mock, patch + +import pytest +import requests + +from pridepy.pdc.client import PDCFile +from pridepy.pdc.downloader import download_pdc_files + + +ABC_MD5 = hashlib.md5(b"abc").hexdigest() + + +def _pdc_file(name="sample.psm", url="https://example.org/sample.psm", md5sum=ABC_MD5, study_id="PDC000109"): + return PDCFile( + study_id=study_id, + file_id="file-1", + file_name=name, + file_format="tsv", + file_size=3, + data_category="Peptide Spectral Matches", + file_type="Text", + file_location="s3://bucket/key", + md5sum=md5sum, + url=url, + ) + + +def _fetcher(files): + def fetch_files(study_id, file_type): + assert study_id == "PDC000109" + assert file_type == "psm" + return files + + return fetch_files + + +def _write_data(_url, target): + with open(target, "wb") as handle: + handle.write(b"abc") + + +class TestPDCDownloader(TestCase): + def test_skip_existing_file_with_matching_checksum(self): + pdc_file = _pdc_file() + with tempfile.TemporaryDirectory() as tmp_dir: + study_dir = os.path.join(tmp_dir, "PDC000109") + os.makedirs(study_dir) + target = os.path.join(study_dir, "sample.psm") + with open(target, "wb") as handle: + handle.write(b"abc") + + with patch("pridepy.pdc.downloader.transport._parallel_download") as mock_download: + stats = download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + skip_if_downloaded_already=True, + checksum_check=True, + fetch_files=_fetcher([pdc_file]), + ) + + mock_download.assert_not_called() + assert stats.skipped == 1 + assert stats.downloaded == 0 + + def test_download_success_moves_part_to_final_file(self): + pdc_file = _pdc_file() + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=_write_data, + ): + stats = download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + fetch_files=_fetcher([pdc_file]), + ) + + target = os.path.join(tmp_dir, "PDC000109", "sample.psm") + assert stats.downloaded == 1 + assert os.path.exists(target) + assert not os.path.exists(target + ".part") + with open(target, "rb") as handle: + assert handle.read() == b"abc" + + def test_csv_can_download_multiple_file_types(self): + calls = [] + + def fetch_files(study_id, file_type): + calls.append((study_id, file_type)) + return [_pdc_file(name=f"{study_id}.{file_type}", study_id=study_id)] + + with tempfile.TemporaryDirectory() as tmp_dir: + csv_path = os.path.join(tmp_dir, "studies.csv") + with open(csv_path, "w", encoding="utf-8") as handle: + handle.write("pdc_id,file-type\nPDC000109,psm\nPDC000110,mzid\n") + + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=_write_data, + ): + stats = download_pdc_files( + accession=csv_path, + file_type=None, + output_folder=tmp_dir, + checksum_check=True, + fetch_files=fetch_files, + ) + + assert calls == [("PDC000109", "psm"), ("PDC000110", "mzid")] + assert stats.studies == 2 + assert stats.total_files == 2 + assert stats.downloaded == 2 + assert os.path.exists(os.path.join(tmp_dir, "PDC000109", "PDC000109.psm")) + assert os.path.exists(os.path.join(tmp_dir, "PDC000110", "PDC000110.mzid")) + + def test_threads_use_multipart_downloader(self): + pdc_file = _pdc_file() + with tempfile.TemporaryDirectory() as tmp_dir: + def fake_multipart(_url, target, threads=1): + assert threads == 4 + _write_data(_url, target) + + with patch( + "pridepy.pdc.downloader.transport._multipart_download", + side_effect=fake_multipart, + ) as mock_multipart: + download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + download_threads=4, + fetch_files=_fetcher([pdc_file]), + ) + + mock_multipart.assert_called_once() + + def test_missing_md5_falls_back_to_size_validation(self): + pdc_file = _pdc_file(md5sum=None) + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=_write_data, + ): + stats = download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + fetch_files=_fetcher([pdc_file]), + ) + + assert stats.downloaded == 1 + + def test_all_empty_matches_fail_and_write_failed_files(self): + with tempfile.TemporaryDirectory() as tmp_dir: + with pytest.raises(RuntimeError, match="Failed to download 1 PDC file"): + download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + fetch_files=lambda _study_id, _file_type: [], + ) + + failed_log = os.path.join(tmp_dir, "failed_files.txt") + with open(failed_log, "r", encoding="utf-8") as handle: + content = handle.read() + assert "PDC000109\t\tNo PDC files matched file_type=psm" in content + + def test_partial_empty_matches_warn_and_continue(self): + pdc_file = _pdc_file() + + def fetch_files(study_id, file_type): + assert file_type == "psm" + if study_id == "PDC_EMPTY": + return [] + assert study_id == "PDC000109" + return [pdc_file] + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=_write_data, + ): + stats = download_pdc_files( + accession="PDC000109,PDC_EMPTY", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + fetch_files=fetch_files, + ) + + assert stats.studies == 2 + assert stats.total_files == 1 + assert stats.downloaded == 1 + assert stats.failed == 0 + assert not os.path.exists(os.path.join(tmp_dir, "failed_files.txt")) + + def test_checksum_failure_writes_failed_files(self): + pdc_file = _pdc_file(md5sum="ffffffffffffffffffffffffffffffff") + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=_write_data, + ): + with pytest.raises(RuntimeError, match="Failed to download 1 PDC file"): + download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + fetch_files=_fetcher([pdc_file]), + ) + + failed_log = os.path.join(tmp_dir, "failed_files.txt") + with open(failed_log, "r", encoding="utf-8") as handle: + content = handle.read() + assert "PDC000109\tsample.psm\tchecksum mismatch" in content + + def test_403_retry_refreshes_signed_url(self): + pdc_file = _pdc_file(url="https://example.org/old") + response = Mock() + response.status_code = 403 + http_error = requests.HTTPError("403 Client Error") + http_error.response = response + seen_urls = [] + + def fake_download(url, target): + seen_urls.append(url) + if url == "https://example.org/old": + raise http_error + _write_data(url, target) + + refresh = Mock(return_value="https://example.org/new") + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.pdc.downloader.transport._parallel_download", + side_effect=fake_download, + ), patch("pridepy.pdc.downloader.time.sleep"): + stats = download_pdc_files( + accession="PDC000109", + file_type="psm", + output_folder=tmp_dir, + checksum_check=True, + retry=True, + fetch_files=_fetcher([pdc_file]), + refresh_url=refresh, + ) + + assert stats.downloaded == 1 + assert seen_urls == ["https://example.org/old", "https://example.org/new"] + refresh.assert_called_once_with("PDC000109", "sample.psm", "psm")