Coverage for python/lsst/daf/butler/direct_butler/_direct_butler.py: 10%
919 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-19 01:32 -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"""Butler top level classes."""
30from __future__ import annotations
32__all__ = (
33 "ButlerValidationError",
34 "DirectButler",
35)
37import collections.abc
38import contextlib
39import io
40import itertools
41import math
42import numbers
43import os
44import uuid
45import warnings
46from collections import Counter, defaultdict
47from collections.abc import Collection, Iterable, Iterator, MutableMapping, Sequence
48from types import EllipsisType
49from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, TextIO, cast
51from deprecated.sphinx import deprecated
52from sqlalchemy.exc import IntegrityError
54from lsst.resources import ResourcePath, ResourcePathExpression
55from lsst.utils.introspection import find_outside_stacklevel, get_class_of
56from lsst.utils.iteration import chunk_iterable
57from lsst.utils.logging import VERBOSE, getLogger
58from lsst.utils.timer import time_this
60from .._butler import Butler, _DeprecatedDefault
61from .._butler_config import ButlerConfig
62from .._butler_instance_options import ButlerInstanceOptions
63from .._butler_metrics import ButlerMetrics
64from .._collection_type import CollectionType
65from .._dataset_existence import DatasetExistence
66from .._dataset_ref import DatasetRef
67from .._dataset_type import DatasetType
68from .._deferredDatasetHandle import DeferredDatasetHandle
69from .._exceptions import DatasetNotFoundError, DimensionValueError, EmptyQueryResultError, ValidationError
70from .._file_dataset import FileDataset
71from .._limited_butler import LimitedButler
72from .._query_all_datasets import QueryAllDatasetsParameters, query_all_datasets
73from .._registry_shim import RegistryShim
74from .._storage_class import StorageClass, StorageClassFactory
75from .._timespan import Timespan
76from ..datastore import Datastore, NullDatastore
77from ..datastores.file_datastore.retrieve_artifacts import ZipIndex, retrieve_and_zip
78from ..datastores.file_datastore.transfer import retrieve_file_transfer_records
79from ..dimensions import DataCoordinate, Dimension, DimensionGroup
80from ..direct_query_driver import DirectQueryDriver
81from ..progress import Progress
82from ..queries import Query
83from ..registry import (
84 ConflictingDefinitionError,
85 DataIdError,
86 MissingDatasetTypeError,
87 RegistryDefaults,
88 _RegistryFactory,
89)
90from ..registry.sql_registry import SqlRegistry
91from ..transfers import RepoExportContext
92from ..utils import transactional
93from ._direct_butler_collections import DirectButlerCollections
95if TYPE_CHECKING:
96 from lsst.resources import ResourceHandleProtocol
98 from .._dataset_provenance import DatasetProvenance
99 from .._dataset_ref import DatasetId
100 from ..datastore import DatasetRefURIs
101 from ..dimensions import DataId, DataIdValue, DimensionElement, DimensionRecord, DimensionUniverse
102 from ..registry import CollectionArgType, Registry
103 from ..transfers import RepoImportBackend
105_LOG = getLogger(__name__)
108class ButlerValidationError(ValidationError):
109 """There is a problem with the Butler configuration."""
111 pass
114class DirectButler(Butler): # numpydoc ignore=PR02
115 """Main entry point for the data access system.
117 Parameters
118 ----------
119 config : `ButlerConfig`
120 The configuration for this Butler instance.
121 registry : `SqlRegistry`
122 The object that manages dataset metadata and relationships.
123 datastore : Datastore
124 The object that manages actual dataset storage.
125 storageClasses : StorageClassFactory
126 An object that maps known storage class names to objects that fully
127 describe them.
129 Notes
130 -----
131 Most users should call the top-level `Butler`.``from_config`` instead of
132 using this constructor directly.
133 """
135 # This is __new__ instead of __init__ because we have to support
136 # instantiation via the legacy constructor Butler.__new__(), which
137 # reads the configuration and selects which subclass to instantiate. The
138 # interaction between __new__ and __init__ is kind of wacky in Python. If
139 # we were using __init__ here, __init__ would be called twice (once when
140 # the DirectButler instance is constructed inside Butler.from_config(), and
141 # a second time with the original arguments to Butler() when the instance
142 # is returned from Butler.__new__()
143 def __new__(
144 cls,
145 *,
146 config: ButlerConfig,
147 registry: SqlRegistry,
148 datastore: Datastore,
149 storageClasses: StorageClassFactory,
150 metrics: ButlerMetrics | None = None,
151 ) -> DirectButler:
152 self = cast(DirectButler, super().__new__(cls))
153 self._config = config
154 self._registry = registry
155 self._datastore = datastore
156 self.storageClasses = storageClasses
157 self._metrics = metrics if metrics is not None else ButlerMetrics()
159 # For execution butler the datastore needs a special
160 # dependency-inversion trick. This is not used by regular butler,
161 # but we do not have a way to distinguish regular butler from execution
162 # butler.
163 self._datastore.set_retrieve_dataset_type_method(self._retrieve_dataset_type)
165 self._registry_shim = RegistryShim(self)
167 return self
169 @classmethod
170 def create_from_config(
171 cls,
172 config: ButlerConfig,
173 *,
174 options: ButlerInstanceOptions,
175 without_datastore: bool = False,
176 ) -> DirectButler:
177 """Construct a Butler instance from a configuration file.
179 Parameters
180 ----------
181 config : `ButlerConfig`
182 The configuration for this Butler instance.
183 options : `ButlerInstanceOptions`
184 Default values and other settings for the Butler instance.
185 without_datastore : `bool`, optional
186 If `True` do not attach a datastore to this butler. Any attempts
187 to use a datastore will fail.
189 Notes
190 -----
191 Most users should call the top-level `Butler`.``from_config``
192 instead of using this function directly.
193 """
194 if "run" in config or "collection" in config:
195 raise ValueError("Passing a run or collection via configuration is no longer supported.")
197 defaults = RegistryDefaults.from_butler_instance_options(options)
198 try:
199 butlerRoot = config.get("root", config.configDir)
200 writeable = options.writeable
201 if writeable is None:
202 writeable = options.run is not None
203 registry = _RegistryFactory(config).from_config(
204 butlerRoot=butlerRoot, writeable=writeable, defaults=defaults
205 )
206 if without_datastore:
207 datastore: Datastore = NullDatastore(None, None)
208 else:
209 datastore = Datastore.fromConfig(
210 config, registry.getDatastoreBridgeManager(), butlerRoot=butlerRoot
211 )
212 # TODO: Once datastore drops dependency on registry we can
213 # construct datastore first and pass opaque tables to registry
214 # constructor.
215 registry.make_datastore_tables(datastore.get_opaque_table_definitions())
216 storageClasses = StorageClassFactory()
217 storageClasses.addFromConfig(config)
219 return DirectButler(
220 config=config,
221 registry=registry,
222 datastore=datastore,
223 storageClasses=storageClasses,
224 metrics=options.metrics,
225 )
226 except Exception:
227 # Failures here usually mean that configuration is incomplete,
228 # just issue an error message which includes config file URI.
229 _LOG.error(f"Failed to instantiate Butler from config {config.configFile}.")
230 raise
232 def clone(
233 self,
234 *,
235 collections: CollectionArgType | None | EllipsisType = ...,
236 run: str | None | EllipsisType = ...,
237 inferDefaults: bool | EllipsisType = ...,
238 dataId: dict[str, str] | EllipsisType = ...,
239 metrics: ButlerMetrics | None = None,
240 ) -> DirectButler:
241 # Docstring inherited
242 defaults = self._registry.defaults.clone(collections, run, inferDefaults, dataId)
243 registry = self._registry.copy(defaults)
245 return DirectButler(
246 registry=registry,
247 config=self._config,
248 datastore=self._datastore.clone(registry.getDatastoreBridgeManager()),
249 storageClasses=self.storageClasses,
250 metrics=metrics,
251 )
253 GENERATION: ClassVar[int] = 3
254 """This is a Generation 3 Butler.
256 This attribute may be removed in the future, once the Generation 2 Butler
257 interface has been fully retired; it should only be used in transitional
258 code.
259 """
261 def _retrieve_dataset_type(self, name: str) -> DatasetType | None:
262 """Return DatasetType defined in registry given dataset type name."""
263 try:
264 return self.get_dataset_type(name)
265 except MissingDatasetTypeError:
266 return None
268 @classmethod
269 def _unpickle(
270 cls,
271 config: ButlerConfig,
272 collections: tuple[str, ...] | None,
273 run: str | None,
274 defaultDataId: dict[str, str],
275 writeable: bool,
276 ) -> DirectButler:
277 """Callable used to unpickle a Butler.
279 We prefer not to use ``Butler.__init__`` directly so we can force some
280 of its many arguments to be keyword-only (note that ``__reduce__``
281 can only invoke callables with positional arguments).
283 Parameters
284 ----------
285 config : `ButlerConfig`
286 Butler configuration, already coerced into a true `ButlerConfig`
287 instance (and hence after any search paths for overrides have been
288 utilized).
289 collections : `tuple` [ `str` ]
290 Names of the default collections to read from.
291 run : `str`, optional
292 Name of the default `~CollectionType.RUN` collection to write to.
293 defaultDataId : `dict` [ `str`, `str` ]
294 Default data ID values.
295 writeable : `bool`
296 Whether the Butler should support write operations.
298 Returns
299 -------
300 butler : `Butler`
301 A new `Butler` instance.
302 """
303 return cls.create_from_config(
304 config=config,
305 options=ButlerInstanceOptions(
306 collections=collections, run=run, writeable=writeable, kwargs=defaultDataId
307 ),
308 )
310 def __reduce__(self) -> tuple:
311 """Support pickling."""
312 return (
313 DirectButler._unpickle,
314 (
315 self._config,
316 self.collections.defaults,
317 self.run,
318 dict(self._registry.defaults.dataId.required),
319 self._registry.isWriteable(),
320 ),
321 )
323 def __str__(self) -> str:
324 return (
325 f"Butler(collections={self.collections}, run={self.run}, "
326 f"datastore='{self._datastore}', registry='{self._registry}')"
327 )
329 def isWriteable(self) -> bool:
330 # Docstring inherited.
331 return self._registry.isWriteable()
333 def _caching_context(self) -> contextlib.AbstractContextManager[None]:
334 """Context manager that enables caching."""
335 return self._registry.caching_context()
337 @contextlib.contextmanager
338 def transaction(self) -> Iterator[None]:
339 """Context manager supporting `Butler` transactions.
341 Transactions can be nested.
342 """
343 with self._registry.transaction(), self._datastore.transaction():
344 yield
346 def _standardizeArgs(
347 self,
348 datasetRefOrType: DatasetRef | DatasetType | str,
349 dataId: DataId | None = None,
350 for_put: bool = True,
351 **kwargs: Any,
352 ) -> tuple[DatasetType, DataId | None]:
353 """Standardize the arguments passed to several Butler APIs.
355 Parameters
356 ----------
357 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
358 When `DatasetRef` the `dataId` should be `None`.
359 Otherwise the `DatasetType` or name thereof.
360 dataId : `dict` or `DataCoordinate`
361 A `dict` of `Dimension` link name, value pairs that label the
362 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
363 should be provided as the second argument.
364 for_put : `bool`, optional
365 If `True` this call is invoked as part of a `Butler.put()`.
366 Otherwise it is assumed to be part of a `Butler.get()`. This
367 parameter is only relevant if there is dataset type
368 inconsistency.
369 **kwargs
370 Additional keyword arguments used to augment or construct a
371 `DataCoordinate`. See `DataCoordinate.standardize`
372 parameters.
374 Returns
375 -------
376 datasetType : `DatasetType`
377 A `DatasetType` instance extracted from ``datasetRefOrType``.
378 dataId : `dict` or `DataId`, optional
379 Argument that can be used (along with ``kwargs``) to construct a
380 `DataId`.
382 Notes
383 -----
384 Butler APIs that conceptually need a DatasetRef also allow passing a
385 `DatasetType` (or the name of one) and a `DataId` (or a dict and
386 keyword arguments that can be used to construct one) separately. This
387 method accepts those arguments and always returns a true `DatasetType`
388 and a `DataId` or `dict`.
390 Standardization of `dict` vs `DataId` is best handled by passing the
391 returned ``dataId`` (and ``kwargs``) to `Registry` APIs, which are
392 generally similarly flexible.
393 """
394 externalDatasetType: DatasetType | None = None
395 internalDatasetType: DatasetType | None = None
396 if isinstance(datasetRefOrType, DatasetRef):
397 if dataId is not None or kwargs:
398 raise ValueError("DatasetRef given, cannot use dataId as well")
399 externalDatasetType = datasetRefOrType.datasetType
400 dataId = datasetRefOrType.dataId
401 else:
402 # Don't check whether DataId is provided, because Registry APIs
403 # can usually construct a better error message when it wasn't.
404 if isinstance(datasetRefOrType, DatasetType):
405 externalDatasetType = datasetRefOrType
406 else:
407 internalDatasetType = self.get_dataset_type(datasetRefOrType)
409 # Check that they are self-consistent
410 if externalDatasetType is not None:
411 internalDatasetType = self.get_dataset_type(externalDatasetType.name)
412 if externalDatasetType != internalDatasetType:
413 # We can allow differences if they are compatible, depending
414 # on whether this is a get or a put. A get requires that
415 # the python type associated with the datastore can be
416 # converted to the user type. A put requires that the user
417 # supplied python type can be converted to the internal
418 # type expected by registry.
419 relevantDatasetType = internalDatasetType
420 if for_put:
421 is_compatible = internalDatasetType.is_compatible_with(externalDatasetType)
422 else:
423 is_compatible = externalDatasetType.is_compatible_with(internalDatasetType)
424 relevantDatasetType = externalDatasetType
425 if not is_compatible:
426 raise ValueError(
427 f"Supplied dataset type ({externalDatasetType}) inconsistent with "
428 f"registry definition ({internalDatasetType})"
429 )
430 # Override the internal definition.
431 internalDatasetType = relevantDatasetType
433 assert internalDatasetType is not None
434 return internalDatasetType, dataId
436 def _rewrite_data_id(
437 self, dataId: DataId | None, datasetType: DatasetType, **kwargs: Any
438 ) -> tuple[DataId | None, dict[str, Any]]:
439 """Rewrite a data ID taking into account dimension records.
441 Take a Data ID and keyword args and rewrite it if necessary to
442 allow the user to specify dimension records rather than dimension
443 primary values.
445 This allows a user to include a dataId dict with keys of
446 ``exposure.day_obs`` and ``exposure.seq_num`` instead of giving
447 the integer exposure ID. It also allows a string to be given
448 for a dimension value rather than the integer ID if that is more
449 convenient. For example, rather than having to specifying the
450 detector with ``detector.full_name``, a string given for ``detector``
451 will be interpreted as the full name and converted to the integer
452 value.
454 Keyword arguments can also use strings for dimensions like detector
455 and exposure but python does not allow them to include ``.`` and
456 so the ``exposure.day_obs`` syntax can not be used in a keyword
457 argument.
459 Parameters
460 ----------
461 dataId : `dict` or `DataCoordinate`
462 A `dict` of `Dimension` link name, value pairs that will label the
463 `DatasetRef` within a Collection.
464 datasetType : `DatasetType`
465 The dataset type associated with this dataId. Required to
466 determine the relevant dimensions.
467 **kwargs
468 Additional keyword arguments used to augment or construct a
469 `DataId`. See `DataId` parameters.
471 Returns
472 -------
473 dataId : `dict` or `DataCoordinate`
474 The, possibly rewritten, dataId. If given a `DataCoordinate` and
475 no keyword arguments, the original dataId will be returned
476 unchanged.
477 **kwargs : `dict`
478 Any unused keyword arguments (would normally be empty dict).
479 """
480 # Process dimension records that are using record information
481 # rather than ids
482 newDataId: dict[str, DataIdValue] = {}
483 byRecord: dict[str, dict[str, Any]] = defaultdict(dict)
485 if isinstance(dataId, DataCoordinate):
486 # Do nothing if we have a DataCoordinate and no kwargs.
487 if not kwargs:
488 return dataId, kwargs
489 # If we have a DataCoordinate with kwargs, we know the
490 # DataCoordinate only has values for real dimensions.
491 newDataId.update(dataId.mapping)
492 elif dataId:
493 # The data is mapping, which means it might have keys like
494 # "exposure.obs_id" (unlike kwargs, because a "." is not allowed in
495 # a keyword parameter).
496 for k, v in dataId.items():
497 if isinstance(k, str) and "." in k:
498 # Someone is using a more human-readable dataId
499 dimensionName, record = k.split(".", 1)
500 byRecord[dimensionName][record] = v
501 else:
502 newDataId[k] = v
504 # Go through the updated dataId and check the type in case someone is
505 # using an alternate key. We have already filtered out the compound
506 # keys dimensions.record format.
507 not_dimensions = {}
509 # Will need to look in the dataId and the keyword arguments
510 # and will remove them if they need to be fixed or are unrecognized.
511 for dataIdDict in (newDataId, kwargs):
512 # Use a list so we can adjust the dict safely in the loop
513 for dimensionName in list(dataIdDict):
514 value = dataIdDict[dimensionName]
515 try:
516 dimension = self.dimensions.dimensions[dimensionName]
517 except KeyError:
518 # This is not a real dimension
519 not_dimensions[dimensionName] = value
520 del dataIdDict[dimensionName]
521 continue
523 # Convert an integral type to an explicit int to simplify
524 # comparisons here
525 if isinstance(value, numbers.Integral):
526 value = int(value)
528 if not isinstance(value, dimension.primaryKey.getPythonType()):
529 for alternate in dimension.alternateKeys:
530 if isinstance(value, alternate.getPythonType()):
531 byRecord[dimensionName][alternate.name] = value
532 del dataIdDict[dimensionName]
533 _LOG.debug(
534 "Converting dimension %s to %s.%s=%s",
535 dimensionName,
536 dimensionName,
537 alternate.name,
538 value,
539 )
540 break
541 else:
542 _LOG.warning(
543 "Type mismatch found for value '%r' provided for dimension %s. "
544 "Could not find matching alternative (primary key has type %s) "
545 "so attempting to use as-is.",
546 value,
547 dimensionName,
548 dimension.primaryKey.getPythonType(),
549 )
551 # By this point kwargs and newDataId should only include valid
552 # dimensions. Merge kwargs in to the new dataId and log if there
553 # are dimensions in both (rather than calling update).
554 for k, v in kwargs.items():
555 if k in newDataId and newDataId[k] != v:
556 _LOG.debug(
557 "Keyword arg %s overriding explicit value in dataId of %s with %s", k, newDataId[k], v
558 )
559 newDataId[k] = v
560 # No need to retain any values in kwargs now.
561 kwargs = {}
563 # If we have some unrecognized dimensions we have to try to connect
564 # them to records in other dimensions. This is made more complicated
565 # by some dimensions having records with clashing names. A mitigation
566 # is that we can tell by this point which dimensions are missing
567 # for the DatasetType but this does not work for calibrations
568 # where additional dimensions can be used to constrain the temporal
569 # axis.
570 if not_dimensions:
571 # Search for all dimensions even if we have been given a value
572 # explicitly. In some cases records are given as well as the
573 # actually dimension and this should not be an error if they
574 # match.
575 mandatoryDimensions = datasetType.dimensions.names # - provided
577 candidateDimensions: set[str] = set()
578 candidateDimensions.update(mandatoryDimensions)
580 # For calibrations we may well be needing temporal dimensions
581 # so rather than always including all dimensions in the scan
582 # restrict things a little. It is still possible for there
583 # to be confusion over day_obs in visit vs exposure for example.
584 # If we are not searching calibration collections things may
585 # fail but they are going to fail anyway because of the
586 # ambiguousness of the dataId...
587 if datasetType.isCalibration():
588 for dim in self.dimensions.dimensions:
589 if dim.temporal:
590 candidateDimensions.add(str(dim))
592 # Look up table for the first association with a dimension
593 guessedAssociation: dict[str, dict[str, Any]] = defaultdict(dict)
595 # Keep track of whether an item is associated with multiple
596 # dimensions.
597 counter: Counter[str] = Counter()
598 assigned: dict[str, set[str]] = defaultdict(set)
600 # Go through the missing dimensions and associate the
601 # given names with records within those dimensions
602 matched_dims = set()
603 for dimensionName in candidateDimensions:
604 dimension = self.dimensions.dimensions[dimensionName]
605 fields = dimension.metadata.names | dimension.uniqueKeys.names
606 for field in not_dimensions:
607 if field in fields:
608 guessedAssociation[dimensionName][field] = not_dimensions[field]
609 counter[dimensionName] += 1
610 assigned[field].add(dimensionName)
611 matched_dims.add(field)
613 # Calculate the fields that matched nothing.
614 never_found = set(not_dimensions) - matched_dims
616 if never_found:
617 raise DimensionValueError(f"Unrecognized keyword args given: {never_found}")
619 # There is a chance we have allocated a single dataId item
620 # to multiple dimensions. Need to decide which should be retained.
621 # For now assume that the most popular alternative wins.
622 # This means that day_obs with seq_num will result in
623 # exposure.day_obs and not visit.day_obs
624 # Also prefer an explicitly missing dimension over an inferred
625 # temporal dimension.
626 for fieldName, assignedDimensions in assigned.items():
627 if len(assignedDimensions) > 1:
628 # Pick the most popular (preferring mandatory dimensions)
629 requiredButMissing = assignedDimensions.intersection(mandatoryDimensions)
630 if requiredButMissing:
631 candidateDimensions = requiredButMissing
632 else:
633 candidateDimensions = assignedDimensions
635 # If this is a choice between visit and exposure and
636 # neither was a required part of the dataset type,
637 # (hence in this branch) always prefer exposure over
638 # visit since exposures are always defined and visits
639 # are defined from exposures.
640 if candidateDimensions == {"exposure", "visit"}:
641 candidateDimensions = {"exposure"}
643 # Select the relevant items and get a new restricted
644 # counter.
645 theseCounts = {k: v for k, v in counter.items() if k in candidateDimensions}
646 duplicatesCounter: Counter[str] = Counter()
647 duplicatesCounter.update(theseCounts)
649 # Choose the most common. If they are equally common
650 # we will pick the one that was found first.
651 # Returns a list of tuples
652 selected = duplicatesCounter.most_common(1)[0][0]
654 _LOG.debug(
655 "Ambiguous dataId entry '%s' associated with multiple dimensions: %s."
656 " Removed ambiguity by choosing dimension %s.",
657 fieldName,
658 ", ".join(assignedDimensions),
659 selected,
660 )
662 for candidateDimension in assignedDimensions:
663 if candidateDimension != selected:
664 del guessedAssociation[candidateDimension][fieldName]
666 # Update the record look up dict with the new associations
667 for dimensionName, values in guessedAssociation.items():
668 if values: # A dict might now be empty
669 _LOG.debug(
670 "Assigned non-dimension dataId keys to dimension %s: %s", dimensionName, values
671 )
672 byRecord[dimensionName].update(values)
674 if byRecord:
675 # Some record specifiers were found so we need to convert
676 # them to the Id form
677 for dimensionName, values in byRecord.items():
678 if dimensionName in newDataId:
679 _LOG.debug(
680 "DataId specified explicit %s dimension value of %s in addition to"
681 " general record specifiers for it of %s. Checking for self-consistency.",
682 dimensionName,
683 newDataId[dimensionName],
684 str(values),
685 )
686 # Get the actual record and compare with these values.
687 # Only query with relevant data ID values.
688 filtered_data_id = {
689 k: v for k, v in newDataId.items() if k in self.dimensions[dimensionName].required
690 }
691 try:
692 recs = self.query_dimension_records(
693 dimensionName,
694 data_id=filtered_data_id,
695 )
696 except (DataIdError, EmptyQueryResultError):
697 raise DimensionValueError(
698 f"Could not find dimension '{dimensionName}'"
699 f" with dataId {filtered_data_id} as part of comparing with"
700 f" record values {byRecord[dimensionName]}"
701 ) from None
702 if len(recs) == 1:
703 errmsg: list[str] = []
704 for k, v in values.items():
705 if (recval := getattr(recs[0], k)) != v:
706 errmsg.append(f"{k} ({recval} != {v})")
707 if errmsg:
708 raise DimensionValueError(
709 f"Dimension {dimensionName} in dataId has explicit value"
710 f" {newDataId[dimensionName]} inconsistent with"
711 f" {dimensionName} dimension record: " + ", ".join(errmsg)
712 )
713 else:
714 # Multiple matches for an explicit dimension
715 # should never happen but let downstream complain.
716 pass
717 continue
719 # Do not use data ID keys in query that aren't relevant.
720 # Otherwise we can have detector queries being constrained
721 # by an exposure ID that doesn't exist and return no matches
722 # for a detector even though it's a good detector name.
723 filtered_data_id = {
724 k: v
725 for k, v in newDataId.items()
726 if k in self.dimensions[dimensionName].minimal_group.names
727 }
729 def _get_attr(obj: Any, attr: str) -> Any:
730 # Used to implement x.exposure.seq_num when given
731 # x and "exposure.seq_num".
732 for component in attr.split("."):
733 obj = getattr(obj, component)
734 return obj
736 with self.query() as q:
737 x = q.expression_factory
738 # Build up a WHERE expression.
739 predicates = tuple(_get_attr(x, f"{dimensionName}.{k}") == v for k, v in values.items())
740 extra_args: dict[str, Any] = {} # For mypy.
741 extra_args.update(filtered_data_id)
742 extra_args.update(kwargs)
743 q = q.where(x.all(*predicates), **extra_args)
744 records = set(q.dimension_records(dimensionName))
746 if len(records) != 1:
747 if len(records) > 1:
748 # visit can have an ambiguous answer without involving
749 # visit_system. The default visit_system is defined
750 # by the instrument.
751 if (
752 dimensionName == "visit"
753 and "visit_system_membership" in self.dimensions
754 and "visit_system" in self.dimensions["instrument"].metadata
755 ):
756 instrument_records = self.query_dimension_records(
757 "instrument",
758 data_id=newDataId,
759 explain=False,
760 **kwargs,
761 )
762 if len(instrument_records) == 1:
763 visit_system = instrument_records[0].visit_system
764 if visit_system is None:
765 # Set to a value that will never match.
766 visit_system = -1
768 # Look up each visit in the
769 # visit_system_membership records.
770 for rec in records:
771 membership = self.query_dimension_records(
772 # Use bind to allow zero results.
773 # This is a fully-specified query.
774 "visit_system_membership",
775 instrument=instrument_records[0].name,
776 visit_system=visit_system,
777 visit=rec.id,
778 explain=False,
779 )
780 if membership:
781 # This record is the right answer.
782 records = {rec}
783 break
785 # The ambiguity may have been resolved so check again.
786 if len(records) > 1:
787 _LOG.debug(
788 "Received %d records from constraints of %s", len(records), str(values)
789 )
790 for r in records:
791 _LOG.debug("- %s", str(r))
792 raise DimensionValueError(
793 f"DataId specification for dimension {dimensionName} is not"
794 f" uniquely constrained to a single dataset by {values}."
795 f" Got {len(records)} results."
796 )
797 else:
798 raise DimensionValueError(
799 f"DataId specification for dimension {dimensionName} matched no"
800 f" records when constrained by {values}"
801 )
803 # Get the primary key from the real dimension object
804 dimension = self.dimensions.dimensions[dimensionName]
805 if not isinstance(dimension, Dimension):
806 raise RuntimeError(
807 f"{dimension.name} is not a true dimension, and cannot be used in data IDs."
808 )
809 newDataId[dimensionName] = getattr(records.pop(), dimension.primaryKey.name)
811 return newDataId, kwargs
813 def _findDatasetRef(
814 self,
815 datasetRefOrType: DatasetRef | DatasetType | str,
816 dataId: DataId | None = None,
817 *,
818 collections: Any = None,
819 predict: bool = False,
820 run: str | None = None,
821 datastore_records: bool = False,
822 timespan: Timespan | None = None,
823 **kwargs: Any,
824 ) -> DatasetRef:
825 """Shared logic for methods that start with a search for a dataset in
826 the registry.
828 Parameters
829 ----------
830 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
831 When `DatasetRef` the `dataId` should be `None`.
832 Otherwise the `DatasetType` or name thereof.
833 dataId : `dict` or `DataCoordinate`, optional
834 A `dict` of `Dimension` link name, value pairs that label the
835 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
836 should be provided as the first argument.
837 collections : Any, optional
838 Collections to be searched, overriding ``self.collections``.
839 Can be any of the types supported by the ``collections`` argument
840 to butler construction.
841 predict : `bool`, optional
842 If `True`, return a newly created `DatasetRef` with a unique
843 dataset ID if finding a reference in the `Registry` fails.
844 Defaults to `False`.
845 run : `str`, optional
846 Run collection name to use for creating `DatasetRef` for predicted
847 datasets. Only used if ``predict`` is `True`.
848 datastore_records : `bool`, optional
849 If `True` add datastore records to returned `DatasetRef`.
850 timespan : `Timespan` or `None`, optional
851 A timespan that the validity range of the dataset must overlap.
852 If not provided and this is a calibration dataset type, an attempt
853 will be made to find the timespan from any temporal coordinate
854 in the data ID.
855 **kwargs
856 Additional keyword arguments used to augment or construct a
857 `DataId`. See `DataId` parameters.
859 Returns
860 -------
861 ref : `DatasetRef`
862 A reference to the dataset identified by the given arguments.
863 This can be the same dataset reference as given if it was
864 resolved.
866 Raises
867 ------
868 LookupError
869 Raised if no matching dataset exists in the `Registry` (and
870 ``predict`` is `False`).
871 ValueError
872 Raised if a resolved `DatasetRef` was passed as an input, but it
873 differs from the one found in the registry.
874 TypeError
875 Raised if no collections were provided.
876 """
877 datasetType, dataId = self._standardizeArgs(datasetRefOrType, dataId, for_put=False, **kwargs)
878 if isinstance(datasetRefOrType, DatasetRef):
879 if collections is not None:
880 warnings.warn("Collections should not be specified with DatasetRef", stacklevel=3)
881 # May need to retrieve datastore records if requested.
882 if datastore_records and datasetRefOrType._datastore_records is None:
883 datasetRefOrType = self._registry.get_datastore_records(datasetRefOrType)
884 return datasetRefOrType
886 dataId, kwargs = self._rewrite_data_id(dataId, datasetType, **kwargs)
888 if datasetType.isCalibration():
889 # Because this is a calibration dataset, first try to make a
890 # standardize the data ID without restricting the dimensions to
891 # those of the dataset type requested, because there may be extra
892 # dimensions that provide temporal information for a validity-range
893 # lookup.
894 dataId = DataCoordinate.standardize(
895 dataId, universe=self.dimensions, defaults=self._registry.defaults.dataId, **kwargs
896 )
897 if timespan is None:
898 if dataId.dimensions.temporal:
899 dataId = self._registry.expandDataId(dataId)
900 # Use the timespan from the data ID to constrain the
901 # calibration lookup, but only if the caller has not
902 # specified an explicit timespan.
903 timespan = dataId.timespan
904 else:
905 # Try an arbitrary timespan. Downstream will fail if this
906 # results in more than one matching dataset.
907 timespan = Timespan(None, None)
908 else:
909 # Standardize the data ID to just the dimensions of the dataset
910 # type instead of letting registry.findDataset do it, so we get the
911 # result even if no dataset is found.
912 dataId = DataCoordinate.standardize(
913 dataId,
914 dimensions=datasetType.dimensions,
915 defaults=self._registry.defaults.dataId,
916 **kwargs,
917 )
918 # Always lookup the DatasetRef, even if one is given, to ensure it is
919 # present in the current collection.
920 ref = self.find_dataset(
921 datasetType,
922 dataId,
923 collections=collections,
924 timespan=timespan,
925 datastore_records=datastore_records,
926 )
927 if ref is None:
928 if predict:
929 if run is None:
930 run = self.run
931 if run is None:
932 raise TypeError("Cannot predict dataset ID/location with run=None.")
933 return DatasetRef(datasetType, dataId, run=run)
934 else:
935 if collections is None:
936 collections = self._registry.defaults.collections
937 raise DatasetNotFoundError(
938 f"Dataset {datasetType.name} with data ID {dataId} "
939 f"could not be found in collections {collections}."
940 )
941 if datasetType != ref.datasetType:
942 # If they differ it is because the user explicitly specified
943 # a compatible dataset type to this call rather than using the
944 # registry definition. The DatasetRef must therefore be recreated
945 # using the user definition such that the expected type is
946 # returned.
947 ref = DatasetRef(
948 datasetType, ref.dataId, run=ref.run, id=ref.id, datastore_records=ref._datastore_records
949 )
951 return ref
953 @transactional
954 def put(
955 self,
956 obj: Any,
957 datasetRefOrType: DatasetRef | DatasetType | str,
958 /,
959 dataId: DataId | None = None,
960 *,
961 run: str | None = None,
962 provenance: DatasetProvenance | None = None,
963 **kwargs: Any,
964 ) -> DatasetRef:
965 """Store and register a dataset.
967 Parameters
968 ----------
969 obj : `object`
970 The dataset.
971 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
972 When `DatasetRef` is provided, ``dataId`` should be `None`.
973 Otherwise the `DatasetType` or name thereof. If a fully resolved
974 `DatasetRef` is given the run and ID are used directly.
975 dataId : `dict` or `DataCoordinate`
976 A `dict` of `Dimension` link name, value pairs that label the
977 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
978 should be provided as the second argument.
979 run : `str`, optional
980 The name of the run the dataset should be added to, overriding
981 ``self.run``. Not used if a resolved `DatasetRef` is provided.
982 provenance : `DatasetProvenance` or `None`, optional
983 Any provenance that should be attached to the serialized dataset.
984 Not supported by all serialization mechanisms.
985 **kwargs
986 Additional keyword arguments used to augment or construct a
987 `DataCoordinate`. See `DataCoordinate.standardize`
988 parameters. Not used if a resolve `DatasetRef` is provided.
990 Returns
991 -------
992 ref : `DatasetRef`
993 A reference to the stored dataset, updated with the correct id if
994 given.
996 Raises
997 ------
998 TypeError
999 Raised if the butler is read-only or if no run has been provided.
1000 """
1001 if isinstance(datasetRefOrType, DatasetRef):
1002 # This is a direct put of predefined DatasetRef.
1003 _LOG.debug("Butler put direct: %s", datasetRefOrType)
1004 if run is not None:
1005 warnings.warn("Run collection is not used for DatasetRef", stacklevel=3)
1007 with self._metrics.instrument_put(_LOG, msg="Dataset put direct"):
1008 # If registry already has a dataset with the same dataset ID,
1009 # dataset type and DataId, then _importDatasets will do
1010 # nothing and just return an original ref. We have to raise in
1011 # this case, there is a datastore check below for that.
1012 self._registry._importDatasets([datasetRefOrType], expand=True)
1013 # Before trying to write to the datastore check that it does
1014 # not know this dataset. This is prone to races, of course.
1015 if self._datastore.knows(datasetRefOrType):
1016 raise ConflictingDefinitionError(
1017 f"Datastore already contains dataset: {datasetRefOrType}"
1018 )
1019 # Try to write dataset to the datastore, if it fails due to a
1020 # race with another write, the content of stored data may be
1021 # unpredictable.
1022 try:
1023 self._datastore.put(obj, datasetRefOrType, provenance=provenance)
1024 except IntegrityError as e:
1025 raise ConflictingDefinitionError(f"Datastore already contains dataset: {e}") from e
1027 return datasetRefOrType
1029 _LOG.debug("Butler put: %s, dataId=%s, run=%s", datasetRefOrType, dataId, run)
1030 if not self.isWriteable():
1031 raise TypeError("Butler is read-only.")
1033 with self._metrics.instrument_put(_LOG, msg="Dataset put with dataID"):
1034 datasetType, dataId = self._standardizeArgs(datasetRefOrType, dataId, **kwargs)
1036 # Handle dimension records in dataId
1037 dataId, kwargs = self._rewrite_data_id(dataId, datasetType, **kwargs)
1039 # Add Registry Dataset entry.
1040 dataId = self._registry.expandDataId(dataId, dimensions=datasetType.dimensions, **kwargs)
1041 (ref,) = self._registry.insertDatasets(datasetType, run=run, dataIds=[dataId])
1042 self._datastore.put(obj, ref, provenance=provenance)
1044 return ref
1046 def getDeferred(
1047 self,
1048 datasetRefOrType: DatasetRef | DatasetType | str,
1049 /,
1050 dataId: DataId | None = None,
1051 *,
1052 parameters: dict | None = None,
1053 collections: Any = None,
1054 storageClass: str | StorageClass | None = None,
1055 timespan: Timespan | None = None,
1056 **kwargs: Any,
1057 ) -> DeferredDatasetHandle:
1058 """Create a `DeferredDatasetHandle` which can later retrieve a dataset,
1059 after an immediate registry lookup.
1061 Parameters
1062 ----------
1063 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
1064 When `DatasetRef` the `dataId` should be `None`.
1065 Otherwise the `DatasetType` or name thereof.
1066 dataId : `dict` or `DataCoordinate`, optional
1067 A `dict` of `Dimension` link name, value pairs that label the
1068 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
1069 should be provided as the first argument.
1070 parameters : `dict`
1071 Additional StorageClass-defined options to control reading,
1072 typically used to efficiently read only a subset of the dataset.
1073 collections : Any, optional
1074 Collections to be searched, overriding ``self.collections``.
1075 Can be any of the types supported by the ``collections`` argument
1076 to butler construction.
1077 storageClass : `StorageClass` or `str`, optional
1078 The storage class to be used to override the Python type
1079 returned by this method. By default the returned type matches
1080 the dataset type definition for this dataset. Specifying a
1081 read `StorageClass` can force a different type to be returned.
1082 This type must be compatible with the original type.
1083 timespan : `Timespan` or `None`, optional
1084 A timespan that the validity range of the dataset must overlap.
1085 If not provided and this is a calibration dataset type, an attempt
1086 will be made to find the timespan from any temporal coordinate
1087 in the data ID.
1088 **kwargs
1089 Additional keyword arguments used to augment or construct a
1090 `DataId`. See `DataId` parameters.
1092 Returns
1093 -------
1094 obj : `DeferredDatasetHandle`
1095 A handle which can be used to retrieve a dataset at a later time.
1097 Raises
1098 ------
1099 LookupError
1100 Raised if no matching dataset exists in the `Registry` or
1101 datastore.
1102 ValueError
1103 Raised if a resolved `DatasetRef` was passed as an input, but it
1104 differs from the one found in the registry.
1105 TypeError
1106 Raised if no collections were provided.
1107 """
1108 if isinstance(datasetRefOrType, DatasetRef):
1109 # Do the quick check first and if that fails, check for artifact
1110 # existence. This is necessary for datastores that are configured
1111 # in trust mode where there won't be a record but there will be
1112 # a file.
1113 if self._datastore.knows(datasetRefOrType) or self._datastore.exists(datasetRefOrType):
1114 ref = datasetRefOrType
1115 else:
1116 raise LookupError(f"Dataset reference {datasetRefOrType} does not exist.")
1117 else:
1118 ref = self._findDatasetRef(
1119 datasetRefOrType, dataId, collections=collections, timespan=timespan, **kwargs
1120 )
1121 return DeferredDatasetHandle(butler=self, ref=ref, parameters=parameters, storageClass=storageClass)
1123 def get(
1124 self,
1125 datasetRefOrType: DatasetRef | DatasetType | str,
1126 /,
1127 dataId: DataId | None = None,
1128 *,
1129 parameters: dict[str, Any] | None = None,
1130 collections: Any = None,
1131 storageClass: StorageClass | str | None = None,
1132 timespan: Timespan | None = None,
1133 **kwargs: Any,
1134 ) -> Any:
1135 """Retrieve a stored dataset.
1137 Parameters
1138 ----------
1139 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
1140 When `DatasetRef` the `dataId` should be `None`.
1141 Otherwise the `DatasetType` or name thereof.
1142 If a resolved `DatasetRef`, the associated dataset
1143 is returned directly without additional querying.
1144 dataId : `dict` or `DataCoordinate`
1145 A `dict` of `Dimension` link name, value pairs that label the
1146 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
1147 should be provided as the first argument.
1148 parameters : `dict`
1149 Additional StorageClass-defined options to control reading,
1150 typically used to efficiently read only a subset of the dataset.
1151 collections : Any, optional
1152 Collections to be searched, overriding ``self.collections``.
1153 Can be any of the types supported by the ``collections`` argument
1154 to butler construction.
1155 storageClass : `StorageClass` or `str`, optional
1156 The storage class to be used to override the Python type
1157 returned by this method. By default the returned type matches
1158 the dataset type definition for this dataset. Specifying a
1159 read `StorageClass` can force a different type to be returned.
1160 This type must be compatible with the original type.
1161 timespan : `Timespan` or `None`, optional
1162 A timespan that the validity range of the dataset must overlap.
1163 If not provided and this is a calibration dataset type, an attempt
1164 will be made to find the timespan from any temporal coordinate
1165 in the data ID.
1166 **kwargs
1167 Additional keyword arguments used to augment or construct a
1168 `DataCoordinate`. See `DataCoordinate.standardize`
1169 parameters.
1171 Returns
1172 -------
1173 obj : `object`
1174 The dataset.
1176 Raises
1177 ------
1178 LookupError
1179 Raised if no matching dataset exists in the `Registry`.
1180 TypeError
1181 Raised if no collections were provided.
1183 Notes
1184 -----
1185 When looking up datasets in a `~CollectionType.CALIBRATION` collection,
1186 this method requires that the given data ID include temporal dimensions
1187 beyond the dimensions of the dataset type itself, in order to find the
1188 dataset with the appropriate validity range. For example, a "bias"
1189 dataset with native dimensions ``{instrument, detector}`` could be
1190 fetched with a ``{instrument, detector, exposure}`` data ID, because
1191 ``exposure`` is a temporal dimension.
1192 """
1193 _LOG.debug("Butler get: %s, dataId=%s, parameters=%s", datasetRefOrType, dataId, parameters)
1194 with self._metrics.instrument_get(_LOG, msg="Retrieved dataset"):
1195 ref = self._findDatasetRef(
1196 datasetRefOrType,
1197 dataId,
1198 collections=collections,
1199 datastore_records=True,
1200 timespan=timespan,
1201 **kwargs,
1202 )
1203 return self._datastore.get(ref, parameters=parameters, storageClass=storageClass)
1205 def getURIs(
1206 self,
1207 datasetRefOrType: DatasetRef | DatasetType | str,
1208 /,
1209 dataId: DataId | None = None,
1210 *,
1211 predict: bool = False,
1212 collections: Any = None,
1213 run: str | None = None,
1214 **kwargs: Any,
1215 ) -> DatasetRefURIs:
1216 """Return the URIs associated with the dataset.
1218 Parameters
1219 ----------
1220 datasetRefOrType : `DatasetRef`, `DatasetType`, or `str`
1221 When `DatasetRef` the `dataId` should be `None`.
1222 Otherwise the `DatasetType` or name thereof.
1223 dataId : `dict` or `DataCoordinate`
1224 A `dict` of `Dimension` link name, value pairs that label the
1225 `DatasetRef` within a Collection. When `None`, a `DatasetRef`
1226 should be provided as the first argument.
1227 predict : `bool`
1228 If `True`, allow URIs to be returned of datasets that have not
1229 been written.
1230 collections : Any, optional
1231 Collections to be searched, overriding ``self.collections``.
1232 Can be any of the types supported by the ``collections`` argument
1233 to butler construction.
1234 run : `str`, optional
1235 Run to use for predictions, overriding ``self.run``.
1236 **kwargs
1237 Additional keyword arguments used to augment or construct a
1238 `DataCoordinate`. See `DataCoordinate.standardize`
1239 parameters.
1241 Returns
1242 -------
1243 uris : `DatasetRefURIs`
1244 The URI to the primary artifact associated with this dataset (if
1245 the dataset was disassembled within the datastore this may be
1246 `None`), and the URIs to any components associated with the dataset
1247 artifact. (can be empty if there are no components).
1248 """
1249 ref = self._findDatasetRef(
1250 datasetRefOrType, dataId, predict=predict, run=run, collections=collections, **kwargs
1251 )
1252 return self._datastore.getURIs(ref, predict)
1254 def get_dataset_type(self, name: str) -> DatasetType:
1255 return self._registry.getDatasetType(name)
1257 def get_dataset(
1258 self,
1259 id: DatasetId | str,
1260 *,
1261 storage_class: str | StorageClass | None = None,
1262 dimension_records: bool = False,
1263 datastore_records: bool = False,
1264 ) -> DatasetRef | None:
1265 id = _to_uuid(id)
1266 ref = self._registry.getDataset(id)
1267 if ref is not None:
1268 if dimension_records:
1269 ref = ref.expanded(
1270 self._registry.expandDataId(ref.dataId, dimensions=ref.datasetType.dimensions)
1271 )
1272 if storage_class:
1273 ref = ref.overrideStorageClass(storage_class)
1274 if datastore_records:
1275 ref = self._registry.get_datastore_records(ref)
1276 return ref
1278 def get_many_datasets(self, ids: Iterable[DatasetId | str]) -> list[DatasetRef]:
1279 uuids = [_to_uuid(id) for id in ids]
1280 return self._registry._managers.datasets.get_dataset_refs(uuids)
1282 def find_dataset(
1283 self,
1284 dataset_type: DatasetType | str,
1285 data_id: DataId | None = None,
1286 *,
1287 collections: str | Sequence[str] | None = None,
1288 timespan: Timespan | None = None,
1289 storage_class: str | StorageClass | None = None,
1290 dimension_records: bool = False,
1291 datastore_records: bool = False,
1292 **kwargs: Any,
1293 ) -> DatasetRef | None:
1294 # Handle any parts of the dataID that are not using primary dimension
1295 # keys.
1296 if isinstance(dataset_type, str):
1297 actual_type = self.get_dataset_type(dataset_type)
1298 else:
1299 actual_type = dataset_type
1301 # Store the component for later.
1302 component_name = actual_type.component()
1303 if actual_type.isComponent():
1304 parent_type = actual_type.makeCompositeDatasetType()
1305 else:
1306 parent_type = actual_type
1308 data_id, kwargs = self._rewrite_data_id(data_id, parent_type, **kwargs)
1310 ref = self.registry.findDataset(
1311 parent_type,
1312 data_id,
1313 collections=collections,
1314 timespan=timespan,
1315 datastore_records=datastore_records,
1316 **kwargs,
1317 )
1318 if ref is not None and dimension_records:
1319 ref = ref.expanded(self._registry.expandDataId(ref.dataId, dimensions=ref.datasetType.dimensions))
1320 if ref is not None and component_name:
1321 ref = ref.makeComponentRef(component_name)
1322 if ref is not None and storage_class is not None:
1323 ref = ref.overrideStorageClass(storage_class)
1325 return ref
1327 def retrieve_artifacts_zip(
1328 self,
1329 refs: Iterable[DatasetRef],
1330 destination: ResourcePathExpression,
1331 overwrite: bool = True,
1332 ) -> ResourcePath:
1333 return retrieve_and_zip(refs, destination, self._datastore.retrieveArtifacts, overwrite)
1335 def retrieveArtifacts(
1336 self,
1337 refs: Iterable[DatasetRef],
1338 destination: ResourcePathExpression,
1339 transfer: str = "auto",
1340 preserve_path: bool = True,
1341 overwrite: bool = False,
1342 ) -> list[ResourcePath]:
1343 # Docstring inherited.
1344 outdir = ResourcePath(destination)
1345 artifact_map = self._datastore.retrieveArtifacts(
1346 refs,
1347 outdir,
1348 transfer=transfer,
1349 preserve_path=preserve_path,
1350 overwrite=overwrite,
1351 write_index=True,
1352 )
1353 return list(artifact_map)
1355 def exists(
1356 self,
1357 dataset_ref_or_type: DatasetRef | DatasetType | str,
1358 /,
1359 data_id: DataId | None = None,
1360 *,
1361 full_check: bool = True,
1362 collections: Any = None,
1363 **kwargs: Any,
1364 ) -> DatasetExistence:
1365 # Docstring inherited.
1366 existence = DatasetExistence.UNRECOGNIZED
1368 if isinstance(dataset_ref_or_type, DatasetRef):
1369 if collections is not None:
1370 warnings.warn("Collections should not be specified with DatasetRef", stacklevel=2)
1371 if data_id is not None:
1372 warnings.warn("A DataID should not be specified with DatasetRef", stacklevel=2)
1373 ref = dataset_ref_or_type
1374 registry_ref = self._registry.getDataset(dataset_ref_or_type.id)
1375 if registry_ref is not None:
1376 existence |= DatasetExistence.RECORDED
1378 if dataset_ref_or_type != registry_ref:
1379 # This could mean that storage classes differ, so we should
1380 # check for that but use the registry ref for the rest of
1381 # the method.
1382 if registry_ref.is_compatible_with(dataset_ref_or_type):
1383 # Use the registry version from now on.
1384 ref = registry_ref
1385 else:
1386 raise ValueError(
1387 f"The ref given to exists() ({ref}) has the same dataset ID as one "
1388 f"in registry but has different incompatible values ({registry_ref})."
1389 )
1390 else:
1391 try:
1392 ref = self._findDatasetRef(dataset_ref_or_type, data_id, collections=collections, **kwargs)
1393 except (LookupError, TypeError):
1394 return existence
1395 existence |= DatasetExistence.RECORDED
1397 if self._datastore.knows(ref):
1398 existence |= DatasetExistence.DATASTORE
1400 if full_check:
1401 if self._datastore.exists(ref):
1402 existence |= DatasetExistence._ARTIFACT
1403 elif existence.value != DatasetExistence.UNRECOGNIZED.value:
1404 # Do not add this flag if we have no other idea about a dataset.
1405 existence |= DatasetExistence(DatasetExistence._ASSUMED)
1407 return existence
1409 def _exists_many(
1410 self,
1411 refs: Iterable[DatasetRef],
1412 /,
1413 *,
1414 full_check: bool = True,
1415 ) -> dict[DatasetRef, DatasetExistence]:
1416 # Docstring inherited.
1417 existence = {ref: DatasetExistence.UNRECOGNIZED for ref in refs}
1419 # Check which refs exist in the registry.
1420 id_map = {ref.id: ref for ref in existence.keys()}
1421 for registry_ref in self.get_many_datasets(id_map.keys()):
1422 # Consistency between the given DatasetRef and the information
1423 # recorded in the registry is not verified.
1424 existence[id_map[registry_ref.id]] |= DatasetExistence.RECORDED
1426 # Ask datastore if it knows about these refs.
1427 knows = self._datastore.knows_these(refs)
1428 for ref, known in knows.items():
1429 if known:
1430 existence[ref] |= DatasetExistence.DATASTORE
1432 if full_check:
1433 mexists = self._datastore.mexists(refs)
1434 for ref, exists in mexists.items():
1435 if exists:
1436 existence[ref] |= DatasetExistence._ARTIFACT
1437 else:
1438 # Do not set this flag if nothing is known about the dataset.
1439 for ref in existence:
1440 if existence[ref] != DatasetExistence.UNRECOGNIZED:
1441 existence[ref] |= DatasetExistence._ASSUMED
1443 return existence
1445 def removeRuns(
1446 self,
1447 names: Iterable[str],
1448 unstore: bool | type[_DeprecatedDefault] = _DeprecatedDefault,
1449 *,
1450 unlink_from_chains: bool = False,
1451 ) -> None:
1452 # Docstring inherited.
1453 if not self.isWriteable():
1454 raise TypeError("Butler is read-only.")
1456 if unstore is not _DeprecatedDefault:
1457 # The value was passed in by a user. Must report it is now
1458 # ignored.
1459 if unstore is True:
1460 msg = "The unstore parameter is deprecated and is now always treated as True. "
1461 else:
1462 msg = "The unstore parameter for removeRuns can no longer be False and is now ignored. "
1463 warnings.warn(
1464 msg + " The parameter will be removed after v30.",
1465 category=FutureWarning,
1466 stacklevel=find_outside_stacklevel("lsst.daf.butler"),
1467 )
1469 names = list(names)
1470 refs: list[DatasetRef] = []
1471 # Map of the chained collections to the RUN children.
1472 parents_to_children: dict[str, set[str]] = defaultdict(set)
1474 with self._caching_context():
1475 # Get information about these RUNs.
1476 collections_info = self.collections.query_info(names, include_parents=unlink_from_chains)
1477 for info in collections_info:
1478 if info.type is not CollectionType.RUN:
1479 raise TypeError(f"The collection type of '{info.name}' is {info.type.name}, not RUN.")
1480 if unlink_from_chains:
1481 if info.parents is None: # For mypy.
1482 raise AssertionError("Internal error: Collection parents required but not received")
1483 for parent in info.parents:
1484 parents_to_children[parent].add(info.name)
1486 # Update the names in case the query unexpectedly had a wildcard.
1487 names = [info.name for info in collections_info]
1489 # Get all the datasets from these runs.
1490 refs = self.query_all_datasets(names, find_first=False, limit=None)
1492 # Call pruneDatasets since we are deliberately removing
1493 # datasets in chunks from the RUN collections rather than
1494 # attempting to remove everything at once.
1495 with time_this(
1496 _LOG,
1497 msg="Removing %d dataset%s from %s",
1498 args=(len(refs), "s" if len(refs) != 1 else "", ", ".join(names)),
1499 ):
1500 self.pruneDatasets(refs, unstore=True, purge=True, disassociate=True)
1502 # Now can remove the actual RUN collection and unlink from chains.
1503 with self._registry.transaction():
1504 # This will fail if caller is not unlinking from chains but the
1505 # RUN is in a chain -- but we have already deleted all the datasets
1506 # by this point.
1507 if unlink_from_chains:
1508 # Use deterministic order for deletions to attempt to minimize
1509 # risk of deadlocks for parallel deletes.
1510 for parent in sorted(parents_to_children):
1511 self.collections.remove_from_chain(parent, sorted(parents_to_children[parent]))
1512 # Sort to avoid potential deadlocks.
1513 for name in sorted(names):
1514 # This should be fast since the collection should be empty.
1515 with time_this(_LOG, msg="Removing RUN collection %s", args=(name,)):
1516 self._registry.removeCollection(name)
1517 _LOG.info("Completely removed the following RUN collections: %s", ", ".join(names))
1519 def pruneDatasets(
1520 self,
1521 refs: Iterable[DatasetRef],
1522 *,
1523 disassociate: bool = True,
1524 unstore: bool = False,
1525 tags: Iterable[str] = (),
1526 purge: bool = False,
1527 ) -> None:
1528 # docstring inherited from LimitedButler
1530 if not self.isWriteable():
1531 raise TypeError("Butler is read-only.")
1532 if purge:
1533 if not disassociate:
1534 raise TypeError("Cannot pass purge=True without disassociate=True.")
1535 if not unstore:
1536 raise TypeError("Cannot pass purge=True without unstore=True.")
1537 elif disassociate:
1538 tags = tuple(tags)
1539 if not tags:
1540 raise TypeError("No tags provided but disassociate=True.")
1541 for tag in tags:
1542 collectionType = self._registry.getCollectionType(tag)
1543 if collectionType is not CollectionType.TAGGED:
1544 raise TypeError(
1545 f"Cannot disassociate from collection '{tag}' "
1546 f"of non-TAGGED type {collectionType.name}."
1547 )
1548 # Transform possibly-single-pass iterable into something we can iterate
1549 # over multiple times.
1550 refs = list(refs)
1551 # Pruning a component of a DatasetRef makes no sense since registry
1552 # doesn't know about components and datastore might not store
1553 # components in a separate file
1554 for ref in refs:
1555 if ref.datasetType.component():
1556 raise ValueError(f"Can not prune a component of a dataset (ref={ref})")
1558 # Chunk the deletions using a size that is reasonably efficient whilst
1559 # also giving reasonable feedback to the user. Chunking also minimizes
1560 # what needs to rollback if there is a failure and should allow
1561 # incremental re-running of the pruning (assuming the query is
1562 # repeated). The only issue will be if the Ctrl-C comes during
1563 # emptyTrash since an admin command would need to run to finish the
1564 # emptying of that chunk.
1565 progress = Progress("lsst.daf.butler.Butler.pruneDatasets", level=_LOG.INFO)
1566 chunk_size = 50_000
1567 n_chunks = math.ceil(len(refs) / chunk_size)
1568 if n_chunks > 1:
1569 _LOG.verbose("Pruning a total of %d datasets", len(refs))
1570 chunk_num = 0
1571 for chunked_refs in progress.wrap(
1572 chunk_iterable(refs, chunk_size=chunk_size), desc="Deleting datasets", total=n_chunks
1573 ):
1574 chunk_num += 1
1575 _LOG.verbose(
1576 "Pruning %d dataset%s in chunk %d/%d",
1577 len(chunked_refs),
1578 "s" if len(chunked_refs) != 1 else "",
1579 chunk_num,
1580 n_chunks,
1581 )
1582 with time_this(
1583 _LOG,
1584 msg="Removing %d datasets for chunk %d/%d",
1585 args=(len(chunked_refs), chunk_num, n_chunks),
1586 ):
1587 self._prune_datasets(
1588 chunked_refs, tags=tags, unstore=unstore, purge=purge, disassociate=disassociate
1589 )
1591 def _prune_datasets(
1592 self,
1593 refs: Collection[DatasetRef],
1594 *,
1595 disassociate: bool = True,
1596 unstore: bool = False,
1597 tags: Iterable[str] = (),
1598 purge: bool = False,
1599 ) -> None:
1600 # We don't need an unreliable Datastore transaction for this, because
1601 # we've been extra careful to ensure that Datastore.trash only involves
1602 # mutating the Registry (it can _look_ at Datastore-specific things,
1603 # but shouldn't change them), and hence all operations here are
1604 # Registry operations.
1605 with self.transaction():
1606 plural = "s" if len(refs) != 1 else ""
1607 if unstore:
1608 with time_this(
1609 _LOG,
1610 msg="Marking %d dataset%s for removal during pruneDatasets",
1611 args=(len(refs), plural),
1612 ):
1613 self._datastore.trash(refs)
1614 if purge:
1615 with time_this(
1616 _LOG, msg="Removing %d pruned dataset%s from registry", args=(len(refs), plural)
1617 ):
1618 self._registry.removeDatasets(refs)
1619 elif disassociate:
1620 assert tags, "Guaranteed by earlier logic in this function."
1621 with time_this(
1622 _LOG, msg="Disassociating %d dataset%ss from tagged collections", args=(len(refs), plural)
1623 ):
1624 for tag in tags:
1625 self._registry.disassociate(tag, refs)
1626 # We've exited the Registry transaction, and apparently committed.
1627 # (if there was an exception, everything rolled back, and it's as if
1628 # nothing happened - and we never get here).
1629 # Datastore artifacts are not yet gone, but they're clearly marked
1630 # as trash, so if we fail to delete now because of (e.g.) filesystem
1631 # problems we can try again later, and if manual administrative
1632 # intervention is required, it's pretty clear what that should entail:
1633 # deleting everything on disk and in private Datastore tables that is
1634 # in the dataset_location_trash table.
1635 if unstore:
1636 # Point of no return for removing artifacts. Restrict the trash
1637 # emptying to the refs that this call trashed.
1638 with time_this(
1639 _LOG,
1640 msg="Attempting to remove artifacts for %d dataset%s associated with pruning",
1641 args=(len(refs), plural),
1642 ):
1643 self._datastore.emptyTrash(refs=refs)
1645 def ingest_zip(
1646 self,
1647 zip_file: ResourcePathExpression,
1648 transfer: str = "auto",
1649 *,
1650 transfer_dimensions: bool = False,
1651 dry_run: bool = False,
1652 ) -> None:
1653 """Ingest a Zip file into this butler.
1655 The Zip file must have been created by `retrieve_artifacts_zip`.
1657 Parameters
1658 ----------
1659 zip_file : `lsst.resources.ResourcePathExpression`
1660 Path to the Zip file.
1661 transfer : `str`, optional
1662 Method to use to transfer the Zip into the datastore.
1663 transfer_dimensions : `bool`, optional
1664 If `True`, dimension record data associated with the new datasets
1665 will be transferred from the Zip, if present.
1666 dry_run : `bool`, optional
1667 If `True` the ingest will be processed without any modifications
1668 made to the target butler and as if the target butler did not
1669 have any of the datasets.
1671 Notes
1672 -----
1673 Run collections and dataset types are created as needed.
1674 """
1675 if not self.isWriteable():
1676 raise TypeError("Butler is read-only.")
1678 zip_path = ResourcePath(zip_file)
1679 index = ZipIndex.from_zip_file(zip_path)
1680 _LOG.verbose(
1681 "Ingesting %s containing %d datasets and %d files.", zip_path, len(index.refs), len(index)
1682 )
1684 # Need to ingest the refs into registry. Re-use the standard ingest
1685 # code by reconstructing FileDataset from the index.
1686 refs = index.refs.to_refs(universe=self.dimensions)
1687 id_to_ref = {ref.id: ref for ref in refs}
1688 datasets: list[FileDataset] = []
1689 processed_ids: set[uuid.UUID] = set()
1690 for path_in_zip, index_info in index.artifact_map.items():
1691 # Disassembled composites need to check this ref isn't already
1692 # included.
1693 unprocessed = {id_ for id_ in index_info.ids if id_ not in processed_ids}
1694 if not unprocessed:
1695 continue
1696 dataset = FileDataset(refs=[id_to_ref[id_] for id_ in unprocessed], path=path_in_zip)
1697 datasets.append(dataset)
1698 processed_ids.update(unprocessed)
1700 # Ingest doesn't create the RUN collections so we have to do that
1701 # here.
1702 #
1703 # Sort by run collection name to ensure Postgres takes locks in the
1704 # same order between different processes, to mitigate an issue
1705 # where Postgres can deadlock due to the unique index on collection
1706 # name. (See DM-47543).
1707 runs = {ref.run for ref in refs}
1708 for run in sorted(runs):
1709 registered = self.collections.register(run)
1710 if registered:
1711 _LOG.verbose("Created RUN collection %s as part of zip ingest", run)
1713 progress = Progress("lsst.daf.butler.Butler.ingest", level=VERBOSE)
1714 import_info = self._prepare_ingest_file_datasets(
1715 datasets, progress, dry_run=dry_run, transfer_dimensions=transfer_dimensions
1716 )
1718 with self.transaction():
1719 # Do not need expanded dataset refs so can ignore the return value.
1720 self._ingest_file_datasets(datasets, import_info, progress, dry_run=dry_run)
1722 try:
1723 self._datastore.ingest_zip(zip_path, transfer=transfer, dry_run=dry_run)
1724 except IntegrityError as e:
1725 raise ConflictingDefinitionError(
1726 f"Datastore already contains one or more datasets: {e}"
1727 ) from e
1729 def _prepare_ingest_file_datasets(
1730 self,
1731 datasets: Sequence[FileDataset],
1732 progress: Progress,
1733 *,
1734 transfer_dimensions: bool = False,
1735 dry_run: bool = False,
1736 ) -> _ImportDatasetsInfo:
1737 # Track DataIDs that are being ingested so we can spot issues early
1738 # with duplication. Retain previous FileDataset so we can report it.
1739 groupedDataIds: MutableMapping[tuple[DatasetType, str], dict[DataCoordinate, FileDataset]] = (
1740 defaultdict(dict)
1741 )
1743 # All the refs we need to import.
1744 refs: list[DatasetRef] = []
1746 for dataset in progress.wrap(datasets, desc="Validating dataIDs"):
1747 for ref in dataset.refs:
1748 group_key = (ref.datasetType, ref.run)
1750 if ref.dataId in groupedDataIds[group_key]:
1751 raise ConflictingDefinitionError(
1752 f"Ingest conflict. Dataset {dataset.path} has same"
1753 " DataId as other ingest dataset"
1754 f" {groupedDataIds[group_key][ref.dataId].path} "
1755 f" ({ref.dataId})"
1756 )
1758 groupedDataIds[group_key][ref.dataId] = dataset
1759 refs.extend(dataset.refs)
1761 # Ensure that dataset types are created and all ref information
1762 # extracted.
1763 import_info = self._prepare_for_import_refs(
1764 self,
1765 refs,
1766 register_dataset_types=True,
1767 dry_run=dry_run,
1768 transfer_dimensions=transfer_dimensions,
1769 )
1770 return import_info
1772 def _ingest_file_datasets(
1773 self,
1774 datasets: Sequence[FileDataset],
1775 import_info: _ImportDatasetsInfo,
1776 progress: Progress,
1777 *,
1778 dry_run: bool = False,
1779 ) -> None:
1780 self._import_dimension_records(import_info.dimension_records, dry_run=dry_run)
1781 imported_refs = self._import_grouped_refs(
1782 import_info.grouped_refs, None, progress, dry_run=dry_run, expand_refs=True
1783 )
1785 # The expanded refs need to be attached back to the original
1786 # FileDatasets for datastore to use.
1787 id_to_ref = {ref.id: ref for ref in imported_refs}
1789 for dataset in progress.wrap(datasets, desc="Re-attaching expanded refs"):
1790 dataset.refs = [id_to_ref[ref.id] for ref in dataset.refs]
1792 def ingest(
1793 self,
1794 *datasets: FileDataset,
1795 transfer: str | None = "auto",
1796 record_validation_info: bool = True,
1797 skip_existing: bool = False,
1798 ) -> None:
1799 # Docstring inherited.
1800 if not datasets:
1801 return
1802 if not self.isWriteable():
1803 raise TypeError("Butler is read-only.")
1804 _LOG.verbose("Ingesting %d file dataset%s.", len(datasets), "" if len(datasets) == 1 else "s")
1805 progress = Progress("lsst.daf.butler.Butler.ingest", level=VERBOSE)
1807 new_datasets, existing_datasets = self._partition_datasets_by_known(datasets)
1808 if existing_datasets:
1809 if skip_existing:
1810 _LOG.info(
1811 "Skipping %d datasets which already exist in the repository.", len(existing_datasets)
1812 )
1813 else:
1814 raise ConflictingDefinitionError(
1815 f"Datastore already contains {len(existing_datasets)} of the given datasets."
1816 f" Example: {existing_datasets[0]}"
1817 )
1819 # We use `datasets` rather `new_datasets` for the Registry
1820 # portion of this, to let it confirm that everything matches the
1821 # existing datasets.
1822 import_info = self._prepare_ingest_file_datasets(datasets, progress)
1824 with self.transaction():
1825 self._ingest_file_datasets(datasets, import_info, progress)
1827 # Bulk-insert everything into Datastore.
1828 # We do not know if any of the registry entries already existed
1829 # (_importDatasets only complains if they exist but differ).
1830 # The _partition_datasets_by_known logic above should catch most
1831 # instances where we attempt to re-ingest files that were already
1832 # ingested, but a concurrent writer could cause a unique constraint
1833 # violation here.
1834 try:
1835 self._datastore.ingest(
1836 *new_datasets, transfer=transfer, record_validation_info=record_validation_info
1837 )
1838 except IntegrityError as e:
1839 raise ConflictingDefinitionError(
1840 f"Datastore already contains one or more datasets: {e}"
1841 ) from e
1843 def _partition_datasets_by_known(
1844 self, datasets: Iterable[FileDataset]
1845 ) -> tuple[list[FileDataset], list[FileDataset]]:
1846 """Divides the given `FileDataset` objects into two groups: those for
1847 which the Datastore already has an entry, and those for which it does
1848 not.
1849 """
1850 new_datasets = []
1851 existing_datasets = []
1853 refs = itertools.chain.from_iterable(dataset.refs for dataset in datasets)
1854 known_refs = self._datastore.knows_these(refs)
1856 for dataset in datasets:
1857 if any(known_refs[ref] for ref in dataset.refs):
1858 existing_datasets.append(dataset)
1859 else:
1860 new_datasets.append(dataset)
1862 return new_datasets, existing_datasets
1864 @contextlib.contextmanager
1865 def export(
1866 self,
1867 *,
1868 directory: str | None = None,
1869 filename: str | None = None,
1870 format: str | None = None,
1871 transfer: str | None = None,
1872 ) -> Iterator[RepoExportContext]:
1873 # Docstring inherited.
1874 if directory is None and transfer is not None:
1875 raise TypeError("Cannot transfer without providing a directory.")
1876 if transfer == "move":
1877 raise TypeError("Transfer may not be 'move': export is read-only")
1878 if format is None:
1879 if filename is None:
1880 raise TypeError("At least one of 'filename' or 'format' must be provided.")
1881 else:
1882 _, format = os.path.splitext(filename)
1883 if not format:
1884 raise ValueError("Please specify a file extension to determine export format.")
1885 format = format[1:] # Strip leading ".""
1886 elif filename is None:
1887 filename = f"export.{format}"
1888 if directory is not None:
1889 filename = os.path.join(directory, filename)
1890 formats = self._config["repo_transfer_formats"]
1891 if format not in formats:
1892 raise ValueError(f"Unknown export format {format!r}, allowed: {','.join(formats.keys())}")
1893 BackendClass = get_class_of(formats[format, "export"])
1894 with open(filename, "w") as stream:
1895 backend = BackendClass(stream, universe=self.dimensions)
1896 try:
1897 helper = RepoExportContext(self, backend=backend, directory=directory, transfer=transfer)
1898 with self._caching_context():
1899 yield helper
1900 except BaseException:
1901 raise
1902 else:
1903 helper._finish()
1905 def import_(
1906 self,
1907 *,
1908 directory: ResourcePathExpression | None = None,
1909 filename: ResourcePathExpression | TextIO | None = None,
1910 format: str | None = None,
1911 transfer: str | None = None,
1912 skip_dimensions: set | None = None,
1913 record_validation_info: bool = True,
1914 without_datastore: bool = False,
1915 ) -> None:
1916 # Docstring inherited.
1917 if not self.isWriteable():
1918 raise TypeError("Butler is read-only.")
1919 if filename is None and format is not None:
1920 filename = ResourcePath(f"export.{format}", forceAbsolute=False)
1921 if directory is not None:
1922 directory = ResourcePath(directory, forceDirectory=True)
1923 # mypy doesn't think this will work but it does in python >= 3.10.
1924 if isinstance(filename, ResourcePathExpression): # type: ignore
1925 filename = ResourcePath(filename, forceAbsolute=False) # type: ignore
1926 if format is None:
1927 format = filename.getExtension()
1928 if not filename.isabs() and directory is not None:
1929 potential = directory.join(filename)
1930 exists_in_cwd = filename.exists()
1931 exists_in_dir = potential.exists()
1932 if exists_in_cwd and exists_in_dir:
1933 _LOG.warning(
1934 "A relative path for filename was specified (%s) which exists relative to cwd. "
1935 "Additionally, the file exists relative to the given search directory (%s). "
1936 "Using the export file in the given directory.",
1937 filename,
1938 potential,
1939 )
1940 # Given they specified an explicit directory and that
1941 # directory has the export file in it, assume that that
1942 # is what was meant despite the file in cwd.
1943 filename = potential
1944 elif exists_in_dir:
1945 filename = potential
1946 elif not exists_in_cwd and not exists_in_dir:
1947 # Raise early.
1948 raise FileNotFoundError(
1949 f"Export file could not be found in {filename.abspath()} or {potential.abspath()}."
1950 )
1951 elif format is None:
1952 format = ".yaml"
1953 BackendClass: type[RepoImportBackend] = get_class_of(
1954 self._config["repo_transfer_formats"][format]["import"]
1955 )
1957 def doImport(importStream: TextIO | ResourceHandleProtocol) -> None:
1958 with self._caching_context():
1959 backend = BackendClass(importStream, self) # type: ignore[call-arg]
1960 backend.register()
1961 with self.transaction():
1962 backend.load(
1963 datastore=self._datastore if not without_datastore else None,
1964 directory=directory,
1965 transfer=transfer,
1966 skip_dimensions=skip_dimensions,
1967 record_validation_info=record_validation_info,
1968 )
1970 if isinstance(filename, ResourcePath):
1971 # We can not use open() here at the moment because of
1972 # DM-38589 since yaml does stream.read(8192) in a loop.
1973 stream = io.StringIO(filename.read().decode())
1974 doImport(stream)
1975 else:
1976 doImport(filename) # type: ignore
1978 def transfer_dimension_records_from(
1979 self, source_butler: LimitedButler | Butler, source_refs: Iterable[DatasetRef]
1980 ) -> None:
1981 # Allowed dimensions in the target butler.
1982 elements = frozenset(element for element in self.dimensions.elements if element.has_own_table)
1984 data_ids = {ref.dataId for ref in source_refs}
1986 dimension_records = self._extract_all_dimension_records_from_data_ids(
1987 source_butler, data_ids, elements
1988 )
1990 # Insert order is important.
1991 for element in self.dimensions.sorted(dimension_records.keys()):
1992 records = [r for r in dimension_records[element].values()]
1993 # Assume that if the record is already present that we can
1994 # use it without having to check that the record metadata
1995 # is consistent.
1996 self._registry.insertDimensionData(element, *records, skip_existing=True)
1997 _LOG.debug("Dimension '%s' -- number of records transferred: %d", element.name, len(records))
1999 def _extract_all_dimension_records_from_data_ids(
2000 self,
2001 source_butler: LimitedButler | Butler,
2002 data_ids: set[DataCoordinate],
2003 allowed_elements: frozenset[DimensionElement],
2004 ) -> dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]:
2005 primary_records = self._extract_dimension_records_from_data_ids(
2006 source_butler, data_ids, allowed_elements
2007 )
2009 can_query = True if isinstance(source_butler, Butler) else False
2011 additional_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2012 for original_element, record_mapping in primary_records.items():
2013 # Get dimensions that depend on this dimension.
2014 populated_by = self.dimensions.get_elements_populated_by(
2015 self.dimensions[original_element.name] # type: ignore
2016 )
2018 for data_id in record_mapping.keys():
2019 for element in populated_by:
2020 if element not in allowed_elements:
2021 continue
2022 if element.name == original_element.name:
2023 continue
2025 if element.name in primary_records:
2026 # If this element has already been stored avoid
2027 # re-finding records since that may lead to additional
2028 # spurious records. e.g. visit is populated_by
2029 # visit_detector_region but querying
2030 # visit_detector_region by visit will return all the
2031 # detectors for this visit -- the visit dataId does not
2032 # constrain this.
2033 # To constrain the query the original dataIds would
2034 # have to be scanned.
2035 continue
2037 if not can_query:
2038 raise RuntimeError(
2039 f"Transferring populated_by records like {element.name} requires a full Butler."
2040 )
2042 records = source_butler.query_dimension_records( # type: ignore
2043 element.name,
2044 explain=False,
2045 **data_id.mapping, # type: ignore
2046 )
2047 for record in records:
2048 additional_records[record.definition].setdefault(record.dataId, record)
2050 # The next step is to walk back through the additional records to
2051 # pick up any missing content (such as visit_definition needing to
2052 # know the exposure). Want to ensure we do not request records we
2053 # already have.
2054 missing_data_ids = set()
2055 for name, record_mapping in additional_records.items():
2056 for data_id in record_mapping.keys():
2057 if data_id not in primary_records[name]:
2058 missing_data_ids.add(data_id)
2060 # Fill out the new records. Assume that these new records do not
2061 # also need to carry over additional populated_by records.
2062 secondary_records = self._extract_dimension_records_from_data_ids(
2063 source_butler, missing_data_ids, allowed_elements
2064 )
2066 # Merge the extra sets of records in with the original.
2067 for name, record_mapping in itertools.chain(additional_records.items(), secondary_records.items()):
2068 primary_records[name].update(record_mapping)
2070 return primary_records
2072 def _extract_dimension_records_from_data_ids(
2073 self,
2074 source_butler: LimitedButler | Butler,
2075 data_ids: set[DataCoordinate],
2076 allowed_elements: frozenset[DimensionElement],
2077 ) -> dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]:
2078 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2080 for data_id in data_ids:
2081 # Need an expanded record, if not expanded that we need a full
2082 # butler with registry (allow mocks with registry too).
2083 if not data_id.hasRecords():
2084 if registry := getattr(source_butler, "registry", None):
2085 data_id = registry.expandDataId(data_id)
2086 else:
2087 raise TypeError("Input butler needs to be a full butler to expand DataId.")
2088 # If this butler doesn't know about a dimension in the source
2089 # butler things will break later.
2090 for element_name in data_id.dimensions.elements:
2091 record = data_id.records[element_name]
2092 if record is not None and record.definition in allowed_elements:
2093 dimension_records[record.definition].setdefault(record.dataId, record)
2095 return dimension_records
2097 def _prepare_for_import_refs(
2098 self,
2099 source_butler: LimitedButler,
2100 source_refs: Iterable[DatasetRef],
2101 *,
2102 register_dataset_types: bool = False,
2103 transfer_dimensions: bool = False,
2104 dry_run: bool = False,
2105 ) -> _ImportDatasetsInfo:
2106 # Docstring inherited.
2107 if not self.isWriteable() and not dry_run:
2108 raise TypeError("Butler is read-only.")
2110 # Will iterate through the refs multiple times so need to convert
2111 # to a list if this isn't a collection.
2112 if not isinstance(source_refs, collections.abc.Collection):
2113 source_refs = list(source_refs)
2115 original_count = len(source_refs)
2116 log_level = _LOG.INFO if original_count > 1 else _LOG.VERBOSE
2117 _LOG.log(
2118 log_level,
2119 "Importing %d dataset%s into %s",
2120 original_count,
2121 "s" if original_count != 1 else "",
2122 str(self),
2123 )
2125 # Importing requires that we group the refs by dimension group and run
2126 # before doing the import.
2127 source_dataset_types = set()
2128 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]] = defaultdict(list)
2129 for ref in source_refs:
2130 grouped_refs[_RefGroup(ref.datasetType.dimensions, ref.run)].append(ref)
2131 source_dataset_types.add(ref.datasetType)
2133 # Check to see if the dataset type in the source butler has
2134 # the same definition in the target butler and register missing
2135 # ones if requested. Registration must happen outside a transaction.
2136 newly_registered_dataset_types = set()
2137 for datasetType in source_dataset_types:
2138 if register_dataset_types:
2139 # Let this raise immediately if inconsistent. Continuing
2140 # on to find additional inconsistent dataset types
2141 # might result in additional unwanted dataset types being
2142 # registered.
2143 try:
2144 if not dry_run and self._registry.registerDatasetType(datasetType):
2145 newly_registered_dataset_types.add(datasetType)
2146 except ConflictingDefinitionError as e:
2147 # Be safe and require that conversions be bidirectional
2148 # when there are storage class mismatches. This is because
2149 # get() will have to support conversion from source to
2150 # target python type (the source formatter will be
2151 # returning source python type) but there also is an
2152 # expectation that people will want to be able to get() in
2153 # the target using the source python type, which will not
2154 # require conversion for transferred datasets but might
2155 # for target-native types. Additionally, butler.get does
2156 # not know that the formatter will return the wrong
2157 # python type and so will always check that the conversion
2158 # works even though it won't need it.
2159 target_dataset_type = self.get_dataset_type(datasetType.name)
2160 target_compatible_with_source = target_dataset_type.is_compatible_with(datasetType)
2161 source_compatible_with_target = datasetType.is_compatible_with(target_dataset_type)
2162 if not (target_compatible_with_source and source_compatible_with_target):
2163 if target_compatible_with_source:
2164 e.add_note(
2165 "Target dataset type storage class is compatible with source "
2166 "but the reverse is not true."
2167 )
2168 elif source_compatible_with_target:
2169 e.add_note(
2170 "Source dataset type storage class is compatible with target "
2171 "but the reverse is not true."
2172 )
2173 else:
2174 e.add_note("If storage classes differ, please register converters.")
2175 raise
2176 else:
2177 # If the dataset type is missing, let it fail immediately.
2178 target_dataset_type = self.get_dataset_type(datasetType.name)
2179 if target_dataset_type != datasetType:
2180 target_compatible_with_source = target_dataset_type.is_compatible_with(datasetType)
2181 source_compatible_with_target = datasetType.is_compatible_with(target_dataset_type)
2182 # Both conversion directions are currently required.
2183 if not (target_compatible_with_source and source_compatible_with_target):
2184 msg = ""
2185 if target_compatible_with_source:
2186 msg = (
2187 "Target storage class is compatible with the source storage class "
2188 "but the reverse is not true."
2189 )
2190 elif source_compatible_with_target:
2191 msg = (
2192 "Source storage class is compatible with the target storage class"
2193 " but the reverse is not true."
2194 )
2195 else:
2196 msg = "If storage classes differ register converters."
2197 if msg:
2198 msg = f"({msg})"
2199 raise ConflictingDefinitionError(
2200 "Source butler dataset type differs from definition"
2201 f" in target butler: {datasetType} !="
2202 f" {target_dataset_type} {msg}"
2203 )
2204 if newly_registered_dataset_types:
2205 # We may have registered some even if there were inconsistencies
2206 # but should let people know (or else remove them again).
2207 _LOG.verbose(
2208 "Registered the following dataset types in the target Butler: %s",
2209 ", ".join(d.name for d in newly_registered_dataset_types),
2210 )
2211 else:
2212 _LOG.verbose("All required dataset types are known to the target Butler")
2214 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2215 if transfer_dimensions:
2216 # Collect all the dimension records for these refs.
2217 # All dimensions are to be copied but the list of valid dimensions
2218 # come from this butler's universe.
2219 elements = frozenset(element for element in self.dimensions.elements if element.has_own_table)
2220 dataIds = {ref.dataId for ref in source_refs}
2221 dimension_records = self._extract_all_dimension_records_from_data_ids(
2222 source_butler, dataIds, elements
2223 )
2224 return _ImportDatasetsInfo(grouped_refs, dimension_records)
2226 def _import_dimension_records(
2227 self,
2228 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]],
2229 *,
2230 dry_run: bool,
2231 ) -> None:
2232 """Import dimension records collected during import pre-process."""
2233 if dimension_records and not dry_run:
2234 _LOG.verbose("Ensuring that dimension records exist for transferred datasets.")
2235 # Order matters.
2236 for element in self.dimensions.sorted(dimension_records.keys()):
2237 records = list(dimension_records[element].values())
2238 # Assume that if the record is already present that we can
2239 # use it without having to check that the record metadata
2240 # is consistent.
2241 self._registry.insertDimensionData(element, *records, skip_existing=True)
2243 def _import_grouped_refs(
2244 self,
2245 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]],
2246 source_butler: LimitedButler | None,
2247 progress: Progress,
2248 *,
2249 dry_run: bool = False,
2250 expand_refs: bool = False,
2251 ) -> list[DatasetRef]:
2252 handled_collections: set[str] = set()
2253 n_to_import = 0
2254 all_imported_refs: list[DatasetRef] = []
2255 # Sort by run collection name to ensure Postgres takes locks in the
2256 # same order between different processes, to mitigate an issue
2257 # where Postgres can deadlock due to the unique index on collection
2258 # name. (See DM-47543).
2259 groups = sorted(grouped_refs.items(), key=lambda item: item[0].run)
2260 for (dimension_group, run), refs_to_import in progress.iter_item_chunks(
2261 groups, desc="Importing to registry by run and dataset type"
2262 ):
2263 if run not in handled_collections:
2264 # May need to create output collection. If source butler
2265 # has a registry, ask for documentation string.
2266 run_doc = None
2267 if source_butler is not None and (registry := getattr(source_butler, "registry", None)):
2268 run_doc = registry.getCollectionDocumentation(run)
2269 if not dry_run:
2270 registered = self.collections.register(run, doc=run_doc)
2271 else:
2272 registered = True
2273 handled_collections.add(run)
2274 if registered:
2275 _LOG.verbose("Creating output run %s", run)
2277 n_refs = len(refs_to_import)
2278 n_to_import += n_refs
2279 _LOG.verbose(
2280 "Importing %d ref%s with dimensions %s into run %s",
2281 n_refs,
2282 "" if n_refs == 1 else "s",
2283 dimension_group.names,
2284 run,
2285 )
2287 # Assume we are using UUIDs and the source refs will match
2288 # those imported.
2289 if not dry_run:
2290 imported_refs = self._registry._importDatasets(refs_to_import, expand=expand_refs)
2291 else:
2292 imported_refs = refs_to_import
2294 all_imported_refs.extend(imported_refs)
2296 assert n_to_import == len(all_imported_refs)
2297 _LOG.verbose("Imported %d datasets into destination butler", n_to_import)
2298 return all_imported_refs
2300 def transfer_from(
2301 self,
2302 source_butler: LimitedButler,
2303 source_refs: Iterable[DatasetRef],
2304 transfer: str = "auto",
2305 skip_missing: bool = True,
2306 register_dataset_types: bool = False,
2307 transfer_dimensions: bool = False,
2308 dry_run: bool = False,
2309 ) -> collections.abc.Collection[DatasetRef]:
2310 # Docstring inherited.
2311 source_refs = list(source_refs)
2312 if not self.isWriteable() and not dry_run:
2313 raise TypeError("Butler is read-only.")
2315 progress = Progress("lsst.daf.butler.Butler.transfer_from", level=VERBOSE)
2317 artifact_existence: dict[ResourcePath, bool] = {}
2318 file_transfer_source = source_butler._file_transfer_source
2319 transfer_records = retrieve_file_transfer_records(
2320 file_transfer_source, source_refs, artifact_existence
2321 )
2322 # In some situations the datastore artifact may be missing and we do
2323 # not want that registry entry to be imported. For example, this can
2324 # happen if a file was removed but the dataset was left in the registry
2325 # for provenance, or if a pipeline task didn't create all of the
2326 # possible files in a QuantumBackedButler.
2327 if skip_missing:
2328 original_ids = {ref.id for ref in source_refs}
2329 missing_ids = original_ids - transfer_records.keys()
2330 if missing_ids:
2331 original_count = len(source_refs)
2332 source_refs = [ref for ref in source_refs if ref.id not in missing_ids]
2333 filtered_count = len(source_refs)
2334 n_missing = original_count - filtered_count
2335 _LOG.verbose(
2336 "%d dataset%s removed because the artifact does not exist. Now have %d.",
2337 n_missing,
2338 "" if n_missing == 1 else "s",
2339 filtered_count,
2340 )
2342 import_info = self._prepare_for_import_refs(
2343 source_butler,
2344 source_refs,
2345 register_dataset_types=register_dataset_types,
2346 dry_run=dry_run,
2347 transfer_dimensions=transfer_dimensions,
2348 )
2350 # Do all the importing in a single transaction.
2351 with self.transaction():
2352 self._import_dimension_records(import_info.dimension_records, dry_run=dry_run)
2353 imported_refs = self._import_grouped_refs(
2354 import_info.grouped_refs, source_butler, progress, dry_run=dry_run
2355 )
2357 # Ask the datastore to transfer. The datastore has to check that
2358 # the source datastore is compatible with the target datastore.
2359 _LOG.verbose("Transferring %d datasets from %s", len(transfer_records), file_transfer_source.name)
2360 accepted, rejected = self._datastore.transfer_from(
2361 transfer_records,
2362 imported_refs,
2363 transfer=transfer,
2364 artifact_existence=artifact_existence,
2365 dry_run=dry_run,
2366 )
2367 if rejected:
2368 # For now, accept the registry entries but not the files.
2369 _LOG.warning(
2370 "%d datasets were rejected and %d accepted for transfer.",
2371 len(rejected),
2372 len(accepted),
2373 )
2375 return imported_refs
2377 def validateConfiguration(
2378 self,
2379 logFailures: bool = False,
2380 datasetTypeNames: Iterable[str] | None = None,
2381 ignore: Iterable[str] | None = None,
2382 ) -> None:
2383 # Docstring inherited.
2384 if datasetTypeNames:
2385 datasetTypes = [self.get_dataset_type(name) for name in datasetTypeNames]
2386 else:
2387 datasetTypes = list(self._registry.queryDatasetTypes())
2389 # filter out anything from the ignore list
2390 if ignore:
2391 ignore = set(ignore)
2392 datasetTypes = [
2393 e for e in datasetTypes if e.name not in ignore and e.nameAndComponent()[0] not in ignore
2394 ]
2395 else:
2396 ignore = set()
2398 # For each datasetType that has an instrument dimension, create
2399 # a DatasetRef for each defined instrument
2400 datasetRefs = []
2402 # Find all the registered instruments (if "instrument" is in the
2403 # universe).
2404 instruments: set[str] = set()
2405 if "instrument" in self.dimensions:
2406 instruments = {rec.name for rec in self.query_dimension_records("instrument", explain=False)}
2408 for datasetType in datasetTypes:
2409 if "instrument" in datasetType.dimensions:
2410 # In order to create a conforming dataset ref, create
2411 # fake DataCoordinate values for the non-instrument
2412 # dimensions. The type of the value does not matter here.
2413 dataId = {dim: 1 for dim in datasetType.dimensions.names if dim != "instrument"}
2415 for instrument in instruments:
2416 datasetRef = DatasetRef(
2417 datasetType,
2418 DataCoordinate.standardize(
2419 dataId, instrument=instrument, dimensions=datasetType.dimensions
2420 ),
2421 run="validate",
2422 )
2423 datasetRefs.append(datasetRef)
2425 entities: list[DatasetType | DatasetRef] = []
2426 entities.extend(datasetTypes)
2427 entities.extend(datasetRefs)
2429 datastoreErrorStr = None
2430 try:
2431 self._datastore.validateConfiguration(entities, logFailures=logFailures)
2432 except ValidationError as e:
2433 datastoreErrorStr = str(e)
2435 # Also check that the LookupKeys used by the datastores match
2436 # registry and storage class definitions
2437 keys = self._datastore.getLookupKeys()
2439 failedNames = set()
2440 failedDataId = set()
2441 for key in keys:
2442 if key.name is not None:
2443 if key.name in ignore:
2444 continue
2446 # skip if specific datasetType names were requested and this
2447 # name does not match
2448 if datasetTypeNames and key.name not in datasetTypeNames:
2449 continue
2451 # See if it is a StorageClass or a DatasetType
2452 if key.name in self.storageClasses:
2453 pass
2454 else:
2455 try:
2456 self.get_dataset_type(key.name)
2457 except KeyError:
2458 if logFailures:
2459 _LOG.critical(
2460 "Key '%s' does not correspond to a DatasetType or StorageClass", key
2461 )
2462 failedNames.add(key)
2463 else:
2464 # Dimensions are checked for consistency when the Butler
2465 # is created and rendezvoused with a universe.
2466 pass
2468 # Check that the instrument is a valid instrument
2469 # Currently only support instrument so check for that
2470 if key.dataId:
2471 dataIdKeys = set(key.dataId)
2472 if {"instrument"} != dataIdKeys:
2473 if logFailures:
2474 _LOG.critical("Key '%s' has unsupported DataId override", key)
2475 failedDataId.add(key)
2476 elif key.dataId["instrument"] not in instruments:
2477 if logFailures:
2478 _LOG.critical("Key '%s' has unknown instrument", key)
2479 failedDataId.add(key)
2481 messages = []
2483 if datastoreErrorStr:
2484 messages.append(datastoreErrorStr)
2486 for failed, msg in (
2487 (failedNames, "Keys without corresponding DatasetType or StorageClass entry: "),
2488 (failedDataId, "Keys with bad DataId entries: "),
2489 ):
2490 if failed:
2491 msg += ", ".join(str(k) for k in failed)
2492 messages.append(msg)
2494 if messages:
2495 raise ValidationError(";\n".join(messages))
2497 @property
2498 @deprecated(
2499 "Please use 'collections' instead. collection_chains will be removed after v28.",
2500 version="v28",
2501 category=FutureWarning,
2502 )
2503 def collection_chains(self) -> DirectButlerCollections:
2504 """Object with methods for modifying collection chains."""
2505 return DirectButlerCollections(self._registry)
2507 @property
2508 def collections(self) -> DirectButlerCollections:
2509 """Object with methods for modifying and inspecting collections."""
2510 return DirectButlerCollections(self._registry)
2512 @property
2513 def run(self) -> str | None:
2514 """Name of the run this butler writes outputs to by default (`str` or
2515 `None`).
2517 This is an alias for ``self.registry.defaults.run``. It cannot be set
2518 directly in isolation, but all defaults may be changed together by
2519 assigning a new `RegistryDefaults` instance to
2520 ``self.registry.defaults``.
2521 """
2522 return self._registry.defaults.run
2524 @property
2525 def registry(self) -> Registry:
2526 """The object that manages dataset metadata and relationships
2527 (`Registry`).
2529 Many operations that don't involve reading or writing butler datasets
2530 are accessible only via `Registry` methods. Eventually these methods
2531 will be replaced by equivalent `Butler` methods.
2532 """
2533 return self._registry_shim
2535 @property
2536 def dimensions(self) -> DimensionUniverse:
2537 # Docstring inherited.
2538 return self._registry.dimensions
2540 def query(self) -> contextlib.AbstractContextManager[Query]:
2541 # Docstring inherited.
2542 return self._registry._query()
2544 def _query_driver(
2545 self,
2546 default_collections: Iterable[str],
2547 default_data_id: DataCoordinate,
2548 ) -> contextlib.AbstractContextManager[DirectQueryDriver]:
2549 """Set up a QueryDriver instance for use with this Butler. Although
2550 this is marked as a private method, it is also used by Butler server.
2551 """
2552 return self._registry._query_driver(default_collections, default_data_id)
2554 @contextlib.contextmanager
2555 def _query_all_datasets_by_page(
2556 self, args: QueryAllDatasetsParameters
2557 ) -> Iterator[Iterator[list[DatasetRef]]]:
2558 with self.query() as query:
2559 pages = query_all_datasets(self, query, args)
2560 yield iter(page.data for page in pages)
2562 def _preload_cache(self, *, load_dimension_record_cache: bool = True) -> None:
2563 """Immediately load caches that are used for common operations."""
2564 self._registry.preload_cache(load_dimension_record_cache=load_dimension_record_cache)
2566 _config: ButlerConfig
2567 """Configuration for this Butler instance."""
2569 _registry: SqlRegistry
2570 """The object that manages dataset metadata and relationships
2571 (`SqlRegistry`).
2573 Most operations that don't involve reading or writing butler datasets are
2574 accessible only via `SqlRegistry` methods.
2575 """
2577 datastore: Datastore
2578 """The object that manages actual dataset storage (`Datastore`).
2580 Direct user access to the datastore should rarely be necessary; the primary
2581 exception is the case where a `Datastore` implementation provides extra
2582 functionality beyond what the base class defines.
2583 """
2585 storageClasses: StorageClassFactory
2586 """An object that maps known storage class names to objects that fully
2587 describe them (`StorageClassFactory`).
2588 """
2590 _registry_shim: RegistryShim
2591 """Shim object to provide a legacy public interface for querying via the
2592 the ``registry`` property.
2593 """
2596class _RefGroup(NamedTuple):
2597 """Key identifying a batch of DatasetRefs to be inserted in
2598 `Butler.transfer_from`.
2599 """
2601 dimensions: DimensionGroup
2602 run: str
2605class _ImportDatasetsInfo(NamedTuple):
2606 """Information extracted from datasets to be imported."""
2608 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]]
2609 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]
2612def _to_uuid(id: DatasetId | str) -> uuid.UUID:
2613 if isinstance(id, uuid.UUID):
2614 return id
2615 else:
2616 return uuid.UUID(id)