From 576f80dba4dcf3f337d698a11f20ef023b30088e Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 22:47:00 +0200 Subject: [PATCH 1/6] Implement retrying download on same URL --- src/datacollective/datasets.py | 23 ++++++ src/datacollective/download.py | 143 ++++++++++++++++++++++++++------- 2 files changed, 139 insertions(+), 27 deletions(-) diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index eeedc40..44d04df 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -21,6 +21,7 @@ ) from datacollective.archive_utils import _extract_archive from datacollective.download import ( + DEFAULT_MAX_DOWNLOAD_ATTEMPTS, DOWNLOAD_SOURCE_DOWNLOAD, DOWNLOAD_SOURCE_LOAD, DOWNLOAD_SOURCE_SAMPLE, @@ -87,6 +88,7 @@ def download_dataset( enable_logging: bool = False, extract: bool = False, overwrite_extracted: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, ) -> Path: """ Download the dataset archive to a local directory and return the archive path. @@ -115,6 +117,10 @@ def download_dataset( extract: Whether to also unpack the archive and return the extracted directory. overwrite_extracted: With `extract=True`, whether to unpack the archive again even if the extracted directory already exists. Does not download the archive again. + max_download_attempts: How many times to request the archive before giving up. + A connection that drops mid-download is retried against the same download + URL, continuing from the bytes already saved, without starting a new download + session. Defaults to 3. Returns: Path to the downloaded dataset archive, or to the extracted directory if @@ -140,6 +146,7 @@ def download_dataset( download_source=DOWNLOAD_SOURCE_DOWNLOAD, extract=extract, overwrite_extracted=overwrite_extracted, + max_download_attempts=max_download_attempts, ) @@ -151,6 +158,7 @@ def download_sample( enable_logging: bool = False, extract: bool = False, overwrite_extracted: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, ) -> Path: """ Download the dataset's sample archive to a local directory and return the archive path. @@ -184,6 +192,10 @@ def download_sample( extract: Whether to also unpack the archive and return the extracted directory. overwrite_extracted: With `extract=True`, whether to unpack the sample again even if the extracted directory already exists. Does not download the sample again. + max_download_attempts: How many times to request the archive before giving up. + A connection that drops mid-download is retried against the same download + URL, continuing from the bytes already saved, without starting a new download + session. Defaults to 3. Returns: Path to the downloaded sample archive, or to the extracted directory if @@ -211,6 +223,7 @@ def download_sample( download_directory=download_directory, show_progress=show_progress, overwrite_existing=overwrite_existing, + max_download_attempts=max_download_attempts, download_source=DOWNLOAD_SOURCE_SAMPLE, ) if not extract: @@ -227,6 +240,7 @@ def _download_archive( download_source: str, extract: bool = False, overwrite_extracted: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, ) -> Path: """Shared body of `download_dataset` and its deprecated alias; the *download_source* sent with the request tells them apart.""" @@ -241,6 +255,7 @@ def _download_archive( download_directory=download_directory, show_progress=show_progress, overwrite_existing=overwrite_existing, + max_download_attempts=max_download_attempts, download_source=download_source, ) if not extract: @@ -276,6 +291,7 @@ def load_dataset( overwrite_existing: bool = False, overwrite_extracted: bool = False, enable_logging: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, return_format: Literal["pandas"] = "pandas", ) -> pd.DataFrame: ... @@ -288,6 +304,7 @@ def load_dataset( overwrite_existing: bool = False, overwrite_extracted: bool = False, enable_logging: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, *, return_format: Literal["hf"], ) -> Dataset | DatasetDict: ... @@ -300,6 +317,7 @@ def load_dataset( overwrite_existing: bool = False, overwrite_extracted: bool = False, enable_logging: bool = False, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, return_format: Literal["pandas", "hf"] = "pandas", ) -> pd.DataFrame | Dataset | DatasetDict: """ @@ -330,6 +348,10 @@ def load_dataset( Only makes sense when overwrite_existing is False. Will check in the download directory for existing extracted files with the default naming of the folder. enable_logging: Whether to enable SDK logging to console and a local log file. + max_download_attempts: How many times to request the archive before giving up. + A connection that drops mid-download is retried against the same download + URL, continuing from the bytes already saved, without starting a new download + session. Defaults to 3. return_format: Format of the returned object. `"pandas"` (default) returns a pandas DataFrame. `"hf"` returns a HuggingFace `Dataset`, or a `DatasetDict` keyed by split name for datasets with multiple splits. @@ -383,6 +405,7 @@ def load_dataset( download_directory=download_directory, show_progress=show_progress, overwrite_existing=overwrite_existing, + max_download_attempts=max_download_attempts, download_source=DOWNLOAD_SOURCE_LOAD, ) diff --git a/src/datacollective/download.py b/src/datacollective/download.py index 15c7313..dad64f2 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -1,7 +1,9 @@ import os +import time from pathlib import Path from typing import Any +import requests from fox_progress_bar import ProgressBar from pydantic import Field @@ -27,6 +29,17 @@ # dataset's sample, and sharing a directory would let one be mistaken for the other. SAMPLES_DIRNAME = "samples" +# A connection that drops while the archive is streaming is retried against the same +# presigned URL, continuing from the bytes already on disk. The URL is still valid for a short +# period of time so there is no need to ask for a new one +DEFAULT_MAX_DOWNLOAD_ATTEMPTS = 3 +DOWNLOAD_RETRY_BACKOFF_SECONDS = 2 +_RETRYABLE_DOWNLOAD_ERRORS = ( + requests.ConnectionError, + requests.Timeout, + requests.exceptions.ChunkedEncodingError, +) + class DownloadPlan(NonEmptyStrModel): download_url: str @@ -44,6 +57,7 @@ def _download_dataset( download_directory: str | None, show_progress: bool, overwrite_existing: bool, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, download_source: str | None = None, ) -> Path: """ @@ -66,6 +80,9 @@ def _download_dataset( show_progress: Show a progress bar when downloading the dataset. overwrite_existing: If True, overwrite any existing dataset archive in the download directory. download_source: Optional context appended to the User-Agent for download analytics. + max_download_attempts: How many times to request the archive before giving up; + a dropped connection is retried against the same download URL, continuing + from the bytes already on disk. Returns: The full path to the downloaded dataset archive. """ @@ -89,6 +106,7 @@ def _download_dataset( download_plan=download_plan, show_progress=show_progress, overwrite_existing=overwrite_existing, + max_download_attempts=max_download_attempts, download_source=download_source, ) @@ -98,6 +116,7 @@ def _download_sample( download_directory: str | None, show_progress: bool, overwrite_existing: bool, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, download_source: str | None = None, ) -> Path: """ @@ -118,6 +137,9 @@ def _download_sample( show_progress: Show a progress bar when downloading the sample. overwrite_existing: If True, download the sample again and replace any existing one. download_source: Optional context appended to the User-Agent for download analytics. + max_download_attempts: How many times to request the archive before giving up; + a dropped connection is retried against the same download URL, continuing + from the bytes already on disk. Returns: The full path to the downloaded sample archive. """ @@ -147,6 +169,7 @@ def _download_sample( download_plan=download_plan, show_progress=show_progress, overwrite_existing=overwrite_existing, + max_download_attempts=max_download_attempts, download_source=download_source, ) @@ -155,6 +178,7 @@ def _run_download_plan( download_plan: DownloadPlan, show_progress: bool, overwrite_existing: bool, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, download_source: str | None = None, ) -> Path: """ @@ -166,6 +190,9 @@ def _run_download_plan( show_progress: Whether to show a progress bar during download. overwrite_existing: If True, remove any existing complete or partial download first. download_source: Optional context appended to the User-Agent for download analytics. + max_download_attempts: How many times to request the archive before giving up; + a dropped connection is retried against the same download URL, continuing + from the bytes already on disk. Returns: The full path to the downloaded archive. @@ -190,6 +217,7 @@ def _run_download_plan( download_plan=download_plan, resume_download_checksum=resume_checksum, show_progress=show_progress, + max_download_attempts=max_download_attempts, download_source=download_source, ) @@ -366,65 +394,126 @@ def _execute_download_plan( download_plan: DownloadPlan, resume_download_checksum: str | None, show_progress: bool, + max_download_attempts: int = DEFAULT_MAX_DOWNLOAD_ATTEMPTS, download_source: str | None = None, ) -> None: """ Execute the download plan, downloading the dataset to a temporary path. + A connection that drops mid-stream is retried up to `max_download_attempts` times + against the same presigned URL, each attempt continuing from the end of the .part + file. Only connection-level failures are retried; an HTTP error, a KeyboardInterrupt + or any other exception fails the download at once. + Args: download_plan: The DownloadPlan object with download details. resume_download_checksum: Provide the checksum to resume a previously interrupted download. show_progress: Whether to show a progress bar during download. download_source: Optional context appended to the User-Agent for download analytics. + max_download_attempts: How many times to request the archive before giving up; + a dropped connection is retried against the same download URL, continuing + from the bytes already on disk. Raises: + ValueError: If `max_download_attempts` is less than 1. DownloadError: If the download fails or is interrupted. """ - headers: dict[str, str] = {} + if max_download_attempts < 1: + raise ValueError( + f"max_download_attempts must be at least 1, got {max_download_attempts}" + ) + previously_downloaded_bytes = 0 if resume_download_checksum and download_plan.tmp_filepath.exists(): previously_downloaded_bytes = download_plan.tmp_filepath.stat().st_size - headers["Range"] = f"bytes={previously_downloaded_bytes}-" progress_bar = None - session_downloaded_bytes = 0 logger.info(f"Downloading dataset: {download_plan.target_filepath}") if show_progress: progress_bar = ProgressBar(download_plan.size_bytes) progress_bar.update(previously_downloaded_bytes) progress_bar._display() + + # Bytes on disk, and hence the offset the next request continues from. + downloaded_bytes = previously_downloaded_bytes try: - with _send_api_request( - method="GET", - url=download_plan.download_url, - stream=True, - timeout=HTTP_TIMEOUT, - extra_headers=headers, - include_auth_headers=False, # Download URL is pre-signed, no auth needed - source_function=download_source, - ) as response: - with open(download_plan.tmp_filepath, "ab") as f: - # Iterate over response in 64KB chunks to avoid using too much memory - for chunk in response.iter_content(chunk_size=1 << 16): - if not chunk: - continue - f.write(chunk) - session_downloaded_bytes += len(chunk) - if progress_bar: - progress_bar.update(len(chunk)) - - if progress_bar: - progress_bar.finish() + for attempt in range(1, max_download_attempts + 1): + try: + _stream_archive_to_disk( + download_plan=download_plan, + offset=downloaded_bytes, + progress_bar=progress_bar, + download_source=download_source, + ) + break + except _RETRYABLE_DOWNLOAD_ERRORS: + downloaded_bytes = _bytes_on_disk(download_plan) + if attempt == max_download_attempts: + raise + wait_seconds = DOWNLOAD_RETRY_BACKOFF_SECONDS * attempt + logger.warning( + f"Connection lost after {downloaded_bytes} of {download_plan.size_bytes} " + f"bytes (attempt {attempt}/{max_download_attempts}). " + f"Retrying in {wait_seconds}s..." + ) + time.sleep(wait_seconds) + + if progress_bar: + progress_bar.finish() except (Exception, KeyboardInterrupt) as e: + downloaded_bytes = _bytes_on_disk(download_plan) raise DownloadError( - session_bytes=session_downloaded_bytes, - total_downloaded_bytes=previously_downloaded_bytes - + session_downloaded_bytes, + session_bytes=downloaded_bytes - previously_downloaded_bytes, + total_downloaded_bytes=downloaded_bytes, total_archive_bytes=download_plan.size_bytes, checksum=download_plan.checksum, ) from e +def _stream_archive_to_disk( + download_plan: DownloadPlan, + offset: int, + progress_bar: ProgressBar | None, + download_source: str | None, +) -> None: + """ + Stream the archive from *offset* to its end and append it to the .part file. + + Args: + download_plan: The DownloadPlan object with download details. + offset: Number of bytes already on disk; a Range header asks for the rest. + progress_bar: Progress bar to advance per chunk, if any. + download_source: Optional context appended to the User-Agent for download analytics. + """ + headers: dict[str, str] = {} + if offset: + headers["Range"] = f"bytes={offset}-" + + with _send_api_request( + method="GET", + url=download_plan.download_url, + stream=True, + timeout=HTTP_TIMEOUT, + extra_headers=headers, + include_auth_headers=False, # Download URL is pre-signed, no auth needed + source_function=download_source, + ) as response: + with open(download_plan.tmp_filepath, "ab") as f: + # Iterate over response in 64KB chunks to avoid using too much memory + for chunk in response.iter_content(chunk_size=1 << 16): + if not chunk: + continue + f.write(chunk) + if progress_bar: + progress_bar.update(len(chunk)) + + +def _bytes_on_disk(download_plan: DownloadPlan) -> int: + """Size of the .part file, or 0 if nothing has been written yet.""" + tmp = download_plan.tmp_filepath + return tmp.stat().st_size if tmp.exists() else 0 + + def _resolve_download_dir(download_directory: str | None) -> Path: """ Resolve and ensure the download directory exists and is writable. From 7f95b2e7530fe409e1fe45881011b51328dda056 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 22:48:14 +0200 Subject: [PATCH 2/6] Update docs --- docs/index.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/docs/index.md b/docs/index.md index ab80bdb..2604214 100644 --- a/docs/index.md +++ b/docs/index.md @@ -415,6 +415,22 @@ or `load_dataset` again. Sample downloads (`download_sample`) do not resume. The API returns no checksum for samples, so an interrupted sample download starts again from the beginning. +**Dropped connections are retried in place.** If the connection drops while an archive +is streaming, the SDK retries the request against the same download URL, continuing from +the bytes already on disk. These retries do not start a new download session, so they do +not count towards the API's download rate limit. By default up to 3 attempts are made; +pass `max_download_attempts` to `download_dataset`, `download_sample` or `load_dataset` +to change this, e.g. for an unreliable connection: + +```python +from datacollective import download_dataset + +archive_path = download_dataset("", max_download_attempts=10) +``` + +Only connection-level failures are retried. HTTP errors and manual cancellation fail +immediately, and a full download can still be resumed later as described above. + ### Automatically check for extracted archives From 965f47436f4464c24edb21dccf4ac1fab5912a9d Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 22:48:19 +0200 Subject: [PATCH 3/6] Update tests --- tests/test_datasets.py | 2 + tests/test_download.py | 173 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 8b575df..81b7b66 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -70,6 +70,7 @@ def test_save_dataset_to_disk_warns_and_reports_its_own_source( "download_directory": "dir", "show_progress": True, "overwrite_existing": False, + "max_download_attempts": 3, "download_source": "save_dataset_to_disk", } ] @@ -100,6 +101,7 @@ def fake_download_sample(**kwargs: object) -> Path: "download_directory": "dir", "show_progress": True, "overwrite_existing": False, + "max_download_attempts": 3, "download_source": "download_sample", } ] diff --git a/tests/test_download.py b/tests/test_download.py index d418142..6530dad 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -4,8 +4,10 @@ from typing import Any import pytest +import requests from datacollective.download import ( + DEFAULT_MAX_DOWNLOAD_ATTEMPTS, DownloadPlan, _download_sample, _get_checksum_filepath, @@ -14,6 +16,7 @@ _get_download_plan, _execute_download_plan, ) +from datacollective.errors import DownloadError def test_resolve_download_dir_prefers_argument(tmp_path: Path, monkeypatch) -> None: @@ -234,6 +237,176 @@ def fake_send_api_request(**kwargs: Any) -> FakeResponse: assert captured["include_auth_headers"] is False +class _FlakyDownloadApi: + """Stands in for `_send_api_request` on the presigned GET: serves *content* but drops the + connection after *fail_after* bytes on each of the first *failures* calls, then honours + the Range header on later calls. Records every call.""" + + def __init__(self, content: bytes, fail_after: int, failures: int) -> None: + self.content = content + self.fail_after = fail_after + self.failures = failures + self.calls: list[dict[str, Any]] = [] + + def __call__(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + api = self + attempt = len(self.calls) + range_header = (kwargs.get("extra_headers") or {}).get("Range") + start = ( + int(range_header.removeprefix("bytes=").rstrip("-")) if range_header else 0 + ) + + class FakeResponse: + def __enter__(self) -> "FakeResponse": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + return None + + def iter_content(self, chunk_size: int): + body = api.content[start:] + if attempt <= api.failures: + yield body[: api.fail_after] + raise requests.ConnectionError("connection reset") + yield body + + return FakeResponse() + + +def test_execute_download_plan_retries_same_url_after_connection_drop( + tmp_path: Path, monkeypatch +) -> None: + plan = _make_download_plan(tmp_path) + api = _FlakyDownloadApi(content=b"0123456789", fail_after=4, failures=1) + sleeps: list[float] = [] + monkeypatch.setattr("datacollective.download._send_api_request", api) + monkeypatch.setattr("datacollective.download.time.sleep", sleeps.append) + + _execute_download_plan( + plan, resume_download_checksum=None, show_progress=False, download_source=None + ) + + assert plan.tmp_filepath.read_bytes() == b"0123456789" + assert len(api.calls) == 2 + # Both attempts hit the same presigned URL; no new download session was requested. + assert {call["url"] for call in api.calls} == {plan.download_url} + assert {call["method"] for call in api.calls} == {"GET"} + assert "Range" not in (api.calls[0]["extra_headers"] or {}) + assert api.calls[1]["extra_headers"]["Range"] == "bytes=4-" + assert sleeps == [2] + + +def test_execute_download_plan_retry_continues_from_resumed_offset( + tmp_path: Path, monkeypatch +) -> None: + plan = _make_download_plan(tmp_path) + plan.tmp_filepath.write_bytes(b"012") + api = _FlakyDownloadApi(content=b"0123456789", fail_after=2, failures=1) + monkeypatch.setattr("datacollective.download._send_api_request", api) + monkeypatch.setattr("datacollective.download.time.sleep", lambda _s: None) + + _execute_download_plan( + plan, resume_download_checksum="abc", show_progress=False, download_source=None + ) + + assert plan.tmp_filepath.read_bytes() == b"0123456789" + assert api.calls[0]["extra_headers"]["Range"] == "bytes=3-" + assert api.calls[1]["extra_headers"]["Range"] == "bytes=5-" + + +def test_execute_download_plan_gives_up_after_max_attempts( + tmp_path: Path, monkeypatch +) -> None: + plan = _make_download_plan(tmp_path) + api = _FlakyDownloadApi(content=b"0123456789", fail_after=1, failures=99) + sleeps: list[float] = [] + monkeypatch.setattr("datacollective.download._send_api_request", api) + monkeypatch.setattr("datacollective.download.time.sleep", sleeps.append) + + with pytest.raises(DownloadError) as excinfo: + _execute_download_plan( + plan, + resume_download_checksum=None, + show_progress=False, + download_source=None, + ) + + assert len(api.calls) == DEFAULT_MAX_DOWNLOAD_ATTEMPTS + assert sleeps == [2, 4] + assert isinstance(excinfo.value.__cause__, requests.ConnectionError) + # One byte landed on disk per attempt. + assert plan.tmp_filepath.read_bytes() == b"012" + + +def test_execute_download_plan_honours_max_download_attempts( + tmp_path: Path, monkeypatch +) -> None: + plan = _make_download_plan(tmp_path) + api = _FlakyDownloadApi(content=b"0123456789", fail_after=1, failures=99) + monkeypatch.setattr("datacollective.download._send_api_request", api) + monkeypatch.setattr("datacollective.download.time.sleep", lambda _s: None) + + with pytest.raises(DownloadError): + _execute_download_plan( + plan, + resume_download_checksum=None, + show_progress=False, + max_download_attempts=1, + download_source=None, + ) + + assert len(api.calls) == 1 + + +def test_execute_download_plan_rejects_zero_attempts(tmp_path: Path) -> None: + plan = _make_download_plan(tmp_path) + with pytest.raises(ValueError, match="max_download_attempts"): + _execute_download_plan( + plan, + resume_download_checksum=None, + show_progress=False, + max_download_attempts=0, + download_source=None, + ) + + +def test_execute_download_plan_does_not_retry_other_errors( + tmp_path: Path, monkeypatch +) -> None: + plan = _make_download_plan(tmp_path) + calls: list[dict[str, Any]] = [] + + class FakeResponse: + def __enter__(self) -> "FakeResponse": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + return None + + def iter_content(self, chunk_size: int): + yield b"abc" + raise requests.HTTPError("403 Forbidden") + + def fake_send_api_request(**kwargs: Any) -> FakeResponse: + calls.append(kwargs) + return FakeResponse() + + monkeypatch.setattr( + "datacollective.download._send_api_request", fake_send_api_request + ) + + with pytest.raises(DownloadError): + _execute_download_plan( + plan, + resume_download_checksum=None, + show_progress=False, + download_source=None, + ) + + assert len(calls) == 1 + + """ Sample downloads """ From c86715356830ca428e2763a760762f79e1fe358b Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 23:00:23 +0200 Subject: [PATCH 4/6] Implement presigned URL persistance --- src/datacollective/download.py | 137 +++++++++++++++++++++++++++++---- 1 file changed, 121 insertions(+), 16 deletions(-) diff --git a/src/datacollective/download.py b/src/datacollective/download.py index dad64f2..dc829fb 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -1,11 +1,13 @@ +import json import os import time +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any import requests from fox_progress_bar import ProgressBar -from pydantic import Field +from pydantic import Field, ValidationError from datacollective.api_utils import ( ENV_DOWNLOAD_PATH, @@ -40,6 +42,11 @@ requests.exceptions.ChunkedEncodingError, ) +# The .checksum sidecar of an interrupted download also keeps the presigned URL and its +# expiry, so a later run can resume from it without opening a new download session (which +# would count towards the API's rate limit). A URL this close to expiring is not reused. +URL_EXPIRY_MARGIN = timedelta(seconds=30) + class DownloadPlan(NonEmptyStrModel): download_url: str @@ -49,6 +56,8 @@ class DownloadPlan(NonEmptyStrModel): # The API returns no checksum for a sample, so a sample download can never be resumed. checksum: str | None checksum_filepath: Path + # When the presigned download URL stops working; None if the API did not say. + expires_at: datetime | None = None def _download_dataset( @@ -66,12 +75,15 @@ def _download_dataset( Flow: 1. Check if the dataset archive already exists in the download directory. If it does and overwrite_existing is False, skip download. - 2. Call the MDC API to get a download session URL and other necessary details. - 3. If overwrite_existing is True, clean up any existing files. - 4. Determine whether to resume download based on existing .checksum and .part files. - If not resuming, write a new .checksum file. - 5. Execute the download plan, downloading to a temporary file. - 6. Once download is complete, rename the temporary file to the target filename and remove the .checksum file. + 2. If an interrupted download left a .part file and a .checksum sidecar whose stored + presigned URL has not expired, resume from that URL without calling the API. + Should the URL be rejected after all, fall through to the next step. + 3. Call the MDC API to get a download session URL and other necessary details. + 4. If overwrite_existing is True, clean up any existing files. + 5. Determine whether to resume download based on existing .checksum and .part files. + 6. Write the .checksum sidecar (checksum, URL, expiry and size) for a later resume. + 7. Execute the download plan, downloading to a temporary file. + 8. Once download is complete, rename the temporary file to the target filename and remove the .checksum file. Args: dataset_id: The unique dataset ID. @@ -95,6 +107,26 @@ def _download_dataset( ) return target_filepath + # Resume an interrupted download with its stored URL, sparing a download-session request + if not overwrite_existing: + stored_plan = _load_stored_download_plan(target_filepath) + if stored_plan is not None: + logger.info("Resuming interrupted download with its stored download URL...") + try: + return _run_download_plan( + download_plan=stored_plan, + show_progress=show_progress, + overwrite_existing=False, + max_download_attempts=max_download_attempts, + download_source=download_source, + ) + except DownloadError as exc: + # The storage refused the stored URL (expired early, revoked, ...): ask + # the API for a fresh one. Any other failure is a real download failure. + if not isinstance(exc.__cause__, (PermissionError, requests.HTTPError)): + raise + logger.info("Stored download URL was rejected. Requesting a new one...") + # Call MDC API to start download session download_plan = _get_download_plan( dataset_id=dataset_id, @@ -209,9 +241,10 @@ def _run_download_plan( resume_checksum = _determine_resume_state(download_plan) - # Write checksum file before starting download (for potential resume later) - if download_plan.checksum and not resume_checksum: - download_plan.checksum_filepath.write_text(download_plan.checksum) + # Persist the resume state before downloading. Rewritten on every run, so a later + # resume holds the newest URL. + if download_plan.checksum: + _write_download_state(download_plan) _execute_download_plan( download_plan=download_plan, @@ -317,6 +350,7 @@ def _build_download_plan( download_url = payload.get("downloadUrl") size_bytes = payload.get("sizeBytes") checksum = payload.get("checksum") + expires_at = payload.get("expiresAt") or None if not download_url or not size_bytes: raise RuntimeError(f"Unexpected response format: {payload}") @@ -334,19 +368,92 @@ def _build_download_plan( size_bytes=size_bytes, checksum=checksum, checksum_filepath=checksum_filepath, + expires_at=expires_at, ) logger.debug( - f"Download plan: full path={target_filepath}, size={size_bytes} bytes, checksum={checksum}", + f"Download plan: full path={target_filepath}, size={size_bytes} bytes, " + f"checksum={checksum}, url expires at {expires_at}", ) return download_plan +def _write_download_state(download_plan: DownloadPlan) -> None: + """ + Write the .checksum sidecar: what a later run needs to resume the .part file. + + Besides the checksum, which tells whether the .part file still belongs to the current + archive, it keeps the presigned URL, its expiry and the archive size, so that a resume + within the URL's lifetime needs no new download session. + """ + state = download_plan.model_dump( + mode="json", include={"checksum", "download_url", "expires_at", "size_bytes"} + ) + download_plan.checksum_filepath.write_text(json.dumps(state)) + + +def _read_download_state(checksum_filepath: Path) -> dict[str, Any]: + """ + Read the .checksum sidecar; {} if there is none or it is unreadable. + + A sidecar written by an older SDK version holds the bare checksum, which is returned + under the "checksum" key so such a download still resumes (via the API). + """ + if not checksum_filepath.exists(): + return {} + text = checksum_filepath.read_text().strip() + try: + state = json.loads(text) + except ValueError: + return {"checksum": text} if text else {} + return state if isinstance(state, dict) else {} + + +def _load_stored_download_plan(target_filepath: Path) -> DownloadPlan | None: + """ + Rebuild the plan of an interrupted download from its sidecar, if it can be resumed + without calling the API: a .part file exists and the sidecar holds a complete state + with a presigned URL that has not expired. Otherwise return None. + """ + tmp_filepath = target_filepath.with_name(target_filepath.name + ".part") + checksum_filepath = _get_checksum_filepath(target_filepath) + if not tmp_filepath.exists(): + return None + + state = _read_download_state(checksum_filepath) + if not state.get("checksum") or not state.get("download_url"): + return None + try: + plan = DownloadPlan( + target_filepath=target_filepath, + tmp_filepath=tmp_filepath, + checksum_filepath=checksum_filepath, + **state, + ) + except (ValidationError, TypeError): + logger.debug(f"Ignoring unreadable download state in `{checksum_filepath}`") + return None + + if not _url_still_valid(plan.expires_at): + logger.debug("Stored download URL has expired") + return None + return plan + + +def _url_still_valid(expires_at: datetime | None) -> bool: + """Whether a presigned URL with this expiry can still be used, with some margin.""" + if expires_at is None: + return False + if expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=timezone.utc) + return datetime.now(timezone.utc) + URL_EXPIRY_MARGIN < expires_at + + def _determine_resume_state(download_plan: DownloadPlan) -> str | None: """ Decide whether a previous partial download can be resumed. A download is resumed only when both the .part and .checksum files exist - and the stored checksum matches the current one. In every other case the + and the checksum stored in the sidecar matches the current one. In every other case the leftover .part / .checksum files are removed and the download starts fresh. A plan without a checksum (a sample) can never be resumed, so any leftover .part file is simply discarded. @@ -368,10 +475,8 @@ def _determine_resume_state(download_plan: DownloadPlan) -> str | None: _cleanup_partial_download(download_plan) return None - checksum_filepath = download_plan.checksum_filepath - stored_checksum = ( - checksum_filepath.read_text().strip() if checksum_filepath.exists() else None - ) + stored = _read_download_state(download_plan.checksum_filepath).get("checksum") + stored_checksum = str(stored).strip() if stored else None if part_exists and stored_checksum: if stored_checksum == download_plan.checksum: From 27acc3397581e762c737be6e8ea87cdc3ce825d8 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 23:05:43 +0200 Subject: [PATCH 5/6] Update docs --- README.md | 2 +- docs/index.md | 32 +++++++++++++++++--------------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 1a06899..1103fce 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,7 @@ dataset_path = download_dataset("your-dataset-id") > `download_dataset` was previously called `save_dataset_to_disk`. The old name still works for backward compatibility, but it is deprecated and new code should use `download_dataset`. > [!TIP] -> **Automatic Resume:** If a download is interrupted (e.g., due to a network error or it gets stopped it manually), the next time you try download the same dataset at the same folder location, we will automatically resume from where the download left off! +> **Automatic Resume:** If a download is interrupted (e.g., due to a network error or it gets stopped manually), the next time you download the same dataset to the same folder, it resumes from where it left off. A connection that drops mid-download is retried automatically (see `max_download_attempts`), and a resume reuses the previous download URL while it is still valid, so neither counts as a new download against the API's rate limit. > [!TIP] > Set `enable_logging=True` to emit detailed SDK logs to the console and a local log file at `~/.mozdata/datacollective.log` with timestamped entries, a per-session id, and retention of 5 backup files at 10 MB each. diff --git a/docs/index.md b/docs/index.md index 2604214..f0969fe 100644 --- a/docs/index.md +++ b/docs/index.md @@ -393,28 +393,20 @@ or `load_dataset` again. **How it works:** 1. When a download starts, the SDK creates a `.checksum` file alongside the partial - download (`.part` file) to track the download state. + download (`.part` file) to track the download state: the archive's checksum, the + presigned download URL and when it expires. 2. If the download is interrupted, both files are preserved. 3. On the next download attempt, the SDK detects the partial download and resumes - from the last byte received. -4. Once the download completes successfully, the temporary files are automatically - cleaned up. + from the last byte received. While the stored download URL is still valid, the + resume reuses it and does not open a new download session, so it does not count + towards the API's download rate limit. Once the URL has expired, or if the storage + rejects it, the SDK requests a new one and still resumes the partial download. +4. Once the download completes successfully, both files are automatically removed. !!! tip You don't need to do anything special to enable resume functionality, it works automatically. Just call the same function again after an interruption. -**Edge cases handled:** - -- If the dataset has been updated since the interrupted download, the SDK detects - the checksum mismatch and starts a fresh download. -- If only partial files exist without proper tracking data, the SDK safely starts - a fresh download. - -!!! note - Sample downloads (`download_sample`) do not resume. The API returns no checksum for - samples, so an interrupted sample download starts again from the beginning. - **Dropped connections are retried in place.** If the connection drops while an archive is streaming, the SDK retries the request against the same download URL, continuing from the bytes already on disk. These retries do not start a new download session, so they do @@ -431,6 +423,16 @@ archive_path = download_dataset("", max_download_attempts=10) Only connection-level failures are retried. HTTP errors and manual cancellation fail immediately, and a full download can still be resumed later as described above. +**Edge cases handled:** + +- If the dataset has been updated since the interrupted download, the SDK detects + the checksum mismatch and starts a fresh download. +- If only partial files exist without proper tracking data, the SDK safely starts + a fresh download. + +!!! note + Sample downloads (`download_sample`) do not resume. The API returns no checksum for + samples, so an interrupted sample download starts again from the beginning. ### Automatically check for extracted archives From ce7c3859b610b5dd765a96000c9446c7b55f785c Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 2 Nov 2026 23:05:54 +0200 Subject: [PATCH 6/6] Update tests --- tests/test_download.py | 275 +++++++++++++++++++++++++++++++++ tests/test_openapi_contract.py | 2 +- 2 files changed, 276 insertions(+), 1 deletion(-) diff --git a/tests/test_download.py b/tests/test_download.py index 6530dad..2be405b 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -1,6 +1,8 @@ from pathlib import Path +import json import logging +from datetime import datetime, timedelta, timezone from typing import Any import pytest @@ -9,7 +11,10 @@ from datacollective.download import ( DEFAULT_MAX_DOWNLOAD_ATTEMPTS, DownloadPlan, + _download_dataset, _download_sample, + _load_stored_download_plan, + _run_download_plan, _get_checksum_filepath, _determine_resume_state, _resolve_download_dir, @@ -407,6 +412,276 @@ def fake_send_api_request(**kwargs: Any) -> FakeResponse: assert len(calls) == 1 +""" +Resuming with the stored presigned URL +""" + + +class _FakeDatasetApi: + """Stands in for `_send_api_request`: answers the download-session POST with *payload* + and any presigned GET with *content* (honouring the Range header), recording every + call. A URL listed in *rejected_urls* answers the GET with a PermissionError instead, + as the storage would an expired or revoked presigned URL.""" + + def __init__( + self, + payload: dict[str, Any], + content: bytes, + rejected_urls: tuple[str, ...] = (), + ) -> None: + self.payload = payload + self.content = content + self.rejected_urls = rejected_urls + self.calls: list[dict[str, Any]] = [] + + @property + def methods(self) -> list[str]: + return [call["method"] for call in self.calls] + + def __call__(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + if kwargs["method"] == "GET" and kwargs["url"] in self.rejected_urls: + raise PermissionError("403 Forbidden") + api = self + range_header = (kwargs.get("extra_headers") or {}).get("Range") + start = ( + int(range_header.removeprefix("bytes=").rstrip("-")) if range_header else 0 + ) + + class FakeResponse: + def json(self) -> dict[str, Any]: + return api.payload + + def __enter__(self) -> "FakeResponse": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + return None + + def iter_content(self, chunk_size: int): + yield api.content[start:] + + return FakeResponse() + + +_CONTENT = b"0123456789" + + +def _future(seconds: int = 3600) -> str: + return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat() + + +def _session_payload( + url: str = "https://example.com/new", checksum: str = "abc" +) -> dict[str, Any]: + return { + "downloadUrl": url, + "sizeBytes": str(len(_CONTENT)), + "checksum": checksum, + "expiresAt": _future(), + } + + +def _write_interrupted_download( + tmp_path: Path, + url: str = "https://example.com/stored", + checksum: str = "abc", + expires_at: str | None = None, + partial: bytes = _CONTENT[:4], +) -> Path: + """Leave behind what an interrupted download of `dataset.tar.gz` leaves: a .part file + and a sidecar with the stored plan. Returns the target path.""" + target = tmp_path / "dataset.tar.gz" + target.with_name(target.name + ".part").write_bytes(partial) + _get_checksum_filepath(target).write_text( + json.dumps( + { + "checksum": checksum, + "download_url": url, + "expires_at": expires_at or _future(), + "size_bytes": len(_CONTENT), + } + ) + ) + return target + + +def _download(tmp_path: Path, **kwargs: Any) -> Path: + return _download_dataset( + dataset_id="dataset-id", + archive_filename="dataset.tar.gz", + download_directory=str(tmp_path), + show_progress=False, + overwrite_existing=kwargs.pop("overwrite_existing", False), + **kwargs, + ) + + +def test_run_download_plan_writes_url_and_expiry_to_sidecar( + tmp_path: Path, monkeypatch +) -> None: + api = _FakeDatasetApi(_session_payload(), _CONTENT) + monkeypatch.setattr("datacollective.download._send_api_request", api) + plan = DownloadPlan.model_validate( + { + **_make_download_plan(tmp_path, checksum="abc").model_dump(), + "size_bytes": len(_CONTENT), + "expires_at": _future(), + } + ) + written: dict[str, Any] = {} + + def capture_then_fail(**kwargs: Any) -> None: + written.update(json.loads(plan.checksum_filepath.read_text())) + raise RuntimeError("stop here") + + monkeypatch.setattr( + "datacollective.download._execute_download_plan", capture_then_fail + ) + + with pytest.raises(RuntimeError, match="stop here"): + _run_download_plan(plan, show_progress=False, overwrite_existing=False) + + assert written == { + "checksum": "abc", + "download_url": plan.download_url, + "expires_at": plan.expires_at.isoformat().replace("+00:00", "Z"), # type: ignore[union-attr] + "size_bytes": len(_CONTENT), + } + + +def test_resume_uses_stored_url_without_calling_the_api( + tmp_path: Path, monkeypatch +) -> None: + target = _write_interrupted_download(tmp_path) + api = _FakeDatasetApi(_session_payload(), _CONTENT) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download(tmp_path) + + assert result == target + assert target.read_bytes() == _CONTENT + assert api.methods == ["GET"] + assert api.calls[0]["url"] == "https://example.com/stored" + assert api.calls[0]["extra_headers"]["Range"] == "bytes=4-" + assert not _get_checksum_filepath(target).exists() + assert not target.with_name(target.name + ".part").exists() + + +def test_resume_asks_the_api_when_stored_url_has_expired( + tmp_path: Path, monkeypatch +) -> None: + target = _write_interrupted_download(tmp_path, expires_at=_future(seconds=10)) + api = _FakeDatasetApi(_session_payload(), _CONTENT) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + _download(tmp_path) + + assert target.read_bytes() == _CONTENT + assert api.methods == ["POST", "GET"] + # The fresh URL is used, and the partial download is still resumed. + assert api.calls[1]["url"] == "https://example.com/new" + assert api.calls[1]["extra_headers"]["Range"] == "bytes=4-" + + +def test_resume_falls_back_to_the_api_when_stored_url_is_rejected( + tmp_path: Path, monkeypatch +) -> None: + target = _write_interrupted_download(tmp_path) + api = _FakeDatasetApi( + _session_payload(), _CONTENT, rejected_urls=("https://example.com/stored",) + ) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + _download(tmp_path) + + assert target.read_bytes() == _CONTENT + assert api.methods == ["GET", "POST", "GET"] + assert api.calls[0]["url"] == "https://example.com/stored" + assert api.calls[2]["url"] == "https://example.com/new" + assert api.calls[2]["extra_headers"]["Range"] == "bytes=4-" + + +def test_resume_with_stored_url_surfaces_other_download_failures( + tmp_path: Path, monkeypatch +) -> None: + _write_interrupted_download(tmp_path) + + def connection_lost(**kwargs: Any) -> Any: + raise requests.ConnectionError("reset") + + monkeypatch.setattr("datacollective.download._send_api_request", connection_lost) + monkeypatch.setattr("datacollective.download.time.sleep", lambda _s: None) + + with pytest.raises(DownloadError): + _download(tmp_path, max_download_attempts=1) + + +def test_overwrite_ignores_stored_url(tmp_path: Path, monkeypatch) -> None: + target = _write_interrupted_download(tmp_path) + api = _FakeDatasetApi(_session_payload(), _CONTENT) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + _download(tmp_path, overwrite_existing=True) + + assert target.read_bytes() == _CONTENT + assert api.methods == ["POST", "GET"] + assert "Range" not in (api.calls[1]["extra_headers"] or {}) + + +def test_legacy_plain_text_sidecar_still_resumes_via_the_api( + tmp_path: Path, monkeypatch +) -> None: + target = tmp_path / "dataset.tar.gz" + target.with_name(target.name + ".part").write_bytes(_CONTENT[:4]) + _get_checksum_filepath(target).write_text("abc") + api = _FakeDatasetApi(_session_payload(), _CONTENT) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + _download(tmp_path) + + assert target.read_bytes() == _CONTENT + assert api.methods == ["POST", "GET"] + assert api.calls[1]["extra_headers"]["Range"] == "bytes=4-" + + +@pytest.mark.parametrize( + "sidecar", + [ + None, + "", + "not json {", + json.dumps(["a", "list"]), + json.dumps({"checksum": "abc"}), + json.dumps({"checksum": "abc", "download_url": "https://x", "size_bytes": 10}), + json.dumps( + { + "checksum": "abc", + "download_url": "https://x", + "size_bytes": 10, + "extra": 1, + } + ), + ], +) +def test_load_stored_download_plan_rejects_incomplete_state( + tmp_path: Path, sidecar: str | None +) -> None: + target = tmp_path / "dataset.tar.gz" + target.with_name(target.name + ".part").write_bytes(b"0123") + if sidecar is not None: + _get_checksum_filepath(target).write_text(sidecar) + + assert _load_stored_download_plan(target) is None + + +def test_load_stored_download_plan_requires_part_file(tmp_path: Path) -> None: + target = _write_interrupted_download(tmp_path) + target.with_name(target.name + ".part").unlink() + + assert _load_stored_download_plan(target) is None + + """ Sample downloads """ diff --git a/tests/test_openapi_contract.py b/tests/test_openapi_contract.py index 31b85ba..83e21a1 100644 --- a/tests/test_openapi_contract.py +++ b/tests/test_openapi_contract.py @@ -318,7 +318,7 @@ def test_dataset_details_fields_exist_in_spec( def test_download_session_fields(schemas: dict[str, Any]) -> None: schema = _resolve(schemas["DownloadSession"], schemas) # download._build_download_plan reads these keys and requires the first two. - assert {"downloadUrl", "sizeBytes"} <= _required(schema) + assert {"downloadUrl", "sizeBytes", "expiresAt"} <= _required(schema) assert "checksum" in _properties(schema, schemas)