Source code for multistorageclient.providers.azure

  1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  2# SPDX-License-Identifier: Apache-2.0
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8# http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15
 16import io
 17import os
 18import tempfile
 19from collections.abc import Callable, Iterator
 20from datetime import datetime, timedelta, timezone
 21from typing import IO, Any, TypeVar
 22from urllib.parse import quote, urlparse
 23
 24from azure.core import MatchConditions
 25from azure.core.exceptions import AzureError, HttpResponseError
 26from azure.identity import DefaultAzureCredential
 27from azure.storage.blob import (
 28    BlobPrefix,
 29    BlobSasPermissions,
 30    BlobServiceClient,
 31    PartialBatchErrorException,
 32    generate_blob_sas,
 33)
 34
 35from ..constants import DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT
 36from ..signers.base import URLSigner
 37from ..telemetry import Telemetry
 38from ..types import (
 39    AWARE_DATETIME_MIN,
 40    Credentials,
 41    CredentialsProvider,
 42    ObjectMetadata,
 43    PreconditionFailedError,
 44    Range,
 45    SignerType,
 46    SymlinkHandling,
 47)
 48from ..utils import safe_makedirs, split_path, validate_attributes
 49from .base import BaseStorageProvider
 50
 51_T = TypeVar("_T")
 52
 53PROVIDER = "azure"
 54AZURE_CONNECTION_STRING_KEY = "connection"
 55AZURE_CREDENTIAL_KEY = "azure_credential"
 56
 57MiB = 1024 * 1024
 58
 59MULTIPART_THRESHOLD = 64 * MiB
 60MULTIPART_CHUNKSIZE = 32 * MiB
 61IO_CHUNKSIZE = 32 * MiB
 62PYTHON_MAX_CONCURRENCY = 8
 63
 64# Azure REST API only returns ``Content-MD5`` for GET ranges up to 4 MiB,
 65# so download chunks must stay at or below this limit when ``validate_content`` is enabled.
 66AZURE_CONTENT_MD5_RANGE_LIMIT_BYTES = 4 * MiB
 67
 68DEFAULT_PRESIGN_EXPIRES_IN = 3600
 69
 70# How long before delegation key expiry we treat the cached key as stale.
 71_DELEGATION_KEY_REFRESH_BUFFER = timedelta(minutes=5)
 72
 73# Azure's maximum allowed delegation key lifetime is 7 days.
 74_DELEGATION_KEY_LIFETIME = timedelta(days=7)
 75
 76
 77def _sas_permissions_for_method(method: str) -> BlobSasPermissions:
 78    """Return the minimal :class:`BlobSasPermissions` needed for *method*."""
 79    m = method.upper()
 80    if m in ("PUT", "POST"):
 81        return BlobSasPermissions(write=True, create=True)
 82    elif m == "DELETE":
 83        return BlobSasPermissions(delete=True)
 84    else:
 85        # GET, HEAD, and any unrecognised method → read-only
 86        return BlobSasPermissions(read=True)
 87
 88
 89def _parse_account_name_from_url(account_url: str) -> str:
 90    """Extract the storage account name from an Azure Blob Storage account URL."""
 91    hostname = urlparse(account_url).hostname
 92    if hostname is None:
 93        raise ValueError(f"Invalid Azure account URL: {account_url!r}")
 94    return hostname.split(".")[0]
 95
 96
 97def _parse_connection_string(conn_str: str) -> dict[str, str]:
 98    """Parse an Azure connection string (``AccountName=foo;AccountKey=bar;...``) into a dict."""
 99    return dict(part.split("=", 1) for part in conn_str.split(";") if "=" in part)
100
101
[docs] 102class AzureURLSigner(URLSigner): 103 """ 104 Generates Azure Blob Storage SAS (Shared Access Signature) URLs. 105 106 Supports two signing paths depending on which credential is provided: 107 108 * **Account key** – uses a static storage account key (parsed from a connection string). 109 * **User delegation key** – uses a time-limited key obtained via Azure Identity (e.g. workload 110 identity, managed identity). Callers are responsible for refreshing the signer when the 111 delegation key approaches expiry; see :py:meth:`AzureBlobStorageProvider._generate_presigned_url`. 112 """ 113 114 def __init__( 115 self, 116 account_name: str, 117 account_url: str, 118 *, 119 account_key: str | None = None, 120 user_delegation_key: Any | None = None, 121 expires_in: int = DEFAULT_PRESIGN_EXPIRES_IN, 122 ) -> None: 123 if account_key is None and user_delegation_key is None: 124 raise ValueError("Either account_key or user_delegation_key must be provided.") 125 self._account_name = account_name 126 self._account_url = account_url.rstrip("/") 127 self._account_key = account_key 128 self._user_delegation_key = user_delegation_key 129 self._expires_in = expires_in 130
[docs] 131 def generate_presigned_url(self, path: str, *, method: str = "GET") -> str: 132 """ 133 Generate a SAS URL for the given blob path. 134 135 :param path: Path in the form ``container/blob/name``. 136 :param method: HTTP method requested by the caller. 137 :return: A fully-qualified SAS URL. 138 """ 139 container_name, blob_name = split_path(path) 140 expiry = datetime.now(timezone.utc) + timedelta(seconds=self._expires_in) 141 142 sas_kwargs: dict[str, Any] = { 143 "account_name": self._account_name, 144 "container_name": container_name, 145 "blob_name": blob_name, 146 "permission": _sas_permissions_for_method(method), 147 "expiry": expiry, 148 } 149 150 if self._account_key is not None: 151 sas_kwargs["account_key"] = self._account_key 152 else: 153 sas_kwargs["user_delegation_key"] = self._user_delegation_key 154 155 sas_token = generate_blob_sas(**sas_kwargs) 156 # Percent-encode the path the same way BlobClient.url does; the SAS itself is signed over the raw names. 157 blob_url = f"{self._account_url}/{quote(container_name)}/{quote(blob_name, safe='~/')}" 158 return f"{blob_url}?{sas_token}"
159 160
[docs] 161class StaticAzureCredentialsProvider(CredentialsProvider): 162 """ 163 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that provides static Azure credentials. 164 """ 165 166 _connection: str 167 168 def __init__(self, connection: str): 169 """ 170 Initializes the :py:class:`StaticAzureCredentialsProvider` with the provided connection string. 171 172 :param connection: The connection string for Azure Blob Storage authentication. 173 """ 174 self._connection = connection 175
[docs] 176 def get_credentials(self) -> Credentials: 177 return Credentials( 178 access_key=self._connection, 179 secret_key="", 180 token=None, 181 expiration=None, 182 custom_fields={AZURE_CONNECTION_STRING_KEY: self._connection}, 183 )
184
[docs] 185 def refresh_credentials(self) -> None: 186 pass
187 188
[docs] 189class DefaultAzureCredentialsProvider(CredentialsProvider): 190 """ 191 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that uses Azure Identity's :py:class:`azure.identity.DefaultAzureCredential` to authenticate with Blob Storage. 192 193 See :py:class:`azure.identity.DefaultAzureCredential` for provider options. 194 """ 195 196 def __init__(self, **kwargs: dict[str, Any]): 197 self._credential = DefaultAzureCredential(**kwargs) 198
[docs] 199 def get_credentials(self) -> Credentials: 200 return Credentials( 201 access_key="", 202 secret_key="", 203 token=None, 204 expiration=None, 205 custom_fields={AZURE_CREDENTIAL_KEY: self._credential}, 206 )
207
[docs] 208 def refresh_credentials(self) -> None: 209 pass
210 211
[docs] 212class AzureBlobStorageProvider(BaseStorageProvider): 213 """ 214 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with Azure Blob Storage. 215 """ 216 217 def __init__( 218 self, 219 endpoint_url: str, 220 base_path: str = "", 221 credentials_provider: CredentialsProvider | None = None, 222 config_dict: dict[str, Any] | None = None, 223 telemetry_provider: Callable[[], Telemetry] | None = None, 224 **kwargs: Any, 225 ): 226 """ 227 Initializes the :py:class:`AzureBlobStorageProvider` with the endpoint URL and optional credentials provider. 228 229 :param endpoint_url: The Azure storage account URL. 230 :param base_path: The root prefix path within the container where all operations will be scoped. 231 :param credentials_provider: The provider to retrieve Azure credentials. 232 :param config_dict: Resolved MSC config. 233 :param telemetry_provider: A function that provides a telemetry instance. 234 :param kwargs: Additional options including: 235 - ``multipart_threshold`` (int): File size threshold (bytes) for switching to parallel chunked transfers. Defaults to 64 MiB. 236 - ``multipart_chunksize`` (int): Block size (bytes) for chunked uploads. Defaults to 32 MiB. 237 - ``io_chunksize`` (int): Chunk size (bytes) for chunked downloads. Defaults to 32 MiB. 238 - ``max_concurrency`` (int): Number of parallel threads for chunked transfers. Defaults to 8. 239 - ``validate_content`` (bool): Opt-in client-side MD5 verification. Defaults to False. 240 """ 241 super().__init__( 242 base_path=base_path, 243 provider_name=PROVIDER, 244 config_dict=config_dict, 245 telemetry_provider=telemetry_provider, 246 ) 247 248 self._account_url = endpoint_url 249 self._credentials_provider = credentials_provider 250 # Cache static connection-string signing material used for per-request signers. 251 self._account_key_signing_material: tuple[str, str] | None = None 252 # Cached delegation key and its expiry for DefaultAzureCredentialsProvider. 253 self._delegation_user_key: Any | None = None 254 self._delegation_signer_expiry: datetime | None = None 255 self._multipart_threshold = int(kwargs.get("multipart_threshold", MULTIPART_THRESHOLD)) 256 self._multipart_chunksize = int(kwargs.get("multipart_chunksize", MULTIPART_CHUNKSIZE)) 257 self._io_chunksize = int(kwargs.get("io_chunksize", IO_CHUNKSIZE)) 258 self._max_concurrency = int(kwargs.get("max_concurrency", PYTHON_MAX_CONCURRENCY)) 259 self._validate_content = kwargs.get("validate_content", False) 260 if not isinstance(self._validate_content, bool): 261 raise TypeError("Option 'validate_content' must be a boolean.") 262 if self._validate_content and self._io_chunksize > AZURE_CONTENT_MD5_RANGE_LIMIT_BYTES: 263 raise ValueError( 264 "Option 'validate_content=True' requires 'io_chunksize' to be " 265 f"<= {AZURE_CONTENT_MD5_RANGE_LIMIT_BYTES} bytes (4 MiB) because Azure only " 266 f"returns Content-MD5 for GET ranges within that limit. Got io_chunksize={self._io_chunksize}." 267 ) 268 269 # https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/storage/azure-storage-blob#optional-configuration 270 client_optional_configuration_keys = { 271 "retry_total", 272 "retry_connect", 273 "retry_read", 274 "retry_status", 275 "connection_timeout", 276 "read_timeout", 277 } 278 self._client_optional_configuration: dict[str, Any] = { 279 key: value for key, value in kwargs.items() if key in client_optional_configuration_keys 280 } 281 if "connection_timeout" not in self._client_optional_configuration: 282 self._client_optional_configuration["connection_timeout"] = DEFAULT_CONNECT_TIMEOUT 283 if "read_timeout" not in self._client_optional_configuration: 284 self._client_optional_configuration["read_timeout"] = DEFAULT_READ_TIMEOUT 285 286 self._transfer_configuration: dict[str, Any] = { 287 "max_single_put_size": self._multipart_threshold, 288 "max_block_size": self._multipart_chunksize, 289 "max_single_get_size": self._multipart_threshold, 290 "max_chunk_get_size": self._io_chunksize, 291 } 292 293 self._blob_service_client = self._create_blob_service_client() 294 295 def _create_blob_service_client(self) -> BlobServiceClient: 296 """ 297 Creates and configures the Azure BlobServiceClient using the current credentials. 298 299 :return: The configured BlobServiceClient. 300 """ 301 combined_config = {**self._client_optional_configuration, **self._transfer_configuration} 302 303 if self._credentials_provider: 304 credentials = self._credentials_provider.get_credentials() 305 306 if isinstance(self._credentials_provider, StaticAzureCredentialsProvider): 307 return BlobServiceClient.from_connection_string( 308 credentials.get_custom_field(AZURE_CONNECTION_STRING_KEY), **combined_config 309 ) 310 elif isinstance(self._credentials_provider, DefaultAzureCredentialsProvider): 311 return BlobServiceClient( 312 account_url=self._account_url, 313 credential=credentials.get_custom_field(AZURE_CREDENTIAL_KEY), 314 **combined_config, 315 ) 316 else: 317 # Fallback to connection string if no built-in credentials provider is provided 318 return BlobServiceClient.from_connection_string(credentials.access_key, **combined_config) 319 else: 320 return BlobServiceClient(account_url=self._account_url, **combined_config) 321 322 def _refresh_blob_service_client_if_needed(self) -> None: 323 """ 324 Refreshes the BlobServiceClient if the current credentials are expired. 325 """ 326 if self._credentials_provider: 327 credentials = self._credentials_provider.get_credentials() 328 if credentials.is_expired(): 329 self._credentials_provider.refresh_credentials() 330 self._blob_service_client = self._create_blob_service_client() 331 332 def _translate_errors( 333 self, 334 func: Callable[[], _T], 335 operation: str, 336 container: str, 337 blob: str, 338 ) -> _T: 339 """ 340 Translates errors like timeouts and client errors. 341 342 :param func: The function that performs the actual Azure Blob Storage operation. 343 :param operation: The type of operation being performed (e.g., "PUT", "GET", "DELETE"). 344 :param container: The name of the Azure container involved in the operation. 345 :param blob: The name of the blob within the Azure container. 346 347 :return The result of the Azure Blob Storage operation, typically the return value of the `func` callable. 348 """ 349 try: 350 return func() 351 except HttpResponseError as error: 352 status_code = error.status_code if error.status_code else -1 353 error_info = f"status_code: {error.status_code}, reason: {error.reason}" 354 if status_code == 404: 355 raise FileNotFoundError(f"Object {container}/{blob} does not exist.") # pylint: disable=raise-missing-from 356 elif status_code == 412: 357 # raised when If-Match or If-Modified fails 358 raise PreconditionFailedError( 359 f"Failed to {operation} object(s) at {container}/{blob}. {error_info}" 360 ) from error 361 else: 362 raise RuntimeError(f"Failed to {operation} object(s) at {container}/{blob}. {error_info}") from error 363 except AzureError as error: 364 error_info = f"message: {error.message}" 365 raise RuntimeError(f"Failed to {operation} object(s) at {container}/{blob}. {error_info}") from error 366 except FileNotFoundError: 367 raise 368 except Exception as error: 369 raise RuntimeError( 370 f"Failed to {operation} object(s) at {container}/{blob}. error_type: {type(error).__name__}, error: {error}" 371 ) from error 372 373 def _put_object( 374 self, 375 path: str, 376 body: bytes, 377 if_match: str | None = None, 378 if_none_match: str | None = None, 379 attributes: dict[str, str] | None = None, 380 ) -> int: 381 """ 382 Uploads an object to Azure Blob Storage. 383 384 :param path: The path to the object to upload. 385 :param body: The content of the object to upload. 386 :param if_match: Optional ETag to match against the object. 387 :param if_none_match: Optional ETag to match against the object. 388 :param attributes: Optional attributes to attach to the object. 389 """ 390 container_name, blob_name = split_path(path) 391 self._refresh_blob_service_client_if_needed() 392 393 def _invoke_api() -> int: 394 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 395 396 kwargs: dict[str, Any] = { 397 "data": body, 398 "overwrite": True, 399 "max_concurrency": self._max_concurrency, 400 "validate_content": self._validate_content, 401 } 402 403 validated_attributes = validate_attributes(attributes) 404 if validated_attributes: 405 kwargs["metadata"] = validated_attributes 406 407 if if_match: 408 kwargs["match_condition"] = MatchConditions.IfNotModified 409 kwargs["etag"] = if_match 410 411 if if_none_match: 412 if if_none_match == "*": 413 raise NotImplementedError("if_none_match='*' is not supported for Azure") 414 kwargs["match_condition"] = MatchConditions.IfModified 415 kwargs["etag"] = if_none_match 416 417 blob_client.upload_blob(**kwargs) 418 419 return len(body) 420 421 return self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 422 423 def _get_object(self, path: str, byte_range: Range | None = None) -> bytes: 424 container_name, blob_name = split_path(path) 425 self._refresh_blob_service_client_if_needed() 426 427 def _invoke_api() -> bytes: 428 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 429 if byte_range: 430 stream = blob_client.download_blob( 431 offset=byte_range.offset, 432 length=byte_range.size, 433 validate_content=self._validate_content, 434 ) 435 else: 436 stream = blob_client.download_blob( 437 max_concurrency=self._max_concurrency, 438 validate_content=self._validate_content, 439 ) 440 return stream.readall() 441 442 return self._translate_errors(_invoke_api, operation="GET", container=container_name, blob=blob_name) 443 444 def _copy_object(self, src_path: str, dest_path: str) -> int: 445 src_container, src_blob = split_path(src_path) 446 dest_container, dest_blob = split_path(dest_path) 447 self._refresh_blob_service_client_if_needed() 448 449 src_object = self._get_object_metadata(src_path) 450 451 def _invoke_api() -> int: 452 src_blob_client = self._blob_service_client.get_blob_client(container=src_container, blob=src_blob) 453 dest_blob_client = self._blob_service_client.get_blob_client(container=dest_container, blob=dest_blob) 454 dest_blob_client.start_copy_from_url(src_blob_client.url) 455 456 return src_object.content_length 457 458 return self._translate_errors(_invoke_api, operation="COPY", container=src_container, blob=src_blob) 459 460 def _delete_object(self, path: str, if_match: str | None = None) -> None: 461 container_name, blob_name = split_path(path) 462 self._refresh_blob_service_client_if_needed() 463 464 def _invoke_api() -> None: 465 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 466 # If if_match is provided, use if_match for conditional deletion 467 if if_match: 468 blob_client.delete_blob(etag=if_match, match_condition=MatchConditions.IfNotModified) 469 else: 470 # No if_match provided, perform unconditional deletion 471 blob_client.delete_blob() 472 473 return self._translate_errors(_invoke_api, operation="DELETE", container=container_name, blob=blob_name) 474 475 def _delete_objects(self, paths: list[str]) -> None: 476 if not paths: 477 return 478 479 by_container: dict[str, list[str]] = {} 480 for p in paths: 481 container_name, blob_name = split_path(p) 482 by_container.setdefault(container_name, []).append(blob_name) 483 self._refresh_blob_service_client_if_needed() 484 485 AZURE_BATCH_LIMIT = 256 486 487 def _invoke_api() -> None: 488 for container_name, blob_names in by_container.items(): 489 container_client = self._blob_service_client.get_container_client(container=container_name) 490 for i in range(0, len(blob_names), AZURE_BATCH_LIMIT): 491 chunk = blob_names[i : i + AZURE_BATCH_LIMIT] 492 try: 493 responses = container_client.delete_blobs(*chunk, raise_on_any_failure=False) 494 except PartialBatchErrorException as error: 495 responses = error.parts 496 for response in responses: 497 status_code = response.status_code 498 if 200 <= status_code < 300 or status_code == 404: 499 continue 500 raise RuntimeError( 501 f"Azure batch delete failed with status_code: {status_code}, response: {response.text()}" 502 ) 503 504 container_desc = "(" + "|".join(by_container) + ")" 505 blob_desc = "(" + "|".join(str(len(blob_names)) for blob_names in by_container.values()) + " keys)" 506 self._translate_errors(_invoke_api, operation="DELETE_MANY", container=container_desc, blob=blob_desc) 507 508 def _is_dir(self, path: str) -> bool: 509 # Ensure the path ends with '/' to mimic a directory 510 path = self._append_delimiter(path) 511 512 container_name, prefix = split_path(path) 513 self._refresh_blob_service_client_if_needed() 514 515 def _invoke_api() -> bool: 516 # List objects with the given prefix 517 container_client = self._blob_service_client.get_container_client(container=container_name) 518 blobs = container_client.walk_blobs(name_starts_with=prefix, delimiter="/") 519 # Check if there are any contents or common prefixes 520 return any(True for _ in blobs) 521 522 return self._translate_errors(_invoke_api, operation="LIST", container=container_name, blob=prefix) 523 524 def _make_symlink(self, path: str, target: str) -> None: 525 container_name, blob_name = split_path(path) 526 target_container, target_key = split_path(target) 527 if container_name != target_container: 528 raise ValueError(f"Cannot create cross-container symlink: '{container_name}' -> '{target_container}'.") 529 relative_target = ObjectMetadata.encode_symlink_target(blob_name, target_key) 530 self._refresh_blob_service_client_if_needed() 531 532 def _invoke_api() -> None: 533 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 534 blob_client.upload_blob( 535 data=b"", 536 overwrite=True, 537 metadata={"msc_symlink_target": relative_target}, 538 ) 539 540 self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 541 542 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata: 543 container_name, blob_name = split_path(path) 544 if path.endswith("/") or (container_name and not blob_name): 545 # If path ends with "/" or empty blob name is provided, then assume it's a "directory", 546 # which metadata is not guaranteed to exist for cases such as 547 # "virtual prefix" that was never explicitly created. 548 if self._is_dir(path): 549 return ObjectMetadata( 550 key=self._append_delimiter(path), 551 type="directory", 552 content_length=0, 553 last_modified=AWARE_DATETIME_MIN, 554 ) 555 else: 556 raise FileNotFoundError(f"Directory {path} does not exist.") 557 else: 558 self._refresh_blob_service_client_if_needed() 559 560 def _invoke_api() -> ObjectMetadata: 561 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 562 properties = blob_client.get_blob_properties() 563 user_metadata = dict(properties.metadata) if properties.metadata else None 564 symlink_target = user_metadata.get("msc_symlink_target") if user_metadata else None 565 return ObjectMetadata( 566 key=path, 567 content_length=properties.size, 568 content_type=properties.content_settings.content_type, 569 last_modified=properties.last_modified, 570 etag=properties.etag.strip('"') if properties.etag else "", 571 metadata=user_metadata, 572 symlink_target=symlink_target, 573 ) 574 575 try: 576 return self._translate_errors(_invoke_api, operation="HEAD", container=container_name, blob=blob_name) 577 except FileNotFoundError: 578 if strict: 579 # If the object does not exist on the given path, we will append a trailing slash and 580 # check if the path is a directory. 581 path = self._append_delimiter(path) 582 if self._is_dir(path): 583 return ObjectMetadata( 584 key=path, 585 type="directory", 586 content_length=0, 587 last_modified=AWARE_DATETIME_MIN, 588 ) 589 raise 590 591 def _list_objects( 592 self, 593 path: str, 594 start_after: str | None = None, 595 end_at: str | None = None, 596 include_directories: bool = False, 597 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW, 598 ) -> Iterator[ObjectMetadata]: 599 container_name, prefix = split_path(path) 600 601 # Get the prefix of the start_after and end_at paths relative to the bucket. 602 if start_after: 603 _, start_after = split_path(start_after) 604 if end_at: 605 _, end_at = split_path(end_at) 606 607 self._refresh_blob_service_client_if_needed() 608 609 def _invoke_api() -> Iterator[ObjectMetadata]: 610 container_client = self._blob_service_client.get_container_client(container=container_name) 611 # Azure has no start key option like other object stores. 612 if include_directories: 613 blobs = container_client.walk_blobs(name_starts_with=prefix, delimiter="/", include=["metadata"]) 614 else: 615 blobs = container_client.list_blobs(name_starts_with=prefix, include=["metadata"]) 616 # Azure guarantees lexicographical order. 617 for blob in blobs: 618 if isinstance(blob, BlobPrefix): 619 prefix_key = blob.name.rstrip("/") 620 # Filter by start_after and end_at if specified 621 if (start_after is None or start_after < prefix_key) and (end_at is None or prefix_key <= end_at): 622 yield ObjectMetadata( 623 key=os.path.join(container_name, prefix_key), 624 type="directory", 625 content_length=0, 626 last_modified=AWARE_DATETIME_MIN, 627 ) 628 elif end_at is not None and end_at < prefix_key: 629 return 630 else: 631 key = blob.name 632 if (start_after is None or start_after < key) and (end_at is None or key <= end_at): 633 if key.endswith("/"): 634 if include_directories: 635 yield ObjectMetadata( 636 key=os.path.join(container_name, key.rstrip("/")), 637 type="directory", 638 content_length=0, 639 last_modified=blob.last_modified, 640 ) 641 else: 642 user_metadata = dict(blob.metadata) if blob.metadata else None 643 symlink_target = user_metadata.get("msc_symlink_target") if user_metadata else None 644 yield ObjectMetadata( 645 key=os.path.join(container_name, key), 646 content_length=blob.size, 647 content_type=blob.content_settings.content_type, 648 last_modified=blob.last_modified, 649 etag=blob.etag.strip('"') if blob.etag else "", 650 symlink_target=symlink_target, 651 ) 652 elif end_at is not None and end_at < key: 653 return 654 655 return self._translate_errors(_invoke_api, operation="LIST", container=container_name, blob=prefix) 656 657 def _generate_presigned_url( 658 self, 659 path: str, 660 *, 661 method: str = "GET", 662 signer_type: SignerType | None = None, 663 signer_options: dict[str, Any] | None = None, 664 ) -> str: 665 """ 666 Generate a SAS URL for a blob in Azure Blob Storage. 667 668 :param path: Path in the form ``container/blob/name``. 669 :param method: HTTP method requested by the caller. 670 :param signer_type: Must be ``None`` or :py:attr:`SignerType.AZURE`. 671 :param signer_options: Optional dict; supports ``expires_in`` (int, seconds). 672 :return: A fully-qualified SAS URL. 673 :raises ValueError: If *signer_type* is not ``None`` / ``SignerType.AZURE``. 674 :raises TypeError: If the configured credential type does not support SAS generation. 675 """ 676 if signer_type is not None and signer_type != SignerType.AZURE: 677 raise ValueError(f"Unsupported signer type for Azure provider: {signer_type!r}") 678 679 options = signer_options or {} 680 expires_in = int(options.get("expires_in", DEFAULT_PRESIGN_EXPIRES_IN)) 681 682 self._refresh_blob_service_client_if_needed() 683 684 if isinstance(self._credentials_provider, StaticAzureCredentialsProvider): 685 # Account key path: cache parsed AccountName + AccountKey, then sign per request. 686 if self._account_key_signing_material is None: 687 conn_str = self._credentials_provider.get_credentials().get_custom_field(AZURE_CONNECTION_STRING_KEY) 688 parsed = _parse_connection_string(conn_str) 689 self._account_key_signing_material = (parsed["AccountName"], parsed["AccountKey"]) 690 account_name, account_key = self._account_key_signing_material 691 signer = AzureURLSigner( 692 account_name=account_name, 693 account_url=self._account_url, 694 account_key=account_key, 695 expires_in=expires_in, 696 ) 697 698 elif isinstance(self._credentials_provider, DefaultAzureCredentialsProvider): 699 # User delegation key path: refresh when the cached key is within the 700 # refresh buffer of its own expiry or has not been fetched yet. 701 now = datetime.now(timezone.utc) 702 if ( 703 self._delegation_user_key is None 704 or self._delegation_signer_expiry is None 705 or now >= self._delegation_signer_expiry - _DELEGATION_KEY_REFRESH_BUFFER 706 ): 707 key_expiry = now + _DELEGATION_KEY_LIFETIME 708 self._delegation_user_key = self._blob_service_client.get_user_delegation_key( 709 key_start_time=now, 710 key_expiry_time=key_expiry, 711 ) 712 self._delegation_signer_expiry = key_expiry 713 signer = AzureURLSigner( 714 account_name=_parse_account_name_from_url(self._account_url), 715 account_url=self._account_url, 716 user_delegation_key=self._delegation_user_key, 717 expires_in=expires_in, 718 ) 719 720 else: 721 raise TypeError( 722 "Azure presigned URLs require StaticAzureCredentialsProvider (connection string) or " 723 "DefaultAzureCredentialsProvider (Azure Identity). " 724 f"Got: {type(self._credentials_provider).__name__!r}" 725 ) 726 727 return signer.generate_presigned_url(path, method=method) 728 729 @property 730 def supports_parallel_listing(self) -> bool: 731 return True 732 733 def _upload_file(self, remote_path: str, f: str | IO, attributes: dict[str, str] | None = None) -> int: 734 container_name, blob_name = split_path(remote_path) 735 file_size: int = 0 736 self._refresh_blob_service_client_if_needed() 737 738 validated_attributes = validate_attributes(attributes) 739 if isinstance(f, str): 740 file_size = os.path.getsize(f) 741 742 if file_size <= self._multipart_threshold: 743 744 def _invoke_api() -> int: 745 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 746 with open(f, "rb") as data: 747 blob_client.upload_blob( 748 data, 749 overwrite=True, 750 metadata=validated_attributes or {}, 751 validate_content=self._validate_content, 752 ) 753 return file_size 754 755 return self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 756 757 def _invoke_api() -> int: 758 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 759 with open(f, "rb") as data: 760 blob_client.upload_blob( 761 data, 762 overwrite=True, 763 metadata=validated_attributes or {}, 764 max_concurrency=self._max_concurrency, 765 validate_content=self._validate_content, 766 ) 767 return file_size 768 769 return self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 770 else: 771 if isinstance(f, io.StringIO): 772 fp: IO = io.BytesIO(f.getvalue().encode("utf-8")) # type: ignore 773 else: 774 fp = f 775 776 fp.seek(0, io.SEEK_END) 777 file_size = fp.tell() 778 fp.seek(0) 779 780 if file_size <= self._multipart_threshold: 781 782 def _invoke_api() -> int: 783 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 784 blob_client.upload_blob( 785 fp, 786 overwrite=True, 787 metadata=validated_attributes or {}, 788 validate_content=self._validate_content, 789 ) 790 return file_size 791 792 return self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 793 794 def _invoke_api() -> int: 795 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 796 blob_client.upload_blob( 797 fp, 798 overwrite=True, 799 metadata=validated_attributes or {}, 800 max_concurrency=self._max_concurrency, 801 validate_content=self._validate_content, 802 ) 803 return file_size 804 805 return self._translate_errors(_invoke_api, operation="PUT", container=container_name, blob=blob_name) 806 807 def _download_file(self, remote_path: str, f: str | IO, metadata: ObjectMetadata | None = None) -> int: 808 if metadata is None: 809 metadata = self._get_object_metadata(remote_path) 810 811 container_name, blob_name = split_path(remote_path) 812 self._refresh_blob_service_client_if_needed() 813 814 if isinstance(f, str): 815 if os.path.dirname(f): 816 safe_makedirs(os.path.dirname(f)) 817 818 if metadata.content_length <= self._multipart_threshold: 819 820 def _invoke_api() -> int: 821 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 822 temp_file_path: str | None = None 823 try: 824 with tempfile.NamedTemporaryFile( 825 mode="wb", delete=False, dir=os.path.dirname(f), prefix="." 826 ) as fp: 827 temp_file_path = fp.name 828 stream = blob_client.download_blob(validate_content=self._validate_content) 829 fp.write(stream.readall()) 830 os.rename(src=temp_file_path, dst=f) 831 except BaseException: 832 if temp_file_path and os.path.exists(temp_file_path): 833 os.unlink(temp_file_path) 834 raise 835 return metadata.content_length 836 837 return self._translate_errors(_invoke_api, operation="GET", container=container_name, blob=blob_name) 838 839 def _invoke_api() -> int: 840 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 841 temp_file_path: str | None = None 842 try: 843 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(f), prefix=".") as fp: 844 temp_file_path = fp.name 845 stream = blob_client.download_blob( 846 max_concurrency=self._max_concurrency, 847 validate_content=self._validate_content, 848 ) 849 stream.readinto(fp) 850 os.rename(src=temp_file_path, dst=f) 851 except BaseException: 852 if temp_file_path and os.path.exists(temp_file_path): 853 os.unlink(temp_file_path) 854 raise 855 return metadata.content_length 856 857 return self._translate_errors(_invoke_api, operation="GET", container=container_name, blob=blob_name) 858 else: 859 if metadata.content_length <= self._multipart_threshold: 860 861 def _invoke_api() -> int: 862 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 863 stream = blob_client.download_blob(validate_content=self._validate_content) 864 if isinstance(f, io.StringIO): 865 f.write(stream.readall().decode("utf-8")) 866 else: 867 f.write(stream.readall()) 868 return metadata.content_length 869 870 return self._translate_errors(_invoke_api, operation="GET", container=container_name, blob=blob_name) 871 872 def _invoke_api() -> int: 873 blob_client = self._blob_service_client.get_blob_client(container=container_name, blob=blob_name) 874 stream = blob_client.download_blob( 875 max_concurrency=self._max_concurrency, 876 validate_content=self._validate_content, 877 ) 878 if isinstance(f, io.StringIO): 879 temp_file_path: str | None = None 880 try: 881 with tempfile.NamedTemporaryFile(mode="wb", delete=False, prefix=".") as tmp: 882 temp_file_path = tmp.name 883 stream.readinto(tmp) 884 with open(temp_file_path, "r", encoding="utf-8") as tmp_read: 885 f.write(tmp_read.read()) 886 finally: 887 if temp_file_path and os.path.exists(temp_file_path): 888 os.unlink(temp_file_path) 889 else: 890 stream.readinto(f) 891 return metadata.content_length 892 893 return self._translate_errors(_invoke_api, operation="GET", container=container_name, blob=blob_name)