Add support of container URL

camptocamp · Jun 17, 2024 · d76d27a · d76d27a
1 parent 67ed280
commit d76d27a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/tilecloud_chain/controller.py b/tilecloud_chain/controller.py
@@ -18,7 +18,7 @@
 import ruamel.yaml
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import DefaultAzureCredential
-from azure.storage.blob import BlobServiceClient, ContentSettings
+from azure.storage.blob import BlobServiceClient, ContainerClient, ContentSettings
 from bottle import jinja2_template
 from PIL import Image
 from prometheus_client import Summary
@@ -104,15 +104,19 @@ def main(args: Optional[list[str]] = None, out: Optional[IO[str]] = None) -> Non
         sys.exit(1)
 
 
-def get_azure_client() -> BlobServiceClient:
+def get_azure_container_client(container: str) -> ContainerClient:
     """Get the Azure blog storage client."""
     if "AZURE_STORAGE_CONNECTION_STRING" in os.environ and os.environ["AZURE_STORAGE_CONNECTION_STRING"]:
-        return BlobServiceClient.from_connection_string(os.environ["AZURE_STORAGE_CONNECTION_STRING"])
+        return BlobServiceClient.from_connection_string(
+            os.environ["AZURE_STORAGE_CONNECTION_STRING"]
+        ).get_container_client(container=container)
+    elif "AZURE_STORAGE_BLOB_CONTAINER_URL" not in os.environ:
+        return ContainerClient.from_container_url(os.environ["AZURE_STORAGE_BLOB_CONTAINER_URL"])
     else:
         return BlobServiceClient(
             account_url=os.environ["AZURE_STORAGE_ACCOUNT_URL"],
             credential=DefaultAzureCredential(),
-        )
+        ).get_container_client(container=container)
 
 
 def _send(
@@ -134,7 +138,8 @@ def _send(
     if cache["type"] == "azure":
         cache_azure = cast(tilecloud_chain.configuration.CacheAzure, cache)
         key_name = os.path.join(f"{cache['folder']}", path)
-        blob = get_azure_client().get_blob_client(container=cache_azure["container"], blob=key_name)
+        container = get_azure_container_client(cache_azure["container"])
+        blob = container.get_blob_client(key_name)
         blob.upload_blob(data, overwrite=True)
 
         blob.upload_blob(
@@ -177,7 +182,9 @@ def _get(path: str, cache: tilecloud_chain.configuration.Cache) -> Optional[byte
         cache_azure = cast(tilecloud_chain.configuration.CacheAzure, cache)
         key_name = os.path.join(f"{cache['folder']}", path)
         try:
-            blob = get_azure_client().get_blob_client(container=cache_azure["container"], blob=key_name)
+            blob = get_azure_container_client(container=cache_azure["container"]).get_blob_client(
+                blob=key_name
+            )
             return blob.download_blob().readall()
         except ResourceNotFoundError:
             return None

diff --git a/tilecloud_chain/server.py b/tilecloud_chain/server.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2013-2023 by Stéphane Brunner
+# Copyright (c) 2013-2024 by Stéphane Brunner
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -53,7 +53,7 @@
 import tilecloud_chain.security
 from tilecloud import Tile, TileCoord
 from tilecloud_chain import TileGeneration, configuration, controller, internal_mapcache
-from tilecloud_chain.controller import get_azure_client
+from tilecloud_chain.controller import get_azure_container_client
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -281,8 +281,8 @@ def _get(
             key_name = os.path.join(cache_azure["folder"], path)
             try:
                 with _GET_TILE.labels(storage="azure").time():
-                    blob = get_azure_client().get_blob_client(
-                        container=cache_azure["container"], blob=key_name
+                    blob = get_azure_container_client(container=cache_azure["container"]).get_blob_client(
+                        blob=key_name
                     )
                 properties = blob.get_blob_properties()
                 data = blob.download_blob().readall()