Coverage for python / lsst / daf / butler / datastores / fileDatastore.py: 8%
1059 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-13 01:20 -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)
1072 # if we are transferring from a local file to a remote location
1073 # it may be more efficient to get the size and checksum of the
1074 # local file rather than the transferred one
1075 if record_validation_info and srcUri.isLocal:
1076 size = srcUri.size()
1077 checksum = self.computeChecksum(srcUri) if self.useChecksum else None
1078 have_sized = True
1080 # Transfer the resource to the destination.
1081 # Allow overwrite of an existing file. This matches the behavior
1082 # of datastore.put() in that it trusts that registry would not
1083 # be asking to overwrite unless registry thought that the
1084 # overwrite was allowed.
1085 tgtLocation.uri.transfer_from(
1086 srcUri, transfer=transfer, transaction=self._transaction, overwrite=True
1087 )
1089 if tgtLocation is None:
1090 # This means we are using direct mode
1091 targetUri = srcUri
1092 targetPath = str(srcUri)
1093 else:
1094 targetUri = tgtLocation.uri
1095 targetPath = tgtLocation.pathInStore.path
1097 # the file should exist in the datastore now
1098 if record_validation_info:
1099 if not have_sized:
1100 size = targetUri.size()
1101 checksum = self.computeChecksum(targetUri) if self.useChecksum else None
1102 else:
1103 # Not recording any file information.
1104 size = -1
1105 checksum = None
1107 return StoredFileInfo(
1108 formatter=formatter,
1109 path=targetPath,
1110 storageClass=ref.datasetType.storageClass,
1111 component=ref.datasetType.component(),
1112 file_size=size,
1113 checksum=checksum,
1114 )
1116 def _prepIngest(self, *datasets: FileDataset, transfer: str | None = None) -> _IngestPrepData:
1117 # Docstring inherited from Datastore._prepIngest.
1118 filtered = []
1120 # Ingest could be given tens of thousands of files. It is not efficient
1121 # to check for the existence of every single file (especially if they
1122 # are remote URIs) but in some transfer modes the files will be checked
1123 # anyhow when they are relocated. For direct or None transfer modes
1124 # it is possible to not know if the file is accessible at all.
1125 # Therefore limit number of files that will be checked (but always
1126 # include the first one).
1127 max_checks = 200
1128 n_datasets = len(datasets)
1129 if n_datasets <= max_checks:
1130 check_every_n = 1
1131 elif transfer in ("direct", None):
1132 check_every_n = int(n_datasets / max_checks + 1) # +1 so that if n < max_checks the answer is 1.
1133 else:
1134 check_every_n = 0
1136 for count, dataset in enumerate(datasets):
1137 acceptable = [ref for ref in dataset.refs if self.constraints.isAcceptable(ref)]
1138 if not acceptable:
1139 continue
1140 else:
1141 dataset.refs = acceptable
1142 if dataset.formatter is None:
1143 dataset.formatter = self.formatterFactory.getFormatterClass(dataset.refs[0])
1144 else:
1145 assert isinstance(dataset.formatter, type | str)
1146 formatter_class = get_class_of(dataset.formatter)
1147 if not issubclass(formatter_class, Formatter | FormatterV2):
1148 raise TypeError(f"Requested formatter {dataset.formatter} is not a Formatter class.")
1149 dataset.formatter = formatter_class
1151 # Decide whether the file should be checked.
1152 check_existence = False
1153 if check_every_n != 0:
1154 # First time through count is 0 so we guarantee to check
1155 # the first file but not necessarily the final one.
1156 check_existence = count % check_every_n == 0
1158 if check_existence:
1159 log.debug(
1160 "Checking file existence: %s (%d/%d) [%s]",
1161 check_existence,
1162 count + 1,
1163 n_datasets,
1164 transfer,
1165 )
1167 dataset.path = self._standardizeIngestPath(
1168 dataset.path, transfer=transfer, check_existence=check_existence
1169 )
1170 filtered.append(dataset)
1171 return _IngestPrepData(filtered)
1173 @transactional
1174 def _finishIngest(
1175 self,
1176 prepData: Datastore.IngestPrepData,
1177 *,
1178 transfer: str | None = None,
1179 record_validation_info: bool = True,
1180 ) -> None:
1181 # Docstring inherited from Datastore._finishIngest.
1182 refsAndInfos = []
1183 progress = Progress("lsst.daf.butler.datastores.FileDatastore.ingest", level=logging.DEBUG)
1184 for dataset in progress.wrap(prepData.datasets, desc="Ingesting dataset files"):
1185 # Do ingest as if the first dataset ref is associated with the file
1186 info = self._extractIngestInfo(
1187 dataset.path,
1188 dataset.refs[0],
1189 formatter=dataset.formatter,
1190 transfer=transfer,
1191 record_validation_info=record_validation_info,
1192 )
1193 refsAndInfos.extend([(ref, info) for ref in dataset.refs])
1195 # In direct mode we can allow repeated ingests of the same thing
1196 # if we are sure that the external dataset is immutable. We use
1197 # UUIDv5 to indicate this. If there is a mix of v4 and v5 they are
1198 # separated.
1199 refs_and_infos_replace = []
1200 refs_and_infos_insert = []
1201 if transfer == "direct":
1202 for entry in refsAndInfos:
1203 if entry[0].id.version == 5:
1204 refs_and_infos_replace.append(entry)
1205 else:
1206 refs_and_infos_insert.append(entry)
1207 else:
1208 refs_and_infos_insert = refsAndInfos
1210 if refs_and_infos_insert:
1211 self._register_datasets(refs_and_infos_insert, insert_mode=DatabaseInsertMode.INSERT)
1212 if refs_and_infos_replace:
1213 self._register_datasets(refs_and_infos_replace, insert_mode=DatabaseInsertMode.REPLACE)
1215 def _calculate_ingested_datastore_name(
1216 self,
1217 srcUri: ResourcePath,
1218 ref: DatasetRef,
1219 formatter: Formatter | FormatterV2 | type[Formatter | FormatterV2] | None = None,
1220 ) -> Location:
1221 """Given a source URI and a DatasetRef, determine the name the
1222 dataset will have inside datastore.
1224 Parameters
1225 ----------
1226 srcUri : `lsst.resources.ResourcePath`
1227 URI to the source dataset file.
1228 ref : `DatasetRef`
1229 Ref associated with the newly-ingested dataset artifact. This
1230 is used to determine the name within the datastore.
1231 formatter : `Formatter` or Formatter class.
1232 Formatter to use for validation. Can be a class or an instance.
1233 No validation of the file extension is performed if the
1234 ``formatter`` is `None`. This can be used if the caller knows
1235 that the source URI and target URI will use the same formatter.
1237 Returns
1238 -------
1239 location : `Location`
1240 Target location for the newly-ingested dataset.
1241 """
1242 # Ingesting a file from outside the datastore.
1243 # This involves a new name.
1244 template = self.templates.getTemplate(ref)
1245 location = self.locationFactory.fromPath(template.format(ref), trusted_path=True)
1247 # Get the extension
1248 ext = srcUri.getExtension()
1250 # Update the destination to include that extension
1251 location.updateExtension(ext)
1253 # Ask the formatter to validate this extension
1254 if formatter is not None:
1255 formatter.validate_extension(location)
1257 return location
1259 def _write_in_memory_to_artifact(
1260 self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None
1261 ) -> StoredFileInfo:
1262 """Write out in memory dataset to datastore.
1264 Parameters
1265 ----------
1266 inMemoryDataset : `object`
1267 Dataset to write to datastore.
1268 ref : `DatasetRef`
1269 Registry information associated with this dataset.
1270 provenance : `DatasetProvenance` or `None`, optional
1271 Any provenance that should be attached to the serialized dataset.
1272 Not supported by all formatters.
1274 Returns
1275 -------
1276 info : `StoredFileInfo`
1277 Information describing the artifact written to the datastore.
1278 """
1279 # May need to coerce the in memory dataset to the correct
1280 # python type, but first we need to make sure the storage class
1281 # reflects the one defined in the data repository.
1282 ref = self._cast_storage_class(ref)
1284 # Confirm that we can accept this dataset
1285 if not self.constraints.isAcceptable(ref):
1286 # Raise rather than use boolean return value.
1287 raise DatasetTypeNotSupportedError(
1288 f"Dataset {ref} has been rejected by this datastore via configuration."
1289 )
1291 location, formatter = self._determine_put_formatter_location(ref)
1293 # The external storage class can differ from the registry storage
1294 # class AND the given in-memory dataset might not match any of the
1295 # storage class definitions.
1296 if formatter.can_accept(inMemoryDataset):
1297 # Do not need to coerce. Must assume that the formatter can handle
1298 # it without further checking of types.
1299 pass
1300 else:
1301 # Coerce to a type that it can accept.
1302 inMemoryDataset = ref.datasetType.storageClass.coerce_type(inMemoryDataset)
1303 required_pytype = ref.datasetType.storageClass.pytype
1305 if not isinstance(inMemoryDataset, required_pytype):
1306 raise TypeError(
1307 f"Inconsistency between supplied object ({type(inMemoryDataset)}) "
1308 f"and storage class type ({required_pytype})"
1309 )
1311 if self._transaction is None:
1312 raise RuntimeError("Attempting to write artifact without transaction enabled")
1314 def _removeFileExists(uri: ResourcePath) -> None:
1315 """Remove a file and do not complain if it is not there.
1317 This is important since a formatter might fail before the file
1318 is written and we should not confuse people by writing spurious
1319 error messages to the log.
1320 """
1321 with contextlib.suppress(FileNotFoundError):
1322 uri.remove()
1324 # Register a callback to try to delete the uploaded data if
1325 # something fails below
1326 uri = location.uri
1327 self._transaction.registerUndo("artifactWrite", _removeFileExists, uri)
1329 # Need to record the specified formatter but if this is a V1 formatter
1330 # we need to convert it to a V2 compatible shim to do the write.
1331 if not isinstance(formatter, Formatter):
1332 formatter_compat = formatter
1333 else:
1334 formatter_compat = FormatterV1inV2(
1335 formatter.file_descriptor,
1336 ref=ref,
1337 formatter=formatter,
1338 write_parameters=formatter.write_parameters,
1339 write_recipes=formatter.write_recipes,
1340 )
1342 assert isinstance(formatter_compat, FormatterV2)
1344 with time_this(log, msg="Writing dataset %s with formatter %s", args=(ref, formatter.name())):
1345 try:
1346 formatter_compat.write(
1347 inMemoryDataset, cache_manager=self.cacheManager, provenance=provenance
1348 )
1349 except Exception as e:
1350 raise RuntimeError(
1351 f"Failed to serialize dataset {ref} of type {get_full_type_name(inMemoryDataset)} "
1352 f"using formatter {formatter.name()}."
1353 ) from e
1355 # URI is needed to resolve what ingest case are we dealing with
1356 return self._extractIngestInfo(uri, ref, formatter=formatter)
1358 def knows(self, ref: DatasetRef) -> bool:
1359 """Check if the dataset is known to the datastore.
1361 Does not check for existence of any artifact.
1363 Parameters
1364 ----------
1365 ref : `DatasetRef`
1366 Reference to the required dataset.
1368 Returns
1369 -------
1370 exists : `bool`
1371 `True` if the dataset is known to the datastore.
1372 """
1373 fileLocations = self._get_dataset_locations_info(ref)
1374 if fileLocations:
1375 return True
1376 return False
1378 def knows_these(self, refs: Iterable[DatasetRef]) -> dict[DatasetRef, bool]:
1379 # Docstring inherited from the base class.
1380 refs = list(refs)
1382 # The records themselves. Could be missing some entries.
1383 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
1385 return {ref: ref.id in records for ref in refs}
1387 def _process_mexists_records(
1388 self,
1389 id_to_ref: dict[DatasetId, DatasetRef],
1390 records: dict[DatasetId, list[StoredFileInfo]],
1391 all_required: bool,
1392 artifact_existence: dict[ResourcePath, bool] | None = None,
1393 ) -> dict[DatasetRef, bool]:
1394 """Check given records for existence.
1396 Helper function for `mexists()`.
1398 Parameters
1399 ----------
1400 id_to_ref : `dict` of [`DatasetId`, `DatasetRef`]
1401 Mapping of the dataset ID to the dataset ref itself.
1402 records : `dict` of [`DatasetId`, `list` of `StoredFileInfo`]
1403 Records as generally returned by
1404 ``_get_stored_records_associated_with_refs``.
1405 all_required : `bool`
1406 Flag to indicate whether existence requires all artifacts
1407 associated with a dataset ID to exist or not for existence.
1408 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1409 Optional mapping of datastore artifact to existence. Updated by
1410 this method with details of all artifacts tested. Can be `None`
1411 if the caller is not interested.
1413 Returns
1414 -------
1415 existence : `dict` of [`DatasetRef`, `bool`]
1416 Mapping from dataset to boolean indicating existence.
1417 """
1418 # The URIs to be checked and a mapping of those URIs to
1419 # the dataset ID.
1420 uris_to_check: list[ResourcePath] = []
1421 location_map: dict[ResourcePath, DatasetId] = {}
1423 location_factory = self.locationFactory
1425 uri_existence: dict[ResourcePath, bool] = {}
1426 for ref_id, infos in records.items():
1427 # Key is the dataset Id, value is list of StoredItemInfo
1428 uris = [info.file_location(location_factory).uri for info in infos]
1429 location_map.update({uri: ref_id for uri in uris})
1431 # Check the local cache directly for a dataset corresponding
1432 # to the remote URI.
1433 if self.cacheManager.file_count > 0:
1434 ref = id_to_ref[ref_id]
1435 for uri, storedFileInfo in zip(uris, infos, strict=True):
1436 check_ref = ref
1437 if not ref.datasetType.isComponent() and (component := storedFileInfo.component):
1438 check_ref = ref.makeComponentRef(component)
1439 if self.cacheManager.known_to_cache(check_ref, uri.getExtension()):
1440 # Proxy for URI existence.
1441 uri_existence[uri] = True
1442 else:
1443 uris_to_check.append(uri)
1444 else:
1445 # Check all of them.
1446 uris_to_check.extend(uris)
1448 if artifact_existence is not None:
1449 # If a URI has already been checked remove it from the list
1450 # and immediately add the status to the output dict.
1451 filtered_uris_to_check = []
1452 for uri in uris_to_check:
1453 if uri in artifact_existence:
1454 uri_existence[uri] = artifact_existence[uri]
1455 else:
1456 filtered_uris_to_check.append(uri)
1457 uris_to_check = filtered_uris_to_check
1459 # Results.
1460 dataset_existence: dict[DatasetRef, bool] = {}
1462 uri_existence.update(ResourcePath.mexists(uris_to_check))
1463 for uri, exists in uri_existence.items():
1464 dataset_id = location_map[uri]
1465 ref = id_to_ref[dataset_id]
1467 # Disassembled composite needs to check all locations.
1468 # all_required indicates whether all need to exist or not.
1469 if ref in dataset_existence:
1470 if all_required:
1471 exists = dataset_existence[ref] and exists
1472 else:
1473 exists = dataset_existence[ref] or exists
1474 dataset_existence[ref] = exists
1476 if artifact_existence is not None:
1477 artifact_existence.update(uri_existence)
1479 return dataset_existence
1481 def mexists(
1482 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1483 ) -> dict[DatasetRef, bool]:
1484 """Check the existence of multiple datasets at once.
1486 Parameters
1487 ----------
1488 refs : iterable of `DatasetRef`
1489 The datasets to be checked.
1490 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1491 Optional mapping of datastore artifact to existence. Updated by
1492 this method with details of all artifacts tested. Can be `None`
1493 if the caller is not interested.
1495 Returns
1496 -------
1497 existence : `dict` of [`DatasetRef`, `bool`]
1498 Mapping from dataset to boolean indicating existence.
1500 Notes
1501 -----
1502 To minimize potentially costly remote existence checks, the local
1503 cache is checked as a proxy for existence. If a file for this
1504 `DatasetRef` does exist no check is done for the actual URI. This
1505 could result in possibly unexpected behavior if the dataset itself
1506 has been removed from the datastore by another process whilst it is
1507 still in the cache.
1508 """
1509 chunk_size = 50_000
1510 dataset_existence: dict[DatasetRef, bool] = {}
1511 log.debug("Checking for the existence of multiple artifacts in datastore in chunks of %d", chunk_size)
1512 n_found_total = 0
1513 n_checked = 0
1514 n_chunks = 0
1515 for chunk in chunk_iterable(refs, chunk_size=chunk_size):
1516 chunk_result = self._mexists(chunk, artifact_existence)
1518 # The log message level and content depend on how many
1519 # datasets we are processing.
1520 n_results = len(chunk_result)
1522 # Use verbose logging to ensure that messages can be seen
1523 # easily if many refs are being checked.
1524 log_threshold = VERBOSE
1525 n_checked += n_results
1527 # This sum can take some time so only do it if we know the
1528 # result is going to be used.
1529 n_found = 0
1530 if log.isEnabledFor(log_threshold):
1531 # Can treat the booleans as 0, 1 integers and sum them.
1532 n_found = sum(chunk_result.values())
1533 n_found_total += n_found
1535 # We are deliberately not trying to count the number of refs
1536 # provided in case it's in the millions. This means there is a
1537 # situation where the number of refs exactly matches the chunk
1538 # size and we will switch to the multi-chunk path even though
1539 # we only have a single chunk.
1540 if n_results < chunk_size and n_chunks == 0:
1541 # Single chunk will be processed so we can provide more detail.
1542 if n_results == 1:
1543 ref = list(chunk_result)[0]
1544 # Use debug logging to be consistent with `exists()`.
1545 log.debug(
1546 "Calling mexists() with single ref that does%s exist (%s).",
1547 "" if chunk_result[ref] else " not",
1548 ref,
1549 )
1550 else:
1551 # Single chunk but multiple files. Summarize.
1552 log.log(
1553 log_threshold,
1554 "Number of datasets found in datastore %s: %d out of %d datasets checked.",
1555 self.name,
1556 n_found,
1557 n_checked,
1558 )
1560 else:
1561 # Use incremental verbose logging when we have multiple chunks.
1562 log.log(
1563 log_threshold,
1564 "Number of datasets found in datastore for chunk %d: %d out of %d checked "
1565 "(running total from all chunks so far: %d found out of %d checked)",
1566 n_chunks,
1567 n_found,
1568 n_results,
1569 n_found_total,
1570 n_checked,
1571 )
1572 dataset_existence.update(chunk_result)
1573 n_chunks += 1
1575 return dataset_existence
1577 def _mexists(
1578 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1579 ) -> dict[DatasetRef, bool]:
1580 """Check the existence of multiple datasets at once.
1582 Parameters
1583 ----------
1584 refs : iterable of `DatasetRef`
1585 The datasets to be checked.
1586 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1587 Optional mapping of datastore artifact to existence. Updated by
1588 this method with details of all artifacts tested. Can be `None`
1589 if the caller is not interested.
1591 Returns
1592 -------
1593 existence : `dict` of [`DatasetRef`, `bool`]
1594 Mapping from dataset to boolean indicating existence.
1595 """
1596 # Make a mapping from refs with the internal storage class to the given
1597 # refs that may have a different one. We'll use the internal refs
1598 # throughout this method and convert back at the very end.
1599 internal_ref_to_input_ref = {self._cast_storage_class(ref): ref for ref in refs}
1601 # Need a mapping of dataset_id to (internal) dataset ref since some
1602 # internal APIs work with dataset_id.
1603 id_to_ref = {ref.id: ref for ref in internal_ref_to_input_ref}
1605 # Set of all IDs we are checking for.
1606 requested_ids = set(id_to_ref.keys())
1608 # The records themselves. Could be missing some entries.
1609 records = self._get_stored_records_associated_with_refs(
1610 id_to_ref.values(), ignore_datastore_records=True
1611 )
1613 dataset_existence = self._process_mexists_records(
1614 id_to_ref, records, True, artifact_existence=artifact_existence
1615 )
1617 # Set of IDs that have been handled.
1618 handled_ids = {ref.id for ref in dataset_existence}
1620 missing_ids = requested_ids - handled_ids
1621 if missing_ids:
1622 dataset_existence.update(
1623 self._mexists_check_expected(
1624 [id_to_ref[missing] for missing in missing_ids], artifact_existence
1625 )
1626 )
1628 return {
1629 internal_ref_to_input_ref[internal_ref]: existence
1630 for internal_ref, existence in dataset_existence.items()
1631 }
1633 def _mexists_check_expected(
1634 self, refs: Sequence[DatasetRef], artifact_existence: dict[ResourcePath, bool] | None = None
1635 ) -> dict[DatasetRef, bool]:
1636 """Check existence of refs that are not known to datastore.
1638 Parameters
1639 ----------
1640 refs : iterable of `DatasetRef`
1641 The datasets to be checked. These are assumed not to be known
1642 to datastore.
1643 artifact_existence : `dict` [`lsst.resources.ResourcePath`, `bool`]
1644 Optional mapping of datastore artifact to existence. Updated by
1645 this method with details of all artifacts tested. Can be `None`
1646 if the caller is not interested.
1648 Returns
1649 -------
1650 existence : `dict` of [`DatasetRef`, `bool`]
1651 Mapping from dataset to boolean indicating existence.
1652 """
1653 dataset_existence: dict[DatasetRef, bool] = {}
1654 if not self.trustGetRequest:
1655 # Must assume these do not exist
1656 for ref in refs:
1657 dataset_existence[ref] = False
1658 else:
1659 log.debug(
1660 "%d datasets were not known to datastore during initial existence check.",
1661 len(refs),
1662 )
1664 # Construct data structure identical to that returned
1665 # by _get_stored_records_associated_with_refs() but using
1666 # guessed names.
1667 records = {}
1668 id_to_ref = {}
1669 for missing_ref in refs:
1670 expected = self._get_expected_dataset_locations_info(missing_ref)
1671 dataset_id = missing_ref.id
1672 records[dataset_id] = [info for _, info in expected]
1673 id_to_ref[dataset_id] = missing_ref
1675 dataset_existence.update(
1676 self._process_mexists_records(
1677 id_to_ref,
1678 records,
1679 False,
1680 artifact_existence=artifact_existence,
1681 )
1682 )
1684 return dataset_existence
1686 def exists(self, ref: DatasetRef) -> bool:
1687 """Check if the dataset exists in the datastore.
1689 Parameters
1690 ----------
1691 ref : `DatasetRef`
1692 Reference to the required dataset.
1694 Returns
1695 -------
1696 exists : `bool`
1697 `True` if the entity exists in the `Datastore`.
1699 Notes
1700 -----
1701 The local cache is checked as a proxy for existence in the remote
1702 object store. It is possible that another process on a different
1703 compute node could remove the file from the object store even
1704 though it is present in the local cache.
1705 """
1706 ref = self._cast_storage_class(ref)
1707 # We cannot trust datastore records from ref, as many unit tests delete
1708 # datasets and check their existence.
1709 fileLocations = self._get_dataset_locations_info(ref, ignore_datastore_records=True)
1711 # if we are being asked to trust that registry might not be correct
1712 # we ask for the expected locations and check them explicitly
1713 if not fileLocations:
1714 if not self.trustGetRequest:
1715 return False
1717 # First check the cache. If it is not found we must check
1718 # the datastore itself. Assume that any component in the cache
1719 # means that the dataset does exist somewhere.
1720 if self.cacheManager.known_to_cache(ref):
1721 return True
1723 # When we are guessing a dataset location we can not check
1724 # for the existence of every component since we can not
1725 # know if every component was written. Instead we check
1726 # for the existence of any of the expected locations.
1727 for location, _ in self._get_expected_dataset_locations_info(ref):
1728 if self._artifact_exists(location):
1729 return True
1730 return False
1732 # All listed artifacts must exist.
1733 for location, storedFileInfo in fileLocations:
1734 # Checking in cache needs the component ref.
1735 check_ref = ref
1736 if not ref.datasetType.isComponent() and (component := storedFileInfo.component):
1737 check_ref = ref.makeComponentRef(component)
1738 if self.cacheManager.known_to_cache(check_ref, location.getExtension()):
1739 continue
1741 if not self._artifact_exists(location):
1742 return False
1744 return True
1746 def getURIs(self, ref: DatasetRef, predict: bool = False) -> DatasetRefURIs:
1747 """Return URIs associated with dataset.
1749 Parameters
1750 ----------
1751 ref : `DatasetRef`
1752 Reference to the required dataset.
1753 predict : `bool`, optional
1754 If the datastore does not know about the dataset, controls whether
1755 it should return a predicted URI or not.
1757 Returns
1758 -------
1759 uris : `DatasetRefURIs`
1760 The URI to the primary artifact associated with this dataset (if
1761 the dataset was disassembled within the datastore this may be
1762 `None`), and the URIs to any components associated with the dataset
1763 artifact. (can be empty if there are no components).
1764 """
1765 many = self.getManyURIs([ref], predict=predict, allow_missing=False)
1766 return many[ref]
1768 def getURI(self, ref: DatasetRef, predict: bool = False) -> ResourcePath:
1769 """URI to the Dataset.
1771 Parameters
1772 ----------
1773 ref : `DatasetRef`
1774 Reference to the required Dataset.
1775 predict : `bool`
1776 If `True`, allow URIs to be returned of datasets that have not
1777 been written.
1779 Returns
1780 -------
1781 uri : `str`
1782 URI pointing to the dataset within the datastore. If the
1783 dataset does not exist in the datastore, and if ``predict`` is
1784 `True`, the URI will be a prediction and will include a URI
1785 fragment "#predicted".
1786 If the datastore does not have entities that relate well
1787 to the concept of a URI the returned URI will be
1788 descriptive. The returned URI is not guaranteed to be obtainable.
1790 Raises
1791 ------
1792 FileNotFoundError
1793 Raised if a URI has been requested for a dataset that does not
1794 exist and guessing is not allowed.
1795 RuntimeError
1796 Raised if a request is made for a single URI but multiple URIs
1797 are associated with this dataset.
1799 Notes
1800 -----
1801 When a predicted URI is requested an attempt will be made to form
1802 a reasonable URI based on file templates and the expected formatter.
1803 """
1804 primary, components = self.getURIs(ref, predict)
1805 if primary is None or components:
1806 raise RuntimeError(
1807 f"Dataset ({ref}) includes distinct URIs for components. Use Datastore.getURIs() instead."
1808 )
1809 return primary
1811 def _predict_URIs(
1812 self,
1813 ref: DatasetRef,
1814 ) -> DatasetRefURIs:
1815 """Predict the URIs of a dataset ref.
1817 Parameters
1818 ----------
1819 ref : `DatasetRef`
1820 Reference to the required Dataset.
1822 Returns
1823 -------
1824 URI : DatasetRefUris
1825 Primary and component URIs. URIs will contain a URI fragment
1826 "#predicted".
1827 """
1828 uris = DatasetRefURIs()
1830 if self.composites.shouldBeDisassembled(ref):
1831 for component, _ in ref.datasetType.storageClass.components.items():
1832 comp_ref = ref.makeComponentRef(component)
1833 comp_location, _ = self._determine_put_formatter_location(comp_ref)
1835 # Add the "#predicted" URI fragment to indicate this is a
1836 # guess
1837 uris.componentURIs[component] = ResourcePath(
1838 comp_location.uri.geturl() + "#predicted", forceDirectory=comp_location.uri.dirLike
1839 )
1841 else:
1842 location, _ = self._determine_put_formatter_location(ref)
1844 # Add the "#predicted" URI fragment to indicate this is a guess
1845 uris.primaryURI = ResourcePath(
1846 location.uri.geturl() + "#predicted", forceDirectory=location.uri.dirLike
1847 )
1849 return uris
1851 def getManyURIs(
1852 self,
1853 refs: Iterable[DatasetRef],
1854 predict: bool = False,
1855 allow_missing: bool = False,
1856 ) -> dict[DatasetRef, DatasetRefURIs]:
1857 # Docstring inherited
1859 uris: dict[DatasetRef, DatasetRefURIs] = {}
1861 records = self._get_stored_records_associated_with_refs(refs)
1862 records_keys = records.keys()
1864 existing_refs = tuple(ref for ref in refs if ref.id in records_keys)
1865 missing_refs = tuple(ref for ref in refs if ref.id not in records_keys)
1867 # Have to handle trustGetRequest mode by checking for the existence
1868 # of the missing refs on disk.
1869 if missing_refs and not predict:
1870 dataset_existence = self._mexists_check_expected(missing_refs, None)
1871 really_missing = set()
1872 not_missing = set()
1873 for ref, exists in dataset_existence.items():
1874 if exists:
1875 not_missing.add(ref)
1876 else:
1877 really_missing.add(ref)
1879 if not_missing:
1880 # Need to recalculate the missing/existing split.
1881 existing_refs = existing_refs + tuple(not_missing)
1882 missing_refs = tuple(really_missing)
1884 for ref in missing_refs:
1885 # if this has never been written then we have to guess
1886 if not predict:
1887 if not allow_missing:
1888 raise FileNotFoundError(f"Dataset {ref} not in this datastore.")
1889 else:
1890 uris[ref] = self._predict_URIs(ref)
1892 for ref in existing_refs:
1893 file_infos = records[ref.id]
1894 file_locations = [(i.file_location(self.locationFactory), i) for i in file_infos]
1895 uris[ref] = self._locations_to_URI(ref, file_locations)
1897 return uris
1899 def _locations_to_URI(
1900 self,
1901 ref: DatasetRef,
1902 file_locations: Sequence[tuple[Location, StoredFileInfo]],
1903 ) -> DatasetRefURIs:
1904 """Convert one or more file locations associated with a DatasetRef
1905 to a DatasetRefURIs.
1907 Parameters
1908 ----------
1909 ref : `DatasetRef`
1910 Reference to the dataset.
1911 file_locations : Sequence[Tuple[Location, StoredFileInfo]]
1912 Each item in the sequence is the location of the dataset within the
1913 datastore and stored information about the file and its formatter.
1914 If there is only one item in the sequence then it is treated as the
1915 primary URI. If there is more than one item then they are treated
1916 as component URIs. If there are no items then an error is raised
1917 unless ``self.trustGetRequest`` is `True`.
1919 Returns
1920 -------
1921 uris: DatasetRefURIs
1922 Represents the primary URI or component URIs described by the
1923 inputs.
1925 Raises
1926 ------
1927 RuntimeError
1928 If no file locations are passed in and ``self.trustGetRequest`` is
1929 `False`.
1930 FileNotFoundError
1931 If the a passed-in URI does not exist, and ``self.trustGetRequest``
1932 is `False`.
1933 RuntimeError
1934 If a passed in `StoredFileInfo`'s ``component`` is `None` (this is
1935 unexpected).
1936 """
1937 guessing = False
1938 uris = DatasetRefURIs()
1940 if not file_locations:
1941 if not self.trustGetRequest:
1942 raise RuntimeError(f"Unexpectedly got no artifacts for dataset {ref}")
1943 file_locations = self._get_expected_dataset_locations_info(ref)
1944 guessing = True
1946 if len(file_locations) == 1:
1947 # No disassembly so this is the primary URI
1948 uris.primaryURI = file_locations[0][0].uri
1949 if guessing and not uris.primaryURI.exists():
1950 raise FileNotFoundError(f"Expected URI ({uris.primaryURI}) does not exist")
1951 else:
1952 for location, file_info in file_locations:
1953 if file_info.component is None:
1954 raise RuntimeError(f"Unexpectedly got no component name for a component at {location}")
1955 if guessing and not location.uri.exists():
1956 # If we are trusting then it is entirely possible for
1957 # some components to be missing. In that case we skip
1958 # to the next component.
1959 if self.trustGetRequest:
1960 continue
1961 raise FileNotFoundError(f"Expected URI ({location.uri}) does not exist")
1962 uris.componentURIs[file_info.component] = location.uri
1964 return uris
1966 def _find_missing_records(
1967 self,
1968 refs: Iterable[DatasetRef],
1969 missing_ids: set[DatasetId],
1970 artifact_existence: dict[ResourcePath, bool] | None = None,
1971 warn_for_missing: bool = True,
1972 ) -> dict[DatasetId, list[StoredFileInfo]]:
1973 if not missing_ids:
1974 return {}
1976 if artifact_existence is None:
1977 artifact_existence = {}
1979 found_records: dict[DatasetId, list[StoredFileInfo]] = defaultdict(list)
1980 id_to_ref = {ref.id: ref for ref in refs if ref.id in missing_ids}
1982 # This should be chunked in case we end up having to check
1983 # the file store since we need some log output to show
1984 # progress.
1985 chunk_size = 50_000
1986 for missing_ids_chunk in chunk_iterable(missing_ids, chunk_size=chunk_size):
1987 records = {}
1988 for missing in missing_ids_chunk:
1989 # Ask the source datastore where the missing artifacts
1990 # should be. An execution butler might not know about the
1991 # artifacts even if they are there.
1992 expected = self._get_expected_dataset_locations_info(id_to_ref[missing])
1993 records[missing] = [info for _, info in expected]
1995 # Call the mexist helper method in case we have not already
1996 # checked these artifacts such that artifact_existence is
1997 # empty. This allows us to benefit from parallelism.
1998 # datastore.mexists() itself does not give us access to the
1999 # derived datastore record.
2000 log.verbose("Checking existence of %d datasets unknown to datastore", len(records))
2001 ref_exists = self._process_mexists_records(
2002 id_to_ref, records, False, artifact_existence=artifact_existence
2003 )
2005 # Now go through the records and propagate the ones that exist.
2006 location_factory = self.locationFactory
2007 for missing, record_list in records.items():
2008 # Skip completely if the ref does not exist.
2009 ref = id_to_ref[missing]
2010 if not ref_exists[ref]:
2011 if warn_for_missing:
2012 log.warning("Asked to transfer dataset %s but no file artifacts exist for it.", ref)
2013 continue
2014 # Check for file artifact to decide which parts of a
2015 # disassembled composite do exist. If there is only a
2016 # single record we don't even need to look because it can't
2017 # be a composite and must exist.
2018 if len(record_list) == 1:
2019 dataset_records = record_list
2020 else:
2021 dataset_records = [
2022 record
2023 for record in record_list
2024 if artifact_existence[record.file_location(location_factory).uri]
2025 ]
2026 assert len(dataset_records) > 0, "Disassembled composite should have had some files."
2028 # Rely on source_records being a defaultdict.
2029 found_records[missing].extend(dataset_records)
2030 log.verbose("Completed scan for missing data files")
2031 return found_records
2033 def retrieveArtifacts(
2034 self,
2035 refs: Iterable[DatasetRef],
2036 destination: ResourcePath,
2037 transfer: str = "auto",
2038 preserve_path: bool = True,
2039 overwrite: bool = False,
2040 write_index: bool = True,
2041 add_prefix: bool = False,
2042 ) -> dict[ResourcePath, ArtifactIndexInfo]:
2043 """Retrieve the file artifacts associated with the supplied refs.
2045 Parameters
2046 ----------
2047 refs : iterable of `DatasetRef`
2048 The datasets for which file artifacts are to be retrieved.
2049 A single ref can result in multiple files. The refs must
2050 be resolved.
2051 destination : `lsst.resources.ResourcePath`
2052 Location to write the file artifacts.
2053 transfer : `str`, optional
2054 Method to use to transfer the artifacts. Must be one of the options
2055 supported by `lsst.resources.ResourcePath.transfer_from()`.
2056 "move" is not allowed.
2057 preserve_path : `bool`, optional
2058 If `True` the full path of the file artifact within the datastore
2059 is preserved. If `False` the final file component of the path
2060 is used.
2061 overwrite : `bool`, optional
2062 If `True` allow transfers to overwrite existing files at the
2063 destination.
2064 write_index : `bool`, optional
2065 If `True` write a file at the top level containing a serialization
2066 of a `ZipIndex` for the downloaded datasets.
2067 add_prefix : `bool`, optional
2068 If `True` and if ``preserve_path`` is `False`, apply a prefix to
2069 the filenames corresponding to some part of the dataset ref ID.
2070 This can be used to guarantee uniqueness.
2072 Returns
2073 -------
2074 artifact_map : `dict` [ `lsst.resources.ResourcePath`, \
2075 `ArtifactIndexInfo` ]
2076 Mapping of retrieved file to associated index information.
2077 """
2078 if not destination.isdir():
2079 raise ValueError(f"Destination location must refer to a directory. Given {destination}")
2081 if transfer == "move":
2082 raise ValueError("Can not move artifacts out of datastore. Use copy instead.")
2084 # Source -> Destination
2085 # This also helps filter out duplicate DatasetRef in the request
2086 # that will map to the same underlying file transfer.
2087 to_transfer: dict[ResourcePath, ResourcePath] = {}
2088 zips_to_transfer: set[ResourcePath] = set()
2090 # Retrieve all the records in bulk indexed by ref.id.
2091 records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2093 # Check for missing records.
2094 known_ids = set(records)
2095 log.debug("Number of datastore records found in database: %d", len(known_ids))
2096 requested_ids = {ref.id for ref in refs}
2097 missing_ids = requested_ids - known_ids
2099 if missing_ids and not self.trustGetRequest:
2100 raise ValueError(f"Number of datasets missing from this datastore: {len(missing_ids)}")
2102 missing_records = self._find_missing_records(refs, missing_ids)
2103 records.update(missing_records)
2105 # One artifact can be used by multiple DatasetRef.
2106 # e.g. DECam.
2107 artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {}
2108 # Sort to ensure that in many refs to one file situation the same
2109 # ref is used for any prefix that might be added.
2110 for ref in sorted(refs):
2111 prefix = str(ref.id)[:8] + "-" if add_prefix else ""
2112 for info in records[ref.id]:
2113 location = info.file_location(self.locationFactory)
2114 source_uri = location.uri
2115 # For DECam/zip we only want to copy once.
2116 # For zip files we need to unpack so that they can be
2117 # zipped up again if needed.
2118 is_zip = source_uri.getExtension() == ".zip" and "zip-path" in source_uri.fragment
2119 # We need to remove fragments for consistency.
2120 cleaned_source_uri = source_uri.replace(fragment="", query="", params="")
2121 if is_zip:
2122 # Assume the DatasetRef definitions are within the Zip
2123 # file itself and so can be dropped from loop.
2124 zips_to_transfer.add(cleaned_source_uri)
2125 elif cleaned_source_uri not in to_transfer:
2126 target_uri = determine_destination_for_retrieved_artifact(
2127 destination, location.pathInStore, preserve_path, prefix
2128 )
2129 to_transfer[cleaned_source_uri] = target_uri
2130 artifact_map[target_uri] = ArtifactIndexInfo.from_single(info.to_simple(), ref.id)
2131 else:
2132 target_uri = to_transfer[cleaned_source_uri]
2133 artifact_map[target_uri].append(ref.id)
2135 # Parallelize the transfer. Re-raise as a single exception if
2136 # a FileExistsError is encountered anywhere.
2137 log.debug("Number of artifacts to transfer to %s: %d", str(destination), len(to_transfer))
2138 try:
2139 ResourcePath.mtransfer(transfer, tuple(to_transfer.items()), overwrite=overwrite)
2140 except* FileExistsError as egroup:
2141 raise FileExistsError(
2142 "Some files already exist in destination directory and overwrite is False"
2143 ) from egroup
2145 # Transfer the Zip files and unpack them.
2146 zipped_artifacts = unpack_zips(zips_to_transfer, requested_ids, destination, preserve_path, overwrite)
2147 artifact_map.update(zipped_artifacts)
2149 if write_index:
2150 index = ZipIndex.from_artifact_map(refs, artifact_map, destination)
2151 index.write_index(destination)
2153 return artifact_map
2155 def ingest_zip(self, zip_path: ResourcePath, transfer: str | None, *, dry_run: bool = False) -> None:
2156 """Ingest an indexed Zip file and contents.
2158 The Zip file must have an index file as created by `retrieveArtifacts`.
2160 Parameters
2161 ----------
2162 zip_path : `lsst.resources.ResourcePath`
2163 Path to the Zip file.
2164 transfer : `str`
2165 Method to use for transferring the Zip file into the datastore.
2166 dry_run : `bool`, optional
2167 If `True` the ingest will be processed without any modifications
2168 made to the target datastore and as if the target datastore did not
2169 have any of the datasets.
2171 Notes
2172 -----
2173 Datastore constraints are bypassed with Zip ingest. A zip file can
2174 contain multiple dataset types. Should the entire Zip be rejected
2175 if one dataset type is in the constraints list?
2177 If any dataset is already present in the datastore the entire ingest
2178 will fail.
2179 """
2180 index = ZipIndex.from_zip_file(zip_path)
2182 # Refs indexed by UUID.
2183 refs = index.refs.to_refs(universe=self.universe)
2184 id_to_ref = {ref.id: ref for ref in refs}
2186 # Any failing constraints trigger entire failure.
2187 if any(not self.constraints.isAcceptable(ref) for ref in refs):
2188 raise DatasetTypeNotSupportedError(
2189 "Some refs in the Zip file are not supported by this datastore"
2190 )
2192 # Transfer the Zip file into the datastore file system.
2193 # There is no RUN as such to use for naming.
2194 # Potentially could use the RUN from the first ref in the index
2195 # There is no requirement that the contents of the Zip files share
2196 # the same RUN.
2197 # Could use the Zip UUID from the index + special "zips/" prefix.
2198 if transfer is None:
2199 # Indicated that the zip file is already in the right place.
2200 if not zip_path.isabs():
2201 tgtLocation = self.locationFactory.fromPath(zip_path.ospath, trusted_path=False)
2202 else:
2203 pathInStore = zip_path.relative_to(self.root)
2204 if pathInStore is None:
2205 raise RuntimeError(
2206 f"Unexpectedly learned that {zip_path} is not within datastore {self.root}"
2207 )
2208 tgtLocation = self.locationFactory.fromPath(pathInStore, trusted_path=True)
2209 elif transfer == "direct":
2210 # Reference in original location.
2211 tgtLocation = None
2212 else:
2213 # Name the zip file based on index contents.
2214 tgtLocation = self.locationFactory.fromPath(index.calculate_zip_file_path_in_store())
2216 # Transfer the Zip file into the datastore.
2217 if not dry_run:
2218 tgtLocation.uri.transfer_from(
2219 zip_path, transfer=transfer, transaction=self._transaction, overwrite=True
2220 )
2221 else:
2222 log.info("Would be copying Zip from %s to %s", zip_path, tgtLocation)
2224 if tgtLocation is None:
2225 path_in_store = str(zip_path)
2226 else:
2227 path_in_store = tgtLocation.pathInStore.path
2229 # Associate each file with a (DatasetRef, StoredFileInfo) tuple.
2230 artifacts: list[tuple[DatasetRef, StoredFileInfo]] = []
2231 for path_in_zip, index_info in index.artifact_map.items():
2232 # Need to modify the info to include the path to the Zip file
2233 # that was previously written to the datastore.
2234 index_info.info.path = f"{path_in_store}#zip-path={path_in_zip}"
2236 info = StoredFileInfo.from_simple(index_info.info)
2237 for id_ in index_info.ids:
2238 artifacts.append((id_to_ref[id_], info))
2240 if not dry_run:
2241 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT)
2242 else:
2243 log.info("Would be registering %d artifacts from Zip into datastore", len(artifacts))
2245 def get(
2246 self,
2247 ref: DatasetRef,
2248 parameters: Mapping[str, Any] | None = None,
2249 storageClass: StorageClass | str | None = None,
2250 ) -> Any:
2251 """Load an InMemoryDataset from the store.
2253 Parameters
2254 ----------
2255 ref : `DatasetRef`
2256 Reference to the required Dataset.
2257 parameters : `dict`
2258 `StorageClass`-specific parameters that specify, for example,
2259 a slice of the dataset to be loaded.
2260 storageClass : `StorageClass` or `str`, optional
2261 The storage class to be used to override the Python type
2262 returned by this method. By default the returned type matches
2263 the dataset type definition for this dataset. Specifying a
2264 read `StorageClass` can force a different type to be returned.
2265 This type must be compatible with the original type.
2267 Returns
2268 -------
2269 inMemoryDataset : `object`
2270 Requested dataset or slice thereof as an InMemoryDataset.
2272 Raises
2273 ------
2274 FileNotFoundError
2275 Requested dataset can not be retrieved.
2276 TypeError
2277 Return value from formatter has unexpected type.
2278 ValueError
2279 Formatter failed to process the dataset.
2280 """
2281 # Supplied storage class for the component being read is either
2282 # from the ref itself or some an override if we want to force
2283 # type conversion.
2284 if storageClass is not None:
2285 ref = ref.overrideStorageClass(storageClass)
2287 allGetInfo = self._prepare_for_direct_get(ref, parameters)
2288 return get_dataset_as_python_object_from_get_info(
2289 allGetInfo, ref=ref, parameters=parameters, cache_manager=self.cacheManager
2290 )
2292 def prepare_get_for_external_client(self, ref: DatasetRef) -> list[DatasetLocationInformation] | None:
2293 # Docstring inherited
2295 locations = self._get_dataset_locations_info(ref)
2296 if len(locations) == 0:
2297 return None
2299 return locations
2301 @transactional
2302 def put(self, inMemoryDataset: Any, ref: DatasetRef, provenance: DatasetProvenance | None = None) -> None:
2303 """Write a InMemoryDataset with a given `DatasetRef` to the store.
2305 Parameters
2306 ----------
2307 inMemoryDataset : `object`
2308 The dataset to store.
2309 ref : `DatasetRef`
2310 Reference to the associated Dataset.
2311 provenance : `DatasetProvenance` or `None`, optional
2312 Any provenance that should be attached to the serialized dataset.
2313 Can be ignored by a formatter or delegate.
2315 Raises
2316 ------
2317 TypeError
2318 Supplied object and storage class are inconsistent.
2319 DatasetTypeNotSupportedError
2320 The associated `DatasetType` is not handled by this datastore.
2322 Notes
2323 -----
2324 If the datastore is configured to reject certain dataset types it
2325 is possible that the put will fail and raise a
2326 `DatasetTypeNotSupportedError`. The main use case for this is to
2327 allow `ChainedDatastore` to put to multiple datastores without
2328 requiring that every datastore accepts the dataset.
2329 """
2330 doDisassembly = self.composites.shouldBeDisassembled(ref)
2331 # doDisassembly = True
2333 artifacts = []
2334 if doDisassembly:
2335 inMemoryDataset = ref.datasetType.storageClass.delegate().add_provenance(
2336 inMemoryDataset, ref, provenance=provenance
2337 )
2338 components = ref.datasetType.storageClass.delegate().disassemble(inMemoryDataset)
2339 if components is None:
2340 raise RuntimeError(
2341 f"Inconsistent configuration: dataset type {ref.datasetType.name} "
2342 f"with storage class {ref.datasetType.storageClass.name} "
2343 "is configured to be disassembled, but cannot be."
2344 )
2345 for component, componentInfo in components.items():
2346 # Don't recurse because we want to take advantage of
2347 # bulk insert -- need a new DatasetRef that refers to the
2348 # same dataset_id but has the component DatasetType
2349 # DatasetType does not refer to the types of components
2350 # So we construct one ourselves.
2351 compRef = ref.makeComponentRef(component)
2352 # Provenance has already been attached above.
2353 storedInfo = self._write_in_memory_to_artifact(componentInfo.component, compRef)
2354 artifacts.append((compRef, storedInfo))
2355 else:
2356 # Write the entire thing out
2357 storedInfo = self._write_in_memory_to_artifact(inMemoryDataset, ref, provenance=provenance)
2358 artifacts.append((ref, storedInfo))
2360 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.INSERT)
2362 @transactional
2363 def put_new(self, in_memory_dataset: Any, ref: DatasetRef) -> Mapping[str, DatasetRef]:
2364 doDisassembly = self.composites.shouldBeDisassembled(ref)
2365 # doDisassembly = True
2367 artifacts = []
2368 if doDisassembly:
2369 components = ref.datasetType.storageClass.delegate().disassemble(in_memory_dataset)
2370 if components is None:
2371 raise RuntimeError(
2372 f"Inconsistent configuration: dataset type {ref.datasetType.name} "
2373 f"with storage class {ref.datasetType.storageClass.name} "
2374 "is configured to be disassembled, but cannot be."
2375 )
2376 for component, componentInfo in components.items():
2377 # Don't recurse because we want to take advantage of
2378 # bulk insert -- need a new DatasetRef that refers to the
2379 # same dataset_id but has the component DatasetType
2380 # DatasetType does not refer to the types of components
2381 # So we construct one ourselves.
2382 compRef = ref.makeComponentRef(component)
2383 storedInfo = self._write_in_memory_to_artifact(componentInfo.component, compRef)
2384 artifacts.append((compRef, storedInfo))
2385 else:
2386 # Write the entire thing out
2387 storedInfo = self._write_in_memory_to_artifact(in_memory_dataset, ref)
2388 artifacts.append((ref, storedInfo))
2390 ref_records: DatasetDatastoreRecords = {self._opaque_table_name: [info for _, info in artifacts]}
2391 ref = ref.replace(datastore_records=ref_records)
2392 return {self.name: ref}
2394 @transactional
2395 def trash(self, ref: DatasetRef | Iterable[DatasetRef], ignore_errors: bool = True) -> None:
2396 # At this point can safely remove these datasets from the cache
2397 # to avoid confusion later on. If they are not trashed later
2398 # the cache will simply be refilled.
2399 self.cacheManager.remove_from_cache(ref)
2401 # If we are in trust mode there will be nothing to move to
2402 # the trash table and we will have to try to delete the file
2403 # immediately.
2404 if self.trustGetRequest:
2405 # Try to keep the logic below for a single file trash.
2406 if isinstance(ref, DatasetRef):
2407 refs = {ref}
2408 else:
2409 # Will recreate ref at the end of this branch.
2410 refs = set(ref)
2412 # Determine which datasets are known to datastore directly.
2413 id_to_ref = {ref.id: ref for ref in refs}
2414 existing_ids = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2415 existing_refs = {id_to_ref[ref_id] for ref_id in existing_ids}
2417 missing = refs - existing_refs
2418 if missing:
2419 # Do an explicit existence check on these refs.
2420 # We only care about the artifacts at this point and not
2421 # the dataset existence.
2422 artifact_existence: dict[ResourcePath, bool] = {}
2423 _ = self.mexists(missing, artifact_existence)
2424 uris = [uri for uri, exists in artifact_existence.items() if exists]
2426 # FUTURE UPGRADE: Implement a parallelized bulk remove.
2427 log.debug("Removing %d artifacts from datastore that are unknown to datastore", len(uris))
2428 for uri in uris:
2429 try:
2430 uri.remove()
2431 except Exception as e:
2432 if ignore_errors:
2433 log.debug("Artifact %s could not be removed: %s", uri, e)
2434 continue
2435 raise
2437 # There is no point asking the code below to remove refs we
2438 # know are missing so update it with the list of existing
2439 # records. Try to retain one vs many logic.
2440 if not existing_refs:
2441 # Nothing more to do since none of the datasets were
2442 # known to the datastore record table.
2443 return
2444 ref = list(existing_refs)
2445 if len(ref) == 1:
2446 ref = ref[0]
2448 # Get file metadata and internal metadata
2449 if not isinstance(ref, DatasetRef):
2450 log.debug("Doing multi-dataset trash in datastore %s", self.name)
2451 # Assumed to be an iterable of refs so bulk mode enabled.
2452 try:
2453 self.bridge.moveToTrash(ref, transaction=self._transaction)
2454 except Exception as e:
2455 if ignore_errors:
2456 log.warning("Unexpected issue moving multiple datasets to trash: %s", e)
2457 else:
2458 raise
2459 return
2461 log.debug("Trashing dataset %s in datastore %s", ref, self.name)
2463 fileLocations = self._get_dataset_locations_info(ref)
2465 if not fileLocations:
2466 err_msg = f"Requested dataset to trash ({ref}) is not known to datastore {self.name}"
2467 if ignore_errors:
2468 log.warning(err_msg)
2469 return
2470 else:
2471 raise FileNotFoundError(err_msg)
2473 for location, _ in fileLocations:
2474 if not self._artifact_exists(location):
2475 err_msg = (
2476 f"Dataset is known to datastore {self.name} but "
2477 f"associated artifact ({location.uri}) is missing"
2478 )
2479 if ignore_errors:
2480 log.warning(err_msg)
2481 return
2482 else:
2483 raise FileNotFoundError(err_msg)
2485 # Mark dataset as trashed
2486 try:
2487 self.bridge.moveToTrash([ref], transaction=self._transaction)
2488 except Exception as e:
2489 if ignore_errors:
2490 log.warning(
2491 "Attempted to mark dataset (%s) to be trashed in datastore %s "
2492 "but encountered an error: %s",
2493 ref,
2494 self.name,
2495 e,
2496 )
2497 pass
2498 else:
2499 raise
2501 def emptyTrash(
2502 self, ignore_errors: bool = True, refs: Collection[DatasetRef] | None = None, dry_run: bool = False
2503 ) -> set[ResourcePath]:
2504 """Remove all datasets from the trash.
2506 Parameters
2507 ----------
2508 ignore_errors : `bool`
2509 If `True` return without error even if something went wrong.
2510 Problems could occur if another process is simultaneously trying
2511 to delete.
2512 refs : `collections.abc.Collection` [ `DatasetRef` ] or `None`
2513 Explicit list of datasets that can be removed from trash. If listed
2514 datasets are not already stored in the trash table they will be
2515 ignored. If `None` every entry in the trash table will be
2516 processed.
2517 dry_run : `bool`, optional
2518 If `True`, the trash table will be queried and results reported
2519 but no artifacts will be removed.
2521 Returns
2522 -------
2523 removed : `set` [ `lsst.resources.ResourcePath` ]
2524 List of artifacts that were removed.
2526 Notes
2527 -----
2528 Will empty the records from the trash tables only if this call finishes
2529 without raising.
2530 """
2531 removed = set()
2532 if refs:
2533 selected_ids = {ref.id for ref in refs}
2534 chunk_size = 50_000
2535 n_chunks = math.ceil(len(selected_ids) / chunk_size)
2536 chunk_num = 0
2537 for chunk in chunk_iterable(selected_ids, chunk_size=chunk_size):
2538 chunk_num += 1
2539 if n_chunks == 1:
2540 log.verbose(
2541 "Emptying datastore trash for %d dataset%s",
2542 len(chunk),
2543 "s" if len(chunk) != 1 else "",
2544 )
2545 else:
2546 log.verbose(
2547 "Emptying datastore trash for chunk %d out of %d of size %d",
2548 chunk_num,
2549 n_chunks,
2550 len(chunk),
2551 )
2552 removed.update(
2553 self._empty_trash_subset(ignore_errors=ignore_errors, selected_ids=chunk, dry_run=dry_run)
2554 )
2555 else:
2556 log.verbose("Emptying all trash in datastore %s", self.name)
2557 removed = self._empty_trash_subset(ignore_errors=ignore_errors, dry_run=dry_run)
2558 log.info(
2559 "%sRemoved %d file artifact%s from datastore %s",
2560 "Would have " if dry_run else "",
2561 len(removed),
2562 "s" if len(removed) != 1 else "",
2563 self.name,
2564 )
2565 return removed
2567 @transactional
2568 def _empty_trash_subset(
2569 self,
2570 *,
2571 ignore_errors: bool = True,
2572 selected_ids: Collection[DatasetId] | None = None,
2573 dry_run: bool = False,
2574 ) -> set[ResourcePath]:
2575 """Empty trash table in transaction.
2577 Parameters
2578 ----------
2579 ignore_errors : `bool`
2580 If `True` return without error even if something went wrong.
2581 Problems could occur if another process is simultaneously trying
2582 to delete.
2583 selected_ids : `collections.abc.collection` [`DatasetId`] or `None`
2584 Explicit list of dataset IDs that can be removed from the trash.
2585 If listed datasets are not already included in the trash table
2586 they will be ignored. If `None` every entry in the trash table
2587 will be processed.
2588 dry_run : `bool`, optional
2589 If `True`, the trash table will be queried and results reported
2590 but no artifacts will be removed.
2592 Returns
2593 -------
2594 removed : `set` [ `lsst.resources.ResourcePath` ]
2595 Artifacts successfully removed.
2597 Notes
2598 -----
2599 Will empty the records from the trash tables only if this call finishes
2600 without raising.
2601 """
2602 # Context manager will empty trash iff we finish it without raising.
2603 # It will also automatically delete the relevant rows from the
2604 # trash table and the records table.
2605 with self.bridge.emptyTrash(
2606 self._table,
2607 record_class=StoredFileInfo,
2608 record_column="path",
2609 selected_ids=selected_ids,
2610 dry_run=dry_run,
2611 ) as trash_data:
2612 # Removing the artifacts themselves requires that the files are
2613 # not also associated with refs that are not to be trashed.
2614 # Therefore need to do a query with the file paths themselves
2615 # and return all the refs associated with them. Can only delete
2616 # a file if the refs to be trashed are the only refs associated
2617 # with the file.
2618 # This requires multiple copies of the trashed items
2619 trashed, artifacts_to_keep = trash_data
2621 # Assume that # in path means there are fragments involved. The
2622 # fragments can not be handled by the emptyTrash bridge call
2623 # so need to be processed independently.
2624 # The generator has to be converted to a list for multiple
2625 # iterations. Clean up the typing so that multiple isinstance
2626 # tests aren't needed later.
2627 trashed_list = [(ref, ninfo) for ref, ninfo in trashed if isinstance(ninfo, StoredFileInfo)]
2629 if artifacts_to_keep is None or any("#" in info[1].path for info in trashed_list):
2630 # The bridge is not helping us so have to work it out
2631 # ourselves. This is not going to be as efficient.
2632 # This mapping does not include the fragments.
2633 if artifacts_to_keep is not None:
2634 # This means we have already checked for non-fragment
2635 # examples so can filter.
2636 paths_to_check = {info.path for _, info in trashed_list if "#" in info.path}
2637 else:
2638 paths_to_check = {info.path for _, info in trashed_list}
2640 path_map = self._refs_associated_with_artifacts(paths_to_check)
2642 for ref, info in trashed_list:
2643 path = info.artifact_path
2644 # For disassembled composites in a Zip it is possible
2645 # for the same path to correspond to the same dataset ref
2646 # multiple times so trap for that.
2647 if ref.id in path_map[path]:
2648 path_map[path].remove(ref.id)
2649 if not path_map[path]:
2650 del path_map[path]
2652 slow_artifacts_to_keep = set(path_map)
2653 if artifacts_to_keep is not None:
2654 artifacts_to_keep.update(slow_artifacts_to_keep)
2655 else:
2656 artifacts_to_keep = slow_artifacts_to_keep
2658 n_direct = 0
2659 artifacts_to_delete: set[ResourcePath] = set()
2660 for ref, info in trashed_list:
2661 # Should not happen for this implementation but need
2662 # to keep mypy happy.
2663 assert info is not None, f"Internal logic error in emptyTrash with ref {ref}."
2665 if info.artifact_path in artifacts_to_keep:
2666 # This is a multi-dataset artifact and we are not
2667 # removing all associated refs.
2668 continue
2670 # Only trashed refs still known to datastore will be returned.
2671 location = info.file_location(self.locationFactory)
2673 if location.pathInStore.isabs():
2674 n_direct += 1
2675 continue
2677 # Strip fragment before storing since it is the artifact
2678 # we are deleting and we do not want repeats for every member
2679 # in a zip.
2680 artifacts_to_delete.add(location.uri.replace(fragment=""))
2682 if n_direct > 0:
2683 s = "s" if n_direct != 1 else ""
2684 log.verbose("Not deleting %d artifact%s using absolute URI%s", n_direct, s, s)
2686 if artifacts_to_keep:
2687 log.verbose(
2688 "%d artifact%s %s not deleted because of association with other datasets",
2689 len(artifacts_to_keep),
2690 "s" if len(artifacts_to_keep) != 1 else "",
2691 "were" if len(artifacts_to_keep) != 1 else "was",
2692 )
2694 if not artifacts_to_delete:
2695 return set()
2697 # Now do the deleting. Special case the log message for a single
2698 # artifact.
2699 if len(artifacts_to_delete) == 1:
2700 log.verbose(
2701 "%s removing file artifact %s from datastore %s",
2702 "Would be" if dry_run else "Now",
2703 list(artifacts_to_delete)[0],
2704 self.name,
2705 )
2706 else:
2707 log.verbose(
2708 "%s removing %d file artifacts from datastore %s",
2709 "Would be" if dry_run else "Now",
2710 len(artifacts_to_delete),
2711 self.name,
2712 )
2714 # For dry-run mode do not attempt to search the file store for
2715 # the artifacts to determine whether they exist or not. Simply
2716 # report that an attempt would be made to delete them. Never
2717 # report direct imports.
2718 if dry_run:
2719 return artifacts_to_delete
2721 # Now remove the actual file artifacts.
2722 remove_result = ResourcePath.mremove(artifacts_to_delete, do_raise=False)
2724 removed: set[ResourcePath] = set()
2725 exceptions: list[Exception] = []
2726 for uri, result in remove_result.items():
2727 if result.exception is None or isinstance(result.exception, FileNotFoundError):
2728 # File not existing is not an error since some other
2729 # process might have been trying to clean it and we do not
2730 # want to raise an error for a situation where the file
2731 # is not there and we do not want it to be there.
2732 removed.add(uri)
2733 else:
2734 exceptions.append(result.exception)
2736 if exceptions:
2737 s_err = "s" if len(exceptions) != 1 else ""
2738 e = ExceptionGroup(f"Error{s_err} removing {len(exceptions)} artifact{s_err}", exceptions)
2739 if ignore_errors:
2740 # Use a debug message here even though it's not
2741 # a good situation. In some cases this can be
2742 # caused by a race between user A and user B
2743 # and neither of them has permissions for the
2744 # other's files. Butler does not know about users
2745 # and trash has no idea what collections these
2746 # files were in (without guessing from a path).
2747 log.debug(
2748 "Encountered %d error%s removing %d artifact%s from datastore %s: %s",
2749 len(exceptions),
2750 s_err,
2751 len(artifacts_to_delete),
2752 "s" if len(artifacts_to_delete) != 1 else "",
2753 self.name,
2754 e,
2755 )
2756 else:
2757 raise e
2758 return removed
2760 @transactional
2761 def transfer_from(
2762 self,
2763 source_records: FileTransferMap,
2764 refs: Collection[DatasetRef],
2765 transfer: str = "auto",
2766 artifact_existence: dict[ResourcePath, bool] | None = None,
2767 dry_run: bool = False,
2768 ) -> tuple[set[DatasetRef], set[DatasetRef]]:
2769 log.verbose("Transferring %d datasets to %s", len(refs), self.name)
2771 # Stop early if "direct" transfer mode is requested. That would
2772 # require that the URI inside the source datastore should be stored
2773 # directly in the target datastore, which seems unlikely to be useful
2774 # since at any moment the source datastore could delete the file.
2775 if transfer in ("direct", "split"):
2776 raise ValueError(
2777 f"Can not transfer from a source datastore using {transfer} mode since"
2778 " those files are controlled by the other datastore."
2779 )
2781 if not refs:
2782 return set(), set()
2784 # Empty existence lookup if none given.
2785 if artifact_existence is None:
2786 artifact_existence = {}
2788 # In order to handle disassembled composites the code works
2789 # at the records level since it can assume that internal APIs
2790 # can be used.
2791 # - If the record already exists in the destination this is assumed
2792 # to be okay.
2793 # - If there is no record but the source and destination URIs are
2794 # identical no transfer is done but the record is added.
2795 # - If the source record refers to an absolute URI currently assume
2796 # that that URI should remain absolute and will be visible to the
2797 # destination butler. May need to have a flag to indicate whether
2798 # the dataset should be transferred. This will only happen if
2799 # the detached Butler has had a local ingest.
2801 # See if we already have these records
2802 log.verbose("Looking up existing datastore records in target %s for %d refs", self.name, len(refs))
2803 target_records = self._get_stored_records_associated_with_refs(refs, ignore_datastore_records=True)
2805 # The artifacts to register
2806 artifacts = []
2808 # Refs that already exist
2809 already_present = []
2811 # Refs that were rejected by this datastore.
2812 rejected = set()
2814 # Refs that were transferred successfully.
2815 accepted = set()
2817 # Record each time we have done a "direct" transfer.
2818 direct_transfers = []
2820 # Keep track of all the file transfers that are required.
2821 from_to: list[tuple[ResourcePath, ResourcePath]] = []
2823 # Now can transfer the artifacts
2824 log.verbose("Transferring artifacts")
2825 for ref in refs:
2826 if not self.constraints.isAcceptable(ref):
2827 # This datastore should not be accepting this dataset.
2828 rejected.add(ref)
2829 continue
2831 accepted.add(ref)
2833 if ref.id in target_records:
2834 # Already have an artifact for this.
2835 already_present.append(ref)
2836 continue
2838 # mypy needs to know these are always resolved refs
2839 for transfer_info in source_records.get(ref.id, []):
2840 info = transfer_info.file_info
2841 source_location = transfer_info.location
2842 target_location = info.file_location(self.locationFactory)
2843 if source_location == target_location and not source_location.pathInStore.isabs():
2844 # Artifact is already in the target location.
2845 # (which is how execution butler currently runs)
2846 pass
2847 else:
2848 if target_location.pathInStore.isabs():
2849 # Just because we can see the artifact when running
2850 # the transfer doesn't mean it will be generally
2851 # accessible to a user of this butler. Need to decide
2852 # what to do about an absolute path.
2853 if transfer == "auto":
2854 # For "auto" transfers we allow the absolute URI
2855 # to be recorded in the target datastore.
2856 direct_transfers.append(source_location)
2857 else:
2858 # The user is explicitly requesting a transfer
2859 # even for an absolute URI. This requires us to
2860 # calculate the target path.
2861 template_ref = ref
2862 if info.component:
2863 template_ref = ref.makeComponentRef(info.component)
2864 target_location = self._calculate_ingested_datastore_name(
2865 source_location.uri,
2866 template_ref,
2867 )
2869 info = info.update(path=target_location.pathInStore.path)
2871 # Need to transfer it to the new location.
2872 from_to.append((source_location.uri, target_location.uri))
2874 artifacts.append((ref, info))
2876 # Do the file transfers in bulk.
2877 # Assume we should always overwrite. If the artifact
2878 # is there this might indicate that a previous transfer
2879 # was interrupted but was not able to be rolled back
2880 # completely (eg pre-emption) so follow Datastore default
2881 # and overwrite. Do not copy if we are in dry-run mode.
2882 if dry_run:
2883 log.info("Would be copying %d file artifacts", len(from_to))
2884 else:
2885 log.verbose("Copying %d file artifacts", len(from_to))
2886 with time_this(log, msg="Transferring datasets into datastore", level=VERBOSE):
2887 ResourcePath.mtransfer(
2888 transfer,
2889 from_to,
2890 overwrite=True,
2891 transaction=self._transaction,
2892 )
2894 if direct_transfers:
2895 log.info(
2896 "Transfer request for an outside-datastore artifact with absolute URI done %d time%s",
2897 len(direct_transfers),
2898 "" if len(direct_transfers) == 1 else "s",
2899 )
2901 # We are overwriting previous datasets that may have already
2902 # existed. We therefore should ensure that we force the
2903 # datastore records to agree. Note that this can potentially lead
2904 # to difficulties if the dataset has previously been ingested
2905 # disassembled and is somehow now assembled, or vice versa.
2906 if not dry_run:
2907 log.verbose("Registering datastore records in database")
2908 self._register_datasets(artifacts, insert_mode=DatabaseInsertMode.REPLACE)
2910 if already_present:
2911 n_skipped = len(already_present)
2912 log.info(
2913 "Skipped transfer of %d dataset%s already present in datastore",
2914 n_skipped,
2915 "" if n_skipped == 1 else "s",
2916 )
2918 log.verbose(
2919 "Finished transfer_from to %s with %d accepted, %d rejected",
2920 self.name,
2921 len(accepted),
2922 len(rejected),
2923 )
2924 return accepted, rejected
2926 def get_file_info_for_transfer(self, dataset_ids: Iterable[DatasetId]) -> FileTransferMap:
2927 source_records = self._get_stored_records_associated_with_refs(
2928 [FakeDatasetRef(id) for id in dataset_ids], ignore_datastore_records=True
2929 )
2930 return self._convert_stored_file_info_to_file_transfer_record(source_records)
2932 def locate_missing_files_for_transfer(
2933 self, refs: Iterable[DatasetRef], artifact_existence: dict[ResourcePath, bool]
2934 ) -> FileTransferMap:
2935 missing_ids = {ref.id for ref in refs}
2936 # Missing IDs can be okay if that datastore has allowed
2937 # gets based on file existence. Should we transfer what we can
2938 # or complain about it and warn?
2939 if not self.trustGetRequest:
2940 return {}
2942 found_records = self._find_missing_records(
2943 refs, missing_ids, artifact_existence, warn_for_missing=False
2944 )
2945 return self._convert_stored_file_info_to_file_transfer_record(found_records)
2947 def _convert_stored_file_info_to_file_transfer_record(
2948 self, info_map: dict[DatasetId, list[StoredFileInfo]]
2949 ) -> FileTransferMap:
2950 output: dict[DatasetId, list[FileTransferRecord]] = {}
2951 for k, file_info_list in info_map.items():
2952 output[k] = [
2953 FileTransferRecord(file_info=info, location=info.file_location(self.locationFactory))
2954 for info in file_info_list
2955 ]
2956 return output
2958 @transactional
2959 def forget(self, refs: Iterable[DatasetRef]) -> None:
2960 # Docstring inherited.
2961 refs = list(refs)
2962 self.bridge.forget(refs)
2963 self._table.delete(["dataset_id"], *[{"dataset_id": ref.id} for ref in refs])
2965 def validateConfiguration(
2966 self, entities: Iterable[DatasetRef | DatasetType | StorageClass], logFailures: bool = False
2967 ) -> None:
2968 """Validate some of the configuration for this datastore.
2970 Parameters
2971 ----------
2972 entities : iterable of `DatasetRef`, `DatasetType`, or `StorageClass`
2973 Entities to test against this configuration. Can be differing
2974 types.
2975 logFailures : `bool`, optional
2976 If `True`, output a log message for every validation error
2977 detected.
2979 Raises
2980 ------
2981 DatastoreValidationError
2982 Raised if there is a validation problem with a configuration.
2983 All the problems are reported in a single exception.
2985 Notes
2986 -----
2987 This method checks that all the supplied entities have valid file
2988 templates and also have formatters defined.
2989 """
2990 templateFailed = None
2991 try:
2992 self.templates.validateTemplates(entities, logFailures=logFailures)
2993 except FileTemplateValidationError as e:
2994 templateFailed = str(e)
2996 formatterFailed = []
2997 for entity in entities:
2998 try:
2999 self.formatterFactory.getFormatterClass(entity)
3000 except KeyError as e:
3001 formatterFailed.append(str(e))
3002 if logFailures:
3003 log.critical("Formatter failure: %s", e)
3005 if templateFailed or formatterFailed:
3006 messages = []
3007 if templateFailed:
3008 messages.append(templateFailed)
3009 if formatterFailed:
3010 messages.append(",".join(formatterFailed))
3011 msg = ";\n".join(messages)
3012 raise DatastoreValidationError(msg)
3014 def getLookupKeys(self) -> set[LookupKey]:
3015 # Docstring is inherited from base class
3016 return (
3017 self.templates.getLookupKeys()
3018 | self.formatterFactory.getLookupKeys()
3019 | self.constraints.getLookupKeys()
3020 )
3022 def validateKey(self, lookupKey: LookupKey, entity: DatasetRef | DatasetType | StorageClass) -> None:
3023 # Docstring is inherited from base class
3024 # The key can be valid in either formatters or templates so we can
3025 # only check the template if it exists
3026 if lookupKey in self.templates:
3027 try:
3028 self.templates[lookupKey].validateTemplate(entity)
3029 except FileTemplateValidationError as e:
3030 raise DatastoreValidationError(e) from e
3032 def export(
3033 self,
3034 refs: Iterable[DatasetRef],
3035 *,
3036 directory: ResourcePathExpression | None = None,
3037 transfer: str | None = "auto",
3038 ) -> Iterable[FileDataset]:
3039 # Docstring inherited from Datastore.export.
3040 if transfer == "auto" and directory is None:
3041 transfer = None
3043 if transfer is not None and transfer != "direct" and directory is None:
3044 raise TypeError(f"Cannot export using transfer mode {transfer} with no export directory given")
3046 if transfer == "move":
3047 raise TypeError("Can not export by moving files out of datastore.")
3049 # Force the directory to be a URI object
3050 directoryUri: ResourcePath | None = None
3051 if directory is not None:
3052 directoryUri = ResourcePath(directory, forceDirectory=True)
3054 if transfer is not None and directoryUri is not None and not directoryUri.exists():
3055 # mypy needs the second test
3056 raise FileNotFoundError(f"Export location {directory} does not exist")
3058 progress = Progress("lsst.daf.butler.datastores.FileDatastore.export", level=logging.DEBUG)
3059 for ref in progress.wrap(refs, "Exporting dataset files"):
3060 fileLocations = self._get_dataset_locations_info(ref)
3061 if not fileLocations:
3062 raise FileNotFoundError(f"Could not retrieve dataset {ref}.")
3063 # For now we can not export disassembled datasets
3064 if len(fileLocations) > 1:
3065 raise NotImplementedError(f"Can not export disassembled datasets such as {ref}")
3066 location, storedFileInfo = fileLocations[0]
3068 pathInStore = location.pathInStore.path
3069 if transfer is None:
3070 # TODO: do we also need to return the readStorageClass somehow?
3071 # We will use the path in store directly. If this is an
3072 # absolute URI, preserve it.
3073 if location.pathInStore.isabs():
3074 pathInStore = str(location.uri)
3075 elif transfer == "direct":
3076 # Use full URIs to the remote store in the export
3077 pathInStore = str(location.uri)
3078 else:
3079 # mypy needs help
3080 assert directoryUri is not None, "directoryUri must be defined to get here"
3081 storeUri = ResourcePath(location.uri, forceDirectory=False)
3083 # if the datastore has an absolute URI to a resource, we
3084 # have two options:
3085 # 1. Keep the absolute URI in the exported YAML
3086 # 2. Allocate a new name in the local datastore and transfer
3087 # it.
3088 # For now go with option 2
3089 if location.pathInStore.isabs():
3090 template = self.templates.getTemplate(ref)
3091 newURI = ResourcePath(template.format(ref), forceAbsolute=False, forceDirectory=False)
3092 pathInStore = str(newURI.updatedExtension(location.pathInStore.getExtension()))
3094 exportUri = directoryUri.join(pathInStore)
3095 exportUri.transfer_from(storeUri, transfer=transfer)
3097 yield FileDataset(refs=[ref], path=pathInStore, formatter=storedFileInfo.formatter)
3099 @staticmethod
3100 def computeChecksum(uri: ResourcePath, algorithm: str = "blake2b", block_size: int = 8192) -> str | None:
3101 """Compute the checksum of the supplied file.
3103 Parameters
3104 ----------
3105 uri : `lsst.resources.ResourcePath`
3106 Name of resource to calculate checksum from.
3107 algorithm : `str`, optional
3108 Name of algorithm to use. Must be one of the algorithms supported
3109 by :py:class`hashlib`.
3110 block_size : `int`
3111 Number of bytes to read from file at one time.
3113 Returns
3114 -------
3115 hexdigest : `str`
3116 Hex digest of the file.
3118 Notes
3119 -----
3120 Currently returns None if the URI is for a remote resource.
3121 """
3122 if algorithm not in hashlib.algorithms_guaranteed:
3123 raise NameError(f"The specified algorithm '{algorithm}' is not supported by hashlib")
3125 if not uri.isLocal:
3126 return None
3128 hasher = hashlib.new(algorithm)
3130 with uri.as_local() as local_uri, open(local_uri.ospath, "rb") as f:
3131 for chunk in iter(lambda: f.read(block_size), b""):
3132 hasher.update(chunk)
3134 return hasher.hexdigest()
3136 def needs_expanded_data_ids(
3137 self,
3138 transfer: str | None,
3139 entity: DatasetRef | DatasetType | StorageClass | None = None,
3140 ) -> bool:
3141 # Docstring inherited.
3142 # This _could_ also use entity to inspect whether the filename template
3143 # involves placeholders other than the required dimensions for its
3144 # dataset type, but that's not necessary for correctness; it just
3145 # enables more optimizations (perhaps only in theory).
3146 return transfer not in ("direct", None)
3148 def import_records(self, data: Mapping[str, DatastoreRecordData]) -> None:
3149 # Docstring inherited from the base class.
3150 record_data = data.get(self.name)
3151 if not record_data:
3152 return
3154 self._bridge.insert(FakeDatasetRef(dataset_id) for dataset_id in record_data.records)
3156 # TODO: Verify that there are no unexpected table names in the dict?
3157 unpacked_records = []
3158 for dataset_id, dataset_data in record_data.records.items():
3159 records = dataset_data.get(self._table.name)
3160 if records:
3161 for info in records:
3162 assert isinstance(info, StoredFileInfo), "Expecting StoredFileInfo records"
3163 unpacked_records.append(info.to_record(dataset_id=dataset_id))
3164 if unpacked_records:
3165 self._table.insert(*unpacked_records, transaction=self._transaction)
3167 def export_records(self, refs: Iterable[DatasetIdRef]) -> Mapping[str, DatastoreRecordData]:
3168 # Docstring inherited from the base class.
3170 # This call to 'bridge.check' filters out "partially deleted" datasets.
3171 # Specifically, ones in the unusual edge state that:
3172 # 1. They have an entry in the registry dataset tables
3173 # 2. They were "trashed" from the datastore, so they are not
3174 # present in the "dataset_location" table.)
3175 # 3. But the trash has not been "emptied", so there are still entries
3176 # in the "opaque" datastore records table.
3177 #
3178 # As far as I can tell, this can only occur in the case of a concurrent
3179 # or aborted call to `Butler.pruneDatasets(unstore=True, purge=False)`.
3180 # Datasets (with or without files existing on disk) can persist in
3181 # this zombie state indefinitely, until someone manually empties
3182 # the trash.
3183 exported_refs = list(self._bridge.check(refs))
3184 ids = {ref.id for ref in exported_refs}
3185 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {id: {} for id in ids}
3186 for row in self._table.fetch(dataset_id=ids):
3187 info: StoredDatastoreItemInfo = StoredFileInfo.from_record(row)
3188 dataset_records = records.setdefault(row["dataset_id"], {})
3189 dataset_records.setdefault(self._table.name, []).append(info)
3191 record_data = DatastoreRecordData(records=records)
3192 return {self.name: record_data}
3194 def export_predicted_records(self, refs: Iterable[DatasetRef]) -> dict[str, DatastoreRecordData]:
3195 # Docstring inherited from the base class.
3196 refs = [self._cast_storage_class(ref) for ref in refs]
3197 records: dict[DatasetId, dict[str, list[StoredDatastoreItemInfo]]] = {}
3198 for ref in refs:
3199 if not self.constraints.isAcceptable(ref):
3200 continue
3201 fileLocations = self._get_expected_dataset_locations_info(ref)
3202 if not fileLocations:
3203 continue
3204 dataset_records = records.setdefault(ref.id, {})
3205 dataset_records.setdefault(self._table.name, [])
3206 for _, storedFileInfo in fileLocations:
3207 dataset_records[self._table.name].append(storedFileInfo)
3209 record_data = DatastoreRecordData(records=records)
3210 return {self.name: record_data}
3212 def set_retrieve_dataset_type_method(self, method: Callable[[str], DatasetType | None] | None) -> None:
3213 # Docstring inherited from the base class.
3214 self._retrieve_dataset_method = method
3216 def _cast_storage_class(self, ref: DatasetRef) -> DatasetRef:
3217 """Update dataset reference to use the storage class from registry."""
3218 if self._retrieve_dataset_method is None:
3219 # We could raise an exception here but unit tests do not define
3220 # this method.
3221 return ref
3222 dataset_type = self._retrieve_dataset_method(ref.datasetType.name)
3223 if dataset_type is not None:
3224 ref = ref.overrideStorageClass(dataset_type.storageClass_name)
3225 return ref
3227 def get_opaque_table_definitions(self) -> Mapping[str, DatastoreOpaqueTable]:
3228 # Docstring inherited from the base class.
3229 return {self._opaque_table_name: DatastoreOpaqueTable(self.makeTableSpec(), StoredFileInfo)}