Coverage for python/lsst/resources/s3utils.py: 79%

150 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-23 02:10 -0700

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__ = ( 

15 "_TooManyRequestsError", 

16 "all_retryable_errors", 

17 "backoff", 

18 "bucketExists", 

19 "clean_test_environment_for_s3", 

20 "getS3Client", 

21 "max_retry_time", 

22 "retryable_client_errors", 

23 "retryable_io_errors", 

24 "s3CheckFileExists", 

25) 

26 

27import functools 

28import os 

29import re 

30import urllib.parse 

31from collections.abc import Callable, Generator 

32from contextlib import contextmanager 

33from http.client import HTTPException, ImproperConnectionState 

34from types import ModuleType 

35from typing import TYPE_CHECKING, Any, NamedTuple, cast 

36from unittest.mock import patch 

37 

38from botocore.client import BaseClient 

39from botocore.exceptions import ClientError 

40from botocore.handlers import validate_bucket_name 

41from urllib3.exceptions import HTTPError, RequestError 

42from urllib3.util import Url, parse_url 

43 

44try: 

45 import boto3 

46except ImportError: 

47 # Hidden from type checkers so that ``boto3`` keeps the type it has when 

48 # the optional dependency is installed. 

49 if not TYPE_CHECKING: 

50 boto3 = None 

51 

52try: 

53 import botocore.config 

54except ImportError: 

55 if not TYPE_CHECKING: 

56 botocore = None 

57 

58 

59from ._resourcePath import ResourcePath 

60from .location import Location 

61from .utils import _get_num_workers 

62 

63# https://pypi.org/project/backoff/ 

64try: 

65 import backoff 

66except ImportError: 

67 

68 class Backoff: 

69 """Mock implementation of the backoff class.""" 

70 

71 @staticmethod 

72 def expo(func: Callable, *args: Any, **kwargs: Any) -> Callable: 

73 return func 

74 

75 @staticmethod 

76 def on_exception(func: Callable, *args: Any, **kwargs: Any) -> Callable: 

77 return func 

78 

79 backoff = cast(ModuleType, Backoff) 

80 

81 

82class _TooManyRequestsError(Exception): 

83 """Private exception that can be used for 429 retry. 

84 

85 botocore refuses to deal with 429 error itself so issues a generic 

86 ClientError. 

87 """ 

88 

89 pass 

90 

91 

92# settings for "backoff" retry decorators. these retries are belt-and- 

93# suspenders along with the retries built into Boto3, to account for 

94# semantic differences in errors between S3-like providers. 

95retryable_io_errors = ( 

96 # http.client 

97 ImproperConnectionState, 

98 HTTPException, 

99 # urllib3.exceptions 

100 RequestError, 

101 HTTPError, 

102 # built-ins 

103 TimeoutError, 

104 ConnectionError, 

105 # private 

106 _TooManyRequestsError, 

107) 

108 

109# Client error can include NoSuchKey so retry may not be the right 

110# thing. This may require more consideration if it is to be used. 

111retryable_client_errors = ( 

112 # botocore.exceptions 

113 ClientError, 

114 # built-ins 

115 PermissionError, 

116) 

117 

118 

119# Combine all errors into an easy package. For now client errors 

120# are not included. 

121all_retryable_errors = retryable_io_errors 

122max_retry_time = 60 

123 

124 

125@contextmanager 

126def clean_test_environment_for_s3() -> Generator[None]: 

127 """Reset S3 environment to ensure that unit tests with a mock S3 can't 

128 accidentally reference real infrastructure, and that site configuration 

129 cannot change how the client behaves. 

130 """ 

131 with patch.dict( 

132 os.environ, 

133 { 

134 "AWS_ACCESS_KEY_ID": "test-access-key", 

135 "AWS_SECRET_ACCESS_KEY": "test-secret-access-key", 

136 "AWS_DEFAULT_REGION": "us-east-1", 

137 }, 

138 ) as patched_environ: 

139 for var in ( 

140 "S3_ENDPOINT_URL", 

141 "AWS_SECURITY_TOKEN", 

142 "AWS_SESSION_TOKEN", 

143 "AWS_PROFILE", 

144 "AWS_SHARED_CREDENTIALS_FILE", 

145 "AWS_CONFIG_FILE", 

146 # A site that turns checksums off changes what an object's 

147 # metadata contains, so let the library defaults apply. 

148 "AWS_REQUEST_CHECKSUM_CALCULATION", 

149 "AWS_RESPONSE_CHECKSUM_VALIDATION", 

150 ): 

151 patched_environ.pop(var, None) 

152 # Clear the cached boto3 S3 client instances. 

153 # This helps us avoid a potential situation where the client could be 

154 # instantiated before moto mocks are installed, which would prevent the 

155 # mocks from taking effect. 

156 _get_s3_client.cache_clear() 

157 yield 

158 

159 

160def getS3Client(profile: str | None = None) -> BaseClient: 

161 """Create a S3 client with AWS (default) or the specified endpoint. 

162 

163 Parameters 

164 ---------- 

165 profile : `str`, optional 

166 The name of an S3 profile describing which S3 service to use. 

167 

168 Returns 

169 ------- 

170 s3client : `botocore.client.S3` 

171 A client of the S3 service. 

172 

173 Notes 

174 ----- 

175 If an explicit profile name is specified, its configuration will be read 

176 from an environment variable named ``LSST_RESOURCES_S3_PROFILE_<profile>`` 

177 if it exists. Note that the name of the profile is case sensitive. This 

178 configuration is specified in the format: ``https://<access key ID>:<secret 

179 key>@<s3 endpoint hostname>``. If the access key ID or secret key values 

180 contain slashes, the slashes must be URI-encoded (replace "/" with "%2F"). 

181 

182 If profile is `None` or the profile environment variable was not set, the 

183 configuration is read from the environment variable ``S3_ENDPOINT_URL``. 

184 If it is not specified, the default AWS endpoint is used. 

185 

186 The access key ID and secret key are optional -- if not specified, they 

187 will be looked up via the `AWS credentials file 

188 <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html>`_. 

189 

190 If the environment variable LSST_DISABLE_BUCKET_VALIDATION exists 

191 and has a value that is not empty, "0", "f", "n", or "false" 

192 (case-insensitive), then bucket name validation is disabled. This 

193 disabling allows Ceph multi-tenancy colon separators to appear in 

194 bucket names. 

195 """ 

196 if boto3 is None: 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 raise ModuleNotFoundError("Could not find boto3. Are you sure it is installed?") 

198 if botocore is None: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true

199 raise ModuleNotFoundError("Could not find botocore. Are you sure it is installed?") 

200 

201 endpoint_config = _get_s3_connection_parameters(profile) 

202 

203 return _get_s3_client(endpoint_config, not _s3_should_validate_bucket()) 

204 

205 

206def _s3_should_validate_bucket() -> bool: 

207 """Indicate whether bucket validation should be enabled. 

208 

209 Returns 

210 ------- 

211 validate : `bool` 

212 If `True` bucket names should be validated. 

213 """ 

214 disable_value = os.environ.get("LSST_DISABLE_BUCKET_VALIDATION", "0") 

215 return bool(re.search(r"^(0|f|n|false)?$", disable_value, re.I)) 

216 

217 

218def _get_s3_connection_parameters(profile: str | None = None) -> _EndpointConfig: 

219 """Calculate the connection details. 

220 

221 Parameters 

222 ---------- 

223 profile : `str`, optional 

224 The name of an S3 profile describing which S3 service to use. 

225 

226 Returns 

227 ------- 

228 config : _EndPointConfig 

229 All the information necessary to connect to the bucket. 

230 """ 

231 endpoint = None 

232 if profile is not None: 

233 var_name = f"LSST_RESOURCES_S3_PROFILE_{profile}" 

234 endpoint = os.environ.get(var_name, None) 

235 if not endpoint: 

236 endpoint = os.environ.get("S3_ENDPOINT_URL", None) 

237 if not endpoint: 

238 endpoint = None # Handle "" 

239 

240 return _parse_endpoint_config(endpoint, profile) 

241 

242 

243def _s3_disable_bucket_validation(client: BaseClient) -> None: 

244 """Disable the bucket name validation in the client. 

245 

246 This removes the ``validate_bucket_name`` handler from the handlers 

247 registered for this client. 

248 

249 Parameters 

250 ---------- 

251 client : `boto3.client` 

252 The client to modify. 

253 """ 

254 client.meta.events.unregister("before-parameter-build.s3", validate_bucket_name) 

255 

256 

257@functools.lru_cache 

258def _get_s3_client(endpoint_config: _EndpointConfig, skip_validation: bool) -> BaseClient: 

259 # Helper function to cache the client for this endpoint 

260 # boto seems to assume it will always have at least 10 available. 

261 max_pool_size = max(_get_num_workers(), 10) 

262 config = botocore.config.Config( 

263 read_timeout=180, 

264 max_pool_connections=max_pool_size, 

265 retries={"mode": "adaptive", "max_attempts": 10}, 

266 ) 

267 

268 session = boto3.Session(profile_name=endpoint_config.profile) 

269 

270 client = session.client( 

271 "s3", 

272 endpoint_url=endpoint_config.endpoint_url, 

273 aws_access_key_id=endpoint_config.access_key_id, 

274 aws_secret_access_key=endpoint_config.secret_access_key, 

275 config=config, 

276 ) 

277 if skip_validation: 

278 _s3_disable_bucket_validation(client) 

279 return client 

280 

281 

282class _EndpointConfig(NamedTuple): 

283 endpoint_url: str | None = None 

284 access_key_id: str | None = None 

285 secret_access_key: str | None = None 

286 profile: str | None = None 

287 

288 

289def _parse_endpoint_config(endpoint: str | None, profile: str | None = None) -> _EndpointConfig: 

290 if not endpoint: 

291 return _EndpointConfig(profile=profile) 

292 

293 parsed = parse_url(endpoint) 

294 

295 # Strip the username/password portion of the URL from the result. 

296 endpoint_url = Url(host=parsed.host, path=parsed.path, port=parsed.port, scheme=parsed.scheme).url 

297 

298 access_key_id = None 

299 secret_access_key = None 

300 if parsed.auth: 

301 split = parsed.auth.split(":") 

302 if len(split) != 2: 

303 raise ValueError("S3 access key and secret not in expected format.") 

304 access_key_id, secret_access_key = split 

305 access_key_id = urllib.parse.unquote(access_key_id) 

306 secret_access_key = urllib.parse.unquote(secret_access_key) 

307 

308 if access_key_id is not None and secret_access_key is not None: 

309 # We already have the necessary configuration for the profile, so do 

310 # not pass the profile to boto3. boto3 will raise an exception if the 

311 # profile is not defined in its configuration file, whether or not it 

312 # needs to read the configuration from it. 

313 profile = None 

314 

315 return _EndpointConfig( 

316 endpoint_url=endpoint_url, 

317 access_key_id=access_key_id, 

318 secret_access_key=secret_access_key, 

319 profile=profile, 

320 ) 

321 

322 

323def s3CheckFileExists( 

324 path: Location | ResourcePath | str, 

325 bucket: str | None = None, 

326 client: BaseClient | None = None, 

327) -> tuple[bool, int]: 

328 """Return if the file exists in the bucket or not. 

329 

330 Parameters 

331 ---------- 

332 path : `Location`, `ResourcePath` or `str` 

333 Location or ResourcePath containing the bucket name and filepath. 

334 bucket : `str`, optional 

335 Name of the bucket in which to look. If provided, path will be assumed 

336 to correspond to be relative to the given bucket. 

337 client : `boto3.client`, optional 

338 S3 Client object to query, if not supplied boto3 will try to resolve 

339 the credentials as in order described in its manual_. 

340 

341 Returns 

342 ------- 

343 exists : `bool` 

344 True if key exists, False otherwise. 

345 size : `int` 

346 Size of the key, if key exists, in bytes, otherwise -1. 

347 

348 Notes 

349 ----- 

350 S3 Paths are sensitive to leading and trailing path separators. 

351 

352 .. _manual: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/\ 

353 configuration.html#configuring-credentials 

354 """ 

355 if boto3 is None: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true

356 raise ModuleNotFoundError("Could not find boto3. Are you sure it is installed?") 

357 

358 if client is None: 

359 client = getS3Client() 

360 

361 if isinstance(path, str): 

362 if bucket is not None: 

363 filepath = path 

364 else: 

365 uri = ResourcePath(path) 

366 bucket = uri.netloc 

367 filepath = uri.relativeToPathRoot 

368 elif isinstance(path, ResourcePath | Location): 368 ↛ 373line 368 didn't jump to line 373 because the condition on line 368 was always true

369 if bucket is None: 

370 bucket = path.netloc 

371 filepath = path.relativeToPathRoot 

372 else: 

373 raise TypeError(f"Unsupported path type: {path!r}.") 

374 

375 try: 

376 obj = client.head_object(Bucket=bucket, Key=filepath) 

377 return (True, obj["ContentLength"]) 

378 except client.exceptions.ClientError as err: 

379 # resource unreachable error means key does not exist 

380 errcode = err.response["ResponseMetadata"]["HTTPStatusCode"] 

381 if errcode == 404: 381 ↛ 389line 381 didn't jump to line 389 because the condition on line 381 was always true

382 return (False, -1) 

383 # head_object returns 404 when object does not exist only when user has 

384 # s3:ListBucket permission. If list permission does not exist a 403 is 

385 # returned. In practical terms this generally means that the file does 

386 # not exist, but it could also mean user lacks s3:GetObject permission: 

387 # https://docs.aws.amazon.com/AmazonS3/latest/API/RESTObjectHEAD.html 

388 # I don't think its possible to discern which case is it with certainty 

389 if errcode == 403: 

390 raise PermissionError( 

391 "Forbidden HEAD operation error occurred. " 

392 "Verify s3:ListBucket and s3:GetObject " 

393 "permissions are granted for your IAM user. " 

394 ) from err 

395 if errcode == 429: 

396 # boto3, incorrectly, does not automatically retry with 429 

397 # so instead we raise an explicit retry exception for backoff. 

398 raise _TooManyRequestsError(str(err)) from err 

399 raise 

400 

401 

402def bucketExists(bucketName: str, client: BaseClient | None = None) -> bool: 

403 """Check if the S3 bucket with the given name actually exists. 

404 

405 Parameters 

406 ---------- 

407 bucketName : `str` 

408 Name of the S3 Bucket. 

409 client : `boto3.client`, optional 

410 S3 Client object to query, if not supplied boto3 will try to resolve 

411 the credentials by calling `getS3Client`. 

412 

413 Returns 

414 ------- 

415 exists : `bool` 

416 True if it exists, False if no Bucket with specified parameters is 

417 found. 

418 """ 

419 if boto3 is None: 419 ↛ 420line 419 didn't jump to line 420 because the condition on line 419 was never true

420 raise ModuleNotFoundError("Could not find boto3. Are you sure it is installed?") 

421 

422 if client is None: 

423 client = getS3Client() 

424 try: 

425 client.get_bucket_location(Bucket=bucketName) 

426 return True 

427 except client.exceptions.NoSuchBucket: 

428 return False 

429 

430 

431def translate_client_error(err: ClientError, uri: ResourcePath) -> None: 

432 """Translate a ClientError into a specialist error if relevant. 

433 

434 Parameters 

435 ---------- 

436 err : `ClientError` 

437 Exception to translate. 

438 uri : `ResourcePath` 

439 The URI of the resource that is resulting in the error. 

440 

441 Raises 

442 ------ 

443 _TooManyRequestsError 

444 Raised if the `ClientError` looks like a 429 retry request. 

445 """ 

446 if "(429)" in str(err): 446 ↛ 450line 446 didn't jump to line 450 because the condition on line 446 was never true

447 # ClientError includes the error code in the message 

448 # but no direct way to access it without looking inside the 

449 # response. 

450 raise _TooManyRequestsError(f"{err} when accessing {uri}") from err 

451 elif "(404)" in str(err): 451 ↛ exitline 451 didn't return from function 'translate_client_error' because the condition on line 451 was always true

452 # Some systems can generate this rather than NoSuchKey. 

453 raise FileNotFoundError(f"Resource not found (permission denied): {uri}")