diff --git a/Makefile b/Makefile index d04c55c6..e35d1973 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ .DEFAULT_GOAL := help -.PHONY: help l lint install-editable +.PHONY: help l lint t test test-cov install-editable help: ## Show this help message @echo "Available commands:" @@ -12,6 +12,16 @@ lint: ## Run linting (pre-commit) on all files l: lint +test: ## Run the test suite + @echo "Running the test suite (pytest)" + pytest + +t: test + +test-cov: ## Run the test suite with coverage report (as in CI) + @echo "Running the test suite with coverage (pytest)" + pytest --cov=. --cov-report html --cov-report term --cov-fail-under=25 + install-editable: ## Install the SDK in editable mode with dev dependencies @echo "Installing the SDK in editable mode" pip install -e ".[dev]" diff --git a/openhexa/sdk/datasets/dataset.py b/openhexa/sdk/datasets/dataset.py index bf0a3ff8..f2d01bf4 100644 --- a/openhexa/sdk/datasets/dataset.py +++ b/openhexa/sdk/datasets/dataset.py @@ -4,14 +4,20 @@ https://github.com/BLSQ/openhexa/wiki/Using-the-OpenHEXA-SDK#working-with-datasets for more information about datasets. """ +import base64 +import io import mimetypes import typing from os import PathLike from pathlib import Path +from urllib.parse import quote import requests -from openhexa.sdk.utils import Iterator, Page, Settings, graphql, read_content +from openhexa.sdk.utils import Iterator, Page, Settings, content_size, graphql, read_content + +AZURE_MAX_SINGLE_PUT_SIZE = 5000 * 1024 * 1024 # 5 GiB +AZURE_BLOCK_SIZE = 64 * 1024 * 1024 # 64 MiB to keep memory low class DatasetFile: @@ -225,6 +231,77 @@ def get_file(self, filename: str) -> DatasetFile: created_at=file["createdAt"], ) + def _generate_upload_url(self, filename: str, content_type: str) -> tuple[str, dict]: + """Ask the backend for a signed upload URL, along with the headers that URL was signed for.""" + data = graphql( + """ + mutation generateDatasetUploadUrl ($input: GenerateDatasetUploadUrlInput!) { + generateDatasetUploadUrl(input: $input) { + uploadUrl + headers + success + errors + } + } + """, + {"input": {"versionId": self.id, "contentType": content_type, "uri": filename}}, + ) + result = data["generateDatasetUploadUrl"] + if result["success"] is False: + self.raise_upload_exception(result["errors"]) + + # The backend tells us which headers the signed URL expects: Azure Blob Storage rejects a "Put Blob" + # request that does not specify the blob type. Other backends only need the content type. + return result["uploadUrl"], result["headers"] or {"Content-Type": content_type} + + def _azure_put(self, upload_url: str, query: str, filename: str, content_type: str, **kwargs) -> str: + """Send a single upload request, and return the URL it was sent to. + + Signed URLs expire after an hour, which a large upload can outlive. When that happens we simply ask + for a fresh one and try again, hence the URL being returned to the caller. + """ + response = requests.put(f"{upload_url}&{query}", verify=Settings.verify_ssl(), **kwargs) + if response.status_code in (401, 403): + upload_url, _ = self._generate_upload_url(filename, content_type) + response = requests.put(f"{upload_url}&{query}", verify=Settings.verify_ssl(), **kwargs) + response.raise_for_status() + + return upload_url + + def _azure_upload_by_blocks(self, content: typing.IO | bytes, upload_url: str, filename: str, content_type: str): + """Upload content to Azure Blob Storage block by block, then commit the blocks as a single blob. + + See https://learn.microsoft.com/en-us/rest/api/storageservices/put-block + """ + if isinstance(content, bytes): + content = io.BytesIO(content) + + block_ids = [] + while block := content.read(AZURE_BLOCK_SIZE): + if isinstance(block, str): + block = block.encode() + # Block ids must be unique within the blob and share the same length once decoded + block_id = base64.b64encode(f"{len(block_ids):032d}".encode()).decode() + upload_url = self._azure_put( + upload_url, + f"comp=block&blockid={quote(block_id, safe='')}", + filename, + content_type, + data=block, + ) + block_ids.append(block_id) + + blocks = "".join(f"{block_id}" for block_id in block_ids) + self._azure_put( + upload_url, + "comp=blocklist", + filename, + content_type, + data=f'{blocks}'.encode(), + # The content type of the blob itself, as opposed to the content type of this request + headers={"x-ms-blob-content-type": content_type}, + ) + def add_file( self, source: str | PathLike[str] | typing.IO | bytes, @@ -243,28 +320,15 @@ def add_file( if mime_type is None: mime_type = "application/octet-stream" - upload_url_result = graphql( - """ - mutation generateDatasetUploadUrl ($input: GenerateDatasetUploadUrlInput!) { - generateDatasetUploadUrl(input: $input) { - uploadUrl - success - errors - } - } - """, - {"input": {"versionId": self.id, "contentType": mime_type, "uri": filename}}, - ) - if upload_url_result["generateDatasetUploadUrl"]["success"] is False: - errors = upload_url_result["generateDatasetUploadUrl"]["errors"] - self.raise_upload_exception(errors) - - upload_url = upload_url_result["generateDatasetUploadUrl"]["uploadUrl"] + upload_url, headers = self._generate_upload_url(filename, mime_type) with read_content(source) as content: - response = requests.put( - upload_url, data=content, headers={"Content-Type": mime_type}, verify=Settings.verify_ssl() - ) - response.raise_for_status() + size = content_size(content) + # upload block by block for very large files (>5GiB) when on Azure + if headers.get("x-ms-blob-type") and (size is None or size > AZURE_MAX_SINGLE_PUT_SIZE): + self._azure_upload_by_blocks(content, upload_url, filename, mime_type) + else: + response = requests.put(upload_url, data=content, headers=headers, verify=Settings.verify_ssl()) + response.raise_for_status() data = graphql( """ @@ -429,7 +493,7 @@ def latest_version(self) -> DatasetVersion | None: if self._latest_version is None: data = graphql( """ - query getLatestVersion($datasetId: ID!) { + query getLatestVersion($datasetId: ID!) { dataset(id: $datasetId) { latestVersion { id diff --git a/openhexa/sdk/utils.py b/openhexa/sdk/utils.py index 9b994ad3..a4d03b36 100644 --- a/openhexa/sdk/utils.py +++ b/openhexa/sdk/utils.py @@ -275,3 +275,21 @@ def read_content(source: str | os.PathLike[str] | typing.IO | bytes): finally: if hasattr(source, "close"): source.close() + + +def content_size(content: typing.IO | bytes) -> int | None: + """Return the size in bytes of a content opened by read_content, or None if it cannot be determined.""" + if isinstance(content, bytes): + return len(content) + try: + return os.fstat(content.fileno()).st_size + except (AttributeError, OSError, ValueError): + pass + # Non-file objects (BytesIO, StringIO...) still expose their size if they are seekable + try: + position = content.tell() + size = content.seek(0, os.SEEK_END) + content.seek(position) + return size + except (AttributeError, OSError, ValueError): + return None diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 20801ce9..5a6c6f92 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,14 +1,49 @@ """Dataset test module.""" +import base64 import os +from io import BytesIO, StringIO from unittest import TestCase -from unittest.mock import patch +from unittest.mock import MagicMock, patch +from urllib.parse import quote from httmock import HTTMock, all_requests, response from openhexa.sdk.datasets import Dataset +from openhexa.sdk.datasets.dataset import DatasetVersion from openhexa.sdk.workspaces import workspace +ENV = {"HEXA_WORKSPACE": "workspace-slug", "HEXA_TOKEN": "token", "HEXA_SERVER_URL": "server"} +UPLOAD_URL = "https://account.blob.core.windows.net/hexa-datasets/dataset_id/version_id/file.csv?sig=signature" +AZURE_HEADERS = {"Content-Type": "application/octet-stream", "x-ms-blob-type": "BlockBlob"} + + +def upload_url_response(upload_url=UPLOAD_URL, headers=None): + """Build a generateDatasetUploadUrl mutation response.""" + return {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "headers": headers, "errors": []}} + + +def version_file_response(): + """Build a createDatasetVersionFile mutation response.""" + return { + "createDatasetVersionFile": { + "success": True, + "errors": [], + "file": { + "id": "file_id", + "filename": "file.csv", + "uri": "file.csv", + "contentType": "application/octet-stream", + "createdAt": "2021-01-01T00:00:00.000Z", + }, + } + } + + +def block_id(index): + """Build the block id the SDK generates for the block at the given index.""" + return base64.b64encode(f"{index:032d}".encode()).decode() + class DatasetTest(TestCase): """Dataset test class.""" @@ -104,3 +139,84 @@ def test_create_dataset_version(self, mock_graphql): self.assertEqual(v.id, "") v = d.create_version("Second version") self.assertEqual(v.id, "") + + @patch.dict(os.environ, ENV) + @patch("openhexa.sdk.datasets.dataset.requests.put") + @patch("openhexa.sdk.datasets.dataset.graphql") + def test_add_file_upload_headers(self, mock_graphql, mock_put): + """The upload request must carry the headers the backend generated the signed URL for.""" + cases = [ + # Azure Blob Storage requires the blob type on top of the content type + (AZURE_HEADERS, AZURE_HEADERS), + # Backends that do not need specific headers fall back to the content type + (None, {"Content-Type": "application/octet-stream"}), + ] + + for returned_headers, expected_headers in cases: + with self.subTest(headers=returned_headers): + mock_graphql.side_effect = [upload_url_response(headers=returned_headers), version_file_response()] + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) + + version.add_file(StringIO("foo,bar"), filename="file.csv") + + self.assertEqual(mock_put.call_args.args[0], UPLOAD_URL) + self.assertEqual(mock_put.call_args.kwargs["headers"], expected_headers) + + @patch.dict(os.environ, ENV) + @patch("openhexa.sdk.datasets.dataset.AZURE_MAX_SINGLE_PUT_SIZE", 8) + @patch("openhexa.sdk.datasets.dataset.AZURE_BLOCK_SIZE", 4) + @patch("openhexa.sdk.datasets.dataset.requests.put") + @patch("openhexa.sdk.datasets.dataset.graphql") + def test_add_file_too_large_for_a_single_request(self, mock_graphql, mock_put): + """Content that does not fit in a single request is staged block by block, then committed.""" + mock_graphql.side_effect = [upload_url_response(headers=AZURE_HEADERS), version_file_response()] + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) + + version.add_file(BytesIO(b"0123456789"), filename="file.csv") + + block_urls = [f"{UPLOAD_URL}&comp=block&blockid={quote(block_id(i), safe='')}" for i in range(3)] + self.assertEqual( + [call.args[0] for call in mock_put.call_args_list], + [*block_urls, f"{UPLOAD_URL}&comp=blocklist"], + ) + self.assertEqual([call.kwargs["data"] for call in mock_put.call_args_list[:3]], [b"0123", b"4567", b"89"]) + + commit = mock_put.call_args_list[-1] + self.assertEqual(commit.kwargs["headers"], {"x-ms-blob-content-type": "application/octet-stream"}) + self.assertEqual( + commit.kwargs["data"], + '' + f"{block_id(0)}{block_id(1)}{block_id(2)}" + "".encode(), + ) + + @patch.dict(os.environ, ENV) + @patch("openhexa.sdk.datasets.dataset.AZURE_MAX_SINGLE_PUT_SIZE", 2) + @patch("openhexa.sdk.datasets.dataset.AZURE_BLOCK_SIZE", 4) + @patch("openhexa.sdk.datasets.dataset.requests.put") + @patch("openhexa.sdk.datasets.dataset.graphql") + def test_add_file_signed_url_expired_during_upload(self, mock_graphql, mock_put): + """A signed URL that expires while blocks are being uploaded is replaced by a fresh one.""" + refreshed_url = UPLOAD_URL.replace("signature", "refreshed_signature") + mock_graphql.side_effect = [ + upload_url_response(headers=AZURE_HEADERS), + upload_url_response(refreshed_url, headers=AZURE_HEADERS), + version_file_response(), + ] + expired_response = MagicMock(status_code=403) + mock_put.side_effect = [expired_response, MagicMock(status_code=201), MagicMock(status_code=201)] + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) + + version.add_file(BytesIO(b"0123"), filename="file.csv") + + block_query = f"comp=block&blockid={quote(block_id(0), safe='')}" + self.assertEqual( + [call.args[0] for call in mock_put.call_args_list], + [ + f"{UPLOAD_URL}&{block_query}", + f"{refreshed_url}&{block_query}", + # The blocks are then committed through the refreshed URL as well + f"{refreshed_url}&comp=blocklist", + ], + ) + expired_response.raise_for_status.assert_not_called()