Coverage for python/lsst/daf/butler/datastores/fileDatastore.py: 8%
1068 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-25 00:34 -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/>.
28"""Generic file-based datastore code."""
30from __future__ import annotations
32__all__ = ("FileDatastore",)
34import contextlib
35import hashlib
36import logging
37import math
38from collections import defaultdict
39from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
40from typing import TYPE_CHECKING, Any, ClassVar, cast
42from sqlalchemy import BigInteger, String
44from lsst.daf.butler import (
45 Config,
46 DatasetDatastoreRecords,
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)
70from lsst.daf.butler.datastore.cache_manager import (
71 AbstractDatastoreCacheManager,
72 DatastoreCacheManager,
73 DatastoreDisabledCacheManager,
74)
75from lsst.daf.butler.datastore.composites import CompositesMap
76from lsst.daf.butler.datastore.file_templates import FileTemplates, FileTemplateValidationError
77from lsst.daf.butler.datastore.generic_base import GenericBaseDatastore
78from lsst.daf.butler.datastore.record_data import DatastoreRecordData
79from lsst.daf.butler.datastore.stored_file_info import StoredDatastoreItemInfo, StoredFileInfo
80from lsst.daf.butler.datastores.file_datastore.get import (
81 DatasetLocationInformation,
82 DatastoreFileGetInformation,
83 generate_datastore_get_information,
84 get_dataset_as_python_object_from_get_info,
85)
86from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import (
87 ArtifactIndexInfo,
88 ZipIndex,
89 determine_destination_for_retrieved_artifact,
90 unpack_zips,
91)
92from lsst.daf.butler.registry.interfaces import (
93 DatabaseInsertMode,
94 DatastoreRegistryBridge,
95 FakeDatasetRef,
96 ReadOnlyDatabaseError,
97)
98from lsst.daf.butler.repo_relocation import replaceRoot
99from lsst.daf.butler.utils import transactional
100from lsst.resources import ResourcePath, ResourcePathExpression
101from lsst.utils.introspection import get_class_of, get_full_type_name
102from lsst.utils.iteration import chunk_iterable
104# For VERBOSE logging usage.
105from lsst.utils.logging import VERBOSE, getLogger
106from lsst.utils.timer import time_this
108from ..datastore import FileTransferMap, FileTransferRecord
110if TYPE_CHECKING:
111 from lsst.daf.butler import DatasetProvenance, LookupKey
112 from lsst.daf.butler.registry.interfaces import DatasetIdRef, DatastoreRegistryBridgeManager
114log = getLogger(__name__)
117class _IngestPrepData(Datastore.IngestPrepData):
118 """Helper class for FileDatastore ingest implementation.
120 Parameters
121 ----------
122 datasets : `~collections.abc.Iterable` of `FileDataset`
123 Files to be ingested by this datastore.
124 """
126 def __init__(self, datasets: Iterable[FileDataset]):
127 super().__init__(ref for dataset in datasets for ref in dataset.refs)
128 self.datasets = datasets
131class FileDatastore(GenericBaseDatastore[StoredFileInfo]):
132 """Generic Datastore for file-based implementations.
134 Should always be sub-classed since key abstract methods are missing.
136 Parameters
137 ----------
138 config : `DatastoreConfig` or `str`
139 Configuration as either a `Config` object or URI to file.
140 bridgeManager : `DatastoreRegistryBridgeManager`
141 Object that manages the interface between `Registry` and datastores.
142 root : `ResourcePath`
143 Root directory URI of this `Datastore`.
144 formatterFactory : `FormatterFactory`
145 Factory for creating instances of formatters.
146 templates : `FileTemplates`
147 File templates that can be used by this `Datastore`.
148 composites : `CompositesMap`
149 Determines whether a dataset should be disassembled on put.
150 trustGetRequest : `bool`
151 Determine whether we can fall back to configuration if a requested
152 dataset is not known to registry.
154 Raises
155 ------
156 ValueError
157 If root location does not exist and ``create`` is `False` in the
158 configuration.
159 """
161 defaultConfigFile: ClassVar[str | None] = None
162 """Path to configuration defaults. Accessed within the ``config`` resource
163 or relative to a search path. Can be None if no defaults specified.
164 """
166 root: ResourcePath
167 """Root directory URI of this `Datastore`."""
169 locationFactory: LocationFactory
170 """Factory for creating locations relative to the datastore root."""
172 formatterFactory: FormatterFactory
173 """Factory for creating instances of formatters."""
175 templates: FileTemplates
176 """File templates that can be used by this `Datastore`."""
178 composites: CompositesMap
179 """Determines whether a dataset should be disassembled on put."""
181 defaultConfigFile = "datastores/fileDatastore.yaml"
182 """Path to configuration defaults. Accessed within the ``config`` resource
183 or relative to a search path. Can be None if no defaults specified.
184 """
186 _retrieve_dataset_method: Callable[[str], DatasetType | None] | None = None
187 """Callable that is used in trusted mode to retrieve registry definition
188 of a named dataset type.
189 """
191 @classmethod
192 def setConfigRoot(cls, root: str, config: Config, full: Config, overwrite: bool = True) -> None:
193 """Set any filesystem-dependent config options for this Datastore to
194 be appropriate for a new empty repository with the given root.
196 Parameters
197 ----------
198 root : `str`
199 URI to the root of the data repository.
200 config : `Config`
201 A `Config` to update. Only the subset understood by
202 this component will be updated. Will not expand
203 defaults.
204 full : `Config`
205 A complete config with all defaults expanded that can be
206 converted to a `DatastoreConfig`. Read-only and will not be
207 modified by this method.
208 Repository-specific options that should not be obtained
209 from defaults when Butler instances are constructed
210 should be copied from ``full`` to ``config``.
211 overwrite : `bool`, optional
212 If `False`, do not modify a value in ``config`` if the value
213 already exists. Default is always to overwrite with the provided
214 ``root``.
216 Notes
217 -----
218 If a keyword is explicitly defined in the supplied ``config`` it
219 will not be overridden by this method if ``overwrite`` is `False`.
220 This allows explicit values set in external configs to be retained.
221 """
222 Config.updateParameters(
223 DatastoreConfig,
224 config,
225 full,
226 toUpdate={"root": root},
227 toCopy=("cls", ("records", "table")),
228 overwrite=overwrite,
229 )
231 @classmethod
232 def makeTableSpec(cls) -> ddl.TableSpec:
233 return ddl.TableSpec(
234 fields=[
235 ddl.FieldSpec(name="dataset_id", dtype=ddl.GUID, primaryKey=True),
236 ddl.FieldSpec(name="path", dtype=String, length=256, nullable=False),
237 ddl.FieldSpec(name="formatter", dtype=String, length=128, nullable=False),
238 ddl.FieldSpec(name="storage_class", dtype=String, length=64, nullable=False),
239 # Use empty string to indicate no component
240 ddl.FieldSpec(name="component", dtype=String, length=32, primaryKey=True),
241 # TODO: should checksum be Base64Bytes instead?
242 ddl.FieldSpec(name="checksum", dtype=String, length=128, nullable=True),
243 ddl.FieldSpec(name="file_size", dtype=BigInteger, nullable=True),
244 ],
245 unique=frozenset(),
246 indexes=[ddl.IndexSpec("path")],
247 )
249 def __init__(
250 self,
251 config: DatastoreConfig,
252 bridgeManager: DatastoreRegistryBridgeManager,
253 root: ResourcePath,
254 formatterFactory: FormatterFactory,
255 templates: FileTemplates,
256 composites: CompositesMap,
257 trustGetRequest: bool,
258 ):
259 super().__init__(config, bridgeManager)
260 self.root = ResourcePath(root)
261 self.formatterFactory = formatterFactory
262 self.templates = templates
263 self.composites = composites
264 self.trustGetRequest = trustGetRequest
266 # Name ourselves either using an explicit name or a name
267 # derived from the (unexpanded) root
268 if "name" in self.config:
269 self.name = self.config["name"]
270 else:
271 # We use the unexpanded root in the name to indicate that this
272 # datastore can be moved without having to update registry.
273 self.name = "{}@{}".format(type(self).__name__, self.config["root"])
275 self.locationFactory = LocationFactory(self.root)
277 self._opaque_table_name = self.config["records", "table"]
278 try:
279 # Storage of paths and formatters, keyed by dataset_id
280 self._table = bridgeManager.opaque.register(self._opaque_table_name, self.makeTableSpec())
281 # Interface to Registry.
282 self._bridge = bridgeManager.register(self.name)
283 except ReadOnlyDatabaseError:
284 # If the database is read only and we just tried and failed to
285 # create a table, it means someone is trying to create a read-only
286 # butler client for an empty repo. That should be okay, as long
287 # as they then try to get any datasets before some other client
288 # creates the table. Chances are they're just validating
289 # configuration.
290 pass
292 # Determine whether checksums should be used - default to False
293 self.useChecksum = self.config.get("checksum", False)
295 # Create a cache manager
296 self.cacheManager: AbstractDatastoreCacheManager
297 if "cached" in self.config:
298 self.cacheManager = DatastoreCacheManager(self.config["cached"], universe=bridgeManager.universe)
299 else:
300 self.cacheManager = DatastoreDisabledCacheManager("", universe=bridgeManager.universe)
302 self.universe = bridgeManager.universe
304 @classmethod
305 def _create_from_config(
306 cls,
307 config: DatastoreConfig,
308 bridgeManager: DatastoreRegistryBridgeManager,
309 butlerRoot: ResourcePathExpression | None,
310 ) -> FileDatastore:
311 if "root" not in config:
312 raise ValueError("No root directory specified in configuration")
314 # Support repository relocation in config
315 # Existence of self.root is checked in subclass
316 root = ResourcePath(replaceRoot(config["root"], butlerRoot), forceDirectory=True, forceAbsolute=True)
318 # Now associate formatters with storage classes
319 formatterFactory = FormatterFactory()
320 formatterFactory.registerFormatters(config["formatters"], universe=bridgeManager.universe)
322 # Read the file naming templates
323 templates = FileTemplates(config["templates"], universe=bridgeManager.universe)
325 # See if composites should be disassembled
326 composites = CompositesMap(config["composites"], universe=bridgeManager.universe)
328 # Determine whether we can fall back to configuration if a
329 # requested dataset is not known to registry
330 trustGetRequest = config.get("trust_get_request", False)
332 self = FileDatastore(
333 config, bridgeManager, root, formatterFactory, templates, composites, trustGetRequest
334 )
336 # Check existence and create directory structure if necessary.
337 #
338 # The concept of a 'root directory' is problematic for some resource
339 # path types that don't necessarily support the concept of a directory
340 # (http, s3, gs... basically anything that isn't a local filesystem or
341 # WebDAV.)
342 # On these resource paths an object representing the
343 # "root" directory may not exist even though files under the root do,
344 # and in a read-only repository we will be unable to create it.
345 # So we only immediately verify the root for local filesystems,
346 # the only case where this check will definitely not give a false
347 # negative.
348 if self.root.isLocal and not self.root.exists():
349 if "create" not in self.config or not self.config["create"]:
350 raise ValueError(f"No valid root and not allowed to create one at: {self.root}")
351 try:
352 self.root.mkdir()
353 except Exception as e:
354 raise ValueError(
355 f"Can not create datastore root '{self.root}', check permissions. Got error: {e}"
356 ) from e
358 return self
360 def clone(self, bridgeManager: DatastoreRegistryBridgeManager) -> Datastore:
361 return FileDatastore(
362 self.config,
363 bridgeManager,
364 self.root,
365 self.formatterFactory,
366 self.templates,
367 self.composites,
368 self.trustGetRequest,
369 )
371 def __str__(self) -> str:
372 return str(self.root)
374 @property
375 def bridge(self) -> DatastoreRegistryBridge:
376 return self._bridge
378 @property
379 def roots(self) -> dict[str, ResourcePath | None]:
380 # Docstring inherited.
381 return {self.name: self.root}
383 def _set_trust_mode(self, mode: bool) -> None:
384 self.trustGetRequest = mode
386 def _artifact_exists(self, location: Location) -> bool:
387 """Check that an artifact exists in this datastore at the specified
388 location.
390 Parameters
391 ----------
392 location : `Location`
393 Expected location of the artifact associated with this datastore.
395 Returns
396 -------
397 exists : `bool`
398 True if the location can be found, false otherwise.
399 """
400 log.debug("Checking if resource exists: %s", location.uri)
401 return location.uri.exists()
403 def addStoredItemInfo(
404 self,
405 refs: Iterable[DatasetRef],
406 infos: Iterable[StoredFileInfo],
407 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT,
408 ) -> None:
409 """Record internal storage information associated with one or more
410 datasets.
412 Parameters
413 ----------
414 refs : sequence of `DatasetRef`
415 The datasets that have been stored.
416 infos : sequence of `StoredDatastoreItemInfo`
417 Metadata associated with the stored datasets.
418 insert_mode : `~lsst.daf.butler.registry.interfaces.DatabaseInsertMode`
419 Mode to use to insert the new records into the table. The
420 options are ``INSERT`` (error if pre-existing), ``REPLACE``
421 (replace content with new values), and ``ENSURE`` (skip if the row
422 already exists).
423 """
424 records = [
425 info.rebase(ref).to_record(dataset_id=ref.id) for ref, info in zip(refs, infos, strict=True)
426 ]
427 match insert_mode:
428 case DatabaseInsertMode.INSERT:
429 self._table.insert(*records, transaction=self._transaction)
430 case DatabaseInsertMode.ENSURE:
431 self._table.ensure(*records, transaction=self._transaction)
432 case DatabaseInsertMode.REPLACE:
433 self._table.replace(*records, transaction=self._transaction)
434 case _:
435 raise ValueError(f"Unknown insert mode of '{insert_mode}'")
437 def getStoredItemsInfo(
438 self, ref: DatasetIdRef, ignore_datastore_records: bool = False
439 ) -> list[StoredFileInfo]:
440 """Retrieve information associated with files stored in this
441 `Datastore` associated with this dataset ref.
443 Parameters
444 ----------
445 ref : `DatasetRef`
446 The dataset that is to be queried.
447 ignore_datastore_records : `bool`
448 If `True` then do not use datastore records stored in refs.
450 Returns
451 -------
452 items : `~collections.abc.Iterable` [`StoredDatastoreItemInfo`]
453 Stored information about the files and associated formatters
454 associated with this dataset. Only one file will be returned
455 if the dataset has not been disassembled. Can return an empty
456 list if no matching datasets can be found.
457 """
458 # Try to get them from the ref first.
459 if ref._datastore_records is not None and not ignore_datastore_records:
460 ref_records = ref._datastore_records.get(self._table.name, [])
461 # Need to make sure they have correct type.
462 for record in ref_records:
463 if not isinstance(record, StoredFileInfo):
464 raise TypeError(f"Datastore record has unexpected type {record.__class__.__name__}")
465 return cast(list[StoredFileInfo], ref_records)
467 # Look for the dataset_id -- there might be multiple matches
468 # if we have disassembled the dataset.
469 records = self._table.fetch(dataset_id=ref.id)
470 return [StoredFileInfo.from_record(record) for record in records]
472 def _register_datasets(
473 self,
474 refsAndInfos: Iterable[tuple[DatasetRef, StoredFileInfo]],
475 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT,
476 ) -> None:
477 """Update registry to indicate that one or more datasets have been
478 stored.
480 Parameters
481 ----------
482 refsAndInfos : sequence `tuple` [`DatasetRef`,
483 `StoredDatastoreItemInfo`]
484 Datasets to register and the internal datastore metadata associated
485 with them.
486 insert_mode : `str`, optional
487 Indicate whether the new records should be new ("insert", default),
488 or allowed to exists ("ensure") or be replaced if already present
489 ("replace").
490 """
491 expandedRefs: list[DatasetRef] = []
492 expandedItemInfos: list[StoredFileInfo] = []
494 for ref, itemInfo in refsAndInfos:
495 expandedRefs.append(ref)
496 expandedItemInfos.append(itemInfo)
498 # Dataset location only cares about registry ID so if we have
499 # disassembled in datastore we have to deduplicate. Since they
500 # will have different datasetTypes we can't use a set
501 registryRefs = {r.id: r for r in expandedRefs}
502 if insert_mode == DatabaseInsertMode.INSERT:
503 self.bridge.insert(registryRefs.values())
504 else:
505 # There are only two columns and all that matters is the
506 # dataset ID.
507 self.bridge.ensure(registryRefs.values())
508 self.addStoredItemInfo(expandedRefs, expandedItemInfos, insert_mode=insert_mode)
510 def _get_stored_records_associated_with_refs(
511 self, refs: Iterable[DatasetIdRef], ignore_datastore_records: bool = False
512 ) -> dict[DatasetId, list[StoredFileInfo]]:
513 """Retrieve all records associated with the provided refs.
515 Parameters
516 ----------
517 refs : iterable of `DatasetIdRef`
518 The refs for which records are to be retrieved.
519 ignore_datastore_records : `bool`
520 If `True` then do not use datastore records stored in refs.
522 Returns
523 -------
524 records : `dict` of [`DatasetId`, `list` of `StoredFileInfo`]
525 The matching records indexed by the ref ID. The number of entries
526 in the dict can be smaller than the number of requested refs.
527 """
528 # Check datastore records in refs first.
529 records_by_ref: defaultdict[DatasetId, list[StoredFileInfo]] = defaultdict(list)
530 refs_with_no_records = []
531 for ref in refs:
532 if ignore_datastore_records or ref._datastore_records is None:
533 refs_with_no_records.append(ref)
534 else:
535 if (ref_records := ref._datastore_records.get(self._table.name)) is not None:
536 # Need to make sure they have correct type.
537 for ref_record in ref_records:
538 if not isinstance(ref_record, StoredFileInfo):
539 raise TypeError(
540 f"Datastore record has unexpected type {ref_record.__class__.__name__}"
541 )
542 records_by_ref[ref.id].append(ref_record)
544 # If there were any refs without datastore records, check opaque table.
545 records = self._table.fetch(dataset_id=[ref.id for ref in refs_with_no_records])
547 # Uniqueness is dataset_id + component so can have multiple records
548 # per ref.
549 for record in records:
550 records_by_ref[record["dataset_id"]].append(StoredFileInfo.from_record(record))
551 return records_by_ref
553 def _refs_associated_with_artifacts(
554 self, paths: Iterable[str | ResourcePath]
555 ) -> dict[str, set[DatasetId]]:
556 """Return paths and associated dataset refs.
558 Parameters
559 ----------
560 paths : `list` of `str` or `lsst.resources.ResourcePath`
561 All the paths to include in search. These are exact matches
562 to the entries in the records table and can include fragments.
564 Returns
565 -------
566 mapping : `dict` of [`str`, `set` [`DatasetId`]]
567 Mapping of each path to a set of associated database IDs.
568 These are artifacts and so any fragments are stripped from the
569 keys.
570 """
571 # Group paths by those that have fragments and those that do not.
572 with_fragment = set()
573 without_fragment = set()
574 for rpath in paths:
575 spath = str(rpath) # Typing says can be ResourcePath so must force to string.
576 if "#" in spath:
577 spath, fragment = spath.rsplit("#", 1)
578 with_fragment.add(spath)
579 else:
580 without_fragment.add(spath)
582 result: dict[str, set[DatasetId]] = defaultdict(set)
583 if without_fragment:
584 records = self._table.fetch(path=without_fragment)
585 for row in records:
586 path = row["path"]
587 result[path].add(row["dataset_id"])
588 if with_fragment:
589 # Do a query per prefix.
590 for path in with_fragment:
591 records = self._table.fetch(path=f"{path}#%")
592 for row in records:
593 # Need to strip fragments before adding to dict.
594 row_path = row["path"]
595 artifact_path = row_path[: row_path.rfind("#")]
596 result[artifact_path].add(row["dataset_id"])
597 return result
599 def _registered_refs_per_artifact(self, pathInStore: ResourcePath) -> set[DatasetId]:
600 """Return all dataset refs associated with the supplied path.
602 Parameters
603 ----------
604 pathInStore : `lsst.resources.ResourcePath`
605 Path of interest in the data store.
607 Returns
608 -------
609 ids : `set` of `int`
610 All `DatasetRef` IDs associated with this path.
611 """
612 records = list(self._table.fetch(path=str(pathInStore)))
613 ids = {r["dataset_id"] for r in records}
614 return ids
616 def removeStoredItemInfo(self, ref: DatasetIdRef) -> None:
617 """Remove information about the file associated with this dataset.
619 Parameters
620 ----------
621 ref : `DatasetRef`
622 The dataset that has been removed.
623 """
624 # Note that this method is actually not used by this implementation,
625 # we depend on bridge to delete opaque records. But there are some
626 # tests that check that this method works, so we keep it for now.
627 self._table.delete(["dataset_id"], {"dataset_id": ref.id})
629 def _get_dataset_locations_info(
630 self, ref: DatasetIdRef, ignore_datastore_records: bool = False
631 ) -> list[DatasetLocationInformation]:
632 r"""Find all the `Location`\ s of the requested dataset in the
633 `Datastore` and the associated stored file information.
635 Parameters
636 ----------
637 ref : `DatasetRef`
638 Reference to the required `Dataset`.
639 ignore_datastore_records : `bool`
640 If `True` then do not use datastore records stored in refs.
642 Returns
643 -------
644 results : `list` [`tuple` [`Location`, `StoredFileInfo` ]]
645 Location of the dataset within the datastore and
646 stored information about each file and its formatter.
647 """
648 # Get the file information (this will fail if no file)
649 records = self.getStoredItemsInfo(ref, ignore_datastore_records)
651 # Use the path to determine the location -- we need to take
652 # into account absolute URIs in the datastore record
653 return [(r.file_location(self.locationFactory), r) for r in records]
655 def _can_remove_dataset_artifact(self, ref: DatasetIdRef, location: Location) -> bool:
656 """Check that there is only one dataset associated with the
657 specified artifact.
659 Parameters
660 ----------
661 ref : `DatasetRef` or `FakeDatasetRef`
662 Dataset to be removed.
663 location : `Location`
664 The location of the artifact to be removed.
666 Returns
667 -------
668 can_remove : `Bool`
669 True if the artifact can be safely removed.
670 """
671 # Can't ever delete absolute URIs.
672 if location.pathInStore.isabs():
673 return False
675 # Get all entries associated with this path
676 allRefs = self._registered_refs_per_artifact(location.pathInStore)
677 if not allRefs:
678 raise RuntimeError(f"Datastore inconsistency error. {location.pathInStore} not in registry")
680 # Remove these refs from all the refs and if there is nothing left
681 # then we can delete
682 remainingRefs = allRefs - {ref.id}
684 if remainingRefs:
685 return False
686 return True
688 def _get_expected_dataset_locations_info(self, ref: DatasetRef) -> list[tuple[Location, StoredFileInfo]]:
689 """Predict the location and related file information of the requested
690 dataset in this datastore.
692 Parameters
693 ----------
694 ref : `DatasetRef`
695 Reference to the required `Dataset`.
697 Returns
698 -------
699 results : `list` [`tuple` [`Location`, `StoredFileInfo` ]]
700 Expected Location of the dataset within the datastore and
701 placeholder information about each file and its formatter.
703 Notes
704 -----
705 Uses the current configuration to determine how we would expect the
706 datastore files to have been written if we couldn't ask registry.
707 This is safe so long as there has been no change to datastore
708 configuration between writing the dataset and wanting to read it.
709 Will not work for files that have been ingested without using the
710 standard file template or default formatter.
711 """
712 # If we have a component ref we always need to ask the questions
713 # of the composite. If the composite is disassembled this routine
714 # should return all components. If the composite was not
715 # disassembled the composite is what is stored regardless of
716 # component request. Note that if the caller has disassembled
717 # a composite there is no way for this guess to know that
718 # without trying both the composite and component ref and seeing
719 # if there is something at the component Location even without
720 # disassembly being enabled.
721 if ref.datasetType.isComponent():
722 ref = ref.makeCompositeRef()
724 # See if the ref is a composite that should be disassembled
725 doDisassembly = self.composites.shouldBeDisassembled(ref)
727 all_info: list[tuple[Location, Formatter | FormatterV2, StorageClass, str | None]] = []
729 if doDisassembly:
730 for component, componentStorage in ref.datasetType.storageClass.components.items():
731 compRef = ref.makeComponentRef(component)
732 location, formatter = self._determine_put_formatter_location(compRef)
733 all_info.append((location, formatter, componentStorage, component))
735 else:
736 # Always use the composite ref if no disassembly
737 location, formatter = self._determine_put_formatter_location(ref)
738 all_info.append((location, formatter, ref.datasetType.storageClass, None))
740 # Convert the list of tuples to have StoredFileInfo as second element
741 return [
742 (
743 location,
744 StoredFileInfo(
745 formatter=formatter,
746 path=location.pathInStore.path,
747 storageClass=storageClass,
748 component=component,
749 checksum=None,
750 file_size=-1,
751 ),
752 )
753 for location, formatter, storageClass, component in all_info
754 ]
756 def _prepare_for_direct_get(
757 self, ref: DatasetRef, parameters: Mapping[str, Any] | None = None
758 ) -> list[DatastoreFileGetInformation]:
759 """Check parameters for ``get`` and obtain formatter and
760 location.
762 Parameters
763 ----------
764 ref : `DatasetRef`
765 Reference to the required Dataset.
766 parameters : `dict`
767 `StorageClass`-specific parameters that specify, for example,
768 a slice of the dataset to be loaded.
770 Returns
771 -------
772 getInfo : `list` [`DatastoreFileGetInformation`]
773 Parameters needed to retrieve each file.
774 """
775 log.debug("Retrieve %s from %s with parameters %s", ref, self.name, parameters)
777 # The storage class we want to use eventually
778 refStorageClass = ref.datasetType.storageClass
780 # For trusted mode need to reset storage class.
781 ref = self._cast_storage_class(ref)
783 # Get file metadata and internal metadata
784 fileLocations = self._get_dataset_locations_info(ref)
785 if not fileLocations:
786 if not self.trustGetRequest:
787 raise FileNotFoundError(f"Could not retrieve dataset {ref}.")
788 # Assume the dataset is where we think it should be
789 fileLocations = self._get_expected_dataset_locations_info(ref)
791 if len(fileLocations) > 1:
792 # If trust is involved it is possible that there will be
793 # components listed here that do not exist in the datastore.
794 # Explicitly check for file artifact existence and filter out any
795 # that are missing.
796 if self.trustGetRequest:
797 fileLocations = [loc for loc in fileLocations if loc[0].uri.exists()]
799 # For now complain only if we have no components at all. One
800 # component is probably a problem but we can punt that to the
801 # assembler.
802 if not fileLocations:
803 raise FileNotFoundError(f"None of the component files for dataset {ref} exist.")
805 return generate_datastore_get_information(
806 fileLocations,
807 readStorageClass=refStorageClass,
808 ref=ref,
809 parameters=parameters,
810 )
812 def _determine_put_formatter_location(
813 self, ref: DatasetRef, provenance: DatasetProvenance | None = None
814 ) -> tuple[Location, Formatter | FormatterV2]:
815 """Calculate the formatter and output location to use for put.
817 Parameters
818 ----------
819 ref : `DatasetRef`
820 Reference to the associated Dataset.
821 provenance : `DatasetProvenance`
822 Any provenance that should be attached to the serialized dataset.
824 Returns
825 -------
826 location : `Location`
827 The location to write the dataset.
828 formatter : `Formatter`
829 The `Formatter` to use to write the dataset.
830 """
831 # Work out output file name
832 try:
833 template = self.templates.getTemplate(ref)
834 except KeyError as e:
835 raise DatasetTypeNotSupportedError(f"Unable to find template for {ref}") from e
837 # Validate the template to protect against filenames from different
838 # dataIds returning the same and causing overwrite confusion.
839 template.validateTemplate(ref)
841 location = self.locationFactory.fromPath(template.format(ref), trusted_path=True)
843 # Get the formatter based on the storage class
844 storageClass = ref.datasetType.storageClass
845 try:
846 formatter = self.formatterFactory.getFormatter(
847 ref,
848 FileDescriptor(location, storageClass=storageClass, component=ref.datasetType.component()),
849 dataId=ref.dataId,
850 ref=ref,
851 provenance=provenance,
852 )
853 except KeyError as e:
854 raise DatasetTypeNotSupportedError(
855 f"Unable to find formatter for {ref} in datastore {self.name}"
856 ) from e
858 # Now that we know the formatter, update the location
859 location = formatter.make_updated_location(location)
861 return location, formatter
863 def _overrideTransferMode(self, *datasets: FileDataset, transfer: str | None = None) -> str | None:
864 # Docstring inherited from base class
865 if transfer != "auto":
866 return transfer
868 # See if the paths are within the datastore or not
869 inside = [self._pathInStore(d.path) is not None for d in datasets]
871 if all(inside):
872 transfer = None
873 elif not any(inside):
874 # Allow ResourcePath to use its own knowledge
875 transfer = "auto"
876 else:
877 # This can happen when importing from a datastore that
878 # has had some datasets ingested using "direct" mode.
879 # Also allow ResourcePath to sort it out but warn about it.
880 # This can happen if you are importing from a datastore
881 # that had some direct transfer datasets.
882 log.warning(
883 "Some datasets are inside the datastore and some are outside. Using 'split' "
884 "transfer mode. This assumes that the files outside the datastore are "
885 "still accessible to the new butler since they will not be copied into "
886 "the target datastore."
887 )
888 transfer = "split"
890 return transfer
892 def _pathInStore(self, path: ResourcePathExpression) -> str | None:
893 """Return path relative to datastore root.
895 Parameters
896 ----------
897 path : `lsst.resources.ResourcePathExpression`
898 Path to dataset. Can be absolute URI. If relative assumed to
899 be relative to the datastore. Returns path in datastore
900 or raises an exception if the path it outside.
902 Returns
903 -------
904 inStore : `str`
905 Path relative to datastore root. Returns `None` if the file is
906 outside the root.
907 """
908 # Relative path will always be relative to datastore
909 pathUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False)
910 return pathUri.relative_to(self.root)
912 def _standardizeIngestPath(
913 self,
914 path: str | ResourcePath,
915 *,
916 transfer: str | None = None,
917 check_existence: bool = False,
918 ) -> str | ResourcePath:
919 """Standardize the path of a to-be-ingested file.
921 Parameters
922 ----------
923 path : `str` or `lsst.resources.ResourcePath`
924 Path of a file to be ingested. This parameter is not expected
925 to be all the types that can be used to construct a
926 `~lsst.resources.ResourcePath`.
927 transfer : `str`, optional
928 How (and whether) the dataset should be added to the datastore.
929 See `ingest` for details of transfer modes.
930 This implementation is provided only so
931 `NotImplementedError` can be raised if the mode is not supported;
932 actual transfers are deferred to `_extractIngestInfo`.
933 check_existence : `bool`, optional
934 If `True` the existence of the file will be checked, otherwise
935 no check will be made.
937 Returns
938 -------
939 path : `str` or `lsst.resources.ResourcePath`
940 New path in what the datastore considers standard form. If an
941 absolute URI was given that will be returned unchanged.
943 Notes
944 -----
945 Subclasses of `FileDatastore` can implement this method instead
946 of `_prepIngest`. It should not modify the data repository or given
947 file in any way.
949 Raises
950 ------
951 NotImplementedError
952 Raised if the datastore does not support the given transfer mode
953 (including the case where ingest is not supported at all).
954 """
955 if transfer not in (None, "direct", "split") + self.root.transferModes:
956 raise NotImplementedError(f"Transfer mode {transfer} not supported.")
958 # A relative URI indicates relative to datastore root
959 srcUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False)
960 if not srcUri.isabs():
961 srcUri = self.root.join(path)
963 if check_existence and not srcUri.exists():
964 raise FileNotFoundError(
965 f"Resource at {srcUri} does not exist; note that paths to ingest "
966 f"are assumed to be relative to {self.root} unless they are absolute."
967 )
969 if transfer is None:
970 relpath = srcUri.relative_to(self.root)
971 if not relpath:
972 raise RuntimeError(
973 f"Transfer is none but source file ({srcUri}) is not within datastore ({self.root})"
974 )
976 # Return the relative path within the datastore for internal
977 # transfer
978 path = relpath
980 return path
982 def _extractIngestInfo(
983 self,
984 path: ResourcePathExpression,
985 ref: DatasetRef,
986 *,
987 formatter: Formatter | FormatterV2 | type[Formatter | FormatterV2],
988 transfer: str | None = None,
989 record_validation_info: bool = True,
990 ) -> StoredFileInfo:
991 """Relocate (if necessary) and extract `StoredFileInfo` from a
992 to-be-ingested file.
994 Parameters
995 ----------
996 path : `lsst.resources.ResourcePathExpression`
997 URI or path of a file to be ingested.
998 ref : `DatasetRef`
999 Reference for the dataset being ingested. Guaranteed to have
1000 ``dataset_id not None`.
1001 formatter : `type` or `Formatter`
1002 `Formatter` subclass to use for this dataset or an instance.
1003 transfer : `str`, optional
1004 How (and whether) the dataset should be added to the datastore.
1005 See `ingest` for details of transfer modes.
1006 record_validation_info : `bool`, optional
1007 If `True`, the default, the datastore can record validation
1008 information associated with the file. If `False` the datastore
1009 will not attempt to track any information such as checksums
1010 or file sizes. This can be useful if such information is tracked
1011 in an external system or if the file is to be compressed in place.
1012 It is up to the datastore whether this parameter is relevant.
1014 Returns
1015 -------
1016 info : `StoredFileInfo`
1017 Internal datastore record for this file. This will be inserted by
1018 the caller; the `_extractIngestInfo` is only responsible for
1019 creating and populating the struct.
1021 Raises
1022 ------
1023 FileNotFoundError
1024 Raised if one of the given files does not exist.
1025 FileExistsError
1026 Raised if transfer is not `None` but the (internal) location the
1027 file would be moved to is already occupied.
1028 """
1029 if self._transaction is None:
1030 raise RuntimeError("Ingest called without transaction enabled")
1032 # Create URI of the source path, do not need to force a relative
1033 # path to absolute.
1034 srcUri = ResourcePath(path, forceAbsolute=False, forceDirectory=False)
1036 # Track whether we have read the size of the source yet
1037 have_sized = False
1039 tgtLocation: Location | None
1040 if transfer is None or transfer == "split":
1041 # A relative path is assumed to be relative to the datastore
1042 # in this context
1043 if not srcUri.isabs():
1044 tgtLocation = self.locationFactory.fromPath(srcUri.ospath, trusted_path=False)
1045 else:
1046 # Work out the path in the datastore from an absolute URI
1047 # This is required to be within the datastore.
1048 pathInStore = srcUri.relative_to(self.root)
1049 if pathInStore is None and transfer is None:
1050 raise RuntimeError(
1051 f"Unexpectedly learned that {srcUri} is not within datastore {self.root}"
1052 )
1053 if pathInStore:
1054 tgtLocation = self.locationFactory.fromPath(pathInStore, trusted_path=True)
1055 elif transfer == "split":
1056 # Outside the datastore but treat that as a direct ingest
1057 # instead.
1058 tgtLocation = None
1059 else:
1060 raise RuntimeError(f"Unexpected transfer mode encountered: {transfer} for URI {srcUri}")
1061 elif transfer == "direct":
1062 # Want to store the full URI to the resource directly in
1063 # datastore. This is useful for referring to permanent archive
1064 # storage for raw data.
1065 # Trust that people know what they are doing.
1066 tgtLocation = None
1067 else:
1068 # Work out the name we want this ingested file to have
1069 # inside the datastore
1070 tgtLocation = self._calculate_ingested_datastore_name(srcUri, ref, formatter)
1071 if not tgtLocation.uri.dirname().exists():
1072 log.debug("Folder %s does not exist yet.", tgtLocation.uri.dirname())
1073 tgtLocation.uri.dirname().mkdir()
1075 # if we are transferring from a local file to a remote location
1076 # it may be more efficient to get the size and checksum of the
1077 # local file rather than the transferred one
1078 if record_validation_info and srcUri.isLocal:
1079 size = srcUri.size()
1080 checksum = self.computeChecksum(srcUri) if self.useChecksum else None
1081 have_sized = True
1083 # Transfer the resource to the destination.
1084 # Allow overwrite of an existing file. This matches the behavior
1085 # of datastore.put() in that it trusts that registry would not
1086 # be asking to overwrite unless registry thought that the
1087 # overwrite was allowed.
1088 tgtLocation.uri.transfer_from(
1089 srcUri, transfer=transfer, transaction=self._transaction, overwrite=True
1090 )
1092 if tgtLocation is None:
1093 # This means we are using direct mode
1094 targetUri = srcUri
1095 targetPath = str(srcUri)
1096 else:
1097 targetUri = tgtLocation.uri
1098 targetPath = tgtLocation.pathInStore.path
1100 # the file should exist in the datastore now
1101 if record_validation_info:
1102 if not have_sized:
1103 size = targetUri.size()
1104 checksum = self.computeChecksum(targetUri) if self.useChecksum else None
1105 else:
1106 # Not recording any file information.
1107 size = -1
1108 checksum = None
1110 return StoredFileInfo(
1111 formatter=formatter,
1112 path=targetPath,
1113 storageClass=ref.datasetType.storageClass,
1114 component=ref.datasetType.component(),
1115 file_size=size,
1116 checksum=checksum,
1117 )
1119 def _prepIngest(self, *datasets: FileDataset, transfer: str | None = None) -> _IngestPrepData:
1120 # Docstring inherited from Datastore._prepIngest.
1121 filtered = []
1123 # Ingest could be given tens of thousands of files. It is not efficient
1124 # to check for the existence of every single file (especially if they
1125 # are remote URIs) but in some transfer modes the files will be checked
1126 # anyhow when they are relocated. For direct or None transfer modes
1127 # it is possible to not know if the file is accessible at all.
1128 # Therefore limit number of files that will be checked (but always
1129 # include the first one).
1130 max_checks = 200
1131 n_datasets = len(datasets)
1132 if n_datasets <= max_checks:
1133 check_every_n = 1
1134 elif transfer in ("direct", None):
1135 check_every_n = int(n_datasets / max_checks + 1) # +1 so that if n < max_checks the answer is 1.
1136 else:
1137 check_every_n = 0
1139 for count, dataset in enumerate(datasets):
1140 acceptable = [ref for ref in dataset.refs if self.constraints.isAcceptable(ref)]
1141 if not acceptable:
1142 continue
1143 else:
1144 dataset.refs = acceptable
1145 if dataset.formatter is None:
1146 dataset.formatter = self.formatterFactory.getFormatterClass(dataset.refs[0])
1147 else:
1148 assert isinstance(dataset.formatter, type | str)
1149 formatter_class = get_class_of(dataset.formatter)
1150 if not issubclass(formatter_class, Formatter | FormatterV2):
1151 raise TypeError(f"Requested formatter {dataset.formatter} is not a Formatter class.")
1152 dataset.formatter = formatter_class
1154 # Decide whether the file should be checked.
1155 check_existence = False
1156 if check_every_n != 0:
1157 # First time through count is 0 so we guarantee to check
1158 # the first file but not necessarily the final one.
1159 check_existence = count % check_every_n == 0
1161 if check_existence:
1162 log.debug(
1163 "Checking file existence: %s (%d/%d) [%s]",
1164 check_existence,
1165 count + 1,
1166 n_datasets,
1167 transfer,
1168 )
1170 dataset.path = self._standardizeIngestPath(
1171 dataset.path, transfer=transfer, check_existence=check_existence
1172 )
1173 filtered.append(dataset)
1174 return _IngestPrepData(filtered)
1176 @transactional
1177 def _finishIngest(
1178 self,
1179 prepData: Datastore.IngestPrepData,
1180 *,
1181 transfer: str | None = None,
1182 record_validation_info: bool = True,
1183 ) -> None:
1184 # Docstring inherited from Datastore._finishIngest.
1185 refsAndInfos = []
1186 progress = Progress("lsst.daf.butler.datastores.FileDatastore.ingest", level=logging.DEBUG)
1187 for dataset in progress.wrap(prepData.datasets, desc="Ingesting dataset files"):
1188 # Do ingest as if the first dataset ref is associated with the file
1189 info = self._extractIngestInfo(
1190 dataset.path,
1191 dataset.refs[0],
1192 formatter=dataset.formatter,
1193 transfer=transfer,
1194 record_validation_info=record_validation_info,
1195 )
1196 refsAndInfos.extend([(ref, info) for ref in dataset.refs])
1198 # In direct mode we can allow repeated ingests of the same thing
1199 # if we are sure that the external dataset is immutable. We use
1200 # UUIDv5 to indicate this. If there is a mix of v4 and v5 they are
1201 # separated.
1202 refs_and_infos_replace = []
1203 refs_and_infos_insert = []
1204 if transfer == "direct":
1205 for entry in refsAndInfos:
1206 if entry[0].id.version == 5:
1207 refs_and_infos_replace.append(entry)
1208 else:
1209 refs_and_infos_insert.append(entry)
1210 else:
1211 refs_and_infos_insert = refsAndInfos
1213 if refs_and_infos_insert:
1214 self._register_datasets(refs_and_infos_insert, insert_mode=DatabaseInsertMode.INSERT)
1215 if refs_and_infos_replace:
1216 self._register_datasets(refs_and_infos_replace, insert_mode=DatabaseInsertMode.REPLACE)
1218 def _calculate_ingested_datastore_name(
1219 self,
1220 srcUri: ResourcePath,
1221 ref: DatasetRef,
1222 formatter: Formatter | FormatterV2 | type[Formatter | FormatterV2] | None = None,
1223 ) -> Location:
1224 """Given a source URI and a DatasetRef, determine the name the
1225 dataset will have inside datastore.
1227 Parameters
1228 ----------
1229 srcUri : `lsst.resources.ResourcePath`
1230 URI to the source dataset file.
1231 ref : `DatasetRef`
1232 Ref associated with the newly-ingested dataset artifact. This
1233 is used to determine the name within the datastore.
1234 formatter : `Formatter` or Formatter class.
1235 Formatter to use for validation. Can be a class or an instance.
1236 No validation of the file extension is performed if the
1237 ``formatter`` is `None`. This can be used if the caller knows
1238 that the source URI and target URI will use the same formatter.
1240 Returns
1241 -------
1242 location : `Location`
1243 Target location for the newly-ingested dataset.
1244 """
1245 # Ingesting a file from outside the datastore.
1246 # This involves a new name.
1247 template = self.templates.getTemplate(ref)
1248 location = self.locationFactory.fromPath(template.format(ref), trusted_path=True)
1250 # Get the extension
1251 ext = srcUri.getExtension()
1253 # Update the destination to include that extension
1254 location.updateExtension(ext)
1256 # Ask the formatter to validate this extension
1257 if formatter is not None:
1258 formatter.validate_extension(location)
1260 return location
1262 def _write_in_memory_to_artifact(
1263 self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None
1264 ) -> StoredFileInfo:
1265 """Write out in memory dataset to datastore.
1267 Parameters
1268 ----------
1269 inMemoryDataset : `object`
1270 Dataset to write to datastore.
1271 ref : `DatasetRef`
1272 Registry information associated with this dataset.
1273 provenance : `DatasetProvenance` or `None`, optional
1274 Any provenance that should be attached to the serialized dataset.
1275 Not supported by all formatters.
1277 Returns
1278 -------
1279 info : `StoredFileInfo`
1280 Information describing the artifact written to the datastore.
1281 """
1282 # May need to coerce the in memory dataset to the correct
1283 # python type, but first we need to make sure the storage class
1284 # reflects the one defined in the data repository.
1285 ref = self._cast_storage_class(ref)
1287 # Confirm that we can accept this dataset
1288 if not self.constraints.isAcceptable(ref):
1289 # Raise rather than use boolean return value.
1290 raise DatasetTypeNotSupportedError(
1291 f"Dataset {ref} has been rejected by this datastore via configuration."
1292 )
1294 location, formatter = self._determine_put_formatter_location(ref)
1296 # The external storage class can differ from the registry storage
1297 # class AND the given in-memory dataset might not match any of the
1298 # storage class definitions.
1299 if formatter.can_accept(inMemoryDataset):
1300 # Do not need to coerce. Must assume that the formatter can handle
1301 # it without further checking of types.
1302 pass
1303 else:
1304 # Coerce to a type that it can accept.
1305 inMemoryDataset = ref.datasetType.storageClass.coerce_type(inMemoryDataset)
1306 required_pytype = ref.datasetType.storageClass.pytype
1308 if not isinstance(inMemoryDataset, required_pytype):
1309 raise TypeError(
1310 f"Inconsistency between supplied object ({type(inMemoryDataset)}) "
1311 f"and storage class type ({required_pytype})"
1312 )
1314 uri = location.uri
1316 if not uri.dirname().exists():
1317 log.debug("Folder %s does not exist yet so creating it.", uri.dirname())
1318 uri.dirname().mkdir()
1320 if self._transaction is None:
1321 raise RuntimeError("Attempting to write artifact without transaction enabled")
1323 def _removeFileExists(uri: ResourcePath) -> None:
1324 """Remove a file and do not complain if it is not there.
1326 This is important since a formatter might fail before the file
1327 is written and we should not confuse people by writing spurious
1328 error messages to the log.
1329 """
1330 with contextlib.suppress(FileNotFoundError):
1331 uri.remove()
1333 # Register a callback to try to delete the uploaded data if
1334 # something fails below
1335 self._transaction.registerUndo("artifactWrite", _removeFileExists, uri)
1337 # Need to record the specified formatter but if this is a V1 formatter
1338 # we need to convert it to a V2 compatible shim to do the write.
1339 if not isinstance(formatter, Formatter):
1340 formatter_compat = formatter
1341 else:
1342 formatter_compat = FormatterV1inV2(
1343 formatter.file_descriptor,
1344 ref=ref,
1345 formatter=formatter,
1346 write_parameters=formatter.write_parameters,
1347 write_recipes=formatter.write_recipes,
1348 )
1350 assert isinstance(formatter_compat, FormatterV2)
1352 with time_this(log, msg="Writing dataset %s with formatter %s", args=(ref, formatter.name())):
1353 try:
1354 formatter_compat.write(
1355 inMemoryDataset, cache_manager=self.cacheManager, provenance=provenance
1356 )
1357 except Exception as e:
1358 raise RuntimeError(
1359 f"Failed to serialize dataset {ref} of type {get_full_type_name(inMemoryDataset)} "
1360 f"using formatter {formatter.name()}."
1361 ) from e
1363 # URI is needed to resolve what ingest case are we dealing with
1364 return self._extractIngestInfo(uri, ref, formatter=formatter)
1366 def knows(self, ref: DatasetRef) -> bool:
1367 """Check if the dataset is known to the datastore.
1369 Does not check for existence of any artifact.
1371 Parameters
1372 ----------
1373 ref : `DatasetRef`
1374 Reference to the required dataset.
1376 Returns
1377 -------
1378 exists : `bool`
1379 `True` if the dataset is known to the datastore.
1380 """
1381 fileLocations = self._get_dataset_locations_info(ref)
1382 if fileLocations:
1383 return True
1384 return False
1386 def knows_these(self, refs: Iterable[DatasetRef]) -> dict[DatasetRef, bool]:
1387 # Docstring inherited from the base class.
1388 refs = list(refs)
1390 # The records themselves. Could be missing some entries.
1391 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
1393 return {ref: ref.id in records for ref in refs}
1395 def _process_mexists_records(
1396 self,
1397 id_to_ref: dict[DatasetId, DatasetRef],
1398 records: dict[DatasetId, list[StoredFileInfo]],
1399 all_required: bool,
1400 artifact_existence: dict[ResourcePath, bool] | None = None,
1401 ) -> dict[DatasetRef, bool]:
1402 """Check given records for existence.
1404 Helper function for `mexists()`.
1406 Parameters
1407 ----------
1408 id_to_ref : `dict` of [`DatasetId`, `DatasetRef`]
1409 Mapping of the dataset ID to the dataset ref itself.
1410 records : `dict` of [`DatasetId`, `list` of `StoredFileInfo`]
1411 Records as generally returned by
1412 ``_get_stored_records_associated_with_refs``.
1413 all_required : `bool`
1414 Flag to indicate whether existence requires all artifacts
1415 associated with a dataset ID to exist or not for existence.
1416 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1417 Optional mapping of datastore artifact to existence. Updated by
1418 this method with details of all artifacts tested. Can be `None`
1419 if the caller is not interested.
1421 Returns
1422 -------
1423 existence : `dict` of [`DatasetRef`, `bool`]
1424 Mapping from dataset to boolean indicating existence.
1425 """
1426 # The URIs to be checked and a mapping of those URIs to
1427 # the dataset ID.
1428 uris_to_check: list[ResourcePath] = []
1429 location_map: dict[ResourcePath, DatasetId] = {}
1431 location_factory = self.locationFactory
1433 uri_existence: dict[ResourcePath, bool] = {}
1434 for ref_id, infos in records.items():
1435 # Key is the dataset Id, value is list of StoredItemInfo
1436 uris = [info.file_location(location_factory).uri for info in infos]
1437 location_map.update({uri: ref_id for uri in uris})
1439 # Check the local cache directly for a dataset corresponding
1440 # to the remote URI.
1441 if self.cacheManager.file_count > 0:
1442 ref = id_to_ref[ref_id]
1443 for uri, storedFileInfo in zip(uris, infos, strict=True):
1444 check_ref = ref
1445 if not ref.datasetType.isComponent() and (component := storedFileInfo.component):
1446 check_ref = ref.makeComponentRef(component)
1447 if self.cacheManager.known_to_cache(check_ref, uri.getExtension()):
1448 # Proxy for URI existence.
1449 uri_existence[uri] = True
1450 else:
1451 uris_to_check.append(uri)
1452 else:
1453 # Check all of them.
1454 uris_to_check.extend(uris)
1456 if artifact_existence is not None:
1457 # If a URI has already been checked remove it from the list
1458 # and immediately add the status to the output dict.
1459 filtered_uris_to_check = []
1460 for uri in uris_to_check:
1461 if uri in artifact_existence:
1462 uri_existence[uri] = artifact_existence[uri]
1463 else:
1464 filtered_uris_to_check.append(uri)
1465 uris_to_check = filtered_uris_to_check
1467 # Results.
1468 dataset_existence: dict[DatasetRef, bool] = {}
1470 uri_existence.update(ResourcePath.mexists(uris_to_check))
1471 for uri, exists in uri_existence.items():
1472 dataset_id = location_map[uri]
1473 ref = id_to_ref[dataset_id]
1475 # Disassembled composite needs to check all locations.
1476 # all_required indicates whether all need to exist or not.
1477 if ref in dataset_existence:
1478 if all_required:
1479 exists = dataset_existence[ref] and exists
1480 else:
1481 exists = dataset_existence[ref] or exists
1482 dataset_existence[ref] = exists
1484 if artifact_existence is not None:
1485 artifact_existence.update(uri_existence)
1487 return dataset_existence
1489 def mexists(
1490 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1491 ) -> dict[DatasetRef, bool]:
1492 """Check the existence of multiple datasets at once.
1494 Parameters
1495 ----------
1496 refs : iterable of `DatasetRef`
1497 The datasets to be checked.
1498 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1499 Optional mapping of datastore artifact to existence. Updated by
1500 this method with details of all artifacts tested. Can be `None`
1501 if the caller is not interested.
1503 Returns
1504 -------
1505 existence : `dict` of [`DatasetRef`, `bool`]
1506 Mapping from dataset to boolean indicating existence.
1508 Notes
1509 -----
1510 To minimize potentially costly remote existence checks, the local
1511 cache is checked as a proxy for existence. If a file for this
1512 `DatasetRef` does exist no check is done for the actual URI. This
1513 could result in possibly unexpected behavior if the dataset itself
1514 has been removed from the datastore by another process whilst it is
1515 still in the cache.
1516 """
1517 chunk_size = 50_000
1518 dataset_existence: dict[DatasetRef, bool] = {}
1519 log.debug("Checking for the existence of multiple artifacts in datastore in chunks of %d", chunk_size)
1520 n_found_total = 0
1521 n_checked = 0
1522 n_chunks = 0
1523 for chunk in chunk_iterable(refs, chunk_size=chunk_size):
1524 chunk_result = self._mexists(chunk, artifact_existence)
1526 # The log message level and content depend on how many
1527 # datasets we are processing.
1528 n_results = len(chunk_result)
1530 # Use verbose logging to ensure that messages can be seen
1531 # easily if many refs are being checked.
1532 log_threshold = VERBOSE
1533 n_checked += n_results
1535 # This sum can take some time so only do it if we know the
1536 # result is going to be used.
1537 n_found = 0
1538 if log.isEnabledFor(log_threshold):
1539 # Can treat the booleans as 0, 1 integers and sum them.
1540 n_found = sum(chunk_result.values())
1541 n_found_total += n_found
1543 # We are deliberately not trying to count the number of refs
1544 # provided in case it's in the millions. This means there is a
1545 # situation where the number of refs exactly matches the chunk
1546 # size and we will switch to the multi-chunk path even though
1547 # we only have a single chunk.
1548 if n_results < chunk_size and n_chunks == 0:
1549 # Single chunk will be processed so we can provide more detail.
1550 if n_results == 1:
1551 ref = list(chunk_result)[0]
1552 # Use debug logging to be consistent with `exists()`.
1553 log.debug(
1554 "Calling mexists() with single ref that does%s exist (%s).",
1555 "" if chunk_result[ref] else " not",
1556 ref,
1557 )
1558 else:
1559 # Single chunk but multiple files. Summarize.
1560 log.log(
1561 log_threshold,
1562 "Number of datasets found in datastore %s: %d out of %d datasets checked.",
1563 self.name,
1564 n_found,
1565 n_checked,
1566 )
1568 else:
1569 # Use incremental verbose logging when we have multiple chunks.
1570 log.log(
1571 log_threshold,
1572 "Number of datasets found in datastore for chunk %d: %d out of %d checked "
1573 "(running total from all chunks so far: %d found out of %d checked)",
1574 n_chunks,
1575 n_found,
1576 n_results,
1577 n_found_total,
1578 n_checked,
1579 )
1580 dataset_existence.update(chunk_result)
1581 n_chunks += 1
1583 return dataset_existence
1585 def _mexists(
1586 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1587 ) -> dict[DatasetRef, bool]:
1588 """Check the existence of multiple datasets at once.
1590 Parameters
1591 ----------
1592 refs : iterable of `DatasetRef`
1593 The datasets to be checked.
1594 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1595 Optional mapping of datastore artifact to existence. Updated by
1596 this method with details of all artifacts tested. Can be `None`
1597 if the caller is not interested.
1599 Returns
1600 -------
1601 existence : `dict` of [`DatasetRef`, `bool`]
1602 Mapping from dataset to boolean indicating existence.
1603 """
1604 # Make a mapping from refs with the internal storage class to the given
1605 # refs that may have a different one. We'll use the internal refs
1606 # throughout this method and convert back at the very end.
1607 internal_ref_to_input_ref = {self._cast_storage_class(ref): ref for ref in refs}
1609 # Need a mapping of dataset_id to (internal) dataset ref since some
1610 # internal APIs work with dataset_id.
1611 id_to_ref = {ref.id: ref for ref in internal_ref_to_input_ref}
1613 # Set of all IDs we are checking for.
1614 requested_ids = set(id_to_ref.keys())
1616 # The records themselves. Could be missing some entries.
1617 records = self._get_stored_records_associated_with_refs(
1618 id_to_ref.values(), ignore_datastore_records=True
1619 )
1621 dataset_existence = self._process_mexists_records(
1622 id_to_ref, records, True, artifact_existence=artifact_existence
1623 )
1625 # Set of IDs that have been handled.
1626 handled_ids = {ref.id for ref in dataset_existence}
1628 missing_ids = requested_ids - handled_ids
1629 if missing_ids:
1630 dataset_existence.update(
1631 self._mexists_check_expected(
1632 [id_to_ref[missing] for missing in missing_ids], artifact_existence
1633 )
1634 )
1636 return {
1637 internal_ref_to_input_ref[internal_ref]: existence
1638 for internal_ref, existence in dataset_existence.items()
1639 }
1641 def _mexists_check_expected(
1642 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1643 ) -> dict[DatasetRef, bool]:
1644 """Check existence of refs that are not known to datastore.
1646 Parameters
1647 ----------
1648 refs : iterable of `DatasetRef`
1649 The datasets to be checked. These are assumed not to be known
1650 to datastore.
1651 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1652 Optional mapping of datastore artifact to existence. Updated by
1653 this method with details of all artifacts tested. Can be `None`
1654 if the caller is not interested.
1656 Returns
1657 -------
1658 existence : `dict` of [`DatasetRef`, `bool`]
1659 Mapping from dataset to boolean indicating existence.
1660 """
1661 dataset_existence: dict[DatasetRef, bool] = {}
1662 if not self.trustGetRequest:
1663 # Must assume these do not exist
1664 for ref in refs:
1665 dataset_existence[ref] = False
1666 else:
1667 log.debug(
1668 "%d datasets were not known to datastore during initial existence check.",
1669 len(refs),
1670 )
1672 # Construct data structure identical to that returned
1673 # by _get_stored_records_associated_with_refs() but using
1674 # guessed names.
1675 records = {}
1676 id_to_ref = {}
1677 for missing_ref in refs:
1678 expected = self._get_expected_dataset_locations_info(missing_ref)
1679 dataset_id = missing_ref.id
1680 records[dataset_id] = [info for _, info in expected]
1681 id_to_ref[dataset_id] = missing_ref
1683 dataset_existence.update(
1684 self._process_mexists_records(
1685 id_to_ref,
1686 records,
1687 False,
1688 artifact_existence=artifact_existence,
1689 )
1690 )
1692 return dataset_existence
1694 def exists(self, ref: DatasetRef) -> bool:
1695 """Check if the dataset exists in the datastore.
1697 Parameters
1698 ----------
1699 ref : `DatasetRef`
1700 Reference to the required dataset.
1702 Returns
1703 -------
1704 exists : `bool`
1705 `True` if the entity exists in the `Datastore`.
1707 Notes
1708 -----
1709 The local cache is checked as a proxy for existence in the remote
1710 object store. It is possible that another process on a different
1711 compute node could remove the file from the object store even
1712 though it is present in the local cache.
1713 """
1714 ref = self._cast_storage_class(ref)
1715 # We cannot trust datastore records from ref, as many unit tests delete
1716 # datasets and check their existence.
1717 fileLocations = self._get_dataset_locations_info(ref, ignore_datastore_records=True)
1719 # if we are being asked to trust that registry might not be correct
1720 # we ask for the expected locations and check them explicitly
1721 if not fileLocations:
1722 if not self.trustGetRequest:
1723 return False
1725 # First check the cache. If it is not found we must check
1726 # the datastore itself. Assume that any component in the cache
1727 # means that the dataset does exist somewhere.
1728 if self.cacheManager.known_to_cache(ref):
1729 return True
1731 # When we are guessing a dataset location we can not check
1732 # for the existence of every component since we can not
1733 # know if every component was written. Instead we check
1734 # for the existence of any of the expected locations.
1735 for location, _ in self._get_expected_dataset_locations_info(ref):
1736 if self._artifact_exists(location):
1737 return True
1738 return False
1740 # All listed artifacts must exist.
1741 for location, storedFileInfo in fileLocations:
1742 # Checking in cache needs the component ref.
1743 check_ref = ref
1744 if not ref.datasetType.isComponent() and (component := storedFileInfo.component):
1745 check_ref = ref.makeComponentRef(component)
1746 if self.cacheManager.known_to_cache(check_ref, location.getExtension()):
1747 continue
1749 if not self._artifact_exists(location):
1750 return False
1752 return True
1754 def getURIs(self, ref: DatasetRef, predict: bool = False) -> DatasetRefURIs:
1755 """Return URIs associated with dataset.
1757 Parameters
1758 ----------
1759 ref : `DatasetRef`
1760 Reference to the required dataset.
1761 predict : `bool`, optional
1762 If the datastore does not know about the dataset, controls whether
1763 it should return a predicted URI or not.
1765 Returns
1766 -------
1767 uris : `DatasetRefURIs`
1768 The URI to the primary artifact associated with this dataset (if
1769 the dataset was disassembled within the datastore this may be
1770 `None`), and the URIs to any components associated with the dataset
1771 artifact. (can be empty if there are no components).
1772 """
1773 many = self.getManyURIs([ref], predict=predict, allow_missing=False)
1774 return many[ref]
1776 def getURI(self, ref: DatasetRef, predict: bool = False) -> ResourcePath:
1777 """URI to the Dataset.
1779 Parameters
1780 ----------
1781 ref : `DatasetRef`
1782 Reference to the required Dataset.
1783 predict : `bool`
1784 If `True`, allow URIs to be returned of datasets that have not
1785 been written.
1787 Returns
1788 -------
1789 uri : `str`
1790 URI pointing to the dataset within the datastore. If the
1791 dataset does not exist in the datastore, and if ``predict`` is
1792 `True`, the URI will be a prediction and will include a URI
1793 fragment "#predicted".
1794 If the datastore does not have entities that relate well
1795 to the concept of a URI the returned URI will be
1796 descriptive. The returned URI is not guaranteed to be obtainable.
1798 Raises
1799 ------
1800 FileNotFoundError
1801 Raised if a URI has been requested for a dataset that does not
1802 exist and guessing is not allowed.
1803 RuntimeError
1804 Raised if a request is made for a single URI but multiple URIs
1805 are associated with this dataset.
1807 Notes
1808 -----
1809 When a predicted URI is requested an attempt will be made to form
1810 a reasonable URI based on file templates and the expected formatter.
1811 """
1812 primary, components = self.getURIs(ref, predict)
1813 if primary is None or components:
1814 raise RuntimeError(
1815 f"Dataset ({ref}) includes distinct URIs for components. Use Datastore.getURIs() instead."
1816 )
1817 return primary
1819 def _predict_URIs(
1820 self,
1821 ref: DatasetRef,
1822 ) -> DatasetRefURIs:
1823 """Predict the URIs of a dataset ref.
1825 Parameters
1826 ----------
1827 ref : `DatasetRef`
1828 Reference to the required Dataset.
1830 Returns
1831 -------
1832 URI : DatasetRefUris
1833 Primary and component URIs. URIs will contain a URI fragment
1834 "#predicted".
1835 """
1836 uris = DatasetRefURIs()
1838 if self.composites.shouldBeDisassembled(ref):
1839 for component, _ in ref.datasetType.storageClass.components.items():
1840 comp_ref = ref.makeComponentRef(component)
1841 comp_location, _ = self._determine_put_formatter_location(comp_ref)
1843 # Add the "#predicted" URI fragment to indicate this is a
1844 # guess
1845 uris.componentURIs[component] = ResourcePath(
1846 comp_location.uri.geturl() + "#predicted", forceDirectory=comp_location.uri.dirLike
1847 )
1849 else:
1850 location, _ = self._determine_put_formatter_location(ref)
1852 # Add the "#predicted" URI fragment to indicate this is a guess
1853 uris.primaryURI = ResourcePath(
1854 location.uri.geturl() + "#predicted", forceDirectory=location.uri.dirLike
1855 )
1857 return uris
1859 def getManyURIs(
1860 self,
1861 refs: Iterable[DatasetRef],
1862 predict: bool = False,
1863 allow_missing: bool = False,
1864 ) -> dict[DatasetRef, DatasetRefURIs]:
1865 # Docstring inherited
1867 uris: dict[DatasetRef, DatasetRefURIs] = {}
1869 records = self._get_stored_records_associated_with_refs(refs)
1870 records_keys = records.keys()
1872 existing_refs = tuple(ref for ref in refs if ref.id in records_keys)
1873 missing_refs = tuple(ref for ref in refs if ref.id not in records_keys)
1875 # Have to handle trustGetRequest mode by checking for the existence
1876 # of the missing refs on disk.
1877 if missing_refs and not predict:
1878 dataset_existence = self._mexists_check_expected(missing_refs, None)
1879 really_missing = set()
1880 not_missing = set()
1881 for ref, exists in dataset_existence.items():
1882 if exists:
1883 not_missing.add(ref)
1884 else:
1885 really_missing.add(ref)
1887 if not_missing:
1888 # Need to recalculate the missing/existing split.
1889 existing_refs = existing_refs + tuple(not_missing)
1890 missing_refs = tuple(really_missing)
1892 for ref in missing_refs:
1893 # if this has never been written then we have to guess
1894 if not predict:
1895 if not allow_missing:
1896 raise FileNotFoundError(f"Dataset {ref} not in this datastore.")
1897 else:
1898 uris[ref] = self._predict_URIs(ref)
1900 for ref in existing_refs:
1901 file_infos = records[ref.id]
1902 file_locations = [(i.file_location(self.locationFactory), i) for i in file_infos]
1903 uris[ref] = self._locations_to_URI(ref, file_locations)
1905 return uris
1907 def _locations_to_URI(
1908 self,
1909 ref: DatasetRef,
1910 file_locations: Sequence[tuple[Location, StoredFileInfo]],
1911 ) -> DatasetRefURIs:
1912 """Convert one or more file locations associated with a DatasetRef
1913 to a DatasetRefURIs.
1915 Parameters
1916 ----------
1917 ref : `DatasetRef`
1918 Reference to the dataset.
1919 file_locations : Sequence[Tuple[Location, StoredFileInfo]]
1920 Each item in the sequence is the location of the dataset within the
1921 datastore and stored information about the file and its formatter.
1922 If there is only one item in the sequence then it is treated as the
1923 primary URI. If there is more than one item then they are treated
1924 as component URIs. If there are no items then an error is raised
1925 unless ``self.trustGetRequest`` is `True`.
1927 Returns
1928 -------
1929 uris: DatasetRefURIs
1930 Represents the primary URI or component URIs described by the
1931 inputs.
1933 Raises
1934 ------
1935 RuntimeError
1936 If no file locations are passed in and ``self.trustGetRequest`` is
1937 `False`.
1938 FileNotFoundError
1939 If the a passed-in URI does not exist, and ``self.trustGetRequest``
1940 is `False`.
1941 RuntimeError
1942 If a passed in `StoredFileInfo`'s ``component`` is `None` (this is
1943 unexpected).
1944 """
1945 guessing = False
1946 uris = DatasetRefURIs()
1948 if not file_locations:
1949 if not self.trustGetRequest:
1950 raise RuntimeError(f"Unexpectedly got no artifacts for dataset {ref}")
1951 file_locations = self._get_expected_dataset_locations_info(ref)
1952 guessing = True
1954 if len(file_locations) == 1:
1955 # No disassembly so this is the primary URI
1956 uris.primaryURI = file_locations[0][0].uri
1957 if guessing and not uris.primaryURI.exists():
1958 raise FileNotFoundError(f"Expected URI ({uris.primaryURI}) does not exist")
1959 else:
1960 for location, file_info in file_locations:
1961 if file_info.component is None:
1962 raise RuntimeError(f"Unexpectedly got no component name for a component at {location}")
1963 if guessing and not location.uri.exists():
1964 # If we are trusting then it is entirely possible for
1965 # some components to be missing. In that case we skip
1966 # to the next component.
1967 if self.trustGetRequest:
1968 continue
1969 raise FileNotFoundError(f"Expected URI ({location.uri}) does not exist")
1970 uris.componentURIs[file_info.component] = location.uri
1972 return uris
1974 def _find_missing_records(
1975 self,
1976 refs: Iterable[DatasetRef],
1977 missing_ids: set[DatasetId],
1978 artifact_existence: dict[ResourcePath, bool] | None = None,
1979 warn_for_missing: bool = True,
1980 ) -> dict[DatasetId, list[StoredFileInfo]]:
1981 if not missing_ids:
1982 return {}
1984 if artifact_existence is None:
1985 artifact_existence = {}
1987 found_records: dict[DatasetId, list[StoredFileInfo]] = defaultdict(list)
1988 id_to_ref = {ref.id: ref for ref in refs if ref.id in missing_ids}
1990 # This should be chunked in case we end up having to check
1991 # the file store since we need some log output to show
1992 # progress.
1993 chunk_size = 50_000
1994 for missing_ids_chunk in chunk_iterable(missing_ids, chunk_size=chunk_size):
1995 records = {}
1996 for missing in missing_ids_chunk:
1997 # Ask the source datastore where the missing artifacts
1998 # should be. An execution butler might not know about the
1999 # artifacts even if they are there.
2000 expected = self._get_expected_dataset_locations_info(id_to_ref[missing])
2001 records[missing] = [info for _, info in expected]
2003 # Call the mexist helper method in case we have not already
2004 # checked these artifacts such that artifact_existence is
2005 # empty. This allows us to benefit from parallelism.
2006 # datastore.mexists() itself does not give us access to the
2007 # derived datastore record.
2008 log.verbose("Checking existence of %d datasets unknown to datastore", len(records))
2009 ref_exists = self._process_mexists_records(
2010 id_to_ref, records, False, artifact_existence=artifact_existence
2011 )
2013 # Now go through the records and propagate the ones that exist.
2014 location_factory = self.locationFactory
2015 for missing, record_list in records.items():
2016 # Skip completely if the ref does not exist.
2017 ref = id_to_ref[missing]
2018 if not ref_exists[ref]:
2019 if warn_for_missing:
2020 log.warning("Asked to transfer dataset %s but no file artifacts exist for it.", ref)
2021 continue
2022 # Check for file artifact to decide which parts of a
2023 # disassembled composite do exist. If there is only a
2024 # single record we don't even need to look because it can't
2025 # be a composite and must exist.
2026 if len(record_list) == 1:
2027 dataset_records = record_list
2028 else:
2029 dataset_records = [
2030 record
2031 for record in record_list
2032 if artifact_existence[record.file_location(location_factory).uri]
2033 ]
2034 assert len(dataset_records) > 0, "Disassembled composite should have had some files."
2036 # Rely on source_records being a defaultdict.
2037 found_records[missing].extend(dataset_records)
2038 log.verbose("Completed scan for missing data files")
2039 return found_records
2041 def retrieveArtifacts(
2042 self,
2043 refs: Iterable[DatasetRef],
2044 destination: ResourcePath,
2045 transfer: str = "auto",
2046 preserve_path: bool = True,
2047 overwrite: bool = False,
2048 write_index: bool = True,
2049 add_prefix: bool = False,
2050 ) -> dict[ResourcePath, ArtifactIndexInfo]:
2051 """Retrieve the file artifacts associated with the supplied refs.
2053 Parameters
2054 ----------
2055 refs : iterable of `DatasetRef`
2056 The datasets for which file artifacts are to be retrieved.
2057 A single ref can result in multiple files. The refs must
2058 be resolved.
2059 destination : `lsst.resources.ResourcePath`
2060 Location to write the file artifacts.
2061 transfer : `str`, optional
2062 Method to use to transfer the artifacts. Must be one of the options
2063 supported by `lsst.resources.ResourcePath.transfer_from()`.
2064 "move" is not allowed.
2065 preserve_path : `bool`, optional
2066 If `True` the full path of the file artifact within the datastore
2067 is preserved. If `False` the final file component of the path
2068 is used.
2069 overwrite : `bool`, optional
2070 If `True` allow transfers to overwrite existing files at the
2071 destination.
2072 write_index : `bool`, optional
2073 If `True` write a file at the top level containing a serialization
2074 of a `ZipIndex` for the downloaded datasets.
2075 add_prefix : `bool`, optional
2076 If `True` and if ``preserve_path`` is `False`, apply a prefix to
2077 the filenames corresponding to some part of the dataset ref ID.
2078 This can be used to guarantee uniqueness.
2080 Returns
2081 -------
2082 artifact_map : `dict` [ `lsst.resources.ResourcePath`, \
2083 `ArtifactIndexInfo` ]
2084 Mapping of retrieved file to associated index information.
2085 """
2086 if not destination.isdir():
2087 raise ValueError(f"Destination location must refer to a directory. Given {destination}")
2089 if transfer == "move":
2090 raise ValueError("Can not move artifacts out of datastore. Use copy instead.")
2092 # Source -> Destination
2093 # This also helps filter out duplicate DatasetRef in the request
2094 # that will map to the same underlying file transfer.
2095 to_transfer: dict[ResourcePath, ResourcePath] = {}
2096 zips_to_transfer: set[ResourcePath] = set()
2098 # Retrieve all the records in bulk indexed by ref.id.
2099 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2101 # Check for missing records.
2102 known_ids = set(records)
2103 log.debug("Number of datastore records found in database: %d", len(known_ids))
2104 requested_ids = {ref.id for ref in refs}
2105 missing_ids = requested_ids - known_ids
2107 if missing_ids and not self.trustGetRequest:
2108 raise ValueError(f"Number of datasets missing from this datastore: {len(missing_ids)}")
2110 missing_records = self._find_missing_records(refs, missing_ids)
2111 records.update(missing_records)
2113 # One artifact can be used by multiple DatasetRef.
2114 # e.g. DECam.
2115 artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {}
2116 # Sort to ensure that in many refs to one file situation the same
2117 # ref is used for any prefix that might be added.
2118 for ref in sorted(refs):
2119 prefix = str(ref.id)[:8] + "-" if add_prefix else ""
2120 for info in records[ref.id]:
2121 location = info.file_location(self.locationFactory)
2122 source_uri = location.uri
2123 # For DECam/zip we only want to copy once.
2124 # For zip files we need to unpack so that they can be
2125 # zipped up again if needed.
2126 is_zip = source_uri.getExtension() == ".zip" and "zip-path" in source_uri.fragment
2127 # We need to remove fragments for consistency.
2128 cleaned_source_uri = source_uri.replace(fragment="", query="", params="")
2129 if is_zip:
2130 # Assume the DatasetRef definitions are within the Zip
2131 # file itself and so can be dropped from loop.
2132 zips_to_transfer.add(cleaned_source_uri)
2133 elif cleaned_source_uri not in to_transfer:
2134 target_uri = determine_destination_for_retrieved_artifact(
2135 destination, location.pathInStore, preserve_path, prefix
2136 )
2137 to_transfer[cleaned_source_uri] = target_uri
2138 artifact_map[target_uri] = ArtifactIndexInfo.from_single(info.to_simple(), ref.id)
2139 else:
2140 target_uri = to_transfer[cleaned_source_uri]
2141 artifact_map[target_uri].append(ref.id)
2143 # Parallelize the transfer. Re-raise as a single exception if
2144 # a FileExistsError is encountered anywhere.
2145 log.debug("Number of artifacts to transfer to %s: %d", str(destination), len(to_transfer))
2146 try:
2147 ResourcePath.mtransfer(transfer, tuple(to_transfer.items()), overwrite=overwrite)
2148 except* FileExistsError as egroup:
2149 raise FileExistsError(
2150 "Some files already exist in destination directory and overwrite is False"
2151 ) from egroup
2153 # Transfer the Zip files and unpack them.
2154 zipped_artifacts = unpack_zips(zips_to_transfer, requested_ids, destination, preserve_path, overwrite)
2155 artifact_map.update(zipped_artifacts)
2157 if write_index:
2158 index = ZipIndex.from_artifact_map(refs, artifact_map, destination)
2159 index.write_index(destination)
2161 return artifact_map
2163 def ingest_zip(self, zip_path: ResourcePath, transfer: str | None, *, dry_run: bool = False) -> None:
2164 """Ingest an indexed Zip file and contents.
2166 The Zip file must have an index file as created by `retrieveArtifacts`.
2168 Parameters
2169 ----------
2170 zip_path : `lsst.resources.ResourcePath`
2171 Path to the Zip file.
2172 transfer : `str`
2173 Method to use for transferring the Zip file into the datastore.
2174 dry_run : `bool`, optional
2175 If `True` the ingest will be processed without any modifications
2176 made to the target datastore and as if the target datastore did not
2177 have any of the datasets.
2179 Notes
2180 -----
2181 Datastore constraints are bypassed with Zip ingest. A zip file can
2182 contain multiple dataset types. Should the entire Zip be rejected
2183 if one dataset type is in the constraints list?
2185 If any dataset is already present in the datastore the entire ingest
2186 will fail.
2187 """
2188 index = ZipIndex.from_zip_file(zip_path)
2190 # Refs indexed by UUID.
2191 refs = index.refs.to_refs(universe=self.universe)
2192 id_to_ref = {ref.id: ref for ref in refs}
2194 # Any failing constraints trigger entire failure.
2195 if any(not self.constraints.isAcceptable(ref) for ref in refs):
2196 raise DatasetTypeNotSupportedError(
2197 "Some refs in the Zip file are not supported by this datastore"
2198 )
2200 # Transfer the Zip file into the datastore file system.
2201 # There is no RUN as such to use for naming.
2202 # Potentially could use the RUN from the first ref in the index
2203 # There is no requirement that the contents of the Zip files share
2204 # the same RUN.
2205 # Could use the Zip UUID from the index + special "zips/" prefix.
2206 if transfer is None:
2207 # Indicated that the zip file is already in the right place.
2208 if not zip_path.isabs():
2209 tgtLocation = self.locationFactory.fromPath(zip_path.ospath, trusted_path=False)
2210 else:
2211 pathInStore = zip_path.relative_to(self.root)
2212 if pathInStore is None:
2213 raise RuntimeError(
2214 f"Unexpectedly learned that {zip_path} is not within datastore {self.root}"
2215 )
2216 tgtLocation = self.locationFactory.fromPath(pathInStore, trusted_path=True)
2217 elif transfer == "direct":
2218 # Reference in original location.
2219 tgtLocation = None
2220 else:
2221 # Name the zip file based on index contents.
2222 tgtLocation = self.locationFactory.fromPath(index.calculate_zip_file_path_in_store())
2223 if not tgtLocation.uri.dirname().exists():
2224 log.debug("Folder %s does not exist yet.", tgtLocation.uri.dirname())
2225 tgtLocation.uri.dirname().mkdir()
2227 # Transfer the Zip file into the datastore.
2228 if not dry_run:
2229 tgtLocation.uri.transfer_from(
2230 zip_path, transfer=transfer, transaction=self._transaction, overwrite=True
2231 )
2232 else:
2233 log.info("Would be copying Zip from %s to %s", zip_path, tgtLocation)
2235 if tgtLocation is None:
2236 path_in_store = str(zip_path)
2237 else:
2238 path_in_store = tgtLocation.pathInStore.path
2240 # Associate each file with a (DatasetRef, StoredFileInfo) tuple.
2241 artifacts: list[tuple[DatasetRef, StoredFileInfo]] = []
2242 for path_in_zip, index_info in index.artifact_map.items():
2243 # Need to modify the info to include the path to the Zip file
2244 # that was previously written to the datastore.
2245 index_info.info.path = f"{path_in_store}#zip-path={path_in_zip}"
2247 info = StoredFileInfo.from_simple(index_info.info)
2248 for id_ in index_info.ids:
2249 artifacts.append((id_to_ref[id_], info))
2251 if not dry_run:
2252 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT)
2253 else:
2254 log.info("Would be registering %d artifacts from Zip into datastore", len(artifacts))
2256 def get(
2257 self,
2258 ref: DatasetRef,
2259 parameters: Mapping[str, Any] | None = None,
2260 storageClass: StorageClass | str | None = None,
2261 ) -> Any:
2262 """Load an InMemoryDataset from the store.
2264 Parameters
2265 ----------
2266 ref : `DatasetRef`
2267 Reference to the required Dataset.
2268 parameters : `dict`
2269 `StorageClass`-specific parameters that specify, for example,
2270 a slice of the dataset to be loaded.
2271 storageClass : `StorageClass` or `str`, optional
2272 The storage class to be used to override the Python type
2273 returned by this method. By default the returned type matches
2274 the dataset type definition for this dataset. Specifying a
2275 read `StorageClass` can force a different type to be returned.
2276 This type must be compatible with the original type.
2278 Returns
2279 -------
2280 inMemoryDataset : `object`
2281 Requested dataset or slice thereof as an InMemoryDataset.
2283 Raises
2284 ------
2285 FileNotFoundError
2286 Requested dataset can not be retrieved.
2287 TypeError
2288 Return value from formatter has unexpected type.
2289 ValueError
2290 Formatter failed to process the dataset.
2291 """
2292 # Supplied storage class for the component being read is either
2293 # from the ref itself or some an override if we want to force
2294 # type conversion.
2295 if storageClass is not None:
2296 ref = ref.overrideStorageClass(storageClass)
2298 allGetInfo = self._prepare_for_direct_get(ref, parameters)
2299 return get_dataset_as_python_object_from_get_info(
2300 allGetInfo, ref=ref, parameters=parameters, cache_manager=self.cacheManager
2301 )
2303 def prepare_get_for_external_client(self, ref: DatasetRef) -> list[DatasetLocationInformation] | None:
2304 # Docstring inherited
2306 locations = self._get_dataset_locations_info(ref)
2307 if len(locations) == 0:
2308 return None
2310 return locations
2312 @transactional
2313 def put(self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None) -> None:
2314 """Write a InMemoryDataset with a given `DatasetRef` to the store.
2316 Parameters
2317 ----------
2318 inMemoryDataset : `object`
2319 The dataset to store.
2320 ref : `DatasetRef`
2321 Reference to the associated Dataset.
2322 provenance : `DatasetProvenance` or `None`, optional
2323 Any provenance that should be attached to the serialized dataset.
2324 Can be ignored by a formatter or delegate.
2326 Raises
2327 ------
2328 TypeError
2329 Supplied object and storage class are inconsistent.
2330 DatasetTypeNotSupportedError
2331 The associated `DatasetType` is not handled by this datastore.
2333 Notes
2334 -----
2335 If the datastore is configured to reject certain dataset types it
2336 is possible that the put will fail and raise a
2337 `DatasetTypeNotSupportedError`. The main use case for this is to
2338 allow `ChainedDatastore` to put to multiple datastores without
2339 requiring that every datastore accepts the dataset.
2340 """
2341 doDisassembly = self.composites.shouldBeDisassembled(ref)
2342 # doDisassembly = True
2344 artifacts = []
2345 if doDisassembly:
2346 inMemoryDataset = ref.datasetType.storageClass.delegate().add_provenance(
2347 inMemoryDataset, ref, provenance=provenance
2348 )
2349 components = ref.datasetType.storageClass.delegate().disassemble(inMemoryDataset)
2350 if components is None:
2351 raise RuntimeError(
2352 f"Inconsistent configuration: dataset type {ref.datasetType.name} "
2353 f"with storage class {ref.datasetType.storageClass.name} "
2354 "is configured to be disassembled, but cannot be."
2355 )
2356 for component, componentInfo in components.items():
2357 # Don't recurse because we want to take advantage of
2358 # bulk insert -- need a new DatasetRef that refers to the
2359 # same dataset_id but has the component DatasetType
2360 # DatasetType does not refer to the types of components
2361 # So we construct one ourselves.
2362 compRef = ref.makeComponentRef(component)
2363 # Provenance has already been attached above.
2364 storedInfo = self._write_in_memory_to_artifact(componentInfo.component, compRef)
2365 artifacts.append((compRef, storedInfo))
2366 else:
2367 # Write the entire thing out
2368 storedInfo = self._write_in_memory_to_artifact(inMemoryDataset, ref, provenance=provenance)
2369 artifacts.append((ref, storedInfo))
2371 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT)
2373 @transactional
2374 def put_new(self, in_memory_dataset: Any, ref: DatasetRef) -> Mapping[str, DatasetRef]:
2375 doDisassembly = self.composites.shouldBeDisassembled(ref)
2376 # doDisassembly = True
2378 artifacts = []
2379 if doDisassembly:
2380 components = ref.datasetType.storageClass.delegate().disassemble(in_memory_dataset)
2381 if components is None:
2382 raise RuntimeError(
2383 f"Inconsistent configuration: dataset type {ref.datasetType.name} "
2384 f"with storage class {ref.datasetType.storageClass.name} "
2385 "is configured to be disassembled, but cannot be."
2386 )
2387 for component, componentInfo in components.items():
2388 # Don't recurse because we want to take advantage of
2389 # bulk insert -- need a new DatasetRef that refers to the
2390 # same dataset_id but has the component DatasetType
2391 # DatasetType does not refer to the types of components
2392 # So we construct one ourselves.
2393 compRef = ref.makeComponentRef(component)
2394 storedInfo = self._write_in_memory_to_artifact(componentInfo.component, compRef)
2395 artifacts.append((compRef, storedInfo))
2396 else:
2397 # Write the entire thing out
2398 storedInfo = self._write_in_memory_to_artifact(in_memory_dataset, ref)
2399 artifacts.append((ref, storedInfo))
2401 ref_records: DatasetDatastoreRecords = {self._opaque_table_name: [info for _, info in artifacts]}
2402 ref = ref.replace(datastore_records=ref_records)
2403 return {self.name: ref}
2405 @transactional
2406 def trash(self, ref: DatasetRef | Iterable[DatasetRef], ignore_errors: bool = True) -> None:
2407 # At this point can safely remove these datasets from the cache
2408 # to avoid confusion later on. If they are not trashed later
2409 # the cache will simply be refilled.
2410 self.cacheManager.remove_from_cache(ref)
2412 # If we are in trust mode there will be nothing to move to
2413 # the trash table and we will have to try to delete the file
2414 # immediately.
2415 if self.trustGetRequest:
2416 # Try to keep the logic below for a single file trash.
2417 if isinstance(ref, DatasetRef):
2418 refs = {ref}
2419 else:
2420 # Will recreate ref at the end of this branch.
2421 refs = set(ref)
2423 # Determine which datasets are known to datastore directly.
2424 id_to_ref = {ref.id: ref for ref in refs}
2425 existing_ids = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2426 existing_refs = {id_to_ref[ref_id] for ref_id in existing_ids}
2428 missing = refs - existing_refs
2429 if missing:
2430 # Do an explicit existence check on these refs.
2431 # We only care about the artifacts at this point and not
2432 # the dataset existence.
2433 artifact_existence: dict[ResourcePath, bool] = {}
2434 _ = self.mexists(missing, artifact_existence)
2435 uris = [uri for uri, exists in artifact_existence.items() if exists]
2437 # FUTURE UPGRADE: Implement a parallelized bulk remove.
2438 log.debug("Removing %d artifacts from datastore that are unknown to datastore", len(uris))
2439 for uri in uris:
2440 try:
2441 uri.remove()
2442 except Exception as e:
2443 if ignore_errors:
2444 log.debug("Artifact %s could not be removed: %s", uri, e)
2445 continue
2446 raise
2448 # There is no point asking the code below to remove refs we
2449 # know are missing so update it with the list of existing
2450 # records. Try to retain one vs many logic.
2451 if not existing_refs:
2452 # Nothing more to do since none of the datasets were
2453 # known to the datastore record table.
2454 return
2455 ref = list(existing_refs)
2456 if len(ref) == 1:
2457 ref = ref[0]
2459 # Get file metadata and internal metadata
2460 if not isinstance(ref, DatasetRef):
2461 log.debug("Doing multi-dataset trash in datastore %s", self.name)
2462 # Assumed to be an iterable of refs so bulk mode enabled.
2463 try:
2464 self.bridge.moveToTrash(ref, transaction=self._transaction)
2465 except Exception as e:
2466 if ignore_errors:
2467 log.warning("Unexpected issue moving multiple datasets to trash: %s", e)
2468 else:
2469 raise
2470 return
2472 log.debug("Trashing dataset %s in datastore %s", ref, self.name)
2474 fileLocations = self._get_dataset_locations_info(ref)
2476 if not fileLocations:
2477 err_msg = f"Requested dataset to trash ({ref}) is not known to datastore {self.name}"
2478 if ignore_errors:
2479 log.warning(err_msg)
2480 return
2481 else:
2482 raise FileNotFoundError(err_msg)
2484 for location, _ in fileLocations:
2485 if not self._artifact_exists(location):
2486 err_msg = (
2487 f"Dataset is known to datastore {self.name} but "
2488 f"associated artifact ({location.uri}) is missing"
2489 )
2490 if ignore_errors:
2491 log.warning(err_msg)
2492 return
2493 else:
2494 raise FileNotFoundError(err_msg)
2496 # Mark dataset as trashed
2497 try:
2498 self.bridge.moveToTrash([ref], transaction=self._transaction)
2499 except Exception as e:
2500 if ignore_errors:
2501 log.warning(
2502 "Attempted to mark dataset (%s) to be trashed in datastore %s "
2503 "but encountered an error: %s",
2504 ref,
2505 self.name,
2506 e,
2507 )
2508 pass
2509 else:
2510 raise
2512 def emptyTrash(
2513 self, ignore_errors: bool = True, refs: Collection[DatasetRef] | None = None, dry_run: bool = False
2514 ) -> set[ResourcePath]:
2515 """Remove all datasets from the trash.
2517 Parameters
2518 ----------
2519 ignore_errors : `bool`
2520 If `True` return without error even if something went wrong.
2521 Problems could occur if another process is simultaneously trying
2522 to delete.
2523 refs : `collections.abc.Collection` [ `DatasetRef` ] or `None`
2524 Explicit list of datasets that can be removed from trash. If listed
2525 datasets are not already stored in the trash table they will be
2526 ignored. If `None` every entry in the trash table will be
2527 processed.
2528 dry_run : `bool`, optional
2529 If `True`, the trash table will be queried and results reported
2530 but no artifacts will be removed.
2532 Returns
2533 -------
2534 removed : `set` [ `lsst.resources.ResourcePath` ]
2535 List of artifacts that were removed.
2537 Notes
2538 -----
2539 Will empty the records from the trash tables only if this call finishes
2540 without raising.
2541 """
2542 removed = set()
2543 if refs:
2544 selected_ids = {ref.id for ref in refs}
2545 chunk_size = 50_000
2546 n_chunks = math.ceil(len(selected_ids) / chunk_size)
2547 chunk_num = 0
2548 for chunk in chunk_iterable(selected_ids, chunk_size=chunk_size):
2549 chunk_num += 1
2550 if n_chunks == 1:
2551 log.verbose(
2552 "Emptying datastore trash for %d dataset%s",
2553 len(chunk),
2554 "s" if len(chunk) != 1 else "",
2555 )
2556 else:
2557 log.verbose(
2558 "Emptying datastore trash for chunk %d out of %d of size %d",
2559 chunk_num,
2560 n_chunks,
2561 len(chunk),
2562 )
2563 removed.update(
2564 self._empty_trash_subset(ignore_errors=ignore_errors, selected_ids=chunk, dry_run=dry_run)
2565 )
2566 else:
2567 log.verbose("Emptying all trash in datastore %s", self.name)
2568 removed = self._empty_trash_subset(ignore_errors=ignore_errors, dry_run=dry_run)
2569 log.info(
2570 "%sRemoved %d file artifact%s from datastore %s",
2571 "Would have " if dry_run else "",
2572 len(removed),
2573 "s" if len(removed) != 1 else "",
2574 self.name,
2575 )
2576 return removed
2578 @transactional
2579 def _empty_trash_subset(
2580 self,
2581 *,
2582 ignore_errors: bool = True,
2583 selected_ids: Collection[DatasetId] | None = None,
2584 dry_run: bool = False,
2585 ) -> set[ResourcePath]:
2586 """Empty trash table in transaction.
2588 Parameters
2589 ----------
2590 ignore_errors : `bool`
2591 If `True` return without error even if something went wrong.
2592 Problems could occur if another process is simultaneously trying
2593 to delete.
2594 selected_ids : `collections.abc.collection` [`DatasetId`] or `None`
2595 Explicit list of dataset IDs that can be removed from the trash.
2596 If listed datasets are not already included in the trash table
2597 they will be ignored. If `None` every entry in the trash table
2598 will be processed.
2599 dry_run : `bool`, optional
2600 If `True`, the trash table will be queried and results reported
2601 but no artifacts will be removed.
2603 Returns
2604 -------
2605 removed : `set` [ `lsst.resources.ResourcePath` ]
2606 Artifacts successfully removed.
2608 Notes
2609 -----
2610 Will empty the records from the trash tables only if this call finishes
2611 without raising.
2612 """
2613 # Context manager will empty trash iff we finish it without raising.
2614 # It will also automatically delete the relevant rows from the
2615 # trash table and the records table.
2616 with self.bridge.emptyTrash(
2617 self._table,
2618 record_class=StoredFileInfo,
2619 record_column="path",
2620 selected_ids=selected_ids,
2621 dry_run=dry_run,
2622 ) as trash_data:
2623 # Removing the artifacts themselves requires that the files are
2624 # not also associated with refs that are not to be trashed.
2625 # Therefore need to do a query with the file paths themselves
2626 # and return all the refs associated with them. Can only delete
2627 # a file if the refs to be trashed are the only refs associated
2628 # with the file.
2629 # This requires multiple copies of the trashed items
2630 trashed, artifacts_to_keep = trash_data
2632 # Assume that # in path means there are fragments involved. The
2633 # fragments can not be handled by the emptyTrash bridge call
2634 # so need to be processed independently.
2635 # The generator has to be converted to a list for multiple
2636 # iterations. Clean up the typing so that multiple isinstance
2637 # tests aren't needed later.
2638 trashed_list = [(ref, ninfo) for ref, ninfo in trashed if isinstance(ninfo, StoredFileInfo)]
2640 if artifacts_to_keep is None or any("#" in info[1].path for info in trashed_list):
2641 # The bridge is not helping us so have to work it out
2642 # ourselves. This is not going to be as efficient.
2643 # This mapping does not include the fragments.
2644 if artifacts_to_keep is not None:
2645 # This means we have already checked for non-fragment
2646 # examples so can filter.
2647 paths_to_check = {info.path for _, info in trashed_list if "#" in info.path}
2648 else:
2649 paths_to_check = {info.path for _, info in trashed_list}
2651 path_map = self._refs_associated_with_artifacts(paths_to_check)
2653 for ref, info in trashed_list:
2654 path = info.artifact_path
2655 # For disassembled composites in a Zip it is possible
2656 # for the same path to correspond to the same dataset ref
2657 # multiple times so trap for that.
2658 if ref.id in path_map[path]:
2659 path_map[path].remove(ref.id)
2660 if not path_map[path]:
2661 del path_map[path]
2663 slow_artifacts_to_keep = set(path_map)
2664 if artifacts_to_keep is not None:
2665 artifacts_to_keep.update(slow_artifacts_to_keep)
2666 else:
2667 artifacts_to_keep = slow_artifacts_to_keep
2669 n_direct = 0
2670 artifacts_to_delete: set[ResourcePath] = set()
2671 for ref, info in trashed_list:
2672 # Should not happen for this implementation but need
2673 # to keep mypy happy.
2674 assert info is not None, f"Internal logic error in emptyTrash with ref {ref}."
2676 if info.artifact_path in artifacts_to_keep:
2677 # This is a multi-dataset artifact and we are not
2678 # removing all associated refs.
2679 continue
2681 # Only trashed refs still known to datastore will be returned.
2682 location = info.file_location(self.locationFactory)
2684 if location.pathInStore.isabs():
2685 n_direct += 1
2686 continue
2688 # Strip fragment before storing since it is the artifact
2689 # we are deleting and we do not want repeats for every member
2690 # in a zip.
2691 artifacts_to_delete.add(location.uri.replace(fragment=""))
2693 if n_direct > 0:
2694 s = "s" if n_direct != 1 else ""
2695 log.verbose("Not deleting %d artifact%s using absolute URI%s", n_direct, s, s)
2697 if artifacts_to_keep:
2698 log.verbose(
2699 "%d artifact%s %s not deleted because of association with other datasets",
2700 len(artifacts_to_keep),
2701 "s" if len(artifacts_to_keep) != 1 else "",
2702 "were" if len(artifacts_to_keep) != 1 else "was",
2703 )
2705 if not artifacts_to_delete:
2706 return set()
2708 # Now do the deleting. Special case the log message for a single
2709 # artifact.
2710 if len(artifacts_to_delete) == 1:
2711 log.verbose(
2712 "%s removing file artifact %s from datastore %s",
2713 "Would be" if dry_run else "Now",
2714 list(artifacts_to_delete)[0],
2715 self.name,
2716 )
2717 else:
2718 log.verbose(
2719 "%s removing %d file artifacts from datastore %s",
2720 "Would be" if dry_run else "Now",
2721 len(artifacts_to_delete),
2722 self.name,
2723 )
2725 # For dry-run mode do not attempt to search the file store for
2726 # the artifacts to determine whether they exist or not. Simply
2727 # report that an attempt would be made to delete them. Never
2728 # report direct imports.
2729 if dry_run:
2730 return artifacts_to_delete
2732 # Now remove the actual file artifacts.
2733 remove_result = ResourcePath.mremove(artifacts_to_delete, do_raise=False)
2735 removed: set[ResourcePath] = set()
2736 exceptions: list[Exception] = []
2737 for uri, result in remove_result.items():
2738 if result.exception is None or isinstance(result.exception, FileNotFoundError):
2739 # File not existing is not an error since some other
2740 # process might have been trying to clean it and we do not
2741 # want to raise an error for a situation where the file
2742 # is not there and we do not want it to be there.
2743 removed.add(uri)
2744 else:
2745 exceptions.append(result.exception)
2747 if exceptions:
2748 s_err = "s" if len(exceptions) != 1 else ""
2749 e = ExceptionGroup(f"Error{s_err} removing {len(exceptions)} artifact{s_err}", exceptions)
2750 if ignore_errors:
2751 # Use a debug message here even though it's not
2752 # a good situation. In some cases this can be
2753 # caused by a race between user A and user B
2754 # and neither of them has permissions for the
2755 # other's files. Butler does not know about users
2756 # and trash has no idea what collections these
2757 # files were in (without guessing from a path).
2758 log.debug(
2759 "Encountered %d error%s removing %d artifact%s from datastore %s: %s",
2760 len(exceptions),
2761 s_err,
2762 len(artifacts_to_delete),
2763 "s" if len(artifacts_to_delete) != 1 else "",
2764 self.name,
2765 e,
2766 )
2767 else:
2768 raise e
2769 return removed
2771 @transactional
2772 def transfer_from(
2773 self,
2774 source_records: FileTransferMap,
2775 refs: Collection[DatasetRef],
2776 transfer: str = "auto",
2777 artifact_existence: dict[ResourcePath, bool] | None = None,
2778 dry_run: bool = False,
2779 ) -> tuple[set[DatasetRef], set[DatasetRef]]:
2780 log.verbose("Transferring %d datasets to %s", len(refs), self.name)
2782 # Stop early if "direct" transfer mode is requested. That would
2783 # require that the URI inside the source datastore should be stored
2784 # directly in the target datastore, which seems unlikely to be useful
2785 # since at any moment the source datastore could delete the file.
2786 if transfer in ("direct", "split"):
2787 raise ValueError(
2788 f"Can not transfer from a source datastore using {transfer} mode since"
2789 " those files are controlled by the other datastore."
2790 )
2792 if not refs:
2793 return set(), set()
2795 # Empty existence lookup if none given.
2796 if artifact_existence is None:
2797 artifact_existence = {}
2799 # In order to handle disassembled composites the code works
2800 # at the records level since it can assume that internal APIs
2801 # can be used.
2802 # - If the record already exists in the destination this is assumed
2803 # to be okay.
2804 # - If there is no record but the source and destination URIs are
2805 # identical no transfer is done but the record is added.
2806 # - If the source record refers to an absolute URI currently assume
2807 # that that URI should remain absolute and will be visible to the
2808 # destination butler. May need to have a flag to indicate whether
2809 # the dataset should be transferred. This will only happen if
2810 # the detached Butler has had a local ingest.
2812 # See if we already have these records
2813 log.verbose("Looking up existing datastore records in target %s for %d refs", self.name, len(refs))
2814 target_records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2816 # The artifacts to register
2817 artifacts = []
2819 # Refs that already exist
2820 already_present = []
2822 # Refs that were rejected by this datastore.
2823 rejected = set()
2825 # Refs that were transferred successfully.
2826 accepted = set()
2828 # Record each time we have done a "direct" transfer.
2829 direct_transfers = []
2831 # Keep track of all the file transfers that are required.
2832 from_to: list[tuple[ResourcePath, ResourcePath]] = []
2834 # Now can transfer the artifacts
2835 log.verbose("Transferring artifacts")
2836 for ref in refs:
2837 if not self.constraints.isAcceptable(ref):
2838 # This datastore should not be accepting this dataset.
2839 rejected.add(ref)
2840 continue
2842 accepted.add(ref)
2844 if ref.id in target_records:
2845 # Already have an artifact for this.
2846 already_present.append(ref)
2847 continue
2849 # mypy needs to know these are always resolved refs
2850 for transfer_info in source_records.get(ref.id, []):
2851 info = transfer_info.file_info
2852 source_location = transfer_info.location
2853 target_location = info.file_location(self.locationFactory)
2854 if source_location == target_location and not source_location.pathInStore.isabs():
2855 # Artifact is already in the target location.
2856 # (which is how execution butler currently runs)
2857 pass
2858 else:
2859 if target_location.pathInStore.isabs():
2860 # Just because we can see the artifact when running
2861 # the transfer doesn't mean it will be generally
2862 # accessible to a user of this butler. Need to decide
2863 # what to do about an absolute path.
2864 if transfer == "auto":
2865 # For "auto" transfers we allow the absolute URI
2866 # to be recorded in the target datastore.
2867 direct_transfers.append(source_location)
2868 else:
2869 # The user is explicitly requesting a transfer
2870 # even for an absolute URI. This requires us to
2871 # calculate the target path.
2872 template_ref = ref
2873 if info.component:
2874 template_ref = ref.makeComponentRef(info.component)
2875 target_location = self._calculate_ingested_datastore_name(
2876 source_location.uri,
2877 template_ref,
2878 )
2880 info = info.update(path=target_location.pathInStore.path)
2882 # Need to transfer it to the new location.
2883 from_to.append((source_location.uri, target_location.uri))
2885 artifacts.append((ref, info))
2887 # Do the file transfers in bulk.
2888 # Assume we should always overwrite. If the artifact
2889 # is there this might indicate that a previous transfer
2890 # was interrupted but was not able to be rolled back
2891 # completely (eg pre-emption) so follow Datastore default
2892 # and overwrite. Do not copy if we are in dry-run mode.
2893 if dry_run:
2894 log.info("Would be copying %d file artifacts", len(from_to))
2895 else:
2896 log.verbose("Copying %d file artifacts", len(from_to))
2897 with time_this(log, msg="Transferring datasets into datastore", level=VERBOSE):
2898 ResourcePath.mtransfer(
2899 transfer,
2900 from_to,
2901 overwrite=True,
2902 transaction=self._transaction,
2903 )
2905 if direct_transfers:
2906 log.info(
2907 "Transfer request for an outside-datastore artifact with absolute URI done %d time%s",
2908 len(direct_transfers),
2909 "" if len(direct_transfers) == 1 else "s",
2910 )
2912 # We are overwriting previous datasets that may have already
2913 # existed. We therefore should ensure that we force the
2914 # datastore records to agree. Note that this can potentially lead
2915 # to difficulties if the dataset has previously been ingested
2916 # disassembled and is somehow now assembled, or vice versa.
2917 if not dry_run:
2918 log.verbose("Registering datastore records in database")
2919 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.REPLACE)
2921 if already_present:
2922 n_skipped = len(already_present)
2923 log.info(
2924 "Skipped transfer of %d dataset%s already present in datastore",
2925 n_skipped,
2926 "" if n_skipped == 1 else "s",
2927 )
2929 log.verbose(
2930 "Finished transfer_from to %s with %d accepted, %d rejected",
2931 self.name,
2932 len(accepted),
2933 len(rejected),
2934 )
2935 return accepted, rejected
2937 def get_file_info_for_transfer(self, dataset_ids: Iterable[DatasetId]) -> FileTransferMap:
2938 source_records = self._get_stored_records_associated_with_refs(
2939 [FakeDatasetRef(id) for id in dataset_ids], ignore_datastore_records=True
2940 )
2941 return self._convert_stored_file_info_to_file_transfer_record(source_records)
2943 def locate_missing_files_for_transfer(
2944 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool]
2945 ) -> FileTransferMap:
2946 missing_ids = {ref.id for ref in refs}
2947 # Missing IDs can be okay if that datastore has allowed
2948 # gets based on file existence. Should we transfer what we can
2949 # or complain about it and warn?
2950 if not self.trustGetRequest:
2951 return {}
2953 found_records = self._find_missing_records(
2954 refs, missing_ids, artifact_existence, warn_for_missing=False
2955 )
2956 return self._convert_stored_file_info_to_file_transfer_record(found_records)
2958 def _convert_stored_file_info_to_file_transfer_record(
2959 self, info_map: dict[DatasetId, list[StoredFileInfo]]
2960 ) -> FileTransferMap:
2961 output: dict[DatasetId, list[FileTransferRecord]] = {}
2962 for k, file_info_list in info_map.items():
2963 output[k] = [
2964 FileTransferRecord(file_info=info, location=info.file_location(self.locationFactory))
2965 for info in file_info_list
2966 ]
2967 return output
2969 @transactional
2970 def forget(self, refs: Iterable[DatasetRef]) -> None:
2971 # Docstring inherited.
2972 refs = list(refs)
2973 self.bridge.forget(refs)
2974 self._table.delete(["dataset_id"], *[{"dataset_id": ref.id} for ref in refs])
2976 def validateConfiguration(
2977 self, entities: Iterable[DatasetRef | DatasetType | StorageClass], logFailures: bool = False
2978 ) -> None:
2979 """Validate some of the configuration for this datastore.
2981 Parameters
2982 ----------
2983 entities : iterable of `DatasetRef`, `DatasetType`, or `StorageClass`
2984 Entities to test against this configuration. Can be differing
2985 types.
2986 logFailures : `bool`, optional
2987 If `True`, output a log message for every validation error
2988 detected.
2990 Raises
2991 ------
2992 DatastoreValidationError
2993 Raised if there is a validation problem with a configuration.
2994 All the problems are reported in a single exception.
2996 Notes
2997 -----
2998 This method checks that all the supplied entities have valid file
2999 templates and also have formatters defined.
3000 """
3001 templateFailed = None
3002 try:
3003 self.templates.validateTemplates(entities, logFailures=logFailures)
3004 except FileTemplateValidationError as e:
3005 templateFailed = str(e)
3007 formatterFailed = []
3008 for entity in entities:
3009 try:
3010 self.formatterFactory.getFormatterClass(entity)
3011 except KeyError as e:
3012 formatterFailed.append(str(e))
3013 if logFailures:
3014 log.critical("Formatter failure: %s", e)
3016 if templateFailed or formatterFailed:
3017 messages = []
3018 if templateFailed:
3019 messages.append(templateFailed)
3020 if formatterFailed:
3021 messages.append(",".join(formatterFailed))
3022 msg = ";\n".join(messages)
3023 raise DatastoreValidationError(msg)
3025 def getLookupKeys(self) -> set[LookupKey]:
3026 # Docstring is inherited from base class
3027 return (
3028 self.templates.getLookupKeys()
3029 | self.formatterFactory.getLookupKeys()
3030 | self.constraints.getLookupKeys()
3031 )
3033 def validateKey(self, lookupKey: LookupKey, entity: DatasetRef | DatasetType | StorageClass) -> None:
3034 # Docstring is inherited from base class
3035 # The key can be valid in either formatters or templates so we can
3036 # only check the template if it exists
3037 if lookupKey in self.templates:
3038 try:
3039 self.templates[lookupKey].validateTemplate(entity)
3040 except FileTemplateValidationError as e:
3041 raise DatastoreValidationError(e) from e
3043 def export(
3044 self,
3045 refs: Iterable[DatasetRef],
3046 *,
3047 directory: ResourcePathExpression | None = None,
3048 transfer: str | None = "auto",
3049 ) -> Iterable[FileDataset]:
3050 # Docstring inherited from Datastore.export.
3051 if transfer == "auto" and directory is None:
3052 transfer = None
3054 if transfer is not None and transfer != "direct" and directory is None:
3055 raise TypeError(f"Cannot export using transfer mode {transfer} with no export directory given")
3057 if transfer == "move":
3058 raise TypeError("Can not export by moving files out of datastore.")
3060 # Force the directory to be a URI object
3061 directoryUri: ResourcePath | None = None
3062 if directory is not None:
3063 directoryUri = ResourcePath(directory, forceDirectory=True)
3065 if transfer is not None and directoryUri is not None and not directoryUri.exists():
3066 # mypy needs the second test
3067 raise FileNotFoundError(f"Export location {directory} does not exist")
3069 progress = Progress("lsst.daf.butler.datastores.FileDatastore.export", level=logging.DEBUG)
3070 for ref in progress.wrap(refs, "Exporting dataset files"):
3071 fileLocations = self._get_dataset_locations_info(ref)
3072 if not fileLocations:
3073 raise FileNotFoundError(f"Could not retrieve dataset {ref}.")
3074 # For now we can not export disassembled datasets
3075 if len(fileLocations) > 1:
3076 raise NotImplementedError(f"Can not export disassembled datasets such as {ref}")
3077 location, storedFileInfo = fileLocations[0]
3079 pathInStore = location.pathInStore.path
3080 if transfer is None:
3081 # TODO: do we also need to return the readStorageClass somehow?
3082 # We will use the path in store directly. If this is an
3083 # absolute URI, preserve it.
3084 if location.pathInStore.isabs():
3085 pathInStore = str(location.uri)
3086 elif transfer == "direct":
3087 # Use full URIs to the remote store in the export
3088 pathInStore = str(location.uri)
3089 else:
3090 # mypy needs help
3091 assert directoryUri is not None, "directoryUri must be defined to get here"
3092 storeUri = ResourcePath(location.uri, forceDirectory=False)
3094 # if the datastore has an absolute URI to a resource, we
3095 # have two options:
3096 # 1. Keep the absolute URI in the exported YAML
3097 # 2. Allocate a new name in the local datastore and transfer
3098 # it.
3099 # For now go with option 2
3100 if location.pathInStore.isabs():
3101 template = self.templates.getTemplate(ref)
3102 newURI = ResourcePath(template.format(ref), forceAbsolute=False, forceDirectory=False)
3103 pathInStore = str(newURI.updatedExtension(location.pathInStore.getExtension()))
3105 exportUri = directoryUri.join(pathInStore)
3106 exportUri.transfer_from(storeUri, transfer=transfer)
3108 yield FileDataset(refs=[ref], path=pathInStore, formatter=storedFileInfo.formatter)
3110 @staticmethod
3111 def computeChecksum(uri: ResourcePath, algorithm: str = "blake2b", block_size: int = 8192) -> str | None:
3112 """Compute the checksum of the supplied file.
3114 Parameters
3115 ----------
3116 uri : `lsst.resources.ResourcePath`
3117 Name of resource to calculate checksum from.
3118 algorithm : `str`, optional
3119 Name of algorithm to use. Must be one of the algorithms supported
3120 by :py:class`hashlib`.
3121 block_size : `int`
3122 Number of bytes to read from file at one time.
3124 Returns
3125 -------
3126 hexdigest : `str`
3127 Hex digest of the file.
3129 Notes
3130 -----
3131 Currently returns None if the URI is for a remote resource.
3132 """
3133 if algorithm not in hashlib.algorithms_guaranteed:
3134 raise NameError(f"The specified algorithm '{algorithm}' is not supported by hashlib")
3136 if not uri.isLocal:
3137 return None
3139 hasher = hashlib.new(algorithm)
3141 with uri.as_local() as local_uri, open(local_uri.ospath, "rb") as f:
3142 for chunk in iter(lambda: f.read(block_size), b""):
3143 hasher.update(chunk)
3145 return hasher.hexdigest()
3147 def needs_expanded_data_ids(
3148 self,
3149 transfer: str | None,
3150 entity: DatasetRef | DatasetType | StorageClass | None = None,
3151 ) -> bool:
3152 # Docstring inherited.
3153 # This _could_ also use entity to inspect whether the filename template
3154 # involves placeholders other than the required dimensions for its
3155 # dataset type, but that's not necessary for correctness; it just
3156 # enables more optimizations (perhaps only in theory).
3157 return transfer not in ("direct", None)
3159 def import_records(self, data: Mapping[str, DatastoreRecordData]) -> None:
3160 # Docstring inherited from the base class.
3161 record_data = data.get(self.name)
3162 if not record_data:
3163 return
3165 self._bridge.insert(FakeDatasetRef(dataset_id) for dataset_id in record_data.records)
3167 # TODO: Verify that there are no unexpected table names in the dict?
3168 unpacked_records = []
3169 for dataset_id, dataset_data in record_data.records.items():
3170 records = dataset_data.get(self._table.name)
3171 if records:
3172 for info in records:
3173 assert isinstance(info, StoredFileInfo), "Expecting StoredFileInfo records"
3174 unpacked_records.append(info.to_record(dataset_id=dataset_id))
3175 if unpacked_records:
3176 self._table.insert(*unpacked_records, transaction=self._transaction)
3178 def export_records(self, refs: Iterable[DatasetIdRef]) -> Mapping[str, DatastoreRecordData]:
3179 # Docstring inherited from the base class.
3180 exported_refs = list(self._bridge.check(refs))
3181 ids = {ref.id for ref in exported_refs}
3182 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {id: {} for id in ids}
3183 for row in self._table.fetch(dataset_id=ids):
3184 info: StoredDatastoreItemInfo = StoredFileInfo.from_record(row)
3185 dataset_records = records.setdefault(row["dataset_id"], {})
3186 dataset_records.setdefault(self._table.name, []).append(info)
3188 record_data = DatastoreRecordData(records=records)
3189 return {self.name: record_data}
3191 def export_predicted_records(self, refs: Iterable[DatasetRef]) -> dict[str, DatastoreRecordData]:
3192 # Docstring inherited from the base class.
3193 refs = [self._cast_storage_class(ref) for ref in refs]
3194 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {}
3195 for ref in refs:
3196 if not self.constraints.isAcceptable(ref):
3197 continue
3198 fileLocations = self._get_expected_dataset_locations_info(ref)
3199 if not fileLocations:
3200 continue
3201 dataset_records = records.setdefault(ref.id, {})
3202 dataset_records.setdefault(self._table.name, [])
3203 for _, storedFileInfo in fileLocations:
3204 dataset_records[self._table.name].append(storedFileInfo)
3206 record_data = DatastoreRecordData(records=records)
3207 return {self.name: record_data}
3209 def set_retrieve_dataset_type_method(self, method: Callable[[str], DatasetType | None] | None) -> None:
3210 # Docstring inherited from the base class.
3211 self._retrieve_dataset_method = method
3213 def _cast_storage_class(self, ref: DatasetRef) -> DatasetRef:
3214 """Update dataset reference to use the storage class from registry."""
3215 if self._retrieve_dataset_method is None:
3216 # We could raise an exception here but unit tests do not define
3217 # this method.
3218 return ref
3219 dataset_type = self._retrieve_dataset_method(ref.datasetType.name)
3220 if dataset_type is not None:
3221 ref = ref.overrideStorageClass(dataset_type.storageClass_name)
3222 return ref
3224 def get_opaque_table_definitions(self) -> Mapping[str, DatastoreOpaqueTable]:
3225 # Docstring inherited from the base class.
3226 return {self._opaque_table_name: DatastoreOpaqueTable(self.makeTableSpec(), StoredFileInfo)}