Coverage for python/lsst/daf/butler/datastores/fileDatastore.py: 86%

1079 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-11-06 01:43 -0800

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"""Generic file-based datastore code.""" 

29 

30from __future__ import annotations 

31 

32__all__ = ("FileDatastore",) 

33 

34import contextlib 

35import hashlib 

36import logging 

37import math 

38from collections import defaultdict 

39from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence 

40from typing import TYPE_CHECKING, Any, ClassVar, cast 

41from uuid import UUID 

42 

43from sqlalchemy import BigInteger, String 

44 

45from lsst.daf.butler import ( 

46 Config, 

47 DatasetId, 

48 DatasetRef, 

49 DatasetType, 

50 DatasetTypeNotSupportedError, 

51 FileDataset, 

52 FileDescriptor, 

53 Formatter, 

54 FormatterFactory, 

55 FormatterV1inV2, 

56 FormatterV2, 

57 Location, 

58 LocationFactory, 

59 Progress, 

60 StorageClass, 

61 ddl, 

62) 

63from lsst.daf.butler.datastore import ( 

64 DatasetRefURIs, 

65 Datastore, 

66 DatastoreConfig, 

67 DatastoreOpaqueTable, 

68 DatastoreValidationError, 

69 DuplicateIdError, 

70) 

71from lsst.daf.butler.datastore.cache_manager import ( 

72 AbstractDatastoreCacheManager, 

73 DatastoreCacheManager, 

74 DatastoreDisabledCacheManager, 

75) 

76from lsst.daf.butler.datastore.composites import CompositesMap 

77from lsst.daf.butler.datastore.file_templates import FileTemplates, FileTemplateValidationError 

78from lsst.daf.butler.datastore.generic_base import GenericBaseDatastore 

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

80from lsst.daf.butler.datastore.stored_file_info import ( 

81 StoredDatastoreItemInfo, 

82 StoredFileInfo, 

83 StoredFileInfoTable, 

84) 

85from lsst.daf.butler.datastores.file_datastore.get import ( 

86 DatasetLocationInformation, 

87 DatastoreFileGetInformation, 

88 generate_datastore_get_information, 

89 get_dataset_as_python_object_from_get_info, 

90) 

91from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ( 

92 ArtifactIndexInfo, 

93 ZipIndex, 

94 determine_destination_for_retrieved_artifact, 

95 unpack_zips, 

96) 

97from lsst.daf.butler.registry.interfaces import ( 

98 DatabaseInsertMode, 

99 DatastoreRegistryBridge, 

100 FakeDatasetRef, 

101 ReadOnlyDatabaseError, 

102) 

103from lsst.daf.butler.repo_relocation import replaceRoot 

104from lsst.daf.butler.utils import transactional 

105from lsst.resources import ResourcePath, ResourcePathExpression 

106from lsst.utils.introspection import get_class_of, get_full_type_name 

107from lsst.utils.iteration import chunk_iterable 

108 

109# For VERBOSE logging usage. 

110from lsst.utils.logging import VERBOSE, getLogger 

111from lsst.utils.timer import time_this 

112 

113from ..datastore import FileTransferMap, FileTransferRecord 

114from .utils import report_duplicate_ids 

115 

116if TYPE_CHECKING: 

117 from lsst.daf.butler import DatasetProvenance, LookupKey 

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

119 

120log = getLogger(__name__) 

121 

122 

123class _IngestPrepData(Datastore.IngestPrepData): 

124 """Helper class for FileDatastore ingest implementation. 

125 

126 Parameters 

127 ---------- 

128 datasets : `~collections.abc.Iterable` of `FileDataset` 

129 Files to be ingested by this datastore. 

130 """ 

131 

132 def __init__(self, datasets: Iterable[FileDataset]): 

133 super().__init__(ref for dataset in datasets for ref in dataset.refs) 

134 self.datasets = datasets 

135 

136 

137class FileDatastore(GenericBaseDatastore[StoredFileInfo]): 

138 """Generic Datastore for file-based implementations. 

139 

140 Should always be sub-classed since key abstract methods are missing. 

141 

142 Parameters 

143 ---------- 

144 config : `DatastoreConfig` or `str` 

145 Configuration as either a `Config` object or URI to file. 

146 bridgeManager : `DatastoreRegistryBridgeManager` 

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

148 root : `lsst.resources.ResourcePath` 

149 Root directory URI of this `Datastore`. 

150 formatterFactory : `FormatterFactory` 

151 Factory for creating instances of formatters. 

152 templates : `FileTemplates` 

153 File templates that can be used by this `Datastore`. 

154 composites : `CompositesMap` 

155 Determines whether a dataset should be disassembled on put. 

156 trustGetRequest : `bool` 

157 Determine whether we can fall back to configuration if a requested 

158 dataset is not known to registry. 

159 

160 Raises 

161 ------ 

162 ValueError 

163 If root location does not exist and ``create`` is `False` in the 

164 configuration. 

165 """ 

166 

167 defaultConfigFile: ClassVar[str | None] = None 

168 """Path to configuration defaults. Accessed within the ``config`` resource 

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

170 """ 

171 

172 root: ResourcePath 

173 """Root directory URI of this `Datastore`.""" 

174 

175 locationFactory: LocationFactory 

176 """Factory for creating locations relative to the datastore root.""" 

177 

178 formatterFactory: FormatterFactory 

179 """Factory for creating instances of formatters.""" 

180 

181 templates: FileTemplates 

182 """File templates that can be used by this `Datastore`.""" 

183 

184 composites: CompositesMap 

185 """Determines whether a dataset should be disassembled on put.""" 

186 

187 defaultConfigFile = "datastores/fileDatastore.yaml" 

188 """Path to configuration defaults. Accessed within the ``config`` resource 

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

190 """ 

191 

192 _retrieve_dataset_method: Callable[[str], DatasetType | None] | None = None 

193 """Callable that is used in trusted mode to retrieve registry definition 

194 of a named dataset type. 

195 """ 

196 

197 @classmethod 

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

199 """Set any filesystem-dependent config options for this Datastore to 

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

201 

202 Parameters 

203 ---------- 

204 root : `str` 

205 URI to the root of the data repository. 

206 config : `Config` 

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

208 this component will be updated. Will not expand 

209 defaults. 

210 full : `Config` 

211 A complete config with all defaults expanded that can be 

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

213 modified by this method. 

214 Repository-specific options that should not be obtained 

215 from defaults when Butler instances are constructed 

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

217 overwrite : `bool`, optional 

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

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

220 ``root``. 

221 

222 Notes 

223 ----- 

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

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

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

227 """ 

228 Config.updateParameters( 

229 DatastoreConfig, 

230 config, 

231 full, 

232 toUpdate={"root": root}, 

233 toCopy=("cls", ("records", "table")), 

234 overwrite=overwrite, 

235 ) 

236 

237 @classmethod 

238 def makeTableSpec(cls) -> ddl.TableSpec: 

239 return ddl.TableSpec( 

240 fields=[ 

241 ddl.FieldSpec(name="dataset_id", dtype=ddl.GUID, primaryKey=True), 

242 ddl.FieldSpec(name="path", dtype=String, length=256, nullable=False), 

243 ddl.FieldSpec(name="formatter", dtype=String, length=128, nullable=False), 

244 ddl.FieldSpec(name="storage_class", dtype=String, length=64, nullable=False), 

245 # Use empty string to indicate no component 

246 ddl.FieldSpec(name="component", dtype=String, length=32, primaryKey=True), 

247 # TODO: should checksum be Base64Bytes instead? 

248 ddl.FieldSpec(name="checksum", dtype=String, length=128, nullable=True), 

249 ddl.FieldSpec(name="file_size", dtype=BigInteger, nullable=True), 

250 ], 

251 unique=frozenset(), 

252 indexes=[ddl.IndexSpec("path")], 

253 ) 

254 

255 def __init__( 

256 self, 

257 config: DatastoreConfig, 

258 bridgeManager: DatastoreRegistryBridgeManager, 

259 root: ResourcePath, 

260 formatterFactory: FormatterFactory, 

261 templates: FileTemplates, 

262 composites: CompositesMap, 

263 trustGetRequest: bool, 

264 ): 

265 super().__init__(config, bridgeManager) 

266 self.root = ResourcePath(root) 

267 self.formatterFactory = formatterFactory 

268 self.templates = templates 

269 self.composites = composites 

270 self.trustGetRequest = trustGetRequest 

271 

272 # Name ourselves either using an explicit name or a name 

273 # derived from the (unexpanded) root 

274 if "name" in self.config: 

275 self.name = self.config["name"] 

276 else: 

277 # We use the unexpanded root in the name to indicate that this 

278 # datastore can be moved without having to update registry. 

279 self.name = "{}@{}".format(type(self).__name__, self.config["root"]) 

280 

281 self.locationFactory = LocationFactory(self.root) 

282 

283 self._opaque_table_name = self.config["records", "table"] 

284 try: 

285 # Storage of paths and formatters, keyed by dataset_id 

286 self._table = bridgeManager.opaque.register(self._opaque_table_name, self.makeTableSpec()) 

287 # Interface to Registry. 

288 self._bridge = bridgeManager.register(self.name) 

289 except ReadOnlyDatabaseError: 

290 # If the database is read only and we just tried and failed to 

291 # create a table, it means someone is trying to create a read-only 

292 # butler client for an empty repo. That should be okay, as long 

293 # as they then try to get any datasets before some other client 

294 # creates the table. Chances are they're just validating 

295 # configuration. 

296 pass 

297 

298 # Determine whether checksums should be used - default to False 

299 self.useChecksum = self.config.get("checksum", False) 

300 

301 # Create a cache manager 

302 self.cacheManager: AbstractDatastoreCacheManager 

303 if "cached" in self.config: 303 ↛ 306line 303 didn't jump to line 306 because the condition on line 303 was always true

304 self.cacheManager = DatastoreCacheManager(self.config["cached"], universe=bridgeManager.universe) 

305 else: 

306 self.cacheManager = DatastoreDisabledCacheManager("", universe=bridgeManager.universe) 

307 

308 self.universe = bridgeManager.universe 

309 

310 @classmethod 

311 def _create_from_config( 

312 cls, 

313 config: DatastoreConfig, 

314 bridgeManager: DatastoreRegistryBridgeManager, 

315 butlerRoot: ResourcePathExpression | None, 

316 ) -> FileDatastore: 

317 if "root" not in config: 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true

318 raise ValueError("No root directory specified in configuration") 

319 

320 # Support repository relocation in config 

321 # Existence of self.root is checked in subclass 

322 root = ResourcePath(replaceRoot(config["root"], butlerRoot), forceDirectory=True, forceAbsolute=True) 

323 

324 # Now associate formatters with storage classes 

325 formatterFactory = FormatterFactory() 

326 formatterFactory.registerFormatters(config["formatters"], universe=bridgeManager.universe) 

327 

328 # Read the file naming templates 

329 templates = FileTemplates(config["templates"], universe=bridgeManager.universe) 

330 

331 # See if composites should be disassembled 

332 composites = CompositesMap(config["composites"], universe=bridgeManager.universe) 

333 

334 # Determine whether we can fall back to configuration if a 

335 # requested dataset is not known to registry 

336 trustGetRequest = config.get("trust_get_request", False) 

337 

338 self = FileDatastore( 

339 config, bridgeManager, root, formatterFactory, templates, composites, trustGetRequest 

340 ) 

341 

342 # Check existence and create directory structure if necessary. 

343 # 

344 # The concept of a 'root directory' is problematic for some resource 

345 # path types that don't necessarily support the concept of a directory 

346 # (http, s3, gs... basically anything that isn't a local filesystem or 

347 # WebDAV.) 

348 # On these resource paths an object representing the 

349 # "root" directory may not exist even though files under the root do, 

350 # and in a read-only repository we will be unable to create it. 

351 # So we only immediately verify the root for local filesystems, 

352 # the only case where this check will definitely not give a false 

353 # negative. 

354 if self.root.isLocal and not self.root.exists(): 

355 if "create" not in self.config or not self.config["create"]: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true

356 raise ValueError(f"No valid root and not allowed to create one at: {self.root}") 

357 try: 

358 self.root.mkdir() 

359 except Exception as e: 

360 raise ValueError( 

361 f"Can not create datastore root '{self.root}', check permissions. Got error: {e}" 

362 ) from e 

363 

364 return self 

365 

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

367 return FileDatastore( 

368 self.config, 

369 bridgeManager, 

370 self.root, 

371 self.formatterFactory, 

372 self.templates, 

373 self.composites, 

374 self.trustGetRequest, 

375 ) 

376 

377 def __str__(self) -> str: 

378 return str(self.root) 

379 

380 @property 

381 def bridge(self) -> DatastoreRegistryBridge: 

382 return self._bridge 

383 

384 @property 

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

386 # Docstring inherited. 

387 return {self.name: self.root} 

388 

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

390 self.trustGetRequest = mode 

391 

392 def _artifact_exists(self, location: Location) -> bool: 

393 """Check that an artifact exists in this datastore at the specified 

394 location. 

395 

396 Parameters 

397 ---------- 

398 location : `Location` 

399 Expected location of the artifact associated with this datastore. 

400 

401 Returns 

402 ------- 

403 exists : `bool` 

404 True if the location can be found, false otherwise. 

405 """ 

406 log.debug("Checking if resource exists: %s", location.uri) 

407 return location.uri.exists() 

408 

409 def addStoredItemInfo( 

410 self, 

411 refs: Iterable[DatasetRef], 

412 infos: Iterable[StoredFileInfo], 

413 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT, 

414 ) -> None: 

415 """Record internal storage information associated with one or more 

416 datasets. 

417 

418 Parameters 

419 ---------- 

420 refs : sequence of `DatasetRef` 

421 The datasets that have been stored. 

422 infos : sequence of `StoredDatastoreItemInfo` 

423 Metadata associated with the stored datasets. 

424 insert_mode : `~lsst.daf.butler.registry.interfaces.DatabaseInsertMode` 

425 Mode to use to insert the new records into the table. The 

426 options are ``INSERT`` (error if pre-existing), ``REPLACE`` 

427 (replace content with new values), and ``ENSURE`` (skip if the row 

428 already exists). 

429 """ 

430 records = [ 

431 info.rebase(ref).to_record(dataset_id=ref.id) for ref, info in zip(refs, infos, strict=True) 

432 ] 

433 match insert_mode: 

434 case DatabaseInsertMode.INSERT: 

435 self._table.insert(*records, transaction=self._transaction) 

436 case DatabaseInsertMode.ENSURE: 436 ↛ 437line 436 didn't jump to line 437 because the pattern on line 436 never matched

437 self._table.ensure(*records, transaction=self._transaction) 

438 case DatabaseInsertMode.REPLACE: 438 ↛ 440line 438 didn't jump to line 440 because the pattern on line 438 always matched

439 self._table.replace(*records, transaction=self._transaction) 

440 case _: 

441 raise ValueError(f"Unknown insert mode of '{insert_mode}'") 

442 

443 def getStoredItemsInfo( 

444 self, ref: DatasetIdRef, ignore_datastore_records: bool = False 

445 ) -> list[StoredFileInfo]: 

446 """Retrieve information associated with files stored in this 

447 `Datastore` associated with this dataset ref. 

448 

449 Parameters 

450 ---------- 

451 ref : `DatasetRef` 

452 The dataset that is to be queried. 

453 ignore_datastore_records : `bool` 

454 If `True` then do not use datastore records stored in refs. 

455 

456 Returns 

457 ------- 

458 items : `~collections.abc.Iterable` [`StoredDatastoreItemInfo`] 

459 Stored information about the files and associated formatters 

460 associated with this dataset. Only one file will be returned 

461 if the dataset has not been disassembled. Can return an empty 

462 list if no matching datasets can be found. 

463 """ 

464 # Try to get them from the ref first. 

465 if ref._datastore_records is not None and not ignore_datastore_records: 

466 ref_records = ref._datastore_records.get(self._table.name, []) 

467 # Need to make sure they have correct type. 

468 for record in ref_records: 

469 if not isinstance(record, StoredFileInfo): 469 ↛ 470line 469 didn't jump to line 470 because the condition on line 469 was never true

470 raise TypeError(f"Datastore record has unexpected type {record.__class__.__name__}") 

471 return cast(list[StoredFileInfo], ref_records) 

472 

473 # Look for the dataset_id -- there might be multiple matches 

474 # if we have disassembled the dataset. 

475 records = self._table.fetch(dataset_id=ref.id) 

476 return [StoredFileInfo.from_record(record) for record in records] 

477 

478 def _register_datasets( 

479 self, 

480 refsAndInfos: Iterable[tuple[DatasetRef, StoredFileInfo]], 

481 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT, 

482 ) -> None: 

483 """Update registry to indicate that one or more datasets have been 

484 stored. 

485 

486 Parameters 

487 ---------- 

488 refsAndInfos : sequence `tuple` [`DatasetRef`, 

489 `StoredDatastoreItemInfo`] 

490 Datasets to register and the internal datastore metadata associated 

491 with them. 

492 insert_mode : `str`, optional 

493 Indicate whether the new records should be new ("insert", default), 

494 or allowed to exists ("ensure") or be replaced if already present 

495 ("replace"). 

496 """ 

497 expandedRefs: list[DatasetRef] = [] 

498 expandedItemInfos: list[StoredFileInfo] = [] 

499 

500 for ref, itemInfo in refsAndInfos: 

501 expandedRefs.append(ref) 

502 expandedItemInfos.append(itemInfo) 

503 

504 # Dataset location only cares about registry ID so if we have 

505 # disassembled in datastore we have to deduplicate. Since they 

506 # will have different datasetTypes we can't use a set 

507 registryRefs = {r.id: r for r in expandedRefs} 

508 if insert_mode == DatabaseInsertMode.INSERT: 

509 self.bridge.insert(registryRefs.values()) 

510 else: 

511 # There are only two columns and all that matters is the 

512 # dataset ID. 

513 self.bridge.ensure(registryRefs.values()) 

514 self.addStoredItemInfo(expandedRefs, expandedItemInfos, insert_mode=insert_mode) 

515 

516 def _get_stored_records_associated_with_refs( 

517 self, refs: Iterable[DatasetIdRef], ignore_datastore_records: bool = False 

518 ) -> dict[DatasetId, list[StoredFileInfo]]: 

519 """Retrieve all records associated with the provided refs. 

520 

521 Parameters 

522 ---------- 

523 refs : `~collections.abc.Iterable` of `DatasetIdRef` 

524 The refs for which records are to be retrieved. 

525 ignore_datastore_records : `bool` 

526 If `True` then do not use datastore records stored in refs. 

527 

528 Returns 

529 ------- 

530 records : `dict` of [`DatasetId`, `list` of `StoredFileInfo`] 

531 The matching records indexed by the ref ID. The number of entries 

532 in the dict can be smaller than the number of requested refs. 

533 """ 

534 # Check datastore records in refs first. 

535 records_by_ref: defaultdict[DatasetId, list[StoredFileInfo]] = defaultdict(list) 

536 refs_with_no_records = [] 

537 for ref in refs: 

538 if ignore_datastore_records or ref._datastore_records is None: 538 ↛ 541line 538 didn't jump to line 541 because the condition on line 538 was always true

539 refs_with_no_records.append(ref) 

540 else: 

541 if (ref_records := ref._datastore_records.get(self._table.name)) is not None: 

542 # Need to make sure they have correct type. 

543 for ref_record in ref_records: 

544 if not isinstance(ref_record, StoredFileInfo): 

545 raise TypeError( 

546 f"Datastore record has unexpected type {ref_record.__class__.__name__}" 

547 ) 

548 records_by_ref[ref.id].append(ref_record) 

549 

550 # If there were any refs without datastore records, check opaque table. 

551 records = self._table.fetch(dataset_id=[ref.id for ref in refs_with_no_records]) 

552 

553 # Uniqueness is dataset_id + component so can have multiple records 

554 # per ref. 

555 for record in records: 

556 records_by_ref[record["dataset_id"]].append(StoredFileInfo.from_record(record)) 

557 return records_by_ref 

558 

559 def _refs_associated_with_artifacts( 

560 self, paths: Iterable[str | ResourcePath] 

561 ) -> dict[str, set[DatasetId]]: 

562 """Return paths and associated dataset refs. 

563 

564 Parameters 

565 ---------- 

566 paths : `list` of `str` or `lsst.resources.ResourcePath` 

567 All the paths to include in search. These are exact matches 

568 to the entries in the records table and can include fragments. 

569 

570 Returns 

571 ------- 

572 mapping : `dict` of [`str`, `set` [`DatasetId`]] 

573 Mapping of each path to a set of associated database IDs. 

574 These are artifacts and so any fragments are stripped from the 

575 keys. 

576 """ 

577 # Group paths by those that have fragments and those that do not. 

578 with_fragment = set() 

579 without_fragment = set() 

580 for rpath in paths: 

581 spath = str(rpath) # Typing says can be ResourcePath so must force to string. 

582 if "#" in spath: 

583 spath, fragment = spath.rsplit("#", 1) 

584 with_fragment.add(spath) 

585 else: 

586 without_fragment.add(spath) 

587 

588 result: dict[str, set[DatasetId]] = defaultdict(set) 

589 if without_fragment: 

590 records = self._table.fetch(path=without_fragment) 

591 for row in records: 

592 path = row["path"] 

593 result[path].add(row["dataset_id"]) 

594 if with_fragment: 

595 # Do a query per prefix. 

596 for path in with_fragment: 

597 records = self._table.fetch(path=f"{path}#%") 

598 for row in records: 

599 # Need to strip fragments before adding to dict. 

600 row_path = row["path"] 

601 artifact_path = row_path[: row_path.rfind("#")] 

602 result[artifact_path].add(row["dataset_id"]) 

603 return result 

604 

605 def _registered_refs_per_artifact(self, pathInStore: ResourcePath) -> set[DatasetId]: 

606 """Return all dataset refs associated with the supplied path. 

607 

608 Parameters 

609 ---------- 

610 pathInStore : `lsst.resources.ResourcePath` 

611 Path of interest in the data store. 

612 

613 Returns 

614 ------- 

615 ids : `set` of `int` 

616 All `DatasetRef` IDs associated with this path. 

617 """ 

618 records = list(self._table.fetch(path=str(pathInStore))) 

619 ids = {r["dataset_id"] for r in records} 

620 return ids 

621 

622 def removeStoredItemInfo(self, ref: DatasetIdRef) -> None: 

623 """Remove information about the file associated with this dataset. 

624 

625 Parameters 

626 ---------- 

627 ref : `DatasetRef` 

628 The dataset that has been removed. 

629 """ 

630 # Note that this method is actually not used by this implementation, 

631 # we depend on bridge to delete opaque records. But there are some 

632 # tests that check that this method works, so we keep it for now. 

633 self._table.delete(["dataset_id"], {"dataset_id": ref.id}) 

634 

635 def _get_dataset_locations_info( 

636 self, ref: DatasetIdRef, ignore_datastore_records: bool = False 

637 ) -> list[DatasetLocationInformation]: 

638 r"""Find all the `Location`\ s of the requested dataset in the 

639 `Datastore` and the associated stored file information. 

640 

641 Parameters 

642 ---------- 

643 ref : `DatasetRef` 

644 Reference to the required `Dataset`. 

645 ignore_datastore_records : `bool` 

646 If `True` then do not use datastore records stored in refs. 

647 

648 Returns 

649 ------- 

650 results : `list` [`tuple` [`Location`, `StoredFileInfo` ]] 

651 Location of the dataset within the datastore and 

652 stored information about each file and its formatter. 

653 """ 

654 # Get the file information (this will fail if no file) 

655 records = self.getStoredItemsInfo(ref, ignore_datastore_records) 

656 

657 # Use the path to determine the location -- we need to take 

658 # into account absolute URIs in the datastore record 

659 return [(r.file_location(self.locationFactory), r) for r in records] 

660 

661 def _can_remove_dataset_artifact(self, ref: DatasetIdRef, location: Location) -> bool: 

662 """Check that there is only one dataset associated with the 

663 specified artifact. 

664 

665 Parameters 

666 ---------- 

667 ref : `DatasetRef` or `FakeDatasetRef` 

668 Dataset to be removed. 

669 location : `Location` 

670 The location of the artifact to be removed. 

671 

672 Returns 

673 ------- 

674 can_remove : `Bool` 

675 True if the artifact can be safely removed. 

676 """ 

677 # Can't ever delete absolute URIs. 

678 if location.pathInStore.isabs(): 

679 return False 

680 

681 # Get all entries associated with this path 

682 allRefs = self._registered_refs_per_artifact(location.pathInStore) 

683 if not allRefs: 

684 raise RuntimeError(f"Datastore inconsistency error. {location.pathInStore} not in registry") 

685 

686 # Remove these refs from all the refs and if there is nothing left 

687 # then we can delete 

688 remainingRefs = allRefs - {ref.id} 

689 

690 if remainingRefs: 

691 return False 

692 return True 

693 

694 def _get_expected_dataset_locations_info(self, ref: DatasetRef) -> list[tuple[Location, StoredFileInfo]]: 

695 """Predict the location and related file information of the requested 

696 dataset in this datastore. 

697 

698 Parameters 

699 ---------- 

700 ref : `DatasetRef` 

701 Reference to the required `Dataset`. 

702 

703 Returns 

704 ------- 

705 results : `list` [`tuple` [`Location`, `StoredFileInfo` ]] 

706 Expected Location of the dataset within the datastore and 

707 placeholder information about each file and its formatter. 

708 

709 Notes 

710 ----- 

711 Uses the current configuration to determine how we would expect the 

712 datastore files to have been written if we couldn't ask registry. 

713 This is safe so long as there has been no change to datastore 

714 configuration between writing the dataset and wanting to read it. 

715 Will not work for files that have been ingested without using the 

716 standard file template or default formatter. 

717 """ 

718 # If we have a component ref we always need to ask the questions 

719 # of the composite. If the composite is disassembled this routine 

720 # should return all components. If the composite was not 

721 # disassembled the composite is what is stored regardless of 

722 # component request. Note that if the caller has disassembled 

723 # a composite there is no way for this guess to know that 

724 # without trying both the composite and component ref and seeing 

725 # if there is something at the component Location even without 

726 # disassembly being enabled. 

727 if ref.datasetType.isComponent(): 727 ↛ 728line 727 didn't jump to line 728 because the condition on line 727 was never true

728 ref = ref.makeCompositeRef() 

729 

730 # See if the ref is a composite that should be disassembled 

731 doDisassembly = self.composites.shouldBeDisassembled(ref) 

732 

733 all_info: list[tuple[Location, Formatter | FormatterV2, StorageClass, str | None]] = [] 

734 

735 if doDisassembly: 

736 for component, componentStorage in ref.datasetType.storageClass.components.items(): 

737 compRef = ref.makeComponentRef(component) 

738 location, formatter = self._determine_put_formatter_location(compRef) 

739 all_info.append((location, formatter, componentStorage, component)) 

740 

741 else: 

742 # Always use the composite ref if no disassembly 

743 location, formatter = self._determine_put_formatter_location(ref) 

744 all_info.append((location, formatter, ref.datasetType.storageClass, None)) 

745 

746 # Convert the list of tuples to have StoredFileInfo as second element 

747 return [ 

748 ( 

749 location, 

750 StoredFileInfo( 

751 formatter=formatter, 

752 path=location.pathInStore.path, 

753 storageClass=storageClass, 

754 component=component, 

755 checksum=None, 

756 file_size=-1, 

757 ), 

758 ) 

759 for location, formatter, storageClass, component in all_info 

760 ] 

761 

762 def _prepare_for_direct_get( 

763 self, ref: DatasetRef, parameters: Mapping[str, Any] | None = None 

764 ) -> list[DatastoreFileGetInformation]: 

765 """Check parameters for ``get`` and obtain formatter and 

766 location. 

767 

768 Parameters 

769 ---------- 

770 ref : `DatasetRef` 

771 Reference to the required Dataset. 

772 parameters : `dict` 

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

774 a slice of the dataset to be loaded. 

775 

776 Returns 

777 ------- 

778 getInfo : `list` [`DatastoreFileGetInformation`] 

779 Parameters needed to retrieve each file. 

780 """ 

781 log.debug("Retrieve %s from %s with parameters %s", ref, self.name, parameters) 

782 

783 # The ref as supplied describes what the caller wants back, including 

784 # any component and read storage class override. Internally the 

785 # composite as defined in the repository is used: that is what a get 

786 # with no overrides would return and it is the ref given to the 

787 # Formatter. Using the registry storage class also resets the storage 

788 # class for trusted mode. 

789 registry_ref = self._cast_storage_class(ref.makeCompositeRef() if ref.isComponent() else ref) 

790 

791 # Get file metadata and internal metadata 

792 fileLocations = self._get_dataset_locations_info(registry_ref) 

793 if not fileLocations: 

794 if not self.trustGetRequest: 

795 raise FileNotFoundError(f"Could not retrieve dataset {ref}.") 

796 # Assume the dataset is where we think it should be 

797 fileLocations = self._get_expected_dataset_locations_info(registry_ref) 

798 

799 if len(fileLocations) > 1: 

800 # If trust is involved it is possible that there will be 

801 # components listed here that do not exist in the datastore. 

802 # Explicitly check for file artifact existence and filter out any 

803 # that are missing. 

804 if self.trustGetRequest: 

805 fileLocations = [loc for loc in fileLocations if loc[0].uri.exists()] 

806 

807 # For now complain only if we have no components at all. One 

808 # component is probably a problem but we can punt that to the 

809 # assembler. 

810 if not fileLocations: 

811 raise FileNotFoundError(f"None of the component files for dataset {ref} exist.") 

812 

813 return generate_datastore_get_information( 

814 fileLocations, 

815 registry_ref=registry_ref, 

816 read_ref=ref, 

817 parameters=parameters, 

818 ) 

819 

820 def _determine_put_formatter_location( 

821 self, ref: DatasetRef, provenance: DatasetProvenance | None = None 

822 ) -> tuple[Location, Formatter | FormatterV2]: 

823 """Calculate the formatter and output location to use for put. 

824 

825 Parameters 

826 ---------- 

827 ref : `DatasetRef` 

828 Reference to the associated Dataset. 

829 provenance : `DatasetProvenance` 

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

831 

832 Returns 

833 ------- 

834 location : `Location` 

835 The location to write the dataset. 

836 formatter : `Formatter` 

837 The `Formatter` to use to write the dataset. 

838 """ 

839 # Work out output file name 

840 try: 

841 template = self.templates.getTemplate(ref) 

842 except KeyError as e: 

843 raise DatasetTypeNotSupportedError(f"Unable to find template for {ref}") from e 

844 

845 # Validate the template to protect against filenames from different 

846 # dataIds returning the same and causing overwrite confusion. 

847 template.validateTemplate(ref) 

848 

849 location = self.locationFactory.fromPath(template.format(ref), trusted_path=True) 

850 

851 # Get the formatter based on the storage class 

852 storageClass = ref.datasetType.storageClass 

853 try: 

854 formatter = self.formatterFactory.getFormatter( 

855 ref, 

856 FileDescriptor(location, storageClass=storageClass, component=ref.datasetType.component()), 

857 dataId=ref.dataId, 

858 ref=ref, 

859 provenance=provenance, 

860 ) 

861 except KeyError as e: 

862 raise DatasetTypeNotSupportedError( 

863 f"Unable to find formatter for {ref} in datastore {self.name}" 

864 ) from e 

865 

866 # Now that we know the formatter, update the location 

867 location = formatter.make_updated_location(location) 

868 

869 return location, formatter 

870 

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

872 # Docstring inherited from base class 

873 if transfer != "auto": 

874 return transfer 

875 

876 # See if the paths are within the datastore or not 

877 inside = [self._pathInStore(d.path) is not None for d in datasets] 

878 

879 if all(inside): 

880 transfer = None 

881 elif not any(inside): 881 ↛ 890line 881 didn't jump to line 890 because the condition on line 881 was always true

882 # Allow ResourcePath to use its own knowledge 

883 transfer = "auto" 

884 else: 

885 # This can happen when importing from a datastore that 

886 # has had some datasets ingested using "direct" mode. 

887 # Also allow ResourcePath to sort it out but warn about it. 

888 # This can happen if you are importing from a datastore 

889 # that had some direct transfer datasets. 

890 log.warning( 

891 "Some datasets are inside the datastore and some are outside. Using 'split' " 

892 "transfer mode. This assumes that the files outside the datastore are " 

893 "still accessible to the new butler since they will not be copied into " 

894 "the target datastore." 

895 ) 

896 transfer = "split" 

897 

898 return transfer 

899 

900 def _pathInStore(self, path: ResourcePathExpression) -> str | None: 

901 """Return path relative to datastore root. 

902 

903 Parameters 

904 ---------- 

905 path : `lsst.resources.ResourcePathExpression` 

906 Path to dataset. Can be absolute URI. If relative assumed to 

907 be relative to the datastore. Returns path in datastore 

908 or raises an exception if the path it outside. 

909 

910 Returns 

911 ------- 

912 inStore : `str` 

913 Path relative to datastore root. Returns `None` if the file is 

914 outside the root. 

915 """ 

916 # Relative path will always be relative to datastore 

917 pathUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False) 

918 return pathUri.relative_to(self.root) 

919 

920 def _standardizeIngestPath( 

921 self, 

922 path: str | ResourcePath, 

923 *, 

924 transfer: str | None = None, 

925 check_existence: bool = False, 

926 ) -> str | ResourcePath: 

927 """Standardize the path of a to-be-ingested file. 

928 

929 Parameters 

930 ---------- 

931 path : `str` or `lsst.resources.ResourcePath` 

932 Path of a file to be ingested. This parameter is not expected 

933 to be all the types that can be used to construct a 

934 `~lsst.resources.ResourcePath`. 

935 transfer : `str`, optional 

936 How (and whether) the dataset should be added to the datastore. 

937 See `ingest` for details of transfer modes. 

938 This implementation is provided only so 

939 `NotImplementedError` can be raised if the mode is not supported; 

940 actual transfers are deferred to `_extractIngestInfo`. 

941 check_existence : `bool`, optional 

942 If `True` the existence of the file will be checked, otherwise 

943 no check will be made. 

944 

945 Returns 

946 ------- 

947 path : `str` or `lsst.resources.ResourcePath` 

948 New path in what the datastore considers standard form. If an 

949 absolute URI was given that will be returned unchanged. 

950 

951 Notes 

952 ----- 

953 Subclasses of `FileDatastore` can implement this method instead 

954 of `_prepIngest`. It should not modify the data repository or given 

955 file in any way. 

956 

957 Raises 

958 ------ 

959 NotImplementedError 

960 Raised if the datastore does not support the given transfer mode 

961 (including the case where ingest is not supported at all). 

962 """ 

963 if transfer not in (None, "direct", "split") + self.root.transferModes: 

964 raise NotImplementedError(f"Transfer mode {transfer} not supported.") 

965 

966 # A relative URI indicates relative to datastore root 

967 srcUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False) 

968 if not srcUri.isabs(): 

969 srcUri = self.root.join(path) 

970 

971 if check_existence and not srcUri.exists(): 

972 raise FileNotFoundError( 

973 f"Resource at {srcUri} does not exist; note that paths to ingest " 

974 f"are assumed to be relative to {self.root} unless they are absolute." 

975 ) 

976 

977 if transfer is None: 

978 relpath = srcUri.relative_to(self.root) 

979 if not relpath: 

980 raise RuntimeError( 

981 f"Transfer is none but source file ({srcUri}) is not within datastore ({self.root})" 

982 ) 

983 

984 # Return the relative path within the datastore for internal 

985 # transfer 

986 path = relpath 

987 

988 return path 

989 

990 def _extractIngestInfo( 

991 self, 

992 path: ResourcePathExpression, 

993 ref: DatasetRef, 

994 *, 

995 formatter: Formatter | FormatterV2 | type[Formatter | FormatterV2], 

996 transfer: str | None = None, 

997 record_validation_info: bool = True, 

998 ) -> StoredFileInfo: 

999 """Relocate (if necessary) and extract `StoredFileInfo` from a 

1000 to-be-ingested file. 

1001 

1002 Parameters 

1003 ---------- 

1004 path : `lsst.resources.ResourcePathExpression` 

1005 URI or path of a file to be ingested. 

1006 ref : `DatasetRef` 

1007 Reference for the dataset being ingested. Guaranteed to have 

1008 ``dataset_id not None`. 

1009 formatter : `type` or `Formatter` 

1010 `Formatter` subclass to use for this dataset or an instance. 

1011 transfer : `str`, optional 

1012 How (and whether) the dataset should be added to the datastore. 

1013 See `ingest` for details of transfer modes. 

1014 record_validation_info : `bool`, optional 

1015 If `True`, the default, the datastore can record validation 

1016 information associated with the file. If `False` the datastore 

1017 will not attempt to track any information such as checksums 

1018 or file sizes. This can be useful if such information is tracked 

1019 in an external system or if the file is to be compressed in place. 

1020 It is up to the datastore whether this parameter is relevant. 

1021 

1022 Returns 

1023 ------- 

1024 info : `StoredFileInfo` 

1025 Internal datastore record for this file. This will be inserted by 

1026 the caller; the `_extractIngestInfo` is only responsible for 

1027 creating and populating the struct. 

1028 

1029 Raises 

1030 ------ 

1031 FileNotFoundError 

1032 Raised if one of the given files does not exist. 

1033 FileExistsError 

1034 Raised if transfer is not `None` but the (internal) location the 

1035 file would be moved to is already occupied. 

1036 """ 

1037 if self._transaction is None: 1037 ↛ 1038line 1037 didn't jump to line 1038 because the condition on line 1037 was never true

1038 raise RuntimeError("Ingest called without transaction enabled") 

1039 

1040 # Create URI of the source path, do not need to force a relative 

1041 # path to absolute. 

1042 srcUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False) 

1043 

1044 # Track whether we have read the size of the source yet 

1045 have_sized = False 

1046 

1047 tgtLocation: Location | None 

1048 if transfer is None or transfer == "split": 

1049 # A relative path is assumed to be relative to the datastore 

1050 # in this context 

1051 if not srcUri.isabs(): 

1052 tgtLocation = self.locationFactory.fromPath(srcUri.ospath, trusted_path=False) 

1053 else: 

1054 # Work out the path in the datastore from an absolute URI 

1055 # This is required to be within the datastore. 

1056 pathInStore = srcUri.relative_to(self.root) 

1057 if pathInStore is None and transfer is None: 1057 ↛ 1058line 1057 didn't jump to line 1058 because the condition on line 1057 was never true

1058 raise RuntimeError( 

1059 f"Unexpectedly learned that {srcUri} is not within datastore {self.root}" 

1060 ) 

1061 if pathInStore: 1061 ↛ 1063line 1061 didn't jump to line 1063 because the condition on line 1061 was always true

1062 tgtLocation = self.locationFactory.fromPath(pathInStore, trusted_path=True) 

1063 elif transfer == "split": 

1064 # Outside the datastore but treat that as a direct ingest 

1065 # instead. 

1066 tgtLocation = None 

1067 else: 

1068 raise RuntimeError(f"Unexpected transfer mode encountered: {transfer} for URI {srcUri}") 

1069 elif transfer == "direct": 

1070 # Want to store the full URI to the resource directly in 

1071 # datastore. This is useful for referring to permanent archive 

1072 # storage for raw data. 

1073 # Trust that people know what they are doing. 

1074 tgtLocation = None 

1075 else: 

1076 # Work out the name we want this ingested file to have 

1077 # inside the datastore 

1078 tgtLocation = self._calculate_ingested_datastore_name(srcUri, ref, formatter) 

1079 

1080 # if we are transferring from a local file to a remote location 

1081 # it may be more efficient to get the size and checksum of the 

1082 # local file rather than the transferred one 

1083 if record_validation_info and srcUri.isLocal: 

1084 size = srcUri.size() 

1085 checksum = self.computeChecksum(srcUri) if self.useChecksum else None 

1086 have_sized = True 

1087 

1088 # Transfer the resource to the destination. 

1089 # Allow overwrite of an existing file. This matches the behavior 

1090 # of datastore.put() in that it trusts that registry would not 

1091 # be asking to overwrite unless registry thought that the 

1092 # overwrite was allowed. 

1093 tgtLocation.uri.transfer_from( 

1094 srcUri, transfer=transfer, transaction=self._transaction, overwrite=True 

1095 ) 

1096 

1097 if tgtLocation is None: 

1098 # This means we are using direct mode 

1099 targetUri = srcUri 

1100 targetPath = str(srcUri) 

1101 else: 

1102 targetUri = tgtLocation.uri 

1103 targetPath = tgtLocation.pathInStore.path 

1104 

1105 # the file should exist in the datastore now 

1106 if record_validation_info: 

1107 if not have_sized: 

1108 size = targetUri.size() 

1109 checksum = self.computeChecksum(targetUri) if self.useChecksum else None 

1110 else: 

1111 # Not recording any file information. 

1112 size = -1 

1113 checksum = None 

1114 

1115 return StoredFileInfo( 

1116 formatter=formatter, 

1117 path=targetPath, 

1118 storageClass=ref.datasetType.storageClass, 

1119 component=ref.datasetType.component(), 

1120 file_size=size, 

1121 checksum=checksum, 

1122 ) 

1123 

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

1125 # Docstring inherited from Datastore._prepIngest. 

1126 filtered = [] 

1127 

1128 # Ingest could be given tens of thousands of files. It is not efficient 

1129 # to check for the existence of every single file (especially if they 

1130 # are remote URIs) but in some transfer modes the files will be checked 

1131 # anyhow when they are relocated. For direct or None transfer modes 

1132 # it is possible to not know if the file is accessible at all. 

1133 # Therefore limit number of files that will be checked (but always 

1134 # include the first one). 

1135 max_checks = 200 

1136 n_datasets = len(datasets) 

1137 if n_datasets <= max_checks: 1137 ↛ 1139line 1137 didn't jump to line 1139 because the condition on line 1137 was always true

1138 check_every_n = 1 

1139 elif transfer in ("direct", None): 

1140 check_every_n = int(n_datasets / max_checks + 1) # +1 so that if n < max_checks the answer is 1. 

1141 else: 

1142 check_every_n = 0 

1143 

1144 dataset_ids: set[DatasetId] = set() 

1145 for count, dataset in enumerate(datasets): 

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

1147 if not acceptable: 

1148 continue 

1149 else: 

1150 dataset.refs = acceptable 

1151 dataset_ids.update(ref.id for ref in acceptable) 

1152 if dataset.formatter is None: 

1153 dataset.formatter = self.formatterFactory.getFormatterClass(dataset.refs[0]) 

1154 else: 

1155 assert isinstance(dataset.formatter, type | str) 

1156 formatter_class = get_class_of(dataset.formatter) 

1157 if not issubclass(formatter_class, Formatter | FormatterV2): 1157 ↛ 1158line 1157 didn't jump to line 1158 because the condition on line 1157 was never true

1158 raise TypeError(f"Requested formatter {dataset.formatter} is not a Formatter class.") 

1159 dataset.formatter = formatter_class 

1160 

1161 # Decide whether the file should be checked. 

1162 check_existence = False 

1163 if check_every_n != 0: 1163 ↛ 1168line 1163 didn't jump to line 1168 because the condition on line 1163 was always true

1164 # First time through count is 0 so we guarantee to check 

1165 # the first file but not necessarily the final one. 

1166 check_existence = count % check_every_n == 0 

1167 

1168 if check_existence: 1168 ↛ 1177line 1168 didn't jump to line 1177 because the condition on line 1168 was always true

1169 log.debug( 

1170 "Checking file existence: %s (%d/%d) [%s]", 

1171 check_existence, 

1172 count + 1, 

1173 n_datasets, 

1174 transfer, 

1175 ) 

1176 

1177 dataset.path = self._standardizeIngestPath( 

1178 dataset.path, transfer=transfer, check_existence=check_existence 

1179 ) 

1180 filtered.append(dataset) 

1181 

1182 # Check that none of dataset IDs is already in this datastore. 

1183 self._check_existing_datasets(dataset_ids, True, transfer) 

1184 

1185 return _IngestPrepData(filtered) 

1186 

1187 def _check_existing_datasets( 

1188 self, dataset_ids: Iterable[DatasetId], allow_raw_reingest: bool, transfer: str | None = None 

1189 ) -> None: 

1190 """Check that none of dataset IDs is already in this datastore. 

1191 

1192 Parameters 

1193 ---------- 

1194 dataset_ids : `~collections.abc.Iterable` [`DatasetId`] 

1195 Dataset IDs to check. 

1196 allow_raw_reingest : `bool` 

1197 If true then allow re-ingest of UUIDv5 datasets in direct mode. 

1198 transfer : `str` or `None` 

1199 Transfer mode, only used when ``allow_raw_reingest`` is true. 

1200 

1201 Raises 

1202 ------ 

1203 DuplicateIdError 

1204 Raised if any dataset ID is already known to the datastore. 

1205 """ 

1206 record_batches = self._table.fetch_batches(dataset_id=sorted(dataset_ids)) 

1207 known_ids = set() 

1208 for records in record_batches: 

1209 for record in records: 

1210 dataset_id: UUID = record["dataset_id"] 

1211 if not (allow_raw_reingest and transfer == "direct" and dataset_id.version == 5): 

1212 known_ids.add(dataset_id) 

1213 if known_ids: 

1214 # Will raise DuplicateIdError. 

1215 report_duplicate_ids(known_ids) 

1216 

1217 @transactional 

1218 def _finishIngest( 

1219 self, 

1220 prepData: Datastore.IngestPrepData, 

1221 *, 

1222 transfer: str | None = None, 

1223 record_validation_info: bool = True, 

1224 ) -> None: 

1225 # Docstring inherited from Datastore._finishIngest. 

1226 refsAndInfos = [] 

1227 progress = Progress("lsst.daf.butler.datastores.FileDatastore.ingest", level=logging.DEBUG) 

1228 for dataset in progress.wrap(prepData.datasets, desc="Ingesting dataset files"): 

1229 # Do ingest as if the first dataset ref is associated with the file 

1230 info = self._extractIngestInfo( 

1231 dataset.path, 

1232 dataset.refs[0], 

1233 formatter=dataset.formatter, 

1234 transfer=transfer, 

1235 record_validation_info=record_validation_info, 

1236 ) 

1237 refsAndInfos.extend([(ref, info) for ref in dataset.refs]) 

1238 

1239 # In direct mode we can allow repeated ingests of the same thing 

1240 # if we are sure that the external dataset is immutable. We use 

1241 # UUIDv5 to indicate this. If there is a mix of v4 and v5 they are 

1242 # separated. 

1243 refs_and_infos_replace = [] 

1244 refs_and_infos_insert = [] 

1245 if transfer == "direct": 

1246 for entry in refsAndInfos: 

1247 if entry[0].id.version == 5: 

1248 refs_and_infos_replace.append(entry) 

1249 else: 

1250 refs_and_infos_insert.append(entry) 

1251 else: 

1252 refs_and_infos_insert = refsAndInfos 

1253 

1254 if refs_and_infos_insert: 

1255 self._register_datasets(refs_and_infos_insert, insert_mode=DatabaseInsertMode.INSERT) 

1256 if refs_and_infos_replace: 

1257 self._register_datasets(refs_and_infos_replace, insert_mode=DatabaseInsertMode.REPLACE) 

1258 

1259 def _calculate_ingested_datastore_name( 

1260 self, 

1261 srcUri: ResourcePath, 

1262 ref: DatasetRef, 

1263 formatter: Formatter | FormatterV2 | type[Formatter | FormatterV2] | None = None, 

1264 ) -> Location: 

1265 """Given a source URI and a DatasetRef, determine the name the 

1266 dataset will have inside datastore. 

1267 

1268 Parameters 

1269 ---------- 

1270 srcUri : `lsst.resources.ResourcePath` 

1271 URI to the source dataset file. 

1272 ref : `DatasetRef` 

1273 Ref associated with the newly-ingested dataset artifact. This 

1274 is used to determine the name within the datastore. 

1275 formatter : `Formatter` or Formatter class. 

1276 Formatter to use for validation. Can be a class or an instance. 

1277 No validation of the file extension is performed if the 

1278 ``formatter`` is `None`. This can be used if the caller knows 

1279 that the source URI and target URI will use the same formatter. 

1280 

1281 Returns 

1282 ------- 

1283 location : `Location` 

1284 Target location for the newly-ingested dataset. 

1285 """ 

1286 # Ingesting a file from outside the datastore. 

1287 # This involves a new name. 

1288 template = self.templates.getTemplate(ref) 

1289 location = self.locationFactory.fromPath(template.format(ref), trusted_path=True) 

1290 

1291 # Get the extension 

1292 ext = srcUri.getExtension() 

1293 

1294 # Update the destination to include that extension 

1295 location.updateExtension(ext) 

1296 

1297 # Ask the formatter to validate this extension 

1298 if formatter is not None: 

1299 formatter.validate_extension(location) 

1300 

1301 return location 

1302 

1303 def _write_in_memory_to_artifact( 

1304 self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None 

1305 ) -> StoredFileInfo: 

1306 """Write out in memory dataset to datastore. 

1307 

1308 Parameters 

1309 ---------- 

1310 inMemoryDataset : `object` 

1311 Dataset to write to datastore. 

1312 ref : `DatasetRef` 

1313 Registry information associated with this dataset. 

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

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

1316 Not supported by all formatters. 

1317 

1318 Returns 

1319 ------- 

1320 info : `StoredFileInfo` 

1321 Information describing the artifact written to the datastore. 

1322 """ 

1323 # May need to coerce the in memory dataset to the correct 

1324 # python type, but first we need to make sure the storage class 

1325 # reflects the one defined in the data repository. 

1326 ref = self._cast_storage_class(ref) 

1327 

1328 # Confirm that we can accept this dataset 

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

1330 # Raise rather than use boolean return value. 

1331 raise DatasetTypeNotSupportedError( 

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

1333 ) 

1334 

1335 location, formatter = self._determine_put_formatter_location(ref) 

1336 

1337 # The external storage class can differ from the registry storage 

1338 # class AND the given in-memory dataset might not match any of the 

1339 # storage class definitions. 

1340 if formatter.can_accept(inMemoryDataset): 

1341 # Do not need to coerce. Must assume that the formatter can handle 

1342 # it without further checking of types. 

1343 pass 

1344 else: 

1345 # Coerce to a type that it can accept. 

1346 inMemoryDataset = ref.datasetType.storageClass.coerce_type(inMemoryDataset) 

1347 required_pytype = ref.datasetType.storageClass.pytype 

1348 

1349 if not isinstance(inMemoryDataset, required_pytype): 1349 ↛ 1350line 1349 didn't jump to line 1350 because the condition on line 1349 was never true

1350 raise TypeError( 

1351 f"Inconsistency between supplied object ({type(inMemoryDataset)}) " 

1352 f"and storage class type ({required_pytype})" 

1353 ) 

1354 

1355 if self._transaction is None: 1355 ↛ 1356line 1355 didn't jump to line 1356 because the condition on line 1355 was never true

1356 raise RuntimeError("Attempting to write artifact without transaction enabled") 

1357 

1358 def _removeFileExists(uri: ResourcePath) -> None: 

1359 """Remove a file and do not complain if it is not there. 

1360 

1361 This is important since a formatter might fail before the file 

1362 is written and we should not confuse people by writing spurious 

1363 error messages to the log. 

1364 """ 

1365 with contextlib.suppress(FileNotFoundError): 

1366 uri.remove() 

1367 

1368 # Register a callback to try to delete the uploaded data if 

1369 # something fails below 

1370 uri = location.uri 

1371 self._transaction.registerUndo("artifactWrite", _removeFileExists, uri) 

1372 

1373 # Need to record the specified formatter but if this is a V1 formatter 

1374 # we need to convert it to a V2 compatible shim to do the write. 

1375 if not isinstance(formatter, Formatter): 

1376 formatter_compat = formatter 

1377 else: 

1378 formatter_compat = FormatterV1inV2( 

1379 formatter.file_descriptor, 

1380 ref=ref, 

1381 formatter=formatter, 

1382 write_parameters=formatter.write_parameters, 

1383 write_recipes=formatter.write_recipes, 

1384 ) 

1385 

1386 assert isinstance(formatter_compat, FormatterV2) 

1387 

1388 with time_this(log, msg="Writing dataset %s with formatter %s", args=(ref, formatter.name())): 

1389 try: 

1390 formatter_compat.write( 

1391 inMemoryDataset, cache_manager=self.cacheManager, provenance=provenance 

1392 ) 

1393 except Exception as e: 

1394 raise RuntimeError( 

1395 f"Failed to serialize dataset {ref} of type {get_full_type_name(inMemoryDataset)} " 

1396 f"using formatter {formatter.name()}." 

1397 ) from e 

1398 

1399 # URI is needed to resolve what ingest case are we dealing with 

1400 return self._extractIngestInfo(uri, ref, formatter=formatter) 

1401 

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

1403 """Check if the dataset is known to the datastore. 

1404 

1405 Does not check for existence of any artifact. 

1406 

1407 Parameters 

1408 ---------- 

1409 ref : `DatasetRef` 

1410 Reference to the required dataset. 

1411 

1412 Returns 

1413 ------- 

1414 exists : `bool` 

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

1416 """ 

1417 fileLocations = self._get_dataset_locations_info(ref, ignore_datastore_records=True) 

1418 if fileLocations: 

1419 return True 

1420 return False 

1421 

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

1423 # Docstring inherited from the base class. 

1424 refs = list(refs) 

1425 

1426 # The records themselves. Could be missing some entries. 

1427 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True) 

1428 

1429 return {ref: ref.id in records for ref in refs} 

1430 

1431 def _process_mexists_records( 

1432 self, 

1433 id_to_ref: dict[DatasetId, DatasetRef], 

1434 records: dict[DatasetId, list[StoredFileInfo]], 

1435 all_required: bool, 

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

1437 ) -> dict[DatasetRef, bool]: 

1438 """Check given records for existence. 

1439 

1440 Helper function for `mexists()`. 

1441 

1442 Parameters 

1443 ---------- 

1444 id_to_ref : `dict` of [`DatasetId`, `DatasetRef`] 

1445 Mapping of the dataset ID to the dataset ref itself. 

1446 records : `dict` of [`DatasetId`, `list` of `StoredFileInfo`] 

1447 Records as generally returned by 

1448 ``_get_stored_records_associated_with_refs``. 

1449 all_required : `bool` 

1450 Flag to indicate whether existence requires all artifacts 

1451 associated with a dataset ID to exist or not for existence. 

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

1453 Optional mapping of datastore artifact to existence. Updated by 

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

1455 if the caller is not interested. 

1456 

1457 Returns 

1458 ------- 

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

1460 Mapping from dataset to boolean indicating existence. 

1461 """ 

1462 # The URIs to be checked and a mapping of those URIs to 

1463 # the dataset ID. 

1464 uris_to_check: list[ResourcePath] = [] 

1465 location_map: dict[ResourcePath, DatasetId] = {} 

1466 

1467 location_factory = self.locationFactory 

1468 

1469 uri_existence: dict[ResourcePath, bool] = {} 

1470 for ref_id, infos in records.items(): 

1471 # Key is the dataset Id, value is list of StoredItemInfo 

1472 uris = [info.file_location(location_factory).uri for info in infos] 

1473 location_map.update({uri: ref_id for uri in uris}) 

1474 

1475 # Check the local cache directly for a dataset corresponding 

1476 # to the remote URI. 

1477 if self.cacheManager.file_count > 0: 

1478 ref = id_to_ref[ref_id] 

1479 for uri, storedFileInfo in zip(uris, infos, strict=True): 

1480 check_ref = ref 

1481 if not ref.datasetType.isComponent() and (component := storedFileInfo.component): 1481 ↛ 1482line 1481 didn't jump to line 1482 because the condition on line 1481 was never true

1482 check_ref = ref.makeComponentRef(component) 

1483 if self.cacheManager.known_to_cache(check_ref, uri.getExtension()): 

1484 # Proxy for URI existence. 

1485 uri_existence[uri] = True 

1486 else: 

1487 uris_to_check.append(uri) 

1488 else: 

1489 # Check all of them. 

1490 uris_to_check.extend(uris) 

1491 

1492 if artifact_existence is not None: 

1493 # If a URI has already been checked remove it from the list 

1494 # and immediately add the status to the output dict. 

1495 filtered_uris_to_check = [] 

1496 for uri in uris_to_check: 

1497 if uri in artifact_existence: 1497 ↛ 1498line 1497 didn't jump to line 1498 because the condition on line 1497 was never true

1498 uri_existence[uri] = artifact_existence[uri] 

1499 else: 

1500 filtered_uris_to_check.append(uri) 

1501 uris_to_check = filtered_uris_to_check 

1502 

1503 # Results. 

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

1505 

1506 uri_existence.update(ResourcePath.mexists(uris_to_check)) 

1507 for uri, exists in uri_existence.items(): 

1508 dataset_id = location_map[uri] 

1509 ref = id_to_ref[dataset_id] 

1510 

1511 # Disassembled composite needs to check all locations. 

1512 # all_required indicates whether all need to exist or not. 

1513 if ref in dataset_existence: 

1514 if all_required: 

1515 exists = dataset_existence[ref] and exists 

1516 else: 

1517 exists = dataset_existence[ref] or exists 

1518 dataset_existence[ref] = exists 

1519 

1520 if artifact_existence is not None: 

1521 artifact_existence.update(uri_existence) 

1522 

1523 return dataset_existence 

1524 

1525 def mexists( 

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

1527 ) -> dict[DatasetRef, bool]: 

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

1529 

1530 Parameters 

1531 ---------- 

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

1533 The datasets to be checked. 

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

1535 Optional mapping of datastore artifact to existence. Updated by 

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

1537 if the caller is not interested. 

1538 

1539 Returns 

1540 ------- 

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

1542 Mapping from dataset to boolean indicating existence. 

1543 

1544 Notes 

1545 ----- 

1546 To minimize potentially costly remote existence checks, the local 

1547 cache is checked as a proxy for existence. If a file for this 

1548 `DatasetRef` does exist no check is done for the actual URI. This 

1549 could result in possibly unexpected behavior if the dataset itself 

1550 has been removed from the datastore by another process whilst it is 

1551 still in the cache. 

1552 """ 

1553 chunk_size = 50_000 

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

1555 log.debug("Checking for the existence of multiple artifacts in datastore in chunks of %d", chunk_size) 

1556 n_found_total = 0 

1557 n_checked = 0 

1558 n_chunks = 0 

1559 for chunk in chunk_iterable(refs, chunk_size=chunk_size): 

1560 chunk_result = self._mexists(chunk, artifact_existence) 

1561 

1562 # The log message level and content depend on how many 

1563 # datasets we are processing. 

1564 n_results = len(chunk_result) 

1565 

1566 # Use verbose logging to ensure that messages can be seen 

1567 # easily if many refs are being checked. 

1568 log_threshold = VERBOSE 

1569 n_checked += n_results 

1570 

1571 # This sum can take some time so only do it if we know the 

1572 # result is going to be used. 

1573 n_found = 0 

1574 if log.isEnabledFor(log_threshold): 

1575 # Can treat the booleans as 0, 1 integers and sum them. 

1576 n_found = sum(chunk_result.values()) 

1577 n_found_total += n_found 

1578 

1579 # We are deliberately not trying to count the number of refs 

1580 # provided in case it's in the millions. This means there is a 

1581 # situation where the number of refs exactly matches the chunk 

1582 # size and we will switch to the multi-chunk path even though 

1583 # we only have a single chunk. 

1584 if n_results < chunk_size and n_chunks == 0: 1584 ↛ 1606line 1584 didn't jump to line 1606 because the condition on line 1584 was always true

1585 # Single chunk will be processed so we can provide more detail. 

1586 if n_results == 1: 

1587 ref = list(chunk_result)[0] 

1588 # Use debug logging to be consistent with `exists()`. 

1589 log.debug( 

1590 "Calling mexists() with single ref that does%s exist (%s).", 

1591 "" if chunk_result[ref] else " not", 

1592 ref, 

1593 ) 

1594 else: 

1595 # Single chunk but multiple files. Summarize. 

1596 log.log( 

1597 log_threshold, 

1598 "Number of datasets found in datastore %s: %d out of %d datasets checked.", 

1599 self.name, 

1600 n_found, 

1601 n_checked, 

1602 ) 

1603 

1604 else: 

1605 # Use incremental verbose logging when we have multiple chunks. 

1606 log.log( 

1607 log_threshold, 

1608 "Number of datasets found in datastore for chunk %d: %d out of %d checked " 

1609 "(running total from all chunks so far: %d found out of %d checked)", 

1610 n_chunks, 

1611 n_found, 

1612 n_results, 

1613 n_found_total, 

1614 n_checked, 

1615 ) 

1616 dataset_existence.update(chunk_result) 

1617 n_chunks += 1 

1618 

1619 return dataset_existence 

1620 

1621 def _mexists( 

1622 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None 

1623 ) -> dict[DatasetRef, bool]: 

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

1625 

1626 Parameters 

1627 ---------- 

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

1629 The datasets to be checked. 

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

1631 Optional mapping of datastore artifact to existence. Updated by 

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

1633 if the caller is not interested. 

1634 

1635 Returns 

1636 ------- 

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

1638 Mapping from dataset to boolean indicating existence. 

1639 """ 

1640 # Make a mapping from refs with the internal storage class to the given 

1641 # refs that may have a different one. We'll use the internal refs 

1642 # throughout this method and convert back at the very end. 

1643 internal_ref_to_input_ref = {self._cast_storage_class(ref): ref for ref in refs} 

1644 

1645 # Need a mapping of dataset_id to (internal) dataset ref since some 

1646 # internal APIs work with dataset_id. 

1647 id_to_ref = {ref.id: ref for ref in internal_ref_to_input_ref} 

1648 

1649 # Set of all IDs we are checking for. 

1650 requested_ids = set(id_to_ref.keys()) 

1651 

1652 # The records themselves. Could be missing some entries. 

1653 records = self._get_stored_records_associated_with_refs( 

1654 id_to_ref.values(), ignore_datastore_records=True 

1655 ) 

1656 

1657 dataset_existence = self._process_mexists_records( 

1658 id_to_ref, records, True, artifact_existence=artifact_existence 

1659 ) 

1660 

1661 # Set of IDs that have been handled. 

1662 handled_ids = {ref.id for ref in dataset_existence} 

1663 

1664 missing_ids = requested_ids - handled_ids 

1665 if missing_ids: 

1666 dataset_existence.update( 

1667 self._mexists_check_expected( 

1668 [id_to_ref[missing] for missing in missing_ids], artifact_existence 

1669 ) 

1670 ) 

1671 

1672 return { 

1673 internal_ref_to_input_ref[internal_ref]: existence 

1674 for internal_ref, existence in dataset_existence.items() 

1675 } 

1676 

1677 def _mexists_check_expected( 

1678 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None 

1679 ) -> dict[DatasetRef, bool]: 

1680 """Check existence of refs that are not known to datastore. 

1681 

1682 Parameters 

1683 ---------- 

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

1685 The datasets to be checked. These are assumed not to be known 

1686 to datastore. 

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

1688 Optional mapping of datastore artifact to existence. Updated by 

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

1690 if the caller is not interested. 

1691 

1692 Returns 

1693 ------- 

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

1695 Mapping from dataset to boolean indicating existence. 

1696 """ 

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

1698 if not self.trustGetRequest: 

1699 # Must assume these do not exist 

1700 for ref in refs: 

1701 dataset_existence[ref] = False 

1702 else: 

1703 log.debug( 

1704 "%d datasets were not known to datastore during initial existence check.", 

1705 len(refs), 

1706 ) 

1707 

1708 # Construct data structure identical to that returned 

1709 # by _get_stored_records_associated_with_refs() but using 

1710 # guessed names. 

1711 records = {} 

1712 id_to_ref = {} 

1713 for missing_ref in refs: 

1714 expected = self._get_expected_dataset_locations_info(missing_ref) 

1715 dataset_id = missing_ref.id 

1716 records[dataset_id] = [info for _, info in expected] 

1717 id_to_ref[dataset_id] = missing_ref 

1718 

1719 dataset_existence.update( 

1720 self._process_mexists_records( 

1721 id_to_ref, 

1722 records, 

1723 False, 

1724 artifact_existence=artifact_existence, 

1725 ) 

1726 ) 

1727 

1728 return dataset_existence 

1729 

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

1731 """Check if the dataset exists in the datastore. 

1732 

1733 Parameters 

1734 ---------- 

1735 ref : `DatasetRef` 

1736 Reference to the required dataset. 

1737 

1738 Returns 

1739 ------- 

1740 exists : `bool` 

1741 `True` if the entity exists in the `Datastore`. 

1742 

1743 Notes 

1744 ----- 

1745 The local cache is checked as a proxy for existence in the remote 

1746 object store. It is possible that another process on a different 

1747 compute node could remove the file from the object store even 

1748 though it is present in the local cache. 

1749 """ 

1750 ref = self._cast_storage_class(ref) 

1751 # We cannot trust datastore records from ref, as many unit tests delete 

1752 # datasets and check their existence. 

1753 fileLocations = self._get_dataset_locations_info(ref, ignore_datastore_records=True) 

1754 

1755 # if we are being asked to trust that registry might not be correct 

1756 # we ask for the expected locations and check them explicitly 

1757 if not fileLocations: 

1758 if not self.trustGetRequest: 

1759 return False 

1760 

1761 # First check the cache. If it is not found we must check 

1762 # the datastore itself. Assume that any component in the cache 

1763 # means that the dataset does exist somewhere. 

1764 if self.cacheManager.known_to_cache(ref): 

1765 return True 

1766 

1767 # When we are guessing a dataset location we can not check 

1768 # for the existence of every component since we can not 

1769 # know if every component was written. Instead we check 

1770 # for the existence of any of the expected locations. 

1771 for location, _ in self._get_expected_dataset_locations_info(ref): 

1772 if self._artifact_exists(location): 

1773 return True 

1774 return False 

1775 

1776 # All listed artifacts must exist. 

1777 for location, storedFileInfo in fileLocations: 

1778 # Checking in cache needs the component ref. 

1779 check_ref = ref 

1780 if not ref.datasetType.isComponent() and (component := storedFileInfo.component): 

1781 check_ref = ref.makeComponentRef(component) 

1782 if self.cacheManager.known_to_cache(check_ref, location.getExtension()): 

1783 continue 

1784 

1785 if not self._artifact_exists(location): 

1786 return False 

1787 

1788 return True 

1789 

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

1791 """Return URIs associated with dataset. 

1792 

1793 Parameters 

1794 ---------- 

1795 ref : `DatasetRef` 

1796 Reference to the required dataset. 

1797 predict : `bool`, optional 

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

1799 it should return a predicted URI or not. 

1800 

1801 Returns 

1802 ------- 

1803 uris : `DatasetRefURIs` 

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

1805 the dataset was disassembled within the datastore this may be 

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

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

1808 """ 

1809 many = self.getManyURIs([ref], predict=predict, allow_missing=False) 

1810 return many[ref] 

1811 

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

1813 """URI to the Dataset. 

1814 

1815 Parameters 

1816 ---------- 

1817 ref : `DatasetRef` 

1818 Reference to the required Dataset. 

1819 predict : `bool` 

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

1821 been written. 

1822 

1823 Returns 

1824 ------- 

1825 uri : `str` 

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

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

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

1829 fragment "#predicted". 

1830 If the datastore does not have entities that relate well 

1831 to the concept of a URI the returned URI will be 

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

1833 

1834 Raises 

1835 ------ 

1836 FileNotFoundError 

1837 Raised if a URI has been requested for a dataset that does not 

1838 exist and guessing is not allowed. 

1839 RuntimeError 

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

1841 are associated with this dataset. 

1842 

1843 Notes 

1844 ----- 

1845 When a predicted URI is requested an attempt will be made to form 

1846 a reasonable URI based on file templates and the expected formatter. 

1847 """ 

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

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

1850 raise RuntimeError( 

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

1852 ) 

1853 return primary 

1854 

1855 def _predict_URIs( 

1856 self, 

1857 ref: DatasetRef, 

1858 ) -> DatasetRefURIs: 

1859 """Predict the URIs of a dataset ref. 

1860 

1861 Parameters 

1862 ---------- 

1863 ref : `DatasetRef` 

1864 Reference to the required Dataset. 

1865 

1866 Returns 

1867 ------- 

1868 URI : DatasetRefUris 

1869 Primary and component URIs. URIs will contain a URI fragment 

1870 "#predicted". 

1871 """ 

1872 uris = DatasetRefURIs() 

1873 

1874 if self.composites.shouldBeDisassembled(ref): 

1875 for component, _ in ref.datasetType.storageClass.components.items(): 

1876 comp_ref = ref.makeComponentRef(component) 

1877 comp_location, _ = self._determine_put_formatter_location(comp_ref) 

1878 

1879 # Add the "#predicted" URI fragment to indicate this is a 

1880 # guess 

1881 uris.componentURIs[component] = ResourcePath( 

1882 comp_location.uri.geturl() + "#predicted", forceDirectory=comp_location.uri.dirLike 

1883 ) 

1884 

1885 else: 

1886 location, _ = self._determine_put_formatter_location(ref) 

1887 

1888 # Add the "#predicted" URI fragment to indicate this is a guess 

1889 uris.primaryURI = ResourcePath( 

1890 location.uri.geturl() + "#predicted", forceDirectory=location.uri.dirLike 

1891 ) 

1892 

1893 return uris 

1894 

1895 def getManyURIs( 

1896 self, 

1897 refs: Iterable[DatasetRef], 

1898 predict: bool = False, 

1899 allow_missing: bool = False, 

1900 ) -> dict[DatasetRef, DatasetRefURIs]: 

1901 # Docstring inherited 

1902 

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

1904 

1905 records = self._get_stored_records_associated_with_refs(refs) 

1906 records_keys = records.keys() 

1907 

1908 existing_refs = tuple(ref for ref in refs if ref.id in records_keys) 

1909 missing_refs = tuple(ref for ref in refs if ref.id not in records_keys) 

1910 

1911 # Have to handle trustGetRequest mode by checking for the existence 

1912 # of the missing refs on disk. 

1913 if missing_refs and not predict: 

1914 dataset_existence = self._mexists_check_expected(missing_refs, None) 

1915 really_missing = set() 

1916 not_missing = set() 

1917 for ref, exists in dataset_existence.items(): 

1918 if exists: 

1919 not_missing.add(ref) 

1920 else: 

1921 really_missing.add(ref) 

1922 

1923 if not_missing: 

1924 # Need to recalculate the missing/existing split. 

1925 existing_refs = existing_refs + tuple(not_missing) 

1926 missing_refs = tuple(really_missing) 

1927 

1928 for ref in missing_refs: 

1929 # if this has never been written then we have to guess 

1930 if not predict: 

1931 if not allow_missing: 

1932 raise FileNotFoundError(f"Dataset {ref} not in this datastore.") 

1933 else: 

1934 uris[ref] = self._predict_URIs(ref) 

1935 

1936 for ref in existing_refs: 

1937 file_infos = records[ref.id] 

1938 file_locations = [(i.file_location(self.locationFactory), i) for i in file_infos] 

1939 uris[ref] = self._locations_to_URI(ref, file_locations) 

1940 

1941 return uris 

1942 

1943 def _locations_to_URI( 

1944 self, 

1945 ref: DatasetRef, 

1946 file_locations: Sequence[tuple[Location, StoredFileInfo]], 

1947 ) -> DatasetRefURIs: 

1948 """Convert one or more file locations associated with a DatasetRef 

1949 to a DatasetRefURIs. 

1950 

1951 Parameters 

1952 ---------- 

1953 ref : `DatasetRef` 

1954 Reference to the dataset. 

1955 file_locations : Sequence[Tuple[Location, StoredFileInfo]] 

1956 Each item in the sequence is the location of the dataset within the 

1957 datastore and stored information about the file and its formatter. 

1958 If there is only one item in the sequence then it is treated as the 

1959 primary URI. If there is more than one item then they are treated 

1960 as component URIs. If there are no items then an error is raised 

1961 unless ``self.trustGetRequest`` is `True`. 

1962 

1963 Returns 

1964 ------- 

1965 uris: DatasetRefURIs 

1966 Represents the primary URI or component URIs described by the 

1967 inputs. 

1968 

1969 Raises 

1970 ------ 

1971 RuntimeError 

1972 If no file locations are passed in and ``self.trustGetRequest`` is 

1973 `False`. 

1974 FileNotFoundError 

1975 If the a passed-in URI does not exist, and ``self.trustGetRequest`` 

1976 is `False`. 

1977 RuntimeError 

1978 If a passed in `StoredFileInfo`'s ``component`` is `None` (this is 

1979 unexpected). 

1980 """ 

1981 guessing = False 

1982 uris = DatasetRefURIs() 

1983 

1984 if not file_locations: 

1985 if not self.trustGetRequest: 1985 ↛ 1986line 1985 didn't jump to line 1986 because the condition on line 1985 was never true

1986 raise RuntimeError(f"Unexpectedly got no artifacts for dataset {ref}") 

1987 file_locations = self._get_expected_dataset_locations_info(ref) 

1988 guessing = True 

1989 

1990 if len(file_locations) == 1: 

1991 # No disassembly so this is the primary URI 

1992 uris.primaryURI = file_locations[0][0].uri 

1993 if guessing and not uris.primaryURI.exists(): 1993 ↛ 1994line 1993 didn't jump to line 1994 because the condition on line 1993 was never true

1994 raise FileNotFoundError(f"Expected URI ({uris.primaryURI}) does not exist") 

1995 else: 

1996 for location, file_info in file_locations: 

1997 if file_info.component is None: 1997 ↛ 1998line 1997 didn't jump to line 1998 because the condition on line 1997 was never true

1998 raise RuntimeError(f"Unexpectedly got no component name for a component at {location}") 

1999 if guessing and not location.uri.exists(): 1999 ↛ 2003line 1999 didn't jump to line 2003 because the condition on line 1999 was never true

2000 # If we are trusting then it is entirely possible for 

2001 # some components to be missing. In that case we skip 

2002 # to the next component. 

2003 if self.trustGetRequest: 

2004 continue 

2005 raise FileNotFoundError(f"Expected URI ({location.uri}) does not exist") 

2006 uris.componentURIs[file_info.component] = location.uri 

2007 

2008 return uris 

2009 

2010 def _find_missing_records( 

2011 self, 

2012 refs: Iterable[DatasetRef], 

2013 missing_ids: set[DatasetId], 

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

2015 warn_for_missing: bool = True, 

2016 ) -> dict[DatasetId, list[StoredFileInfo]]: 

2017 if not missing_ids: 

2018 return {} 

2019 

2020 if artifact_existence is None: 2020 ↛ 2021line 2020 didn't jump to line 2021 because the condition on line 2020 was never true

2021 artifact_existence = {} 

2022 

2023 found_records: dict[DatasetId, list[StoredFileInfo]] = defaultdict(list) 

2024 id_to_ref = {ref.id: ref for ref in refs if ref.id in missing_ids} 

2025 

2026 # This should be chunked in case we end up having to check 

2027 # the file store since we need some log output to show 

2028 # progress. 

2029 chunk_size = 50_000 

2030 for missing_ids_chunk in chunk_iterable(missing_ids, chunk_size=chunk_size): 

2031 records = {} 

2032 for missing in missing_ids_chunk: 

2033 # Ask the source datastore where the missing artifacts 

2034 # should be. An execution butler might not know about the 

2035 # artifacts even if they are there. 

2036 expected = self._get_expected_dataset_locations_info(id_to_ref[missing]) 

2037 records[missing] = [info for _, info in expected] 

2038 

2039 # Call the mexist helper method in case we have not already 

2040 # checked these artifacts such that artifact_existence is 

2041 # empty. This allows us to benefit from parallelism. 

2042 # datastore.mexists() itself does not give us access to the 

2043 # derived datastore record. 

2044 log.verbose("Checking existence of %d datasets unknown to datastore", len(records)) 

2045 ref_exists = self._process_mexists_records( 

2046 id_to_ref, records, False, artifact_existence=artifact_existence 

2047 ) 

2048 

2049 # Now go through the records and propagate the ones that exist. 

2050 location_factory = self.locationFactory 

2051 for missing, record_list in records.items(): 

2052 # Skip completely if the ref does not exist. 

2053 ref = id_to_ref[missing] 

2054 if not ref_exists[ref]: 

2055 if warn_for_missing: 2055 ↛ 2056line 2055 didn't jump to line 2056 because the condition on line 2055 was never true

2056 log.warning("Asked to transfer dataset %s but no file artifacts exist for it.", ref) 

2057 continue 

2058 # Check for file artifact to decide which parts of a 

2059 # disassembled composite do exist. If there is only a 

2060 # single record we don't even need to look because it can't 

2061 # be a composite and must exist. 

2062 if len(record_list) == 1: 

2063 dataset_records = record_list 

2064 else: 

2065 dataset_records = [ 

2066 record 

2067 for record in record_list 

2068 if artifact_existence[record.file_location(location_factory).uri] 

2069 ] 

2070 assert len(dataset_records) > 0, "Disassembled composite should have had some files." 

2071 

2072 # Rely on source_records being a defaultdict. 

2073 found_records[missing].extend(dataset_records) 

2074 log.verbose("Completed scan for missing data files") 

2075 return found_records 

2076 

2077 def retrieveArtifacts( 

2078 self, 

2079 refs: Iterable[DatasetRef], 

2080 destination: ResourcePath, 

2081 transfer: str = "auto", 

2082 preserve_path: bool = True, 

2083 overwrite: bool = False, 

2084 write_index: bool = True, 

2085 add_prefix: bool = False, 

2086 ) -> dict[ResourcePath, ArtifactIndexInfo]: 

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

2088 

2089 Parameters 

2090 ---------- 

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

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

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

2094 be resolved. 

2095 destination : `lsst.resources.ResourcePath` 

2096 Location to write the file artifacts. 

2097 transfer : `str`, optional 

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

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

2100 "move" is not allowed. 

2101 preserve_path : `bool`, optional 

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

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

2104 is used. 

2105 overwrite : `bool`, optional 

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

2107 destination. 

2108 write_index : `bool`, optional 

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

2110 of a `ZipIndex` for the downloaded datasets. 

2111 add_prefix : `bool`, optional 

2112 If `True` and if ``preserve_path`` is `False`, apply a prefix to 

2113 the filenames corresponding to some part of the dataset ref ID. 

2114 This can be used to guarantee uniqueness. 

2115 

2116 Returns 

2117 ------- 

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

2119 `ArtifactIndexInfo` ] 

2120 Mapping of retrieved file to associated index information. 

2121 """ 

2122 if not destination.isdir(): 

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

2124 

2125 if transfer == "move": 

2126 raise ValueError("Can not move artifacts out of datastore. Use copy instead.") 

2127 

2128 # Source -> Destination 

2129 # This also helps filter out duplicate DatasetRef in the request 

2130 # that will map to the same underlying file transfer. 

2131 to_transfer: dict[ResourcePath, ResourcePath] = {} 

2132 zips_to_transfer: set[ResourcePath] = set() 

2133 

2134 # Retrieve all the records in bulk indexed by ref.id. 

2135 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True) 

2136 

2137 # Check for missing records. 

2138 known_ids = set(records) 

2139 log.debug("Number of datastore records found in database: %d", len(known_ids)) 

2140 requested_ids = {ref.id for ref in refs} 

2141 missing_ids = requested_ids - known_ids 

2142 

2143 if missing_ids and not self.trustGetRequest: 2143 ↛ 2144line 2143 didn't jump to line 2144 because the condition on line 2143 was never true

2144 raise ValueError(f"Number of datasets missing from this datastore: {len(missing_ids)}") 

2145 

2146 missing_records = self._find_missing_records(refs, missing_ids) 

2147 records.update(missing_records) 

2148 

2149 # One artifact can be used by multiple DatasetRef. 

2150 # e.g. DECam. 

2151 artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {} 

2152 # Sort to ensure that in many refs to one file situation the same 

2153 # ref is used for any prefix that might be added. 

2154 for ref in sorted(refs): 

2155 prefix = str(ref.id)[:8] + "-" if add_prefix else "" 

2156 for info in records[ref.id]: 

2157 location = info.file_location(self.locationFactory) 

2158 source_uri = location.uri 

2159 # For DECam/zip we only want to copy once. 

2160 # For zip files we need to unpack so that they can be 

2161 # zipped up again if needed. 

2162 is_zip = source_uri.getExtension() == ".zip" and "zip-path" in source_uri.fragment 

2163 # We need to remove fragments for consistency. 

2164 cleaned_source_uri = source_uri.replace(fragment="", query="", params="") 

2165 if is_zip: 2165 ↛ 2168line 2165 didn't jump to line 2168 because the condition on line 2165 was never true

2166 # Assume the DatasetRef definitions are within the Zip 

2167 # file itself and so can be dropped from loop. 

2168 zips_to_transfer.add(cleaned_source_uri) 

2169 elif cleaned_source_uri not in to_transfer: 2169 ↛ 2176line 2169 didn't jump to line 2176 because the condition on line 2169 was always true

2170 target_uri = determine_destination_for_retrieved_artifact( 

2171 destination, location.pathInStore, preserve_path, prefix 

2172 ) 

2173 to_transfer[cleaned_source_uri] = target_uri 

2174 artifact_map[target_uri] = ArtifactIndexInfo.from_single(info.to_simple(), ref.id) 

2175 else: 

2176 target_uri = to_transfer[cleaned_source_uri] 

2177 artifact_map[target_uri].append(ref.id) 

2178 

2179 # Parallelize the transfer. Re-raise as a single exception if 

2180 # a FileExistsError is encountered anywhere. 

2181 log.debug("Number of artifacts to transfer to %s: %d", str(destination), len(to_transfer)) 

2182 try: 

2183 ResourcePath.mtransfer(transfer, tuple(to_transfer.items()), overwrite=overwrite) 

2184 except* FileExistsError as egroup: 

2185 raise FileExistsError( 

2186 "Some files already exist in destination directory and overwrite is False" 

2187 ) from egroup 

2188 

2189 # Transfer the Zip files and unpack them. 

2190 zipped_artifacts = unpack_zips(zips_to_transfer, requested_ids, destination, preserve_path, overwrite) 

2191 artifact_map.update(zipped_artifacts) 

2192 

2193 if write_index: 

2194 index = ZipIndex.from_artifact_map(refs, artifact_map, destination) 

2195 index.write_index(destination) 

2196 

2197 return artifact_map 

2198 

2199 def ingest_zip( 

2200 self, 

2201 zip_path: ResourcePath, 

2202 transfer: str | None, 

2203 *, 

2204 dry_run: bool = False, 

2205 ) -> None: 

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

2207 

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

2209 

2210 Parameters 

2211 ---------- 

2212 zip_path : `lsst.resources.ResourcePath` 

2213 Path to the Zip file. 

2214 transfer : `str` 

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

2216 dry_run : `bool`, optional 

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

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

2219 have any of the datasets. 

2220 

2221 Notes 

2222 ----- 

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

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

2225 if one dataset type is in the constraints list? 

2226 

2227 If any dataset is already present in the datastore the entire ingest 

2228 will fail. 

2229 """ 

2230 index = ZipIndex.from_zip_file(zip_path) 

2231 

2232 # Refs indexed by UUID. 

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

2234 id_to_ref = {ref.id: ref for ref in refs} 

2235 

2236 # Any failing constraints trigger entire failure. 

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

2238 raise DatasetTypeNotSupportedError( 

2239 "Some refs in the Zip file are not supported by this datastore" 

2240 ) 

2241 

2242 # Check that none of dataset IDs is already in this datastore. Note 

2243 # that unlike ingest() this method does not differentiate between 

2244 # UUIDv5 and all others. 

2245 if not dry_run: 

2246 self._check_existing_datasets(id_to_ref, False) 

2247 

2248 # Transfer the Zip file into the datastore file system. 

2249 # There is no RUN as such to use for naming. 

2250 # Potentially could use the RUN from the first ref in the index 

2251 # There is no requirement that the contents of the Zip files share 

2252 # the same RUN. 

2253 # Could use the Zip UUID from the index + special "zips/" prefix. 

2254 if transfer is None: 2254 ↛ 2256line 2254 didn't jump to line 2256 because the condition on line 2254 was never true

2255 # Indicated that the zip file is already in the right place. 

2256 if not zip_path.isabs(): 

2257 tgtLocation = self.locationFactory.fromPath(zip_path.ospath, trusted_path=False) 

2258 else: 

2259 pathInStore = zip_path.relative_to(self.root) 

2260 if pathInStore is None: 

2261 raise RuntimeError( 

2262 f"Unexpectedly learned that {zip_path} is not within datastore {self.root}" 

2263 ) 

2264 tgtLocation = self.locationFactory.fromPath(pathInStore, trusted_path=True) 

2265 elif transfer == "direct": 2265 ↛ 2267line 2265 didn't jump to line 2267 because the condition on line 2265 was never true

2266 # Reference in original location. 

2267 tgtLocation = None 

2268 else: 

2269 # Name the zip file based on index contents. 

2270 tgtLocation = self.locationFactory.fromPath(index.calculate_zip_file_path_in_store()) 

2271 

2272 # Transfer the Zip file into the datastore. 

2273 if not dry_run: 

2274 tgtLocation.uri.transfer_from( 

2275 zip_path, transfer=transfer, transaction=self._transaction, overwrite=True 

2276 ) 

2277 else: 

2278 log.info("Would be copying Zip from %s to %s", zip_path, tgtLocation) 

2279 

2280 if tgtLocation is None: 2280 ↛ 2281line 2280 didn't jump to line 2281 because the condition on line 2280 was never true

2281 path_in_store = str(zip_path) 

2282 else: 

2283 path_in_store = tgtLocation.pathInStore.path 

2284 

2285 # Associate each file with a (DatasetRef, StoredFileInfo) tuple. 

2286 artifacts: list[tuple[DatasetRef, StoredFileInfo]] = [] 

2287 for path_in_zip, index_info in index.artifact_map.items(): 

2288 # Need to modify the info to include the path to the Zip file 

2289 # that was previously written to the datastore. 

2290 index_info.info.path = f"{path_in_store}#zip-path={path_in_zip}" 

2291 

2292 info = StoredFileInfo.from_simple(index_info.info) 

2293 for id_ in index_info.ids: 

2294 artifacts.append((id_to_ref[id_], info)) 

2295 

2296 if not dry_run: 

2297 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT) 

2298 else: 

2299 log.info("Would be registering %d artifacts from Zip into datastore", len(artifacts)) 

2300 

2301 def get( 

2302 self, 

2303 ref: DatasetRef, 

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

2305 storageClass: StorageClass | str | None = None, 

2306 ) -> Any: 

2307 """Load an InMemoryDataset from the store. 

2308 

2309 Parameters 

2310 ---------- 

2311 ref : `DatasetRef` 

2312 Reference to the required Dataset. 

2313 parameters : `dict` 

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

2315 a slice of the dataset to be loaded. 

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

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

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

2319 the dataset type definition for this dataset. Specifying a 

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

2321 This type must be compatible with the original type. 

2322 

2323 Returns 

2324 ------- 

2325 inMemoryDataset : `object` 

2326 Requested dataset or slice thereof as an InMemoryDataset. 

2327 

2328 Raises 

2329 ------ 

2330 FileNotFoundError 

2331 Requested dataset can not be retrieved. 

2332 TypeError 

2333 Return value from formatter has unexpected type. 

2334 ValueError 

2335 Formatter failed to process the dataset. 

2336 """ 

2337 # Supplied storage class for the component being read is either 

2338 # from the ref itself or some an override if we want to force 

2339 # type conversion. 

2340 if storageClass is not None: 

2341 ref = ref.overrideStorageClass(storageClass) 

2342 

2343 allGetInfo = self._prepare_for_direct_get(ref, parameters) 

2344 return get_dataset_as_python_object_from_get_info( 

2345 allGetInfo, ref=ref, parameters=parameters, cache_manager=self.cacheManager 

2346 ) 

2347 

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

2349 # Docstring inherited 

2350 

2351 locations = self._get_dataset_locations_info(ref) 

2352 if len(locations) == 0: 2352 ↛ 2355line 2352 didn't jump to line 2355 because the condition on line 2352 was always true

2353 return None 

2354 

2355 return locations 

2356 

2357 @transactional 

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

2359 """Write a InMemoryDataset with a given `DatasetRef` to the store. 

2360 

2361 Parameters 

2362 ---------- 

2363 inMemoryDataset : `object` 

2364 The dataset to store. 

2365 ref : `DatasetRef` 

2366 Reference to the associated Dataset. 

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

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

2369 Can be ignored by a formatter or delegate. 

2370 

2371 Raises 

2372 ------ 

2373 TypeError 

2374 Supplied object and storage class are inconsistent. 

2375 DatasetTypeNotSupportedError 

2376 The associated `DatasetType` is not handled by this datastore. 

2377 

2378 Notes 

2379 ----- 

2380 If the datastore is configured to reject certain dataset types it 

2381 is possible that the put will fail and raise a 

2382 `DatasetTypeNotSupportedError`. The main use case for this is to 

2383 allow `ChainedDatastore` to put to multiple datastores without 

2384 requiring that every datastore accepts the dataset. 

2385 """ 

2386 # Before modifying anything make sure that dataset ID is not already 

2387 # registered. The check must be done before calling 

2388 # _write_in_memory_to_artifact because it defines undo callback that 

2389 # deletes existing artifact. This does not protect against concurrent 

2390 # inserts by multiple clients, e.g. when doing concurrent imports. 

2391 records = self._table.fetch(dataset_id=ref.id) 

2392 if next(records, None): 

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

2394 

2395 doDisassembly = self.composites.shouldBeDisassembled(ref) 

2396 # doDisassembly = True 

2397 

2398 artifacts = [] 

2399 if doDisassembly: 

2400 inMemoryDataset = ref.datasetType.storageClass.delegate().add_provenance( 

2401 inMemoryDataset, ref, provenance=provenance 

2402 ) 

2403 components = ref.datasetType.storageClass.delegate().disassemble(inMemoryDataset) 

2404 if components is None: 2404 ↛ 2405line 2404 didn't jump to line 2405 because the condition on line 2404 was never true

2405 raise RuntimeError( 

2406 f"Inconsistent configuration: dataset type {ref.datasetType.name} " 

2407 f"with storage class {ref.datasetType.storageClass.name} " 

2408 "is configured to be disassembled, but cannot be." 

2409 ) 

2410 for component, componentInfo in components.items(): 

2411 # Don't recurse because we want to take advantage of 

2412 # bulk insert -- need a new DatasetRef that refers to the 

2413 # same dataset_id but has the component DatasetType 

2414 # DatasetType does not refer to the types of components 

2415 # So we construct one ourselves. 

2416 compRef = ref.makeComponentRef(component) 

2417 # Provenance has already been attached above. 

2418 storedInfo = self._write_in_memory_to_artifact(componentInfo.component, compRef) 

2419 artifacts.append((compRef, storedInfo)) 

2420 else: 

2421 # Write the entire thing out 

2422 storedInfo = self._write_in_memory_to_artifact(inMemoryDataset, ref, provenance=provenance) 

2423 artifacts.append((ref, storedInfo)) 

2424 

2425 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT) 

2426 

2427 @transactional 

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

2429 # At this point can safely remove these datasets from the cache 

2430 # to avoid confusion later on. If they are not trashed later 

2431 # the cache will simply be refilled. 

2432 self.cacheManager.remove_from_cache(ref) 

2433 

2434 # If we are in trust mode there will be nothing to move to 

2435 # the trash table and we will have to try to delete the file 

2436 # immediately. 

2437 if self.trustGetRequest: 

2438 # Try to keep the logic below for a single file trash. 

2439 if isinstance(ref, DatasetRef): 

2440 refs = {ref} 

2441 else: 

2442 # Will recreate ref at the end of this branch. 

2443 refs = set(ref) 

2444 

2445 # Determine which datasets are known to datastore directly. 

2446 id_to_ref = {ref.id: ref for ref in refs} 

2447 existing_ids = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True) 

2448 existing_refs = {id_to_ref[ref_id] for ref_id in existing_ids} 

2449 

2450 missing = refs - existing_refs 

2451 if missing: 

2452 # Do an explicit existence check on these refs. 

2453 # We only care about the artifacts at this point and not 

2454 # the dataset existence. 

2455 artifact_existence: dict[ResourcePath, bool] = {} 

2456 _ = self.mexists(missing, artifact_existence) 

2457 uris = [uri for uri, exists in artifact_existence.items() if exists] 

2458 

2459 # FUTURE UPGRADE: Implement a parallelized bulk remove. 

2460 log.debug("Removing %d artifacts from datastore that are unknown to datastore", len(uris)) 

2461 for uri in uris: 

2462 try: 

2463 uri.remove() 

2464 except Exception as e: 

2465 if ignore_errors: 

2466 log.debug("Artifact %s could not be removed: %s", uri, e) 

2467 continue 

2468 raise 

2469 

2470 # There is no point asking the code below to remove refs we 

2471 # know are missing so update it with the list of existing 

2472 # records. Try to retain one vs many logic. 

2473 if not existing_refs: 

2474 # Nothing more to do since none of the datasets were 

2475 # known to the datastore record table. 

2476 return 

2477 ref = list(existing_refs) 

2478 if len(ref) == 1: 

2479 ref = ref[0] 

2480 

2481 # Get file metadata and internal metadata 

2482 if not isinstance(ref, DatasetRef): 

2483 log.debug("Doing multi-dataset trash in datastore %s", self.name) 

2484 # Assumed to be an iterable of refs so bulk mode enabled. 

2485 try: 

2486 self.bridge.moveToTrash(ref, transaction=self._transaction) 

2487 except Exception as e: 

2488 if ignore_errors: 

2489 log.warning("Unexpected issue moving multiple datasets to trash: %s", e) 

2490 else: 

2491 raise 

2492 return 

2493 

2494 log.debug("Trashing dataset %s in datastore %s", ref, self.name) 

2495 

2496 fileLocations = self._get_dataset_locations_info(ref) 

2497 

2498 if not fileLocations: 

2499 err_msg = f"Requested dataset to trash ({ref}) is not known to datastore {self.name}" 

2500 if ignore_errors: 

2501 log.warning(err_msg) 

2502 return 

2503 else: 

2504 raise FileNotFoundError(err_msg) 

2505 

2506 for location, _ in fileLocations: 

2507 if not self._artifact_exists(location): 2507 ↛ 2508line 2507 didn't jump to line 2508 because the condition on line 2507 was never true

2508 err_msg = ( 

2509 f"Dataset is known to datastore {self.name} but " 

2510 f"associated artifact ({location.uri}) is missing" 

2511 ) 

2512 if ignore_errors: 

2513 log.warning(err_msg) 

2514 return 

2515 else: 

2516 raise FileNotFoundError(err_msg) 

2517 

2518 # Mark dataset as trashed 

2519 try: 

2520 self.bridge.moveToTrash([ref], transaction=self._transaction) 

2521 except Exception as e: 

2522 if ignore_errors: 

2523 log.warning( 

2524 "Attempted to mark dataset (%s) to be trashed in datastore %s " 

2525 "but encountered an error: %s", 

2526 ref, 

2527 self.name, 

2528 e, 

2529 ) 

2530 pass 

2531 else: 

2532 raise 

2533 

2534 def emptyTrash( 

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

2536 ) -> set[ResourcePath]: 

2537 """Remove all datasets from the trash. 

2538 

2539 Parameters 

2540 ---------- 

2541 ignore_errors : `bool` 

2542 If `True` return without error even if something went wrong. 

2543 Problems could occur if another process is simultaneously trying 

2544 to delete. 

2545 refs : `collections.abc.Collection` [ `DatasetRef` ] or `None` 

2546 Explicit list of datasets that can be removed from trash. If listed 

2547 datasets are not already stored in the trash table they will be 

2548 ignored. If `None` every entry in the trash table will be 

2549 processed. 

2550 dry_run : `bool`, optional 

2551 If `True`, the trash table will be queried and results reported 

2552 but no artifacts will be removed. 

2553 

2554 Returns 

2555 ------- 

2556 removed : `set` [ `lsst.resources.ResourcePath` ] 

2557 List of artifacts that were removed. 

2558 

2559 Notes 

2560 ----- 

2561 Will empty the records from the trash tables only if this call finishes 

2562 without raising. 

2563 """ 

2564 removed = set() 

2565 if refs: 

2566 selected_ids = {ref.id for ref in refs} 

2567 chunk_size = 50_000 

2568 n_chunks = math.ceil(len(selected_ids) / chunk_size) 

2569 chunk_num = 0 

2570 for chunk in chunk_iterable(selected_ids, chunk_size=chunk_size): 

2571 chunk_num += 1 

2572 if n_chunks == 1: 2572 ↛ 2579line 2572 didn't jump to line 2579 because the condition on line 2572 was always true

2573 log.verbose( 

2574 "Emptying datastore trash for %d dataset%s", 

2575 len(chunk), 

2576 "s" if len(chunk) != 1 else "", 

2577 ) 

2578 else: 

2579 log.verbose( 

2580 "Emptying datastore trash for chunk %d out of %d of size %d", 

2581 chunk_num, 

2582 n_chunks, 

2583 len(chunk), 

2584 ) 

2585 removed.update( 

2586 self._empty_trash_subset(ignore_errors=ignore_errors, selected_ids=chunk, dry_run=dry_run) 

2587 ) 

2588 else: 

2589 log.verbose("Emptying all trash in datastore %s", self.name) 

2590 removed = self._empty_trash_subset(ignore_errors=ignore_errors, dry_run=dry_run) 

2591 log.info( 

2592 "%sRemoved %d file artifact%s from datastore %s", 

2593 "Would have " if dry_run else "", 

2594 len(removed), 

2595 "s" if len(removed) != 1 else "", 

2596 self.name, 

2597 ) 

2598 return removed 

2599 

2600 @transactional 

2601 def _empty_trash_subset( 

2602 self, 

2603 *, 

2604 ignore_errors: bool = True, 

2605 selected_ids: Collection[DatasetId] | None = None, 

2606 dry_run: bool = False, 

2607 ) -> set[ResourcePath]: 

2608 """Empty trash table in transaction. 

2609 

2610 Parameters 

2611 ---------- 

2612 ignore_errors : `bool` 

2613 If `True` return without error even if something went wrong. 

2614 Problems could occur if another process is simultaneously trying 

2615 to delete. 

2616 selected_ids : `collections.abc.collection` [`DatasetId`] or `None` 

2617 Explicit list of dataset IDs that can be removed from the trash. 

2618 If listed datasets are not already included in the trash table 

2619 they will be ignored. If `None` every entry in the trash table 

2620 will be processed. 

2621 dry_run : `bool`, optional 

2622 If `True`, the trash table will be queried and results reported 

2623 but no artifacts will be removed. 

2624 

2625 Returns 

2626 ------- 

2627 removed : `set` [ `lsst.resources.ResourcePath` ] 

2628 Artifacts successfully removed. 

2629 

2630 Notes 

2631 ----- 

2632 Will empty the records from the trash tables only if this call finishes 

2633 without raising. 

2634 """ 

2635 # Context manager will empty trash iff we finish it without raising. 

2636 # It will also automatically delete the relevant rows from the 

2637 # trash table and the records table. 

2638 with self.bridge.emptyTrash( 

2639 self._table, 

2640 record_class=StoredFileInfo, 

2641 record_column="path", 

2642 selected_ids=selected_ids, 

2643 dry_run=dry_run, 

2644 ) as trash_data: 

2645 # Removing the artifacts themselves requires that the files are 

2646 # not also associated with refs that are not to be trashed. 

2647 # Therefore need to do a query with the file paths themselves 

2648 # and return all the refs associated with them. Can only delete 

2649 # a file if the refs to be trashed are the only refs associated 

2650 # with the file. 

2651 # This requires multiple copies of the trashed items 

2652 trashed, artifacts_to_keep = trash_data 

2653 

2654 # Assume that # in path means there are fragments involved. The 

2655 # fragments can not be handled by the emptyTrash bridge call 

2656 # so need to be processed independently. 

2657 # The generator has to be converted to a list for multiple 

2658 # iterations. Clean up the typing so that multiple isinstance 

2659 # tests aren't needed later. 

2660 trashed_list = [(ref, ninfo) for ref, ninfo in trashed if isinstance(ninfo, StoredFileInfo)] 

2661 

2662 if artifacts_to_keep is None or any("#" in info[1].path for info in trashed_list): 

2663 # The bridge is not helping us so have to work it out 

2664 # ourselves. This is not going to be as efficient. 

2665 # This mapping does not include the fragments. 

2666 if artifacts_to_keep is not None: 

2667 # This means we have already checked for non-fragment 

2668 # examples so can filter. 

2669 paths_to_check = {info.path for _, info in trashed_list if "#" in info.path} 

2670 else: 

2671 paths_to_check = {info.path for _, info in trashed_list} 

2672 

2673 path_map = self._refs_associated_with_artifacts(paths_to_check) 

2674 

2675 for ref, info in trashed_list: 

2676 path = info.artifact_path 

2677 # For disassembled composites in a Zip it is possible 

2678 # for the same path to correspond to the same dataset ref 

2679 # multiple times so trap for that. 

2680 if ref.id in path_map[path]: 

2681 path_map[path].remove(ref.id) 

2682 if not path_map[path]: 

2683 del path_map[path] 

2684 

2685 slow_artifacts_to_keep = set(path_map) 

2686 if artifacts_to_keep is not None: 

2687 artifacts_to_keep.update(slow_artifacts_to_keep) 

2688 else: 

2689 artifacts_to_keep = slow_artifacts_to_keep 

2690 

2691 n_direct = 0 

2692 artifacts_to_delete: set[ResourcePath] = set() 

2693 for ref, info in trashed_list: 

2694 # Should not happen for this implementation but need 

2695 # to keep mypy happy. 

2696 assert info is not None, f"Internal logic error in emptyTrash with ref {ref}." 

2697 

2698 if info.artifact_path in artifacts_to_keep: 

2699 # This is a multi-dataset artifact and we are not 

2700 # removing all associated refs. 

2701 continue 

2702 

2703 # Only trashed refs still known to datastore will be returned. 

2704 location = info.file_location(self.locationFactory) 

2705 

2706 if location.pathInStore.isabs(): 2706 ↛ 2707line 2706 didn't jump to line 2707 because the condition on line 2706 was never true

2707 n_direct += 1 

2708 continue 

2709 

2710 # Strip fragment before storing since it is the artifact 

2711 # we are deleting and we do not want repeats for every member 

2712 # in a zip. 

2713 artifacts_to_delete.add(location.uri.replace(fragment="")) 

2714 

2715 if n_direct > 0: 2715 ↛ 2716line 2715 didn't jump to line 2716 because the condition on line 2715 was never true

2716 s = "s" if n_direct != 1 else "" 

2717 log.verbose("Not deleting %d artifact%s using absolute URI%s", n_direct, s, s) 

2718 

2719 if artifacts_to_keep: 

2720 log.verbose( 

2721 "%d artifact%s %s not deleted because of association with other datasets", 

2722 len(artifacts_to_keep), 

2723 "s" if len(artifacts_to_keep) != 1 else "", 

2724 "were" if len(artifacts_to_keep) != 1 else "was", 

2725 ) 

2726 

2727 if not artifacts_to_delete: 

2728 return set() 

2729 

2730 # Now do the deleting. Special case the log message for a single 

2731 # artifact. 

2732 if len(artifacts_to_delete) == 1: 

2733 log.verbose( 

2734 "%s removing file artifact %s from datastore %s", 

2735 "Would be" if dry_run else "Now", 

2736 list(artifacts_to_delete)[0], 

2737 self.name, 

2738 ) 

2739 else: 

2740 log.verbose( 

2741 "%s removing %d file artifacts from datastore %s", 

2742 "Would be" if dry_run else "Now", 

2743 len(artifacts_to_delete), 

2744 self.name, 

2745 ) 

2746 

2747 # For dry-run mode do not attempt to search the file store for 

2748 # the artifacts to determine whether they exist or not. Simply 

2749 # report that an attempt would be made to delete them. Never 

2750 # report direct imports. 

2751 if dry_run: 

2752 return artifacts_to_delete 

2753 

2754 # Now remove the actual file artifacts. 

2755 remove_result = ResourcePath.mremove(artifacts_to_delete, do_raise=False) 

2756 

2757 removed: set[ResourcePath] = set() 

2758 exceptions: list[Exception] = [] 

2759 for uri, result in remove_result.items(): 

2760 if result.exception is None or isinstance(result.exception, FileNotFoundError): 2760 ↛ 2767line 2760 didn't jump to line 2767 because the condition on line 2760 was always true

2761 # File not existing is not an error since some other 

2762 # process might have been trying to clean it and we do not 

2763 # want to raise an error for a situation where the file 

2764 # is not there and we do not want it to be there. 

2765 removed.add(uri) 

2766 else: 

2767 exceptions.append(result.exception) 

2768 

2769 if exceptions: 2769 ↛ 2770line 2769 didn't jump to line 2770 because the condition on line 2769 was never true

2770 s_err = "s" if len(exceptions) != 1 else "" 

2771 e = ExceptionGroup(f"Error{s_err} removing {len(exceptions)} artifact{s_err}", exceptions) 

2772 if ignore_errors: 

2773 # Use a debug message here even though it's not 

2774 # a good situation. In some cases this can be 

2775 # caused by a race between user A and user B 

2776 # and neither of them has permissions for the 

2777 # other's files. Butler does not know about users 

2778 # and trash has no idea what collections these 

2779 # files were in (without guessing from a path). 

2780 log.debug( 

2781 "Encountered %d error%s removing %d artifact%s from datastore %s: %s", 

2782 len(exceptions), 

2783 s_err, 

2784 len(artifacts_to_delete), 

2785 "s" if len(artifacts_to_delete) != 1 else "", 

2786 self.name, 

2787 e, 

2788 ) 

2789 else: 

2790 raise e 

2791 return removed 

2792 

2793 @transactional 

2794 def transfer_from( 

2795 self, 

2796 source_records: FileTransferMap, 

2797 refs: Collection[DatasetRef], 

2798 transfer: str = "auto", 

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

2800 dry_run: bool = False, 

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

2802 log.verbose("Transferring %d datasets to %s", len(refs), self.name) 

2803 

2804 # Stop early if "direct" transfer mode is requested. That would 

2805 # require that the URI inside the source datastore should be stored 

2806 # directly in the target datastore, which seems unlikely to be useful 

2807 # since at any moment the source datastore could delete the file. 

2808 if transfer in ("direct", "split"): 

2809 raise ValueError( 

2810 f"Can not transfer from a source datastore using {transfer} mode since" 

2811 " those files are controlled by the other datastore." 

2812 ) 

2813 

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

2815 return set(), set() 

2816 

2817 # Empty existence lookup if none given. 

2818 if artifact_existence is None: 

2819 artifact_existence = {} 

2820 

2821 # In order to handle disassembled composites the code works 

2822 # at the records level since it can assume that internal APIs 

2823 # can be used. 

2824 # - If the record already exists in the destination this is assumed 

2825 # to be okay. 

2826 # - If there is no record but the source and destination URIs are 

2827 # identical no transfer is done but the record is added. 

2828 # - If the source record refers to an absolute URI currently assume 

2829 # that that URI should remain absolute and will be visible to the 

2830 # destination butler. May need to have a flag to indicate whether 

2831 # the dataset should be transferred. This will only happen if 

2832 # the detached Butler has had a local ingest. 

2833 

2834 # See if we already have these records 

2835 log.verbose("Looking up existing datastore records in target %s for %d refs", self.name, len(refs)) 

2836 target_records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True) 

2837 

2838 # The artifacts to register 

2839 artifacts = [] 

2840 

2841 # Refs that already exist 

2842 already_present = [] 

2843 

2844 # Refs that were rejected by this datastore. 

2845 rejected = set() 

2846 

2847 # Refs that were transferred successfully. 

2848 accepted = set() 

2849 

2850 # Record each time we have done a "direct" transfer. 

2851 direct_transfers = [] 

2852 

2853 # Keep track of all the file transfers that are required. 

2854 from_to: list[tuple[ResourcePath, ResourcePath]] = [] 

2855 

2856 # Now can transfer the artifacts 

2857 log.verbose("Transferring artifacts") 

2858 for ref in refs: 

2859 if not self.constraints.isAcceptable(ref): 2859 ↛ 2861line 2859 didn't jump to line 2861 because the condition on line 2859 was never true

2860 # This datastore should not be accepting this dataset. 

2861 rejected.add(ref) 

2862 continue 

2863 

2864 accepted.add(ref) 

2865 

2866 if ref.id in target_records: 

2867 # Already have an artifact for this. 

2868 already_present.append(ref) 

2869 continue 

2870 

2871 # mypy needs to know these are always resolved refs 

2872 for transfer_info in source_records.get(ref.id, []): 

2873 info = transfer_info.file_info 

2874 source_location = transfer_info.location 

2875 target_location = info.file_location(self.locationFactory) 

2876 if transfer == "unsafe_direct": 

2877 # Use the existing file from the source location in place, 

2878 # by recording the absolute URI in the target DB. This is 

2879 # "unsafe" because the file could be deleted from the 

2880 # source Butler at any time, leaving a dangling reference. 

2881 source_location = source_location.toAbsolute() 

2882 direct_transfers.append(source_location) 

2883 info = info.update(path=str(source_location.uri)) 

2884 elif source_location == target_location and not source_location.pathInStore.isabs(): 2884 ↛ 2887line 2884 didn't jump to line 2887 because the condition on line 2884 was never true

2885 # Artifact is already in the target location. 

2886 # (which is how execution butler currently runs) 

2887 pass 

2888 else: 

2889 if target_location.pathInStore.isabs(): 

2890 # Just because we can see the artifact when running 

2891 # the transfer doesn't mean it will be generally 

2892 # accessible to a user of this butler. Need to decide 

2893 # what to do about an absolute path. 

2894 if transfer == "auto": 

2895 # For "auto" transfers we allow the absolute URI 

2896 # to be recorded in the target datastore. 

2897 direct_transfers.append(source_location) 

2898 else: 

2899 # The user is explicitly requesting a transfer 

2900 # even for an absolute URI. This requires us to 

2901 # calculate the target path. 

2902 template_ref = ref 

2903 if info.component: 2903 ↛ 2904line 2903 didn't jump to line 2904 because the condition on line 2903 was never true

2904 template_ref = ref.makeComponentRef(info.component) 

2905 target_location = self._calculate_ingested_datastore_name( 

2906 source_location.uri, 

2907 template_ref, 

2908 ) 

2909 

2910 info = info.update(path=target_location.pathInStore.path) 

2911 

2912 # Need to transfer it to the new location. 

2913 from_to.append((source_location.uri, target_location.uri)) 

2914 

2915 artifacts.append((ref, info)) 

2916 

2917 # Do the file transfers in bulk. 

2918 # Assume we should always overwrite. If the artifact 

2919 # is there this might indicate that a previous transfer 

2920 # was interrupted but was not able to be rolled back 

2921 # completely (eg pre-emption) so follow Datastore default 

2922 # and overwrite. Do not copy if we are in dry-run mode. 

2923 if dry_run: 

2924 log.info("Would be copying %d file artifacts", len(from_to)) 

2925 else: 

2926 log.verbose("Copying %d file artifacts", len(from_to)) 

2927 with time_this(log, msg="Transferring datasets into datastore", level=VERBOSE): 

2928 ResourcePath.mtransfer( 

2929 transfer, 

2930 from_to, 

2931 overwrite=True, 

2932 transaction=self._transaction, 

2933 ) 

2934 

2935 if direct_transfers: 

2936 log.info( 

2937 "Transfer request for an outside-datastore artifact with absolute URI done %d time%s", 

2938 len(direct_transfers), 

2939 "" if len(direct_transfers) == 1 else "s", 

2940 ) 

2941 

2942 # We are overwriting previous datasets that may have already 

2943 # existed. We therefore should ensure that we force the 

2944 # datastore records to agree. Note that this can potentially lead 

2945 # to difficulties if the dataset has previously been ingested 

2946 # disassembled and is somehow now assembled, or vice versa. 

2947 if not dry_run: 

2948 log.verbose("Registering datastore records in database") 

2949 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.REPLACE) 

2950 

2951 if already_present: 

2952 n_skipped = len(already_present) 

2953 log.info( 

2954 "Skipped transfer of %d dataset%s already present in datastore", 

2955 n_skipped, 

2956 "" if n_skipped == 1 else "s", 

2957 ) 

2958 

2959 log.verbose( 

2960 "Finished transfer_from to %s with %d accepted, %d rejected", 

2961 self.name, 

2962 len(accepted), 

2963 len(rejected), 

2964 ) 

2965 return accepted, rejected 

2966 

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

2968 source_records = self._get_stored_records_associated_with_refs( 

2969 [FakeDatasetRef(id) for id in dataset_ids], ignore_datastore_records=True 

2970 ) 

2971 return self._convert_stored_file_info_to_file_transfer_record(source_records) 

2972 

2973 def locate_missing_files_for_transfer( 

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

2975 ) -> FileTransferMap: 

2976 missing_ids = {ref.id for ref in refs} 

2977 # Missing IDs can be okay if that datastore has allowed 

2978 # gets based on file existence. Should we transfer what we can 

2979 # or complain about it and warn? 

2980 if not self.trustGetRequest: 

2981 return {} 

2982 

2983 found_records = self._find_missing_records( 

2984 refs, missing_ids, artifact_existence, warn_for_missing=False 

2985 ) 

2986 return self._convert_stored_file_info_to_file_transfer_record(found_records) 

2987 

2988 def _convert_stored_file_info_to_file_transfer_record( 

2989 self, info_map: dict[DatasetId, list[StoredFileInfo]] 

2990 ) -> FileTransferMap: 

2991 output: dict[DatasetId, list[FileTransferRecord]] = {} 

2992 for k, file_info_list in info_map.items(): 

2993 output[k] = [ 

2994 FileTransferRecord(file_info=info, location=info.file_location(self.locationFactory)) 

2995 for info in file_info_list 

2996 ] 

2997 return output 

2998 

2999 @transactional 

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

3001 # Docstring inherited. 

3002 refs = list(refs) 

3003 self.bridge.forget(refs) 

3004 self._table.delete(["dataset_id"], *[{"dataset_id": ref.id} for ref in refs]) 

3005 

3006 def validateConfiguration( 

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

3008 ) -> None: 

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

3010 

3011 Parameters 

3012 ---------- 

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

3014 | `StorageClass`] 

3015 Entities to test against this configuration. Can be differing 

3016 types. 

3017 logFailures : `bool`, optional 

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

3019 detected. 

3020 

3021 Returns 

3022 ------- 

3023 None 

3024 

3025 Raises 

3026 ------ 

3027 DatastoreValidationError 

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

3029 All the problems are reported in a single exception. 

3030 

3031 Notes 

3032 ----- 

3033 This method checks that all the supplied entities have valid file 

3034 templates and also have formatters defined. 

3035 """ 

3036 templateFailed = None 

3037 try: 

3038 self.templates.validateTemplates(entities, logFailures=logFailures) 

3039 except FileTemplateValidationError as e: 

3040 templateFailed = str(e) 

3041 

3042 formatterFailed = [] 

3043 for entity in entities: 

3044 try: 

3045 self.formatterFactory.getFormatterClass(entity) 

3046 except KeyError as e: 

3047 formatterFailed.append(str(e)) 

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

3049 log.critical("Formatter failure: %s", e) 

3050 

3051 if templateFailed or formatterFailed: 

3052 messages = [] 

3053 if templateFailed: 3053 ↛ 3054line 3053 didn't jump to line 3054 because the condition on line 3053 was never true

3054 messages.append(templateFailed) 

3055 if formatterFailed: 3055 ↛ 3057line 3055 didn't jump to line 3057 because the condition on line 3055 was always true

3056 messages.append(",".join(formatterFailed)) 

3057 msg = ";\n".join(messages) 

3058 raise DatastoreValidationError(msg) 

3059 

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

3061 # Docstring is inherited from base class 

3062 return ( 

3063 self.templates.getLookupKeys() 

3064 | self.formatterFactory.getLookupKeys() 

3065 | self.constraints.getLookupKeys() 

3066 ) 

3067 

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

3069 # Docstring is inherited from base class 

3070 # The key can be valid in either formatters or templates so we can 

3071 # only check the template if it exists 

3072 if lookupKey in self.templates: 

3073 try: 

3074 self.templates[lookupKey].validateTemplate(entity) 

3075 except FileTemplateValidationError as e: 

3076 raise DatastoreValidationError(e) from e 

3077 

3078 def export( 

3079 self, 

3080 refs: Iterable[DatasetRef], 

3081 *, 

3082 directory: ResourcePathExpression | None = None, 

3083 transfer: str | None = "auto", 

3084 ) -> Iterable[FileDataset]: 

3085 # Docstring inherited from Datastore.export. 

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

3087 transfer = None 

3088 

3089 if transfer is not None and transfer != "direct" and directory is None: 

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

3091 

3092 if transfer == "move": 

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

3094 

3095 # Force the directory to be a URI object 

3096 directoryUri: ResourcePath | None = None 

3097 if directory is not None: 

3098 directoryUri = ResourcePath(directory, forceDirectory=True) 

3099 

3100 if transfer is not None and directoryUri is not None and not directoryUri.exists(): 3100 ↛ 3102line 3100 didn't jump to line 3102 because the condition on line 3100 was never true

3101 # mypy needs the second test 

3102 raise FileNotFoundError(f"Export location {directory} does not exist") 

3103 

3104 progress = Progress("lsst.daf.butler.datastores.FileDatastore.export", level=logging.DEBUG) 

3105 for ref in progress.wrap(refs, "Exporting dataset files"): 

3106 fileLocations = self._get_dataset_locations_info(ref) 

3107 if not fileLocations: 

3108 raise FileNotFoundError(f"Could not retrieve dataset {ref}.") 

3109 # For now we can not export disassembled datasets 

3110 if len(fileLocations) > 1: 

3111 raise NotImplementedError(f"Can not export disassembled datasets such as {ref}") 

3112 location, storedFileInfo = fileLocations[0] 

3113 

3114 pathInStore = location.pathInStore.path 

3115 if transfer is None: 

3116 # TODO: do we also need to return the readStorageClass somehow? 

3117 # We will use the path in store directly. If this is an 

3118 # absolute URI, preserve it. 

3119 if location.pathInStore.isabs(): 3119 ↛ 3120line 3119 didn't jump to line 3120 because the condition on line 3119 was never true

3120 pathInStore = str(location.uri) 

3121 elif transfer == "direct": 

3122 # Use full URIs to the remote store in the export 

3123 pathInStore = str(location.uri) 

3124 else: 

3125 # mypy needs help 

3126 assert directoryUri is not None, "directoryUri must be defined to get here" 

3127 storeUri = ResourcePath(location.uri, forceDirectory=False) 

3128 

3129 # if the datastore has an absolute URI to a resource, we 

3130 # have two options: 

3131 # 1. Keep the absolute URI in the exported YAML 

3132 # 2. Allocate a new name in the local datastore and transfer 

3133 # it. 

3134 # For now go with option 2 

3135 if location.pathInStore.isabs(): 3135 ↛ 3136line 3135 didn't jump to line 3136 because the condition on line 3135 was never true

3136 template = self.templates.getTemplate(ref) 

3137 newURI = ResourcePath(template.format(ref), forceAbsolute=False, forceDirectory=False) 

3138 pathInStore = str(newURI.updatedExtension(location.pathInStore.getExtension())) 

3139 

3140 exportUri = directoryUri.join(pathInStore) 

3141 exportUri.transfer_from(storeUri, transfer=transfer) 

3142 

3143 yield FileDataset(refs=[ref], path=pathInStore, formatter=storedFileInfo.formatter) 

3144 

3145 @staticmethod 

3146 def computeChecksum(uri: ResourcePath, algorithm: str = "blake2b", block_size: int = 8192) -> str | None: 

3147 """Compute the checksum of the supplied file. 

3148 

3149 Parameters 

3150 ---------- 

3151 uri : `lsst.resources.ResourcePath` 

3152 Name of resource to calculate checksum from. 

3153 algorithm : `str`, optional 

3154 Name of algorithm to use. Must be one of the algorithms supported 

3155 by :py:class`hashlib`. 

3156 block_size : `int` 

3157 Number of bytes to read from file at one time. 

3158 

3159 Returns 

3160 ------- 

3161 hexdigest : `str` 

3162 Hex digest of the file. 

3163 

3164 Notes 

3165 ----- 

3166 Currently returns None if the URI is for a remote resource. 

3167 """ 

3168 if algorithm not in hashlib.algorithms_guaranteed: 3168 ↛ 3169line 3168 didn't jump to line 3169 because the condition on line 3168 was never true

3169 raise NameError(f"The specified algorithm '{algorithm}' is not supported by hashlib") 

3170 

3171 if not uri.isLocal: 3171 ↛ 3172line 3171 didn't jump to line 3172 because the condition on line 3171 was never true

3172 return None 

3173 

3174 hasher = hashlib.new(algorithm) 

3175 

3176 with uri.as_local() as local_uri, open(local_uri.ospath, "rb") as f: 

3177 for chunk in iter(lambda: f.read(block_size), b""): 

3178 hasher.update(chunk) 

3179 

3180 return hasher.hexdigest() 

3181 

3182 def needs_expanded_data_ids( 

3183 self, 

3184 transfer: str | None, 

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

3186 ) -> bool: 

3187 # Docstring inherited. 

3188 # This _could_ also use entity to inspect whether the filename template 

3189 # involves placeholders other than the required dimensions for its 

3190 # dataset type, but that's not necessary for correctness; it just 

3191 # enables more optimizations (perhaps only in theory). 

3192 return transfer not in ("direct", None) 

3193 

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

3195 # Docstring inherited from the base class. 

3196 record_data = data.get(self.name) 

3197 if not record_data: 3197 ↛ 3198line 3197 didn't jump to line 3198 because the condition on line 3197 was never true

3198 return 

3199 

3200 self._bridge.insert(FakeDatasetRef(dataset_id) for dataset_id in record_data.records) 

3201 

3202 # TODO: Verify that there are no unexpected table names in the dict? 

3203 unpacked_records = [] 

3204 for dataset_id, dataset_data in record_data.records.items(): 

3205 records = dataset_data.get(self._table.name) 

3206 if records: 3206 ↛ 3204line 3206 didn't jump to line 3204 because the condition on line 3206 was always true

3207 for info in records: 

3208 assert isinstance(info, StoredFileInfo), "Expecting StoredFileInfo records" 

3209 unpacked_records.append(info.to_record(dataset_id=dataset_id)) 

3210 if unpacked_records: 

3211 self._table.insert(*unpacked_records, transaction=self._transaction) 

3212 

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

3214 # Docstring inherited from the base class. 

3215 

3216 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {} 

3217 for batch in self._export_rows([ref.id for ref in refs]): 

3218 for row in batch: 

3219 info: StoredDatastoreItemInfo = StoredFileInfo.from_record(row) 

3220 dataset_records = records.setdefault(row["dataset_id"], {}) 

3221 dataset_records.setdefault(self._table.name, []).append(info) 

3222 

3223 record_data = DatastoreRecordData(records=records) 

3224 return {self.name: record_data} 

3225 

3226 def _export_rows(self, datasets: Collection[DatasetId]) -> Iterator[Sequence[Mapping[str, Any]]]: 

3227 # This call to 'bridge.check' filters out "partially deleted" datasets. 

3228 # Specifically, ones in the unusual edge state that: 

3229 # 1. They have an entry in the registry dataset tables 

3230 # 2. They were "trashed" from the datastore, so they are not 

3231 # present in the "dataset_location" table.) 

3232 # 3. But the trash has not been "emptied", so there are still entries 

3233 # in the "opaque" datastore records table. 

3234 # 

3235 # As far as I can tell, this can only occur in the case of a concurrent 

3236 # or aborted call to `Butler.pruneDatasets(unstore=True, purge=False)`. 

3237 # Datasets (with or without files existing on disk) can persist in 

3238 # this zombie state indefinitely, until someone manually empties 

3239 # the trash. 

3240 found_ids = self._bridge.check(datasets) 

3241 return self._table.fetch_batches(dataset_id=found_ids) 

3242 

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

3244 # Docstring inherited from the base class. 

3245 

3246 tables: list[DatastoreRecordTable] = [] 

3247 for batch in self._export_rows(datasets): 

3248 file_info = StoredFileInfoTable.from_records(batch) 

3249 tables.append(DatastoreRecordTable.from_stored_file_info_table(self.name, file_info)) 

3250 return DatastoreRecordTable.combine(tables) 

3251 

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

3253 # Docstring inherited from the base class. 

3254 

3255 records = table.to_stored_file_info_table().to_records() 

3256 dataset_ids = [FakeDatasetRef(record["dataset_id"]) for record in records] 

3257 if len(records) > 0: 3257 ↛ exitline 3257 didn't return from function 'import_table' because the condition on line 3257 was always true

3258 self._bridge.insert(dataset_ids) 

3259 self._table.insert(*records, transaction=self._transaction) 

3260 

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

3262 # Docstring inherited from the base class. 

3263 refs = [self._cast_storage_class(ref) for ref in refs] 

3264 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {} 

3265 for ref in refs: 

3266 if not self.constraints.isAcceptable(ref): 3266 ↛ 3267line 3266 didn't jump to line 3267 because the condition on line 3266 was never true

3267 continue 

3268 fileLocations = self._get_expected_dataset_locations_info(ref) 

3269 if not fileLocations: 3269 ↛ 3270line 3269 didn't jump to line 3270 because the condition on line 3269 was never true

3270 continue 

3271 dataset_records = records.setdefault(ref.id, {}) 

3272 dataset_records.setdefault(self._table.name, []) 

3273 for _, storedFileInfo in fileLocations: 

3274 dataset_records[self._table.name].append(storedFileInfo) 

3275 

3276 record_data = DatastoreRecordData(records=records) 

3277 return {self.name: record_data} 

3278 

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

3280 # Docstring inherited from the base class. 

3281 self._retrieve_dataset_method = method 

3282 

3283 def _cast_storage_class(self, ref: DatasetRef) -> DatasetRef: 

3284 """Update dataset reference to use the storage class from registry.""" 

3285 if self._retrieve_dataset_method is None: 

3286 # We could raise an exception here but unit tests do not define 

3287 # this method. 

3288 return ref 

3289 dataset_type = self._retrieve_dataset_method(ref.datasetType.name) 

3290 if dataset_type is not None: 3290 ↛ 3292line 3290 didn't jump to line 3292 because the condition on line 3290 was always true

3291 ref = ref.overrideStorageClass(dataset_type.storageClass_name) 

3292 return ref 

3293 

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

3295 # Docstring inherited from the base class. 

3296 return {self._opaque_table_name: DatastoreOpaqueTable(self.makeTableSpec(), StoredFileInfo)}