diff --git a/docs/index.md b/docs/index.md index 643af90..62819fd 100644 --- a/docs/index.md +++ b/docs/index.md @@ -113,6 +113,34 @@ dataset = download_dataset("your-dataset-id") The files will be stored under `MDC_DOWNLOAD_PATH` (default `~/.mozdata/datasets`). +### Download a dataset's sample + +Some publishers provide a **sample**: a small excerpt of the dataset for evaluation before +downloading the full archive. Use `download_sample` to fetch it: + +```python +from datacollective import download_sample + +sample_path = download_sample("your-dataset-id") +``` + +The sample archive is saved under `samples//` inside `MDC_DOWNLOAD_PATH`, +apart from the full dataset archives. A sample may share a filename with the full archive, +so keeping them in separate directories ensures one is never mistaken for the other. + +The function returns the path to the sample archive. The archive is not extracted; use +`load_dataset` for the full dataset if you want a loaded DataFrame. + +If the sample is already present, it is returned without downloading again. Pass +`overwrite_existing=True` to download it again and replace the existing file. + +!!! note + Downloading a sample requires an API key whose organization has agreed to the dataset's + terms on the MDC platform, just like downloading the full dataset. Datasets that + publish no sample raise `FileNotFoundError`. To find datasets with a sample, pass + `has_sample=True` to `list_datasets` (see + [Browse and search the dataset catalog](#browse-and-search-the-dataset-catalog)). + ## Programmatic submissions and uploads The SDK supports creating dataset submissions and uploading files with resumable uploads. @@ -368,6 +396,10 @@ or `load_dataset` again. - 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/__init__.py b/src/datacollective/__init__.py index 065b8a4..46a4d97 100644 --- a/src/datacollective/__init__.py +++ b/src/datacollective/__init__.py @@ -2,6 +2,7 @@ from datacollective.datasets import ( download_dataset, + download_sample, get_dataset_details, list_dataset_filters, list_datasets, @@ -33,6 +34,7 @@ __all__ = [ "download_dataset", + "download_sample", "save_dataset_to_disk", "load_dataset", "get_dataset_details", diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index a2d1597..e1dd2a9 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -23,8 +23,10 @@ from datacollective.download import ( DOWNLOAD_SOURCE_DOWNLOAD, DOWNLOAD_SOURCE_LOAD, + DOWNLOAD_SOURCE_SAMPLE, DOWNLOAD_SOURCE_SAVE, _download_dataset, + _download_sample, ) from datacollective.hf_utils import _convert_to_hf, _require_datasets from datacollective.logging_utils import ( @@ -123,6 +125,63 @@ def download_dataset( ) +def download_sample( + dataset_id: str, + download_directory: str | None = None, + show_progress: bool = True, + overwrite_existing: bool = False, + enable_logging: bool = False, +) -> Path: + """ + Download the dataset's sample archive to a local directory and return the archive path. + The sample is a small excerpt of the dataset that its publisher provides for evaluation. + + The archive is saved under `samples//` inside the download directory, so it + can never be mistaken for the full dataset archive even when the two share a filename. + Skips download if the sample the API currently serves is already there, unless + `overwrite_existing=True`. A sample the publisher has since replaced is downloaded under + its new filename, leaving the earlier file in place. + + Unlike `download_dataset`, an interrupted sample download starts again rather than + resuming: the API returns no checksum for samples to match a partial file against. + + Requires an API key whose organization has agreed to the dataset's terms on the MDC + platform. Buying the dataset is not required. + + Args: + dataset_id: The dataset ID (as shown in MDC platform) or slug. + download_directory: Directory under which to save the sample archive. + If None or empty, falls back to env MDC_DOWNLOAD_PATH or default. + show_progress: Whether to show a progress bar during download. + overwrite_existing: Whether to download the sample again, replacing the existing one. + enable_logging: Whether to enable SDK logging to console and a local log file. + + Returns: + Path to the downloaded sample archive. + + Raises: + ValueError: If dataset_id is empty. + FileNotFoundError: If the dataset does not exist, or publishes no sample (404). + ResourceRemovedError: If the dataset is no longer available (410). + PermissionError: If the dataset's terms have not been agreed to or the dataset is + private to another organization (403), or download directory is not writable. + RuntimeError: If rate limit is exceeded (429) or unexpected response format. + requests.HTTPError: For other non-2xx responses. + """ + _enable_logging(enable_logging) + logger.info(f"Downloading sample of dataset {dataset_id}") + + dataset_details = get_dataset_details(dataset_id) + + return _download_sample( + dataset_id=dataset_details.id, + download_directory=download_directory, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=DOWNLOAD_SOURCE_SAMPLE, + ) + + def _download_archive( dataset_id: str, download_directory: str | None, diff --git a/src/datacollective/download.py b/src/datacollective/download.py index e8c9b01..15c7313 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -20,6 +20,12 @@ DOWNLOAD_SOURCE_DOWNLOAD = "download_dataset" DOWNLOAD_SOURCE_SAVE = "save_dataset_to_disk" DOWNLOAD_SOURCE_LOAD = "load_dataset" +DOWNLOAD_SOURCE_SAMPLE = "download_sample" + +# Samples live under /samples// rather than next to the full +# archives: a sample may carry the same filename as its dataset's archive, or as another +# dataset's sample, and sharing a directory would let one be mistaken for the other. +SAMPLES_DIRNAME = "samples" class DownloadPlan(NonEmptyStrModel): @@ -27,9 +33,8 @@ class DownloadPlan(NonEmptyStrModel): target_filepath: Path tmp_filepath: Path size_bytes: int = Field(..., gt=0) - checksum: ( - str | None - ) # We allow None if the API does not return a checksum, but that's generally unexpected. + # The API returns no checksum for a sample, so a sample download can never be resumed. + checksum: str | None checksum_filepath: Path @@ -80,6 +85,93 @@ def _download_dataset( download_source=download_source, ) + return _run_download_plan( + download_plan=download_plan, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=download_source, + ) + + +def _download_sample( + dataset_id: str, + download_directory: str | None, + show_progress: bool, + overwrite_existing: bool, + download_source: str | None = None, +) -> Path: + """ + Download a dataset's sample archive from MDC to a local directory. + + Follows the same flow as `_download_dataset`, with two differences: + + - The sample's filename is only known once the API is called, so the API is asked first + and the existence check is done on the exact filename it returns. Samples are kept in a + directory of their own, `/samples//`, since a sample's + filename may collide with the full archive's or with another dataset's sample. + - The API returns no checksum for a sample, so an interrupted download cannot be matched + to the current sample and starts again instead of resuming. + + Args: + dataset_id: The unique dataset ID. + download_directory: The base directory to download into. + 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. + Returns: + The full path to the downloaded sample archive. + """ + sample_dir = ( + _resolve_download_dir(download_directory) / SAMPLES_DIRNAME / dataset_id + ) + + download_plan = _get_sample_download_plan( + dataset_id=dataset_id, + sample_dir=sample_dir, + download_source=download_source, + ) + target_filepath = download_plan.target_filepath + + if target_filepath.exists() and not overwrite_existing: + logger.info( + f"Skipping download. Sample archive already exists at `{target_filepath}`" + ) + return target_filepath + + # Created only now, so a refused request leaves no empty directory behind. + sample_dir.mkdir(parents=True, exist_ok=True) + + # With no checksum in the plan, `_determine_resume_state` discards any leftover + # .part file and the download starts again, which is all a sample can do. + return _run_download_plan( + download_plan=download_plan, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=download_source, + ) + + +def _run_download_plan( + download_plan: DownloadPlan, + show_progress: bool, + overwrite_existing: bool, + download_source: str | None = None, +) -> Path: + """ + Carry out a download plan: clean up when overwriting, resume where possible, download to + the temporary file, then move it into place. + + Args: + download_plan: The DownloadPlan object with download details. + 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. + + Returns: + The full path to the downloaded archive. + """ + target_filepath = download_plan.target_filepath + # If overwriting, clean up any existing complete or partial download files if overwrite_existing: logger.info( @@ -138,6 +230,62 @@ def _get_download_plan( ) payload: dict[str, Any] = resp.json() + return _build_download_plan(payload, target_filepath) + + +def _get_sample_download_plan( + dataset_id: str, + sample_dir: Path, + download_source: str | None = None, +) -> DownloadPlan: + """ + Send a POST request to the API to receive the download details for a dataset's sample. + + Args: + dataset_id: The dataset ID (as shown in MDC platform). + sample_dir: Resolved, absolute directory to save the sample in. + download_source: Optional context appended to the User-Agent for download analytics. + + Returns: + a DownloadPlan object + + Raises: + FileNotFoundError: If the dataset does not exist or publishes no sample (404). + ResourceRemovedError: If the dataset is no longer available (410). + PermissionError: If access is denied (403). + RuntimeError: If rate limit is exceeded (429) or unexpected response format. + requests.HTTPError: For other non-2xx responses. + """ + sample_url = f"{_get_api_url()}/datasets/{dataset_id}/sample-download" + resp = _send_api_request( + method="POST", url=sample_url, source_function=download_source + ) + + payload: dict[str, Any] = resp.json() + filename = str(payload.get("filename") or "") + # The filename becomes a path on disk, so accept a bare, visible name only. + if Path(filename).name != filename or not filename or filename.startswith("."): + raise RuntimeError(f"Unexpected sample filename in response: {filename!r}") + + return _build_download_plan(payload, sample_dir / filename) + + +def _build_download_plan( + payload: dict[str, Any], target_filepath: Path +) -> DownloadPlan: + """ + Build a DownloadPlan from a download response of the MDC API. + + Args: + payload: The JSON body of the response. + target_filepath: Resolved, absolute full file path to save the download. + + Returns: + a DownloadPlan object + + Raises: + RuntimeError: If the response lacks a download URL or size. + """ download_url = payload.get("downloadUrl") size_bytes = payload.get("sizeBytes") checksum = payload.get("checksum") @@ -172,6 +320,8 @@ def _determine_resume_state(download_plan: DownloadPlan) -> str | None: 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 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. Args: download_plan: The DownloadPlan object with download details. @@ -180,6 +330,16 @@ def _determine_resume_state(download_plan: DownloadPlan) -> str | None: The checksum to use for resumption, or None if starting fresh. """ part_exists = download_plan.tmp_filepath.exists() + + if download_plan.checksum is None: + if part_exists: + logger.info( + "Partial download found but the download cannot be resumed. " + "Starting fresh download..." + ) + _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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 7b71ac7..1461884 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -73,3 +73,33 @@ def test_save_dataset_to_disk_warns_and_reports_its_own_source( "download_source": "save_dataset_to_disk", } ] + + +def test_download_sample_resolves_id_and_reports_its_own_source( + monkeypatch: MonkeyPatch, +) -> None: + calls: list[dict[str, object]] = [] + + def fake_download_sample(**kwargs: object) -> Path: + calls.append(kwargs) + return Path("sample.tar.gz") + + monkeypatch.setattr( + datasets, + "get_dataset_details", + lambda dataset_id: DatasetDetails(id="canonical-id"), + ) + monkeypatch.setattr(datasets, "_download_sample", fake_download_sample) + + result = datasets.download_sample("a-dataset-slug", download_directory="dir") + + assert result == Path("sample.tar.gz") + assert calls == [ + { + "dataset_id": "canonical-id", + "download_directory": "dir", + "show_progress": True, + "overwrite_existing": False, + "download_source": "download_sample", + } + ] diff --git a/tests/test_download.py b/tests/test_download.py index f19fc0c..d418142 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -3,8 +3,11 @@ import logging from typing import Any +import pytest + from datacollective.download import ( DownloadPlan, + _download_sample, _get_checksum_filepath, _determine_resume_state, _resolve_download_dir, @@ -229,3 +232,186 @@ def fake_send_api_request(**kwargs: Any) -> FakeResponse: assert captured["url"] == plan.download_url assert captured["source_function"] == "download_dataset" assert captured["include_auth_headers"] is False + + +""" +Sample downloads +""" + + +class _FakeSampleApi: + """Stands in for `_send_api_request`: answers the sample-download POST with *payload* + and the presigned GET with *content*, recording every call.""" + + def __init__(self, payload: dict[str, Any], content: bytes = b"sample") -> None: + self.payload = payload + self.content = content + self.calls: list[dict[str, Any]] = [] + + def __call__(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + api = self + + 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 + + return FakeResponse() + + +def _sample_payload(filename: str = "sample.tar.gz") -> dict[str, Any]: + return { + "downloadUrl": "https://example.com/sample", + "filename": filename, + "sizeBytes": "6", + "contentType": "application/gzip", + "expiresAt": "2026-08-27T12:00:00.000Z", + } + + +def test_download_sample_saves_under_samples_dir(tmp_path: Path, monkeypatch) -> None: + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", + str(tmp_path), + show_progress=False, + overwrite_existing=False, + download_source="download_sample", + ) + + assert result == tmp_path / "samples" / "ds" / "sample.tar.gz" + assert result.read_bytes() == b"sample" + assert sorted(p.name for p in result.parent.iterdir()) == ["sample.tar.gz"] + + post, get = api.calls + assert post["method"] == "POST" + assert post["url"].endswith("/datasets/ds/sample-download") + assert post["source_function"] == "download_sample" + assert post.get("include_auth_headers", True) is True + assert get["method"] == "GET" + assert get["url"] == "https://example.com/sample" + assert get["include_auth_headers"] is False + + +def test_download_sample_skips_download_when_already_downloaded( + tmp_path: Path, monkeypatch +) -> None: + existing = tmp_path / "samples" / "ds" / "sample.tar.gz" + existing.parent.mkdir(parents=True) + existing.write_bytes(b"cached") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result == existing + assert result.read_bytes() == b"cached" + assert [call["method"] for call in api.calls] == ["POST"] # no GET of the file + + +def test_download_sample_fetches_replaced_sample_without_overwrite( + tmp_path: Path, monkeypatch +) -> None: + """The existence check is on the exact filename the API returns, so an earlier sample + with another name does not mask a sample the publisher has since replaced.""" + sample_dir = tmp_path / "samples" / "ds" + sample_dir.mkdir(parents=True) + (sample_dir / "old-sample.tar.gz").write_bytes(b"old") + (sample_dir / "notes.txt").write_bytes(b"mine") + api = _FakeSampleApi(_sample_payload("new-sample.tar.gz")) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result == sample_dir / "new-sample.tar.gz" + assert result.read_bytes() == b"sample" + assert (sample_dir / "old-sample.tar.gz").read_bytes() == b"old" + assert (sample_dir / "notes.txt").read_bytes() == b"mine" + + +def test_download_sample_restarts_interrupted_download( + tmp_path: Path, monkeypatch, caplog +) -> None: + partial = tmp_path / "samples" / "ds" / "sample.tar.gz.part" + partial.parent.mkdir(parents=True) + partial.write_bytes(b"stale") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + with caplog.at_level(logging.INFO, logger="datacollective.download"): + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result.read_bytes() == b"sample" + assert not partial.exists() + assert "Starting fresh download" in caplog.text + + +def test_download_sample_overwrite_replaces_existing_sample( + tmp_path: Path, monkeypatch +) -> None: + sample_dir = tmp_path / "samples" / "ds" + sample_dir.mkdir(parents=True) + (sample_dir / "sample.tar.gz").write_bytes(b"old") + (sample_dir / "notes.txt").write_bytes(b"mine") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=True + ) + + assert result == sample_dir / "sample.tar.gz" + assert result.read_bytes() == b"sample" + assert [call["method"] for call in api.calls] == ["POST", "GET"] + # Only the sample itself is replaced; other files in the directory are left alone. + assert (sample_dir / "notes.txt").read_bytes() == b"mine" + + +@pytest.mark.parametrize( + "filename", ["../escape.tar.gz", "nested/sample.tar.gz", "", ".hidden.tar.gz"] +) +def test_download_sample_rejects_unsafe_filename( + tmp_path: Path, monkeypatch, filename: str +) -> None: + api = _FakeSampleApi(_sample_payload(filename)) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + with pytest.raises(RuntimeError, match="Unexpected sample filename"): + _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert len(api.calls) == 1 # nothing downloaded + + +def test_download_sample_refused_leaves_no_directory( + tmp_path: Path, monkeypatch +) -> None: + def refuse(**kwargs: Any) -> Any: + raise PermissionError("You must agree to the terms of use") + + monkeypatch.setattr("datacollective.download._send_api_request", refuse) + + with pytest.raises(PermissionError): + _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert not (tmp_path / "samples").exists() diff --git a/tests/test_openapi_contract.py b/tests/test_openapi_contract.py index da70906..31b85ba 100644 --- a/tests/test_openapi_contract.py +++ b/tests/test_openapi_contract.py @@ -50,6 +50,7 @@ SDK_OPERATIONS = { ("get", "/datasets/{datasetId}"), ("post", "/datasets/{datasetId}/download"), + ("post", "/datasets/{datasetId}/sample-download"), ("post", "/submissions"), ("get", "/submissions/{submissionId}"), ("patch", "/submissions/{submissionId}"), @@ -321,6 +322,12 @@ def test_download_session_fields(schemas: dict[str, Any]) -> None: assert "checksum" in _properties(schema, schemas) +def test_sample_download_fields(schemas: dict[str, Any]) -> None: + schema = _resolve(schemas["SampleDownload"], schemas) + # download._get_sample_download_plan reads these keys and requires all three. + assert {"downloadUrl", "sizeBytes", "filename"} <= _required(schema) + + """ Check upload session spec """