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