diff --git a/python/lightning_sdk/cli/dataset/download.py b/python/lightning_sdk/cli/dataset/download.py index 78937ade..33b1e352 100644 --- a/python/lightning_sdk/cli/dataset/download.py +++ b/python/lightning_sdk/cli/dataset/download.py @@ -23,27 +23,39 @@ "--target_path", default=".", type=click.Path(file_okay=False, dir_okay=True), - help="Local directory to download the dataset zip into.", + help="Local directory to download the dataset into.", +) +@click.option( + "--unzip", + is_flag=True, + default=False, + help="Extract a dataset stored as exactly one .zip artifact.", ) def download_dataset_cmd( name: str, cluster_id: Optional[str] = None, target_path: str = ".", + unzip: bool = False, ) -> None: - """Download a dataset version as a zip file. + """Download a dataset version. NAME must be a Lightning path: // or ///. If no version specified, defaults to most recent version. + By default, files are downloaded into a directory. Pass --unzip to + explicitly extract a version stored as exactly one ZIP artifact. + Usage: lightning dataset download my-org/my-teamspace/my-dataset lightning dataset download my-org/my-teamspace/my-dataset/v3 lightning dataset download my-org/my-teamspace/my-dataset/v3 --target-path ./data + lightning dataset download my-org/my-teamspace/my-dataset --unzip """ info = download_dataset( name=name, target_path=target_path, cluster_id=cluster_id, + unzip=unzip, ) click.echo(f"Downloaded dataset '{info.name}' version '{info.version}' to {info.path}") diff --git a/python/lightning_sdk/datasets.py b/python/lightning_sdk/datasets.py index 8fa5350e..55237b4f 100644 --- a/python/lightning_sdk/datasets.py +++ b/python/lightning_sdk/datasets.py @@ -39,7 +39,7 @@ class DownloadedDatasetInfo: """Metadata returned after a successful dataset download. Attributes: - path: The local file path to the downloaded zip file. + path: The local directory containing the dataset. name: The dataset name. version: The version tag that was downloaded (e.g. ``"v1"``). """ @@ -123,25 +123,34 @@ def download_dataset( cluster_id: Optional[str] = None, num_workers: int = _DEFAULT_DOWNLOAD_WORKERS, part_size: int = _DEFAULT_DOWNLOAD_PART_SIZE, + unzip: bool = False, ) -> DownloadedDatasetInfo: - """Download a dataset version as a zip file from Lightning Datasets. + """Download a dataset version from Lightning Datasets. - Files download concurrently, each fetched via chunked HTTP-Range requests. + By default, files are downloaded into a directory. Set ``unzip=True`` to + safely extract a dataset version stored as exactly one ZIP artifact. Files + download concurrently, each fetched via chunked HTTP-Range requests. Args: name: Lightning path to dataset in the format ``//`` or ``///``. If no version specified, defaults to the most recent version. - target_path: Local directory to download the dataset zip into. + target_path: Local directory to download the dataset into. Defaults to the current directory (``"."``). cluster_id: Optional cluster ID to download from. num_workers: number of concurrent download threads (default 16). part_size: byte-range part size for splitting large files (default 64 MB). + unzip: Extract a version stored as exactly one ZIP artifact into a + directory. This is never enabled automatically. Returns: DownloadedDatasetInfo: Metadata about the downloaded dataset, including the local path, dataset name, and version. + + Raises: + ValueError: If ``unzip=True`` is used with a version that is not exactly + one ZIP artifact. """ org_name, ts_name, dataset_name, version = _parse_dataset_path(name) teamspace = _get_teamspace(name=ts_name, organization=org_name) @@ -154,18 +163,19 @@ def download_dataset( import os - zip_path = os.path.join(target_path, f"{dataset_name}_{version}.zip") + output_path = os.path.join(target_path, f"{dataset_name}_{version}") _download_dataset_version( project_id=project_id, dataset_name=dataset_name, version=version, - target_path=zip_path, + target_path=output_path, cluster_id=cluster_id, dataset_id=dataset_id, num_workers=num_workers, part_size=part_size, + unzip=unzip, ) - return DownloadedDatasetInfo(path=zip_path, name=dataset_name, version=version) + return DownloadedDatasetInfo(path=output_path, name=dataset_name, version=version) def list_datasets(name: str) -> list: diff --git a/python/lightning_sdk/lightning_cloud/utils/dataset.py b/python/lightning_sdk/lightning_cloud/utils/dataset.py index 83e88cdc..0a1bfc62 100644 --- a/python/lightning_sdk/lightning_cloud/utils/dataset.py +++ b/python/lightning_sdk/lightning_cloud/utils/dataset.py @@ -1,27 +1,28 @@ -from typing import Callable, List, Optional, Union import concurrent.futures import json import math import os +import shutil +import tempfile +import zipfile from concurrent.futures import ThreadPoolExecutor from functools import partial -from pathlib import Path +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import Callable, List, Optional, Union import backoff import requests from tqdm.auto import tqdm -from lightning_sdk.lightning_cloud.rest_client import LightningClient -from lightning_sdk.lightning_cloud.openapi.api_client import ApiClient -from lightning_sdk.lightning_cloud import env - from lightning_sdk.api.utils import ( - _BYTES_PER_MB, _MAX_BATCH_SIZE, _MAX_SIZE_MULTI_PART_CHUNK, _MAX_WORKERS, _SIZE_LIMIT_SINGLE_PART, ) +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.openapi.api_client import ApiClient +from lightning_sdk.lightning_cloud.rest_client import LightningClient # Total upload concurrency budget, split across files x within-file parts. _DEFAULT_UPLOAD_WORKERS = 16 @@ -32,6 +33,65 @@ _DEFAULT_DOWNLOAD_PART_SIZE = 64 * 1024 * 1024 # 64 MiB byte-range parts +def _normalize_relative_path(path: str, source: str, strip_leading_slashes: bool = False) -> str: + """Normalize a relative path and reject paths that could escape a destination.""" + if not isinstance(path, str) or not path or "\x00" in path: + raise ValueError(f"Unsafe {source} path {path!r}: expected a non-empty relative path.") + + if strip_leading_slashes: + path = path.lstrip("/") + + # Dataset and ZIP paths use forward slashes, but validate backslashes as + # separators too so the same input cannot become unsafe on Windows. + normalized = path.replace("\\", "/") + posix_path = PurePosixPath(normalized) + windows_path = PureWindowsPath(path) + if ( + posix_path.is_absolute() + or windows_path.is_absolute() + or windows_path.drive + or any(part == ".." for part in posix_path.parts) + ): + raise ValueError( + f"Unsafe {source} path {path!r}: absolute paths and parent-directory traversal are not allowed." + ) + + relative_path = str(posix_path) + if relative_path in {"", "."}: + raise ValueError(f"Unsafe {source} path {path!r}: expected a non-empty relative path.") + return relative_path + + +def _safe_destination_path(destination: str, relative_path: str, source: str) -> Path: + """Return a path contained by destination, accounting for existing symlinks.""" + root = Path(destination).resolve() + target = root.joinpath(*PurePosixPath(relative_path).parts) + try: + target.resolve(strict=False).relative_to(root) + except ValueError as ex: + raise ValueError(f"Unsafe {source} path {relative_path!r}: path escapes the destination.") from ex + return target + + +def _extract_zip_safely(archive_path: str, destination: str) -> None: + """Extract a ZIP after validating every member path.""" + with zipfile.ZipFile(archive_path) as archive: + members = [] + for member in archive.infolist(): + relative_path = _normalize_relative_path(member.filename, "ZIP member") + target = _safe_destination_path(destination, relative_path, "ZIP member") + members.append((member, target)) + + Path(destination).mkdir(parents=True, exist_ok=True) + for member, target in members: + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + continue + target.parent.mkdir(parents=True, exist_ok=True) + with archive.open(member) as source_file, target.open("wb") as target_file: + shutil.copyfileobj(source_file, target_file) + + def _parse_dataset_path(name: str) -> tuple: """Parse a dataset path like 'org/teamspace/dataset/version' into components. @@ -61,9 +121,7 @@ def _parse_dataset_path(name: str) -> tuple: return org_name, ts_name, dataset_name, version -def _resolve_dataset_id_and_version( - project_id: str, dataset_name: str, version: Optional[str] = None -) -> tuple: +def _resolve_dataset_id_and_version(project_id: str, dataset_name: str, version: Optional[str] = None) -> tuple: """List datasets once and return ``(dataset_id, resolved_version)``. Combines name->id and default-version resolution into a single API round-trip @@ -125,8 +183,13 @@ def _download_dataset_files( def _to_entry(indexed_file: tuple) -> tuple: index, file_info = indexed_file + raw_path = file_info.get("filepath") + if raw_path is None: + raw_path = f"file_{index}" + relative_path = _normalize_relative_path(raw_path, "dataset filepath", strip_leading_slashes=True) return ( - file_info.get("filepath", f"file_{index}").lstrip("/"), + relative_path, + _safe_destination_path(dest_dir, relative_path, "dataset filepath"), file_info.get("url"), int(file_info.get("size") or 0), ) @@ -135,21 +198,19 @@ def _to_entry(indexed_file: tuple) -> tuple: total_bytes = 0 urls = {} tasks = [] - # Lazily normalize each entry while pre-allocating its file and building the - # flat byte-range task list, so files_list is traversed only once. - for rel, url, size in entries: + # Pre-allocate each downloadable file and build the flat byte-range task list. + for rel, local_path, url, size in entries: if not url: continue total_bytes += size urls[rel] = url - local_path = os.path.join(dest_dir, rel) - os.makedirs(os.path.dirname(local_path) or dest_dir, exist_ok=True) - with open(local_path, "wb") as f: + local_path.parent.mkdir(parents=True, exist_ok=True) + with local_path.open("wb") as f: if size: f.truncate(size) for start in range(0, size, part_size): end = min(start + part_size, size) - 1 - tasks.append((rel, local_path, start, end)) + tasks.append((rel, str(local_path), start, end)) if not urls: return @@ -193,6 +254,7 @@ def _download_part(task: tuple) -> None: finally: pbar.close() + def _download_dataset_version( project_id: str, dataset_name: str, @@ -202,29 +264,28 @@ def _download_dataset_version( dataset_id: Optional[str] = None, num_workers: int = _DEFAULT_DOWNLOAD_WORKERS, part_size: int = _DEFAULT_DOWNLOAD_PART_SIZE, -): - """ - Download a dataset version as a zip file from the API. + unzip: bool = False, +) -> None: + """Download a dataset version from the API. Fetches presigned file URLs from the files endpoint and downloads the files concurrently, each fetched with chunked HTTP-Range requests (see - ``_download_dataset_files``), then packages them into a zip archive at - ``target_path``. + ``_download_dataset_files``). By default files are written into + ``target_path``. With ``unzip``, the version must contain one ZIP artifact, + which is safely extracted into ``target_path``. Args: project_id: The project ID. dataset_name: The dataset name. version: The dataset version to download. - target_path: Local file path where the downloaded zip will be saved. + target_path: Destination directory. cluster_id: Optional cluster ID. dataset_id: The resolved dataset ID; when provided, skips an extra datasets-list API round-trip to resolve the name. num_workers: number of concurrent download threads (default 16). part_size: byte-range part size for splitting large files (default 64 MB). + unzip: Extract a dataset stored as exactly one ZIP artifact. """ - import shutil - import tempfile - cloud_url = env.LIGHTNING_CLOUD_URL client = LightningClient(retry=False) api_client: ApiClient = client.api_client @@ -262,13 +323,26 @@ def _download_dataset_version( try: files_data = json.loads(resp.data) if resp.data else {} - except json.JSONDecodeError: - raise ValueError(f"Failed to parse response from {files_url}: {resp.data[:200]}") + except json.JSONDecodeError as ex: + raise ValueError(f"Failed to parse response from {files_url}: {resp.data[:200]}") from ex files_list = files_data.get("files", []) if not files_list: raise ValueError(f"No files found for dataset '{dataset_name}' version '{version}' in project '{project_id}'.") + zip_relative_path = None + if unzip: + if len(files_list) != 1: + raise ValueError("`unzip=True` requires the dataset version to contain exactly one .zip artifact.") + raw_path = files_list[0].get("filepath") + if raw_path is None: + raw_path = "file_0" + zip_relative_path = _normalize_relative_path(raw_path, "dataset filepath", strip_leading_slashes=True) + if not zip_relative_path.lower().endswith(".zip"): + raise ValueError("`unzip=True` requires the dataset version to contain exactly one .zip artifact.") + if not files_list[0].get("url"): + raise ValueError("The dataset ZIP artifact does not have a download URL.") + # Re-fetch a fresh presigned URL for one file (used when a URL expires mid-download). def _refresh_file(rel_path: str) -> dict: r = api_client.request( @@ -279,22 +353,30 @@ def _refresh_file(rel_path: str) -> dict: _preload_content=True, ) for f in (json.loads(r.data) if r.data else {}).get("files", []): - if f.get("filepath", "").lstrip("/") == rel_path: + filepath = f.get("filepath") + if ( + filepath is not None + and _normalize_relative_path(filepath, "dataset filepath", strip_leading_slashes=True) == rel_path + ): return {"url": f.get("url"), "size": int(f.get("size") or 0)} return {"url": None, "size": 0} - # Download in parallel into a temp dir, then package into a zip archive. - tmp_dir = tempfile.mkdtemp() + # ZIP extraction stages the archive in a temporary directory. + # Raw directory downloads write directly to the requested output directory. + staging = unzip + dest_dir = tempfile.mkdtemp(prefix="lightning-dataset-") if staging else target_path + if not staging: + os.makedirs(dest_dir, exist_ok=True) try: _download_dataset_files( - files_list, tmp_dir, project_id, _refresh_file, num_workers=num_workers, part_size=part_size + files_list, dest_dir, project_id, _refresh_file, num_workers=num_workers, part_size=part_size ) - base = target_path - if base.endswith(".zip"): - base = base[:-4] - shutil.make_archive(base, "zip", tmp_dir) + if unzip: + archive_path = _safe_destination_path(dest_dir, zip_relative_path, "dataset filepath") + _extract_zip_safely(str(archive_path), target_path) finally: - shutil.rmtree(tmp_dir, ignore_errors=True) + if staging: + shutil.rmtree(dest_dir, ignore_errors=True) def _create_dataset( diff --git a/python/tests/cli/dataset/test_download.py b/python/tests/cli/dataset/test_download.py index 4563eb3a..0c0e25f1 100644 --- a/python/tests/cli/dataset/test_download.py +++ b/python/tests/cli/dataset/test_download.py @@ -1,3 +1,10 @@ +from types import SimpleNamespace +from unittest import mock + +import pytest +from click.testing import CliRunner + +from lightning_sdk.cli.dataset.download import download_dataset_cmd from tests.cli.help import assert_help_contains, mock_command_logging @@ -6,8 +13,9 @@ def test_dataset_download_help() -> None: assert_help_contains( "lightning dataset download --help", "Usage: lightning dataset download", - "Download a dataset version as a zip file.", + "Download a dataset version.", "NAME must be a Lightning path:", + "--unzip", ) @@ -16,7 +24,7 @@ def test_datasets_download_help() -> None: assert_help_contains( "lightning datasets download --help", "Usage: lightning datasets download", - "Download a dataset version as a zip file.", + "Download a dataset version.", ) @@ -38,3 +46,28 @@ def test_dataset_download_invalid_name() -> None: with pytest.raises(ValueError, match="NAME must be a Lightning path"): command_text("lightning dataset download badname") + + +@pytest.mark.parametrize( + ("args", "unzip"), + [ + ([], False), + (["--unzip"], True), + ], +) +def test_dataset_download_unzip_flag_passthrough(args, unzip) -> None: + info = SimpleNamespace(name="my-dataset", version="v1", path="/tmp/my-dataset_v1") + with mock.patch("lightning_sdk.cli.dataset.download.download_dataset", return_value=info) as download: + result = CliRunner().invoke( + download_dataset_cmd, + ["my-org/my-teamspace/my-dataset", *args], + ) + + assert result.exit_code == 0, result.output + download.assert_called_once_with( + name="my-org/my-teamspace/my-dataset", + target_path=".", + cluster_id=None, + unzip=unzip, + ) + assert "Downloaded dataset 'my-dataset' version 'v1' to /tmp/my-dataset_v1" in result.output diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index dd501520..1446a2d4 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -1248,6 +1248,7 @@ def test_download_dataset_version( mock_lightning_client, internal_teamspace_api_list_mocker, internal_user_api_mocker, + tmp_path, ): import json @@ -1276,42 +1277,37 @@ def fake_get(url, headers=None, stream=None, **kwargs): captured["headers"] = headers return _FakeRangeResponse(b"data") - target_path = "/tmp/test_dataset_download.zip" - try: - with ( - mock.patch("requests.get", side_effect=fake_get) as mock_get, - mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), - mock.patch("shutil.make_archive") as mock_make_archive, - ): - _download_dataset_version( - project_id="proj-1", - dataset_name="ds-1", - version="3", - target_path=target_path, - cluster_id="aws-us-east", - ) - - mock_api_client.request.assert_any_call( - "GET", - "https://lightning.ai/v1/projects/proj-1/lit-datasets", - headers=mock.ANY, - _preload_content=True, - ) - mock_api_client.request.assert_any_call( - "GET", - "https://lightning.ai/v1/projects/proj-1/lit-datasets/ds-1/versions/3/files", - query_params={"clusterId": "aws-us-east"}, - headers=mock.ANY, - _preload_content=True, + target_path = tmp_path / "test_dataset_download" + with ( + mock.patch("requests.get", side_effect=fake_get) as mock_get, + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + ): + _download_dataset_version( + project_id="proj-1", + dataset_name="ds-1", + version="3", + target_path=str(target_path), + cluster_id="aws-us-east", ) - # parallel/chunked download issues a Range GET against the presigned URL - assert mock_get.call_count == 1 - assert captured["url"] == "https://presigned.example.com/data.csv" - assert captured["headers"]["Range"] == "bytes=0-3" - mock_make_archive.assert_called_once_with("/tmp/test_dataset_download", "zip", mock.ANY) - finally: - if os.path.exists(target_path): - os.unlink(target_path) + + mock_api_client.request.assert_any_call( + "GET", + "https://lightning.ai/v1/projects/proj-1/lit-datasets", + headers=mock.ANY, + _preload_content=True, + ) + mock_api_client.request.assert_any_call( + "GET", + "https://lightning.ai/v1/projects/proj-1/lit-datasets/ds-1/versions/3/files", + query_params={"clusterId": "aws-us-east"}, + headers=mock.ANY, + _preload_content=True, + ) + # parallel/chunked download issues a Range GET against the presigned URL + assert mock_get.call_count == 1 + assert captured["url"] == "https://presigned.example.com/data.csv" + assert captured["headers"]["Range"] == "bytes=0-3" + assert (target_path / "data.csv").read_bytes() == b"data" @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://lightning.ai"}) @@ -1321,6 +1317,7 @@ def test_download_dataset_version_no_token_no_cluster( mock_lightning_client, internal_teamspace_api_list_mocker, internal_user_api_mocker, + tmp_path, ): import json @@ -1342,37 +1339,32 @@ def test_download_dataset_version_no_token_no_cluster( mock_client_instance.api_client = mock_api_client mock_lightning_client.return_value = mock_client_instance - target_path = "/tmp/test_ds_no_token.zip" - try: - with ( - mock.patch("requests.get", side_effect=lambda *a, **k: _FakeRangeResponse(b"data")), - mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), - mock.patch("shutil.make_archive") as mock_make_archive, - ): - _download_dataset_version( - project_id="proj-1", - dataset_name="ds-2", - version="1", - target_path=target_path, - ) - - mock_api_client.request.assert_any_call( - "GET", - "https://lightning.ai/v1/projects/proj-1/lit-datasets", - headers=mock.ANY, - _preload_content=True, - ) - mock_api_client.request.assert_any_call( - "GET", - "https://lightning.ai/v1/projects/proj-1/lit-datasets/ds-2/versions/1/files", - query_params={}, - headers=mock.ANY, - _preload_content=True, + target_path = tmp_path / "test_ds_no_token" + with ( + mock.patch("requests.get", side_effect=lambda *a, **k: _FakeRangeResponse(b"data")), + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + ): + _download_dataset_version( + project_id="proj-1", + dataset_name="ds-2", + version="1", + target_path=str(target_path), ) - mock_make_archive.assert_called_once_with("/tmp/test_ds_no_token", "zip", mock.ANY) - finally: - if os.path.exists(target_path): - os.unlink(target_path) + + mock_api_client.request.assert_any_call( + "GET", + "https://lightning.ai/v1/projects/proj-1/lit-datasets", + headers=mock.ANY, + _preload_content=True, + ) + mock_api_client.request.assert_any_call( + "GET", + "https://lightning.ai/v1/projects/proj-1/lit-datasets/ds-2/versions/1/files", + query_params={}, + headers=mock.ANY, + _preload_content=True, + ) + assert (target_path / "data.csv").read_bytes() == b"data" @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://lightning.ai"}) diff --git a/python/tests/test_dataset_download.py b/python/tests/test_dataset_download.py new file mode 100644 index 00000000..a4a6876b --- /dev/null +++ b/python/tests/test_dataset_download.py @@ -0,0 +1,276 @@ +import io +import json +import zipfile +from concurrent.futures import Future +from types import SimpleNamespace +from unittest import mock + +import pytest + +from lightning_sdk.datasets import download_dataset +from lightning_sdk.lightning_cloud.utils.dataset import ( + _download_dataset_files, + _download_dataset_version, +) + + +class _RangeResponse: + def __init__(self, payload: bytes): + self.payload = payload + self.status_code = 200 + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def raise_for_status(self): + return None + + def iter_content(self, chunk_size=None): + return iter([self.payload]) + + +class _SyncExecutor: + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def submit(self, function, *args, **kwargs): + future = Future() + try: + future.set_result(function(*args, **kwargs)) + except Exception as ex: + future.set_exception(ex) + return future + + +def _mock_files_api(files): + response = mock.MagicMock() + response.data = json.dumps({"files": files}) + api_client = mock.MagicMock() + api_client.request.return_value = response + api_client.default_headers = {} + client = mock.MagicMock() + client.api_client = api_client + return mock.patch("lightning_sdk.lightning_cloud.utils.dataset.LightningClient", return_value=client) + + +def _range_get(payloads): + def get(url, headers=None, stream=None): + start, end = (int(value) for value in headers["Range"].split("=")[1].split("-")) + return _RangeResponse(payloads[url][start : end + 1]) + + return get + + +def _zip_payload(files): + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as archive: + for path, content in files.items(): + archive.writestr(path, content) + return output.getvalue() + + +@pytest.mark.parametrize("unzip", [False, True]) +def test_download_dataset_output_mode_and_parallel_options(tmp_path, unzip): + teamspace = SimpleNamespace(id="project-id") + with ( + mock.patch("lightning_sdk.datasets._get_teamspace", return_value=teamspace), + mock.patch("lightning_sdk.datasets.raise_access_error_if_not_allowed"), + mock.patch( + "lightning_sdk.datasets._resolve_dataset_id_and_version", + return_value=("dataset-id", "v7"), + ), + mock.patch("lightning_sdk.datasets._download_dataset_version") as download_version, + ): + result = download_dataset( + "my-org/my-teamspace/my-dataset", + target_path=str(tmp_path), + num_workers=3, + part_size=1024, + unzip=unzip, + ) + + expected_path = str(tmp_path / "my-dataset_v7") + assert result.path == expected_path + download_version.assert_called_once_with( + project_id="project-id", + dataset_name="my-dataset", + version="v7", + target_path=expected_path, + cluster_id=None, + dataset_id="dataset-id", + num_workers=3, + part_size=1024, + unzip=unzip, + ) + + +def test_download_dataset_version_defaults_to_directory(tmp_path): + payload = b"raw dataset contents" + files = [{"filepath": "/nested/data.txt", "url": "https://files/data", "size": len(payload)}] + output = tmp_path / "dataset_v1" + + with ( + _mock_files_api(files), + mock.patch("requests.get", side_effect=_range_get({"https://files/data": payload})), + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + ): + _download_dataset_version( + project_id="project-id", + dataset_name="dataset", + version="v1", + target_path=str(output), + dataset_id="dataset-id", + ) + + assert output.is_dir() + assert (output / "nested" / "data.txt").read_bytes() == payload + assert not (tmp_path / "dataset_v1.zip").exists() + + +def test_download_dataset_version_never_auto_extracts(tmp_path): + payload = _zip_payload({"nested/data.txt": b"still archived"}) + files = [{"filepath": "dataset.zip", "url": "https://files/archive", "size": len(payload)}] + output = tmp_path / "dataset_v1" + + with ( + _mock_files_api(files), + mock.patch("requests.get", side_effect=_range_get({"https://files/archive": payload})), + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + ): + _download_dataset_version( + project_id="project-id", + dataset_name="dataset", + version="v1", + target_path=str(output), + dataset_id="dataset-id", + ) + + assert (output / "dataset.zip").read_bytes() == payload + assert not (output / "nested").exists() + + +def test_download_dataset_version_explicit_unzip(tmp_path): + payload = _zip_payload({"nested/data.txt": b"extracted contents"}) + files = [{"filepath": "dataset.zip", "url": "https://files/archive", "size": len(payload)}] + output = tmp_path / "dataset_v1" + staging = tmp_path / "staging" + + with ( + _mock_files_api(files), + mock.patch("requests.get", side_effect=_range_get({"https://files/archive": payload})), + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + mock.patch( + "lightning_sdk.lightning_cloud.utils.dataset.tempfile.mkdtemp", + return_value=str(staging), + ), + ): + _download_dataset_version( + project_id="project-id", + dataset_name="dataset", + version="v1", + target_path=str(output), + dataset_id="dataset-id", + unzip=True, + ) + + assert (output / "nested" / "data.txt").read_bytes() == b"extracted contents" + assert not (output / "dataset.zip").exists() + assert not staging.exists() + + +@pytest.mark.parametrize( + "files", + [ + [{"filepath": "data.txt", "url": "https://files/data", "size": 1}], + [ + {"filepath": "one.zip", "url": "https://files/one", "size": 1}, + {"filepath": "two.zip", "url": "https://files/two", "size": 1}, + ], + ], +) +def test_download_dataset_version_unzip_requires_one_zip(tmp_path, files): + with ( + _mock_files_api(files), + mock.patch("requests.get") as get, + pytest.raises(ValueError, match="exactly one .zip artifact"), + ): + _download_dataset_version( + project_id="project-id", + dataset_name="dataset", + version="v1", + target_path=str(tmp_path / "output"), + dataset_id="dataset-id", + unzip=True, + ) + + get.assert_not_called() + + +def test_download_dataset_files_skips_missing_url(tmp_path): + files = [{"filepath": "missing.bin", "size": 4}] + with mock.patch("requests.get") as get: + _download_dataset_files( + files, + str(tmp_path), + "project-id", + refresh_file=lambda _: {"url": None, "size": 0}, + ) + + get.assert_not_called() + assert not (tmp_path / "missing.bin").exists() + + +@pytest.mark.parametrize("filepath", ["../escape.txt", "/../escape.txt", "C:\\temp\\escape.txt"]) +def test_download_dataset_files_rejects_unsafe_remote_paths(tmp_path, filepath): + files = [{"filepath": filepath, "url": "https://files/data", "size": 4}] + with ( + mock.patch("requests.get") as get, + pytest.raises(ValueError, match="Unsafe dataset filepath"), + ): + _download_dataset_files( + files, + str(tmp_path / "output"), + "project-id", + refresh_file=lambda _: {"url": None, "size": 0}, + ) + + get.assert_not_called() + + +def test_download_dataset_version_rejects_zip_member_traversal(tmp_path): + payload = _zip_payload({"../escape.txt": b"escaped"}) + files = [{"filepath": "dataset.zip", "url": "https://files/archive", "size": len(payload)}] + output = tmp_path / "output" + staging = tmp_path / "staging" + + with ( + _mock_files_api(files), + mock.patch("requests.get", side_effect=_range_get({"https://files/archive": payload})), + mock.patch("concurrent.futures.ThreadPoolExecutor", _SyncExecutor), + mock.patch( + "lightning_sdk.lightning_cloud.utils.dataset.tempfile.mkdtemp", + return_value=str(staging), + ), + pytest.raises(ValueError, match="Unsafe ZIP member path"), + ): + _download_dataset_version( + project_id="project-id", + dataset_name="dataset", + version="v1", + target_path=str(output), + dataset_id="dataset-id", + unzip=True, + ) + + assert not (tmp_path / "escape.txt").exists() + assert not output.exists() + assert not staging.exists()