fix: Support Azure backend for dataset files (HEXA-1781)

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]"

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"<Latest>{block_id}</Latest>" for block_id in block_ids)
+        self._azure_put(
+            upload_url,
+            "comp=blocklist",
+            filename,
+            content_type,
+            data=f'<?xml version="1.0" encoding="utf-8"?><BlockList>{blocks}</BlockList>'.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

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

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, "<newVersionId>")
         v = d.create_version("Second version")
         self.assertEqual(v.id, "<newVersionId>")
+
+    @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"],
+            '<?xml version="1.0" encoding="utf-8"?><BlockList>'
+            f"<Latest>{block_id(0)}</Latest><Latest>{block_id(1)}</Latest><Latest>{block_id(2)}</Latest>"
+            "</BlockList>".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()