Add blob_exists method

data_safe_haven/external/api/azure_api.py

@@ -60,7 +60,7 @@
     StorageAccountKey,
     StorageAccountListKeysResult,
 )
-from azure.storage.blob import BlobServiceClient
+from azure.storage.blob import BlobClient, BlobServiceClient
 from azure.storage.filedatalake import DataLakeServiceClient
 
 from data_safe_haven.exceptions import (
@@ -80,6 +80,32 @@ def __init__(
         super().__init__(subscription_name)
         self.logger = NonLoggingSingleton() if disable_logging else LoggingSingleton()
 
+    def blob_client(
+        self,
+        resource_group_name: str,
+        storage_account_name: str,
+        storage_container_name: str,
+        blob_name: str,
+    ) -> BlobClient:
+        """Construct a client for a blob which may exist or not"""
+        # Connect to Azure client
+        storage_account_keys = self.get_storage_account_keys(
+            resource_group_name, storage_account_name
+        )
+        blob_service_client = BlobServiceClient.from_connection_string(
+            f"DefaultEndpointsProtocol=https;AccountName={storage_account_name};AccountKey={storage_account_keys[0].value};EndpointSuffix=core.windows.net"
+        )
+
+        if not isinstance(blob_service_client, BlobServiceClient):
+            msg = f"Could not connect to storage account '{storage_account_name}'."
+            raise DataSafeHavenAzureError(msg)
+
+        blob_client = blob_service_client.get_blob_client(
+            container=storage_container_name, blob=blob_name
+        )
+
+        return blob_client
+
     def compile_desired_state(
         self,
         automation_account_name: str,
@@ -181,20 +207,13 @@ def download_blob(
             DataSafeHavenAzureError if the blob could not be downloaded
         """
         try:
-            # Connect to Azure client
-            storage_account_keys = self.get_storage_account_keys(
-                resource_group_name, storage_account_name
-            )
-            blob_service_client = BlobServiceClient.from_connection_string(
-                f"DefaultEndpointsProtocol=https;AccountName={storage_account_name};AccountKey={storage_account_keys[0].value};EndpointSuffix=core.windows.net"
+            blob_client = self.blob_client(
+                resource_group_name,
+                storage_account_name,
+                storage_container_name,
+                blob_name,
             )
-            if not isinstance(blob_service_client, BlobServiceClient):
-                msg = f"Could not connect to storage account '{storage_account_name}'."
-                raise DataSafeHavenAzureError(msg)
             # Download the requested file
-            blob_client = blob_service_client.get_blob_client(
-                container=storage_container_name, blob=blob_name
-            )
             return str(blob_client.download_blob(encoding="utf-8").readall())
         except Exception as exc:
             msg = f"Blob file '{blob_name}' could not be downloaded from '{storage_account_name}'\n{exc}."
@@ -1232,24 +1251,34 @@ def upload_blob(
             DataSafeHavenAzureError if the blob could not be uploaded
         """
         try:
-            # Connect to Azure client
-            storage_account_keys = self.get_storage_account_keys(
-                resource_group_name, storage_account_name
-            )
-            blob_service_client = BlobServiceClient.from_connection_string(
-                f"DefaultEndpointsProtocol=https;AccountName={storage_account_name};AccountKey={storage_account_keys[0].value};EndpointSuffix=core.windows.net"
+            blob_client = self.blob_client(
+                resource_group_name,
+                storage_account_name,
+                storage_container_name,
+                blob_name,
             )
-            if not isinstance(blob_service_client, BlobServiceClient):
-                msg = f"Could not connect to storage account '{storage_account_name}'."
-                raise DataSafeHavenAzureError(msg)
             # Upload the created file
-            blob_client = blob_service_client.get_blob_client(
-                container=storage_container_name, blob=blob_name
-            )
             blob_client.upload_blob(blob_data, overwrite=True)
             self.logger.info(
                 f"Uploaded file [green]{blob_name}[/] to blob storage.",
             )
         except Exception as exc:
             msg = f"Blob file '{blob_name}' could not be uploaded to '{storage_account_name}'\n{exc}."
             raise DataSafeHavenAzureError(msg) from exc
+
+    def blob_exists(
+        self,
+        blob_name: str,
+        resource_group_name: str,
+        storage_account_name: str,
+        storage_container_name: str,
+    ) -> None:
+        blob_client = self.blob_client(
+            resource_group_name, storage_account_name, storage_container_name, blob_name
+        )
+        # Upload the created file
+        exists = blob_client.exists()
+        response = "exists" if exists else "does not exist"
+        self.logger.info(
+            f"File [green]{blob_name}[/] {response} in blob storage.",
+        )

tests/external/api/azure_api.py

@@ -24,6 +24,25 @@ def get_key(self, key_name):
     )
 
 
+@fixture
+def mock_blob_client(monkeypatch):  # noqa: ARG001
+    class MockBlobClient:
+        def __init__(
+            self,
+            resource_group_name,  # noqa: ARG002
+            storage_account_name,  # noqa: ARG002
+            storage_container_name,  # noqa: ARG002
+            blob_name,
+        ):
+            self.blob_name = blob_name
+
+        def exists(self):
+            if self.blob_name == "exists":
+                return True
+            else:
+                return False
+
+
 class TestAzureApi:
     def test_get_keyvault_key(self, mock_key_client):  # noqa: ARG002
         api = AzureApi("subscription name")
@@ -36,3 +55,15 @@ def test_get_keyvault_key_missing(self, mock_key_client):  # noqa: ARG002
             DataSafeHavenAzureError, match="Failed to retrieve key does not exist"
         ):
             api.get_keyvault_key("does not exist", "key vault name")
+
+    def test_blob_exists(self, mock_blob_client):  # noqa: ARG002
+        api = AzureApi("subscription name")
+        assert api.blob_exists(
+            "resource_group", "storage_account", "storage_container", "exists"
+        )
+
+    def test_blob_does_not_exist(self, mock_blob_client):  # noqa: ARG002
+        api = AzureApi("subscription name")
+        assert not api.blob_exists(
+            "resource_group", "storage_account", "storage_container", "abc.txt"
+        )