|
| 1 | +"""Azure Blob Storage adapter implementing the StorageProvider contract. |
| 2 | +
|
| 3 | +The StorageProvider interface is bucket/key oriented; in Azure terms the |
| 4 | +``bucket`` argument maps to a *container* and ``key`` maps to a *blob name*. |
| 5 | +
|
| 6 | +Authentication resolves in this order (first match wins): |
| 7 | +
|
| 8 | +1. ``AZURE_STORAGE_CONNECTION_STRING`` |
| 9 | +2. ``AZURE_STORAGE_ACCOUNT_NAME`` + ``AZURE_STORAGE_ACCOUNT_KEY`` (shared key) |
| 10 | +3. ``AZURE_STORAGE_ACCOUNT_NAME`` + ``DefaultAzureCredential`` |
| 11 | + (managed identity / workload identity — recommended in production) |
| 12 | +
|
| 13 | +Signed URLs are issued as SAS tokens. Shared-key auth signs with the account |
| 14 | +key directly; identity-based auth requests a short-lived *user delegation key* |
| 15 | +from Azure AD and signs with that instead. |
| 16 | +
|
| 17 | +This adapter requires the optional ``storage-azure`` extra:: |
| 18 | +
|
| 19 | + uv sync --extra storage-azure |
| 20 | +""" |
| 21 | + |
| 22 | +from datetime import UTC, datetime, timedelta |
| 23 | + |
| 24 | +from azure.core.exceptions import ResourceNotFoundError |
| 25 | +from azure.storage.blob import ( |
| 26 | + BlobSasPermissions, |
| 27 | + ContentSettings, |
| 28 | + generate_blob_sas, |
| 29 | +) |
| 30 | +from azure.storage.blob.aio import BlobServiceClient |
| 31 | + |
| 32 | +from src.lib.config import Settings, settings |
| 33 | +from src.lib.storage.base import StorageProvider |
| 34 | + |
| 35 | + |
| 36 | +class StorageConfigurationError(RuntimeError): |
| 37 | + """Raised when Azure Blob Storage credentials are not configured.""" |
| 38 | + |
| 39 | + |
| 40 | +def _parse_connection_string(connection_string: str) -> dict[str, str]: |
| 41 | + """Parse a ``key=value;...`` Azure connection string into a dict.""" |
| 42 | + parts: dict[str, str] = {} |
| 43 | + for segment in connection_string.split(";"): |
| 44 | + segment = segment.strip() |
| 45 | + if not segment: |
| 46 | + continue |
| 47 | + key, sep, value = segment.partition("=") |
| 48 | + if sep: |
| 49 | + parts[key.strip()] = value.strip() |
| 50 | + return parts |
| 51 | + |
| 52 | + |
| 53 | +class AzureBlobStorageProvider(StorageProvider): |
| 54 | + """StorageProvider backed by Azure Blob Storage (async SDK).""" |
| 55 | + |
| 56 | + def __init__( |
| 57 | + self, |
| 58 | + client: BlobServiceClient, |
| 59 | + account_name: str, |
| 60 | + account_key: str | None = None, |
| 61 | + ) -> None: |
| 62 | + """Wrap an already-configured async ``BlobServiceClient``. |
| 63 | +
|
| 64 | + Args: |
| 65 | + client: Async Azure ``BlobServiceClient``. |
| 66 | + account_name: Storage account name (used for SAS generation). |
| 67 | + account_key: Shared key. When ``None``, signed URLs are issued via |
| 68 | + a user delegation key (managed/workload identity). |
| 69 | + """ |
| 70 | + self._client = client |
| 71 | + self._account_name = account_name |
| 72 | + self._account_key = account_key |
| 73 | + |
| 74 | + @classmethod |
| 75 | + def from_settings( |
| 76 | + cls, config: Settings | None = None |
| 77 | + ) -> "AzureBlobStorageProvider": |
| 78 | + """Build an adapter from application settings. |
| 79 | +
|
| 80 | + Raises: |
| 81 | + StorageConfigurationError: When no usable credentials are found. |
| 82 | + """ |
| 83 | + config = config or settings |
| 84 | + |
| 85 | + if config.AZURE_STORAGE_CONNECTION_STRING: |
| 86 | + parsed = _parse_connection_string(config.AZURE_STORAGE_CONNECTION_STRING) |
| 87 | + account_name = parsed.get("AccountName") |
| 88 | + if not account_name: |
| 89 | + msg = "AZURE_STORAGE_CONNECTION_STRING is missing AccountName" |
| 90 | + raise StorageConfigurationError(msg) |
| 91 | + client = BlobServiceClient.from_connection_string( |
| 92 | + config.AZURE_STORAGE_CONNECTION_STRING |
| 93 | + ) |
| 94 | + return cls(client, account_name, account_key=parsed.get("AccountKey")) |
| 95 | + |
| 96 | + account_name = config.AZURE_STORAGE_ACCOUNT_NAME |
| 97 | + if not account_name: |
| 98 | + msg = ( |
| 99 | + "Azure storage requires AZURE_STORAGE_CONNECTION_STRING or " |
| 100 | + "AZURE_STORAGE_ACCOUNT_NAME" |
| 101 | + ) |
| 102 | + raise StorageConfigurationError(msg) |
| 103 | + |
| 104 | + account_url = ( |
| 105 | + f"https://{account_name}.blob.{config.AZURE_STORAGE_ENDPOINT_SUFFIX}" |
| 106 | + ) |
| 107 | + |
| 108 | + if config.AZURE_STORAGE_ACCOUNT_KEY: |
| 109 | + client = BlobServiceClient( |
| 110 | + account_url, credential=config.AZURE_STORAGE_ACCOUNT_KEY |
| 111 | + ) |
| 112 | + return cls( |
| 113 | + client, account_name, account_key=config.AZURE_STORAGE_ACCOUNT_KEY |
| 114 | + ) |
| 115 | + |
| 116 | + # Fall back to managed/workload identity. azure-identity is imported |
| 117 | + # lazily so connection-string / shared-key deployments don't need it. |
| 118 | + try: |
| 119 | + from azure.identity.aio import DefaultAzureCredential |
| 120 | + except ImportError as exc: # pragma: no cover - optional dependency |
| 121 | + msg = ( |
| 122 | + "Managed-identity Azure auth requires azure-identity; install " |
| 123 | + "the storage-azure extra or set AZURE_STORAGE_ACCOUNT_KEY" |
| 124 | + ) |
| 125 | + raise StorageConfigurationError(msg) from exc |
| 126 | + |
| 127 | + client = BlobServiceClient(account_url, credential=DefaultAzureCredential()) |
| 128 | + return cls(client, account_name, account_key=None) |
| 129 | + |
| 130 | + async def upload( |
| 131 | + self, |
| 132 | + bucket: str, |
| 133 | + key: str, |
| 134 | + data: bytes, |
| 135 | + content_type: str | None = None, |
| 136 | + ) -> str: |
| 137 | + blob_client = self._client.get_blob_client(container=bucket, blob=key) |
| 138 | + content_settings = ( |
| 139 | + ContentSettings(content_type=content_type) if content_type else None |
| 140 | + ) |
| 141 | + await blob_client.upload_blob( |
| 142 | + data, overwrite=True, content_settings=content_settings |
| 143 | + ) |
| 144 | + return blob_client.url |
| 145 | + |
| 146 | + async def download(self, bucket: str, key: str) -> bytes: |
| 147 | + blob_client = self._client.get_blob_client(container=bucket, blob=key) |
| 148 | + try: |
| 149 | + stream = await blob_client.download_blob() |
| 150 | + except ResourceNotFoundError as exc: |
| 151 | + raise FileNotFoundError(f"{bucket}/{key}") from exc |
| 152 | + return await stream.readall() |
| 153 | + |
| 154 | + async def delete(self, bucket: str, key: str) -> None: |
| 155 | + blob_client = self._client.get_blob_client(container=bucket, blob=key) |
| 156 | + try: |
| 157 | + await blob_client.delete_blob() |
| 158 | + except ResourceNotFoundError: |
| 159 | + # Idempotent: deleting a missing blob is a no-op. |
| 160 | + return |
| 161 | + |
| 162 | + async def get_signed_url( |
| 163 | + self, bucket: str, key: str, expires_in: int = 3600 |
| 164 | + ) -> str: |
| 165 | + start = datetime.now(UTC) |
| 166 | + expiry = start + timedelta(seconds=expires_in) |
| 167 | + permission = BlobSasPermissions(read=True) |
| 168 | + |
| 169 | + if self._account_key: |
| 170 | + sas_token = generate_blob_sas( |
| 171 | + account_name=self._account_name, |
| 172 | + container_name=bucket, |
| 173 | + blob_name=key, |
| 174 | + account_key=self._account_key, |
| 175 | + permission=permission, |
| 176 | + expiry=expiry, |
| 177 | + start=start, |
| 178 | + ) |
| 179 | + else: |
| 180 | + user_delegation_key = await self._client.get_user_delegation_key( |
| 181 | + key_start_time=start, key_expiry_time=expiry |
| 182 | + ) |
| 183 | + sas_token = generate_blob_sas( |
| 184 | + account_name=self._account_name, |
| 185 | + container_name=bucket, |
| 186 | + blob_name=key, |
| 187 | + user_delegation_key=user_delegation_key, |
| 188 | + permission=permission, |
| 189 | + expiry=expiry, |
| 190 | + start=start, |
| 191 | + ) |
| 192 | + |
| 193 | + blob_client = self._client.get_blob_client(container=bucket, blob=key) |
| 194 | + return f"{blob_client.url}?{sas_token}" |
| 195 | + |
| 196 | + async def aclose(self) -> None: |
| 197 | + """Close the underlying client and its transport.""" |
| 198 | + await self._client.close() |
| 199 | + |
| 200 | + async def __aenter__(self) -> "AzureBlobStorageProvider": |
| 201 | + return self |
| 202 | + |
| 203 | + async def __aexit__(self, *_exc: object) -> None: |
| 204 | + await self.aclose() |
0 commit comments