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 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: 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)