Allow for resuming downloads without counting for rate limiting

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.

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("<dataset-id>", 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
 

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: