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