diff --git a/pridepy/download/pride.py b/pridepy/download/pride.py index 80f9348..d75f478 100644 --- a/pridepy/download/pride.py +++ b/pridepy/download/pride.py @@ -53,6 +53,14 @@ class PrideProvider(Provider): ARCHIVE_FTP_URL_PREFIX: ClassVar[str] = "ftp://ftp.pride.ebi.ac.uk/" ARCHIVE_HTTPS_URL_PREFIX: ClassVar[str] = "https://ftp.pride.ebi.ac.uk/" S3_URL: ClassVar[str] = "https://hh.fire.sdo.ebi.ac.uk" + # EBI-internal FIRE S3 endpoint. Only reachable from within EBI + # infrastructure (compute/login nodes), where it is markedly faster and + # more reliable than the public FTP/HTTPS/Globus paths — which is the + # whole point of the ``fire`` protocol. Overridable via the + # ``PRIDEPY_FIRE_ENDPOINT`` env var for sites with a different host. + FIRE_S3_URL: ClassVar[str] = os.environ.get( + "PRIDEPY_FIRE_ENDPOINT", "https://hl.fire.sdo.ebi.ac.uk" + ) S3_BUCKET: ClassVar[str] = "pride-public" PROTOCOL_ORDER: ClassVar[List[str]] = ["aspera", "s3", "ftp", "globus"] @@ -138,7 +146,16 @@ def get_submitted_file_path_prefix(self, accession): def _protocol_sequence(protocol: str) -> List[str]: """ Build the ordered list of protocols to try for a requested download mode. + + ``fire`` (the EBI-internal FIRE S3 endpoint) is only ever tried when it + is explicitly requested — it is never folded into another protocol's + fallback chain, because it is unreachable outside EBI and would just + waste retries there. When requested, it is tried first and then falls + back to the public protocols so a run started outside EBI still + completes. """ + if protocol == "fire": + return ["fire"] + PrideProvider.PROTOCOL_ORDER if protocol not in PrideProvider.PROTOCOL_ORDER: return [] return [protocol] + [p for p in PrideProvider.PROTOCOL_ORDER if p != protocol] @@ -218,7 +235,9 @@ def _get_download_url(file_record: Dict, protocol: str) -> str: PrideProvider.ARCHIVE_HTTPS_URL_PREFIX, 1, ) - if protocol == "s3": + if protocol in ("s3", "fire"): + # Both S3 modes derive the object key from the FTP path; they + # differ only in the FIRE endpoint host (external vs EBI-internal). return ftp_url raise ValueError(f"Unsupported protocol: {protocol}") @@ -511,17 +530,26 @@ def download_files_from_globus( @staticmethod def download_files_from_s3( - file_list_json: List[Dict], output_folder: str, skip_if_downloaded_already + file_list_json: List[Dict], + output_folder: str, + skip_if_downloaded_already, + endpoint_url: Optional[str] = None, ): """ - Download files using S3 transfer URL with a progress bar and retry logic. + Download files from a FIRE S3 endpoint with a progress bar and retry logic. + :param file_list_json: file list in JSON format :param output_folder: folder to download the files :param skip_if_downloaded_already: Boolean value to skip the download if the file has already been downloaded. + :param endpoint_url: FIRE S3 endpoint. Defaults to the public + :attr:`S3_URL` (``hh.fire``); pass :attr:`FIRE_S3_URL` (``hl.fire``) + for the EBI-internal ``fire`` protocol. """ if not os.path.isdir(output_folder): os.makedirs(output_folder, exist_ok=True) + endpoint_url = endpoint_url or PrideProvider.S3_URL + # Retry and timeout config retry_config = Config( retries={"max_attempts": 5, "mode": "standard"}, @@ -533,19 +561,21 @@ def download_files_from_s3( s3_resource = boto3.resource( "s3", config=retry_config, - endpoint_url=PrideProvider.S3_URL, + endpoint_url=endpoint_url, ) bucket = s3_resource.Bucket(PrideProvider.S3_BUCKET) + logging.info( + "Downloading %d file(s) from FIRE S3 endpoint %s (bucket %s)", + len(file_list_json), endpoint_url, PrideProvider.S3_BUCKET, + ) failed: List[str] = [] for file in file_list_json: try: - # Determine S3 or FTP path - download_url = ( - file["publicFileLocations"][0]["value"] - if file["publicFileLocations"][0]["name"] == "FTP Protocol" - else file["publicFileLocations"][1]["value"] - ) + # Resolve the canonical FTP URL, then map it to the S3 object + # key: pride-public mirrors the archive layout, so the key is + # the archive-relative path (YYYY/MM/ACCESSION/filename). + download_url = PrideProvider._get_download_url(file, "ftp") ftp_base_url = "ftp://ftp.pride.ebi.ac.uk/pride/data/archive/" s3_path = download_url.replace(ftp_base_url, "") @@ -731,11 +761,14 @@ def _batch_download_by_protocol( download_threads=download_threads, ) return - if protocol == "s3": + if protocol in ("s3", "fire"): PrideProvider.download_files_from_s3( file_list, output_folder, skip_if_downloaded_already=skip_if_downloaded_already, + endpoint_url=( + PrideProvider.FIRE_S3_URL if protocol == "fire" else PrideProvider.S3_URL + ), ) return raise ValueError(f"Unsupported protocol: {protocol}") @@ -925,9 +958,9 @@ def _download_files_batch( :param aspera_maximum_bandwidth: parameter in Aspera sets the maximum bandwidth for the transfer. :param skip_if_downloaded_already: Boolean value to skip the download if the file has already been downloaded. """ - protocols_supported = ["ftp", "aspera", "globus", "s3"] + protocols_supported = ["ftp", "aspera", "globus", "s3", "fire"] if protocol not in protocols_supported: - logging.error("Protocol should be one of ftp, aspera, globus, s3") + logging.error("Protocol should be one of ftp, aspera, globus, s3, fire") return os.makedirs(output_folder, exist_ok=True) @@ -943,8 +976,13 @@ def _download_files_batch( protocol_sequence = PrideProvider._protocol_sequence(protocol) primary_protocol = protocol_sequence[0] - # Retry with the primary protocol first, then fall back to others - fallback_sequence = protocol_sequence + # Retry with the primary protocol first, then fall back to others. + # ``fire`` is excluded from the per-file fallback: a FIRE failure is an + # endpoint-level condition (the EBI-internal host is unreachable), so it + # fails identically for every file. Re-attempting it per file in Phase 2 + # would just burn retries on connections that cannot succeed — send those + # files straight to the public fallback protocols instead. + fallback_sequence = [p for p in protocol_sequence if p != "fire"] # Phase 1: batch download with the requested protocol. Reuses a single # FTP/S3 connection for all files (the previous behaviour) instead of diff --git a/pridepy/pridepy.py b/pridepy/pridepy.py index 17723cf..bfeb27a 100644 --- a/pridepy/pridepy.py +++ b/pridepy/pridepy.py @@ -6,7 +6,9 @@ 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) +PROTOCOL_CHOICES = click.Choice( + ["ftp", "aspera", "globus", "s3", "fire"], case_sensitive=False +) @click.group() @@ -25,7 +27,10 @@ def main(): "--protocol", default="ftp", type=PROTOCOL_CHOICES, - help="Protocol to use for download: ftp, aspera, globus, s3. Default is ftp with fallback enabled.", + help="Protocol to use for download: ftp, aspera, globus, s3, fire. " + "'fire' uses the EBI-internal FIRE S3 endpoint (only reachable inside EBI " + "infrastructure; falls back to the public protocols elsewhere). " + "Default is ftp with fallback enabled.", ) @click.option( "-o", @@ -119,7 +124,10 @@ def download_all_public_raw_files( "--protocol", default="ftp", type=PROTOCOL_CHOICES, - help="Protocol to use for download: ftp, aspera, globus, s3. Default is ftp with fallback enabled.", + help="Protocol to use for download: ftp, aspera, globus, s3, fire. " + "'fire' uses the EBI-internal FIRE S3 endpoint (only reachable inside EBI " + "infrastructure; falls back to the public protocols elsewhere). " + "Default is ftp with fallback enabled.", ) @click.option( "-o", @@ -232,7 +240,10 @@ def download_all_public_category_files( "--protocol", default="ftp", type=PROTOCOL_CHOICES, - help="Protocol to use for download: ftp, aspera, globus, s3. Default is ftp with fallback enabled.", + help="Protocol to use for download: ftp, aspera, globus, s3, fire. " + "'fire' uses the EBI-internal FIRE S3 endpoint (only reachable inside EBI " + "infrastructure; falls back to the public protocols elsewhere). " + "Default is ftp with fallback enabled.", ) @click.option("-f", "--file-name", required=True, help="fileName to be downloaded") @click.option( @@ -553,7 +564,10 @@ def _read_url_arguments(url_list_path, urls_csv=None): "--protocol", default="ftp", type=PROTOCOL_CHOICES, - help="Protocol to use for download: ftp, aspera, globus, s3. Default is ftp with fallback enabled.", + help="Protocol to use for download: ftp, aspera, globus, s3, fire. " + "'fire' uses the EBI-internal FIRE S3 endpoint (only reachable inside EBI " + "infrastructure; falls back to the public protocols elsewhere). " + "Default is ftp with fallback enabled.", ) @click.option( "-F", diff --git a/pridepy/tests/test_review_fixes.py b/pridepy/tests/test_review_fixes.py index 242941c..2458d40 100644 --- a/pridepy/tests/test_review_fixes.py +++ b/pridepy/tests/test_review_fixes.py @@ -142,6 +142,68 @@ def test_s3_batch_raises_when_a_file_fails(self): records, tmp_dir, skip_if_downloaded_already=False ) + def test_fire_protocol_uses_internal_endpoint_and_derives_key(self): + """`fire` must hit the EBI-internal FIRE endpoint and map the FTP path + to the correct pride-public S3 object key.""" + records = [_pride_record("a.raw", accession="PXD002137", date="2015/08")] + mock_obj = Mock() + mock_obj.content_length = 10 + mock_bucket = Mock() + mock_bucket.Object.return_value = mock_obj + mock_resource = Mock() + mock_resource.Bucket.return_value = mock_bucket + with tempfile.TemporaryDirectory() as tmp_dir: + with patch( + "pridepy.download.pride.boto3.resource", return_value=mock_resource + ) as mock_boto: + PrideProvider._batch_download_by_protocol( + records, + tmp_dir, + protocol="fire", + skip_if_downloaded_already=False, + aspera_maximum_bandwidth="100M", + ) + # boto3.resource was created against the internal hl.fire endpoint. + assert mock_boto.call_args.kwargs["endpoint_url"] == PrideProvider.FIRE_S3_URL + assert "hl.fire.sdo.ebi.ac.uk" in PrideProvider.FIRE_S3_URL + # The object key is the archive-relative path, no ftp:// prefix. + mock_bucket.Object.assert_called_once_with("2015/08/PXD002137/a.raw") + + def test_fire_protocol_sequence_requested_only(self): + """`fire` is tried first when requested, then the public protocols; + it is never folded into another protocol's fallback chain.""" + assert PrideProvider._protocol_sequence("fire") == [ + "fire", "aspera", "s3", "ftp", "globus", + ] + assert "fire" not in PrideProvider._protocol_sequence("ftp") + assert "fire" not in PrideProvider._protocol_sequence("s3") + + def test_fire_not_retried_in_phase2_fallback(self): + """After a FIRE (endpoint-level) failure, the per-file Phase-2 fallback + must go straight to the public protocols and never re-attempt fire.""" + record = _pride_record("a.raw", accession="PXD002137", date="2015/08") + captured = {} + + def _fake_fallback(*, file_record, output_folder, protocol_sequence, **kw): + captured["seq"] = protocol_sequence + return True # pretend a public protocol succeeded + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch.object(PrideProvider, "_batch_download_by_protocol"), \ + patch("pridepy.download.pride._provider_util.validate_download", + return_value=(False, "missing")), \ + patch.object(PrideProvider, "_download_with_fallback", + side_effect=_fake_fallback): + PrideProvider._download_files_batch( + file_list_json=[record], + accession="PXD002137", + output_folder=tmp_dir, + skip_if_downloaded_already=False, + protocol="fire", + ) + assert captured["seq"] == ["aspera", "s3", "ftp", "globus"] + assert "fire" not in captured["seq"] + def test_download_files_forwards_protocol_to_get_download_url(self): seen = []