Coverage for python/lsst/daf/butler/direct_butler/_direct_butler.py: 10%
912 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:23 -0700
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,
1260 *,
1261 storage_class: str | StorageClass | None = None,
1262 dimension_records: bool = False,
1263 datastore_records: bool = False,
1264 ) -> DatasetRef | None:
1265 ref = self._registry.getDataset(id)
1266 if ref is not None:
1267 if dimension_records:
1268 ref = ref.expanded(
1269 self._registry.expandDataId(ref.dataId, dimensions=ref.datasetType.dimensions)
1270 )
1271 if storage_class:
1272 ref = ref.overrideStorageClass(storage_class)
1273 if datastore_records:
1274 ref = self._registry.get_datastore_records(ref)
1275 return ref
1277 def find_dataset(
1278 self,
1279 dataset_type: DatasetType | str,
1280 data_id: DataId | None = None,
1281 *,
1282 collections: str | Sequence[str] | None = None,
1283 timespan: Timespan | None = None,
1284 storage_class: str | StorageClass | None = None,
1285 dimension_records: bool = False,
1286 datastore_records: bool = False,
1287 **kwargs: Any,
1288 ) -> DatasetRef | None:
1289 # Handle any parts of the dataID that are not using primary dimension
1290 # keys.
1291 if isinstance(dataset_type, str):
1292 actual_type = self.get_dataset_type(dataset_type)
1293 else:
1294 actual_type = dataset_type
1296 # Store the component for later.
1297 component_name = actual_type.component()
1298 if actual_type.isComponent():
1299 parent_type = actual_type.makeCompositeDatasetType()
1300 else:
1301 parent_type = actual_type
1303 data_id, kwargs = self._rewrite_data_id(data_id, parent_type, **kwargs)
1305 ref = self._registry.findDataset(
1306 parent_type,
1307 data_id,
1308 collections=collections,
1309 timespan=timespan,
1310 datastore_records=datastore_records,
1311 **kwargs,
1312 )
1313 if ref is not None and dimension_records:
1314 ref = ref.expanded(self._registry.expandDataId(ref.dataId, dimensions=ref.datasetType.dimensions))
1315 if ref is not None and component_name:
1316 ref = ref.makeComponentRef(component_name)
1317 if ref is not None and storage_class is not None:
1318 ref = ref.overrideStorageClass(storage_class)
1320 return ref
1322 def retrieve_artifacts_zip(
1323 self,
1324 refs: Iterable[DatasetRef],
1325 destination: ResourcePathExpression,
1326 overwrite: bool = True,
1327 ) -> ResourcePath:
1328 return retrieve_and_zip(refs, destination, self._datastore.retrieveArtifacts, overwrite)
1330 def retrieveArtifacts(
1331 self,
1332 refs: Iterable[DatasetRef],
1333 destination: ResourcePathExpression,
1334 transfer: str = "auto",
1335 preserve_path: bool = True,
1336 overwrite: bool = False,
1337 ) -> list[ResourcePath]:
1338 # Docstring inherited.
1339 outdir = ResourcePath(destination)
1340 artifact_map = self._datastore.retrieveArtifacts(
1341 refs,
1342 outdir,
1343 transfer=transfer,
1344 preserve_path=preserve_path,
1345 overwrite=overwrite,
1346 write_index=True,
1347 )
1348 return list(artifact_map)
1350 def exists(
1351 self,
1352 dataset_ref_or_type: DatasetRef | DatasetType | str,
1353 /,
1354 data_id: DataId | None = None,
1355 *,
1356 full_check: bool = True,
1357 collections: Any = None,
1358 **kwargs: Any,
1359 ) -> DatasetExistence:
1360 # Docstring inherited.
1361 existence = DatasetExistence.UNRECOGNIZED
1363 if isinstance(dataset_ref_or_type, DatasetRef):
1364 if collections is not None:
1365 warnings.warn("Collections should not be specified with DatasetRef", stacklevel=2)
1366 if data_id is not None:
1367 warnings.warn("A DataID should not be specified with DatasetRef", stacklevel=2)
1368 ref = dataset_ref_or_type
1369 registry_ref = self._registry.getDataset(dataset_ref_or_type.id)
1370 if registry_ref is not None:
1371 existence |= DatasetExistence.RECORDED
1373 if dataset_ref_or_type != registry_ref:
1374 # This could mean that storage classes differ, so we should
1375 # check for that but use the registry ref for the rest of
1376 # the method.
1377 if registry_ref.is_compatible_with(dataset_ref_or_type):
1378 # Use the registry version from now on.
1379 ref = registry_ref
1380 else:
1381 raise ValueError(
1382 f"The ref given to exists() ({ref}) has the same dataset ID as one "
1383 f"in registry but has different incompatible values ({registry_ref})."
1384 )
1385 else:
1386 try:
1387 ref = self._findDatasetRef(dataset_ref_or_type, data_id, collections=collections, **kwargs)
1388 except (LookupError, TypeError):
1389 return existence
1390 existence |= DatasetExistence.RECORDED
1392 if self._datastore.knows(ref):
1393 existence |= DatasetExistence.DATASTORE
1395 if full_check:
1396 if self._datastore.exists(ref):
1397 existence |= DatasetExistence._ARTIFACT
1398 elif existence.value != DatasetExistence.UNRECOGNIZED.value:
1399 # Do not add this flag if we have no other idea about a dataset.
1400 existence |= DatasetExistence(DatasetExistence._ASSUMED)
1402 return existence
1404 def _exists_many(
1405 self,
1406 refs: Iterable[DatasetRef],
1407 /,
1408 *,
1409 full_check: bool = True,
1410 ) -> dict[DatasetRef, DatasetExistence]:
1411 # Docstring inherited.
1412 existence = {ref: DatasetExistence.UNRECOGNIZED for ref in refs}
1414 # Registry does not have a bulk API to check for a ref.
1415 for ref in refs:
1416 registry_ref = self._registry.getDataset(ref.id)
1417 if registry_ref is not None:
1418 # It is possible, albeit unlikely, that the given ref does
1419 # not match the one in registry even though the UUID matches.
1420 # When checking a single ref we raise, but it's impolite to
1421 # do that when potentially hundreds of refs are being checked.
1422 # We could change the API to only accept UUIDs and that would
1423 # remove the ability to even check and remove the worry
1424 # about differing storage classes. Given the ongoing discussion
1425 # on refs vs UUIDs and whether to raise or have a new
1426 # private flag, treat this as a private API for now.
1427 existence[ref] |= DatasetExistence.RECORDED
1429 # Ask datastore if it knows about these refs.
1430 knows = self._datastore.knows_these(refs)
1431 for ref, known in knows.items():
1432 if known:
1433 existence[ref] |= DatasetExistence.DATASTORE
1435 if full_check:
1436 mexists = self._datastore.mexists(refs)
1437 for ref, exists in mexists.items():
1438 if exists:
1439 existence[ref] |= DatasetExistence._ARTIFACT
1440 else:
1441 # Do not set this flag if nothing is known about the dataset.
1442 for ref in existence:
1443 if existence[ref] != DatasetExistence.UNRECOGNIZED:
1444 existence[ref] |= DatasetExistence._ASSUMED
1446 return existence
1448 def removeRuns(
1449 self,
1450 names: Iterable[str],
1451 unstore: bool | type[_DeprecatedDefault] = _DeprecatedDefault,
1452 *,
1453 unlink_from_chains: bool = False,
1454 ) -> None:
1455 # Docstring inherited.
1456 if not self.isWriteable():
1457 raise TypeError("Butler is read-only.")
1459 if unstore is not _DeprecatedDefault:
1460 # The value was passed in by a user. Must report it is now
1461 # ignored.
1462 if unstore is True:
1463 msg = "The unstore parameter is deprecated and is now always treated as True. "
1464 else:
1465 msg = "The unstore parameter for removeRuns can no longer be False and is now ignored. "
1466 warnings.warn(
1467 msg + " The parameter will be removed after v30.",
1468 category=FutureWarning,
1469 stacklevel=find_outside_stacklevel("lsst.daf.butler"),
1470 )
1472 names = list(names)
1473 refs: list[DatasetRef] = []
1474 # Map of the chained collections to the RUN children.
1475 parents_to_children: dict[str, set[str]] = defaultdict(set)
1477 with self._caching_context():
1478 # Get information about these RUNs.
1479 collections_info = self.collections.query_info(names, include_parents=unlink_from_chains)
1480 for info in collections_info:
1481 if info.type is not CollectionType.RUN:
1482 raise TypeError(f"The collection type of '{info.name}' is {info.type.name}, not RUN.")
1483 if unlink_from_chains:
1484 if info.parents is None: # For mypy.
1485 raise AssertionError("Internal error: Collection parents required but not received")
1486 for parent in info.parents:
1487 parents_to_children[parent].add(info.name)
1489 # Update the names in case the query unexpectedly had a wildcard.
1490 names = [info.name for info in collections_info]
1492 # Get all the datasets from these runs.
1493 refs = self.query_all_datasets(names, find_first=False, limit=None)
1495 # Call pruneDatasets since we are deliberately removing
1496 # datasets in chunks from the RUN collections rather than
1497 # attempting to remove everything at once.
1498 with time_this(
1499 _LOG,
1500 msg="Removing %d dataset%s from %s",
1501 args=(len(refs), "s" if len(refs) != 1 else "", ", ".join(names)),
1502 ):
1503 self.pruneDatasets(refs, unstore=True, purge=True, disassociate=True)
1505 # Now can remove the actual RUN collection and unlink from chains.
1506 with self._registry.transaction():
1507 # This will fail if caller is not unlinking from chains but the
1508 # RUN is in a chain -- but we have already deleted all the datasets
1509 # by this point.
1510 if unlink_from_chains:
1511 # Use deterministic order for deletions to attempt to minimize
1512 # risk of deadlocks for parallel deletes.
1513 for parent in sorted(parents_to_children):
1514 self.collections.remove_from_chain(parent, sorted(parents_to_children[parent]))
1515 # Sort to avoid potential deadlocks.
1516 for name in sorted(names):
1517 # This should be fast since the collection should be empty.
1518 with time_this(_LOG, msg="Removing RUN collection %s", args=(name,)):
1519 self._registry.removeCollection(name)
1520 _LOG.info("Completely removed the following RUN collections: %s", ", ".join(names))
1522 def pruneDatasets(
1523 self,
1524 refs: Iterable[DatasetRef],
1525 *,
1526 disassociate: bool = True,
1527 unstore: bool = False,
1528 tags: Iterable[str] = (),
1529 purge: bool = False,
1530 ) -> None:
1531 # docstring inherited from LimitedButler
1533 if not self.isWriteable():
1534 raise TypeError("Butler is read-only.")
1535 if purge:
1536 if not disassociate:
1537 raise TypeError("Cannot pass purge=True without disassociate=True.")
1538 if not unstore:
1539 raise TypeError("Cannot pass purge=True without unstore=True.")
1540 elif disassociate:
1541 tags = tuple(tags)
1542 if not tags:
1543 raise TypeError("No tags provided but disassociate=True.")
1544 for tag in tags:
1545 collectionType = self._registry.getCollectionType(tag)
1546 if collectionType is not CollectionType.TAGGED:
1547 raise TypeError(
1548 f"Cannot disassociate from collection '{tag}' "
1549 f"of non-TAGGED type {collectionType.name}."
1550 )
1551 # Transform possibly-single-pass iterable into something we can iterate
1552 # over multiple times.
1553 refs = list(refs)
1554 # Pruning a component of a DatasetRef makes no sense since registry
1555 # doesn't know about components and datastore might not store
1556 # components in a separate file
1557 for ref in refs:
1558 if ref.datasetType.component():
1559 raise ValueError(f"Can not prune a component of a dataset (ref={ref})")
1561 # Chunk the deletions using a size that is reasonably efficient whilst
1562 # also giving reasonable feedback to the user. Chunking also minimizes
1563 # what needs to rollback if there is a failure and should allow
1564 # incremental re-running of the pruning (assuming the query is
1565 # repeated). The only issue will be if the Ctrl-C comes during
1566 # emptyTrash since an admin command would need to run to finish the
1567 # emptying of that chunk.
1568 progress = Progress("lsst.daf.butler.Butler.pruneDatasets", level=_LOG.INFO)
1569 chunk_size = 50_000
1570 n_chunks = math.ceil(len(refs) / chunk_size)
1571 if n_chunks > 1:
1572 _LOG.verbose("Pruning a total of %d datasets", len(refs))
1573 chunk_num = 0
1574 for chunked_refs in progress.wrap(
1575 chunk_iterable(refs, chunk_size=chunk_size), desc="Deleting datasets", total=n_chunks
1576 ):
1577 chunk_num += 1
1578 _LOG.verbose(
1579 "Pruning %d dataset%s in chunk %d/%d",
1580 len(chunked_refs),
1581 "s" if len(chunked_refs) != 1 else "",
1582 chunk_num,
1583 n_chunks,
1584 )
1585 with time_this(
1586 _LOG,
1587 msg="Removing %d datasets for chunk %d/%d",
1588 args=(len(chunked_refs), chunk_num, n_chunks),
1589 ):
1590 self._prune_datasets(
1591 chunked_refs, tags=tags, unstore=unstore, purge=purge, disassociate=disassociate
1592 )
1594 def _prune_datasets(
1595 self,
1596 refs: Collection[DatasetRef],
1597 *,
1598 disassociate: bool = True,
1599 unstore: bool = False,
1600 tags: Iterable[str] = (),
1601 purge: bool = False,
1602 ) -> None:
1603 # We don't need an unreliable Datastore transaction for this, because
1604 # we've been extra careful to ensure that Datastore.trash only involves
1605 # mutating the Registry (it can _look_ at Datastore-specific things,
1606 # but shouldn't change them), and hence all operations here are
1607 # Registry operations.
1608 with self.transaction():
1609 plural = "s" if len(refs) != 1 else ""
1610 if unstore:
1611 with time_this(
1612 _LOG,
1613 msg="Marking %d dataset%s for removal during pruneDatasets",
1614 args=(len(refs), plural),
1615 ):
1616 self._datastore.trash(refs)
1617 if purge:
1618 with time_this(
1619 _LOG, msg="Removing %d pruned dataset%s from registry", args=(len(refs), plural)
1620 ):
1621 self._registry.removeDatasets(refs)
1622 elif disassociate:
1623 assert tags, "Guaranteed by earlier logic in this function."
1624 with time_this(
1625 _LOG, msg="Disassociating %d dataset%ss from tagged collections", args=(len(refs), plural)
1626 ):
1627 for tag in tags:
1628 self._registry.disassociate(tag, refs)
1629 # We've exited the Registry transaction, and apparently committed.
1630 # (if there was an exception, everything rolled back, and it's as if
1631 # nothing happened - and we never get here).
1632 # Datastore artifacts are not yet gone, but they're clearly marked
1633 # as trash, so if we fail to delete now because of (e.g.) filesystem
1634 # problems we can try again later, and if manual administrative
1635 # intervention is required, it's pretty clear what that should entail:
1636 # deleting everything on disk and in private Datastore tables that is
1637 # in the dataset_location_trash table.
1638 if unstore:
1639 # Point of no return for removing artifacts. Restrict the trash
1640 # emptying to the refs that this call trashed.
1641 with time_this(
1642 _LOG,
1643 msg="Attempting to remove artifacts for %d dataset%s associated with pruning",
1644 args=(len(refs), plural),
1645 ):
1646 self._datastore.emptyTrash(refs=refs)
1648 def ingest_zip(
1649 self,
1650 zip_file: ResourcePathExpression,
1651 transfer: str = "auto",
1652 *,
1653 transfer_dimensions: bool = False,
1654 dry_run: bool = False,
1655 ) -> None:
1656 """Ingest a Zip file into this butler.
1658 The Zip file must have been created by `retrieve_artifacts_zip`.
1660 Parameters
1661 ----------
1662 zip_file : `lsst.resources.ResourcePathExpression`
1663 Path to the Zip file.
1664 transfer : `str`, optional
1665 Method to use to transfer the Zip into the datastore.
1666 transfer_dimensions : `bool`, optional
1667 If `True`, dimension record data associated with the new datasets
1668 will be transferred from the Zip, if present.
1669 dry_run : `bool`, optional
1670 If `True` the ingest will be processed without any modifications
1671 made to the target butler and as if the target butler did not
1672 have any of the datasets.
1674 Notes
1675 -----
1676 Run collections and dataset types are created as needed.
1677 """
1678 if not self.isWriteable():
1679 raise TypeError("Butler is read-only.")
1681 zip_path = ResourcePath(zip_file)
1682 index = ZipIndex.from_zip_file(zip_path)
1683 _LOG.verbose(
1684 "Ingesting %s containing %d datasets and %d files.", zip_path, len(index.refs), len(index)
1685 )
1687 # Need to ingest the refs into registry. Re-use the standard ingest
1688 # code by reconstructing FileDataset from the index.
1689 refs = index.refs.to_refs(universe=self.dimensions)
1690 id_to_ref = {ref.id: ref for ref in refs}
1691 datasets: list[FileDataset] = []
1692 processed_ids: set[uuid.UUID] = set()
1693 for path_in_zip, index_info in index.artifact_map.items():
1694 # Disassembled composites need to check this ref isn't already
1695 # included.
1696 unprocessed = {id_ for id_ in index_info.ids if id_ not in processed_ids}
1697 if not unprocessed:
1698 continue
1699 dataset = FileDataset(refs=[id_to_ref[id_] for id_ in unprocessed], path=path_in_zip)
1700 datasets.append(dataset)
1701 processed_ids.update(unprocessed)
1703 # Ingest doesn't create the RUN collections so we have to do that
1704 # here.
1705 #
1706 # Sort by run collection name to ensure Postgres takes locks in the
1707 # same order between different processes, to mitigate an issue
1708 # where Postgres can deadlock due to the unique index on collection
1709 # name. (See DM-47543).
1710 runs = {ref.run for ref in refs}
1711 for run in sorted(runs):
1712 registered = self.collections.register(run)
1713 if registered:
1714 _LOG.verbose("Created RUN collection %s as part of zip ingest", run)
1716 progress = Progress("lsst.daf.butler.Butler.ingest", level=VERBOSE)
1717 import_info = self._prepare_ingest_file_datasets(
1718 datasets, progress, dry_run=dry_run, transfer_dimensions=transfer_dimensions
1719 )
1721 with self.transaction():
1722 # Do not need expanded dataset refs so can ignore the return value.
1723 self._ingest_file_datasets(datasets, import_info, progress, dry_run=dry_run)
1725 try:
1726 self._datastore.ingest_zip(zip_path, transfer=transfer, dry_run=dry_run)
1727 except IntegrityError as e:
1728 raise ConflictingDefinitionError(
1729 f"Datastore already contains one or more datasets: {e}"
1730 ) from e
1732 def _prepare_ingest_file_datasets(
1733 self,
1734 datasets: Sequence[FileDataset],
1735 progress: Progress,
1736 *,
1737 transfer_dimensions: bool = False,
1738 dry_run: bool = False,
1739 ) -> _ImportDatasetsInfo:
1740 # Track DataIDs that are being ingested so we can spot issues early
1741 # with duplication. Retain previous FileDataset so we can report it.
1742 groupedDataIds: MutableMapping[tuple[DatasetType, str], dict[DataCoordinate, FileDataset]] = (
1743 defaultdict(dict)
1744 )
1746 # All the refs we need to import.
1747 refs: list[DatasetRef] = []
1749 for dataset in progress.wrap(datasets, desc="Validating dataIDs"):
1750 for ref in dataset.refs:
1751 group_key = (ref.datasetType, ref.run)
1753 if ref.dataId in groupedDataIds[group_key]:
1754 raise ConflictingDefinitionError(
1755 f"Ingest conflict. Dataset {dataset.path} has same"
1756 " DataId as other ingest dataset"
1757 f" {groupedDataIds[group_key][ref.dataId].path} "
1758 f" ({ref.dataId})"
1759 )
1761 groupedDataIds[group_key][ref.dataId] = dataset
1762 refs.extend(dataset.refs)
1764 # Ensure that dataset types are created and all ref information
1765 # extracted.
1766 import_info = self._prepare_for_import_refs(
1767 self,
1768 refs,
1769 register_dataset_types=True,
1770 dry_run=dry_run,
1771 transfer_dimensions=transfer_dimensions,
1772 )
1773 return import_info
1775 def _ingest_file_datasets(
1776 self,
1777 datasets: Sequence[FileDataset],
1778 import_info: _ImportDatasetsInfo,
1779 progress: Progress,
1780 *,
1781 dry_run: bool = False,
1782 ) -> None:
1783 self._import_dimension_records(import_info.dimension_records, dry_run=dry_run)
1784 imported_refs = self._import_grouped_refs(
1785 import_info.grouped_refs, None, progress, dry_run=dry_run, expand_refs=True
1786 )
1788 # The expanded refs need to be attached back to the original
1789 # FileDatasets for datastore to use.
1790 id_to_ref = {ref.id: ref for ref in imported_refs}
1792 for dataset in progress.wrap(datasets, desc="Re-attaching expanded refs"):
1793 dataset.refs = [id_to_ref[ref.id] for ref in dataset.refs]
1795 def ingest(
1796 self,
1797 *datasets: FileDataset,
1798 transfer: str | None = "auto",
1799 record_validation_info: bool = True,
1800 skip_existing: bool = False,
1801 ) -> None:
1802 # Docstring inherited.
1803 if not datasets:
1804 return
1805 if not self.isWriteable():
1806 raise TypeError("Butler is read-only.")
1807 _LOG.verbose("Ingesting %d file dataset%s.", len(datasets), "" if len(datasets) == 1 else "s")
1808 progress = Progress("lsst.daf.butler.Butler.ingest", level=VERBOSE)
1810 new_datasets, existing_datasets = self._partition_datasets_by_known(datasets)
1811 if existing_datasets:
1812 if skip_existing:
1813 _LOG.info(
1814 "Skipping %d datasets which already exist in the repository.", len(existing_datasets)
1815 )
1816 else:
1817 raise ConflictingDefinitionError(
1818 f"Datastore already contains {len(existing_datasets)} of the given datasets."
1819 f" Example: {existing_datasets[0]}"
1820 )
1822 # We use `datasets` rather `new_datasets` for the Registry
1823 # portion of this, to let it confirm that everything matches the
1824 # existing datasets.
1825 import_info = self._prepare_ingest_file_datasets(datasets, progress)
1827 with self.transaction():
1828 self._ingest_file_datasets(datasets, import_info, progress)
1830 # Bulk-insert everything into Datastore.
1831 # We do not know if any of the registry entries already existed
1832 # (_importDatasets only complains if they exist but differ).
1833 # The _partition_datasets_by_known logic above should catch most
1834 # instances where we attempt to re-ingest files that were already
1835 # ingested, but a concurrent writer could cause a unique constraint
1836 # violation here.
1837 try:
1838 self._datastore.ingest(
1839 *new_datasets, transfer=transfer, record_validation_info=record_validation_info
1840 )
1841 except IntegrityError as e:
1842 raise ConflictingDefinitionError(
1843 f"Datastore already contains one or more datasets: {e}"
1844 ) from e
1846 def _partition_datasets_by_known(
1847 self, datasets: Iterable[FileDataset]
1848 ) -> tuple[list[FileDataset], list[FileDataset]]:
1849 """Divides the given `FileDataset` objects into two groups: those for
1850 which the Datastore already has an entry, and those for which it does
1851 not.
1852 """
1853 new_datasets = []
1854 existing_datasets = []
1856 refs = itertools.chain.from_iterable(dataset.refs for dataset in datasets)
1857 known_refs = self._datastore.knows_these(refs)
1859 for dataset in datasets:
1860 if any(known_refs[ref] for ref in dataset.refs):
1861 existing_datasets.append(dataset)
1862 else:
1863 new_datasets.append(dataset)
1865 return new_datasets, existing_datasets
1867 @contextlib.contextmanager
1868 def export(
1869 self,
1870 *,
1871 directory: str | None = None,
1872 filename: str | None = None,
1873 format: str | None = None,
1874 transfer: str | None = None,
1875 ) -> Iterator[RepoExportContext]:
1876 # Docstring inherited.
1877 if directory is None and transfer is not None:
1878 raise TypeError("Cannot transfer without providing a directory.")
1879 if transfer == "move":
1880 raise TypeError("Transfer may not be 'move': export is read-only")
1881 if format is None:
1882 if filename is None:
1883 raise TypeError("At least one of 'filename' or 'format' must be provided.")
1884 else:
1885 _, format = os.path.splitext(filename)
1886 if not format:
1887 raise ValueError("Please specify a file extension to determine export format.")
1888 format = format[1:] # Strip leading ".""
1889 elif filename is None:
1890 filename = f"export.{format}"
1891 if directory is not None:
1892 filename = os.path.join(directory, filename)
1893 formats = self._config["repo_transfer_formats"]
1894 if format not in formats:
1895 raise ValueError(f"Unknown export format {format!r}, allowed: {','.join(formats.keys())}")
1896 BackendClass = get_class_of(formats[format, "export"])
1897 with open(filename, "w") as stream:
1898 backend = BackendClass(stream, universe=self.dimensions)
1899 try:
1900 helper = RepoExportContext(self, backend=backend, directory=directory, transfer=transfer)
1901 with self._caching_context():
1902 yield helper
1903 except BaseException:
1904 raise
1905 else:
1906 helper._finish()
1908 def import_(
1909 self,
1910 *,
1911 directory: ResourcePathExpression | None = None,
1912 filename: ResourcePathExpression | TextIO | None = None,
1913 format: str | None = None,
1914 transfer: str | None = None,
1915 skip_dimensions: set | None = None,
1916 record_validation_info: bool = True,
1917 without_datastore: bool = False,
1918 ) -> None:
1919 # Docstring inherited.
1920 if not self.isWriteable():
1921 raise TypeError("Butler is read-only.")
1922 if filename is None and format is not None:
1923 filename = ResourcePath(f"export.{format}", forceAbsolute=False)
1924 if directory is not None:
1925 directory = ResourcePath(directory, forceDirectory=True)
1926 # mypy doesn't think this will work but it does in python >= 3.10.
1927 if isinstance(filename, ResourcePathExpression): # type: ignore
1928 filename = ResourcePath(filename, forceAbsolute=False) # type: ignore
1929 if format is None:
1930 format = filename.getExtension()
1931 if not filename.isabs() and directory is not None:
1932 potential = directory.join(filename)
1933 exists_in_cwd = filename.exists()
1934 exists_in_dir = potential.exists()
1935 if exists_in_cwd and exists_in_dir:
1936 _LOG.warning(
1937 "A relative path for filename was specified (%s) which exists relative to cwd. "
1938 "Additionally, the file exists relative to the given search directory (%s). "
1939 "Using the export file in the given directory.",
1940 filename,
1941 potential,
1942 )
1943 # Given they specified an explicit directory and that
1944 # directory has the export file in it, assume that that
1945 # is what was meant despite the file in cwd.
1946 filename = potential
1947 elif exists_in_dir:
1948 filename = potential
1949 elif not exists_in_cwd and not exists_in_dir:
1950 # Raise early.
1951 raise FileNotFoundError(
1952 f"Export file could not be found in {filename.abspath()} or {potential.abspath()}."
1953 )
1954 elif format is None:
1955 format = ".yaml"
1956 BackendClass: type[RepoImportBackend] = get_class_of(
1957 self._config["repo_transfer_formats"][format]["import"]
1958 )
1960 def doImport(importStream: TextIO | ResourceHandleProtocol) -> None:
1961 with self._caching_context():
1962 backend = BackendClass(importStream, self) # type: ignore[call-arg]
1963 backend.register()
1964 with self.transaction():
1965 backend.load(
1966 datastore=self._datastore if not without_datastore else None,
1967 directory=directory,
1968 transfer=transfer,
1969 skip_dimensions=skip_dimensions,
1970 record_validation_info=record_validation_info,
1971 )
1973 if isinstance(filename, ResourcePath):
1974 # We can not use open() here at the moment because of
1975 # DM-38589 since yaml does stream.read(8192) in a loop.
1976 stream = io.StringIO(filename.read().decode())
1977 doImport(stream)
1978 else:
1979 doImport(filename) # type: ignore
1981 def transfer_dimension_records_from(
1982 self, source_butler: LimitedButler | Butler, source_refs: Iterable[DatasetRef]
1983 ) -> None:
1984 # Allowed dimensions in the target butler.
1985 elements = frozenset(element for element in self.dimensions.elements if element.has_own_table)
1987 data_ids = {ref.dataId for ref in source_refs}
1989 dimension_records = self._extract_all_dimension_records_from_data_ids(
1990 source_butler, data_ids, elements
1991 )
1993 # Insert order is important.
1994 for element in self.dimensions.sorted(dimension_records.keys()):
1995 records = [r for r in dimension_records[element].values()]
1996 # Assume that if the record is already present that we can
1997 # use it without having to check that the record metadata
1998 # is consistent.
1999 self._registry.insertDimensionData(element, *records, skip_existing=True)
2000 _LOG.debug("Dimension '%s' -- number of records transferred: %d", element.name, len(records))
2002 def _extract_all_dimension_records_from_data_ids(
2003 self,
2004 source_butler: LimitedButler | Butler,
2005 data_ids: set[DataCoordinate],
2006 allowed_elements: frozenset[DimensionElement],
2007 ) -> dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]:
2008 primary_records = self._extract_dimension_records_from_data_ids(
2009 source_butler, data_ids, allowed_elements
2010 )
2012 can_query = True if isinstance(source_butler, Butler) else False
2014 additional_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2015 for original_element, record_mapping in primary_records.items():
2016 # Get dimensions that depend on this dimension.
2017 populated_by = self.dimensions.get_elements_populated_by(
2018 self.dimensions[original_element.name] # type: ignore
2019 )
2021 for data_id in record_mapping.keys():
2022 for element in populated_by:
2023 if element not in allowed_elements:
2024 continue
2025 if element.name == original_element.name:
2026 continue
2028 if element.name in primary_records:
2029 # If this element has already been stored avoid
2030 # re-finding records since that may lead to additional
2031 # spurious records. e.g. visit is populated_by
2032 # visit_detector_region but querying
2033 # visit_detector_region by visit will return all the
2034 # detectors for this visit -- the visit dataId does not
2035 # constrain this.
2036 # To constrain the query the original dataIds would
2037 # have to be scanned.
2038 continue
2040 if not can_query:
2041 raise RuntimeError(
2042 f"Transferring populated_by records like {element.name} requires a full Butler."
2043 )
2045 records = source_butler.query_dimension_records( # type: ignore
2046 element.name,
2047 explain=False,
2048 **data_id.mapping, # type: ignore
2049 )
2050 for record in records:
2051 additional_records[record.definition].setdefault(record.dataId, record)
2053 # The next step is to walk back through the additional records to
2054 # pick up any missing content (such as visit_definition needing to
2055 # know the exposure). Want to ensure we do not request records we
2056 # already have.
2057 missing_data_ids = set()
2058 for name, record_mapping in additional_records.items():
2059 for data_id in record_mapping.keys():
2060 if data_id not in primary_records[name]:
2061 missing_data_ids.add(data_id)
2063 # Fill out the new records. Assume that these new records do not
2064 # also need to carry over additional populated_by records.
2065 secondary_records = self._extract_dimension_records_from_data_ids(
2066 source_butler, missing_data_ids, allowed_elements
2067 )
2069 # Merge the extra sets of records in with the original.
2070 for name, record_mapping in itertools.chain(additional_records.items(), secondary_records.items()):
2071 primary_records[name].update(record_mapping)
2073 return primary_records
2075 def _extract_dimension_records_from_data_ids(
2076 self,
2077 source_butler: LimitedButler | Butler,
2078 data_ids: set[DataCoordinate],
2079 allowed_elements: frozenset[DimensionElement],
2080 ) -> dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]:
2081 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2083 for data_id in data_ids:
2084 # Need an expanded record, if not expanded that we need a full
2085 # butler with registry (allow mocks with registry too).
2086 if not data_id.hasRecords():
2087 if registry := getattr(source_butler, "registry", None):
2088 data_id = registry.expandDataId(data_id)
2089 else:
2090 raise TypeError("Input butler needs to be a full butler to expand DataId.")
2091 # If this butler doesn't know about a dimension in the source
2092 # butler things will break later.
2093 for element_name in data_id.dimensions.elements:
2094 record = data_id.records[element_name]
2095 if record is not None and record.definition in allowed_elements:
2096 dimension_records[record.definition].setdefault(record.dataId, record)
2098 return dimension_records
2100 def _prepare_for_import_refs(
2101 self,
2102 source_butler: LimitedButler,
2103 source_refs: Iterable[DatasetRef],
2104 *,
2105 register_dataset_types: bool = False,
2106 transfer_dimensions: bool = False,
2107 dry_run: bool = False,
2108 ) -> _ImportDatasetsInfo:
2109 # Docstring inherited.
2110 if not self.isWriteable():
2111 raise TypeError("Butler is read-only.")
2113 # Will iterate through the refs multiple times so need to convert
2114 # to a list if this isn't a collection.
2115 if not isinstance(source_refs, collections.abc.Collection):
2116 source_refs = list(source_refs)
2118 original_count = len(source_refs)
2119 log_level = _LOG.INFO if original_count > 1 else _LOG.VERBOSE
2120 _LOG.log(
2121 log_level,
2122 "Importing %d dataset%s into %s",
2123 original_count,
2124 "s" if original_count != 1 else "",
2125 str(self),
2126 )
2128 # Importing requires that we group the refs by dimension group and run
2129 # before doing the import.
2130 source_dataset_types = set()
2131 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]] = defaultdict(list)
2132 for ref in source_refs:
2133 grouped_refs[_RefGroup(ref.datasetType.dimensions, ref.run)].append(ref)
2134 source_dataset_types.add(ref.datasetType)
2136 # Check to see if the dataset type in the source butler has
2137 # the same definition in the target butler and register missing
2138 # ones if requested. Registration must happen outside a transaction.
2139 newly_registered_dataset_types = set()
2140 for datasetType in source_dataset_types:
2141 if register_dataset_types:
2142 # Let this raise immediately if inconsistent. Continuing
2143 # on to find additional inconsistent dataset types
2144 # might result in additional unwanted dataset types being
2145 # registered.
2146 try:
2147 if not dry_run and self._registry.registerDatasetType(datasetType):
2148 newly_registered_dataset_types.add(datasetType)
2149 except ConflictingDefinitionError as e:
2150 # Be safe and require that conversions be bidirectional
2151 # when there are storage class mismatches. This is because
2152 # get() will have to support conversion from source to
2153 # target python type (the source formatter will be
2154 # returning source python type) but there also is an
2155 # expectation that people will want to be able to get() in
2156 # the target using the source python type, which will not
2157 # require conversion for transferred datasets but might
2158 # for target-native types. Additionally, butler.get does
2159 # not know that the formatter will return the wrong
2160 # python type and so will always check that the conversion
2161 # works even though it won't need it.
2162 target_dataset_type = self.get_dataset_type(datasetType.name)
2163 target_compatible_with_source = target_dataset_type.is_compatible_with(datasetType)
2164 source_compatible_with_target = datasetType.is_compatible_with(target_dataset_type)
2165 if not (target_compatible_with_source and source_compatible_with_target):
2166 if target_compatible_with_source:
2167 e.add_note(
2168 "Target dataset type storage class is compatible with source "
2169 "but the reverse is not true."
2170 )
2171 elif source_compatible_with_target:
2172 e.add_note(
2173 "Source dataset type storage class is compatible with target "
2174 "but the reverse is not true."
2175 )
2176 else:
2177 e.add_note("If storage classes differ, please register converters.")
2178 raise
2179 else:
2180 # If the dataset type is missing, let it fail immediately.
2181 target_dataset_type = self.get_dataset_type(datasetType.name)
2182 if target_dataset_type != datasetType:
2183 target_compatible_with_source = target_dataset_type.is_compatible_with(datasetType)
2184 source_compatible_with_target = datasetType.is_compatible_with(target_dataset_type)
2185 # Both conversion directions are currently required.
2186 if not (target_compatible_with_source and source_compatible_with_target):
2187 msg = ""
2188 if target_compatible_with_source:
2189 msg = (
2190 "Target storage class is compatible with the source storage class "
2191 "but the reverse is not true."
2192 )
2193 elif source_compatible_with_target:
2194 msg = (
2195 "Source storage class is compatible with the target storage class"
2196 " but the reverse is not true."
2197 )
2198 else:
2199 msg = "If storage classes differ register converters."
2200 if msg:
2201 msg = f"({msg})"
2202 raise ConflictingDefinitionError(
2203 "Source butler dataset type differs from definition"
2204 f" in target butler: {datasetType} !="
2205 f" {target_dataset_type} {msg}"
2206 )
2207 if newly_registered_dataset_types:
2208 # We may have registered some even if there were inconsistencies
2209 # but should let people know (or else remove them again).
2210 _LOG.verbose(
2211 "Registered the following dataset types in the target Butler: %s",
2212 ", ".join(d.name for d in newly_registered_dataset_types),
2213 )
2214 else:
2215 _LOG.verbose("All required dataset types are known to the target Butler")
2217 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]] = defaultdict(dict)
2218 if transfer_dimensions:
2219 # Collect all the dimension records for these refs.
2220 # All dimensions are to be copied but the list of valid dimensions
2221 # come from this butler's universe.
2222 elements = frozenset(element for element in self.dimensions.elements if element.has_own_table)
2223 dataIds = {ref.dataId for ref in source_refs}
2224 dimension_records = self._extract_all_dimension_records_from_data_ids(
2225 source_butler, dataIds, elements
2226 )
2227 return _ImportDatasetsInfo(grouped_refs, dimension_records)
2229 def _import_dimension_records(
2230 self,
2231 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]],
2232 *,
2233 dry_run: bool,
2234 ) -> None:
2235 """Import dimension records collected during import pre-process."""
2236 if dimension_records and not dry_run:
2237 _LOG.verbose("Ensuring that dimension records exist for transferred datasets.")
2238 # Order matters.
2239 for element in self.dimensions.sorted(dimension_records.keys()):
2240 records = list(dimension_records[element].values())
2241 # Assume that if the record is already present that we can
2242 # use it without having to check that the record metadata
2243 # is consistent.
2244 self._registry.insertDimensionData(element, *records, skip_existing=True)
2246 def _import_grouped_refs(
2247 self,
2248 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]],
2249 source_butler: LimitedButler | None,
2250 progress: Progress,
2251 *,
2252 dry_run: bool = False,
2253 expand_refs: bool = False,
2254 ) -> list[DatasetRef]:
2255 handled_collections: set[str] = set()
2256 n_to_import = 0
2257 all_imported_refs: list[DatasetRef] = []
2258 # Sort by run collection name to ensure Postgres takes locks in the
2259 # same order between different processes, to mitigate an issue
2260 # where Postgres can deadlock due to the unique index on collection
2261 # name. (See DM-47543).
2262 groups = sorted(grouped_refs.items(), key=lambda item: item[0].run)
2263 for (dimension_group, run), refs_to_import in progress.iter_item_chunks(
2264 groups, desc="Importing to registry by run and dataset type"
2265 ):
2266 if run not in handled_collections:
2267 # May need to create output collection. If source butler
2268 # has a registry, ask for documentation string.
2269 run_doc = None
2270 if source_butler is not None and (registry := getattr(source_butler, "registry", None)):
2271 run_doc = registry.getCollectionDocumentation(run)
2272 if not dry_run:
2273 registered = self.collections.register(run, doc=run_doc)
2274 else:
2275 registered = True
2276 handled_collections.add(run)
2277 if registered:
2278 _LOG.verbose("Creating output run %s", run)
2280 n_refs = len(refs_to_import)
2281 n_to_import += n_refs
2282 _LOG.verbose(
2283 "Importing %d ref%s with dimensions %s into run %s",
2284 n_refs,
2285 "" if n_refs == 1 else "s",
2286 dimension_group.names,
2287 run,
2288 )
2290 # Assume we are using UUIDs and the source refs will match
2291 # those imported.
2292 if not dry_run:
2293 imported_refs = self._registry._importDatasets(refs_to_import, expand=expand_refs)
2294 else:
2295 imported_refs = refs_to_import
2297 all_imported_refs.extend(imported_refs)
2299 assert n_to_import == len(all_imported_refs)
2300 _LOG.verbose("Imported %d datasets into destination butler", n_to_import)
2301 return all_imported_refs
2303 def transfer_from(
2304 self,
2305 source_butler: LimitedButler,
2306 source_refs: Iterable[DatasetRef],
2307 transfer: str = "auto",
2308 skip_missing: bool = True,
2309 register_dataset_types: bool = False,
2310 transfer_dimensions: bool = False,
2311 dry_run: bool = False,
2312 ) -> collections.abc.Collection[DatasetRef]:
2313 # Docstring inherited.
2314 source_refs = list(source_refs)
2315 if not self.isWriteable():
2316 raise TypeError("Butler is read-only.")
2318 progress = Progress("lsst.daf.butler.Butler.transfer_from", level=VERBOSE)
2320 artifact_existence: dict[ResourcePath, bool] = {}
2321 file_transfer_source = source_butler._file_transfer_source
2322 transfer_records = retrieve_file_transfer_records(
2323 file_transfer_source, source_refs, artifact_existence
2324 )
2325 # In some situations the datastore artifact may be missing and we do
2326 # not want that registry entry to be imported. For example, this can
2327 # happen if a file was removed but the dataset was left in the registry
2328 # for provenance, or if a pipeline task didn't create all of the
2329 # possible files in a QuantumBackedButler.
2330 if skip_missing:
2331 original_ids = {ref.id for ref in source_refs}
2332 missing_ids = original_ids - transfer_records.keys()
2333 if missing_ids:
2334 original_count = len(source_refs)
2335 source_refs = [ref for ref in source_refs if ref.id not in missing_ids]
2336 filtered_count = len(source_refs)
2337 n_missing = original_count - filtered_count
2338 _LOG.verbose(
2339 "%d dataset%s removed because the artifact does not exist. Now have %d.",
2340 n_missing,
2341 "" if n_missing == 1 else "s",
2342 filtered_count,
2343 )
2345 import_info = self._prepare_for_import_refs(
2346 source_butler,
2347 source_refs,
2348 register_dataset_types=register_dataset_types,
2349 dry_run=dry_run,
2350 transfer_dimensions=transfer_dimensions,
2351 )
2353 # Do all the importing in a single transaction.
2354 with self.transaction():
2355 self._import_dimension_records(import_info.dimension_records, dry_run=dry_run)
2356 imported_refs = self._import_grouped_refs(
2357 import_info.grouped_refs, source_butler, progress, dry_run=dry_run
2358 )
2360 # Ask the datastore to transfer. The datastore has to check that
2361 # the source datastore is compatible with the target datastore.
2362 _LOG.verbose("Transferring %d datasets from %s", len(transfer_records), file_transfer_source.name)
2363 accepted, rejected = self._datastore.transfer_from(
2364 transfer_records,
2365 imported_refs,
2366 transfer=transfer,
2367 artifact_existence=artifact_existence,
2368 dry_run=dry_run,
2369 )
2370 if rejected:
2371 # For now, accept the registry entries but not the files.
2372 _LOG.warning(
2373 "%d datasets were rejected and %d accepted for transfer.",
2374 len(rejected),
2375 len(accepted),
2376 )
2378 return imported_refs
2380 def validateConfiguration(
2381 self,
2382 logFailures: bool = False,
2383 datasetTypeNames: Iterable[str] | None = None,
2384 ignore: Iterable[str] | None = None,
2385 ) -> None:
2386 # Docstring inherited.
2387 if datasetTypeNames:
2388 datasetTypes = [self.get_dataset_type(name) for name in datasetTypeNames]
2389 else:
2390 datasetTypes = list(self._registry.queryDatasetTypes())
2392 # filter out anything from the ignore list
2393 if ignore:
2394 ignore = set(ignore)
2395 datasetTypes = [
2396 e for e in datasetTypes if e.name not in ignore and e.nameAndComponent()[0] not in ignore
2397 ]
2398 else:
2399 ignore = set()
2401 # For each datasetType that has an instrument dimension, create
2402 # a DatasetRef for each defined instrument
2403 datasetRefs = []
2405 # Find all the registered instruments (if "instrument" is in the
2406 # universe).
2407 instruments: set[str] = set()
2408 if "instrument" in self.dimensions:
2409 instruments = {rec.name for rec in self.query_dimension_records("instrument", explain=False)}
2411 for datasetType in datasetTypes:
2412 if "instrument" in datasetType.dimensions:
2413 # In order to create a conforming dataset ref, create
2414 # fake DataCoordinate values for the non-instrument
2415 # dimensions. The type of the value does not matter here.
2416 dataId = {dim: 1 for dim in datasetType.dimensions.names if dim != "instrument"}
2418 for instrument in instruments:
2419 datasetRef = DatasetRef(
2420 datasetType,
2421 DataCoordinate.standardize(
2422 dataId, instrument=instrument, dimensions=datasetType.dimensions
2423 ),
2424 run="validate",
2425 )
2426 datasetRefs.append(datasetRef)
2428 entities: list[DatasetType | DatasetRef] = []
2429 entities.extend(datasetTypes)
2430 entities.extend(datasetRefs)
2432 datastoreErrorStr = None
2433 try:
2434 self._datastore.validateConfiguration(entities, logFailures=logFailures)
2435 except ValidationError as e:
2436 datastoreErrorStr = str(e)
2438 # Also check that the LookupKeys used by the datastores match
2439 # registry and storage class definitions
2440 keys = self._datastore.getLookupKeys()
2442 failedNames = set()
2443 failedDataId = set()
2444 for key in keys:
2445 if key.name is not None:
2446 if key.name in ignore:
2447 continue
2449 # skip if specific datasetType names were requested and this
2450 # name does not match
2451 if datasetTypeNames and key.name not in datasetTypeNames:
2452 continue
2454 # See if it is a StorageClass or a DatasetType
2455 if key.name in self.storageClasses:
2456 pass
2457 else:
2458 try:
2459 self.get_dataset_type(key.name)
2460 except KeyError:
2461 if logFailures:
2462 _LOG.critical(
2463 "Key '%s' does not correspond to a DatasetType or StorageClass", key
2464 )
2465 failedNames.add(key)
2466 else:
2467 # Dimensions are checked for consistency when the Butler
2468 # is created and rendezvoused with a universe.
2469 pass
2471 # Check that the instrument is a valid instrument
2472 # Currently only support instrument so check for that
2473 if key.dataId:
2474 dataIdKeys = set(key.dataId)
2475 if {"instrument"} != dataIdKeys:
2476 if logFailures:
2477 _LOG.critical("Key '%s' has unsupported DataId override", key)
2478 failedDataId.add(key)
2479 elif key.dataId["instrument"] not in instruments:
2480 if logFailures:
2481 _LOG.critical("Key '%s' has unknown instrument", key)
2482 failedDataId.add(key)
2484 messages = []
2486 if datastoreErrorStr:
2487 messages.append(datastoreErrorStr)
2489 for failed, msg in (
2490 (failedNames, "Keys without corresponding DatasetType or StorageClass entry: "),
2491 (failedDataId, "Keys with bad DataId entries: "),
2492 ):
2493 if failed:
2494 msg += ", ".join(str(k) for k in failed)
2495 messages.append(msg)
2497 if messages:
2498 raise ValidationError(";\n".join(messages))
2500 @property
2501 @deprecated(
2502 "Please use 'collections' instead. collection_chains will be removed after v28.",
2503 version="v28",
2504 category=FutureWarning,
2505 )
2506 def collection_chains(self) -> DirectButlerCollections:
2507 """Object with methods for modifying collection chains."""
2508 return DirectButlerCollections(self._registry)
2510 @property
2511 def collections(self) -> DirectButlerCollections:
2512 """Object with methods for modifying and inspecting collections."""
2513 return DirectButlerCollections(self._registry)
2515 @property
2516 def run(self) -> str | None:
2517 """Name of the run this butler writes outputs to by default (`str` or
2518 `None`).
2520 This is an alias for ``self.registry.defaults.run``. It cannot be set
2521 directly in isolation, but all defaults may be changed together by
2522 assigning a new `RegistryDefaults` instance to
2523 ``self.registry.defaults``.
2524 """
2525 return self._registry.defaults.run
2527 @property
2528 def registry(self) -> Registry:
2529 """The object that manages dataset metadata and relationships
2530 (`Registry`).
2532 Many operations that don't involve reading or writing butler datasets
2533 are accessible only via `Registry` methods. Eventually these methods
2534 will be replaced by equivalent `Butler` methods.
2535 """
2536 return self._registry_shim
2538 @property
2539 def dimensions(self) -> DimensionUniverse:
2540 # Docstring inherited.
2541 return self._registry.dimensions
2543 def query(self) -> contextlib.AbstractContextManager[Query]:
2544 # Docstring inherited.
2545 return self._registry._query()
2547 def _query_driver(
2548 self,
2549 default_collections: Iterable[str],
2550 default_data_id: DataCoordinate,
2551 ) -> contextlib.AbstractContextManager[DirectQueryDriver]:
2552 """Set up a QueryDriver instance for use with this Butler. Although
2553 this is marked as a private method, it is also used by Butler server.
2554 """
2555 return self._registry._query_driver(default_collections, default_data_id)
2557 @contextlib.contextmanager
2558 def _query_all_datasets_by_page(
2559 self, args: QueryAllDatasetsParameters
2560 ) -> Iterator[Iterator[list[DatasetRef]]]:
2561 with self.query() as query:
2562 pages = query_all_datasets(self, query, args)
2563 yield iter(page.data for page in pages)
2565 def _preload_cache(self, *, load_dimension_record_cache: bool = True) -> None:
2566 """Immediately load caches that are used for common operations."""
2567 self._registry.preload_cache(load_dimension_record_cache=load_dimension_record_cache)
2569 _config: ButlerConfig
2570 """Configuration for this Butler instance."""
2572 _registry: SqlRegistry
2573 """The object that manages dataset metadata and relationships
2574 (`SqlRegistry`).
2576 Most operations that don't involve reading or writing butler datasets are
2577 accessible only via `SqlRegistry` methods.
2578 """
2580 datastore: Datastore
2581 """The object that manages actual dataset storage (`Datastore`).
2583 Direct user access to the datastore should rarely be necessary; the primary
2584 exception is the case where a `Datastore` implementation provides extra
2585 functionality beyond what the base class defines.
2586 """
2588 storageClasses: StorageClassFactory
2589 """An object that maps known storage class names to objects that fully
2590 describe them (`StorageClassFactory`).
2591 """
2593 _registry_shim: RegistryShim
2594 """Shim object to provide a legacy public interface for querying via the
2595 the ``registry`` property.
2596 """
2599class _RefGroup(NamedTuple):
2600 """Key identifying a batch of DatasetRefs to be inserted in
2601 `Butler.transfer_from`.
2602 """
2604 dimensions: DimensionGroup
2605 run: str
2608class _ImportDatasetsInfo(NamedTuple):
2609 """Information extracted from datasets to be imported."""
2611 grouped_refs: defaultdict[_RefGroup, list[DatasetRef]]
2612 dimension_records: dict[DimensionElement, dict[DataCoordinate, DimensionRecord]]