From 46e53b2b9a14e37e3113eebbb39ee62572920451 Mon Sep 17 00:00:00 2001 From: Bram Jans Date: Thu, 20 Aug 2026 21:56:32 +0200 Subject: [PATCH 1/5] fix: Add Azure-specific headers when adding files... MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit to datasets. Azure's Put Blob REST API requires the header x-ms-blob-type: BlockBlob on every blob-creating PUT. We didn't send this, resulting in error: 400 MissingRequiredHeader — An HTTP header that's mandatory for this request is not specified. --- openhexa/sdk/datasets/dataset.py | 18 ++++++-- tests/test_dataset.py | 76 ++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 3 deletions(-) diff --git a/openhexa/sdk/datasets/dataset.py b/openhexa/sdk/datasets/dataset.py index bf0a3ff8..fa546657 100644 --- a/openhexa/sdk/datasets/dataset.py +++ b/openhexa/sdk/datasets/dataset.py @@ -8,12 +8,22 @@ import typing from os import PathLike from pathlib import Path +from urllib.parse import urlparse import requests from openhexa.sdk.utils import Iterator, Page, Settings, graphql, read_content +def is_azure_blob_url(url: str) -> bool: + """Return whether the given URL points to an Azure Blob Storage endpoint. + + Azure Blob Storage endpoints look like .blob.core.windows.net. The domain varies across + Azure clouds (public, China, US Gov, Germany), but always contains ".blob.core.". + """ + return ".blob.core." in urlparse(url).netloc + + class DatasetFile: """Represent a single file within a dataset. Files are attached to dataset through versions.""" @@ -260,10 +270,12 @@ def add_file( self.raise_upload_exception(errors) upload_url = upload_url_result["generateDatasetUploadUrl"]["uploadUrl"] + headers = {"Content-Type": mime_type} + if is_azure_blob_url(upload_url): + # The Azure Blob Storage "Put Blob" API rejects requests that do not specify the blob type + headers["x-ms-blob-type"] = "BlockBlob" with read_content(source) as content: - response = requests.put( - upload_url, data=content, headers={"Content-Type": mime_type}, verify=Settings.verify_ssl() - ) + response = requests.put(upload_url, data=content, headers=headers, verify=Settings.verify_ssl()) response.raise_for_status() data = graphql( diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 20801ce9..cec74c4c 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,12 +1,14 @@ """Dataset test module.""" import os +from io import StringIO from unittest import TestCase from unittest.mock import patch 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 @@ -104,3 +106,77 @@ 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, + { + "HEXA_WORKSPACE": "workspace-slug", + "HEXA_TOKEN": "token", + "HEXA_SERVER_URL": "server", + }, + ) + @patch("openhexa.sdk.datasets.dataset.requests.put") + @patch("openhexa.sdk.datasets.dataset.graphql") + def test_add_file_to_azure_blob_storage(self, mock_graphql, mock_put): + """Uploads to Azure Blob Storage must specify the blob type, on top of the content type.""" + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) + upload_url = "https://account.blob.core.windows.net/hexa-datasets/version_id/file.csv?sig=signature" + mock_graphql.side_effect = [ + {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, + { + "createDatasetVersionFile": { + "success": True, + "errors": [], + "file": { + "id": "file_id", + "filename": "file.csv", + "uri": "file.csv", + "contentType": "text/csv", + "createdAt": "2021-01-01T00:00:00.000Z", + }, + } + }, + ] + + 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"], + {"Content-Type": "application/octet-stream", "x-ms-blob-type": "BlockBlob"}, + ) + + @patch.dict( + os.environ, + { + "HEXA_WORKSPACE": "workspace-slug", + "HEXA_TOKEN": "token", + "HEXA_SERVER_URL": "server", + }, + ) + @patch("openhexa.sdk.datasets.dataset.requests.put") + @patch("openhexa.sdk.datasets.dataset.graphql") + def test_add_file_to_gcs(self, mock_graphql, mock_put): + """Uploads to other storage backends must not carry Azure-specific headers.""" + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) + upload_url = "https://storage.googleapis.com/hexa-datasets/version_id/file.csv?X-Goog-Signature=signature" + mock_graphql.side_effect = [ + {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, + { + "createDatasetVersionFile": { + "success": True, + "errors": [], + "file": { + "id": "file_id", + "filename": "file.csv", + "uri": "file.csv", + "contentType": "text/csv", + "createdAt": "2021-01-01T00:00:00.000Z", + }, + } + }, + ] + + version.add_file(StringIO("foo,bar"), filename="file.csv") + + self.assertEqual(mock_put.call_args.kwargs["headers"], {"Content-Type": "application/octet-stream"}) From ba2e1d54766cc18c792d92b5be83dbf37d6a4965 Mon Sep 17 00:00:00 2001 From: Bram Jans Date: Thu, 20 Aug 2026 21:59:51 +0200 Subject: [PATCH 2/5] fix: Improve test --- tests/test_dataset.py | 90 +++++++++++++++---------------------------- 1 file changed, 32 insertions(+), 58 deletions(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index cec74c4c..cd8dc8c7 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -117,66 +117,40 @@ def test_create_dataset_version(self, mock_graphql): ) @patch("openhexa.sdk.datasets.dataset.requests.put") @patch("openhexa.sdk.datasets.dataset.graphql") - def test_add_file_to_azure_blob_storage(self, mock_graphql, mock_put): - """Uploads to Azure Blob Storage must specify the blob type, on top of the content type.""" - version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) - upload_url = "https://account.blob.core.windows.net/hexa-datasets/version_id/file.csv?sig=signature" - mock_graphql.side_effect = [ - {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, - { - "createDatasetVersionFile": { - "success": True, - "errors": [], - "file": { - "id": "file_id", - "filename": "file.csv", - "uri": "file.csv", - "contentType": "text/csv", - "createdAt": "2021-01-01T00:00:00.000Z", - }, - } - }, + def test_add_file_upload_headers(self, mock_graphql, mock_put): + """Uploads to Azure Blob Storage must specify the blob type, other backends must not receive it.""" + cases = [ + ( + "https://account.blob.core.windows.net/hexa-datasets/version_id/file.csv?sig=signature", + {"Content-Type": "application/octet-stream", "x-ms-blob-type": "BlockBlob"}, + ), + ( + "https://storage.googleapis.com/hexa-datasets/version_id/file.csv?X-Goog-Signature=signature", + {"Content-Type": "application/octet-stream"}, + ), ] - 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"], - {"Content-Type": "application/octet-stream", "x-ms-blob-type": "BlockBlob"}, - ) - - @patch.dict( - os.environ, - { - "HEXA_WORKSPACE": "workspace-slug", - "HEXA_TOKEN": "token", - "HEXA_SERVER_URL": "server", - }, - ) - @patch("openhexa.sdk.datasets.dataset.requests.put") - @patch("openhexa.sdk.datasets.dataset.graphql") - def test_add_file_to_gcs(self, mock_graphql, mock_put): - """Uploads to other storage backends must not carry Azure-specific headers.""" - version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) - upload_url = "https://storage.googleapis.com/hexa-datasets/version_id/file.csv?X-Goog-Signature=signature" - mock_graphql.side_effect = [ - {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, - { - "createDatasetVersionFile": { - "success": True, - "errors": [], - "file": { - "id": "file_id", - "filename": "file.csv", - "uri": "file.csv", - "contentType": "text/csv", - "createdAt": "2021-01-01T00:00:00.000Z", + for upload_url, expected_headers in cases: + with self.subTest(upload_url=upload_url): + mock_graphql.side_effect = [ + {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, + { + "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", + }, + } }, - } - }, - ] + ] + version = DatasetVersion(dataset=None, id="version_id", name="v1", created_at=None) - version.add_file(StringIO("foo,bar"), filename="file.csv") + version.add_file(StringIO("foo,bar"), filename="file.csv") - self.assertEqual(mock_put.call_args.kwargs["headers"], {"Content-Type": "application/octet-stream"}) + self.assertEqual(mock_put.call_args.args[0], upload_url) + self.assertEqual(mock_put.call_args.kwargs["headers"], expected_headers) From 56f6a81eae10469438a23ee54e0a4fb65f713cbf Mon Sep 17 00:00:00 2001 From: Bram Jans Date: Fri, 21 Aug 2026 09:53:06 +0200 Subject: [PATCH 3/5] Add test and test-cov to Makefile --- Makefile | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) 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]" From 49ad8833d1fd273d60eba0ee6ed9a73bf8d7f648 Mon Sep 17 00:00:00 2001 From: Bram Jans Date: Fri, 21 Aug 2026 09:56:59 +0200 Subject: [PATCH 4/5] fix: Dataset files: use block uploads for Azure Azure has a hard limit on 5GiB for a single file upload PUT. Since dataset files can be larger, we need to do them block per block to avoid hitting the limit. This implements this to upload in 64MiB blocks. --- openhexa/sdk/datasets/dataset.py | 124 ++++++++++++++++++++------- openhexa/sdk/utils.py | 18 ++++ tests/test_dataset.py | 142 ++++++++++++++++++++++--------- 3 files changed, 213 insertions(+), 71 deletions(-) diff --git a/openhexa/sdk/datasets/dataset.py b/openhexa/sdk/datasets/dataset.py index fa546657..11ce446d 100644 --- a/openhexa/sdk/datasets/dataset.py +++ b/openhexa/sdk/datasets/dataset.py @@ -4,24 +4,25 @@ 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 urlparse +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 - -def is_azure_blob_url(url: str) -> bool: - """Return whether the given URL points to an Azure Blob Storage endpoint. - - Azure Blob Storage endpoints look like .blob.core.windows.net. The domain varies across - Azure clouds (public, China, US Gov, Germany), but always contains ".blob.core.". - """ - return ".blob.core." in urlparse(url).netloc +# Azure Blob Storage refuses to store more than 5000 MiB through a single "Put Blob" request. Larger +# files have to be uploaded block by block, then committed as a whole. +# See https://learn.microsoft.com/en-us/rest/api/storageservices/put-blob +AZURE_MAX_SINGLE_PUT_SIZE = 5000 * 1024 * 1024 +# A block can hold up to 4000 MiB and a blob can hold up to 50.000 blocks. We use smaller blocks to keep +# the memory footprint of an upload low: 64 MiB blocks still allow for blobs of about 3 TiB. +AZURE_BLOCK_SIZE = 64 * 1024 * 1024 class DatasetFile: @@ -235,6 +236,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 _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 _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._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._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, @@ -253,30 +325,16 @@ 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"] - headers = {"Content-Type": mime_type} - if is_azure_blob_url(upload_url): - # The Azure Blob Storage "Put Blob" API rejects requests that do not specify the blob type - headers["x-ms-blob-type"] = "BlockBlob" + upload_url, headers = self._generate_upload_url(filename, mime_type) with read_content(source) as content: - response = requests.put(upload_url, data=content, headers=headers, verify=Settings.verify_ssl()) - response.raise_for_status() + size = content_size(content) + # An Azure Blob Storage upload has to be split in blocks when the file is too large for a single + # request, or when we cannot tell its size (a request without a Content-Length is rejected). + if headers.get("x-ms-blob-type") and (size is None or size > AZURE_MAX_SINGLE_PUT_SIZE): + self._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( """ 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 cd8dc8c7..5a6c6f92 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,9 +1,11 @@ """Dataset test module.""" +import base64 import os -from io import StringIO +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 @@ -11,6 +13,37 @@ 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.""" @@ -107,50 +140,83 @@ def test_create_dataset_version(self, mock_graphql): v = d.create_version("Second version") self.assertEqual(v.id, "") - @patch.dict( - os.environ, - { - "HEXA_WORKSPACE": "workspace-slug", - "HEXA_TOKEN": "token", - "HEXA_SERVER_URL": "server", - }, - ) + @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): - """Uploads to Azure Blob Storage must specify the blob type, other backends must not receive it.""" + """The upload request must carry the headers the backend generated the signed URL for.""" cases = [ - ( - "https://account.blob.core.windows.net/hexa-datasets/version_id/file.csv?sig=signature", - {"Content-Type": "application/octet-stream", "x-ms-blob-type": "BlockBlob"}, - ), - ( - "https://storage.googleapis.com/hexa-datasets/version_id/file.csv?X-Goog-Signature=signature", - {"Content-Type": "application/octet-stream"}, - ), + # 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 upload_url, expected_headers in cases: - with self.subTest(upload_url=upload_url): - mock_graphql.side_effect = [ - {"generateDatasetUploadUrl": {"success": True, "uploadUrl": upload_url, "errors": []}}, - { - "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", - }, - } - }, - ] + 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.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() From e35445b913087eefc94f03945f1c1f972827a565 Mon Sep 17 00:00:00 2001 From: Bram Jans Date: Fri, 21 Aug 2026 13:53:14 +0200 Subject: [PATCH 5/5] Cleanup --- openhexa/sdk/datasets/dataset.py | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/openhexa/sdk/datasets/dataset.py b/openhexa/sdk/datasets/dataset.py index 11ce446d..f2d01bf4 100644 --- a/openhexa/sdk/datasets/dataset.py +++ b/openhexa/sdk/datasets/dataset.py @@ -16,13 +16,8 @@ from openhexa.sdk.utils import Iterator, Page, Settings, content_size, graphql, read_content -# Azure Blob Storage refuses to store more than 5000 MiB through a single "Put Blob" request. Larger -# files have to be uploaded block by block, then committed as a whole. -# See https://learn.microsoft.com/en-us/rest/api/storageservices/put-blob -AZURE_MAX_SINGLE_PUT_SIZE = 5000 * 1024 * 1024 -# A block can hold up to 4000 MiB and a blob can hold up to 50.000 blocks. We use smaller blocks to keep -# the memory footprint of an upload low: 64 MiB blocks still allow for blobs of about 3 TiB. -AZURE_BLOCK_SIZE = 64 * 1024 * 1024 +AZURE_MAX_SINGLE_PUT_SIZE = 5000 * 1024 * 1024 # 5 GiB +AZURE_BLOCK_SIZE = 64 * 1024 * 1024 # 64 MiB to keep memory low class DatasetFile: @@ -259,7 +254,7 @@ def _generate_upload_url(self, filename: str, content_type: str) -> tuple[str, d # 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 _put(self, upload_url: str, query: str, filename: str, content_type: str, **kwargs) -> str: + 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 @@ -273,7 +268,7 @@ def _put(self, upload_url: str, query: str, filename: str, content_type: str, ** return upload_url - def _upload_by_blocks(self, content: typing.IO | bytes, upload_url: str, filename: str, content_type: str): + 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 @@ -287,7 +282,7 @@ def _upload_by_blocks(self, content: typing.IO | bytes, upload_url: str, filenam 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._put( + upload_url = self._azure_put( upload_url, f"comp=block&blockid={quote(block_id, safe='')}", filename, @@ -297,7 +292,7 @@ def _upload_by_blocks(self, content: typing.IO | bytes, upload_url: str, filenam block_ids.append(block_id) blocks = "".join(f"{block_id}" for block_id in block_ids) - self._put( + self._azure_put( upload_url, "comp=blocklist", filename, @@ -328,10 +323,9 @@ def add_file( upload_url, headers = self._generate_upload_url(filename, mime_type) with read_content(source) as content: size = content_size(content) - # An Azure Blob Storage upload has to be split in blocks when the file is too large for a single - # request, or when we cannot tell its size (a request without a Content-Length is rejected). + # 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._upload_by_blocks(content, upload_url, filename, mime_type) + 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() @@ -499,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