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