From 8c52fa813bc452a5e5ec92db27d08d547e6b56c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20Brunner?= <stephane.brunner@camptocamp.com>
Date: Tue, 2 Jul 2024 14:23:07 +0200
Subject: [PATCH] Fix support of container URL

---
 tilecloud_chain/__init__.py   | 24 +++++++++++++++++++++++-
 tilecloud_chain/controller.py | 22 ++--------------------
 tilecloud_chain/server.py     |  9 +++++++--
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/tilecloud_chain/__init__.py b/tilecloud_chain/__init__.py
index 78e9179d3..d31a50a40 100644
--- a/tilecloud_chain/__init__.py
+++ b/tilecloud_chain/__init__.py
@@ -31,6 +31,8 @@
 import c2cwsgiutils.setup_process
 import jsonschema_validator
 import psycopg2
+from azure.identity import DefaultAzureCredential
+from azure.storage.blob import BlobServiceClient, ContainerClient, ContentSettings
 from c2cwsgiutils import sentry
 from PIL import Image
 from prometheus_client import Counter, Summary
@@ -379,6 +381,26 @@ def filter(self, record: Any) -> bool:
         return True
 
 
+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"]
+        ).get_container_client(container=container)
+    elif "AZURE_STORAGE_BLOB_CONTAINER_URL" in os.environ:
+        container_client = ContainerClient.from_container_url(os.environ["AZURE_STORAGE_BLOB_CONTAINER_URL"])
+        if os.environ.get("AZURE_STORAGE_BLOB_VALIDATE_CONTAINER_NAME", "true").lower() == "true":
+            assert (
+                container == container_client.container_name
+            ), f"Container name mismatch: {container} != {container_client.container_name}"
+        return container_client
+    else:
+        return BlobServiceClient(
+            account_url=os.environ["AZURE_STORAGE_ACCOUNT_URL"],
+            credential=DefaultAzureCredential(),
+        ).get_container_client(container=container)
+
+
 class TileGeneration:
     """Base class of all the tile generation."""
 
@@ -802,9 +824,9 @@ def get_store(
             cache_azure = cast(tilecloud_chain.configuration.CacheAzureTyped, cache)
             # on Azure
             cache_tilestore = AzureStorageBlobTileStore(
-                container=cache_azure["container"],
                 tilelayout=layout,
                 cache_control=cache_azure.get("cache_control"),
+                client=get_azure_container_client(cache_azure["container"]),
             )
         elif cache["type"] == "mbtiles":
             metadata = {}
diff --git a/tilecloud_chain/controller.py b/tilecloud_chain/controller.py
index 96827c448..02f5ac76a 100644
--- a/tilecloud_chain/controller.py
+++ b/tilecloud_chain/controller.py
@@ -17,8 +17,7 @@
 import requests
 import ruamel.yaml
 from azure.core.exceptions import ResourceNotFoundError
-from azure.identity import DefaultAzureCredential
-from azure.storage.blob import BlobServiceClient, ContainerClient, ContentSettings
+from azure.storage.blob import ContentSettings
 from bottle import jinja2_template
 from PIL import Image
 from prometheus_client import Summary
@@ -32,6 +31,7 @@
     TileGeneration,
     add_common_options,
     configuration,
+    get_azure_container_client,
     get_queue_store,
     get_tile_matrix_identifier,
 )
@@ -104,24 +104,6 @@ def main(args: Optional[list[str]] = None, out: Optional[IO[str]] = None) -> Non
         sys.exit(1)
 
 
-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"]
-        ).get_container_client(container=container)
-    elif "AZURE_STORAGE_BLOB_CONTAINER_URL" in os.environ:
-        container_client = ContainerClient.from_container_url(os.environ["AZURE_STORAGE_BLOB_CONTAINER_URL"])
-        if os.environ.get("AZURE_STORAGE_BLOB_VALIDATE_CONTAINER_NAME", "true").lower() == "true":
-            assert container == container_client.container_name
-        return container_client
-    else:
-        return BlobServiceClient(
-            account_url=os.environ["AZURE_STORAGE_ACCOUNT_URL"],
-            credential=DefaultAzureCredential(),
-        ).get_container_client(container=container)
-
-
 def _send(
     data: Union[bytes, str], path: str, mime_type: str, cache: tilecloud_chain.configuration.Cache
 ) -> None:
diff --git a/tilecloud_chain/server.py b/tilecloud_chain/server.py
index 00fd85258..076e8bdd9 100644
--- a/tilecloud_chain/server.py
+++ b/tilecloud_chain/server.py
@@ -52,8 +52,13 @@
 import tilecloud_chain.configuration
 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_container_client
+from tilecloud_chain import (
+    TileGeneration,
+    configuration,
+    controller,
+    get_azure_container_client,
+    internal_mapcache,
+)
 
 _LOGGER = logging.getLogger(__name__)