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 codecs
17import io
18import logging
19import os
20import tempfile
21from collections.abc import Callable, Iterator
22from typing import IO, Any, TypeVar
23
24import boto3
25import botocore
26from boto3.exceptions import S3UploadFailedError
27from boto3.s3.transfer import TransferConfig
28from botocore.credentials import RefreshableCredentials
29from botocore.exceptions import ClientError, IncompleteReadError, ReadTimeoutError, ResponseStreamingError
30from botocore.session import get_session
31from dateutil.parser import parse as dateutil_parse
32
33from multistorageclient_rust import RustClient, RustClientError, RustRetryableError
34
35from ..constants import DEFAULT_CONNECT_TIMEOUT, DEFAULT_MAX_POOL_CONNECTIONS, DEFAULT_READ_TIMEOUT
36from ..rust_utils import parse_retry_config, run_async_rust_client_method
37from ..signers import CloudFrontURLSigner, URLSigner
38from ..telemetry import Telemetry
39from ..types import (
40 AWARE_DATETIME_MIN,
41 Credentials,
42 CredentialsProvider,
43 ObjectMetadata,
44 PreconditionFailedError,
45 Range,
46 RetryableError,
47 SignerType,
48 SuffixRange,
49 SymlinkHandling,
50)
51from ..utils import (
52 safe_makedirs,
53 split_path,
54 validate_attributes,
55)
56from .base import BaseStorageProvider, _translate_iterator_errors
57
58logger = logging.getLogger(__name__)
59
60_T = TypeVar("_T")
61
62MiB = 1024 * 1024
63
64# Python and Rust share the same multipart_threshold to keep the code simple.
65MULTIPART_THRESHOLD = 64 * MiB
66MULTIPART_CHUNKSIZE = 32 * MiB
67IO_CHUNKSIZE = 32 * MiB
68PYTHON_MAX_CONCURRENCY = 8
69MAX_POOL_CONNECTIONS = DEFAULT_MAX_POOL_CONNECTIONS
70
71PROVIDER = "s3"
72
73EXPRESS_ONEZONE_STORAGE_CLASS = "EXPRESS_ONEZONE"
74
75# Source: https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html
76SUPPORTED_CHECKSUM_ALGORITHMS = frozenset({"CRC32", "CRC32C", "SHA1", "SHA256", "CRC64NVME"})
77
78
[docs]
79class StaticS3CredentialsProvider(CredentialsProvider):
80 """
81 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that provides static S3 credentials.
82 """
83
84 _access_key: str
85 _secret_key: str
86 _session_token: str | None
87
88 def __init__(self, access_key: str, secret_key: str, session_token: str | None = None):
89 """
90 Initializes the :py:class:`StaticS3CredentialsProvider` with the provided access key, secret key, and optional
91 session token.
92
93 :param access_key: The access key for S3 authentication.
94 :param secret_key: The secret key for S3 authentication.
95 :param session_token: An optional session token for temporary credentials.
96 """
97 self._access_key = access_key
98 self._secret_key = secret_key
99 self._session_token = session_token
100
[docs]
101 def get_credentials(self) -> Credentials:
102 return Credentials(
103 access_key=self._access_key,
104 secret_key=self._secret_key,
105 token=self._session_token,
106 expiration=None,
107 )
108
[docs]
109 def refresh_credentials(self) -> None:
110 pass
111
112
113DEFAULT_PRESIGN_EXPIRES_IN = 3600
114
115_S3_METHOD_MAPPING: dict[str, str] = {
116 "GET": "get_object",
117 "PUT": "put_object",
118}
119
120
[docs]
121class S3URLSigner(URLSigner):
122 """Generates pre-signed URLs using the boto3 S3 client.
123
124 When the underlying credentials are temporary (STS, IAM role, EC2 instance
125 profile), the effective URL lifetime is the **shorter** of ``expires_in``
126 and the remaining credential lifetime — boto3 will not warn if the
127 credential expires before ``expires_in``.
128
129 See https://docs.aws.amazon.com/AmazonS3/latest/userguide/using-presigned-url.html
130 """
131
132 def __init__(self, s3_client: Any, bucket: str, expires_in: int = DEFAULT_PRESIGN_EXPIRES_IN) -> None:
133 self._s3_client = s3_client
134 self._bucket = bucket
135 self._expires_in = expires_in
136
[docs]
137 def generate_presigned_url(self, path: str, *, method: str = "GET") -> str:
138 client_method = _S3_METHOD_MAPPING.get(method.upper())
139 if client_method is None:
140 raise ValueError(f"Unsupported method for S3 presigning: {method!r}")
141 return self._s3_client.generate_presigned_url(
142 ClientMethod=client_method,
143 Params={"Bucket": self._bucket, "Key": path},
144 ExpiresIn=self._expires_in,
145 )
146
147
[docs]
148class S3StorageProvider(BaseStorageProvider):
149 """
150 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with Amazon S3 or S3-compatible object stores.
151 """
152
153 def __init__(
154 self,
155 region_name: str = "",
156 endpoint_url: str = "",
157 base_path: str = "",
158 credentials_provider: CredentialsProvider | None = None,
159 profile_name: str | None = None,
160 config_dict: dict[str, Any] | None = None,
161 telemetry_provider: Callable[[], Telemetry] | None = None,
162 verify: bool | str | None = None,
163 **kwargs: Any,
164 ) -> None:
165 """
166 Initializes the :py:class:`S3StorageProvider` with the region, endpoint URL, and optional credentials provider.
167
168 :param region_name: The AWS region where the S3 bucket is located.
169 :param endpoint_url: The custom endpoint URL for the S3 service.
170 :param base_path: The root prefix path within the S3 bucket where all operations will be scoped.
171 :param credentials_provider: The provider to retrieve S3 credentials.
172 :param profile_name: AWS shared configuration + credentials files profile. For :py:class:`boto3.session.Session`. Ignored if ``credentials_provider`` is set.
173 :param config_dict: Resolved MSC config.
174 :param telemetry_provider: A function that provides a telemetry instance.
175 :param verify: Controls SSL certificate verification.
176 Can be ``True`` (verify using system CA bundle, default), ``False`` (skip verification), or a string path to a custom CA certificate bundle.
177 :param request_checksum_calculation: For :py:class:`botocore.config.Config`.
178 When the underlying S3 client should calculate request checksums.
179 :param response_checksum_validation: For :py:class:`botocore.config.Config`.
180 When the underlying S3 client should validate response checksums.
181 :param max_pool_connections: For :py:class:`botocore.config.Config`.
182 The maximum number of connections to keep in a connection pool.
183 :param connect_timeout: For :py:class:`botocore.config.Config`.
184 The time in seconds till a timeout exception is thrown when attempting to make a connection.
185 :param read_timeout: For :py:class:`botocore.config.Config`.
186 The time in seconds till a timeout exception is thrown when attempting to read from a connection.
187 :param retries: For :py:class:`botocore.config.Config`.
188 A dictionary for configuration related to retry behavior.
189 :param s3: For :py:class:`botocore.config.Config`.
190 A dictionary of S3 specific configurations.
191 :param checksum_algorithm: Upload-only object integrity algorithm.
192 One of ``"CRC32"``, ``"CRC32C"``, ``"SHA1"``, ``"SHA256"``, ``"CRC64NVME"`` (case-insensitive).
193 When ``rust_client`` is enabled, only ``"SHA256"`` is accepted.
194 :param multipart_threshold: For :py:class:`boto3.s3.transfer.TransferConfig`.
195 The transfer size threshold for which multipart uploads, downloads, and copies will automatically be triggered.
196 :param max_concurrency: For :py:class:`boto3.s3.transfer.TransferConfig`.
197 The maximum number of threads that will be making requests to perform a transfer.
198 If ``use_threads`` is set to ``False``, the value provided is ignored as the transfer will only ever use the current thread.
199 :param multipart_chunksize: For :py:class:`boto3.s3.transfer.TransferConfig`.
200 The partition size of each part for a multipart transfer.
201 :param io_chunksize: For :py:class:`boto3.s3.transfer.TransferConfig`.
202 The max size of each chunk in the ``io`` queue. Currently, this is the size used when read is called on the downloaded stream as well.
203 Note: This value is ignored when resolved transfer manager type is CRTTransferManager.
204 """
205 super().__init__(
206 base_path=base_path,
207 provider_name=PROVIDER,
208 config_dict=config_dict,
209 telemetry_provider=telemetry_provider,
210 )
211
212 self._region_name = region_name
213 self._endpoint_url = endpoint_url
214 self._credentials_provider = credentials_provider
215 self._profile_name = profile_name
216 self._verify = verify
217
218 self._signature_version = kwargs.get("signature_version", "s3v4")
219 self._s3_client = self._create_s3_client(
220 request_checksum_calculation=kwargs.get("request_checksum_calculation"),
221 response_checksum_validation=kwargs.get("response_checksum_validation"),
222 max_pool_connections=kwargs.get("max_pool_connections", MAX_POOL_CONNECTIONS),
223 connect_timeout=kwargs.get("connect_timeout", DEFAULT_CONNECT_TIMEOUT),
224 read_timeout=kwargs.get("read_timeout", DEFAULT_READ_TIMEOUT),
225 retries=kwargs.get("retries"),
226 s3=kwargs.get("s3"),
227 )
228 self._transfer_config = TransferConfig(
229 multipart_threshold=int(kwargs.get("multipart_threshold", MULTIPART_THRESHOLD)),
230 max_concurrency=int(kwargs.get("max_concurrency", PYTHON_MAX_CONCURRENCY)),
231 multipart_chunksize=int(kwargs.get("multipart_chunksize", MULTIPART_CHUNKSIZE)),
232 io_chunksize=int(kwargs.get("io_chunksize", IO_CHUNKSIZE)),
233 use_threads=True,
234 )
235 self._multipart_threshold = self._transfer_config.multipart_threshold
236
237 self._signer_cache: dict[tuple, URLSigner] = {}
238
239 self._checksum_algorithm: str | None = self._validate_checksum_algorithm(kwargs.get("checksum_algorithm"))
240
241 self._rust_client = None
242 if "rust_client" in kwargs:
243 # Inherit the rust client options from the kwargs
244 rust_client_options = kwargs["rust_client"]
245 if "max_pool_connections" in kwargs:
246 rust_client_options["max_pool_connections"] = kwargs["max_pool_connections"]
247 if "max_concurrency" in kwargs:
248 rust_client_options["max_concurrency"] = kwargs["max_concurrency"]
249 if "multipart_chunksize" in kwargs:
250 rust_client_options["multipart_chunksize"] = kwargs["multipart_chunksize"]
251 if "read_timeout" in kwargs:
252 rust_client_options["read_timeout"] = kwargs["read_timeout"]
253 if "connect_timeout" in kwargs:
254 rust_client_options["connect_timeout"] = kwargs["connect_timeout"]
255 if "checksum_algorithm" in kwargs:
256 rust_client_options["checksum_algorithm"] = kwargs["checksum_algorithm"]
257 if self._signature_version == "UNSIGNED":
258 rust_client_options["skip_signature"] = True
259
260 # Rust client only supports SHA256 today (object_store limitation).
261 rust_checksum = rust_client_options.get("checksum_algorithm")
262 if rust_checksum is not None:
263 if str(rust_checksum).upper() != "SHA256":
264 raise ValueError(
265 f"Rust client only supports checksum_algorithm='SHA256' today "
266 f"(object_store limitation), got '{rust_checksum!r}'. "
267 f"Disable rust_client or use SHA256."
268 )
269 rust_client_options["checksum_algorithm"] = "sha256"
270 self._rust_client = self._create_rust_client(rust_client_options)
271
272 @staticmethod
273 def _validate_checksum_algorithm(value: str | None) -> str | None:
274 """
275 Normalizes the optional ``checksum_algorithm`` to uppercase, or returns ``None`` if disabled.
276 """
277 if value is None:
278 return None
279 if not isinstance(value, str) or not value:
280 raise ValueError(
281 f"checksum_algorithm must be one of {sorted(SUPPORTED_CHECKSUM_ALGORITHMS)} or None, got {value!r}."
282 )
283 normalized = value.upper()
284 if normalized not in SUPPORTED_CHECKSUM_ALGORITHMS:
285 raise ValueError(
286 f"Unsupported checksum_algorithm '{value}'. Supported: {sorted(SUPPORTED_CHECKSUM_ALGORITHMS)}."
287 )
288 return normalized
289
290 def _is_directory_bucket(self, bucket: str) -> bool:
291 """
292 Determines if the bucket is a directory bucket based on bucket name.
293 """
294 # S3 Express buckets have a specific naming convention
295 return "--x-s3" in bucket
296
297 def _create_s3_client(
298 self,
299 request_checksum_calculation: str | None = None,
300 response_checksum_validation: str | None = None,
301 max_pool_connections: int = MAX_POOL_CONNECTIONS,
302 connect_timeout: float | None = None,
303 read_timeout: float | None = None,
304 retries: dict[str, Any] | None = None,
305 s3: dict[str, Any] | None = None,
306 ):
307 """
308 Creates and configures the boto3 S3 client, using refreshable credentials if possible.
309
310 :param request_checksum_calculation: For :py:class:`botocore.config.Config`. When the underlying S3 client should calculate request checksums. See the equivalent option in the `AWS configuration file <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-a-configuration-file>`_.
311 :param response_checksum_validation: For :py:class:`botocore.config.Config`. When the underlying S3 client should validate response checksums. See the equivalent option in the `AWS configuration file <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-a-configuration-file>`_.
312 :param max_pool_connections: For :py:class:`botocore.config.Config`. The maximum number of connections to keep in a connection pool.
313 :param connect_timeout: For :py:class:`botocore.config.Config`. The time in seconds till a timeout exception is thrown when attempting to make a connection.
314 :param read_timeout: For :py:class:`botocore.config.Config`. The time in seconds till a timeout exception is thrown when attempting to read from a connection.
315 :param retries: For :py:class:`botocore.config.Config`. A dictionary for configuration related to retry behavior.
316 :param s3: For :py:class:`botocore.config.Config`. A dictionary of S3 specific configurations.
317
318 :return: The configured S3 client.
319 """
320 options = {
321 # https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
322 "config": boto3.session.Config( # pyright: ignore [reportAttributeAccessIssue]
323 max_pool_connections=max_pool_connections,
324 connect_timeout=connect_timeout,
325 read_timeout=read_timeout,
326 retries=retries or {"mode": "standard"},
327 request_checksum_calculation=request_checksum_calculation,
328 response_checksum_validation=response_checksum_validation,
329 s3=s3,
330 ),
331 }
332
333 if self._region_name:
334 options["region_name"] = self._region_name
335
336 if self._endpoint_url:
337 options["endpoint_url"] = self._endpoint_url
338
339 if self._verify is not None:
340 options["verify"] = self._verify
341
342 if self._credentials_provider:
343 creds = self._fetch_credentials()
344 if creds.get("expiry_time"):
345 # Use RefreshableCredentials if expiry_time provided.
346 refreshable_credentials = RefreshableCredentials.create_from_metadata(
347 metadata=creds, refresh_using=self._fetch_credentials, method="custom-refresh"
348 )
349
350 botocore_session = get_session()
351 botocore_session._credentials = refreshable_credentials
352
353 boto3_session = boto3.Session(botocore_session=botocore_session)
354
355 return boto3_session.client("s3", **options)
356 else:
357 # Add static credentials to the options dictionary
358 options["aws_access_key_id"] = creds["access_key"]
359 options["aws_secret_access_key"] = creds["secret_key"]
360 if creds["token"]:
361 options["aws_session_token"] = creds["token"]
362
363 if self._signature_version:
364 signature_config = botocore.config.Config( # pyright: ignore[reportAttributeAccessIssue]
365 signature_version=botocore.UNSIGNED
366 if self._signature_version == "UNSIGNED"
367 else self._signature_version
368 )
369 options["config"] = options["config"].merge(signature_config)
370
371 # Fallback to standard credential chain.
372 session = boto3.Session(profile_name=self._profile_name)
373 return session.client("s3", **options)
374
375 def _create_rust_client(self, rust_client_options: dict[str, Any] | None = None):
376 """
377 Creates and configures the rust client, using refreshable credentials if possible.
378 """
379 configs = dict(rust_client_options) if rust_client_options else {}
380
381 # Extract and parse retry configuration
382 retry_config = parse_retry_config(configs)
383
384 if self._region_name and "region_name" not in configs:
385 configs["region_name"] = self._region_name
386
387 if self._endpoint_url and "endpoint_url" not in configs:
388 configs["endpoint_url"] = self._endpoint_url
389
390 if self._profile_name and "profile_name" not in configs:
391 configs["profile_name"] = self._profile_name
392
393 # If the user specifies a bucket, use it. Otherwise, use the base path.
394 if "bucket" not in configs:
395 bucket, _ = split_path(self._base_path)
396 configs["bucket"] = bucket
397
398 if "max_pool_connections" not in configs:
399 configs["max_pool_connections"] = MAX_POOL_CONNECTIONS
400
401 return RustClient(
402 provider=PROVIDER,
403 configs=configs,
404 credentials_provider=self._credentials_provider,
405 retry=retry_config,
406 )
407
408 def _fetch_credentials(self) -> dict:
409 """
410 Refreshes the S3 client if the current credentials are expired.
411 """
412 if not self._credentials_provider:
413 raise RuntimeError("Cannot fetch credentials if no credential provider configured.")
414 self._credentials_provider.refresh_credentials()
415 credentials = self._credentials_provider.get_credentials()
416 return {
417 "access_key": credentials.access_key,
418 "secret_key": credentials.secret_key,
419 "token": credentials.token,
420 "expiry_time": credentials.expiration,
421 }
422
423 def _translate_errors(
424 self,
425 func: Callable[[], _T],
426 operation: str,
427 bucket: str,
428 key: str,
429 ) -> _T:
430 """
431 Translates errors like timeouts and client errors.
432
433 :param func: The function that performs the actual S3 operation.
434 :param operation: The type of operation being performed (e.g., "PUT", "GET", "DELETE").
435 :param bucket: The name of the S3 bucket involved in the operation.
436 :param key: The key of the object within the S3 bucket.
437
438 :return: The result of the S3 operation, typically the return value of the `func` callable.
439 """
440 try:
441 return func()
442 except ClientError as error:
443 status_code = error.response["ResponseMetadata"]["HTTPStatusCode"]
444 request_id = error.response["ResponseMetadata"].get("RequestId")
445 host_id = error.response["ResponseMetadata"].get("HostId")
446 error_code = error.response["Error"]["Code"]
447 error_info = f"request_id: {request_id}, host_id: {host_id}, status_code: {status_code}"
448
449 if status_code == 404:
450 if error_code == "NoSuchUpload":
451 error_message = error.response["Error"]["Message"]
452 raise RetryableError(f"Multipart upload failed for {bucket}/{key}: {error_message}") from error
453 raise FileNotFoundError(f"Object {bucket}/{key} does not exist. {error_info}") # pylint: disable=raise-missing-from
454 elif status_code == 403:
455 raise PermissionError(
456 f"Permission denied to {operation} object(s) at {bucket}/{key}. {error_info}"
457 ) from error
458 elif status_code == 412: # Precondition Failed
459 raise PreconditionFailedError(
460 f"ETag mismatch for {operation} operation on {bucket}/{key}. {error_info}"
461 ) from error
462 elif status_code == 429:
463 raise RetryableError(
464 f"Too many request to {operation} object(s) at {bucket}/{key}. {error_info}"
465 ) from error
466 elif status_code == 503:
467 raise RetryableError(
468 f"Service unavailable when {operation} object(s) at {bucket}/{key}. {error_info}"
469 ) from error
470 elif status_code == 501:
471 raise NotImplementedError(
472 f"Operation {operation} not implemented for object(s) at {bucket}/{key}. {error_info}"
473 ) from error
474 elif status_code == 408:
475 # 408 Request Timeout is from Google Cloud Storage
476 raise RetryableError(
477 f"Request timeout when {operation} object(s) at {bucket}/{key}. {error_info}"
478 ) from error
479 else:
480 raise RuntimeError(
481 f"Failed to {operation} object(s) at {bucket}/{key}. {error_info}, "
482 f"error_type: {type(error).__name__}"
483 ) from error
484 except RustClientError as error:
485 message = error.args[0]
486 status_code = error.args[1]
487 if status_code == 404:
488 raise FileNotFoundError(f"Object {bucket}/{key} does not exist. {message}") from error
489 elif status_code == 403:
490 raise PermissionError(
491 f"Permission denied to {operation} object(s) at {bucket}/{key}. {message}"
492 ) from error
493 else:
494 raise RetryableError(
495 f"Failed to {operation} object(s) at {bucket}/{key}. {message}. status_code: {status_code}"
496 ) from error
497 except (ReadTimeoutError, IncompleteReadError, ResponseStreamingError) as error:
498 raise RetryableError(
499 f"Failed to {operation} object(s) at {bucket}/{key} due to network timeout or incomplete read. "
500 f"error_type: {type(error).__name__}"
501 ) from error
502 except RustRetryableError as error:
503 raise RetryableError(
504 f"Failed to {operation} object(s) at {bucket}/{key} due to retryable error from Rust. "
505 f"error_type: {type(error).__name__}"
506 ) from error
507 except Exception as error:
508 if (
509 isinstance(error, S3UploadFailedError)
510 and isinstance(error.__context__, ClientError)
511 and error.__context__.response["ResponseMetadata"]["HTTPStatusCode"] == 403
512 ):
513 raise PermissionError(
514 f"Permission denied to {operation} object(s) at {bucket}/{key}. {error}"
515 ) from error
516 raise RuntimeError(
517 f"Failed to {operation} object(s) at {bucket}/{key}, error type: {type(error).__name__}, error: {error}"
518 ) from error
519
520 def _put_object(
521 self,
522 path: str,
523 body: bytes,
524 if_match: str | None = None,
525 if_none_match: str | None = None,
526 attributes: dict[str, str] | None = None,
527 content_type: str | None = None,
528 ) -> int:
529 """
530 Uploads an object to the specified S3 path.
531
532 :param path: The S3 path where the object will be uploaded.
533 :param body: The content of the object as bytes.
534 :param if_match: Optional If-Match header value. Use "*" to only upload if the object doesn't exist.
535 :param if_none_match: Optional If-None-Match header value. Use "*" to only upload if the object doesn't exist.
536 :param attributes: Optional attributes to attach to the object.
537 :param content_type: Optional Content-Type header value.
538 """
539 bucket, key = split_path(path)
540
541 def _invoke_api() -> int:
542 kwargs = {"Bucket": bucket, "Key": key, "Body": body}
543 if content_type:
544 kwargs["ContentType"] = content_type
545 if self._is_directory_bucket(bucket):
546 kwargs["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS
547 if if_match:
548 kwargs["IfMatch"] = if_match
549 if if_none_match:
550 kwargs["IfNoneMatch"] = if_none_match
551 validated_attributes = validate_attributes(attributes)
552 if validated_attributes:
553 kwargs["Metadata"] = validated_attributes
554
555 # TODO(NGCDP-5804): Add support to update ContentType header in Rust client
556 rust_unsupported_feature_keys = {"StorageClass", "IfMatch", "IfNoneMatch", "ContentType"}
557 if (
558 self._rust_client
559 # Rust client doesn't support creating objects with trailing /, see https://github.com/apache/arrow-rs/issues/7026
560 and not path.endswith("/")
561 and all(key not in kwargs for key in rust_unsupported_feature_keys)
562 ):
563 rust_attributes = {"attributes": validated_attributes} if validated_attributes else {}
564 run_async_rust_client_method(self._rust_client, "put", key, body, **rust_attributes)
565 else:
566 if self._checksum_algorithm:
567 kwargs["ChecksumAlgorithm"] = self._checksum_algorithm
568 self._s3_client.put_object(**kwargs)
569
570 return len(body)
571
572 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
573
574 def _supports_suffix_range(self) -> bool:
575 return self._rust_client is None
576
577 def _get_object(self, path: str, byte_range: Range | SuffixRange | None = None) -> bytes:
578 bucket, key = split_path(path)
579
580 if isinstance(byte_range, SuffixRange):
581 suffix_range = byte_range
582
583 def _invoke_suffix_api() -> bytes | None:
584 try:
585 response = self._s3_client.get_object(Bucket=bucket, Key=key, Range=f"bytes=-{suffix_range.size}")
586 except ClientError as error:
587 # Raised for empty objects, or by servers without suffix range support.
588 if error.response["ResponseMetadata"]["HTTPStatusCode"] == 416:
589 return None
590 raise
591 # Servers may ignore the range and return the whole object.
592 return response["Body"].read()[-suffix_range.size :]
593
594 return self._get_object_suffix(
595 path,
596 suffix_range,
597 lambda: self._translate_errors(_invoke_suffix_api, operation="GET", bucket=bucket, key=key),
598 )
599
600 def _invoke_api() -> bytes:
601 if byte_range:
602 bytes_range = f"bytes={byte_range.offset}-{byte_range.offset + byte_range.size - 1}"
603 if self._rust_client:
604 response = run_async_rust_client_method(
605 self._rust_client,
606 "get",
607 key,
608 byte_range,
609 )
610 return response
611 else:
612 response = self._s3_client.get_object(Bucket=bucket, Key=key, Range=bytes_range)
613 else:
614 if self._rust_client:
615 response = run_async_rust_client_method(self._rust_client, "get", key)
616 return response
617 else:
618 response = self._s3_client.get_object(Bucket=bucket, Key=key)
619
620 return response["Body"].read()
621
622 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
623
624 def _copy_object(self, src_path: str, dest_path: str) -> int:
625 src_bucket, src_key = split_path(src_path)
626 dest_bucket, dest_key = split_path(dest_path)
627
628 src_object = self._get_object_metadata(src_path)
629
630 def _invoke_api() -> int:
631 self._s3_client.copy(
632 CopySource={"Bucket": src_bucket, "Key": src_key},
633 Bucket=dest_bucket,
634 Key=dest_key,
635 Config=self._transfer_config,
636 )
637
638 return src_object.content_length
639
640 return self._translate_errors(_invoke_api, operation="COPY", bucket=dest_bucket, key=dest_key)
641
642 def _delete_object(self, path: str, if_match: str | None = None) -> None:
643 bucket, key = split_path(path)
644
645 def _invoke_api() -> None:
646 # Delete conditionally when if_match (etag) is provided; otherwise delete unconditionally
647 if if_match:
648 self._s3_client.delete_object(Bucket=bucket, Key=key, IfMatch=if_match)
649 else:
650 self._s3_client.delete_object(Bucket=bucket, Key=key)
651
652 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
653
654 def _delete_objects(self, paths: list[str]) -> None:
655 if not paths:
656 return
657
658 by_bucket: dict[str, list[str]] = {}
659 for p in paths:
660 bucket, key = split_path(p)
661 by_bucket.setdefault(bucket, []).append(key)
662
663 S3_BATCH_LIMIT = 1000
664
665 def _invoke_api() -> None:
666 all_errors: list[str] = []
667 for bucket, keys in by_bucket.items():
668 for i in range(0, len(keys), S3_BATCH_LIMIT):
669 chunk = keys[i : i + S3_BATCH_LIMIT]
670 response = self._s3_client.delete_objects(
671 Bucket=bucket, Delete={"Objects": [{"Key": k} for k in chunk]}
672 )
673 errors = response.get("Errors") or []
674 for e in errors:
675 all_errors.append(f"{bucket}/{e.get('Key', '?')}: {e.get('Code', '')} {e.get('Message', '')}")
676 if all_errors:
677 raise RuntimeError(f"DeleteObjects reported errors: {'; '.join(all_errors)}")
678
679 bucket_desc = "(" + "|".join(by_bucket) + ")"
680 key_desc = "(" + "|".join(str(len(keys)) for keys in by_bucket.values()) + " keys)"
681 self._translate_errors(_invoke_api, operation="DELETE_MANY", bucket=bucket_desc, key=key_desc)
682
683 def _is_dir(self, path: str) -> bool:
684 # Ensure the path ends with '/' to mimic a directory
685 path = self._append_delimiter(path)
686
687 bucket, key = split_path(path)
688
689 def _invoke_api() -> bool:
690 # List objects with the given prefix
691 response = self._s3_client.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1, Delimiter="/")
692
693 # Check if there are any contents or common prefixes
694 return bool(response.get("Contents", []) or response.get("CommonPrefixes", []))
695
696 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=key)
697
698 def _make_symlink(self, path: str, target: str) -> None:
699 bucket, key = split_path(path)
700 target_bucket, target_key = split_path(target)
701 if bucket != target_bucket:
702 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket}' -> '{target_bucket}'.")
703 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
704
705 def _invoke_api() -> None:
706 self._s3_client.put_object(
707 Bucket=bucket,
708 Key=key,
709 Body=b"",
710 Metadata={"msc-symlink-target": relative_target},
711 )
712
713 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
714
715 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
716 bucket, key = split_path(path)
717 if path.endswith("/") or (bucket and not key):
718 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
719 # which metadata is not guaranteed to exist for cases such as
720 # "virtual prefix" that was never explicitly created.
721 if self._is_dir(path):
722 return ObjectMetadata(
723 key=path,
724 type="directory",
725 content_length=0,
726 last_modified=AWARE_DATETIME_MIN,
727 )
728 else:
729 raise FileNotFoundError(f"Directory {path} does not exist.")
730 else:
731
732 def _invoke_api() -> ObjectMetadata:
733 response = self._s3_client.head_object(Bucket=bucket, Key=key)
734 user_metadata = response.get("Metadata")
735 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
736
737 return ObjectMetadata(
738 key=path,
739 type="file",
740 content_length=response["ContentLength"],
741 content_type=response.get("ContentType"),
742 last_modified=response.get("LastModified", AWARE_DATETIME_MIN),
743 etag=response["ETag"].strip('"') if "ETag" in response else None,
744 storage_class=response.get("StorageClass"),
745 metadata=user_metadata,
746 symlink_target=symlink_target,
747 )
748
749 try:
750 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
751 except FileNotFoundError:
752 if strict:
753 # If the object does not exist on the given path, we will append a trailing slash and
754 # check if the path is a directory.
755 path = self._append_delimiter(path)
756 if self._is_dir(path):
757 return ObjectMetadata(
758 key=path,
759 type="directory",
760 content_length=0,
761 last_modified=AWARE_DATETIME_MIN,
762 )
763 raise
764
765 def _list_objects(
766 self,
767 path: str,
768 start_after: str | None = None,
769 end_at: str | None = None,
770 include_directories: bool = False,
771 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
772 ) -> Iterator[ObjectMetadata]:
773 bucket, prefix = split_path(path)
774
775 # Get the prefix of the start_after and end_at paths relative to the bucket.
776 if start_after:
777 _, start_after = split_path(start_after)
778 if end_at:
779 _, end_at = split_path(end_at)
780
781 def _invoke_api() -> Iterator[ObjectMetadata]:
782 paginator = self._s3_client.get_paginator("list_objects_v2")
783 if include_directories:
784 page_iterator = paginator.paginate(
785 Bucket=bucket, Prefix=prefix, Delimiter="/", StartAfter=(start_after or "")
786 )
787 else:
788 page_iterator = paginator.paginate(Bucket=bucket, Prefix=prefix, StartAfter=(start_after or ""))
789
790 for page in page_iterator:
791 # A page holds both CommonPrefixes and Contents for the same key window, so a prefix past
792 # end_at must not stop the listing before the page's objects have been processed.
793 past_end_at = False
794 for item in page.get("CommonPrefixes", []):
795 prefix_key = item["Prefix"].rstrip("/")
796 # Filter by start_after and end_at - S3's StartAfter doesn't filter CommonPrefixes
797 if (start_after is None or start_after < prefix_key) and (end_at is None or prefix_key <= end_at):
798 yield ObjectMetadata(
799 key=os.path.join(bucket, prefix_key),
800 type="directory",
801 content_length=0,
802 last_modified=AWARE_DATETIME_MIN,
803 )
804 elif end_at is not None and end_at < prefix_key:
805 past_end_at = True
806 break
807
808 # S3 guarantees lexicographical order for general purpose buckets (for
809 # normal S3) but not directory buckets (for S3 Express One Zone).
810 for response_object in page.get("Contents", []):
811 key = response_object["Key"]
812 if end_at is None or key <= end_at:
813 if key.endswith("/"):
814 if include_directories:
815 yield ObjectMetadata(
816 key=os.path.join(bucket, key.rstrip("/")),
817 type="directory",
818 content_length=0,
819 last_modified=response_object["LastModified"],
820 )
821 else:
822 symlink_target = None
823 if response_object.get("Size", 0) == 0:
824 try:
825 meta = self._get_object_metadata(os.path.join(bucket, key))
826 symlink_target = meta.symlink_target
827 except Exception:
828 symlink_target = None
829 yield ObjectMetadata(
830 key=os.path.join(bucket, key),
831 type="file",
832 content_length=response_object["Size"],
833 last_modified=response_object["LastModified"],
834 etag=response_object["ETag"].strip('"'),
835 storage_class=response_object.get("StorageClass"),
836 symlink_target=symlink_target,
837 )
838 else:
839 return
840
841 if past_end_at:
842 return
843
844 return _translate_iterator_errors(
845 _invoke_api(), lambda f: self._translate_errors(f, operation="LIST", bucket=bucket, key=prefix)
846 )
847
848 @property
849 def supports_parallel_listing(self) -> bool:
850 """
851 S3 supports parallel listing via delimiter-based prefix discovery.
852
853 Note: Directory bucket handling is done in list_objects_recursive().
854 """
855 return True
856
[docs]
857 def list_objects_recursive(
858 self,
859 path: str = "",
860 start_after: str | None = None,
861 end_at: str | None = None,
862 max_workers: int = 32,
863 look_ahead: int = 2,
864 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
865 ) -> Iterator[ObjectMetadata]:
866 """
867 List all objects recursively using parallel prefix discovery for improved performance.
868
869 For S3, uses the Rust client's list_recursive when available for maximum performance.
870 Falls back to Python implementation otherwise.
871
872 Returns files only (no directories), in lexicographic order.
873
874 :param symlink_handling: How to handle symbolic links during listing.
875 """
876 if (start_after is not None) and (end_at is not None) and not (start_after < end_at):
877 raise ValueError(f"start_after ({start_after}) must be before end_at ({end_at})!")
878
879 full_path = self._prepend_base_path(path)
880 bucket, _ = split_path(full_path)
881
882 if self._is_directory_bucket(bucket):
883 yield from self.list_objects(
884 path, start_after, end_at, include_directories=False, symlink_handling=symlink_handling
885 )
886 return
887
888 if self._rust_client:
889 yield from self._emit_metrics(
890 operation=BaseStorageProvider._Operation.LIST,
891 f=lambda: self._list_objects_recursive_rust(path, full_path, bucket, start_after, end_at, max_workers),
892 )
893 else:
894 yield from super().list_objects_recursive(
895 path, start_after, end_at, max_workers, look_ahead, symlink_handling
896 )
897
898 def _list_objects_recursive_rust(
899 self,
900 path: str,
901 full_path: str,
902 bucket: str,
903 start_after: str | None,
904 end_at: str | None,
905 max_workers: int,
906 ) -> Iterator[ObjectMetadata]:
907 """
908 Use Rust client's list_recursive for parallel listing.
909
910 The Rust client already handles parallel listing internally.
911 Returns files only in lexicographic order.
912 """
913 _, prefix = split_path(full_path)
914
915 def _invoke_api() -> Iterator[ObjectMetadata]:
916 result = run_async_rust_client_method(
917 self._rust_client,
918 "list_recursive",
919 [prefix] if prefix else [""],
920 max_concurrency=max_workers,
921 )
922
923 start_after_full = self._prepend_base_path(start_after) if start_after else None
924 end_at_full = self._prepend_base_path(end_at) if end_at else None
925
926 for obj in result.objects:
927 full_key = os.path.join(bucket, obj.key)
928
929 if start_after_full and full_key <= start_after_full:
930 continue
931 if end_at_full and full_key > end_at_full:
932 break
933
934 relative_key = full_key.removeprefix(self._base_path).lstrip("/")
935
936 yield ObjectMetadata(
937 key=relative_key,
938 content_length=obj.content_length,
939 last_modified=dateutil_parse(obj.last_modified),
940 type="file" if obj.object_type == "object" else obj.object_type,
941 etag=obj.etag,
942 )
943
944 yield from _translate_iterator_errors(
945 _invoke_api(), lambda f: self._translate_errors(f, operation="LIST", bucket=bucket, key=prefix)
946 )
947
948 def _upload_file(
949 self,
950 remote_path: str,
951 f: str | IO,
952 attributes: dict[str, str] | None = None,
953 content_type: str | None = None,
954 ) -> int:
955 file_size: int = 0
956
957 if isinstance(f, str):
958 bucket, key = split_path(remote_path)
959 file_size = os.path.getsize(f)
960
961 # Upload small files
962 if file_size <= self._multipart_threshold:
963 if self._rust_client and not content_type and not self._is_directory_bucket(bucket):
964 validated_attributes = validate_attributes(attributes)
965 rust_attributes = {"attributes": validated_attributes} if validated_attributes else {}
966 self._translate_errors(
967 lambda: run_async_rust_client_method(self._rust_client, "upload", f, key, **rust_attributes),
968 operation="PUT",
969 bucket=bucket,
970 key=key,
971 )
972 else:
973 with open(f, "rb") as fp:
974 self._put_object(remote_path, fp.read(), attributes=attributes, content_type=content_type)
975 return file_size
976
977 # Upload large files using TransferConfig
978 def _invoke_api() -> int:
979 extra_args = {}
980 if content_type:
981 extra_args["ContentType"] = content_type
982 if self._is_directory_bucket(bucket):
983 extra_args["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS
984 validated_attributes = validate_attributes(attributes)
985 if validated_attributes:
986 extra_args["Metadata"] = validated_attributes
987 if self._rust_client and "ContentType" not in extra_args and "StorageClass" not in extra_args:
988 rust_attributes = {"attributes": validated_attributes} if validated_attributes else {}
989 run_async_rust_client_method(
990 self._rust_client, "upload_multipart_from_file", f, key, **rust_attributes
991 )
992 else:
993 if self._checksum_algorithm:
994 extra_args["ChecksumAlgorithm"] = self._checksum_algorithm
995 self._s3_client.upload_file(
996 Filename=f,
997 Bucket=bucket,
998 Key=key,
999 Config=self._transfer_config,
1000 ExtraArgs=extra_args,
1001 )
1002
1003 return file_size
1004
1005 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
1006 else:
1007 # Upload small files
1008 f.seek(0, io.SEEK_END)
1009 file_size = f.tell()
1010 f.seek(0)
1011
1012 if file_size <= self._multipart_threshold:
1013 if isinstance(f, io.StringIO):
1014 self._put_object(
1015 remote_path, f.read().encode("utf-8"), attributes=attributes, content_type=content_type
1016 )
1017 else:
1018 self._put_object(remote_path, f.read(), attributes=attributes, content_type=content_type)
1019 return file_size
1020
1021 # Upload large files using TransferConfig
1022 bucket, key = split_path(remote_path)
1023
1024 def _invoke_api() -> int:
1025 extra_args = {}
1026 if content_type:
1027 extra_args["ContentType"] = content_type
1028 if self._is_directory_bucket(bucket):
1029 extra_args["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS
1030 validated_attributes = validate_attributes(attributes)
1031 if validated_attributes:
1032 extra_args["Metadata"] = validated_attributes
1033
1034 if (
1035 self._rust_client
1036 and isinstance(f, io.BytesIO)
1037 and "ContentType" not in extra_args
1038 and "StorageClass" not in extra_args
1039 ):
1040 data = f.getbuffer()
1041 rust_attributes = {"attributes": validated_attributes} if validated_attributes else {}
1042 try:
1043 run_async_rust_client_method(
1044 self._rust_client, "upload_multipart_from_bytes", key, data, **rust_attributes
1045 )
1046 finally:
1047 data.release()
1048 else:
1049 if self._checksum_algorithm:
1050 extra_args["ChecksumAlgorithm"] = self._checksum_algorithm
1051 self._s3_client.upload_fileobj(
1052 Fileobj=f,
1053 Bucket=bucket,
1054 Key=key,
1055 Config=self._transfer_config,
1056 ExtraArgs=extra_args,
1057 )
1058
1059 return file_size
1060
1061 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
1062
1063 def _download_file(self, remote_path: str, f: str | IO, metadata: ObjectMetadata | None = None) -> int:
1064 """
1065 Downloads an object to a local path or file-like object.
1066
1067 Python downloads to paths stage data in a temporary file and replace the destination only after the transfer
1068 succeeds. Transfer or rename failures preserve the existing destination and trigger best-effort cleanup of the
1069 staging file. Cleanup failures are logged without masking the original error. Rust downloads use the Rust
1070 client's destination and resumable state handling.
1071
1072 :param remote_path: The object path to download.
1073 :param f: The destination file path or writable file-like object.
1074 :param metadata: Optional object metadata to avoid an additional lookup.
1075 :return: The downloaded object's content length.
1076 """
1077 if metadata is None:
1078 metadata = self._get_object_metadata(remote_path)
1079
1080 if isinstance(f, str):
1081 bucket, key = split_path(remote_path)
1082 if os.path.dirname(f):
1083 safe_makedirs(os.path.dirname(f))
1084
1085 # Download small files
1086 if metadata.content_length <= self._multipart_threshold:
1087 if self._rust_client:
1088 self._translate_errors(
1089 lambda: run_async_rust_client_method(self._rust_client, "download", key, f),
1090 operation="GET",
1091 bucket=bucket,
1092 key=key,
1093 )
1094 else:
1095 temp_file_path: str | None = None
1096 try:
1097 with tempfile.NamedTemporaryFile(
1098 mode="wb", delete=False, dir=os.path.dirname(f), prefix="."
1099 ) as fp:
1100 temp_file_path = fp.name
1101 fp.write(self._get_object(remote_path))
1102 os.rename(src=temp_file_path, dst=f)
1103 finally:
1104 if temp_file_path:
1105 try:
1106 os.unlink(temp_file_path)
1107 except FileNotFoundError:
1108 pass
1109 except OSError as error:
1110 logger.warning("Could not remove temporary file %s: %s", temp_file_path, error)
1111 return metadata.content_length
1112
1113 # Download large files using TransferConfig
1114 def _invoke_api() -> int:
1115 if self._rust_client:
1116 run_async_rust_client_method(self._rust_client, "download_multipart_to_file", key, f)
1117 else:
1118 temp_file_path: str | None = None
1119 try:
1120 with tempfile.NamedTemporaryFile(
1121 mode="wb", delete=False, dir=os.path.dirname(f), prefix="."
1122 ) as fp:
1123 temp_file_path = fp.name
1124 self._s3_client.download_fileobj(
1125 Bucket=bucket,
1126 Key=key,
1127 Fileobj=fp,
1128 Config=self._transfer_config,
1129 )
1130 os.rename(src=temp_file_path, dst=f)
1131 finally:
1132 if temp_file_path:
1133 try:
1134 os.unlink(temp_file_path)
1135 except FileNotFoundError:
1136 pass
1137 except OSError as error:
1138 logger.warning("Could not remove temporary file %s: %s", temp_file_path, error)
1139
1140 return metadata.content_length
1141
1142 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
1143 else:
1144 # Download small files
1145 if metadata.content_length <= self._multipart_threshold:
1146 response = self._get_object(remote_path)
1147 # Python client returns `bytes`, but Rust client returns an object that implements the buffer protocol,
1148 # so we need to check whether `.decode()` is available.
1149 if isinstance(f, io.StringIO):
1150 if hasattr(response, "decode"):
1151 f.write(response.decode("utf-8"))
1152 else:
1153 f.write(codecs.decode(memoryview(response), "utf-8"))
1154 else:
1155 f.write(response)
1156 return metadata.content_length
1157
1158 # Download large files using TransferConfig
1159 bucket, key = split_path(remote_path)
1160
1161 def _invoke_api() -> int:
1162 self._s3_client.download_fileobj(
1163 Bucket=bucket,
1164 Key=key,
1165 Fileobj=f,
1166 Config=self._transfer_config,
1167 )
1168
1169 return metadata.content_length
1170
1171 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
1172
1173 def _generate_presigned_url(
1174 self,
1175 path: str,
1176 *,
1177 method: str = "GET",
1178 signer_type: SignerType | None = None,
1179 signer_options: dict[str, Any] | None = None,
1180 ) -> str:
1181 options = signer_options or {}
1182 bucket, key = split_path(path)
1183
1184 if signer_type is None or signer_type == SignerType.S3:
1185 expires_in = int(options.get("expires_in", DEFAULT_PRESIGN_EXPIRES_IN))
1186 cache_key: tuple = (SignerType.S3, bucket, expires_in)
1187 if cache_key not in self._signer_cache:
1188 self._signer_cache[cache_key] = S3URLSigner(self._s3_client, bucket, expires_in=expires_in)
1189 elif signer_type == SignerType.CLOUDFRONT:
1190 cache_key = (SignerType.CLOUDFRONT, frozenset(options.items()))
1191 if cache_key not in self._signer_cache:
1192 self._signer_cache[cache_key] = CloudFrontURLSigner(**options)
1193 else:
1194 raise ValueError(f"Unsupported signer type for S3 provider: {signer_type!r}")
1195
1196 return self._signer_cache[cache_key].generate_presigned_url(key, method=method)