Feat: azure backend to save cogs with dask

Author: wietzesuijker

odc/geo/cog/_az.py

@@ -0,0 +1,191 @@
+import base64
+from typing import Any, Union
+
+from azure.storage.blob import BlobBlock, BlobServiceClient
+from dask.delayed import Delayed
+
+from ._mpu import mpu_write
+from ._multipart import MultiPartUploadBase
+
+
+class AzureLimits:
+    """
+    Common Azure writer settings.
+    """
+
+    @property
+    def min_write_sz(self) -> int:
+        # Azure minimum write size for blocks (default is 4 MiB)
+        return 4 * (1 << 20)
+
+    @property
+    def max_write_sz(self) -> int:
+        # Azure maximum write size for blocks (default is 100 MiB)
+        return 100 * (1 << 20)
+
+    @property
+    def min_part(self) -> int:
+        return 1
+
+    @property
+    def max_part(self) -> int:
+        # Azure supports up to 50,000 blocks per blob
+        return 50_000
+
+
+class MultiPartUpload(AzureLimits, MultiPartUploadBase):
+    def __init__(
+        self, account_url: str, container: str, blob: str, credential: Any = None
+    ):
+        """
+        Initialise Azure multipart upload.
+
+        :param account_url: URL of the Azure storage account.
+        :param container: Name of the container.
+        :param blob: Name of the blob.
+        :param credential: Authentication credentials (e.g., SAS token or key).
+        """
+        self.account_url = account_url
+        self.container = container
+        self.blob = blob
+        self.credential = credential
+
+        # Initialise Azure Blob service client
+        self.blob_service_client = BlobServiceClient(
+            account_url=account_url, credential=credential
+        )
+        self.container_client = self.blob_service_client.get_container_client(container)
+        self.blob_client = self.container_client.get_blob_client(blob)
+
+        self.block_ids: list[str] = []
+
+    def initiate(self, **kwargs) -> str:
+        """
+        Initialise the upload. No-op for Azure.
+        """
+        return "azure-block-upload"
+
+    def write_part(self, part: int, data: bytes) -> dict[str, Any]:
+        """
+        Stage a block in Azure.
+
+        :param part: Part number (unique).
+        :param data: Data for this part.
+        :return: A dictionary containing part information.
+        """
+        block_id = base64.b64encode(f"block-{part}".encode()).decode()
+        self.blob_client.stage_block(block_id=block_id, data=data)
+        self.block_ids.append(block_id)
+        return {"PartNumber": part, "BlockId": block_id}
+
+    def finalise(self, parts: list[dict[str, Any]]) -> str:
+        """
+        Commit the block list to finalise the upload.
+
+        :param parts: List of uploaded parts metadata.
+        :return: The ETag of the finalised blob.
+        """
+        block_list = [BlobBlock(block_id=part["BlockId"]) for part in parts]
+        self.blob_client.commit_block_list(block_list)
+        return self.blob_client.get_blob_properties().etag
+
+    def cancel(self):
+        """
+        Cancel the upload by clearing the block list.
+        """
+        self.block_ids.clear()
+
+    @property
+    def url(self) -> str:
+        """
+        Get the Azure blob URL.
+
+        :return: The full URL of the blob.
+        """
+        return self.blob_client.url
+
+    @property
+    def started(self) -> bool:
+        """
+        Check if any blocks have been staged.
+
+        :return: True if blocks have been staged, False otherwise.
+        """
+        return bool(self.block_ids)
+
+    def writer(self, kw: dict[str, Any], client: Any = None):
+        """
+        Return a stateless writer compatible with Dask.
+        """
+        return DelayedAzureWriter(self, kw)
+
+    def upload(
+        self,
+        chunks: Union["dask.bag.Bag", list["dask.bag.Bag"]],
+        *,
+        mk_header: Any = None,
+        mk_footer: Any = None,
+        user_kw: dict[str, Any] = None,
+        writes_per_chunk: int = 1,
+        spill_sz: int = 20 * (1 << 20),
+        client: Any = None,
+        **kw,
+    ) -> "Delayed":
+        """
+        Upload chunks to Azure Blob Storage with multipart uploads.
+
+        :param chunks: Dask bag of chunks to upload.
+        :param mk_header: Function to create header data.
+        :param mk_footer: Function to create footer data.
+        :param user_kw: User-provided metadata for the upload.
+        :param writes_per_chunk: Number of writes per chunk.
+        :param spill_sz: Spill size for buffering data.
+        :param client: Dask client for distributed execution.
+        :return: A Dask delayed object representing the finalised upload.
+        """
+        write = self.writer(kw, client=client) if spill_sz else None
+        return mpu_write(
+            chunks,
+            write,
+            mk_header=mk_header,
+            mk_footer=mk_footer,
+            user_kw=user_kw,
+            writes_per_chunk=writes_per_chunk,
+            spill_sz=spill_sz,
+            dask_name_prefix="azure-finalise",
+        )
+
+
+class DelayedAzureWriter(AzureLimits):
+    """
+    Dask-compatible writer for Azure Blob Storage multipart uploads.
+    """
+
+    def __init__(self, mpu: MultiPartUpload, kw: dict[str, Any]):
+        """
+        Initialise the Azure writer.
+
+        :param mpu: MultiPartUpload instance.
+        :param kw: Additional parameters for the writer.
+        """
+        self.mpu = mpu
+        self.kw = kw  # Additional metadata like ContentType
+
+    def __call__(self, part: int, data: bytes) -> dict[str, Any]:
+        """
+        Write a single part to Azure Blob Storage.
+
+        :param part: Part number.
+        :param data: Chunk data.
+        :return: Metadata for the written part.
+        """
+        return self.mpu.write_part(part, data)
+
+    def finalise(self, parts: list[dict[str, Any]]) -> str:
+        """
+        Finalise the upload by committing the block list.
+
+        :param parts: List of uploaded parts metadata.
+        :return: ETag of the finalised blob.
+        """
+        return self.mpu.finalise(parts)

tests/test_az.py

@@ -0,0 +1,74 @@
+"""Tests for the Azure MultiPartUpload class."""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+from odc.geo.cog._az import MultiPartUpload
+
+
+def test_mpu_init():
+    """Basic test for the MultiPartUpload class."""
+    account_url = "https://account_name.blob.core.windows.net"
+    mpu = MultiPartUpload(account_url, "container", "some.blob", None)
+    if mpu.account_url != account_url:
+        raise AssertionError(f"mpu.account_url should be '{account_url}'.")
+    if mpu.container != "container":
+        raise AssertionError("mpu.container should be 'container'.")
+    if mpu.blob != "some.blob":
+        raise AssertionError("mpu.blob should be 'some.blob'.")
+    if mpu.credential is not None:
+        raise AssertionError("mpu.credential should be 'None'.")
+
+
+class TestMultiPartUpload(unittest.TestCase):
+    """Test the MultiPartUpload class."""
+
+    @patch("odc.geo.cog._az.BlobServiceClient")
+    def test_azure_multipart_upload(self, mock_blob_service_client):
+        """Test the MultiPartUpload class."""
+        # Arrange - mock the Azure Blob SDK
+        # Mock the blob client and its methods
+        mock_blob_client = MagicMock()
+        mock_container_client = MagicMock()
+        mcc = mock_container_client
+        mock_blob_service_client.return_value.get_container_client.return_value = mcc
+        mock_container_client.get_blob_client.return_value = mock_blob_client
+
+        # Simulate return values for Azure Blob SDK methods
+        mock_blob_client.get_blob_properties.return_value.etag = "mock-etag"
+
+        # Test parameters
+        account_url = "https://mockaccount.blob.core.windows.net"
+        container = "mock-container"
+        blob = "mock-blob"
+        credential = "mock-sas-token"
+
+        # Act - create an instance of MultiPartUpload and call its methods
+        azure_upload = MultiPartUpload(account_url, container, blob, credential)
+        upload_id = azure_upload.initiate()
+        part1 = azure_upload.write_part(1, b"first chunk of data")
+        part2 = azure_upload.write_part(2, b"second chunk of data")
+        etag = azure_upload.finalise([part1, part2])
+
+        # Assert - check the results
+        # Check that the initiate method behaves as expected
+        self.assertEqual(upload_id, "azure-block-upload")
+
+        # Verify the calls to Azure Blob SDK methods
+        mock_blob_service_client.assert_called_once_with(
+            account_url=account_url, credential=credential
+        )
+        mock_blob_client.stage_block.assert_any_call(
+            part1["BlockId"], b"first chunk of data"
+        )
+        mock_blob_client.stage_block.assert_any_call(
+            part2["BlockId"], b"second chunk of data"
+        )
+        mock_blob_client.commit_block_list.assert_called_once()
+        self.assertEqual(etag, "mock-etag")
+
+        # Verify block list passed during finalise
+        block_list = mock_blob_client.commit_block_list.call_args[0][0]
+        self.assertEqual(len(block_list), 2)
+        self.assertEqual(block_list[0].id, part1["BlockId"])
+        self.assertEqual(block_list[1].id, part2["BlockId"])