Repository: opendatacube/odc-geo
Pull request #195 (2024-12-13 02:17:13)

feat: save COGs with dask to azure
This PR introduces Azure Blob Storage support for saving COGs with Dask, complementing the existing AWS S3 support.

- Implements a MultiPartUpload class in _az.py to handle Azure-specific multipart upload operations, including block staging, finalization, and cancellation.
- Adds MultiPartUploadBase in _multipart.py to provide a shared interface for multipart uploads, applicable to both Azure and S3 implementations.
- Updates save_cog_with_dask to recognize Azure URLs and integrate the Azure multipart upload workflow.
- Adds unit tests in tests/test_az.py using mocked Azure SDK to verify functionality specific to Azure Blob Storage.

Here’s a simple test I ran to verify the Azure implementation:

```python
import dask.distributed
import planetary_computer as pc
from odc.geo.cog._tifffile import save_cog_with_dask
from odc.stac import configure_rio, stac_load
from pystac_client import Client

# Setup Dask client for efficient processing
client = dask.distributed.Client()
configure_rio(cloud_defaults=True, client=client)

# Query STAC API for Sentinel-2 data
catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")
query = catalog.search(
    collections=["sentinel-2-l2a"],
    datetime="2019-06",
    query={"s2:mgrs_tile": dict(eq="06VVN")},
)
items = list(query.items())
print(f"Found {len(items)} datasets")

# Load Sentinel-2 data with specific bands
resolution = 40 
xx = stac_load(
    items,
    bands=["SCL"],
    resolution=resolution,
    chunks={"x": 2048, "y": 2048},
    patch_url=pc.sign,
    dtype="uint16",
    nodata=0,
)

# Select a single time slice for saving
to_save = xx.SCL.isel(time=3)

# Azure Blob Storage details
blob_name = f"SCL-{to_save.time.dt.strftime('%Y%m%d').item()}.tif"
azure_url = f"az://{container}/{blob_name}"


result = save_cog_with_dask(
    to_save,
    dst=f"az://{container}/{blob_name}",
    compression="DEFLATE",
    predictor=2,
    overview_resampling="nearest",
    azure={
        "account_url": f"https://{storage_account}.blob.core.windows.net",
        "credential": sas_token,  # can be blob specific or container wide
    },
    blocksize=[512, 512],
)

# Compute the result to trigger the actual upload
result.compute()
print(f"COG successfully uploaded to: {azure_url}")
```

While implementing this, I noticed a couple of quirks in the original code:

- Spelling Inconsistencies: Words like finalise vs. finalize were used. I  used **`ise`** rather than `ize`.
- Typing Updates: The code inconsistently uses Dict/List from typing vs. built-ins (dict/list). I've changed it to Dict and List since the dev-env.yml specifies python=3.8 (but I hadn't used python 3.8 in ages).

As this is my first PR here, I look forward to feedback and suggestions for improvement!

Kirill888: @wietzesuijker thanks for the PR.

Can I ask you to undo `dict -> Dict`, `list -> List` changes, we no longer aim to support 3.8 going forward as most of our dependencies do not support it either. This also blows up review effort. When making changes like this please try to separate "clean up changes" from "actual changes" and use separate commits for each separate logical unit of work, example:

- fix spelling
- mypy harmonization
- add doc comments
- implement azure back-end


I'll need to find some time to review this properly and to remember what was the original plan for extending to other blob stores. At this point I'm not convinced we need a base class, but my memory of this code is hazy. In the meantime please

- undo mypy changes
- separate changes to doc comments from the "real change"
- make sure to run `black` formatter, or do the pre-commit thing
- push new change history into this PR with `--force-with-lease` option

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) Report
Attention: Patch coverage is `39.31034%` with `88 lines` in your changes missing coverage. Please review.
> Project coverage is 94.18%. Comparing base [(`796a99c`)](https://app.codecov.io/gh/opendatacube/odc-geo/commit/796a99cd292047ce8598bba629f59b58195481cd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) to head [(`a54072f`)](https://app.codecov.io/gh/opendatacube/odc-geo/commit/a54072f93f67563eaab7dcc20ed5d51806be4116?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).
> Report is 2 commits behind head on develop.

| [Files with missing lines](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) | Patch % | Lines |
|---|---|---|
| [odc/geo/cog/\_az.py](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&filepath=odc%2Fgeo%2Fcog%2F_az.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-b2RjL2dlby9jb2cvX2F6LnB5) | 4.91% | [58 Missing :warning: ](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) |
| [odc/geo/cog/\_tifffile.py](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&filepath=odc%2Fgeo%2Fcog%2F_tifffile.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-b2RjL2dlby9jb2cvX3RpZmZmaWxlLnB5) | 59.52% | [17 Missing :warning: ](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) |
| [odc/geo/cog/\_multipart.py](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&filepath=odc%2Fgeo%2Fcog%2F_multipart.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-b2RjL2dlby9jb2cvX211bHRpcGFydC5weQ==) | 73.33% | [8 Missing :warning: ](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) |
| [odc/geo/cog/\_s3.py](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&filepath=odc%2Fgeo%2Fcog%2F_s3.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-b2RjL2dlby9jb2cvX3MzLnB5) | 58.33% | [5 Missing :warning: ](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #195      +/-   ##
===========================================
- Coverage    95.44%   94.18%   -1.27%     
===========================================
  Files           31       33       +2     
  Lines         5534     5639     +105     
===========================================
+ Hits          5282     5311      +29     
- Misses         252      328      +76     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/opendatacube/odc-geo/pull/195?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).

Kirill888: ### Optional Dependency Handling

Example:

https://github.com/opendatacube/odc-geo/blob/62b7c5ebe41a3c7fc19f267daf06246e437ff830/odc/geo/_xr_interop.py#L55-L61 

The `have` object is here:

https://github.com/opendatacube/odc-geo/blob/62b7c5ebe41a3c7fc19f267daf06246e437ff830/odc/geo/_interop.py#L22-L28

Add `azure` to that, then use it instead of try imports.

### Tests

We use `pytest` for tests, here is example of skipping whole file when needed dependencies are not present

https://github.com/opendatacube/odc-geo/blob/62b7c5ebe41a3c7fc19f267daf06246e437ff830/tests/test_converters.py#L6-L9

### Packaging

Add optional `azure` feature here:

https://github.com/opendatacube/odc-geo/blob/62b7c5ebe41a3c7fc19f267daf06246e437ff830/setup.cfg#L54-L60


### Run checks locally

`mypy` and `pylint` checks are failing, run them locally or setup actions in your fork, actions on the PR require manual approval  from maintainers.

### Simplifying Review process

There is now quite a bit of cleanup changes here, let's have those in a separate PR that I can merge, then we can focus on actual change without getting bogged down by all the mypy harmonization changes.


wietzesuijker: Added [a PR](https://github.com/opendatacube/odc-geo/pull/196) for the cleanup. I'll give the rest another push tomorrow (UK time). Hopefully, that will give us enough time to review before the holiday season; if not, we can review again in 2025 :). 
Kirill888: @wietzesuijker other PR is merged, and some other changes, please rebase this against main.
wietzesuijker: Thanks @Kirill888! I didn't get as much done as I'd liked. I updated the test for az and started using `have.azure` and `have.s3`. Could you review the code to see if I'm on the right track? I might need to git rebase and rewrite the history, but I pushed that issue into tomorrow as I was working on this when it wasn't merged yet.
Thanks again.
Kirill888: @wietzesuijker there is currently a conflict after merging your other PR, can you please rebase this branch against `develop` branch: `git pull --rebase upstream/develop` or something along these lines.
wietzesuijker: @Kirill888 I've added some changes to fix non-related mypy errors in https://github.com/opendatacube/odc-geo/pull/195/commits/3b0223c16bbf6459f3f12e39fcbe25b491572b8a and squashed that with your commit.
Also, pylint complained about duplicate code, so I pushed https://github.com/opendatacube/odc-geo/pull/195/commits/e1b6187771bbb1a2413624ef9a43d46c82bbf17f.

The tests pass (after a few tries), and the basic example in the description of the PR still works. 
If you have any other suggestions, let me know.

Patch:
diff --git a/odc/geo/_interop.py b/odc/geo/_interop.py
index e7c61d4..f25ea5d 100644
--- a/odc/geo/_interop.py
+++ b/odc/geo/_interop.py
@@ -43,6 +43,14 @@ class _LibChecker:
     def tifffile(self) -> bool:
         return self._check("tifffile")
 
+    @property
+    def azure(self) -> bool:
+        return self._check("azure.storage.blob")
+
+    @property
+    def botocore(self) -> bool:
+        return self._check("botocore")
+
     @staticmethod
     def _check(lib_name: str) -> bool:
         return importlib.util.find_spec(lib_name) is not None
diff --git a/odc/geo/cog/_az.py b/odc/geo/cog/_az.py
new file mode 100644
index 0000000..78c5e21
--- /dev/null
+++ b/odc/geo/cog/_az.py
@@ -0,0 +1,168 @@
+import base64
+from typing import Any
+
+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 AzMultiPartUpload(AzureLimits, MultiPartUploadBase):
+    """
+    Azure Blob Storage multipart upload.
+    """
+
+    # pylint: disable=too-many-instance-attributes
+
+    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
+        # pylint: disable=import-outside-toplevel,import-error
+        from azure.storage.blob import BlobServiceClient
+
+        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.
+        """
+        # pylint: disable=import-outside-toplevel,import-error
+        from azure.storage.blob import BlobBlock
+
+        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, other: str = ""):
+        """
+        Cancel the upload by clearing the block list.
+        """
+        assert other == ""
+        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 dask_name_prefix(self) -> str:
+        """Return the Dask name prefix for Azure."""
+        return "azure-finalise"
+
+
+class DelayedAzureWriter(AzureLimits):
+    """
+    Dask-compatible writer for Azure Blob Storage multipart uploads.
+    """
+
+    def __init__(self, mpu: AzMultiPartUpload, kw: dict[str, Any]):
+        """
+        Initialise the Azure writer.
+
+        :param mpu: AzMultiPartUpload 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)
diff --git a/odc/geo/cog/_mpu.py b/odc/geo/cog/_mpu.py
index f1776d8..ddc453a 100644
--- a/odc/geo/cog/_mpu.py
+++ b/odc/geo/cog/_mpu.py
@@ -495,3 +495,49 @@ def _finalizer_dask_op(
 
     _, rr = _root.flush(write, leftPartId=1, finalise=True)
     return rr
+
+
+def get_mpu_kwargs(
+    mk_header=None,
+    mk_footer=None,
+    user_kw=None,
+    writes_per_chunk=1,
+    spill_sz=20 * (1 << 20),
+    client=None,
+) -> dict:
+    """
+    Construct shared keyword arguments for multipart uploads.
+    """
+    return {
+        "mk_header": mk_header,
+        "mk_footer": mk_footer,
+        "user_kw": user_kw,
+        "writes_per_chunk": writes_per_chunk,
+        "spill_sz": spill_sz,
+        "client": client,
+    }
+
+
+def mpu_upload(
+    chunks: Union[dask.bag.Bag, list[dask.bag.Bag]],
+    *,
+    writer: Any,
+    dask_name_prefix: str,
+    **kw,
+) -> "Delayed":
+    """Shared logic for multipart uploads to storage services."""
+    client = kw.pop("client", None)
+    writer_kw = dict(kw)
+    if client is not None:
+        writer_kw["client"] = client
+    spill_sz = kw.get("spill_sz", 20 * (1 << 20))
+    if spill_sz:
+        write = writer(writer_kw)
+    else:
+        write = None
+    return mpu_write(
+        chunks,
+        write,
+        dask_name_prefix=dask_name_prefix,
+        **kw,  # everything else remains
+    )
diff --git a/odc/geo/cog/_multipart.py b/odc/geo/cog/_multipart.py
new file mode 100644
index 0000000..c7376bf
--- /dev/null
+++ b/odc/geo/cog/_multipart.py
@@ -0,0 +1,87 @@
+"""
+Multipart upload interface.
+
+Defines the `MultiPartUploadBase` class for implementing multipart upload functionality.
+This interface standardises methods for initiating, uploading, and finalising
+multipart uploads across storage backends.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Union, TYPE_CHECKING
+
+from dask.delayed import Delayed
+from ._mpu import get_mpu_kwargs, mpu_upload
+
+if TYPE_CHECKING:
+    # pylint: disable=import-outside-toplevel,import-error
+    import dask.bag
+
+
+class MultiPartUploadBase(ABC):
+    """Abstract base class for multipart upload."""
+
+    @abstractmethod
+    def initiate(self, **kwargs) -> str:
+        """Initiate a multipart upload and return an identifier."""
+
+    @abstractmethod
+    def write_part(self, part: int, data: bytes) -> dict[str, Any]:
+        """Upload a single part."""
+
+    @abstractmethod
+    def finalise(self, parts: list[dict[str, Any]]) -> str:
+        """Finalise the upload with a list of parts."""
+
+    @abstractmethod
+    def cancel(self, other: str = ""):
+        """Cancel the multipart upload."""
+
+    @property
+    @abstractmethod
+    def url(self) -> str:
+        """Return the URL of the upload target."""
+
+    @property
+    @abstractmethod
+    def started(self) -> bool:
+        """Check if the multipart upload has been initiated."""
+
+    @abstractmethod
+    def writer(self, kw: dict[str, Any], *, client: Any = None) -> Any:
+        """
+        Return a Dask-compatible writer for multipart uploads.
+
+        :param kw: Additional parameters for the writer.
+        :param client: Dask client for distributed execution.
+        """
+
+    @abstractmethod
+    def dask_name_prefix(self) -> str:
+        """Return the dask name prefix specific to the backend."""
+
+    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 = None,
+        writes_per_chunk: int = 1,
+        spill_sz: int = 20 * (1 << 20),
+        client: Any = None,
+    ) -> Delayed:
+        """High-level upload that calls mpu_upload under the hood."""
+        kwargs = get_mpu_kwargs(
+            mk_header=mk_header,
+            mk_footer=mk_footer,
+            user_kw=user_kw,
+            writes_per_chunk=writes_per_chunk,
+            spill_sz=spill_sz,
+            client=client,
+        )
+        return mpu_upload(
+            chunks,
+            writer=self.writer,
+            dask_name_prefix=self.dask_name_prefix(),
+            **kwargs,
+        )
diff --git a/odc/geo/cog/_s3.py b/odc/geo/cog/_s3.py
index 6da5cf8..0ae6981 100644
--- a/odc/geo/cog/_s3.py
+++ b/odc/geo/cog/_s3.py
@@ -9,14 +9,16 @@ from typing import TYPE_CHECKING, Any, Optional
 
 from cachetools import cached
 
-from ._mpu import PartsWriter, SomeData, mpu_write
+from ._mpu import PartsWriter, SomeData
+from ._multipart import MultiPartUploadBase
 
 if TYPE_CHECKING:
     import dask.bag
-    import distributed
     from botocore.credentials import ReadOnlyCredentials
     from dask.delayed import Delayed
 
+    import distributed
+
 _state: dict[str, Any] = {}
 
 
@@ -68,7 +70,7 @@ class S3Limits:
         return 10_000
 
 
-class MultiPartUpload(S3Limits):
+class S3MultiPartUpload(S3Limits, MultiPartUploadBase):
     """
     Dask to S3 dumper.
     """
@@ -195,31 +197,9 @@ class MultiPartUpload(S3Limits):
             writer.prep_client(client)
         return writer
 
-    # pylint: disable=too-many-arguments
-    def upload(
-        self,
-        chunks: "dask.bag.Bag" | list["dask.bag.Bag"],
-        *,
-        mk_header: Any = None,
-        mk_footer: Any = None,
-        user_kw: dict[str, Any] | None = None,
-        writes_per_chunk: int = 1,
-        spill_sz: int = 20 * (1 << 20),
-        client: Any = None,
-        **kw,
-    ) -> "Delayed":
-        """Upload chunks to S3 with multipart uploads."""
-        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="s3finalise",
-        )
+    def dask_name_prefix(self) -> str:
+        """Return the Dask name prefix for S3."""
+        return "s3finalise"
 
 
 def _safe_get(v, timeout=0.1):
@@ -236,7 +216,7 @@ class DelayedS3Writer(S3Limits):
 
     # pylint: disable=import-outside-toplevel,import-error
 
-    def __init__(self, mpu: MultiPartUpload, kw: dict[str, Any]):
+    def __init__(self, mpu: S3MultiPartUpload, kw: dict[str, Any]):
         self.mpu = mpu
         self.kw = kw  # mostly ContentType= kinda thing
         self._shared_var: Optional["distributed.Variable"] = None
@@ -262,7 +242,7 @@ class DelayedS3Writer(S3Limits):
             self._shared_var = Variable(self._build_name("MPUpload"), client)
         return self._shared_var
 
-    def _ensure_init(self, final_write: bool = False) -> MultiPartUpload:
+    def _ensure_init(self, final_write: bool = False) -> S3MultiPartUpload:
         # pylint: disable=too-many-return-statements
         mpu = self.mpu
         if mpu.started:
diff --git a/odc/geo/cog/_tifffile.py b/odc/geo/cog/_tifffile.py
index 3406629..3ab86e2 100644
--- a/odc/geo/cog/_tifffile.py
+++ b/odc/geo/cog/_tifffile.py
@@ -11,6 +11,7 @@ import itertools
 from functools import partial
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from urllib.parse import urlparse
 from xml.sax.saxutils import escape as xml_escape
 
 import numpy as np
@@ -23,7 +24,8 @@ from ..math import resolve_nodata
 from ..types import Shape2d, SomeNodata, Unset, shape_
 from ._mpu import mpu_write
 from ._mpu_fs import MPUFileSink
-from ._s3 import MultiPartUpload, s3_parse_url
+from ._multipart import MultiPartUploadBase
+
 from ._shared import (
     GDAL_COMP,
     GEOTIFF_TAGS,
@@ -626,15 +628,16 @@ def save_cog_with_dask(
     bigtiff: bool = True,
     overview_resampling: Union[int, str] = "nearest",
     aws: Optional[dict[str, Any]] = None,
+    azure: Optional[dict[str, Any]] = None,
     client: Any = None,
     stats: bool | int = True,
     **kw,
 ) -> Any:
     """
-    Save a Cloud Optimized GeoTIFF to S3 or file with Dask.
+    Save a Cloud Optimized GeoTIFF to S3, Azure Blob Storage, or file with Dask.
 
     :param xx: Pixels as :py:class:`xarray.DataArray` backed by Dask
-    :param dst: S3 url or a file path on shared storage
+    :param dst: S3, Azure URL, or file path
     :param compression: Compression to use, default is ``DEFLATE``
     :param level: Compression "level", depends on chosen compression
     :param predictor: TIFF predictor setting
@@ -643,6 +646,7 @@ def save_cog_with_dask(
     :param blocksize: Configure blocksizes for main and overview images
     :param bigtiff: Generate BigTIFF by default, set to ``False`` to disable
     :param aws: Configure AWS write access
+    :param azure: Azure credentials/config
     :param client: Dask client
     :param stats: Set to ``False`` to disable stats computation
 
@@ -653,13 +657,12 @@ def save_cog_with_dask(
 
     from ..xr import ODCExtensionDa
 
-    if aws is None:
-        aws = {}
+    aws = aws or {}
+    azure = azure or {}
 
-    upload_params = {k: kw.pop(k) for k in ["writes_per_chunk", "spill_sz"] if k in kw}
-    upload_params.update(
-        {k: aws.pop(k) for k in ["writes_per_chunk", "spill_sz"] if k in aws}
-    )
+    upload_params = {
+        k: kw.pop(k, None) for k in ["writes_per_chunk", "spill_sz"] if k in kw
+    }
     parts_base = kw.pop("parts_base", None)
 
     # Normalise compression settings and remove GDAL compat options from kw
@@ -699,7 +702,7 @@ def save_cog_with_dask(
 
     if band_names and len(band_names) != meta.nsamples:
         raise ValueError(
-            f"Found {len(band_names)} band names ({band_names}) but there are {meta.nsamples} bands."
+            f"Found {len(band_names)} band names ({band_names}), expected {meta.nsamples} bands."
         )
 
     layers = _pyramids_from_cog_metadata(xx, meta, resampling=overview_resampling)
@@ -731,19 +734,37 @@ def save_cog_with_dask(
             "_stats": _stats,
         }
 
-    tiles_write_order = _tiles[::-1]
-    if len(tiles_write_order) > 4:
-        tiles_write_order = [
-            dask.bag.concat(tiles_write_order[:4]),
-            *tiles_write_order[4:],
-        ]
+    # Determine output type and initiate uploader
+    parsed_url = urlparse(dst)
+    if parsed_url.scheme == "s3":
+        if have.botocore:
+            from ._s3 import S3MultiPartUpload, s3_parse_url
 
-    bucket, key = s3_parse_url(dst)
-    if not bucket:
-        # assume disk output
+            bucket, key = s3_parse_url(dst)
+            uploader: MultiPartUploadBase = S3MultiPartUpload(bucket, key, **aws)
+        else:
+            raise RuntimeError("Please install `boto3` to use S3")
+    elif parsed_url.scheme == "az":
+        if have.azure:
+            from ._az import AzMultiPartUpload
+
+            assert azure is not None
+            assert "account_url" in azure
+            assert "credential" in azure
+
+            uploader = AzMultiPartUpload(
+                account_url=azure["account_url"],
+                container=parsed_url.netloc,
+                blob=parsed_url.path.lstrip("/"),
+                credential=azure["credential"],
+            )
+        else:
+            raise RuntimeError("Please install `azure-storage-blob` to use Azure")
+    else:
+        # Assume local disk
         write = MPUFileSink(dst, parts_base=parts_base)
         return mpu_write(
-            tiles_write_order,
+            _tiles[::-1],
             write,
             mk_header=_patch_hdr,
             user_kw={
@@ -755,15 +776,15 @@ def save_cog_with_dask(
             **upload_params,
         )
 
-    upload_params["ContentType"] = (
-        "image/tiff;application=geotiff;profile=cloud-optimized"
-    )
+    # Upload tiles
+    tiles_write_order = _tiles[::-1]  # Reverse tiles for writing
+    if len(tiles_write_order) > 4:  # Optimize for larger datasets
+        tiles_write_order = [
+            dask.bag.concat(tiles_write_order[:4]),
+            *tiles_write_order[4:],
+        ]
 
-    cleanup = aws.pop("cleanup", False)
-    s3_sink = MultiPartUpload(bucket, key, **aws)
-    if cleanup:
-        s3_sink.cancel("all")
-    return s3_sink.upload(
+    return uploader.upload(
         tiles_write_order,
         mk_header=_patch_hdr,
         user_kw={
diff --git a/odc/geo/gcp.py b/odc/geo/gcp.py
index 56fade3..f72f7dd 100644
--- a/odc/geo/gcp.py
+++ b/odc/geo/gcp.py
@@ -102,9 +102,16 @@ class GCPMapping:
 
     def points(self) -> Tuple[Geometry, Geometry]:
         """Return multipoint geometries for (Pixel, World)."""
+        pix_points: list[tuple[float, float]] = [
+            (float(p[0]), float(p[1])) for p in self._pix.tolist()
+        ]
+        wld_points: list[tuple[float, float]] = [
+            (float(p[0]), float(p[1])) for p in self._wld.tolist()
+        ]
+
         return (
-            multipoint(self._pix.tolist(), None),
-            multipoint(self._wld.tolist(), self.crs),
+            multipoint(pix_points, None),
+            multipoint(wld_points, self.crs),
         )
 
     def __dask_tokenize__(self):
diff --git a/odc/geo/geom.py b/odc/geo/geom.py
index 043ac6c..f1e8985 100644
--- a/odc/geo/geom.py
+++ b/odc/geo/geom.py
@@ -350,13 +350,16 @@ class BoundingBox(Sequence[float]):
         ny = y1 - y0
         pts = quasi_random_r2(n, offset=offset)
         s = numpy.asarray([nx, ny], dtype="float32")
-        edge_pts = []
+        edge_pts: list[tuple[float, float]] = []
 
         if with_edges:
             sample_density = numpy.sqrt(n / (nx * ny))
             n_side = int(numpy.round(sample_density * min(nx, ny))) + 1
             n_side = max(2, n_side)
-            edge_pts = self.boundary(n_side).coords[:-1]
+            edge_pts = [
+                (float(ep[0]), float(ep[1]))
+                for ep in list(self.boundary(n_side).coords[:-1])
+            ]
             if padding is None:
                 padding = 0.3 * min(nx, ny) / (n_side - 1)
 
@@ -368,7 +371,11 @@ class BoundingBox(Sequence[float]):
         pts[:, 0] += x0
         pts[:, 1] += y0
 
-        return multipoint(pts.tolist() + edge_pts, self.crs)
+        coords: list[tuple[float, float]] = [
+            (float(p[0]), float(p[1])) for p in pts.tolist()
+        ] + edge_pts
+
+        return multipoint(coords, self.crs)
 
 
 def wrap_shapely(method):
diff --git a/odc/geo/roi.py b/odc/geo/roi.py
index 7e4ee42..50759ed 100644
--- a/odc/geo/roi.py
+++ b/odc/geo/roi.py
@@ -284,7 +284,7 @@ class VariableSizedTiles:
     @property
     def chunks(self) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
         """Dask compatible chunk rerpesentation."""
-        y, x = (tuple(np.diff(idx).tolist()) for idx in self._offsets)
+        y, x = (tuple(map(int, np.diff(idx))) for idx in self._offsets)
         return (y, x)
 
     def locate(self, pix: SomeIndex2d) -> Tuple[int, int]:
diff --git a/setup.cfg b/setup.cfg
index 4f09ae5..a19ab78 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -54,10 +54,14 @@ tiff =
 s3 =
     boto3
 
+az =
+    azure-storage-blob
+
 all =
     %(warp)s
     %(tiff)s
     %(s3)s
+    %(az)s
 
 test =
     pytest


Test patch:
diff --git a/tests/test_az.py b/tests/test_az.py
new file mode 100644
index 0000000..9462f09
--- /dev/null
+++ b/tests/test_az.py
@@ -0,0 +1,80 @@
+"""Tests for the Azure AzMultiPartUpload class."""
+
+import base64
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+pytest.importorskip("azure.storage.blob")
+from odc.geo.cog._az import AzMultiPartUpload  # noqa: E402
+
+
+@pytest.fixture
+def azure_mpu():
+    """Fixture for initializing AzMultiPartUpload."""
+    account_url = "https://account_name.blob.core.windows.net"
+    return AzMultiPartUpload(account_url, "container", "some.blob", None)
+
+
+def test_mpu_init(azure_mpu):
+    """Basic test for AzMultiPartUpload initialization."""
+    assert azure_mpu.account_url == "https://account_name.blob.core.windows.net"
+    assert azure_mpu.container == "container"
+    assert azure_mpu.blob == "some.blob"
+    assert azure_mpu.credential is None
+
+
+@patch("odc.geo.cog._az.BlobServiceClient")
+def test_azure_multipart_upload(mock_blob_service_client):
+    """Test the full Azure AzMultiPartUpload functionality."""
+    # Mock Azure Blob SDK client structure
+    mock_blob_client = MagicMock()
+    mock_container_client = MagicMock()
+    mock_blob_service_client.return_value.get_container_client.return_value = (
+        mock_container_client
+    )
+    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"
+
+    # Create an instance of AzMultiPartUpload and call its methods
+    azure_upload = AzMultiPartUpload(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])
+
+    # Define block IDs
+    block_id1 = base64.b64encode(b"block-1").decode("utf-8")
+    block_id2 = base64.b64encode(b"block-2").decode("utf-8")
+
+    # Verify the results
+    assert upload_id == "azure-block-upload"
+    assert etag == "mock-etag"
+
+    # Verify BlobServiceClient instantiation
+    mock_blob_service_client.assert_called_once_with(
+        account_url=account_url, credential=credential
+    )
+
+    # Verify stage_block calls
+    mock_blob_client.stage_block.assert_any_call(
+        block_id=block_id1, data=b"first chunk of data"
+    )
+    mock_blob_client.stage_block.assert_any_call(
+        block_id=block_id2, data=b"second chunk of data"
+    )
+
+    # Verify commit_block_list was called correctly
+    block_list = mock_blob_client.commit_block_list.call_args[0][0]
+    assert len(block_list) == 2
+    assert block_list[0].id == block_id1
+    assert block_list[1].id == block_id2
+    mock_blob_client.commit_block_list.assert_called_once()
diff --git a/tests/test_s3.py b/tests/test_s3.py
index 8349bd8..d04a462 100644
--- a/tests/test_s3.py
+++ b/tests/test_s3.py
@@ -1,9 +1,29 @@
-from odc.geo.cog._s3 import MultiPartUpload
+"""Tests for odc.geo.cog._s3."""
 
-# TODO: moto
+import unittest
 
+from odc.geo.cog._s3 import S3MultiPartUpload
 
+# Conditional import for S3 support
+try:
+    from odc.geo.cog._s3 import S3MultiPartUpload
+
+    HAVE_S3 = True
+except ImportError:
+    S3MultiPartUpload = None
+    HAVE_S3 = False
+
+
+def require_s3(test_func):
+    """Decorator to skip tests if s3 dependencies are not installed."""
+    return unittest.skipUnless(HAVE_S3, "s3 dependencies are not installed")(test_func)
+
+
+@require_s3
 def test_s3_mpu():
-    mpu = MultiPartUpload("bucket", "file.dat")
-    assert mpu.bucket == "bucket"
-    assert mpu.key == "file.dat"
+    """Test S3MultiPartUpload class initialization."""
+    mpu = S3MultiPartUpload("bucket", "file.dat")
+    if mpu.bucket != "bucket":
+        raise ValueError("Invalid bucket")
+    if mpu.key != "file.dat":
+        raise ValueError("Invalid key")