diff --git a/README.md b/README.md index b9c7b75..36daf6a 100644 --- a/README.md +++ b/README.md @@ -158,6 +158,9 @@ print(response) For predefined licenses, pass `licenseAbbreviation=License.` and leave `licenseUrl` and `license` unset. For custom licenses, pass a custom string to `license` and optionally include `licenseUrl` and `licenseAbbreviation`. +> [!TIP] +> To also attach an optional sample of your dataset, pass `sample_file_path="/path/to/dataset-sample.tar.gz"` to `create_submission_with_upload`, or upload it separately with `upload_sample_file(file_path=..., submission_id=...)`. + > [!TIP] > To upload a new `.tar.gz` version to an already approved dataset, call `upload_dataset_file(file_path=..., submission_id=...)` directly. Find the submission under **Profile → Uploads**, open the approved dataset, and copy the value after `/profile/submissions/` in the URL. Note that this value is the submission ID, which is different from the public dataset ID. diff --git a/docs/demo_upload.py b/docs/demo_upload.py index d0d08a4..0a08169 100644 --- a/docs/demo_upload.py +++ b/docs/demo_upload.py @@ -44,6 +44,9 @@ response = create_submission_with_upload( file_path="example_dataset.tar.gz", submission=submission, + enable_logging=True, + # Optional: a small, representative excerpt of the dataset + # sample_file_path="example_dataset_sample.tar.gz", ) print(response) diff --git a/docs/index.md b/docs/index.md index 19eabb3..2a1db39 100644 --- a/docs/index.md +++ b/docs/index.md @@ -174,6 +174,9 @@ print(response) For predefined licenses, pass `licenseAbbreviation=License.` and leave `licenseUrl` and `license` unset. For a custom license, pass a custom string to `license` and optionally include `licenseUrl` and `licenseAbbreviation`. +> [!TIP] +> To also attach an optional sample of your dataset, pass `sample_file_path="/path/to/dataset-sample.tar.gz"` to `create_submission_with_upload`, or upload it separately with `upload_sample_file(file_path=..., submission_id=...)`. + > [!TIP] > To upload a new `.tar.gz` version to an already approved and published dataset, call `upload_dataset_file(file_path=..., submission_id=...)` directly. Get the submission ID from **Profile → Uploads** by opening the approved dataset and copying the value after `/profile/submissions/` in the URL. This submission ID is different from the dataset ID. diff --git a/docs/upload.md b/docs/upload.md index 1249b2a..dff4a7f 100644 --- a/docs/upload.md +++ b/docs/upload.md @@ -9,7 +9,8 @@ The SDK provides a complete workflow for uploading datasets: 1. **Create a draft submission** - Initialize a new dataset submission 2. **Update submission metadata** - Add required metadata fields to the submission 3. **Upload the dataset file** - Upload your archive using resumable multipart uploads -4. **Submit for review** - Finalize the submission for review +4. **Upload a sample file** (optional) - Upload a small, representative excerpt of the dataset +5. **Submit for review** - Finalize the submission for review Additionally, the functionality of uploading a dataset file can be used independently to upload a new archive version to an already approved and published dataset submission. Check the [Upload a New File Version to an Approved Dataset](#upload-a-new-file-version-to-an-approved-dataset) section for more details. @@ -98,6 +99,18 @@ print(response) For predefined licenses, pass `licenseAbbreviation=License.` and leave `licenseUrl` and `license` unset. For a custom license, pass a custom string to `license` and optionally include `licenseUrl` and `licenseAbbreviation`. +To also attach an **optional** sample file, pass `sample_file_path`: + +```python +response = create_submission_with_upload( + file_path="/path/to/dataset.tar.gz", + submission=submission, + sample_file_path="/path/to/dataset-sample.tar.gz", +) +``` + +See [Uploading a Sample File](#uploading-a-sample-file) for details. + ### Visibility `visibility` controls who can access the dataset and must be one of the `Visibility` enum values: @@ -108,6 +121,56 @@ For predefined licenses, pass `licenseAbbreviation=License.` and leave `l | `Visibility.PRIVATE` | Everyone | Your organization & Approved requesters only | | `Visibility.RESTRICTED` | Your organization | Your organization (via SDK) | +## Uploading a Sample File + +A **sample file** is a small, representative excerpt of your dataset that reviewers and +potential users can inspect without downloading the full archive. It is **optional**: a +submission can be submitted for review without one, and uploading a sample never replaces +the dataset file itself. + +Sample files are uploaded exactly like the dataset archive — resumable multipart upload of +a `.tar.gz` archive (`application/gzip`) — through the submission's sample endpoints. + +Either pass `sample_file_path` to `create_submission_with_upload`: + +```python +from datacollective import create_submission_with_upload + +response = create_submission_with_upload( + file_path="/path/to/dataset.tar.gz", + submission=submission, + sample_file_path="/path/to/dataset-sample.tar.gz", + # Optional, defaults to `.mdc-sample-upload.json` next to the sample + sample_state_path="/custom/path/sample-upload-state.json", +) + +print(response["submission"]["sampleFileReferenceId"]) +``` + +Or upload it on its own with `upload_sample_file`, using the submission ID (this works for +draft submissions as well as already approved ones): + +```python +from datacollective import upload_sample_file + +upload_state = upload_sample_file( + file_path="/path/to/dataset-sample.tar.gz", + submission_id=submission_id, +) + +print(f"Sample upload complete! File Upload ID: {upload_state.fileUploadId}") +``` + +`upload_sample_file` accepts the same arguments as `upload_dataset_file` +(`state_path`, `show_progress`, `enable_logging`, `part_size`) and is equally resumable. +Its state file uses a `.mdc-sample-upload.json` suffix, so a sample upload and a dataset +upload never overwrite each other's resume state. + +> [!NOTE] +> When using `create_submission_with_upload`, the dataset archive is uploaded first and the +> sample file right after, before the submission is sent for review. A missing +> `sample_file_path` raises `FileNotFoundError` up front, before anything is uploaded. + ## Upload a New File Version to an Approved Dataset Use `upload_dataset_file` when the dataset already exists on the platform and is already in the **Published / Approved** state. @@ -239,7 +302,20 @@ print(f"Upload complete! File Upload ID: {upload_state.fileUploadId}") > [!TIP] > You can also find your submission ID by going to your [Uploads](https://mozilladatacollective.com/profile/uploads) in your profile, click on the dataset submission of your choice, and the URL will contain the submission ID (e.g., `https://mozilladatacollective.com/submissions/cmmjpewijXXXXXXXXX`). -### Step 4: Submit for Review +### Step 4 (Optional): Upload a Sample File + +```python +from datacollective import upload_sample_file + +sample_state = upload_sample_file( + file_path="/path/to/your/dataset-sample.tar.gz", + submission_id=submission_id, +) + +print(f"Sample upload complete! File Upload ID: {sample_state.fileUploadId}") +``` + +### Step 5: Submit for Review ```python from datacollective import DatasetSubmission, submit_submission @@ -373,4 +449,5 @@ For detailed API documentation, see the [API Reference](api.md) section. - [`create_submission_draft`](api.md) - Create a draft submission - [`update_submission`](api.md) - Update submission metadata - [`upload_dataset_file`](api.md) - Upload a file to a submission +- [`upload_sample_file`](api.md) - Upload an optional sample file to a submission - [`submit_submission`](api.md) - Submit a draft for review diff --git a/src/datacollective/__init__.py b/src/datacollective/__init__.py index 54f0297..2dc4188 100644 --- a/src/datacollective/__init__.py +++ b/src/datacollective/__init__.py @@ -23,7 +23,7 @@ submit_submission, update_submission, ) -from datacollective.upload import upload_dataset_file +from datacollective.upload import upload_dataset_file, upload_sample_file __all__ = [ "download_dataset", @@ -35,6 +35,7 @@ "submit_submission", "create_submission_with_upload", "upload_dataset_file", + "upload_sample_file", "DatasetDetails", "DatasetSubmission", "License", diff --git a/src/datacollective/models.py b/src/datacollective/models.py index 80ad5a9..656eb9a 100644 --- a/src/datacollective/models.py +++ b/src/datacollective/models.py @@ -231,6 +231,10 @@ class DatasetSubmission(NonEmptyStrModel, Dataset): None, description="Identifier for the associated file upload, if any. Generated by the API when a file is uploaded.", ) + sampleFileReferenceId: str | None = Field( + None, + description="Identifier for the associated sample file, if any. Generated by the API when a sample file is uploaded.", + ) exclusivityOptOutAt: str | None = Field( None, description="Timestamp when exclusivity opt-out was set, if applicable." ) diff --git a/src/datacollective/submissions.py b/src/datacollective/submissions.py index 87108ea..49e69d0 100644 --- a/src/datacollective/submissions.py +++ b/src/datacollective/submissions.py @@ -1,5 +1,6 @@ from __future__ import annotations +from pathlib import Path from typing import Any from datacollective.api_utils import _get_api_url, _send_api_request @@ -17,7 +18,7 @@ _validate_final_submission_fields, SUBMIT_FIELDS, ) -from datacollective.upload import upload_dataset_file +from datacollective.upload import upload_dataset_file, upload_sample_file from datacollective.upload_utils import _resolve_upload_state, DEFAULT_PART_SIZE logger = get_logger(__name__) @@ -104,6 +105,8 @@ def create_submission_with_upload( state_path: str | None = None, enable_logging: bool = False, part_size: int = DEFAULT_PART_SIZE, + sample_file_path: str | None = None, + sample_state_path: str | None = None, ) -> dict[str, Any]: """ Single point function to create a submission, update metadata, upload a file, and submit for review. @@ -116,6 +119,9 @@ def create_submission_with_upload( enable_logging: Whether to enable detailed logging during the process. part_size: Multipart part size in bytes. Ignored when resuming an existing upload, which keeps the part size recorded in its state file. + sample_file_path: Optional path to a sample archive to upload alongside the + dataset archive. A sample file is not required to submit a dataset. + sample_state_path: Optional path to persist the sample upload state. """ _enable_logging(enable_logging) @@ -123,6 +129,10 @@ def create_submission_with_upload( _validate_final_submission_fields(submission, require_file_upload_id=False) + # Fail fast on a missing sample file, before uploading the dataset archive + if sample_file_path and not Path(sample_file_path).exists(): + raise FileNotFoundError(f"Sample file not found: `{sample_file_path}`") + state_file, existing_upload_state = _resolve_upload_state(file_path, state_path) if existing_upload_state: @@ -158,6 +168,16 @@ def create_submission_with_upload( part_size=part_size, ) + if sample_file_path: + logger.info("Uploading sample file...") + upload_sample_file( + file_path=sample_file_path, + submission_id=submission_id, + state_path=sample_state_path, + enable_logging=enable_logging, + part_size=part_size, + ) + # The uploaded file is linked to the submission automatically when the # multipart upload completes (the upload was started with `submissionId`), # so `fileUploadId` is not sent on the metadata PATCH. We still record it on diff --git a/src/datacollective/upload.py b/src/datacollective/upload.py index 6d72ae7..d5ee24a 100644 --- a/src/datacollective/upload.py +++ b/src/datacollective/upload.py @@ -51,6 +51,76 @@ def upload_dataset_file( part_size: Multipart part size in bytes. Ignored when resuming an existing upload, which keeps the part size recorded in its state file. """ + return _upload_file( + file_path=file_path, + submission_id=submission_id, + state_path=state_path, + show_progress=show_progress, + enable_logging=enable_logging, + part_size=part_size, + is_sample=False, + ) + + +def upload_sample_file( + file_path: str, + submission_id: str, + state_path: str | None = None, + show_progress: bool = True, + enable_logging: bool = False, + part_size: int = DEFAULT_PART_SIZE, +) -> UploadState: + """ + Upload an **optional** sample file for a dataset submission. + + A sample file is a small, representative excerpt of the dataset that + users can inspect without downloading the full archive. It is uploaded + exactly like the dataset archive (resumable multipart upload, + `application/gzip` MIME type) but through the submission's sample endpoints, + and it does not replace the dataset file. + + Args: + file_path: Path to the sample archive on disk. + submission_id: Dataset submission ID (not the dataset ID). + state_path: Optional path to persist upload state. Defaults to + `.mdc-sample-upload.json` alongside the archive. + enable_logging: Whether to enable detailed logging during the upload. + show_progress: Whether to show a progress bar during upload. + part_size: Multipart part size in bytes. Ignored when resuming an + existing upload, which keeps the part size recorded in its state file. + """ + return _upload_file( + file_path=file_path, + submission_id=submission_id, + state_path=state_path, + show_progress=show_progress, + enable_logging=enable_logging, + part_size=part_size, + is_sample=True, + ) + + +def _upload_file( + file_path: str, + submission_id: str, + state_path: str | None, + show_progress: bool, + enable_logging: bool, + part_size: int, + is_sample: bool, +) -> UploadState: + """ + Shared multipart upload function for the dataset archive and the sample file. + + Args: + file_path: Path to the archive on disk. + submission_id: Dataset submission ID (not the dataset ID). + state_path: Optional path to persist upload state. + show_progress: Whether to show a progress bar during upload. + enable_logging: Whether to enable detailed logging during the upload. + part_size: Multipart part size in bytes. + is_sample: Whether to upload the file as the submission's sample file. + """ path = Path(file_path) _enable_logging(enable_logging) @@ -65,7 +135,9 @@ def upload_dataset_file( final_filename = path.name - state_file = Path(state_path) if state_path else _default_state_path(path) + state_file = ( + Path(state_path) if state_path else _default_state_path(path, is_sample) + ) state = _load_or_create_state( state_file=state_file, @@ -73,6 +145,7 @@ def upload_dataset_file( final_filename=final_filename, file_size=file_size, part_size=part_size, + is_sample=is_sample, ) expected_parts = _expected_parts(state.fileSize, state.partSize) @@ -83,7 +156,7 @@ def upload_dataset_file( f"Resuming: {len(parts_by_number)}/{expected_parts} parts already uploaded." ) - logger.info(f"Uploading file: {final_filename}") + logger.info(f"Uploading: {final_filename}") progress_bar = _init_progress_bar( show_progress=show_progress, @@ -122,7 +195,14 @@ def upload_dataset_file( logger.info("Completing upload...") - _complete_upload(state.fileUploadId, state.uploadId, state.parts, state.checksum) + _complete_upload( + state.fileUploadId, + state.uploadId, + state.parts, + state.checksum, + state.submissionId, + state.isSample, + ) logger.info(f"Upload complete. File upload ID: {state.fileUploadId}") diff --git a/src/datacollective/upload_utils.py b/src/datacollective/upload_utils.py index cc420da..b436d8f 100644 --- a/src/datacollective/upload_utils.py +++ b/src/datacollective/upload_utils.py @@ -29,6 +29,11 @@ DEFAULT_PART_SIZE = 10 * 1024 * 1024 # 10 MB default part size to upload chunk by chunk DEFAULT_MIME_TYPE = "application/gzip" +# Suffixes of the resumable state files, kept separate so that uploading a +# dataset archive and a sample file never share the same state file. +STATE_FILE_SUFFIX = ".mdc-upload.json" +SAMPLE_STATE_FILE_SUFFIX = ".mdc-sample-upload.json" + # Storage requires every part except the last to be at least 5 MB MINIMUM_PART_SIZE = 5 * 1024 * 1024 # Storage caps a multipart upload at 10.000 presigned parts @@ -81,6 +86,7 @@ class UploadState(NonEmptyStrModel): mimeType: str parts: list[UploadPart] = Field(default_factory=list) checksum: str | None = None + isSample: bool = False class PresignedPartUrl(NonEmptyStrModel): @@ -108,8 +114,30 @@ class _CompleteUploadPayload(NonEmptyStrModel): checksum: str +def _upload_base_url(submission_id: str, is_sample: bool) -> str: + """ + Base URL of the multipart upload endpoints. + + The dataset archive and the optional sample file are uploaded the same way, + only through different endpoints. + + Args: + submission_id: Dataset submission ID. + is_sample: Whether the upload targets the sample file endpoints instead + of the dataset archive ones. + """ + if is_sample: + return f"{_get_api_url()}/submissions/{submission_id}/sample" + return f"{_get_api_url()}/uploads" + + def _initiate_upload( - submission_id: str, filename: str, file_size: int, mime_type: str, part_size: int + submission_id: str, + filename: str, + file_size: int, + mime_type: str, + part_size: int, + is_sample: bool = False, ) -> UploadSession: """ Start a multipart upload for a dataset submission. @@ -120,6 +148,7 @@ def _initiate_upload( file_size: Size of the file in bytes. mime_type: MIME type for the file. part_size: Multipart part size in bytes to use for this upload. + is_sample: Whether to upload the file as the submission's sample file. """ payload = _UploadInitiatePayload( submissionId=submission_id, @@ -127,7 +156,7 @@ def _initiate_upload( fileSize=file_size, mimeType=mime_type, ) - url = f"{_get_api_url()}/uploads" + url = _upload_base_url(submission_id, is_sample) resp = _send_api_request("POST", url, json_body=payload.model_dump()) data = dict(resp.json()) session_payload = { @@ -141,16 +170,24 @@ def _initiate_upload( raise RuntimeError("Upload initiation did not return expected fields") from exc -def _get_presigned_part_url(file_upload_id: str, part_number: int) -> PresignedPartUrl: +def _get_presigned_part_url( + file_upload_id: str, + part_number: int, + submission_id: str, + is_sample: bool = False, +) -> PresignedPartUrl: """ Request a presigned URL for a specific multipart part. Args: file_upload_id: File upload ID. part_number: 1-based multipart part number. + submission_id: Dataset submission ID. + is_sample: Whether the part belongs to a sample file upload. """ request = _PresignedPartRequest(fileUploadId=file_upload_id, partNumber=part_number) - url = f"{_get_api_url()}/uploads/{request.fileUploadId}/parts/{request.partNumber}" + base_url = _upload_base_url(submission_id, is_sample) + url = f"{base_url}/{request.fileUploadId}/parts/{request.partNumber}" resp = _send_api_request("GET", url) data = dict(resp.json()) presigned_url = data.get("url") or data.get("presignedUrl") @@ -167,6 +204,8 @@ def _complete_upload( upload_id: str | None, parts: list[UploadPart], checksum: str, + submission_id: str, + is_sample: bool = False, ) -> dict[str, Any]: """ Complete a multipart upload and persist the checksum. @@ -178,7 +217,8 @@ def _complete_upload( checksum=checksum, ) - url = f"{_get_api_url()}/uploads/{request.fileUploadId}" + base_url = _upload_base_url(submission_id, is_sample) + url = f"{base_url}/{request.fileUploadId}" payload = { "parts": [part.model_dump() for part in request.parts], "checksum": request.checksum, @@ -205,8 +245,9 @@ def _save_upload_state(path: Path, state: UploadState) -> None: path.write_text(json.dumps(state.model_dump(), indent=2)) -def _default_state_path(file_path: Path) -> Path: - return file_path.with_name(file_path.name + ".mdc-upload.json") +def _default_state_path(file_path: Path, is_sample: bool = False) -> Path: + suffix = SAMPLE_STATE_FILE_SUFFIX if is_sample else STATE_FILE_SUFFIX + return file_path.with_name(file_path.name + suffix) def _load_or_create_state( @@ -215,10 +256,13 @@ def _load_or_create_state( final_filename: str, file_size: int, part_size: int, + is_sample: bool = False, ) -> UploadState: state = _load_upload_state(state_file) if state: - if not _state_matches(state, submission_id, final_filename, file_size): + if not _state_matches( + state, submission_id, final_filename, file_size, is_sample + ): logger.warning( "Upload state does not match file or submission. Restarting upload." ) @@ -231,7 +275,12 @@ def _load_or_create_state( f"Initiating upload for '{final_filename}' ({_format_bytes(file_size)})..." ) session = _initiate_upload( - submission_id, final_filename, file_size, DEFAULT_MIME_TYPE, part_size + submission_id, + final_filename, + file_size, + DEFAULT_MIME_TYPE, + part_size, + is_sample, ) state = UploadState( submissionId=submission_id, @@ -243,6 +292,7 @@ def _load_or_create_state( mimeType=DEFAULT_MIME_TYPE, parts=[], checksum=None, + isSample=is_sample, ) _save_upload_state(state_file, state) @@ -250,13 +300,18 @@ def _load_or_create_state( def _state_matches( - state: UploadState, submission_id: str, filename: str, file_size: int + state: UploadState, + submission_id: str, + filename: str, + file_size: int, + is_sample: bool = False, ) -> bool: return ( state.fileSize == file_size and state.filename == filename and state.submissionId == submission_id and state.mimeType == DEFAULT_MIME_TYPE + and state.isSample == is_sample ) @@ -297,7 +352,12 @@ def _upload_missing_parts( if part_number in parts_by_number: continue - presigned = _get_presigned_part_url(state.fileUploadId, part_number) + presigned = _get_presigned_part_url( + state.fileUploadId, + part_number, + state.submissionId, + state.isSample, + ) response = _upload_part_with_retry(presigned.url, chunk) etag = _extract_etag(response) parts_by_number[part_number] = etag @@ -342,10 +402,12 @@ def _upload_part(presigned_url: str, payload: bytes) -> requests.Response: def _resolve_upload_state( - file_path: str, state_path: str | None + file_path: str, state_path: str | None, is_sample: bool = False ) -> tuple[Path, Any | None]: state_file = ( - Path(state_path) if state_path else _default_state_path(Path(file_path)) + Path(state_path) + if state_path + else _default_state_path(Path(file_path), is_sample) ) return state_file, _load_upload_state(state_file) diff --git a/tests/e2e/test_upload_e2e.py b/tests/e2e/test_upload_e2e.py index 5ca5f9c..9989048 100644 --- a/tests/e2e/test_upload_e2e.py +++ b/tests/e2e/test_upload_e2e.py @@ -3,7 +3,7 @@ from pathlib import Path from datacollective.submissions import create_submission_with_upload -from datacollective.upload import upload_dataset_file +from datacollective.upload import upload_dataset_file, upload_sample_file from tests.e2e.conftest import sample_dataset_submission, skip_if_rate_limited @@ -37,6 +37,66 @@ def test_create_submission_with_upload( ) +def test_create_submission_with_upload_including_sample_file( + tmp_path: Path, + live_api_env: None, + example_dataset_archive_path: Path, +) -> None: + name = f"python-sdk-e2e-sample-{datetime.now().strftime('%H:%M - %d/%m/%Y')}" + state_path = tmp_path / "sample-flow-upload-state.json" + sample_state_path = tmp_path / "sample-flow-sample-upload-state.json" + submission = sample_dataset_submission(name=name) + + try: + response = create_submission_with_upload( + file_path=str(example_dataset_archive_path), + submission=submission, + state_path=str(state_path), + enable_logging=True, + # The archive doubles as the sample file for this live test + sample_file_path=str(example_dataset_archive_path), + sample_state_path=str(sample_state_path), + ) + except Exception as exc: # noqa: BLE001 + skip_if_rate_limited(exc) + else: + submission_payload = response.get("submission", {}) + + assert isinstance(submission_payload, dict) + assert submission_payload.get("fileUploadId") + assert submission_payload.get("sampleFileReferenceId") + assert not sample_state_path.exists(), ( + "Sample upload state should be cleaned up after success" + ) + + +def test_upload_sample_file_for_approved_submission( + tmp_path: Path, + live_api_env: None, + example_dataset_archive_path: Path, + approved_dataset_submission_id: str, +) -> None: + state_path = tmp_path / "approved-sample-upload-state.json" + upload_state = None + + try: + upload_state = upload_sample_file( + file_path=str(example_dataset_archive_path), + submission_id=approved_dataset_submission_id, + state_path=str(state_path), + show_progress=False, + ) + except Exception as exc: # noqa: BLE001 + skip_if_rate_limited(exc) + + assert upload_state is not None + assert upload_state.isSample is True + assert upload_state.submissionId == approved_dataset_submission_id + assert upload_state.fileUploadId + assert upload_state.checksum + assert not state_path.exists(), "Upload state should be cleaned up after success" + + def test_upload_dataset_file_updates_approved_dataset_version( tmp_path: Path, live_api_env: None, diff --git a/tests/test_submissions.py b/tests/test_submissions.py index 5ed2b3e..11930d2 100644 --- a/tests/test_submissions.py +++ b/tests/test_submissions.py @@ -121,3 +121,118 @@ def test_create_submission_with_upload_rejects_missing_required_metadata_before_ assert "`longDescription`" in message assert "either `licenseAbbreviation` or `license`" in message assert "`fileUploadId`" not in message + + +def _stub_submission_workflow( + monkeypatch, sample_uploads: list[dict[str, object]] +) -> None: + """Stub out every network step of `create_submission_with_upload`.""" + + class FakeUploadState: + fileUploadId = "file-upload-id" + + monkeypatch.setattr( + submissions_module, + "create_submission_draft", + lambda submission: {"submission": {"id": "submission-id"}}, + ) + monkeypatch.setattr( + submissions_module, + "update_submission", + lambda submission_id, submission: {"submission": {"id": submission_id}}, + ) + monkeypatch.setattr( + submissions_module, + "upload_dataset_file", + lambda **kwargs: FakeUploadState(), + ) + + def fake_upload_sample_file(**kwargs: object) -> FakeUploadState: + sample_uploads.append(kwargs) + return FakeUploadState() + + monkeypatch.setattr( + submissions_module, "upload_sample_file", fake_upload_sample_file + ) + monkeypatch.setattr( + submissions_module, + "submit_submission", + lambda submission_id, submission: { + "submission": {"id": submission_id, "status": "submitted"} + }, + ) + + +def test_create_submission_with_upload_uploads_optional_sample_file( + tmp_path: Path, monkeypatch +) -> None: + archive_path = tmp_path / "dataset.tar.gz" + archive_path.write_bytes(bytearray(b"dataset-payload")) + sample_path = tmp_path / "sample.tar.gz" + sample_path.write_bytes(bytearray(b"sample-payload")) + sample_state_path = tmp_path / "sample-upload-state.json" + sample_uploads: list[dict[str, object]] = [] + _stub_submission_workflow(monkeypatch, sample_uploads) + + response = submissions_module.create_submission_with_upload( + file_path=str(archive_path), + submission=_build_complete_submission(file_upload_id=None), + sample_file_path=str(sample_path), + sample_state_path=str(sample_state_path), + ) + + assert response["submission"]["status"] == "submitted" + assert sample_uploads == [ + { + "file_path": str(sample_path), + "submission_id": "submission-id", + "state_path": str(sample_state_path), + "enable_logging": False, + "part_size": submissions_module.DEFAULT_PART_SIZE, + } + ] + + +def test_create_submission_with_upload_skips_sample_upload_by_default( + tmp_path: Path, monkeypatch +) -> None: + archive_path = tmp_path / "dataset.tar.gz" + archive_path.write_bytes(bytearray(b"dataset-payload")) + sample_uploads: list[dict[str, object]] = [] + _stub_submission_workflow(monkeypatch, sample_uploads) + + submissions_module.create_submission_with_upload( + file_path=str(archive_path), + submission=_build_complete_submission(file_upload_id=None), + ) + + assert sample_uploads == [] + + +def test_create_submission_with_upload_rejects_missing_sample_file_before_upload( + tmp_path: Path, monkeypatch +) -> None: + archive_path = tmp_path / "dataset.tar.gz" + archive_path.write_bytes(bytearray(b"dataset-payload")) + + monkeypatch.setattr( + submissions_module, + "create_submission_draft", + lambda submission: pytest.fail( + "create_submission_with_upload should fail before creating a draft" + ), + ) + monkeypatch.setattr( + submissions_module, + "upload_dataset_file", + lambda *args, **kwargs: pytest.fail( + "create_submission_with_upload should fail before uploading" + ), + ) + + with pytest.raises(FileNotFoundError, match="Sample file not found"): + submissions_module.create_submission_with_upload( + file_path=str(archive_path), + submission=_build_complete_submission(file_upload_id=None), + sample_file_path=str(tmp_path / "missing-sample.tar.gz"), + ) diff --git a/tests/test_upload_state.py b/tests/test_upload_state.py index 667ddde..53faede 100644 --- a/tests/test_upload_state.py +++ b/tests/test_upload_state.py @@ -3,6 +3,7 @@ from datacollective.upload import ( upload_dataset_file, + upload_sample_file, ) @@ -19,3 +20,18 @@ def test_upload_dataset_file_rejects_empty_file(tmp_path: Path) -> None: with pytest.raises(ValueError, match="non-empty file"): upload_dataset_file(str(empty_file), submission_id="submission") + + +def test_upload_sample_file_rejects_missing_file(tmp_path: Path) -> None: + missing_file = tmp_path / "missing-sample.tar.gz" + + with pytest.raises(FileNotFoundError, match="File not found"): + upload_sample_file(str(missing_file), submission_id="submission") + + +def test_upload_sample_file_rejects_empty_file(tmp_path: Path) -> None: + empty_file = tmp_path / "empty-sample.tar.gz" + empty_file.write_bytes(bytearray()) + + with pytest.raises(ValueError, match="non-empty file"): + upload_sample_file(str(empty_file), submission_id="submission") diff --git a/tests/test_upload_utils.py b/tests/test_upload_utils.py index 77aa3ba..b2e876d 100644 --- a/tests/test_upload_utils.py +++ b/tests/test_upload_utils.py @@ -2,18 +2,65 @@ import pytest +import datacollective.upload_utils as upload_utils_module from datacollective.models import UploadPart from datacollective.upload_utils import ( DEFAULT_PART_SIZE, MAX_UPLOAD_PARTS, UploadState, + _complete_upload, + _default_state_path, + _get_presigned_part_url, + _initiate_upload, _save_upload_state, _load_upload_state, _ensure_part_size_is_valid, _expected_parts, + _state_matches, ) +def _build_state(**overrides: object) -> UploadState: + values: dict[str, object] = { + "submissionId": "submission", + "fileUploadId": "file-upload", + "uploadId": "upload-id", + "fileSize": 1024, + "partSize": DEFAULT_PART_SIZE, + "filename": "dataset.tar.gz", + "mimeType": "application/gzip", + } + values.update(overrides) + return UploadState(**values) # type: ignore[arg-type] + + +@pytest.fixture +def captured_requests(monkeypatch) -> list[dict[str, object]]: + """Record the upload API calls instead of sending them.""" + requests: list[dict[str, object]] = [] + + class FakeResponse: + def json(self) -> dict[str, object]: + return { + "fileUploadId": "file-upload", + "uploadId": "upload-id", + "url": "https://storage.example.test/part", + "partNumber": 1, + } + + def fake_send_api_request( + method: str, url: str, json_body: dict[str, object] | None = None + ) -> FakeResponse: + requests.append({"method": method, "url": url, "json_body": json_body}) + return FakeResponse() + + monkeypatch.setattr( + upload_utils_module, "_get_api_url", lambda: "https://api.example.test" + ) + monkeypatch.setattr(upload_utils_module, "_send_api_request", fake_send_api_request) + return requests + + def test_upload_state_round_trip(tmp_path: Path) -> None: state_path = tmp_path / "upload-state.json" state = UploadState( @@ -75,3 +122,101 @@ def test_expected_parts_rounds_up_for_remainder() -> None: # A trailing partial chunk must get its own part. assert _expected_parts(file_size=250, part_size=100) == 3 assert _expected_parts(file_size=200, part_size=100) == 2 + + +def test_dataset_upload_uses_uploads_endpoints( + captured_requests: list[dict[str, object]], +) -> None: + _initiate_upload( + "submission", "dataset.tar.gz", 1024, "application/gzip", DEFAULT_PART_SIZE + ) + _get_presigned_part_url("file-upload", 1, "submission") + _complete_upload( + "file-upload", + "upload-id", + [UploadPart(partNumber=1, etag="etag-1")], + "abc123", + "submission", + ) + + assert [request["url"] for request in captured_requests] == [ + "https://api.example.test/uploads", + "https://api.example.test/uploads/file-upload/parts/1", + "https://api.example.test/uploads/file-upload", + ] + + +def test_sample_upload_uses_submission_sample_endpoints( + captured_requests: list[dict[str, object]], +) -> None: + _initiate_upload( + "submission", + "sample.tar.gz", + 1024, + "application/gzip", + DEFAULT_PART_SIZE, + is_sample=True, + ) + _get_presigned_part_url("file-upload", 2, "submission", is_sample=True) + _complete_upload( + "file-upload", + "upload-id", + [UploadPart(partNumber=1, etag="etag-1")], + "abc123", + "submission", + is_sample=True, + ) + + assert [(request["method"], request["url"]) for request in captured_requests] == [ + ("POST", "https://api.example.test/submissions/submission/sample"), + ( + "GET", + "https://api.example.test/submissions/submission/sample/file-upload/parts/2", + ), + ("POST", "https://api.example.test/submissions/submission/sample/file-upload"), + ] + + +def test_sample_upload_initiate_payload_matches_dataset_upload( + captured_requests: list[dict[str, object]], +) -> None: + _initiate_upload( + "submission", + "sample.tar.gz", + 1024, + "application/gzip", + DEFAULT_PART_SIZE, + is_sample=True, + ) + + assert captured_requests[0]["json_body"] == { + "submissionId": "submission", + "filename": "sample.tar.gz", + "fileSize": 1024, + "mimeType": "application/gzip", + } + + +def test_sample_upload_uses_a_separate_default_state_file() -> None: + archive = Path("/data/dataset.tar.gz") + + dataset_state_path = _default_state_path(archive) + sample_state_path = _default_state_path(archive, is_sample=True) + + assert dataset_state_path.name == "dataset.tar.gz.mdc-upload.json" + assert sample_state_path.name == "dataset.tar.gz.mdc-sample-upload.json" + assert dataset_state_path != sample_state_path + + +def test_state_matches_rejects_state_for_the_other_upload_kind() -> None: + sample_state = _build_state(isSample=True) + + assert _state_matches( + sample_state, "submission", "dataset.tar.gz", 1024, is_sample=True + ) + assert not _state_matches(sample_state, "submission", "dataset.tar.gz", 1024) + + +def test_upload_state_defaults_to_a_dataset_upload() -> None: + # State files written before sample uploads existed must still load. + assert _build_state().isSample is False