Coverage for python/lsst/daf/butler/datastores/chainedDatastore.py: 91%

524 statements  

« prev     ^ index     » next       coverage.py v7.17.2, created at 2026-11-06 09:46 +0000

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

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

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

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28"""Chained datastore.""" 

29 

30from __future__ import annotations 

31 

32__all__ = ("ChainedDatastore",) 

33 

34import itertools 

35import time 

36import warnings 

37from collections.abc import Callable, Collection, Iterable, Mapping, Sequence 

38from typing import TYPE_CHECKING, Any 

39 

40from lsst.daf.butler import DatasetRef, DatasetTypeNotSupportedError, DimensionUniverse, FileDataset 

41from lsst.daf.butler.datastore import ( 

42 DatasetRefURIs, 

43 Datastore, 

44 DatastoreConfig, 

45 DatastoreOpaqueTable, 

46 DatastoreValidationError, 

47) 

48from lsst.daf.butler.datastore.constraints import Constraints 

49from lsst.daf.butler.datastore.record_data import DatastoreRecordData, DatastoreRecordTable 

50from lsst.daf.butler.datastores.file_datastore.get import DatasetLocationInformation 

51from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ArtifactIndexInfo, ZipIndex 

52from lsst.resources import ResourcePath 

53from lsst.utils import doImportType 

54from lsst.utils.introspection import get_full_type_name 

55from lsst.utils.logging import getLogger 

56 

57from .._dataset_ref import DatasetId 

58from ..datastore import DuplicateIdError, FileTransferMap 

59from .utils import report_duplicate_ids 

60 

61if TYPE_CHECKING: 

62 from lsst.daf.butler import Config, DatasetProvenance, DatasetType, LookupKey, StorageClass 

63 from lsst.daf.butler.registry.interfaces import DatasetIdRef, DatastoreRegistryBridgeManager 

64 from lsst.resources import ResourcePathExpression 

65 

66log = getLogger(__name__) 

67 

68 

69class _IngestPrepData(Datastore.IngestPrepData): 

70 """Helper class for ChainedDatastore ingest implementation. 

71 

72 Parameters 

73 ---------- 

74 children : `list` of `tuple` 

75 Pairs of `Datastore`, `IngestPrepData` for all child datastores. 

76 """ 

77 

78 def __init__(self, children: list[tuple[Datastore, Datastore.IngestPrepData, set[ResourcePath]]]): 

79 super().__init__(itertools.chain.from_iterable(data.refs.values() for _, data, _ in children)) 

80 self.children = children 

81 

82 

83class ChainedDatastore(Datastore): 

84 """Chained Datastores to allow read and writes from multiple datastores. 

85 

86 A ChainedDatastore is configured with multiple datastore configurations. 

87 A ``put()`` is always sent to each datastore. A ``get()`` 

88 operation is sent to each datastore in turn and the first datastore 

89 to return a valid dataset is used. 

90 

91 Parameters 

92 ---------- 

93 config : `DatastoreConfig` or `str` 

94 Configuration. This configuration must include a ``datastores`` field 

95 as a sequence of datastore configurations. The order in this sequence 

96 indicates the order to use for read operations. 

97 bridgeManager : `DatastoreRegistryBridgeManager` 

98 Object that manages the interface between `Registry` and datastores. 

99 datastores : `list` [`Datastore`] 

100 All the child datastores known to this datastore. 

101 

102 Notes 

103 ----- 

104 ChainedDatastore never supports `None` or ``"move"`` as an ingest transfer 

105 mode. It supports ``"copy"``, ``"symlink"``, ``"relsymlink"`` 

106 and ``"hardlink"`` if and only if all its child datastores do. 

107 """ 

108 

109 defaultConfigFile = "datastores/chainedDatastore.yaml" 

110 """Path to configuration defaults. Accessed within the ``configs`` resource 

111 or relative to a search path. Can be None if no defaults specified. 

112 """ 

113 

114 containerKey = "datastores" 

115 """Key to specify where child datastores are configured.""" 

116 

117 datastores: list[Datastore] 

118 """All the child datastores known to this datastore.""" 

119 

120 datastoreConstraints: Sequence[Constraints | None] 

121 """Constraints to be applied to each of the child datastores.""" 

122 

123 universe: DimensionUniverse 

124 """Dimension universe associated with the butler.""" 

125 

126 @classmethod 

127 def setConfigRoot(cls, root: str, config: Config, full: Config, overwrite: bool = True) -> None: 

128 """Set any filesystem-dependent config options for child Datastores to 

129 be appropriate for a new empty repository with the given root. 

130 

131 Parameters 

132 ---------- 

133 root : `str` 

134 Filesystem path to the root of the data repository. 

135 config : `Config` 

136 A `Config` to update. Only the subset understood by 

137 this component will be updated. Will not expand 

138 defaults. 

139 full : `Config` 

140 A complete config with all defaults expanded that can be 

141 converted to a `DatastoreConfig`. Read-only and will not be 

142 modified by this method. 

143 Repository-specific options that should not be obtained 

144 from defaults when Butler instances are constructed 

145 should be copied from ``full`` to ``config``. 

146 overwrite : `bool`, optional 

147 If `False`, do not modify a value in ``config`` if the value 

148 already exists. Default is always to overwrite with the provided 

149 ``root``. 

150 

151 Notes 

152 ----- 

153 If a keyword is explicitly defined in the supplied ``config`` it 

154 will not be overridden by this method if ``overwrite`` is `False`. 

155 This allows explicit values set in external configs to be retained. 

156 """ 

157 # Extract the part of the config we care about updating 

158 datastoreConfig = DatastoreConfig(config, mergeDefaults=False) 

159 

160 # And the subset of the full config that we can use for reference. 

161 # Do not bother with defaults because we are told this already has 

162 # them. 

163 fullDatastoreConfig = DatastoreConfig(full, mergeDefaults=False) 

164 

165 # Loop over each datastore config and pass the subsets to the 

166 # child datastores to process. 

167 

168 containerKey = cls.containerKey 

169 for idx, (child, fullChild) in enumerate( 

170 zip(datastoreConfig[containerKey], fullDatastoreConfig[containerKey], strict=True) 

171 ): 

172 childConfig = DatastoreConfig(child, mergeDefaults=False) 

173 fullChildConfig = DatastoreConfig(fullChild, mergeDefaults=False) 

174 datastoreClass = doImportType(fullChildConfig["cls"]) 

175 if not issubclass(datastoreClass, Datastore): 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true

176 raise TypeError(f"Imported child class {fullChildConfig['cls']} is not a Datastore") 

177 newroot = f"{root}/{datastoreClass.__qualname__}_{idx}" 

178 datastoreClass.setConfigRoot(newroot, childConfig, fullChildConfig, overwrite=overwrite) 

179 

180 # Reattach to parent 

181 datastoreConfig[containerKey, idx] = childConfig 

182 

183 # Reattach modified datastore config to parent 

184 # If this has a datastore key we attach there, otherwise we assume 

185 # this information goes at the top of the config hierarchy. 

186 if DatastoreConfig.component in config: 

187 config[DatastoreConfig.component] = datastoreConfig 

188 else: 

189 config.update(datastoreConfig) 

190 

191 return 

192 

193 def __init__( 

194 self, 

195 config: DatastoreConfig, 

196 bridgeManager: DatastoreRegistryBridgeManager, 

197 datastores: list[Datastore], 

198 ): 

199 super().__init__(config, bridgeManager) 

200 

201 self.datastores = list(datastores) 

202 

203 # Name ourself based on our children 

204 if self.datastores: 204 ↛ 209line 204 didn't jump to line 209 because the condition on line 204 was always true

205 # We must set the names explicitly 

206 self._names = [d.name for d in self.datastores] 

207 childNames = ",".join(self.names) 

208 else: 

209 childNames = f"(empty@{time.time()})" 

210 self._names = [childNames] 

211 self.name = f"{type(self).__qualname__}[{childNames}]" 

212 

213 # We declare we are ephemeral if all our child datastores declare 

214 # they are ephemeral 

215 self.isEphemeral = all(d.isEphemeral for d in self.datastores) 

216 

217 # per-datastore override constraints 

218 if "datastore_constraints" in self.config: 

219 overrides = self.config["datastore_constraints"] 

220 

221 if len(overrides) != len(self.datastores): 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true

222 raise DatastoreValidationError( 

223 f"Number of registered datastores ({len(self.datastores)})" 

224 " differs from number of constraints overrides" 

225 f" {len(overrides)}" 

226 ) 

227 

228 self.datastoreConstraints = [ 

229 Constraints(c.get("constraints"), universe=bridgeManager.universe) for c in overrides 

230 ] 

231 

232 else: 

233 self.datastoreConstraints = (None,) * len(self.datastores) 

234 

235 self.universe = bridgeManager.universe 

236 

237 log.debug("Created %s (%s)", self.name, ("ephemeral" if self.isEphemeral else "permanent")) 

238 

239 @classmethod 

240 def _create_from_config( 

241 cls, 

242 config: DatastoreConfig, 

243 bridgeManager: DatastoreRegistryBridgeManager, 

244 butlerRoot: ResourcePathExpression | None, 

245 ) -> ChainedDatastore: 

246 # Scan for child datastores and instantiate them with the same registry 

247 datastores = [] 

248 for c in config["datastores"]: 

249 c = DatastoreConfig(c) 

250 datastoreType = doImportType(c["cls"]) 

251 if not issubclass(datastoreType, Datastore): 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true

252 raise TypeError(f"Imported child class {c['cls']} is not a Datastore") 

253 datastore = datastoreType._create_from_config(c, bridgeManager, butlerRoot=butlerRoot) 

254 log.debug("Creating child datastore %s", datastore.name) 

255 datastores.append(datastore) 

256 

257 return ChainedDatastore(config, bridgeManager, datastores) 

258 

259 def clone(self, bridgeManager: DatastoreRegistryBridgeManager) -> Datastore: 

260 datastores = [ds.clone(bridgeManager) for ds in self.datastores] 

261 return ChainedDatastore(self.config, bridgeManager, datastores) 

262 

263 @property 

264 def names(self) -> tuple[str, ...]: 

265 return tuple(self._names) 

266 

267 @property 

268 def roots(self) -> dict[str, ResourcePath | None]: 

269 # Docstring inherited. 

270 roots = {} 

271 for datastore in self.datastores: 

272 roots.update(datastore.roots) 

273 return roots 

274 

275 def __str__(self) -> str: 

276 chainName = ", ".join(str(ds) for ds in self.datastores) 

277 return chainName 

278 

279 def _set_trust_mode(self, mode: bool) -> None: 

280 for datastore in self.datastores: 

281 datastore._set_trust_mode(mode) 

282 

283 def knows(self, ref: DatasetRef) -> bool: 

284 """Check if the dataset is known to any of the datastores. 

285 

286 Does not check for existence of any artifact. 

287 

288 Parameters 

289 ---------- 

290 ref : `DatasetRef` 

291 Reference to the required dataset. 

292 

293 Returns 

294 ------- 

295 exists : `bool` 

296 `True` if the dataset is known to the datastore. 

297 """ 

298 for datastore in self.datastores: 

299 if datastore.knows(ref): 

300 log.debug("%s known to datastore %s", ref, datastore.name) 

301 return True 

302 return False 

303 

304 def knows_these(self, refs: Iterable[DatasetRef]) -> dict[DatasetRef, bool]: 

305 # Docstring inherited from the base class. 

306 refs_known: dict[DatasetRef, bool] = {} 

307 for datastore in self.datastores: 

308 refs_known.update(datastore.knows_these(refs)) 

309 

310 # No need to check in next datastore for refs that are known. 

311 # We only update entries that were initially False. 

312 refs = [ref for ref, known in refs_known.items() if not known] 

313 

314 return refs_known 

315 

316 def mexists( 

317 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None 

318 ) -> dict[DatasetRef, bool]: 

319 """Check the existence of multiple datasets at once. 

320 

321 Parameters 

322 ---------- 

323 refs : `~collections.abc.Iterable` of `DatasetRef` 

324 The datasets to be checked. 

325 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`] 

326 Optional mapping of datastore artifact to existence. Updated by 

327 this method with details of all artifacts tested. Can be `None` 

328 if the caller is not interested. 

329 

330 Returns 

331 ------- 

332 existence : `dict` of [`DatasetRef`, `bool`] 

333 Mapping from dataset to boolean indicating existence in any 

334 of the child datastores. 

335 """ 

336 dataset_existence: dict[DatasetRef, bool] = {} 

337 for datastore in self.datastores: 

338 dataset_existence.update(datastore.mexists(refs, artifact_existence=artifact_existence)) 

339 

340 # For next datastore no point asking about ones we know 

341 # exist already. No special exemption for ephemeral datastores. 

342 refs = [ref for ref, exists in dataset_existence.items() if not exists] 

343 

344 return dataset_existence 

345 

346 def exists(self, ref: DatasetRef) -> bool: 

347 """Check if the dataset exists in one of the datastores. 

348 

349 Parameters 

350 ---------- 

351 ref : `DatasetRef` 

352 Reference to the required dataset. 

353 

354 Returns 

355 ------- 

356 exists : `bool` 

357 `True` if the entity exists in one of the child datastores. 

358 """ 

359 for datastore in self.datastores: 

360 if datastore.exists(ref): 

361 log.debug("Found %s in datastore %s", ref, datastore.name) 

362 return True 

363 return False 

364 

365 def get( 

366 self, 

367 ref: DatasetRef, 

368 parameters: Mapping[str, Any] | None = None, 

369 storageClass: StorageClass | str | None = None, 

370 ) -> Any: 

371 """Load an InMemoryDataset from the store. 

372 

373 The dataset is returned from the first datastore that has 

374 the dataset. 

375 

376 Parameters 

377 ---------- 

378 ref : `DatasetRef` 

379 Reference to the required Dataset. 

380 parameters : `dict` 

381 `StorageClass`-specific parameters that specify, for example, 

382 a slice of the dataset to be loaded. 

383 storageClass : `StorageClass` or `str`, optional 

384 The storage class to be used to override the Python type 

385 returned by this method. By default the returned type matches 

386 the dataset type definition for this dataset. Specifying a 

387 read `StorageClass` can force a different type to be returned. 

388 This type must be compatible with the original type. 

389 

390 Returns 

391 ------- 

392 inMemoryDataset : `object` 

393 Requested dataset or slice thereof as an InMemoryDataset. 

394 

395 Raises 

396 ------ 

397 FileNotFoundError 

398 Requested dataset can not be retrieved. 

399 TypeError 

400 Return value from formatter has unexpected type. 

401 ValueError 

402 Formatter failed to process the dataset. 

403 """ 

404 for datastore in self.datastores: 

405 try: 

406 inMemoryObject = datastore.get(ref, parameters, storageClass=storageClass) 

407 log.debug("Found dataset %s in datastore %s", ref, datastore.name) 

408 return inMemoryObject 

409 except FileNotFoundError: 

410 pass 

411 

412 raise FileNotFoundError(f"Dataset {ref} could not be found in any of the datastores") 

413 

414 def prepare_get_for_external_client(self, ref: DatasetRef) -> list[DatasetLocationInformation] | None: 

415 datastore = self._get_matching_datastore(ref) 

416 if datastore is None: 416 ↛ 419line 416 didn't jump to line 419 because the condition on line 416 was always true

417 return None 

418 

419 return datastore.prepare_get_for_external_client(ref) 

420 

421 def _get_matching_datastore(self, ref: DatasetRef) -> Datastore | None: 

422 """Return the first child datastore that owns the specified dataset.""" 

423 for datastore in self.datastores: 

424 if datastore.knows(ref): 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true

425 return datastore 

426 

427 return None 

428 

429 def put(self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None) -> None: 

430 """Write a InMemoryDataset with a given `DatasetRef` to each 

431 datastore. 

432 

433 The put() to child datastores can fail with 

434 `DatasetTypeNotSupportedError`. The put() for this datastore will be 

435 deemed to have succeeded so long as at least one child datastore 

436 accepted the inMemoryDataset. 

437 

438 Parameters 

439 ---------- 

440 inMemoryDataset : `object` 

441 The dataset to store. 

442 ref : `DatasetRef` 

443 Reference to the associated Dataset. 

444 provenance : `DatasetProvenance` or `None`, optional 

445 Any provenance that should be attached to the serialized dataset. 

446 Not supported by all serialization mechanisms. 

447 

448 Raises 

449 ------ 

450 TypeError 

451 Supplied object and storage class are inconsistent. 

452 DatasetTypeNotSupportedError 

453 All datastores reported `DatasetTypeNotSupportedError`. 

454 """ 

455 log.debug("Put %s", ref) 

456 

457 # Before trying to call put() on any child datastore check if any of 

458 # them already have this dataset ID, this is to avoid potential 

459 # inconsistencies when one of the children has that dataset already but 

460 # others do not. 

461 if self.knows(ref): 

462 raise DuplicateIdError(f"Attempt to override existing dataset: {ref.id}") 

463 

464 # Confirm that we can accept this dataset 

465 if not self.constraints.isAcceptable(ref): 

466 # Raise rather than use boolean return value. 

467 raise DatasetTypeNotSupportedError( 

468 f"Dataset {ref} has been rejected by this datastore via configuration." 

469 ) 

470 

471 isPermanent = False 

472 nsuccess = 0 

473 npermanent = 0 

474 nephemeral = 0 

475 for datastore, constraints in zip(self.datastores, self.datastoreConstraints, strict=True): 

476 if ( 

477 constraints is not None and not constraints.isAcceptable(ref) 

478 ) or not datastore.constraints.isAcceptable(ref): 

479 log.debug("Datastore %s skipping put via configuration for ref %s", datastore.name, ref) 

480 continue 

481 

482 if datastore.isEphemeral: 

483 nephemeral += 1 

484 else: 

485 npermanent += 1 

486 try: 

487 datastore.put(inMemoryDataset, ref, provenance=provenance) 

488 nsuccess += 1 

489 if not datastore.isEphemeral: 

490 isPermanent = True 

491 except DatasetTypeNotSupportedError: 

492 pass 

493 

494 if nsuccess == 0: 

495 raise DatasetTypeNotSupportedError(f"None of the chained datastores supported ref {ref}") 

496 

497 if not isPermanent and npermanent > 0: 497 ↛ 498line 497 didn't jump to line 498 because the condition on line 497 was never true

498 warnings.warn(f"Put of {ref} only succeeded in ephemeral databases", stacklevel=2) 

499 

500 if self._transaction is not None: 

501 self._transaction.registerUndo("put", self.remove, ref) 

502 

503 def _overrideTransferMode(self, *datasets: Any, transfer: str | None = None) -> str | None: 

504 # Docstring inherited from base class. 

505 if transfer != "auto": 

506 return transfer 

507 # Ask each datastore what they think auto means 

508 transfers = {d._overrideTransferMode(*datasets, transfer=transfer) for d in self.datastores} 

509 

510 # Remove any untranslated "auto" values 

511 transfers.discard(transfer) 

512 

513 if len(transfers) == 1: 513 ↛ 514line 513 didn't jump to line 514 because the condition on line 513 was never true

514 return transfers.pop() 

515 if not transfers: 515 ↛ 519line 515 didn't jump to line 519 because the condition on line 515 was always true

516 # Everything reported "auto" 

517 return transfer 

518 

519 raise RuntimeError( 

520 "Chained datastore does not yet support different transfer modes" 

521 f" from 'auto' in each child datastore (wanted {transfers})" 

522 ) 

523 

524 def _prepIngest(self, *datasets: FileDataset, transfer: str | None = None) -> _IngestPrepData: 

525 # Docstring inherited from Datastore._prepIngest. 

526 def isDatasetAcceptable(dataset: FileDataset, *, name: str, constraints: Constraints) -> bool: 

527 acceptable = [ref for ref in dataset.refs if constraints.isAcceptable(ref)] 

528 if not acceptable: 

529 log.debug( 

530 "Datastore %s skipping ingest via configuration for refs %s", 

531 name, 

532 ", ".join(str(ref) for ref in dataset.refs), 

533 ) 

534 return False 

535 else: 

536 return True 

537 

538 # Filter down to just datasets the chained datastore's own 

539 # configuration accepts. 

540 okForParent: list[FileDataset] = [ 

541 dataset 

542 for dataset in datasets 

543 if isDatasetAcceptable(dataset, name=self.name, constraints=self.constraints) 

544 ] 

545 

546 # Iterate over nested datastores and call _prepIngest on each. 

547 # Save the results to a list: 

548 children: list[tuple[Datastore, Datastore.IngestPrepData, set[ResourcePath]]] = [] 

549 # ...and remember whether all of the failures are due to 

550 # NotImplementedError being raised. 

551 allFailuresAreNotImplementedError = True 

552 for datastore, constraints in zip(self.datastores, self.datastoreConstraints, strict=True): 

553 okForChild: list[FileDataset] 

554 if constraints is not None: 

555 okForChild = [ 

556 dataset 

557 for dataset in okForParent 

558 if isDatasetAcceptable(dataset, name=datastore.name, constraints=constraints) 

559 ] 

560 else: 

561 okForChild = okForParent 

562 try: 

563 prepDataForChild = datastore._prepIngest(*okForChild, transfer=transfer) 

564 except NotImplementedError: 

565 log.debug( 

566 "Skipping ingest for datastore %s because transfer mode %s is not supported.", 

567 datastore.name, 

568 transfer, 

569 ) 

570 continue 

571 allFailuresAreNotImplementedError = False 

572 if okForChild: 

573 # Do not store for later if a datastore has rejected 

574 # everything. 

575 # Include the source paths if this is a "move". It's clearer 

576 # to find the paths now rather than try to infer how 

577 # each datastore has stored them in the internal prep class. 

578 paths = ( 

579 {ResourcePath(dataset.path, forceDirectory=False) for dataset in okForChild} 

580 if transfer == "move" 

581 else set() 

582 ) 

583 children.append((datastore, prepDataForChild, paths)) 

584 if allFailuresAreNotImplementedError: 

585 raise NotImplementedError(f"No child datastore supports transfer mode {transfer}.") 

586 return _IngestPrepData(children=children) 

587 

588 def _finishIngest( 

589 self, 

590 prepData: _IngestPrepData, 

591 *, 

592 transfer: str | None = None, 

593 record_validation_info: bool = True, 

594 ) -> None: 

595 # Docstring inherited from Datastore._finishIngest. 

596 # For "move" we must use "copy" and then delete the input 

597 # data at the end. This has no rollback option if the ingest 

598 # subsequently fails. If there is only one active datastore 

599 # accepting any files we can leave it as "move" 

600 actual_transfer: str | None 

601 if transfer == "move" and len(prepData.children) > 1: 

602 actual_transfer = "copy" 

603 else: 

604 actual_transfer = transfer 

605 to_be_deleted: set[ResourcePath] = set() 

606 for datastore, prepDataForChild, paths in prepData.children: 

607 datastore._finishIngest( 

608 prepDataForChild, transfer=actual_transfer, record_validation_info=record_validation_info 

609 ) 

610 to_be_deleted.update(paths) 

611 if actual_transfer != transfer: 

612 # These datasets were copied but now need to be deleted. 

613 # This can not be rolled back. 

614 for uri in to_be_deleted: 

615 uri.remove() 

616 

617 def getManyURIs( 

618 self, 

619 refs: Iterable[DatasetRef], 

620 predict: bool = False, 

621 allow_missing: bool = False, 

622 ) -> dict[DatasetRef, DatasetRefURIs]: 

623 # Docstring inherited 

624 

625 uris: dict[DatasetRef, DatasetRefURIs] = {} 

626 ephemeral_uris: dict[DatasetRef, DatasetRefURIs] = {} 

627 

628 missing_refs = set(refs) 

629 

630 # If predict is True we don't want to predict a dataset in the first 

631 # datastore if it actually exists in a later datastore, so in that 

632 # case check all datastores with predict=False first, and then try 

633 # again with predict=True. 

634 for p in (False, True) if predict else (False,): 

635 if not missing_refs: 

636 break 

637 for datastore in self.datastores: 

638 try: 

639 got_uris = datastore.getManyURIs(missing_refs, p, allow_missing=True) 

640 except NotImplementedError: 

641 # some datastores may not implement generating URIs 

642 continue 

643 if datastore.isEphemeral: 

644 # Only use these as last resort so do not constrain 

645 # subsequent queries. 

646 ephemeral_uris.update(got_uris) 

647 continue 

648 

649 missing_refs -= got_uris.keys() 

650 uris.update(got_uris) 

651 if not missing_refs: 

652 break 

653 

654 if missing_refs and ephemeral_uris: 

655 ephemeral_refs = missing_refs.intersection(ephemeral_uris.keys()) 

656 for ref in ephemeral_refs: 

657 uris[ref] = ephemeral_uris[ref] 

658 missing_refs.remove(ref) 

659 

660 if missing_refs and not allow_missing: 

661 raise FileNotFoundError(f"Dataset(s) {missing_refs} not in this datastore.") 

662 

663 return uris 

664 

665 def getURIs(self, ref: DatasetRef, predict: bool = False) -> DatasetRefURIs: 

666 """Return URIs associated with dataset. 

667 

668 Parameters 

669 ---------- 

670 ref : `DatasetRef` 

671 Reference to the required dataset. 

672 predict : `bool`, optional 

673 If the datastore does not know about the dataset, controls whether 

674 it should return a predicted URI or not. 

675 

676 Returns 

677 ------- 

678 uris : `DatasetRefURIs` 

679 The URI to the primary artifact associated with this dataset (if 

680 the dataset was disassembled within the datastore this may be 

681 `None`), and the URIs to any components associated with the dataset 

682 artifact. (can be empty if there are no components). 

683 

684 Notes 

685 ----- 

686 The returned URI is from the first datastore in the list that has 

687 the dataset with preference given to the first dataset coming from 

688 a permanent datastore. If no datastores have the dataset and prediction 

689 is allowed, the predicted URI for the first datastore in the list will 

690 be returned. 

691 """ 

692 log.debug("Requesting URIs for %s", ref) 

693 predictedUri: DatasetRefURIs | None = None 

694 predictedEphemeralUri: DatasetRefURIs | None = None 

695 firstEphemeralUri: DatasetRefURIs | None = None 

696 for datastore in self.datastores: 

697 if datastore.exists(ref): 

698 if not datastore.isEphemeral: 

699 uri = datastore.getURIs(ref) 

700 log.debug("Retrieved non-ephemeral URI: %s", uri) 

701 return uri 

702 elif not firstEphemeralUri: 

703 firstEphemeralUri = datastore.getURIs(ref) 

704 elif predict: 

705 if not predictedUri and not datastore.isEphemeral: 

706 predictedUri = datastore.getURIs(ref, predict) 

707 elif not predictedEphemeralUri and datastore.isEphemeral: 

708 predictedEphemeralUri = datastore.getURIs(ref, predict) 

709 

710 if firstEphemeralUri: 

711 log.debug("Retrieved ephemeral URI: %s", firstEphemeralUri) 

712 return firstEphemeralUri 

713 

714 if predictedUri: 

715 log.debug("Retrieved predicted URI: %s", predictedUri) 

716 return predictedUri 

717 

718 if predictedEphemeralUri: 

719 log.debug("Retrieved predicted ephemeral URI: %s", predictedEphemeralUri) 

720 return predictedEphemeralUri 

721 

722 raise FileNotFoundError(f"Dataset {ref} not in any datastore") 

723 

724 def getURI(self, ref: DatasetRef, predict: bool = False) -> ResourcePath: 

725 """URI to the Dataset. 

726 

727 The returned URI is from the first datastore in the list that has 

728 the dataset with preference given to the first dataset coming from 

729 a permanent datastore. If no datastores have the dataset and prediction 

730 is allowed, the predicted URI for the first datastore in the list will 

731 be returned. 

732 

733 Parameters 

734 ---------- 

735 ref : `DatasetRef` 

736 Reference to the required Dataset. 

737 predict : `bool` 

738 If `True`, allow URIs to be returned of datasets that have not 

739 been written. 

740 

741 Returns 

742 ------- 

743 uri : `lsst.resources.ResourcePath` 

744 URI pointing to the dataset within the datastore. If the 

745 dataset does not exist in the datastore, and if ``predict`` is 

746 `True`, the URI will be a prediction and will include a URI 

747 fragment "#predicted". 

748 

749 Notes 

750 ----- 

751 If the datastore does not have entities that relate well 

752 to the concept of a URI the returned URI string will be 

753 descriptive. The returned URI is not guaranteed to be obtainable. 

754 

755 Raises 

756 ------ 

757 FileNotFoundError 

758 A URI has been requested for a dataset that does not exist and 

759 guessing is not allowed. 

760 RuntimeError 

761 Raised if a request is made for a single URI but multiple URIs 

762 are associated with this dataset. 

763 """ 

764 log.debug("Requesting URI for %s", ref) 

765 primary, components = self.getURIs(ref, predict) 

766 if primary is None or components: 766 ↛ 767line 766 didn't jump to line 767 because the condition on line 766 was never true

767 raise RuntimeError( 

768 f"Dataset ({ref}) includes distinct URIs for components. Use Datastore.getURIs() instead." 

769 ) 

770 return primary 

771 

772 def retrieveArtifacts( 

773 self, 

774 refs: Iterable[DatasetRef], 

775 destination: ResourcePath, 

776 transfer: str = "auto", 

777 preserve_path: bool = True, 

778 overwrite: bool = False, 

779 write_index: bool = True, 

780 add_prefix: bool = False, 

781 ) -> dict[ResourcePath, ArtifactIndexInfo]: 

782 """Retrieve the file artifacts associated with the supplied refs. 

783 

784 Parameters 

785 ---------- 

786 refs : `~collections.abc.Iterable` of `DatasetRef` 

787 The datasets for which file artifacts are to be retrieved. 

788 A single ref can result in multiple files. The refs must 

789 be resolved. 

790 destination : `lsst.resources.ResourcePath` 

791 Location to write the file artifacts. 

792 transfer : `str`, optional 

793 Method to use to transfer the artifacts. Must be one of the options 

794 supported by `lsst.resources.ResourcePath.transfer_from`. 

795 "move" is not allowed. 

796 preserve_path : `bool`, optional 

797 If `True` the full path of the file artifact within the datastore 

798 is preserved. If `False` the final file component of the path 

799 is used. 

800 overwrite : `bool`, optional 

801 If `True` allow transfers to overwrite existing files at the 

802 destination. 

803 write_index : `bool`, optional 

804 If `True` write a file at the top level containing a serialization 

805 of a `ZipIndex` for the downloaded datasets. 

806 add_prefix : `bool`, optional 

807 Add a prefix based on the DatasetId. Only used if ``preserve_path`` 

808 is `False`. 

809 

810 Returns 

811 ------- 

812 artifact_map : `dict` [ `lsst.resources.ResourcePath`, \ 

813 `ArtifactIndexInfo` ] 

814 Mapping of retrieved file to associated index information. 

815 """ 

816 if not destination.isdir(): 

817 raise ValueError(f"Destination location must refer to a directory. Given {destination}") 

818 

819 # Using getURIs is not feasible since it becomes difficult to 

820 # determine the path within the datastore later on. For now 

821 # follow getURIs implementation approach. 

822 

823 pending = set(refs) 

824 

825 # There is a question as to whether an exception should be raised 

826 # early if some of the refs are missing, or whether files should be 

827 # transferred until a problem is hit. Prefer to complain up front. 

828 # Use the datastore integer as primary key. 

829 grouped_by_datastore: dict[int, set[DatasetRef]] = {} 

830 

831 for number, datastore in enumerate(self.datastores): 

832 if datastore.isEphemeral: 

833 # In the future we will want to distinguish in-memory from 

834 # caching datastore since using an on-disk local 

835 # cache is exactly what we should be doing. 

836 continue 

837 try: 

838 # Checking file existence is expensive. Have the option 

839 # of checking whether the datastore knows of these datasets 

840 # instead, which is fast but can potentially lead to 

841 # retrieveArtifacts failing. 

842 knows = datastore.knows_these(pending) 

843 datastore_refs = {ref for ref, exists in knows.items() if exists} 

844 except NotImplementedError: 

845 # Some datastores may not support retrieving artifacts 

846 continue 

847 

848 if datastore_refs: 

849 grouped_by_datastore[number] = datastore_refs 

850 

851 # Remove these from the pending list so that we do not bother 

852 # looking for them any more. 

853 pending = pending - datastore_refs 

854 

855 if pending: 855 ↛ 856line 855 didn't jump to line 856 because the condition on line 855 was never true

856 raise RuntimeError(f"Some datasets were not found in any datastores: {pending}") 

857 

858 # Now do the transfer. 

859 merged_artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {} 

860 for number, datastore_refs in grouped_by_datastore.items(): 

861 artifact_map = self.datastores[number].retrieveArtifacts( 

862 datastore_refs, 

863 destination, 

864 transfer=transfer, 

865 preserve_path=preserve_path, 

866 overwrite=overwrite, 

867 write_index=False, # Disable index writing regardless. 

868 add_prefix=add_prefix, 

869 ) 

870 merged_artifact_map.update(artifact_map) 

871 

872 if write_index: 872 ↛ 876line 872 didn't jump to line 876 because the condition on line 872 was always true

873 index = ZipIndex.from_artifact_map(refs, merged_artifact_map, destination) 

874 index.write_index(destination) 

875 

876 return merged_artifact_map 

877 

878 def ingest_zip(self, zip_path: ResourcePath, transfer: str | None, *, dry_run: bool = False) -> None: 

879 """Ingest an indexed Zip file and contents. 

880 

881 The Zip file must have an index file as created by `retrieveArtifacts`. 

882 

883 Parameters 

884 ---------- 

885 zip_path : `lsst.resources.ResourcePath` 

886 Path to the Zip file. 

887 transfer : `str` 

888 Method to use for transferring the Zip file into the datastore. 

889 dry_run : `bool`, optional 

890 If `True` the ingest will be processed without any modifications 

891 made to the target datastore and as if the target datastore did not 

892 have any of the datasets. 

893 

894 Notes 

895 ----- 

896 Datastore constraints are bypassed with Zip ingest. A zip file can 

897 contain multiple dataset types. Should the entire Zip be rejected 

898 if one dataset type is in the constraints list? If configured to 

899 reject everything, ingest should not be attempted. 

900 

901 The Zip file is given to each datastore in turn, ignoring datastores 

902 where it is not supported. Is deemed successful if any of the 

903 datastores accept the file. 

904 """ 

905 index = ZipIndex.from_zip_file(zip_path) 

906 refs = index.refs.to_refs(self.universe) 

907 

908 # For now raise if any refs are not supported. 

909 # Being selective will require that we return the ingested refs 

910 # to the caller so that registry can be modified to remove the 

911 # entries. 

912 if any(not self.constraints.isAcceptable(ref) for ref in refs): 912 ↛ 913line 912 didn't jump to line 913 because the condition on line 912 was never true

913 raise DatasetTypeNotSupportedError( 

914 "Some of the refs in the given Zip file are not acceptable to this datastore." 

915 ) 

916 

917 # Before trying to call ingest_zip() on any child datastore check if 

918 # any of them already have these datasets, this is to avoid potential 

919 # inconsistencies when one of the children has that dataset already but 

920 # others do not. 

921 known_ids = [] 

922 if not dry_run: 922 ↛ 924line 922 didn't jump to line 924 because the condition on line 922 was always true

923 known_ids = [ref.id for ref, known in self.knows_these(refs).items() if known] 

924 if known_ids: 924 ↛ 926line 924 didn't jump to line 926 because the condition on line 924 was never true

925 # Will raise DuplicateIdError. 

926 report_duplicate_ids(known_ids) 

927 

928 n_success = 0 

929 final_exception: Exception | None = None 

930 for number, (datastore, constraints) in enumerate( 

931 zip(self.datastores, self.datastoreConstraints, strict=True) 

932 ): 

933 if datastore.isEphemeral: 

934 continue 

935 

936 # There can be constraints for the datastore in the configuration 

937 # of the chaining, or constraints in the configuration of the 

938 # datastore itself. 

939 if any( 939 ↛ 944line 939 didn't jump to line 944 because the condition on line 939 was never true

940 (constraints is not None and not constraints.isAcceptable(ref)) 

941 or not datastore.constraints.isAcceptable(ref) 

942 for ref in refs 

943 ): 

944 log.debug("Datastore %s skipping zip ingest due to constraints", datastore.name) 

945 continue 

946 try: 

947 datastore.ingest_zip(zip_path, transfer=transfer, dry_run=dry_run) 

948 except NotImplementedError: 

949 continue 

950 except Exception as e: 

951 final_exception = e 

952 else: 

953 n_success += 1 

954 

955 if n_success: 955 ↛ 957line 955 didn't jump to line 957 because the condition on line 955 was always true

956 return 

957 if final_exception: 

958 raise final_exception 

959 raise RuntimeError("Ingest was not successful in any datastores.") 

960 

961 def remove(self, ref: DatasetRef) -> None: 

962 """Indicate to the datastore that a dataset can be removed. 

963 

964 The dataset will be removed from each datastore. The dataset is 

965 not required to exist in every child datastore. 

966 

967 Parameters 

968 ---------- 

969 ref : `DatasetRef` 

970 Reference to the required dataset. 

971 

972 Raises 

973 ------ 

974 FileNotFoundError 

975 Attempt to remove a dataset that does not exist. Raised if none 

976 of the child datastores removed the dataset. 

977 """ 

978 log.debug("Removing %s", ref) 

979 self.trash(ref, ignore_errors=False) 

980 self.emptyTrash(ignore_errors=False, refs=[ref]) 

981 

982 def forget(self, refs: Iterable[DatasetRef]) -> None: 

983 for datastore in tuple(self.datastores): 

984 datastore.forget(refs) 

985 

986 def trash(self, ref: DatasetRef | Iterable[DatasetRef], ignore_errors: bool = True) -> None: 

987 if isinstance(ref, DatasetRef): 

988 ref_label = str(ref) 

989 else: 

990 ref_label = "bulk datasets" 

991 

992 log.debug("Trashing %s", ref_label) 

993 

994 counter = 0 

995 for datastore in self.datastores: 

996 try: 

997 datastore.trash(ref, ignore_errors=ignore_errors) 

998 counter += 1 

999 except FileNotFoundError: 

1000 pass 

1001 

1002 if counter == 0: 

1003 err_msg = f"Could not mark for removal from any child datastore: {ref_label}" 

1004 if ignore_errors: 1004 ↛ 1005line 1004 didn't jump to line 1005 because the condition on line 1004 was never true

1005 log.warning(err_msg) 

1006 else: 

1007 raise FileNotFoundError(err_msg) 

1008 

1009 def emptyTrash( 

1010 self, ignore_errors: bool = True, refs: Collection[DatasetRef] | None = None, dry_run: bool = False 

1011 ) -> set[ResourcePath]: 

1012 removed = set() 

1013 for datastore in self.datastores: 

1014 removed.update(datastore.emptyTrash(ignore_errors=ignore_errors, refs=refs, dry_run=dry_run)) 

1015 return removed 

1016 

1017 def transfer(self, inputDatastore: Datastore, ref: DatasetRef) -> None: 

1018 """Retrieve a dataset from an input `Datastore`, 

1019 and store the result in this `Datastore`. 

1020 

1021 Parameters 

1022 ---------- 

1023 inputDatastore : `Datastore` 

1024 The external `Datastore` from which to retrieve the Dataset. 

1025 ref : `DatasetRef` 

1026 Reference to the required dataset in the input data store. 

1027 

1028 Returns 

1029 ------- 

1030 results : `list` 

1031 List containing the return value from the ``put()`` to each 

1032 child datastore. 

1033 """ 

1034 assert inputDatastore is not self # unless we want it for renames? 

1035 inMemoryDataset = inputDatastore.get(ref) 

1036 self.put(inMemoryDataset, ref) 

1037 

1038 def validateConfiguration( 

1039 self, entities: Iterable[DatasetRef | DatasetType | StorageClass], logFailures: bool = False 

1040 ) -> None: 

1041 """Validate some of the configuration for this datastore. 

1042 

1043 Parameters 

1044 ---------- 

1045 entities : `~collections.abc.Iterable` [`DatasetRef` | `DatasetType` \ 

1046 | `StorageClass`] 

1047 Entities to test against this configuration. Can be differing 

1048 types. 

1049 logFailures : `bool`, optional 

1050 If `True`, output a log message for every validation error 

1051 detected. 

1052 

1053 Returns 

1054 ------- 

1055 None 

1056 

1057 Raises 

1058 ------ 

1059 DatastoreValidationError 

1060 Raised if there is a validation problem with a configuration. 

1061 All the problems are reported in a single exception. 

1062 

1063 Notes 

1064 ----- 

1065 This method checks each datastore in turn. 

1066 """ 

1067 # Need to catch each of the datastore outputs and ensure that 

1068 # all are tested. 

1069 failures = [] 

1070 for datastore in self.datastores: 

1071 try: 

1072 datastore.validateConfiguration(entities, logFailures=logFailures) 

1073 except DatastoreValidationError as e: 

1074 if logFailures: 1074 ↛ 1076line 1074 didn't jump to line 1076 because the condition on line 1074 was always true

1075 log.critical("Datastore %s failed validation", datastore.name) 

1076 failures.append(f"Datastore {self.name}: {e}") 

1077 

1078 if failures: 

1079 msg = ";\n".join(failures) 

1080 raise DatastoreValidationError(msg) 

1081 

1082 def validateKey(self, lookupKey: LookupKey, entity: DatasetRef | DatasetType | StorageClass) -> None: 

1083 # Docstring is inherited from base class 

1084 failures = [] 

1085 for datastore in self.datastores: 

1086 try: 

1087 datastore.validateKey(lookupKey, entity) 

1088 except DatastoreValidationError as e: 

1089 failures.append(f"Datastore {self.name}: {e}") 

1090 

1091 if failures: 

1092 msg = ";\n".join(failures) 

1093 raise DatastoreValidationError(msg) 

1094 

1095 def getLookupKeys(self) -> set[LookupKey]: 

1096 # Docstring is inherited from base class 

1097 keys = set() 

1098 for datastore in self.datastores: 

1099 keys.update(datastore.getLookupKeys()) 

1100 

1101 keys.update(self.constraints.getLookupKeys()) 

1102 for p in self.datastoreConstraints: 

1103 if p is not None: 1103 ↛ 1102line 1103 didn't jump to line 1102 because the condition on line 1103 was always true

1104 keys.update(p.getLookupKeys()) 

1105 

1106 return keys 

1107 

1108 def needs_expanded_data_ids( 

1109 self, 

1110 transfer: str | None, 

1111 entity: DatasetRef | DatasetType | StorageClass | None = None, 

1112 ) -> bool: 

1113 # Docstring inherited. 

1114 # We can't safely use `self.datastoreConstraints` with `entity` to 

1115 # check whether a child datastore would even want to ingest this 

1116 # dataset, because we don't want to filter out datastores that might 

1117 # need an expanded data ID based in incomplete information (e.g. we 

1118 # pass a StorageClass, but the constraint dispatches on DatasetType). 

1119 # So we pessimistically check if any datastore would need an expanded 

1120 # data ID for this transfer mode. 

1121 return any(datastore.needs_expanded_data_ids(transfer, entity) for datastore in self.datastores) 

1122 

1123 def import_records(self, data: Mapping[str, DatastoreRecordData]) -> None: 

1124 # Docstring inherited from the base class. 

1125 

1126 for datastore in self.datastores: 

1127 datastore.import_records(data) 

1128 

1129 def export_records(self, refs: Iterable[DatasetIdRef]) -> Mapping[str, DatastoreRecordData]: 

1130 # Docstring inherited from the base class. 

1131 

1132 all_records: dict[str, DatastoreRecordData] = {} 

1133 

1134 # Merge all sub-datastore records into one structure 

1135 for datastore in self.datastores: 

1136 sub_records = datastore.export_records(refs) 

1137 for name, record_data in sub_records.items(): 

1138 # All datastore names must be unique in a chain. 

1139 if name in all_records: 1139 ↛ 1140line 1139 didn't jump to line 1140 because the condition on line 1139 was never true

1140 raise ValueError("Non-unique datastore name found in datastore {datastore}") 

1141 all_records[name] = record_data 

1142 

1143 return all_records 

1144 

1145 def export_table(self, datasets: Collection[DatasetId]) -> DatastoreRecordTable: 

1146 # Docstring inherited from the base class. 

1147 

1148 return DatastoreRecordTable.combine( 

1149 [datastore.export_table(datasets) for datastore in self.datastores] 

1150 ) 

1151 

1152 def import_table(self, table: DatastoreRecordTable) -> None: 

1153 # Docstring inherited from the base class. 

1154 

1155 table.validate_datastore_names(self.names) 

1156 for datastore in self.datastores: 

1157 datastore.import_table(table.filter_by_datastore_name(datastore.name)) 

1158 

1159 def export_predicted_records(self, refs: Iterable[DatasetRef]) -> dict[str, DatastoreRecordData]: 

1160 # Docstring inherited from the base class. 

1161 

1162 all_records: dict[str, DatastoreRecordData] = {} 

1163 

1164 # Filter out datasets that this datastore is not allowed to contain. 

1165 refs = [ref for ref in refs if self.constraints.isAcceptable(ref)] 

1166 

1167 # Merge all sub-datastore records into one structure 

1168 for datastore in self.datastores: 

1169 sub_records = datastore.export_predicted_records(refs) 

1170 for name, record_data in sub_records.items(): 

1171 # All datastore names must be unique in a chain. 

1172 if name in all_records: 1172 ↛ 1173line 1172 didn't jump to line 1173 because the condition on line 1172 was never true

1173 raise ValueError("Non-unique datastore name found in datastore {datastore}") 

1174 all_records[name] = record_data 

1175 

1176 return all_records 

1177 

1178 def export( 

1179 self, 

1180 refs: Iterable[DatasetRef], 

1181 *, 

1182 directory: ResourcePathExpression | None = None, 

1183 transfer: str | None = "auto", 

1184 ) -> Iterable[FileDataset]: 

1185 # Docstring inherited from Datastore.export. 

1186 if transfer == "auto" and directory is None: 

1187 transfer = None 

1188 

1189 if transfer is not None and directory is None: 

1190 raise TypeError(f"Cannot export using transfer mode {transfer} with no export directory given") 

1191 

1192 if transfer == "move": 

1193 raise TypeError("Can not export by moving files out of datastore.") 

1194 

1195 # Exporting from a chain has the potential for a dataset to be 

1196 # in one or more of the datastores in the chain. We only need one 

1197 # of them since we assume the datasets are the same in all (but 

1198 # the file format could be different of course since that is a 

1199 # per-datastore configuration). 

1200 # We also do not know whether any of the datastores in the chain 

1201 # support file export. 

1202 

1203 # Ensure we have an ordered sequence that is not an iterator or set. 

1204 if not isinstance(refs, Sequence): 

1205 refs = list(refs) 

1206 

1207 # If any of the datasets are missing entirely we need to raise early 

1208 # before we try to run the export. This can be a little messy but is 

1209 # better than exporting files from the first datastore and then finding 

1210 # that one is missing but is not in the second datastore either. 

1211 known = [datastore.knows_these(refs) for datastore in self.datastores] 

1212 refs_known: set[DatasetRef] = set() 

1213 for known_to_this in known: 

1214 refs_known.update({ref for ref, knows_this in known_to_this.items() if knows_this}) 

1215 missing_count = len(refs) - len(refs_known) 

1216 if missing_count: 

1217 raise FileNotFoundError(f"Not all datasets known to this datastore. Missing {missing_count}") 

1218 

1219 # To allow us to slot each result into the right place after 

1220 # asking each datastore, create a dict with the index. 

1221 ref_positions = {ref: i for i, ref in enumerate(refs)} 

1222 

1223 # Presize the final export list. 

1224 exported: list[FileDataset | None] = [None] * len(refs) 

1225 

1226 # The order of the returned dataset has to match the order of the 

1227 # given refs, even if they are all from different datastores. 

1228 for i, datastore in enumerate(self.datastores): 

1229 known_to_this = known[i] 

1230 filtered = [ref for ref, knows in known_to_this.items() if knows and ref in ref_positions] 

1231 

1232 try: 

1233 this_export = datastore.export(filtered, directory=directory, transfer=transfer) 

1234 except NotImplementedError: 

1235 # Try the next datastore. 

1236 continue 

1237 

1238 for ref, export in zip(filtered, this_export, strict=True): 

1239 # Get the position and also delete it from the list. 

1240 exported[ref_positions.pop(ref)] = export 

1241 

1242 # Every dataset should be accounted for because of the earlier checks 

1243 # but make sure that we did fill all the slots to appease mypy. 

1244 for i, dataset in enumerate(exported): 

1245 if dataset is None: 1245 ↛ 1246line 1245 didn't jump to line 1246 because the condition on line 1245 was never true

1246 raise FileNotFoundError(f"Failed to export dataset {refs[i]}.") 

1247 yield dataset 

1248 

1249 def get_file_info_for_transfer(self, dataset_ids: Iterable[DatasetId]) -> FileTransferMap: 

1250 unassigned_ids = set(dataset_ids) 

1251 output: FileTransferMap = {} 

1252 found_acceptable_datastore = False 

1253 for datastore in self.datastores: 

1254 try: 

1255 found = datastore.get_file_info_for_transfer(unassigned_ids) 

1256 found_acceptable_datastore = True 

1257 output.update(found) 

1258 unassigned_ids -= found.keys() 

1259 except NotImplementedError: 

1260 pass 

1261 

1262 if not found_acceptable_datastore: 

1263 types = {get_full_type_name(d) for d in self.datastores} 

1264 raise TypeError( 

1265 "ChainedDatastore had no datastores able to provide file transfer information." 

1266 f" Had {','.join(types)}" 

1267 ) 

1268 

1269 return output 

1270 

1271 def locate_missing_files_for_transfer( 

1272 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool] 

1273 ) -> FileTransferMap: 

1274 missing_refs = {ref.id: ref for ref in refs} 

1275 output: FileTransferMap = {} 

1276 for datastore in self.datastores: 

1277 # Have to check each datastore in turn. If we do not do 

1278 # this warnings will be issued further down for datasets 

1279 # that are in one and not the other. The existence cache 

1280 # will prevent repeat checks. 

1281 

1282 found = datastore.locate_missing_files_for_transfer(missing_refs.values(), artifact_existence) 

1283 output.update(found) 

1284 for id in found.keys(): 

1285 missing_refs.pop(id) 

1286 log.debug("Adding %d missing refs to list for transfer from %s", len(found), datastore.name) 

1287 

1288 return output 

1289 

1290 def transfer_from( 

1291 self, 

1292 source_records: FileTransferMap, 

1293 refs: Collection[DatasetRef], 

1294 transfer: str = "auto", 

1295 artifact_existence: dict[ResourcePath, bool] | None = None, 

1296 dry_run: bool = False, 

1297 ) -> tuple[set[DatasetRef], set[DatasetRef]]: 

1298 # Docstring inherited 

1299 if not refs: 1299 ↛ 1300line 1299 didn't jump to line 1300 because the condition on line 1299 was never true

1300 return set(), set() 

1301 

1302 # Assume that each child datastore knows how to look inside a chained 

1303 # datastore for compatible datastores (and so there is no need to 

1304 # unpack the source datastores here). 

1305 # Need to decide if a ref accepted by one datastore should be sent to 

1306 # later datastores (as is done in put()). More efficient to filter out 

1307 # accepted datasets. 

1308 if artifact_existence is None: 

1309 artifact_existence = {} 

1310 available_refs = set(refs) 

1311 accepted: set[DatasetRef] = set() 

1312 rejected: set[DatasetRef] = set() 

1313 nsuccess = 0 

1314 

1315 log.debug("Initiating transfer to chained datastore %s", self.name) 

1316 for datastore in self.datastores: 

1317 # Rejections from this datastore might be acceptances in the next. 

1318 # We add them all up but then recalculate at the end. 

1319 if not available_refs: 

1320 break 

1321 log.verbose("Transferring %d datasets to %s from chain", len(available_refs), datastore.name) 

1322 

1323 try: 

1324 current_accepted, current_rejected = datastore.transfer_from( 

1325 source_records, 

1326 available_refs, 

1327 transfer=transfer, 

1328 artifact_existence=artifact_existence, 

1329 dry_run=dry_run, 

1330 ) 

1331 except (TypeError, NotImplementedError): 

1332 # The datastores were incompatible. 

1333 continue 

1334 else: 

1335 nsuccess += 1 

1336 

1337 # Do not send accepted refs to later datastores. 

1338 available_refs -= current_accepted 

1339 

1340 accepted.update(current_accepted) 

1341 rejected.update(current_rejected) 

1342 

1343 if nsuccess == 0: 

1344 raise TypeError("None of the child datastores could accept file transfers") 

1345 

1346 # It's not rejected if some other datastore accepted it. 

1347 rejected -= accepted 

1348 log.verbose( 

1349 "Finished transfer_from to %s with %d accepted, %d rejected from %d requested.", 

1350 self.name, 

1351 len(accepted), 

1352 len(rejected), 

1353 len(refs), 

1354 ) 

1355 

1356 return accepted, rejected 

1357 

1358 def get_opaque_table_definitions(self) -> Mapping[str, DatastoreOpaqueTable]: 

1359 # Docstring inherited from the base class. 

1360 tables: dict[str, DatastoreOpaqueTable] = {} 

1361 for datastore in self.datastores: 

1362 tables.update(datastore.get_opaque_table_definitions()) 

1363 return tables 

1364 

1365 def set_retrieve_dataset_type_method(self, method: Callable[[str], DatasetType | None] | None) -> None: 

1366 # Docstring inherited from the base class. 

1367 for datastore in self.datastores: 

1368 datastore.set_retrieve_dataset_type_method(method)