Add Azure Blob Storage backend

.github/workflows/ci.yml

@@ -32,7 +32,7 @@ jobs:
           python -m pip install pre-commit codecov
 
       - name: Install module
-        run: python -m pip install -e .["test"]
+        run: python -m pip install -e .["test,az"]
 
       - name: Run pre-commit
         run: pre-commit run --all-files

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 21.10b0
+    rev: 22.3.0
     hooks:
       - id: black
         language_version: python
@@ -12,7 +12,7 @@ repos:
         language_version: python
 
   - repo: https://github.com/PyCQA/flake8
-    rev: 3.8.3
+    rev: 4.0.1
     hooks:
       - id: flake8
         language_version: python

cogeo_mosaic/backends/__init__.py

@@ -3,6 +3,7 @@
 from typing import Any
 from urllib.parse import urlparse
 
+from cogeo_mosaic.backends.az import ABSBackend
 from cogeo_mosaic.backends.base import BaseBackend
 from cogeo_mosaic.backends.dynamodb import DynamoDBBackend
 from cogeo_mosaic.backends.file import FileBackend
@@ -34,6 +35,10 @@ def MosaicBackend(input: str, *args: Any, **kwargs: Any) -> BaseBackend:
     elif parsed.scheme == "gs":
         return GCSBackend(input, *args, **kwargs)
 
+    # `https://{storageaccount}.blob.core.windows.net/{container}/{key}`
+    elif parsed.scheme == "https" and parsed.netloc.endswith(".blob.core.windows.net"):
+        return ABSBackend(input, *args, **kwargs)
+
     # `dynamodb://{region}/{table}:{mosaic}`
     elif parsed.scheme == "dynamodb":
         return DynamoDBBackend(input, *args, **kwargs)

cogeo_mosaic/backends/az.py

@@ -0,0 +1,110 @@
+"""cogeo-mosaic Azure Blob Storage backend."""
+
+import json
+from typing import Any
+from urllib.parse import urlparse
+
+import attr
+from cachetools import TTLCache, cached
+from cachetools.keys import hashkey
+
+from cogeo_mosaic.backends.base import BaseBackend
+from cogeo_mosaic.backends.utils import _compress_gz_json, _decompress_gz
+from cogeo_mosaic.cache import cache_config
+from cogeo_mosaic.errors import _HTTP_EXCEPTIONS, MosaicError, MosaicExistsError
+from cogeo_mosaic.mosaic import MosaicJSON
+
+try:
+    from azure.core.exceptions import HttpResponseError
+    from azure.identity import DefaultAzureCredential
+    from azure.storage.blob import BlobServiceClient
+except ImportError:
+    HttpResponseError = None
+    DefaultAzureCredential = None
+    BlobServiceClient = None
+
+
+@attr.s
+class ABSBackend(BaseBackend):
+    """Azure Blob Storage Backend Adapter"""
+
+    client: Any = attr.ib(default=None)
+    account_url: str = attr.ib(init=False)
+    container: str = attr.ib(init=False)
+    key: str = attr.ib(init=False)
+
+    _backend_name = "Azure Blob Storage"
+
+    def __attrs_post_init__(self):
+        """Post Init: parse path and create client."""
+        assert (
+            HttpResponseError is not None
+        ), "'azure-identity' and 'azure-storage-blob' must be installed to use ABSBackend"
+
+        az_credentials = DefaultAzureCredential()
+
+        parsed = urlparse(self.input)
+        self.account_url = "https://%s" % parsed.netloc
+        self.container = parsed.path.split("/")[1]
+        self.key = parsed.path.strip("/%s" % self.container)
+        self.client = self.client or BlobServiceClient(
+            account_url=self.account_url, credential=az_credentials
+        )
+        super().__attrs_post_init__()
+
+    def write(self, overwrite: bool = False, **kwargs: Any):
+        """Write mosaicjson document to Azure Blob Storage."""
+        if not overwrite and self._head_object(self.key, self.container):
+            raise MosaicExistsError("Mosaic file already exist, use `overwrite=True`.")
+
+        mosaic_doc = self.mosaic_def.dict(exclude_none=True)
+        if self.key.endswith(".gz"):
+            body = _compress_gz_json(mosaic_doc)
+        else:
+            body = json.dumps(mosaic_doc).encode("utf-8")
+
+        self._put_object(self.key, self.container, body, **kwargs)
+
+    @cached(
+        TTLCache(maxsize=cache_config.maxsize, ttl=cache_config.ttl),
+        key=lambda self: hashkey(self.input),
+    )
+    def _read(self) -> MosaicJSON:  # type: ignore
+        """Get mosaicjson document."""
+        body = self._get_object(self.key, self.container)
+        self._file_byte_size = len(body)
+
+        if self.key.endswith(".gz"):
+            body = _decompress_gz(body)
+
+        return MosaicJSON(**json.loads(body))
+
+    def _get_object(self, key: str, container: str) -> bytes:
+        try:
+            container_client = self.client.get_container_client(container)
+            blob_client = container_client.get_blob_client(key)
+            response = blob_client.download_blob().readall()
+        except HttpResponseError as e:
+            exc = _HTTP_EXCEPTIONS.get(e.status_code, MosaicError)
+            raise exc(e.reason) from e
+
+        return response
+
+    def _put_object(self, key: str, container: str, body: bytes, **kwargs) -> str:
+        try:
+            container_client = self.client.get_container_client(container)
+            blob_client = container_client.get_blob_client(key)
+            blob_client.upload_blob(body)
+        except HttpResponseError as e:
+            exc = _HTTP_EXCEPTIONS.get(e.status_code, MosaicError)
+            raise exc(e.reason) from e
+
+        return key
+
+    def _head_object(self, key: str, container: str) -> bool:
+        try:
+            container_client = self.client.get_container_client(container)
+            blob_client = container_client.get_blob_client(key)
+            return blob_client.exists()
+        except HttpResponseError:
+            return False

docs/advanced/backends.md

@@ -10,6 +10,8 @@ Starting in version `3.0.0`, we introduced specific `backend` to abstract mosaic
 
 - **GCSBackend** (`gs://`)
 
+- **ABSBackend** (`https://{storageaccount}.blob.core.windows.net/{container}/{key}`)
+
 - **DynamoDBBackend** (`dynamodb://{region}/{table_name}`). If `region` is not passed, it reads the value of the `AWS_REGION` environment variable. If that environment variable does not exist, it falls back to `us-east-1`. If you choose not to pass a `region`, you still need three `/` before the table name, like so `dynamodb:///{table_name}`.
 
 - **SQLiteBackend** (`sqlite:///{file.db}:{mosaic_name}`)
@@ -87,6 +89,18 @@ permissions to let cogeo-mosaic access the files. For example:
   following condition: `resource.type == "storage.googleapis.com/Object"
   && resource.name.startsWith("projects/_/buckets/mybucket")`
 
+## ABS Backend
+
+The ABS Backend allows read and write operations from/to Azure Blob Storage.
+
+The backend uses DefaultAzureCredential for authorization, see
+[here](https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential)
+for details on how to authenticate.
+
+When using this backend is necessary to set the appropriate roles and IAM
+permissions to let cogeo-mosaic access the files. See the Azure standard
+roles; Storage Blob Data Contributor/Owner/Reader.
+
 ## STAC Backend
 
 The STACBackend is a **read-only** backend, meaning it can't be used to write a file. This backend will POST to the input url looking for STAC items which will then be used to create the mosaicJSON in memory.

setup.py

@@ -20,6 +20,7 @@
 
 extra_reqs = {
     "aws": ["boto3"],
+    "az": ["azure-identity", "azure-storage-blob"],
     "gcp": ["google-cloud-storage"],
     "test": ["pytest", "pytest-cov"],
     "dev": ["pytest", "pytest-cov", "pre-commit"],
@@ -28,7 +29,7 @@
 
 setup(
     name="cogeo-mosaic",
-    description=u"Create mosaicJSON.",
+    description="Create mosaicJSON.",
     long_description=readme,
     long_description_content_type="text/markdown",
     python_requires=">=3.7",
@@ -42,7 +43,7 @@
         "Topic :: Scientific/Engineering :: GIS",
     ],
     keywords="COG Mosaic GIS",
-    author=u"Vincent Sarago",
+    author="Vincent Sarago",
     author_email="vincent@developmentseed.org",
     url="https://github.com/developmentseed/cogeo-mosaic",
     license="MIT",

tests/test_backends.py

@@ -18,6 +18,7 @@
 from rio_tiler.errors import PointOutsideBounds
 
 from cogeo_mosaic.backends import GCSBackend, MosaicBackend
+from cogeo_mosaic.backends.az import ABSBackend
 from cogeo_mosaic.backends.dynamodb import DynamoDBBackend
 from cogeo_mosaic.backends.file import FileBackend
 from cogeo_mosaic.backends.memory import MemoryBackend
@@ -428,6 +429,117 @@ def test_gs_backend(session):
     session.reset_mock()
 
 
+@patch("cogeo_mosaic.backends.az.BlobServiceClient")
+def test_abs_backend(session):
+    """Test ABS backend."""
+    with open(mosaic_gz, "rb") as f:
+        session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.return_value.readall.return_value = (
+            f.read()
+        )
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.upload_blob.return_value = (
+        True
+    )
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.return_value = (
+        False
+    )
+
+    with MosaicBackend(
+        "https://storage_account.blob.core.windows.net/container/mymosaic.json.gz"
+    ) as mosaic:
+        assert mosaic._backend_name == "Azure Blob Storage"
+        assert isinstance(mosaic, ABSBackend)
+        assert (
+            mosaic.mosaicid
+            == "24d43802c19ef67cc498c327b62514ecf70c2bbb1bbc243dda1ee075"
+        )
+        assert mosaic.quadkey_zoom == 7
+        assert list(
+            mosaic.mosaic_def.dict(exclude_none=True, exclude={"tiles"}).keys()
+        ) == [
+            "mosaicjson",
+            "version",
+            "minzoom",
+            "maxzoom",
+            "quadkey_zoom",
+            "bounds",
+            "center",
+        ]
+        assert mosaic.assets_for_tile(150, 182, 9) == ["cog1.tif", "cog2.tif"]
+        assert mosaic.assets_for_point(-73, 45) == ["cog1.tif", "cog2.tif"]
+        session.return_value.get_container_client.assert_called_once_with("container")
+        session.return_value.get_container_client.return_value.get_blob_client.assert_called_once_with(
+            "mymosaic.json.gz"
+        )
+
+        session.return_value.get_container_client.return_value.get_blob_client.return_value.upload_blob.assert_not_called()
+        session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.assert_not_called()
+        session.reset_mock()
+
+    with open(mosaic_gz, "rb") as f:
+        session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.return_value.readall.return_value = (
+            f.read()
+        )
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.upload_blob.return_value = (
+        True
+    )
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.return_value = (
+        False
+    )
+
+    with MosaicBackend(
+        "https://storage_account.blob.core.windows.net/container/mymosaic.json.gz",
+        mosaic_def=mosaic_content,
+    ) as mosaic:
+        assert isinstance(mosaic, ABSBackend)
+        mosaic.write()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.assert_not_called()
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.assert_called_once()
+    session.reset_mock()
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.return_value = (
+        False
+    )
+    with MosaicBackend(
+        "https://storage_account.blob.core.windows.net/container/00000.json",
+        mosaic_def=mosaic_content,
+    ) as mosaic:
+        assert isinstance(mosaic, ABSBackend)
+        mosaic.write()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.assert_not_called()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.assert_called_once()
+    session.reset_mock()
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.return_value = (
+        True
+    )
+    with MosaicBackend(
+        "https://storage_account.blob.core.windows.net/container/00000.json",
+        mosaic_def=mosaic_content,
+    ) as mosaic:
+        assert isinstance(mosaic, ABSBackend)
+        with pytest.raises(MosaicExistsError):
+            mosaic.write()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.assert_not_called()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.assert_called_once()
+    session.reset_mock()
+
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.return_value = (
+        True
+    )
+    with MosaicBackend(
+        "https://storage_account.blob.core.windows.net/container/00000.json",
+        mosaic_def=mosaic_content,
+    ) as mosaic:
+        assert isinstance(mosaic, ABSBackend)
+        mosaic.write(overwrite=True)
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.download_blob.assert_not_called()
+    session.return_value.get_container_client.return_value.get_blob_client.return_value.exists.assert_not_called()
+    session.reset_mock()
+
+
 class MockMeta(object):
     """Mock Meta."""