Coverage for python/lsst/resources/s3.py: 21%

339 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-19 01:29 -0800

1# This file is part of lsst-resources. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ("S3ResourcePath",) 

15 

16import concurrent.futures 

17import contextlib 

18import io 

19import logging 

20import os 

21import re 

22import sys 

23import threading 

24from collections import defaultdict 

25from collections.abc import Iterable, Iterator 

26from functools import cache, cached_property 

27from typing import IO, TYPE_CHECKING, cast 

28 

29from botocore.exceptions import ClientError 

30 

31from lsst.utils.iteration import chunk_iterable 

32from lsst.utils.timer import time_this 

33 

34from ._resourceHandles._baseResourceHandle import ResourceHandleProtocol 

35from ._resourceHandles._s3ResourceHandle import S3ResourceHandle 

36from ._resourcePath import _EXECUTOR_TYPE, MBulkResult, ResourcePath, _get_executor_class, _patch_environ 

37from .s3utils import ( 

38 _get_s3_connection_parameters, 

39 _s3_disable_bucket_validation, 

40 _s3_should_validate_bucket, 

41 all_retryable_errors, 

42 backoff, 

43 bucketExists, 

44 getS3Client, 

45 max_retry_time, 

46 retryable_io_errors, 

47 s3CheckFileExists, 

48 translate_client_error, 

49) 

50from .utils import _get_num_workers 

51 

52try: 

53 from boto3.s3.transfer import TransferConfig # type: ignore 

54except ImportError: 

55 TransferConfig = None 

56 

57try: 

58 import s3fs 

59 from fsspec.spec import AbstractFileSystem 

60except ImportError: 

61 s3fs = None 

62 AbstractFileSystem = type 

63 

64if TYPE_CHECKING: 

65 with contextlib.suppress(ImportError): 

66 import boto3 

67 

68 from .utils import TransactionProtocol 

69 

70 

71log = logging.getLogger(__name__) 

72 

73 

74class ProgressPercentage: 

75 """Progress bar for S3 file uploads. 

76 

77 Parameters 

78 ---------- 

79 file : `ResourcePath` 

80 Resource that is relevant to the progress percentage. The size of this 

81 resource will be used to determine progress. The name will be used 

82 in the log messages unless overridden by ``file_for_msg``. 

83 file_for_msg : `ResourcePath` or `None`, optional 

84 Resource name to include in log messages in preference to ``file``. 

85 msg : `str`, optional 

86 Message text to be included in every progress log message. 

87 """ 

88 

89 log_level = logging.DEBUG 

90 """Default log level to use when issuing a message.""" 

91 

92 def __init__(self, file: ResourcePath, file_for_msg: ResourcePath | None = None, msg: str = ""): 

93 self._filename = file 

94 self._file_for_msg = str(file_for_msg) if file_for_msg is not None else str(file) 

95 self._size = file.size() 

96 self._seen_so_far = 0 

97 self._lock = threading.Lock() 

98 self._msg = msg 

99 

100 def __call__(self, bytes_amount: int) -> None: 

101 # To simplify, assume this is hooked up to a single filename 

102 with self._lock: 

103 self._seen_so_far += bytes_amount 

104 percentage = (100 * self._seen_so_far) // self._size 

105 log.log( 

106 self.log_level, 

107 "%s %s %s / %s (%s%%)", 

108 self._msg, 

109 self._file_for_msg, 

110 self._seen_so_far, 

111 self._size, 

112 percentage, 

113 ) 

114 

115 

116@cache 

117def _parse_string_to_maybe_bool(maybe_bool_str: str) -> bool | None: 

118 """Map a string to either a boolean value or None. 

119 

120 Parameters 

121 ---------- 

122 maybe_bool_str : `str` 

123 The value to parse 

124 

125 Results 

126 ------- 

127 maybe_bool : `bool` or `None` 

128 The parsed value. 

129 """ 

130 if maybe_bool_str.lower() in ["t", "true", "yes", "y", "1"]: 

131 maybe_bool = True 

132 elif maybe_bool_str.lower() in ["f", "false", "no", "n", "0"]: 

133 maybe_bool = False 

134 elif maybe_bool_str.lower() in ["none", ""]: 

135 maybe_bool = None 

136 else: 

137 raise ValueError(f'Value of "{maybe_bool_str}" is not True, False, or None.') 

138 

139 return maybe_bool 

140 

141 

142class S3ResourcePath(ResourcePath): 

143 """S3 URI resource path implementation class. 

144 

145 Notes 

146 ----- 

147 In order to configure the behavior of instances of this class, the 

148 environment variable is inspected: 

149 

150 - LSST_S3_USE_THREADS: May be True, False, or None. Sets whether threading 

151 is used for downloads, with a value of None defaulting to boto's default 

152 value. Users may wish to set it to False when the downloads will be started 

153 within threads other than python's main thread. 

154 """ 

155 

156 use_threads: bool | None = None 

157 """Explicitly turn on or off threading in use of boto's download_fileobj. 

158 Setting this to None results in boto's default behavior.""" 

159 

160 @cached_property 

161 def _environ_use_threads(self) -> bool | None: 

162 try: 

163 use_threads_str = os.environ["LSST_S3_USE_THREADS"] 

164 except KeyError: 

165 use_threads_str = "None" 

166 

167 use_threads = _parse_string_to_maybe_bool(use_threads_str) 

168 

169 return use_threads 

170 

171 @contextlib.contextmanager 

172 def _use_threads_temp_override(self, multithreaded: bool) -> Iterator: 

173 """Temporarily override the value of use_threads.""" 

174 original = self.use_threads 

175 self.use_threads = multithreaded 

176 yield 

177 self.use_threads = original 

178 

179 @property 

180 def _transfer_config(self) -> TransferConfig: 

181 if self.use_threads is None: 

182 self.use_threads = self._environ_use_threads 

183 

184 if self.use_threads is None: 

185 transfer_config = TransferConfig() 

186 else: 

187 transfer_config = TransferConfig(use_threads=self.use_threads) 

188 

189 return transfer_config 

190 

191 @property 

192 def client(self) -> boto3.client: 

193 """Client object to address remote resource.""" 

194 return getS3Client(self._profile) 

195 

196 @property 

197 def _profile(self) -> str | None: 

198 """Profile name to use for looking up S3 credentials and endpoint.""" 

199 return self._uri.username 

200 

201 @property 

202 def _bucket(self) -> str: 

203 """S3 bucket where the files are stored.""" 

204 # Notionally the bucket is stored in the 'hostname' part of the URI. 

205 # However, Ceph S3 uses a "multi-tenant" syntax for bucket names in the 

206 # form 'tenant:bucket'. The part after the colon is parsed as the port 

207 # portion of the URI, and urllib throws an exception if you try to read 

208 # a non-integer port value. So manually split off this portion of the 

209 # URI. 

210 split = self._uri.netloc.split("@") 

211 num_components = len(split) 

212 if num_components == 2: 

213 # There is a profile@ portion of the URL, so take the second half. 

214 bucket = split[1] 

215 elif num_components == 1: 

216 # There is no profile@, so take the whole netloc. 

217 bucket = split[0] 

218 else: 

219 raise ValueError(f"Unexpected extra '@' in S3 URI: '{str(self)}'") 

220 

221 if not bucket: 

222 raise ValueError(f"S3 URI does not include bucket name: '{str(self)}'") 

223 

224 return bucket 

225 

226 @classmethod 

227 def _mexists( 

228 cls, uris: Iterable[ResourcePath], *, num_workers: int | None = None 

229 ) -> dict[ResourcePath, bool]: 

230 # Force client to be created for each profile before creating threads. 

231 profiles = set[str | None]() 

232 for path in uris: 

233 if path.scheme == "s3": 

234 path = cast(S3ResourcePath, path) 

235 profiles.add(path._profile) 

236 for profile in profiles: 

237 getS3Client(profile) 

238 

239 return super()._mexists(uris, num_workers=num_workers) 

240 

241 @classmethod 

242 def _mremove(cls, uris: Iterable[ResourcePath]) -> dict[ResourcePath, MBulkResult]: 

243 # Delete multiple objects in one API call. 

244 # Must group by profile and bucket. 

245 grouped_uris: dict[tuple[str | None, str], list[S3ResourcePath]] = defaultdict(list) 

246 for uri in uris: 

247 uri = cast(S3ResourcePath, uri) 

248 grouped_uris[uri._profile, uri._bucket].append(uri) 

249 

250 results: dict[ResourcePath, MBulkResult] = {} 

251 for related_uris in grouped_uris.values(): 

252 # API requires no more than 1000 per call. 

253 chunk_num = 0 

254 chunks: list[tuple[ResourcePath, ...]] = [] 

255 key_to_uri: dict[str, ResourcePath] = {} 

256 for chunk in chunk_iterable(related_uris, chunk_size=1_000): 

257 for uri in chunk: 

258 key = uri.relativeToPathRoot 

259 key_to_uri[key] = uri 

260 # Default to assuming everything worked. 

261 results[uri] = MBulkResult(True, None) 

262 chunk_num += 1 

263 chunks.append(chunk) 

264 

265 # Bulk remove. 

266 with time_this( 

267 log, 

268 msg="Bulk delete; %d chunk%s; totalling %d dataset%s", 

269 args=( 

270 len(chunks), 

271 "s" if len(chunks) != 1 else "", 

272 len(related_uris), 

273 "s" if len(related_uris) != 1 else "", 

274 ), 

275 ): 

276 errored = cls._mremove_select(chunks) 

277 

278 # Update with error information. 

279 results.update(errored) 

280 

281 return results 

282 

283 @classmethod 

284 def _mremove_select(cls, chunks: list[tuple[ResourcePath, ...]]) -> dict[ResourcePath, MBulkResult]: 

285 if len(chunks) == 1: 

286 # Do the removal directly without futures. 

287 return cls._delete_objects_wrapper(chunks[0]) 

288 pool_executor_class = _get_executor_class() 

289 if issubclass(pool_executor_class, concurrent.futures.ProcessPoolExecutor): 

290 # Patch the environment to make it think there is only one worker 

291 # for each subprocess. 

292 with _patch_environ({"LSST_RESOURCES_NUM_WORKERS": "1"}): 

293 return cls._mremove_with_pool(pool_executor_class, chunks) 

294 else: 

295 return cls._mremove_with_pool(pool_executor_class, chunks) 

296 

297 @classmethod 

298 def _mremove_with_pool( 

299 cls, 

300 pool_executor_class: _EXECUTOR_TYPE, 

301 chunks: list[tuple[ResourcePath, ...]], 

302 *, 

303 num_workers: int | None = None, 

304 ) -> dict[ResourcePath, MBulkResult]: 

305 # Different name because different API to base class. 

306 # No need to make more workers than we have chunks. 

307 max_workers = num_workers if num_workers is not None else min(len(chunks), _get_num_workers()) 

308 results: dict[ResourcePath, MBulkResult] = {} 

309 with pool_executor_class(max_workers=max_workers) as remove_executor: 

310 future_remove = { 

311 remove_executor.submit(cls._delete_objects_wrapper, chunk): i 

312 for i, chunk in enumerate(chunks) 

313 } 

314 for future in concurrent.futures.as_completed(future_remove): 

315 try: 

316 results.update(future.result()) 

317 except Exception as e: 

318 # The chunk utterly failed. 

319 chunk = chunks[future_remove[future]] 

320 for uri in chunk: 

321 results[uri] = MBulkResult(False, e) 

322 return results 

323 

324 @classmethod 

325 def _delete_objects_wrapper(cls, uris: tuple[ResourcePath, ...]) -> dict[ResourcePath, MBulkResult]: 

326 """Convert URIs to keys and call low-level API.""" 

327 if not uris: 

328 return {} 

329 keys: list[dict[str, str]] = [] 

330 key_to_uri: dict[str, ResourcePath] = {} 

331 for uri in uris: 

332 key = uri.relativeToPathRoot 

333 key_to_uri[key] = uri 

334 keys.append({"Key": key}) 

335 

336 first_uri = cast(S3ResourcePath, uris[0]) 

337 results = cls._delete_related_objects(first_uri.client, first_uri._bucket, keys) 

338 

339 # Remap error object keys to uris. 

340 return {key_to_uri[key]: result for key, result in results.items()} 

341 

342 @classmethod 

343 @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time) 

344 def _delete_related_objects( 

345 cls, client: boto3.client, bucket: str, keys: list[dict[str, str]] 

346 ) -> dict[str, MBulkResult]: 

347 # Delete multiple objects from the same bucket, allowing for backoff 

348 # retry. 

349 response = client.delete_objects(Bucket=bucket, Delete={"Objects": keys, "Quiet": True}) 

350 # Use Quiet mode so we assume everything worked unless told otherwise. 

351 # Only returning errors -- indexed by Key name. 

352 errors: dict[str, MBulkResult] = {} 

353 for errored_key in response.get("Errors", []): 

354 errors[errored_key["Key"]] = MBulkResult( 

355 False, ClientError({"Error": errored_key}, f"delete_objects: {errored_key['Key']}") 

356 ) 

357 return errors 

358 

359 @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time) 

360 def exists(self) -> bool: 

361 """Check that the S3 resource exists.""" 

362 if self.is_root: 

363 # Only check for the bucket since the path is irrelevant 

364 return bucketExists(self._bucket, self.client) 

365 exists, _ = s3CheckFileExists(self, bucket=self._bucket, client=self.client) 

366 return exists 

367 

368 @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time) 

369 def size(self) -> int: 

370 """Return the size of the resource in bytes.""" 

371 if self.dirLike: 

372 return 0 

373 exists, sz = s3CheckFileExists(self, bucket=self._bucket, client=self.client) 

374 if not exists: 

375 raise FileNotFoundError(f"Resource {self} does not exist") 

376 return sz 

377 

378 @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time) 

379 def remove(self) -> None: 

380 """Remove the resource.""" 

381 # https://github.com/boto/boto3/issues/507 - there is no 

382 # way of knowing if the file was actually deleted except 

383 # for checking all the keys again, reponse is HTTP 204 OK 

384 # response all the time 

385 try: 

386 self.client.delete_object(Bucket=self._bucket, Key=self.relativeToPathRoot) 

387 except (self.client.exceptions.NoSuchKey, self.client.exceptions.NoSuchBucket) as err: 

388 raise FileNotFoundError("No such resource: {self}") from err 

389 

390 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

391 def read(self, size: int = -1) -> bytes: 

392 args = {} 

393 if size > 0: 

394 args["Range"] = f"bytes=0-{size - 1}" 

395 try: 

396 response = self.client.get_object(Bucket=self._bucket, Key=self.relativeToPathRoot, **args) 

397 except (self.client.exceptions.NoSuchKey, self.client.exceptions.NoSuchBucket) as err: 

398 raise FileNotFoundError(f"No such resource: {self}") from err 

399 except ClientError as err: 

400 translate_client_error(err, self) 

401 raise 

402 with time_this(log, msg="Read from %s", args=(self,)): 

403 body = response["Body"].read() 

404 response["Body"].close() 

405 return body 

406 

407 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

408 def write(self, data: bytes, overwrite: bool = True) -> None: 

409 if not overwrite and self.exists(): 

410 raise FileExistsError(f"Remote resource {self} exists and overwrite has been disabled") 

411 with time_this(log, msg="Write to %s", args=(self,)): 

412 self.client.put_object(Bucket=self._bucket, Key=self.relativeToPathRoot, Body=data) 

413 

414 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

415 def mkdir(self) -> None: 

416 """Write a directory key to S3.""" 

417 if not bucketExists(self._bucket, self.client): 

418 raise ValueError(f"Bucket {self._bucket} does not exist for {self}!") 

419 

420 if not self.dirLike: 

421 raise NotADirectoryError(f"Can not create a 'directory' for file-like URI {self}") 

422 

423 # don't create S3 key when root is at the top-level of an Bucket 

424 if self.path != "/": 

425 self.client.put_object(Bucket=self._bucket, Key=self.relativeToPathRoot) 

426 

427 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

428 def _download_file( 

429 self, local_file: IO | ResourceHandleProtocol, progress: ProgressPercentage | None 

430 ) -> None: 

431 """Download the remote resource to a local file. 

432 

433 Helper routine for _as_local to allow backoff without regenerating 

434 the temporary file. 

435 """ 

436 try: 

437 self.client.download_fileobj( 

438 self._bucket, 

439 self.relativeToPathRoot, 

440 local_file, 

441 Callback=progress, 

442 Config=self._transfer_config, 

443 ) 

444 except ( 

445 self.client.exceptions.NoSuchKey, 

446 self.client.exceptions.NoSuchBucket, 

447 ) as err: 

448 raise FileNotFoundError(f"No such resource: {self}") from err 

449 except ClientError as err: 

450 translate_client_error(err, self) 

451 raise 

452 

453 def to_fsspec(self) -> tuple[AbstractFileSystem, str]: 

454 """Return an abstract file system and path that can be used by fsspec. 

455 

456 Returns 

457 ------- 

458 fs : `fsspec.spec.AbstractFileSystem` 

459 A file system object suitable for use with the returned path. 

460 path : `str` 

461 A path that can be opened by the file system object. 

462 """ 

463 if s3fs is None: 

464 raise ImportError("s3fs is not available") 

465 # Must remove the profile from the URL and form it again. 

466 endpoint_config = _get_s3_connection_parameters(self._profile) 

467 s3 = s3fs.S3FileSystem( 

468 profile=endpoint_config.profile, 

469 endpoint_url=endpoint_config.endpoint_url, 

470 key=endpoint_config.access_key_id, 

471 secret=endpoint_config.secret_access_key, 

472 ) 

473 if not _s3_should_validate_bucket(): 

474 # Accessing the s3 property forces the boto client to be 

475 # constructed and cached and allows the validation to be removed. 

476 _s3_disable_bucket_validation(s3.s3) 

477 

478 return s3, f"{self._bucket}/{self.relativeToPathRoot}" 

479 

480 @contextlib.contextmanager 

481 def _as_local( 

482 self, multithreaded: bool = True, tmpdir: ResourcePath | None = None 

483 ) -> Iterator[ResourcePath]: 

484 """Download object from S3 and place in temporary directory. 

485 

486 Parameters 

487 ---------- 

488 multithreaded : `bool`, optional 

489 If `True` the transfer will be allowed to attempt to improve 

490 throughput by using parallel download streams. This may of no 

491 effect if the URI scheme does not support parallel streams or 

492 if a global override has been applied. If `False` parallel 

493 streams will be disabled. 

494 tmpdir : `ResourcePath` or `None`, optional 

495 Explicit override of the temporary directory to use for remote 

496 downloads. 

497 

498 Returns 

499 ------- 

500 local_uri : `ResourcePath` 

501 A URI to a local POSIX file corresponding to a local temporary 

502 downloaded copy of the resource. 

503 """ 

504 with ( 

505 ResourcePath.temporary_uri(prefix=tmpdir, suffix=self.getExtension(), delete=True) as tmp_uri, 

506 self._use_threads_temp_override(multithreaded), 

507 time_this(log, msg="Downloading %s to local file", args=(self,)), 

508 ): 

509 progress = ( 

510 ProgressPercentage(self, msg="Downloading:") 

511 if log.isEnabledFor(ProgressPercentage.log_level) 

512 else None 

513 ) 

514 with tmp_uri.open("wb") as tmpFile: 

515 self._download_file(tmpFile, progress) 

516 yield tmp_uri 

517 

518 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

519 def _upload_file(self, local_file: ResourcePath, progress: ProgressPercentage | None) -> None: 

520 """Upload a local file with backoff. 

521 

522 Helper method to wrap file uploading in backoff for transfer_from. 

523 """ 

524 try: 

525 self.client.upload_file( 

526 local_file.ospath, 

527 self._bucket, 

528 self.relativeToPathRoot, 

529 Callback=progress, 

530 Config=self._transfer_config, 

531 ) 

532 except self.client.exceptions.NoSuchBucket as err: 

533 raise NotADirectoryError(f"Target does not exist: {err}") from err 

534 except ClientError as err: 

535 translate_client_error(err, self) 

536 raise 

537 

538 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

539 def _copy_from(self, src: S3ResourcePath) -> None: 

540 copy_source = { 

541 "Bucket": src._bucket, 

542 "Key": src.relativeToPathRoot, 

543 } 

544 try: 

545 self.client.copy_object(CopySource=copy_source, Bucket=self._bucket, Key=self.relativeToPathRoot) 

546 except (self.client.exceptions.NoSuchKey, self.client.exceptions.NoSuchBucket) as err: 

547 raise FileNotFoundError(f"No such resource to transfer: {src} -> {self}") from err 

548 except ClientError as err: 

549 translate_client_error(err, self) 

550 raise 

551 

552 def transfer_from( 

553 self, 

554 src: ResourcePath, 

555 transfer: str = "copy", 

556 overwrite: bool = False, 

557 transaction: TransactionProtocol | None = None, 

558 multithreaded: bool = True, 

559 ) -> None: 

560 """Transfer the current resource to an S3 bucket. 

561 

562 Parameters 

563 ---------- 

564 src : `ResourcePath` 

565 Source URI. 

566 transfer : `str` 

567 Mode to use for transferring the resource. Supports the following 

568 options: copy. 

569 overwrite : `bool`, optional 

570 Allow an existing file to be overwritten. Defaults to `False`. 

571 transaction : `~lsst.resources.utils.TransactionProtocol`, optional 

572 Currently unused. 

573 multithreaded : `bool`, optional 

574 If `True` the transfer will be allowed to attempt to improve 

575 throughput by using parallel download streams. This may of no 

576 effect if the URI scheme does not support parallel streams or 

577 if a global override has been applied. If `False` parallel 

578 streams will be disabled. 

579 """ 

580 # Fail early to prevent delays if remote resources are requested 

581 if transfer not in self.transferModes: 

582 raise ValueError(f"Transfer mode '{transfer}' not supported by URI scheme {self.scheme}") 

583 

584 # Existence checks cost time so do not call this unless we know 

585 # that debugging is enabled. 

586 if log.isEnabledFor(logging.DEBUG): 

587 log.debug( 

588 "Transferring %s [exists: %s] -> %s [exists: %s] (transfer=%s)", 

589 src, 

590 src.exists(), 

591 self, 

592 self.exists(), 

593 transfer, 

594 ) 

595 

596 # Short circuit if the URIs are identical immediately. 

597 if self == src: 

598 log.debug( 

599 "Target and destination URIs are identical: %s, returning immediately." 

600 " No further action required.", 

601 self, 

602 ) 

603 return 

604 

605 if not overwrite and self.exists(): 

606 raise FileExistsError(f"Destination path '{self}' already exists.") 

607 

608 if transfer == "auto": 

609 transfer = self.transferDefault 

610 

611 timer_msg = "Transfer from %s to %s" 

612 timer_args = (src, self) 

613 

614 if isinstance(src, type(self)) and self.client == src.client: 

615 # Looks like an S3 remote uri so we can use direct copy. 

616 # This only works if the source and destination are using the same 

617 # S3 endpoint and profile. 

618 with time_this(log, msg=timer_msg, args=timer_args): 

619 self._copy_from(src) 

620 

621 else: 

622 # Use local file and upload it 

623 with src.as_local(multithreaded=multithreaded) as local_uri: 

624 progress = ( 

625 ProgressPercentage(local_uri, file_for_msg=src, msg="Uploading:") 

626 if log.isEnabledFor(ProgressPercentage.log_level) 

627 else None 

628 ) 

629 with ( 

630 time_this(log, msg=timer_msg, args=timer_args), 

631 self._use_threads_temp_override(multithreaded), 

632 ): 

633 self._upload_file(local_uri, progress) 

634 

635 # This was an explicit move requested from a remote resource 

636 # try to remove that resource 

637 if transfer == "move": 

638 # Transactions do not work here 

639 src.remove() 

640 

641 @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time) 

642 def walk( 

643 self, file_filter: str | re.Pattern | None = None 

644 ) -> Iterator[list | tuple[ResourcePath, list[str], list[str]]]: 

645 """Walk the directory tree returning matching files and directories. 

646 

647 Parameters 

648 ---------- 

649 file_filter : `str` or `re.Pattern`, optional 

650 Regex to filter out files from the list before it is returned. 

651 

652 Yields 

653 ------ 

654 dirpath : `ResourcePath` 

655 Current directory being examined. 

656 dirnames : `list` of `str` 

657 Names of subdirectories within dirpath. 

658 filenames : `list` of `str` 

659 Names of all the files within dirpath. 

660 """ 

661 # We pretend that S3 uses directories and files and not simply keys 

662 if not (self.isdir() or self.is_root): 

663 raise ValueError(f"Can not walk a non-directory URI: {self}") 

664 

665 if isinstance(file_filter, str): 

666 file_filter = re.compile(file_filter) 

667 

668 s3_paginator = self.client.get_paginator("list_objects_v2") 

669 

670 # Limit each query to a single "directory" to match os.walk 

671 # We could download all keys at once with no delimiter and work 

672 # it out locally but this could potentially lead to large memory 

673 # usage for millions of keys. It will also make the initial call 

674 # to this method potentially very slow. If making this method look 

675 # like os.walk was not required, we could query all keys with 

676 # pagination and return them in groups of 1000, but that would 

677 # be a different interface since we can't guarantee we would get 

678 # them all grouped properly across the 1000 limit boundary. 

679 prefix = self.relativeToPathRoot if not self.is_root else "" 

680 prefix_len = len(prefix) 

681 dirnames = [] 

682 filenames = [] 

683 files_there = False 

684 

685 for page in s3_paginator.paginate(Bucket=self._bucket, Prefix=prefix, Delimiter="/"): 

686 # All results are returned as full key names and we must 

687 # convert them back to the root form. The prefix is fixed 

688 # and delimited so that is a simple trim 

689 

690 # Directories are reported in the CommonPrefixes result 

691 # which reports the entire key and must be stripped. 

692 found_dirs = [dir["Prefix"][prefix_len:] for dir in page.get("CommonPrefixes", ())] 

693 dirnames.extend(found_dirs) 

694 

695 found_files = [file["Key"][prefix_len:] for file in page.get("Contents", ())] 

696 if found_files: 

697 files_there = True 

698 if file_filter is not None: 

699 found_files = [f for f in found_files if file_filter.search(f)] 

700 

701 filenames.extend(found_files) 

702 

703 # Directories do not exist so we can't test for them. If no files 

704 # or directories were found though, this means that it effectively 

705 # does not exist and we should match os.walk() behavior and return 

706 # immediately. 

707 if not dirnames and not files_there: 

708 return 

709 else: 

710 yield self, dirnames, filenames 

711 

712 for dir in dirnames: 

713 new_uri = self.join(dir) 

714 yield from new_uri.walk(file_filter) 

715 

716 @contextlib.contextmanager 

717 def _openImpl( 

718 self, 

719 mode: str = "r", 

720 *, 

721 encoding: str | None = None, 

722 ) -> Iterator[ResourceHandleProtocol]: 

723 with S3ResourceHandle(mode, log, self) as handle: 

724 if "b" in mode: 

725 yield handle 

726 else: 

727 if encoding is None: 

728 encoding = sys.getdefaultencoding() 

729 # cast because the protocol is compatible, but does not have 

730 # BytesIO in the inheritance tree 

731 with io.TextIOWrapper(cast(io.BytesIO, handle), encoding=encoding, write_through=True) as sub: 

732 yield sub 

733 

734 def generate_presigned_get_url(self, *, expiration_time_seconds: int) -> str: 

735 # Docstring inherited 

736 return self._generate_presigned_url("get_object", expiration_time_seconds) 

737 

738 def generate_presigned_put_url(self, *, expiration_time_seconds: int) -> str: 

739 # Docstring inherited 

740 return self._generate_presigned_url("put_object", expiration_time_seconds) 

741 

742 def _generate_presigned_url(self, method: str, expiration_time_seconds: int) -> str: 

743 url = self.client.generate_presigned_url( 

744 method, 

745 Params={"Bucket": self._bucket, "Key": self.relativeToPathRoot}, 

746 ExpiresIn=expiration_time_seconds, 

747 ) 

748 if self.fragment: 

749 resource = ResourcePath(url) 

750 url = str(resource.replace(fragment=self.fragment)) 

751 return url