Coverage for python/lsst/daf/butler/registry/sql_registry.py: 20%
452 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:23 -0700
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30from .. import ddl
32__all__ = ("SqlRegistry",)
34import contextlib
35import logging
36import warnings
37from collections import defaultdict
38from collections.abc import Iterable, Iterator, Mapping, Sequence
39from typing import TYPE_CHECKING, Any, cast
41import sqlalchemy
43from lsst.resources import ResourcePathExpression
44from lsst.utils.iteration import ensure_iterable
46from .._collection_type import CollectionType
47from .._column_tags import DatasetColumnTag
48from .._config import Config
49from .._dataset_ref import DatasetId, DatasetIdGenEnum, DatasetRef
50from .._dataset_type import DatasetType
51from .._exceptions import (
52 CalibrationLookupError,
53 DataIdValueError,
54 DimensionNameError,
55 InconsistentDataIdError,
56)
57from .._storage_class import StorageClassFactory
58from .._timespan import Timespan
59from ..dimensions import (
60 DataCoordinate,
61 DataId,
62 DimensionConfig,
63 DimensionDataAttacher,
64 DimensionElement,
65 DimensionGroup,
66 DimensionRecord,
67 DimensionUniverse,
68)
69from ..dimensions.record_cache import DimensionRecordCache
70from ..direct_query_driver import DirectQueryDriver
71from ..progress import Progress
72from ..queries import Query
73from ..registry import (
74 CollectionExpressionError,
75 CollectionSummary,
76 CollectionTypeError,
77 ConflictingDefinitionError,
78 NoDefaultCollectionError,
79 OrphanedRecordError,
80 RegistryConfig,
81 RegistryDefaults,
82 queries,
83)
84from ..registry.interfaces import ChainedCollectionRecord, ReadOnlyDatabaseError, RunRecord
85from ..registry.managers import RegistryManagerInstances, RegistryManagerTypes
86from ..registry.wildcards import CollectionWildcard, DatasetTypeWildcard
87from ..utils import transactional
89if TYPE_CHECKING:
90 from .._butler_config import ButlerConfig
91 from ..datastore._datastore import DatastoreOpaqueTable
92 from ..datastore.stored_file_info import StoredDatastoreItemInfo
93 from ..registry._registry import CollectionArgType
94 from ..registry.interfaces import (
95 CollectionRecord,
96 Database,
97 DatastoreRegistryBridgeManager,
98 ObsCoreTableManager,
99 )
102_LOG = logging.getLogger(__name__)
105class SqlRegistry:
106 """Butler Registry implementation that uses SQL database as backend.
108 Parameters
109 ----------
110 database : `Database`
111 Database instance to store Registry.
112 defaults : `RegistryDefaults`
113 Default collection search path and/or output `~CollectionType.RUN`
114 collection.
115 managers : `RegistryManagerInstances`
116 All the managers required for this registry.
117 """
119 defaultConfigFile: str | None = None
120 """Path to configuration defaults. Accessed within the ``configs`` resource
121 or relative to a search path. Can be None if no defaults specified.
122 """
124 @classmethod
125 def forceRegistryConfig(
126 cls, config: ButlerConfig | RegistryConfig | Config | str | None
127 ) -> RegistryConfig:
128 """Force the supplied config to a `RegistryConfig`.
130 Parameters
131 ----------
132 config : `RegistryConfig`, `Config` or `str` or `None`
133 Registry configuration, if missing then default configuration will
134 be loaded from registry.yaml.
136 Returns
137 -------
138 registry_config : `RegistryConfig`
139 A registry config.
140 """
141 if not isinstance(config, RegistryConfig):
142 if isinstance(config, str | Config) or config is None:
143 config = RegistryConfig(config)
144 else:
145 raise ValueError(f"Incompatible Registry configuration: {config}")
146 return config
148 @classmethod
149 def createFromConfig(
150 cls,
151 config: RegistryConfig | str | None = None,
152 dimensionConfig: DimensionConfig | str | None = None,
153 butlerRoot: ResourcePathExpression | None = None,
154 ) -> SqlRegistry:
155 """Create registry database and return `SqlRegistry` instance.
157 This method initializes database contents, database must be empty
158 prior to calling this method.
160 Parameters
161 ----------
162 config : `RegistryConfig` or `str`, optional
163 Registry configuration, if missing then default configuration will
164 be loaded from registry.yaml.
165 dimensionConfig : `DimensionConfig` or `str`, optional
166 Dimensions configuration, if missing then default configuration
167 will be loaded from dimensions.yaml.
168 butlerRoot : convertible to `lsst.resources.ResourcePath`, optional
169 Path to the repository root this `SqlRegistry` will manage.
171 Returns
172 -------
173 registry : `SqlRegistry`
174 A new `SqlRegistry` instance.
175 """
176 config = cls.forceRegistryConfig(config)
177 config.replaceRoot(butlerRoot)
179 if isinstance(dimensionConfig, str):
180 dimensionConfig = DimensionConfig(dimensionConfig)
181 elif dimensionConfig is None:
182 dimensionConfig = DimensionConfig()
183 elif not isinstance(dimensionConfig, DimensionConfig):
184 raise TypeError(f"Incompatible Dimension configuration type: {type(dimensionConfig)}")
186 DatabaseClass = config.getDatabaseClass()
187 database = DatabaseClass.fromUri(
188 config.connectionString,
189 origin=config.get("origin", 0),
190 namespace=config.get("namespace"),
191 allow_temporary_tables=config.areTemporaryTablesAllowed,
192 )
193 managerTypes = RegistryManagerTypes.fromConfig(config)
194 managers = managerTypes.makeRepo(database, dimensionConfig)
195 return cls(database, RegistryDefaults(), managers)
197 @classmethod
198 def fromConfig(
199 cls,
200 config: ButlerConfig | RegistryConfig | Config | str,
201 butlerRoot: ResourcePathExpression | None = None,
202 writeable: bool = True,
203 defaults: RegistryDefaults | None = None,
204 ) -> SqlRegistry:
205 """Create `Registry` subclass instance from `config`.
207 Registry database must be initialized prior to calling this method.
209 Parameters
210 ----------
211 config : `ButlerConfig`, `RegistryConfig`, `Config` or `str`
212 Registry configuration.
213 butlerRoot : `lsst.resources.ResourcePathExpression`, optional
214 Path to the repository root this `Registry` will manage.
215 writeable : `bool`, optional
216 If `True` (default) create a read-write connection to the database.
217 defaults : `RegistryDefaults`, optional
218 Default collection search path and/or output `~CollectionType.RUN`
219 collection.
221 Returns
222 -------
223 registry : `SqlRegistry`
224 A new `SqlRegistry` subclass instance.
225 """
226 config = cls.forceRegistryConfig(config)
227 config.replaceRoot(butlerRoot)
228 DatabaseClass = config.getDatabaseClass()
229 database = DatabaseClass.fromUri(
230 config.connectionString,
231 origin=config.get("origin", 0),
232 namespace=config.get("namespace"),
233 writeable=writeable,
234 allow_temporary_tables=config.areTemporaryTablesAllowed,
235 )
236 managerTypes = RegistryManagerTypes.fromConfig(config)
237 with database.session():
238 managers = managerTypes.loadRepo(database)
239 if defaults is None:
240 defaults = RegistryDefaults()
241 return cls(database, defaults, managers)
243 def __init__(
244 self,
245 database: Database,
246 defaults: RegistryDefaults,
247 managers: RegistryManagerInstances,
248 ):
249 self._db = database
250 self._managers = managers
251 self.storageClasses = StorageClassFactory()
252 # This is public to SqlRegistry's internal-to-daf_butler callers, but
253 # it is intentionally not part of RegistryShim.
254 self.dimension_record_cache = DimensionRecordCache(
255 self._managers.dimensions.universe,
256 fetch=self._managers.dimensions.fetch_cache_dict,
257 )
258 # Intentionally invoke property setter to initialize defaults. This
259 # can only be done after most of the rest of Registry has already been
260 # initialized, and must be done before the property getter is used.
261 self.defaults = defaults
262 # TODO: This is currently initialized by `make_datastore_tables`,
263 # eventually we'll need to do it during construction.
264 # The mapping is indexed by the opaque table name.
265 self._datastore_record_classes: Mapping[str, type[StoredDatastoreItemInfo]] = {}
267 def __str__(self) -> str:
268 return str(self._db)
270 def __repr__(self) -> str:
271 return f"SqlRegistry({self._db!r}, {self.dimensions!r})"
273 def isWriteable(self) -> bool:
274 """Return `True` if this registry allows write operations, and `False`
275 otherwise.
276 """
277 return self._db.isWriteable()
279 def copy(self, defaults: RegistryDefaults | None = None) -> SqlRegistry:
280 """Create a new `SqlRegistry` backed by the same data repository
281 as this one and sharing a database connection pool with it, but with
282 independent defaults and database sessions.
284 Parameters
285 ----------
286 defaults : `~lsst.daf.butler.registry.RegistryDefaults`, optional
287 Default collections and data ID values for the new registry. If
288 not provided, ``self.defaults`` will be used (but future changes
289 to either registry's defaults will not affect the other).
291 Returns
292 -------
293 copy : `SqlRegistry`
294 A new `SqlRegistry` instance with its own defaults.
295 """
296 if defaults is None:
297 # No need to copy, because `RegistryDefaults` is immutable; we
298 # effectively copy on write.
299 defaults = self.defaults
300 db = self._db.clone()
301 result = SqlRegistry(db, defaults, self._managers.clone(db))
302 result._datastore_record_classes = dict(self._datastore_record_classes)
303 result.dimension_record_cache.load_from(self.dimension_record_cache)
304 return result
306 @property
307 def dimensions(self) -> DimensionUniverse:
308 """Definitions of all dimensions recognized by this `Registry`
309 (`DimensionUniverse`).
310 """
311 return self._managers.dimensions.universe
313 @property
314 def defaults(self) -> RegistryDefaults:
315 """Default collection search path and/or output `~CollectionType.RUN`
316 collection (`~lsst.daf.butler.registry.RegistryDefaults`).
318 This is an immutable struct whose components may not be set
319 individually, but the entire struct can be set by assigning to this
320 property.
321 """
322 return self._defaults
324 @defaults.setter
325 def defaults(self, value: RegistryDefaults) -> None:
326 if value.run is not None:
327 self.registerRun(value.run)
328 value.finish(self)
329 self._defaults = value
331 def refresh(self) -> None:
332 """Refresh all in-memory state by querying the database.
334 This may be necessary to enable querying for entities added by other
335 registry instances after this one was constructed.
336 """
337 self.dimension_record_cache.reset()
338 with self._db.transaction():
339 self._managers.refresh()
341 def refresh_collection_summaries(self) -> None:
342 """Refresh content of the collection summary tables in the database.
344 This only cleans dataset type summaries, we may want to add cleanup of
345 governor summaries later.
346 """
347 for dataset_type in self.queryDatasetTypes():
348 self._managers.datasets.refresh_collection_summaries(dataset_type)
350 def caching_context(self) -> contextlib.AbstractContextManager[None]:
351 """Return context manager that enables caching.
353 Returns
354 -------
355 manager
356 A context manager that enables client-side caching. Entering
357 the context returns `None`.
358 """
359 return self._managers.caching_context_manager()
361 @contextlib.contextmanager
362 def transaction(self, *, savepoint: bool = False) -> Iterator[None]:
363 """Return a context manager that represents a transaction.
365 Parameters
366 ----------
367 savepoint : `bool`
368 Whether to issue a SAVEPOINT in the database.
370 Yields
371 ------
372 `None`
373 """
374 with self._db.transaction(savepoint=savepoint):
375 yield
377 def resetConnectionPool(self) -> None:
378 """Reset SQLAlchemy connection pool for `SqlRegistry` database.
380 This operation is useful when using registry with fork-based
381 multiprocessing. To use registry across fork boundary one has to make
382 sure that there are no currently active connections (no session or
383 transaction is in progress) and connection pool is reset using this
384 method. This method should be called by the child process immediately
385 after the fork.
386 """
387 self._db._engine.dispose()
389 def registerOpaqueTable(self, tableName: str, spec: ddl.TableSpec) -> None:
390 """Add an opaque (to the `Registry`) table for use by a `Datastore` or
391 other data repository client.
393 Opaque table records can be added via `insertOpaqueData`, retrieved via
394 `fetchOpaqueData`, and removed via `deleteOpaqueData`.
396 Parameters
397 ----------
398 tableName : `str`
399 Logical name of the opaque table. This may differ from the
400 actual name used in the database by a prefix and/or suffix.
401 spec : `ddl.TableSpec`
402 Specification for the table to be added.
403 """
404 self._managers.opaque.register(tableName, spec)
406 @transactional
407 def insertOpaqueData(self, tableName: str, *data: dict) -> None:
408 """Insert records into an opaque table.
410 Parameters
411 ----------
412 tableName : `str`
413 Logical name of the opaque table. Must match the name used in a
414 previous call to `registerOpaqueTable`.
415 *data
416 Each additional positional argument is a dictionary that represents
417 a single row to be added.
418 """
419 self._managers.opaque[tableName].insert(*data)
421 def fetchOpaqueData(self, tableName: str, **where: Any) -> Iterator[Mapping[str, Any]]:
422 """Retrieve records from an opaque table.
424 Parameters
425 ----------
426 tableName : `str`
427 Logical name of the opaque table. Must match the name used in a
428 previous call to `registerOpaqueTable`.
429 **where
430 Additional keyword arguments are interpreted as equality
431 constraints that restrict the returned rows (combined with AND);
432 keyword arguments are column names and values are the values they
433 must have.
435 Yields
436 ------
437 row : `dict`
438 A dictionary representing a single result row.
439 """
440 yield from self._managers.opaque[tableName].fetch(**where)
442 @transactional
443 def deleteOpaqueData(self, tableName: str, **where: Any) -> None:
444 """Remove records from an opaque table.
446 Parameters
447 ----------
448 tableName : `str`
449 Logical name of the opaque table. Must match the name used in a
450 previous call to `registerOpaqueTable`.
451 **where
452 Additional keyword arguments are interpreted as equality
453 constraints that restrict the deleted rows (combined with AND);
454 keyword arguments are column names and values are the values they
455 must have.
456 """
457 self._managers.opaque[tableName].delete(where.keys(), where)
459 def registerCollection(
460 self, name: str, type: CollectionType = CollectionType.TAGGED, doc: str | None = None
461 ) -> bool:
462 """Add a new collection if one with the given name does not exist.
464 Parameters
465 ----------
466 name : `str`
467 The name of the collection to create.
468 type : `CollectionType`
469 Enum value indicating the type of collection to create.
470 doc : `str`, optional
471 Documentation string for the collection.
473 Returns
474 -------
475 registered : `bool`
476 Boolean indicating whether the collection was already registered
477 or was created by this call.
479 Notes
480 -----
481 This method cannot be called within transactions, as it needs to be
482 able to perform its own transaction to be concurrent.
483 """
484 _, registered = self._managers.collections.register(name, type, doc=doc)
485 return registered
487 def getCollectionType(self, name: str) -> CollectionType:
488 """Return an enumeration value indicating the type of the given
489 collection.
491 Parameters
492 ----------
493 name : `str`
494 The name of the collection.
496 Returns
497 -------
498 type : `CollectionType`
499 Enum value indicating the type of this collection.
501 Raises
502 ------
503 lsst.daf.butler.registry.MissingCollectionError
504 Raised if no collection with the given name exists.
505 """
506 return self._managers.collections.find(name).type
508 def get_collection_record(self, name: str) -> CollectionRecord:
509 """Return the record for this collection.
511 Parameters
512 ----------
513 name : `str`
514 Name of the collection for which the record is to be retrieved.
516 Returns
517 -------
518 record : `CollectionRecord`
519 The record for this collection.
520 """
521 return self._managers.collections.find(name)
523 def registerRun(self, name: str, doc: str | None = None) -> bool:
524 """Add a new run if one with the given name does not exist.
526 Parameters
527 ----------
528 name : `str`
529 The name of the run to create.
530 doc : `str`, optional
531 Documentation string for the collection.
533 Returns
534 -------
535 registered : `bool`
536 Boolean indicating whether a new run was registered. `False`
537 if it already existed.
539 Notes
540 -----
541 This method cannot be called within transactions, as it needs to be
542 able to perform its own transaction to be concurrent.
543 """
544 _, registered = self._managers.collections.register(name, CollectionType.RUN, doc=doc)
545 return registered
547 @transactional
548 def removeCollection(self, name: str) -> None:
549 """Remove the given collection from the registry.
551 Parameters
552 ----------
553 name : `str`
554 The name of the collection to remove.
556 Raises
557 ------
558 lsst.daf.butler.registry.MissingCollectionError
559 Raised if no collection with the given name exists.
560 sqlalchemy.exc.IntegrityError
561 Raised if the database rows associated with the collection are
562 still referenced by some other table, such as a dataset in a
563 datastore (for `~CollectionType.RUN` collections only) or a
564 `~CollectionType.CHAINED` collection of which this collection is
565 a child.
567 Notes
568 -----
569 If this is a `~CollectionType.RUN` collection, all datasets and quanta
570 in it will removed from the `Registry` database. This requires that
571 those datasets be removed (or at least trashed) from any datastores
572 that hold them first.
574 A collection may not be deleted as long as it is referenced by a
575 `~CollectionType.CHAINED` collection; the ``CHAINED`` collection must
576 be deleted or redefined first.
577 """
578 self._managers.collections.remove(name)
580 def getCollectionChain(self, parent: str) -> tuple[str, ...]:
581 """Return the child collections in a `~CollectionType.CHAINED`
582 collection.
584 Parameters
585 ----------
586 parent : `str`
587 Name of the chained collection. Must have already been added via
588 a call to `Registry.registerCollection`.
590 Returns
591 -------
592 children : `~collections.abc.Sequence` [ `str` ]
593 An ordered sequence of collection names that are searched when the
594 given chained collection is searched.
596 Raises
597 ------
598 lsst.daf.butler.registry.MissingCollectionError
599 Raised if ``parent`` does not exist in the `Registry`.
600 lsst.daf.butler.registry.CollectionTypeError
601 Raised if ``parent`` does not correspond to a
602 `~CollectionType.CHAINED` collection.
603 """
604 record = self._managers.collections.find(parent)
605 if record.type is not CollectionType.CHAINED:
606 raise CollectionTypeError(f"Collection '{parent}' has type {record.type.name}, not CHAINED.")
607 assert isinstance(record, ChainedCollectionRecord)
608 return record.children
610 @transactional
611 def setCollectionChain(self, parent: str, children: Any, *, flatten: bool = False) -> None:
612 """Define or redefine a `~CollectionType.CHAINED` collection.
614 Parameters
615 ----------
616 parent : `str`
617 Name of the chained collection. Must have already been added via
618 a call to `Registry.registerCollection`.
619 children : collection expression
620 An expression defining an ordered search of child collections,
621 generally an iterable of `str`; see
622 :ref:`daf_butler_collection_expressions` for more information.
623 flatten : `bool`, optional
624 If `True` (`False` is default), recursively flatten out any nested
625 `~CollectionType.CHAINED` collections in ``children`` first.
627 Raises
628 ------
629 lsst.daf.butler.registry.MissingCollectionError
630 Raised when any of the given collections do not exist in the
631 `Registry`.
632 lsst.daf.butler.registry.CollectionTypeError
633 Raised if ``parent`` does not correspond to a
634 `~CollectionType.CHAINED` collection.
635 CollectionCycleError
636 Raised if the given collections contains a cycle.
638 Notes
639 -----
640 If this function is called within a call to ``Butler.transaction``, it
641 will hold a lock that prevents other processes from modifying the
642 parent collection until the end of the transaction. Keep these
643 transactions short.
644 """
645 children = CollectionWildcard.from_expression(children).require_ordered()
646 if flatten:
647 children = self.queryCollections(children, flattenChains=True)
649 self._managers.collections.update_chain(parent, list(children), allow_use_in_caching_context=True)
651 def getCollectionParentChains(self, collection: str) -> set[str]:
652 """Return the CHAINED collections that directly contain the given one.
654 Parameters
655 ----------
656 collection : `str`
657 Name of the collection.
659 Returns
660 -------
661 chains : `set` of `str`
662 Set of `~CollectionType.CHAINED` collection names.
663 """
664 return self._managers.collections.getParentChains(self._managers.collections.find(collection).key)
666 def getCollectionDocumentation(self, collection: str) -> str | None:
667 """Retrieve the documentation string for a collection.
669 Parameters
670 ----------
671 collection : `str`
672 Name of the collection.
674 Returns
675 -------
676 docs : `str` or `None`
677 Docstring for the collection with the given name.
678 """
679 return self._managers.collections.getDocumentation(self._managers.collections.find(collection).key)
681 def setCollectionDocumentation(self, collection: str, doc: str | None) -> None:
682 """Set the documentation string for a collection.
684 Parameters
685 ----------
686 collection : `str`
687 Name of the collection.
688 doc : `str` or `None`
689 Docstring for the collection with the given name; will replace any
690 existing docstring. Passing `None` will remove any existing
691 docstring.
692 """
693 self._managers.collections.setDocumentation(self._managers.collections.find(collection).key, doc)
695 def getCollectionSummary(self, collection: str) -> CollectionSummary:
696 """Return a summary for the given collection.
698 Parameters
699 ----------
700 collection : `str`
701 Name of the collection for which a summary is to be retrieved.
703 Returns
704 -------
705 summary : `~lsst.daf.butler.registry.CollectionSummary`
706 Summary of the dataset types and governor dimension values in
707 this collection.
708 """
709 record = self._managers.collections.find(collection)
710 return self._managers.datasets.getCollectionSummary(record)
712 def registerDatasetType(self, datasetType: DatasetType) -> bool:
713 """Add a new `DatasetType` to the Registry.
715 It is not an error to register the same `DatasetType` twice.
717 Parameters
718 ----------
719 datasetType : `DatasetType`
720 The `DatasetType` to be added.
722 Returns
723 -------
724 inserted : `bool`
725 `True` if ``datasetType`` was inserted, `False` if an identical
726 existing `DatasetType` was found. Note that in either case the
727 DatasetType is guaranteed to be defined in the Registry
728 consistently with the given definition.
730 Raises
731 ------
732 ValueError
733 Raised if the dimensions or storage class are invalid.
734 lsst.daf.butler.registry.ConflictingDefinitionError
735 Raised if this `DatasetType` is already registered with a different
736 definition.
738 Notes
739 -----
740 This method cannot be called within transactions, as it needs to be
741 able to perform its own transaction to be concurrent.
742 """
743 return self._managers.datasets.register_dataset_type(datasetType)
745 def removeDatasetType(self, name: str | tuple[str, ...]) -> None:
746 """Remove the named `DatasetType` from the registry.
748 .. warning::
750 Registry implementations can cache the dataset type definitions.
751 This means that deleting the dataset type definition may result in
752 unexpected behavior from other butler processes that are active
753 that have not seen the deletion.
755 Parameters
756 ----------
757 name : `str` or `tuple` [`str`]
758 Name of the type to be removed or tuple containing a list of type
759 names to be removed. Wildcards are allowed.
761 Raises
762 ------
763 lsst.daf.butler.registry.OrphanedRecordError
764 Raised if an attempt is made to remove the dataset type definition
765 when there are already datasets associated with it.
767 Notes
768 -----
769 If the dataset type is not registered the method will return without
770 action.
771 """
772 for datasetTypeExpression in ensure_iterable(name):
773 # Catch any warnings from the caller specifying a component
774 # dataset type. This will result in an error later but the
775 # warning could be confusing when the caller is not querying
776 # anything.
777 with warnings.catch_warnings():
778 warnings.simplefilter("ignore", category=FutureWarning)
779 datasetTypes = list(self.queryDatasetTypes(datasetTypeExpression))
780 if not datasetTypes:
781 _LOG.info("Dataset type %r not defined", datasetTypeExpression)
782 else:
783 for datasetType in datasetTypes:
784 self._managers.datasets.remove_dataset_type(datasetType.name)
785 _LOG.info("Removed dataset type %r", datasetType.name)
787 def getDatasetType(self, name: str) -> DatasetType:
788 """Get the `DatasetType`.
790 Parameters
791 ----------
792 name : `str`
793 Name of the type.
795 Returns
796 -------
797 type : `DatasetType`
798 The `DatasetType` associated with the given name.
800 Raises
801 ------
802 lsst.daf.butler.registry.MissingDatasetTypeError
803 Raised if the requested dataset type has not been registered.
805 Notes
806 -----
807 This method handles component dataset types automatically, though most
808 other registry operations do not.
809 """
810 parent_name, component = DatasetType.splitDatasetTypeName(name)
811 parent_dataset_type = self._managers.datasets.get_dataset_type(parent_name)
812 if component is None:
813 return parent_dataset_type
814 else:
815 return parent_dataset_type.makeComponentDatasetType(component)
817 def supportsIdGenerationMode(self, mode: DatasetIdGenEnum) -> bool:
818 """Test whether the given dataset ID generation mode is supported by
819 `insertDatasets`.
821 Parameters
822 ----------
823 mode : `DatasetIdGenEnum`
824 Enum value for the mode to test.
826 Returns
827 -------
828 supported : `bool`
829 Whether the given mode is supported.
830 """
831 return True
833 def findDataset(
834 self,
835 datasetType: DatasetType | str,
836 dataId: DataId | None = None,
837 *,
838 collections: CollectionArgType | None = None,
839 timespan: Timespan | None = None,
840 datastore_records: bool = False,
841 **kwargs: Any,
842 ) -> DatasetRef | None:
843 """Find a dataset given its `DatasetType` and data ID.
845 This can be used to obtain a `DatasetRef` that permits the dataset to
846 be read from a `Datastore`. If the dataset is a component and can not
847 be found using the provided dataset type, a dataset ref for the parent
848 will be returned instead but with the correct dataset type.
850 Parameters
851 ----------
852 datasetType : `DatasetType` or `str`
853 A `DatasetType` or the name of one. If this is a `DatasetType`
854 instance, its storage class will be respected and propagated to
855 the output, even if it differs from the dataset type definition
856 in the registry, as long as the storage classes are convertible.
857 dataId : `dict` or `DataCoordinate`, optional
858 A `dict`-like object containing the `Dimension` links that identify
859 the dataset within a collection.
860 collections : collection expression, optional
861 An expression that fully or partially identifies the collections to
862 search for the dataset; see
863 :ref:`daf_butler_collection_expressions` for more information.
864 Defaults to ``self.defaults.collections``.
865 timespan : `Timespan`, optional
866 A timespan that the validity range of the dataset must overlap.
867 If not provided, any `~CollectionType.CALIBRATION` collections
868 matched by the ``collections`` argument will not be searched.
869 datastore_records : `bool`, optional
870 Whether to attach datastore records to the `DatasetRef`.
871 **kwargs
872 Additional keyword arguments passed to
873 `DataCoordinate.standardize` to convert ``dataId`` to a true
874 `DataCoordinate` or augment an existing one.
876 Returns
877 -------
878 ref : `DatasetRef`
879 A reference to the dataset, or `None` if no matching Dataset
880 was found.
882 Raises
883 ------
884 lsst.daf.butler.registry.NoDefaultCollectionError
885 Raised if ``collections`` is `None` and
886 ``self.defaults.collections`` is `None`.
887 LookupError
888 Raised if one or more data ID keys are missing.
889 lsst.daf.butler.registry.MissingDatasetTypeError
890 Raised if the dataset type does not exist.
891 lsst.daf.butler.registry.MissingCollectionError
892 Raised if any of ``collections`` does not exist in the registry.
894 Notes
895 -----
896 This method simply returns `None` and does not raise an exception even
897 when the set of collections searched is intrinsically incompatible with
898 the dataset type, e.g. if ``datasetType.isCalibration() is False``, but
899 only `~CollectionType.CALIBRATION` collections are being searched.
900 This may make it harder to debug some lookup failures, but the behavior
901 is intentional; we consider it more important that failed searches are
902 reported consistently, regardless of the reason, and that adding
903 additional collections that do not contain a match to the search path
904 never changes the behavior.
906 This method handles component dataset types automatically, though most
907 other registry operations do not.
908 """
909 if collections is None:
910 if not self.defaults.collections:
911 raise NoDefaultCollectionError(
912 "No collections provided to findDataset, and no defaults from registry construction."
913 )
914 collections = self.defaults.collections
915 backend = queries.SqlQueryBackend(self._db, self._managers, self.dimension_record_cache)
916 with backend.caching_context():
917 collection_wildcard = CollectionWildcard.from_expression(collections, require_ordered=True)
918 if collection_wildcard.empty():
919 return None
920 matched_collections = backend.resolve_collection_wildcard(collection_wildcard)
921 resolved_dataset_type = backend.resolve_single_dataset_type_wildcard(datasetType)
922 dataId = DataCoordinate.standardize(
923 dataId,
924 dimensions=resolved_dataset_type.dimensions,
925 universe=self.dimensions,
926 defaults=self.defaults.dataId,
927 **kwargs,
928 )
929 governor_constraints = {name: {cast(str, dataId[name])} for name in dataId.dimensions.governors}
930 (filtered_collections,) = backend.filter_dataset_collections(
931 [resolved_dataset_type],
932 matched_collections,
933 governor_constraints=governor_constraints,
934 ).values()
935 if not filtered_collections:
936 return None
937 if timespan is None:
938 filtered_collections = [
939 collection_record
940 for collection_record in filtered_collections
941 if collection_record.type is not CollectionType.CALIBRATION
942 ]
943 if filtered_collections:
944 requested_columns = {"dataset_id", "run", "collection"}
945 with backend.context() as context:
946 predicate = context.make_data_coordinate_predicate(
947 dataId.subset(resolved_dataset_type.dimensions), full=False
948 )
949 if timespan is not None:
950 requested_columns.add("timespan")
951 predicate = predicate.logical_and(
952 context.make_timespan_overlap_predicate(
953 DatasetColumnTag(resolved_dataset_type.name, "timespan"), timespan
954 )
955 )
956 relation = backend.make_dataset_query_relation(
957 resolved_dataset_type, filtered_collections, requested_columns, context
958 ).with_rows_satisfying(predicate)
959 rows = list(context.fetch_iterable(relation))
960 else:
961 rows = []
962 if not rows:
963 return None
964 elif len(rows) == 1:
965 best_row = rows[0]
966 else:
967 rank_by_collection_key = {record.key: n for n, record in enumerate(filtered_collections)}
968 collection_tag = DatasetColumnTag(resolved_dataset_type.name, "collection")
969 row_iter = iter(rows)
970 best_row = next(row_iter)
971 best_rank = rank_by_collection_key[best_row[collection_tag]]
972 have_tie = False
973 for row in row_iter:
974 if (rank := rank_by_collection_key[row[collection_tag]]) < best_rank:
975 best_row = row
976 best_rank = rank
977 have_tie = False
978 elif rank == best_rank:
979 have_tie = True
980 assert timespan is not None, "Rank ties should be impossible given DB constraints."
981 if have_tie:
982 raise CalibrationLookupError(
983 f"Ambiguous calibration lookup for {resolved_dataset_type.name} in collections "
984 f"{collection_wildcard.strings} with timespan {timespan}."
985 )
986 reader = queries.DatasetRefReader(
987 resolved_dataset_type,
988 translate_collection=lambda k: self._managers.collections[k].name,
989 )
990 ref = reader.read(best_row, data_id=dataId)
991 if datastore_records:
992 ref = self.get_datastore_records(ref)
994 return ref
996 @transactional
997 def insertDatasets(
998 self,
999 datasetType: DatasetType | str,
1000 dataIds: Iterable[DataId],
1001 run: str | None = None,
1002 expand: bool = True,
1003 idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE,
1004 ) -> list[DatasetRef]:
1005 """Insert one or more datasets into the `Registry`.
1007 This always adds new datasets; to associate existing datasets with
1008 a new collection, use ``associate``.
1010 Parameters
1011 ----------
1012 datasetType : `DatasetType` or `str`
1013 A `DatasetType` or the name of one.
1014 dataIds : `~collections.abc.Iterable` of `dict` or `DataCoordinate`
1015 Dimension-based identifiers for the new datasets.
1016 run : `str`, optional
1017 The name of the run that produced the datasets. Defaults to
1018 ``self.defaults.run``.
1019 expand : `bool`, optional
1020 If `True` (default), expand data IDs as they are inserted. This is
1021 necessary in general to allow datastore to generate file templates,
1022 but it may be disabled if the caller can guarantee this is
1023 unnecessary.
1024 idGenerationMode : `DatasetIdGenEnum`, optional
1025 Specifies option for generating dataset IDs. By default unique IDs
1026 are generated for each inserted dataset.
1028 Returns
1029 -------
1030 refs : `list` of `DatasetRef`
1031 Resolved `DatasetRef` instances for all given data IDs (in the same
1032 order).
1034 Raises
1035 ------
1036 lsst.daf.butler.registry.DatasetTypeError
1037 Raised if ``datasetType`` is not known to registry.
1038 lsst.daf.butler.registry.CollectionTypeError
1039 Raised if ``run`` collection type is not `~CollectionType.RUN`.
1040 lsst.daf.butler.registry.NoDefaultCollectionError
1041 Raised if ``run`` is `None` and ``self.defaults.run`` is `None`.
1042 lsst.daf.butler.registry.ConflictingDefinitionError
1043 If a dataset with the same dataset type and data ID as one of those
1044 given already exists in ``run``.
1045 lsst.daf.butler.registry.MissingCollectionError
1046 Raised if ``run`` does not exist in the registry.
1047 """
1048 datasetType = self._managers.datasets.conform_exact_dataset_type(datasetType)
1049 if run is None:
1050 if self.defaults.run is None:
1051 raise NoDefaultCollectionError(
1052 "No run provided to insertDatasets, and no default from registry construction."
1053 )
1054 run = self.defaults.run
1055 runRecord = self._managers.collections.find(run)
1056 if runRecord.type is not CollectionType.RUN:
1057 raise CollectionTypeError(
1058 f"Given collection is of type {runRecord.type.name}; RUN collection required."
1059 )
1060 assert isinstance(runRecord, RunRecord)
1062 expandedDataIds = [
1063 DataCoordinate.standardize(dataId, dimensions=datasetType.dimensions) for dataId in dataIds
1064 ]
1065 if expand:
1066 _LOG.debug("Expanding %d data IDs", len(expandedDataIds))
1067 expandedDataIds = self.expand_data_ids(expandedDataIds)
1068 _LOG.debug("Finished expanding data IDs")
1070 try:
1071 refs = list(
1072 self._managers.datasets.insert(datasetType.name, runRecord, expandedDataIds, idGenerationMode)
1073 )
1074 if self._managers.obscore:
1075 self._managers.obscore.add_datasets(refs)
1076 except sqlalchemy.exc.IntegrityError as err:
1077 raise ConflictingDefinitionError(
1078 "A database constraint failure was triggered by inserting "
1079 f"one or more datasets of type {datasetType} into "
1080 f"collection '{run}'. "
1081 "This probably means a dataset with the same data ID "
1082 "and dataset type already exists, but it may also mean a "
1083 "dimension row is missing."
1084 ) from err
1085 return refs
1087 @transactional
1088 def _importDatasets(
1089 self,
1090 datasets: Iterable[DatasetRef],
1091 expand: bool = True,
1092 ) -> list[DatasetRef]:
1093 """Import one or more datasets into the `Registry`.
1095 Difference from `insertDatasets` method is that this method accepts
1096 `DatasetRef` instances which should already be resolved and have a
1097 dataset ID. If registry supports globally-unique dataset IDs (e.g.
1098 `uuid.UUID`) then datasets which already exist in the registry will be
1099 ignored if imported again.
1101 Parameters
1102 ----------
1103 datasets : `~collections.abc.Iterable` of `DatasetRef`
1104 Datasets to be inserted. All `DatasetRef` instances must have
1105 identical ``run`` attributes. ``run``
1106 attribute can be `None` and defaults to ``self.defaults.run``.
1107 Datasets can specify ``id`` attribute which will be used for
1108 inserted datasets.
1109 Datasets can be of multiple dataset types, but all the dataset
1110 types must have the same set of dimensions.
1111 expand : `bool`, optional
1112 If `True` (default), expand data IDs as they are inserted. This is
1113 necessary in general, but it may be disabled if the caller can
1114 guarantee this is unnecessary.
1116 Returns
1117 -------
1118 refs : `list` of `DatasetRef`
1119 Resolved `DatasetRef` instances for all given data IDs (in the same
1120 order). If any of ``datasets`` has an ID which already exists in
1121 the database then it will not be inserted or updated, but a
1122 resolved `DatasetRef` will be returned for it in any case.
1124 Raises
1125 ------
1126 lsst.daf.butler.registry.NoDefaultCollectionError
1127 Raised if ``run`` is `None` and ``self.defaults.run`` is `None`.
1128 lsst.daf.butler.registry.DatasetTypeError
1129 Raised if a dataset type is not known to registry.
1130 lsst.daf.butler.registry.ConflictingDefinitionError
1131 If a dataset with the same dataset type and data ID as one of those
1132 given already exists in ``run``.
1133 lsst.daf.butler.registry.MissingCollectionError
1134 Raised if ``run`` does not exist in the registry.
1136 Notes
1137 -----
1138 This method is considered package-private and internal to Butler
1139 implementation. Clients outside daf_butler package should not use this
1140 method.
1141 """
1142 datasets = list(datasets)
1143 if not datasets:
1144 # nothing to do
1145 return []
1147 # find run name
1148 runs = {dataset.run for dataset in datasets}
1149 if len(runs) != 1:
1150 raise ValueError(f"Multiple run names in input datasets: {runs}")
1151 run = runs.pop()
1153 runRecord = self._managers.collections.find(run)
1154 if runRecord.type is not CollectionType.RUN:
1155 raise CollectionTypeError(
1156 f"Given collection '{runRecord.name}' is of type {runRecord.type.name};"
1157 " RUN collection required."
1158 )
1159 assert isinstance(runRecord, RunRecord)
1161 if expand:
1162 _LOG.debug("Expanding %d data IDs", len(datasets))
1163 datasets = self.expand_refs(datasets)
1164 _LOG.debug("Finished expanding data IDs")
1166 try:
1167 self._managers.datasets.import_(runRecord, datasets)
1168 if self._managers.obscore:
1169 self._managers.obscore.add_datasets(datasets)
1170 except sqlalchemy.exc.IntegrityError as err:
1171 raise ConflictingDefinitionError(
1172 "A database constraint failure was triggered by inserting "
1173 f"one or more datasets into collection '{run}'. "
1174 "This probably means a dataset with the same data ID "
1175 "and dataset type already exists, but it may also mean a "
1176 "dimension row is missing."
1177 ) from err
1178 return datasets
1180 def getDataset(self, id: DatasetId) -> DatasetRef | None:
1181 """Retrieve a Dataset entry.
1183 Parameters
1184 ----------
1185 id : `DatasetId`
1186 The unique identifier for the dataset.
1188 Returns
1189 -------
1190 ref : `DatasetRef` or `None`
1191 A ref to the Dataset, or `None` if no matching Dataset
1192 was found.
1193 """
1194 return self._managers.datasets.getDatasetRef(id)
1196 @transactional
1197 def removeDatasets(self, refs: Iterable[DatasetRef]) -> None:
1198 """Remove datasets from the Registry.
1200 The datasets will be removed unconditionally from all collections.
1201 `Datastore` records will *not* be deleted; the caller is responsible
1202 for ensuring that the dataset has already been removed from all
1203 Datastores.
1205 Parameters
1206 ----------
1207 refs : `~collections.abc.Iterable` [`DatasetRef`]
1208 References to the datasets to be removed. Should be considered
1209 invalidated upon return.
1211 Raises
1212 ------
1213 lsst.daf.butler.registry.OrphanedRecordError
1214 Raised if any dataset is still present in any `Datastore`.
1215 """
1216 try:
1217 self._managers.datasets.delete(refs)
1218 except sqlalchemy.exc.IntegrityError as err:
1219 raise OrphanedRecordError(
1220 "One or more datasets is still present in one or more Datastores."
1221 ) from err
1223 @transactional
1224 def associate(self, collection: str, refs: Iterable[DatasetRef]) -> None:
1225 """Add existing datasets to a `~CollectionType.TAGGED` collection.
1227 If a DatasetRef with the same exact ID is already in a collection
1228 nothing is changed. If a `DatasetRef` with the same `DatasetType` and
1229 data ID but with different ID exists in the collection,
1230 `~lsst.daf.butler.registry.ConflictingDefinitionError` is raised.
1232 Parameters
1233 ----------
1234 collection : `str`
1235 Indicates the collection the datasets should be associated with.
1236 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
1237 An iterable of resolved `DatasetRef` instances that already exist
1238 in this `Registry`.
1240 Raises
1241 ------
1242 lsst.daf.butler.registry.ConflictingDefinitionError
1243 If a Dataset with the given `DatasetRef` already exists in the
1244 given collection.
1245 lsst.daf.butler.registry.MissingCollectionError
1246 Raised if ``collection`` does not exist in the registry.
1247 lsst.daf.butler.registry.CollectionTypeError
1248 Raise adding new datasets to the given ``collection`` is not
1249 allowed.
1250 """
1251 progress = Progress("lsst.daf.butler.Registry.associate", level=logging.DEBUG)
1252 collectionRecord = self._managers.collections.find(collection)
1253 for datasetType, refsForType in progress.iter_item_chunks(
1254 DatasetRef.iter_by_type(refs), desc="Associating datasets by type"
1255 ):
1256 try:
1257 self._managers.datasets.associate(datasetType, collectionRecord, refsForType)
1258 if self._managers.obscore:
1259 # If a TAGGED collection is being monitored by ObsCore
1260 # manager then we may need to save the dataset.
1261 self._managers.obscore.associate(refsForType, collectionRecord)
1262 except sqlalchemy.exc.IntegrityError as err:
1263 raise ConflictingDefinitionError(
1264 f"Constraint violation while associating dataset of type {datasetType.name} with "
1265 f"collection {collection}. This probably means that one or more datasets with the same "
1266 "dataset type and data ID already exist in the collection, but it may also indicate "
1267 "that the datasets do not exist."
1268 ) from err
1270 @transactional
1271 def disassociate(self, collection: str, refs: Iterable[DatasetRef]) -> None:
1272 """Remove existing datasets from a `~CollectionType.TAGGED` collection.
1274 ``collection`` and ``ref`` combinations that are not currently
1275 associated are silently ignored.
1277 Parameters
1278 ----------
1279 collection : `str`
1280 The collection the datasets should no longer be associated with.
1281 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
1282 An iterable of resolved `DatasetRef` instances that already exist
1283 in this `Registry`.
1285 Raises
1286 ------
1287 lsst.daf.butler.AmbiguousDatasetError
1288 Raised if any of the given dataset references is unresolved.
1289 lsst.daf.butler.registry.MissingCollectionError
1290 Raised if ``collection`` does not exist in the registry.
1291 lsst.daf.butler.registry.CollectionTypeError
1292 Raise adding new datasets to the given ``collection`` is not
1293 allowed.
1294 """
1295 progress = Progress("lsst.daf.butler.Registry.disassociate", level=logging.DEBUG)
1296 collectionRecord = self._managers.collections.find(collection)
1297 for datasetType, refsForType in progress.iter_item_chunks(
1298 DatasetRef.iter_by_type(refs), desc="Disassociating datasets by type"
1299 ):
1300 self._managers.datasets.disassociate(datasetType, collectionRecord, refsForType)
1301 if self._managers.obscore:
1302 self._managers.obscore.disassociate(refsForType, collectionRecord)
1304 @transactional
1305 def certify(self, collection: str, refs: Iterable[DatasetRef], timespan: Timespan) -> None:
1306 """Associate one or more datasets with a calibration collection and a
1307 validity range within it.
1309 Parameters
1310 ----------
1311 collection : `str`
1312 The name of an already-registered `~CollectionType.CALIBRATION`
1313 collection.
1314 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
1315 Datasets to be associated.
1316 timespan : `Timespan`
1317 The validity range for these datasets within the collection.
1319 Raises
1320 ------
1321 lsst.daf.butler.AmbiguousDatasetError
1322 Raised if any of the given `DatasetRef` instances is unresolved.
1323 lsst.daf.butler.registry.ConflictingDefinitionError
1324 Raised if the collection already contains a different dataset with
1325 the same `DatasetType` and data ID and an overlapping validity
1326 range.
1327 DatasetTypeError
1328 Raised if ``ref.datasetType.isCalibration() is False`` for any ref.
1329 CollectionTypeError
1330 Raised if
1331 ``collection.type is not CollectionType.CALIBRATION``.
1332 """
1333 progress = Progress("lsst.daf.butler.Registry.certify", level=logging.DEBUG)
1334 collectionRecord = self._managers.collections.find(collection)
1335 for datasetType, refsForType in progress.iter_item_chunks(
1336 DatasetRef.iter_by_type(refs), desc="Certifying datasets by type"
1337 ):
1338 self._managers.datasets.certify(
1339 datasetType,
1340 collectionRecord,
1341 refsForType,
1342 timespan,
1343 context=queries.SqlQueryContext(self._db, self._managers.column_types),
1344 )
1346 @transactional
1347 def decertify(
1348 self,
1349 collection: str,
1350 datasetType: str | DatasetType,
1351 timespan: Timespan,
1352 *,
1353 dataIds: Iterable[DataId] | None = None,
1354 ) -> None:
1355 """Remove or adjust datasets to clear a validity range within a
1356 calibration collection.
1358 Parameters
1359 ----------
1360 collection : `str`
1361 The name of an already-registered `~CollectionType.CALIBRATION`
1362 collection.
1363 datasetType : `str` or `DatasetType`
1364 Name or `DatasetType` instance for the datasets to be decertified.
1365 timespan : `Timespan`, optional
1366 The validity range to remove datasets from within the collection.
1367 Datasets that overlap this range but are not contained by it will
1368 have their validity ranges adjusted to not overlap it, which may
1369 split a single dataset validity range into two.
1370 dataIds : iterable [`dict` or `DataCoordinate`], optional
1371 Data IDs that should be decertified within the given validity range
1372 If `None`, all data IDs for ``self.datasetType`` will be
1373 decertified.
1375 Raises
1376 ------
1377 DatasetTypeError
1378 Raised if ``datasetType.isCalibration() is False``.
1379 CollectionTypeError
1380 Raised if
1381 ``collection.type is not CollectionType.CALIBRATION``.
1382 """
1383 collectionRecord = self._managers.collections.find(collection)
1384 if isinstance(datasetType, str):
1385 datasetType = self.getDatasetType(datasetType)
1386 standardizedDataIds = None
1387 if dataIds is not None:
1388 standardizedDataIds = [
1389 DataCoordinate.standardize(d, dimensions=datasetType.dimensions) for d in dataIds
1390 ]
1391 self._managers.datasets.decertify(
1392 datasetType,
1393 collectionRecord,
1394 timespan,
1395 data_ids=standardizedDataIds,
1396 context=queries.SqlQueryContext(self._db, self._managers.column_types),
1397 )
1399 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager:
1400 """Return an object that allows a new `Datastore` instance to
1401 communicate with this `Registry`.
1403 Returns
1404 -------
1405 manager : `~.interfaces.DatastoreRegistryBridgeManager`
1406 Object that mediates communication between this `Registry` and its
1407 associated datastores.
1408 """
1409 return self._managers.datastores
1411 def getDatasetLocations(self, ref: DatasetRef) -> Iterable[str]:
1412 """Retrieve datastore locations for a given dataset.
1414 Parameters
1415 ----------
1416 ref : `DatasetRef`
1417 A reference to the dataset for which to retrieve storage
1418 information.
1420 Returns
1421 -------
1422 datastores : `~collections.abc.Iterable` [ `str` ]
1423 All the matching datastores holding this dataset.
1425 Raises
1426 ------
1427 lsst.daf.butler.AmbiguousDatasetError
1428 Raised if ``ref.id`` is `None`.
1429 """
1430 return self._managers.datastores.findDatastores(ref)
1432 def expandDataId(
1433 self,
1434 dataId: DataId | None = None,
1435 *,
1436 dimensions: Iterable[str] | DimensionGroup | None = None,
1437 records: Mapping[str, DimensionRecord | None] | None = None,
1438 withDefaults: bool = True,
1439 **kwargs: Any,
1440 ) -> DataCoordinate:
1441 """Expand a dimension-based data ID to include additional information.
1443 Parameters
1444 ----------
1445 dataId : `DataCoordinate` or `dict`, optional
1446 Data ID to be expanded; augmented and overridden by ``kwargs``.
1447 dimensions : `~collections.abc.Iterable` [ `str` ], \
1448 `DimensionGroup`, optional
1449 The dimensions to be identified by the new `DataCoordinate`.
1450 If not provided, will be inferred from the keys of ``dataId`` and
1451 ``**kwargs``, and ``universe`` must be provided unless ``dataId``
1452 is already a `DataCoordinate`.
1453 records : `~collections.abc.Mapping` [`str`, `DimensionRecord`], \
1454 optional
1455 Dimension record data to use before querying the database for that
1456 data, keyed by element name.
1457 withDefaults : `bool`, optional
1458 Utilize ``self.defaults.dataId`` to fill in missing governor
1459 dimension key-value pairs. Defaults to `True` (i.e. defaults are
1460 used).
1461 **kwargs
1462 Additional keywords are treated like additional key-value pairs for
1463 ``dataId``, extending and overriding.
1465 Returns
1466 -------
1467 expanded : `DataCoordinate`
1468 A data ID that includes full metadata for all of the dimensions it
1469 identifies, i.e. guarantees that ``expanded.hasRecords()`` and
1470 ``expanded.hasFull()`` both return `True`.
1472 Raises
1473 ------
1474 lsst.daf.butler.registry.DataIdError
1475 Raised when ``dataId`` or keyword arguments specify unknown
1476 dimensions or values, or when a resulting data ID contains
1477 contradictory key-value pairs, according to dimension
1478 relationships.
1480 Notes
1481 -----
1482 This method cannot be relied upon to reject invalid data ID values
1483 for dimensions that do actually not have any record columns. For
1484 efficiency reasons the records for these dimensions (which have only
1485 dimension key values that are given by the caller) may be constructed
1486 directly rather than obtained from the registry database.
1487 """
1488 if not withDefaults:
1489 defaults = None
1490 else:
1491 defaults = self.defaults.dataId
1492 standardized = DataCoordinate.standardize(
1493 dataId,
1494 dimensions=dimensions,
1495 universe=self.dimensions,
1496 defaults=defaults,
1497 **kwargs,
1498 )
1499 if standardized.hasRecords():
1500 return standardized
1501 if records is None:
1502 records = {}
1503 else:
1504 records = dict(records)
1505 if isinstance(dataId, DataCoordinate) and dataId.hasRecords():
1506 for element_name in dataId.dimensions.elements:
1507 records[element_name] = dataId.records[element_name]
1508 keys: dict[str, str | int] = dict(standardized.mapping)
1509 for element_name in standardized.dimensions.lookup_order:
1510 element = self.dimensions[element_name]
1511 record = records.get(element_name, ...) # Use ... to mean not found; None might mean NULL
1512 if record is ...:
1513 if element_name in self.dimensions.dimensions.names and keys.get(element_name) is None:
1514 raise DimensionNameError(f"No value or null value for dimension {element_name}.")
1515 else:
1516 record = self._managers.dimensions.fetch_one(
1517 element_name,
1518 DataCoordinate.standardize(keys, dimensions=element.minimal_group),
1519 self.dimension_record_cache,
1520 )
1521 records[element_name] = record
1522 if record is not None:
1523 for d in element.implied:
1524 value = getattr(record, d.name)
1525 if keys.setdefault(d.name, value) != value:
1526 raise InconsistentDataIdError(
1527 f"Data ID {standardized} has {d.name}={keys[d.name]!r}, "
1528 f"but {element_name} implies {d.name}={value!r}."
1529 )
1530 else:
1531 if element_name in standardized.dimensions.names:
1532 raise DataIdValueError(
1533 f"Could not fetch record for dimension {element.name} via keys {keys}."
1534 )
1535 if element.defines_relationships:
1536 raise InconsistentDataIdError(
1537 f"Could not fetch record for element {element_name} via keys {keys}, "
1538 "but it is marked as defining relationships; this means one or more dimensions are "
1539 "have inconsistent values.",
1540 )
1541 return DataCoordinate.standardize(keys, dimensions=standardized.dimensions).expanded(records=records)
1543 def expand_data_ids(self, data_ids: Iterable[DataCoordinate]) -> list[DataCoordinate]:
1544 output = list(data_ids)
1546 grouped_by_dimensions: defaultdict[DimensionGroup, list[int]] = defaultdict(list)
1547 for i, data_id in enumerate(data_ids):
1548 if not data_id.hasRecords():
1549 grouped_by_dimensions[data_id.dimensions].append(i)
1551 if not grouped_by_dimensions:
1552 # All given DataCoordinate values are already expanded.
1553 return output
1555 attacher = DimensionDataAttacher(
1556 cache=self.dimension_record_cache,
1557 dimensions=DimensionGroup.union(*grouped_by_dimensions.keys(), universe=self.dimensions),
1558 )
1559 with self._query() as query:
1560 for dimensions, indexes in grouped_by_dimensions.items():
1561 expanded = attacher.attach(dimensions, (output[index] for index in indexes), query)
1562 for index, data_id in zip(indexes, expanded):
1563 output[index] = data_id
1565 return output
1567 def expand_refs(self, dataset_refs: list[DatasetRef]) -> list[DatasetRef]:
1568 expanded_ids = self.expand_data_ids([ref.dataId for ref in dataset_refs])
1569 return [ref.expanded(data_id) for ref, data_id in zip(dataset_refs, expanded_ids)]
1571 def insertDimensionData(
1572 self,
1573 element: DimensionElement | str,
1574 *data: Mapping[str, Any] | DimensionRecord,
1575 conform: bool = True,
1576 replace: bool = False,
1577 skip_existing: bool = False,
1578 ) -> None:
1579 """Insert one or more dimension records into the database.
1581 Parameters
1582 ----------
1583 element : `DimensionElement` or `str`
1584 The `DimensionElement` or name thereof that identifies the table
1585 records will be inserted into.
1586 *data : `dict` or `DimensionRecord`
1587 One or more records to insert.
1588 conform : `bool`, optional
1589 If `False` (`True` is default) perform no checking or conversions,
1590 and assume that ``element`` is a `DimensionElement` instance and
1591 ``data`` is a one or more `DimensionRecord` instances of the
1592 appropriate subclass.
1593 replace : `bool`, optional
1594 If `True` (`False` is default), replace existing records in the
1595 database if there is a conflict.
1596 skip_existing : `bool`, optional
1597 If `True` (`False` is default), skip insertion if a record with
1598 the same primary key values already exists. Unlike
1599 `syncDimensionData`, this will not detect when the given record
1600 differs from what is in the database, and should not be used when
1601 this is a concern.
1602 """
1603 if isinstance(element, str):
1604 element = self.dimensions[element]
1605 if conform:
1606 records = [
1607 row if isinstance(row, DimensionRecord) else element.RecordClass(**row) for row in data
1608 ]
1609 else:
1610 # Ignore typing since caller said to trust them with conform=False.
1611 records = data # type: ignore
1612 if element.name in self.dimension_record_cache:
1613 self.dimension_record_cache.reset()
1614 self._managers.dimensions.insert(
1615 element,
1616 *records,
1617 replace=replace,
1618 skip_existing=skip_existing,
1619 )
1621 def syncDimensionData(
1622 self,
1623 element: DimensionElement | str,
1624 row: Mapping[str, Any] | DimensionRecord,
1625 conform: bool = True,
1626 update: bool = False,
1627 ) -> bool | dict[str, Any]:
1628 """Synchronize the given dimension record with the database, inserting
1629 if it does not already exist and comparing values if it does.
1631 Parameters
1632 ----------
1633 element : `DimensionElement` or `str`
1634 The `DimensionElement` or name thereof that identifies the table
1635 records will be inserted into.
1636 row : `dict` or `DimensionRecord`
1637 The record to insert.
1638 conform : `bool`, optional
1639 If `False` (`True` is default) perform no checking or conversions,
1640 and assume that ``element`` is a `DimensionElement` instance and
1641 ``data`` is a `DimensionRecord` instances of the appropriate
1642 subclass.
1643 update : `bool`, optional
1644 If `True` (`False` is default), update the existing record in the
1645 database if there is a conflict.
1647 Returns
1648 -------
1649 inserted_or_updated : `bool` or `dict`
1650 `True` if a new row was inserted, `False` if no changes were
1651 needed, or a `dict` mapping updated column names to their old
1652 values if an update was performed (only possible if
1653 ``update=True``).
1655 Raises
1656 ------
1657 lsst.daf.butler.registry.ConflictingDefinitionError
1658 Raised if the record exists in the database (according to primary
1659 key lookup) but is inconsistent with the given one.
1660 """
1661 if conform:
1662 if isinstance(element, str):
1663 element = self.dimensions[element]
1664 record = row if isinstance(row, DimensionRecord) else element.RecordClass(**row)
1665 else:
1666 # Ignore typing since caller said to trust them with conform=False.
1667 record = row # type: ignore
1668 if record.definition.name in self.dimension_record_cache:
1669 self.dimension_record_cache.reset()
1670 return self._managers.dimensions.sync(record, update=update)
1672 def queryDatasetTypes(
1673 self,
1674 expression: Any = ...,
1675 *,
1676 missing: list[str] | None = None,
1677 ) -> Iterable[DatasetType]:
1678 """Iterate over the dataset types whose names match an expression.
1680 Parameters
1681 ----------
1682 expression : dataset type expression, optional
1683 An expression that fully or partially identifies the dataset types
1684 to return, such as a `str`, `re.Pattern`, or iterable thereof.
1685 ``...`` can be used to return all dataset types, and is the
1686 default. See :ref:`daf_butler_dataset_type_expressions` for more
1687 information.
1688 missing : `list` of `str`, optional
1689 String dataset type names that were explicitly given (i.e. not
1690 regular expression patterns) but not found will be appended to this
1691 list, if it is provided.
1693 Returns
1694 -------
1695 dataset_types : `~collections.abc.Iterable` [ `DatasetType`]
1696 An `~collections.abc.Iterable` of `DatasetType` instances whose
1697 names match ``expression``.
1699 Raises
1700 ------
1701 lsst.daf.butler.registry.DatasetTypeExpressionError
1702 Raised when ``expression`` is invalid.
1703 """
1704 wildcard = DatasetTypeWildcard.from_expression(expression)
1705 return self._managers.datasets.resolve_wildcard(wildcard, missing=missing)
1707 def queryCollections(
1708 self,
1709 expression: Any = ...,
1710 datasetType: DatasetType | None = None,
1711 collectionTypes: Iterable[CollectionType] | CollectionType = CollectionType.all(),
1712 flattenChains: bool = False,
1713 includeChains: bool | None = None,
1714 ) -> Sequence[str]:
1715 """Iterate over the collections whose names match an expression.
1717 Parameters
1718 ----------
1719 expression : collection expression, optional
1720 An expression that identifies the collections to return, such as
1721 a `str` (for full matches or partial matches via globs),
1722 `re.Pattern` (for partial matches), or iterable thereof. ``...``
1723 can be used to return all collections, and is the default.
1724 See :ref:`daf_butler_collection_expressions` for more information.
1725 datasetType : `DatasetType`, optional
1726 If provided, only yield collections that may contain datasets of
1727 this type. This is a conservative approximation in general; it may
1728 yield collections that do not have any such datasets.
1729 collectionTypes : `~collections.abc.Set` [`CollectionType`] or \
1730 `CollectionType`, optional
1731 If provided, only yield collections of these types.
1732 flattenChains : `bool`, optional
1733 If `True` (`False` is default), recursively yield the child
1734 collections of matching `~CollectionType.CHAINED` collections.
1735 includeChains : `bool`, optional
1736 If `True`, yield records for matching `~CollectionType.CHAINED`
1737 collections. Default is the opposite of ``flattenChains``: include
1738 either CHAINED collections or their children, but not both.
1740 Returns
1741 -------
1742 collections : `~collections.abc.Sequence` [ `str` ]
1743 The names of collections that match ``expression``.
1745 Raises
1746 ------
1747 lsst.daf.butler.registry.CollectionExpressionError
1748 Raised when ``expression`` is invalid.
1750 Notes
1751 -----
1752 The order in which collections are returned is unspecified, except that
1753 the children of a `~CollectionType.CHAINED` collection are guaranteed
1754 to be in the order in which they are searched. When multiple parent
1755 `~CollectionType.CHAINED` collections match the same criteria, the
1756 order in which the two lists appear is unspecified, and the lists of
1757 children may be incomplete if a child has multiple parents.
1758 """
1759 # Right now the datasetTypes argument is completely ignored, but that
1760 # is consistent with its [lack of] guarantees. DM-24939 or a follow-up
1761 # ticket will take care of that.
1762 if datasetType is not None:
1763 warnings.warn(
1764 "The datasetType parameter should no longer be used. It has"
1765 " never had any effect. Will be removed after v28",
1766 FutureWarning,
1767 )
1768 try:
1769 wildcard = CollectionWildcard.from_expression(expression)
1770 except TypeError as exc:
1771 raise CollectionExpressionError(f"Invalid collection expression '{expression}'") from exc
1772 collectionTypes = ensure_iterable(collectionTypes)
1773 return [
1774 record.name
1775 for record in self._managers.collections.resolve_wildcard(
1776 wildcard,
1777 collection_types=frozenset(collectionTypes),
1778 flatten_chains=flattenChains,
1779 include_chains=includeChains,
1780 )
1781 ]
1783 @contextlib.contextmanager
1784 def _query(self) -> Iterator[Query]:
1785 """Context manager returning a `Query` object used for construction
1786 and execution of complex queries.
1787 """
1788 with self._query_driver(self.defaults.collections, self.defaults.dataId) as driver:
1789 yield Query(driver)
1791 @contextlib.contextmanager
1792 def _query_driver(
1793 self,
1794 default_collections: Iterable[str],
1795 default_data_id: DataCoordinate,
1796 ) -> Iterator[DirectQueryDriver]:
1797 """Set up a `QueryDriver` instance for query execution."""
1798 # Query internals do repeated lookups of the same collections, so it
1799 # benefits from the collection record cache.
1800 with self._managers.caching_context.enable_collection_record_cache():
1801 driver = DirectQueryDriver(
1802 self._db,
1803 self.dimensions,
1804 self._managers,
1805 self.dimension_record_cache,
1806 default_collections=default_collections,
1807 default_data_id=default_data_id,
1808 )
1809 with driver:
1810 yield driver
1812 def get_datastore_records(self, ref: DatasetRef) -> DatasetRef:
1813 """Retrieve datastore records for given ref.
1815 Parameters
1816 ----------
1817 ref : `DatasetRef`
1818 Dataset reference for which to retrieve its corresponding datastore
1819 records.
1821 Returns
1822 -------
1823 updated_ref : `DatasetRef`
1824 Dataset reference with filled datastore records.
1826 Notes
1827 -----
1828 If this method is called with the dataset ref that is not known to the
1829 registry then the reference with an empty set of records is returned.
1830 """
1831 datastore_records: dict[str, list[StoredDatastoreItemInfo]] = {}
1832 for opaque, record_class in self._datastore_record_classes.items():
1833 records = self.fetchOpaqueData(opaque, dataset_id=ref.id)
1834 datastore_records[opaque] = [record_class.from_record(record) for record in records]
1835 return ref.replace(datastore_records=datastore_records)
1837 def store_datastore_records(self, refs: Mapping[str, DatasetRef]) -> None:
1838 """Store datastore records for given refs.
1840 Parameters
1841 ----------
1842 refs : `~collections.abc.Mapping` [`str`, `DatasetRef`]
1843 Mapping of a datastore name to dataset reference stored in that
1844 datastore, reference must include datastore records.
1845 """
1846 for datastore_name, ref in refs.items():
1847 # Store ref IDs in the bridge table.
1848 bridge = self._managers.datastores.register(datastore_name)
1849 bridge.insert([ref])
1851 # store records in opaque tables
1852 assert ref._datastore_records is not None, "Dataset ref must have datastore records"
1853 for table_name, records in ref._datastore_records.items():
1854 opaque_table = self._managers.opaque.get(table_name)
1855 assert opaque_table is not None, f"Unexpected opaque table name {table_name}"
1856 opaque_table.insert(*(record.to_record(dataset_id=ref.id) for record in records))
1858 def make_datastore_tables(self, tables: Mapping[str, DatastoreOpaqueTable]) -> None:
1859 """Create opaque tables used by datastores.
1861 Parameters
1862 ----------
1863 tables : `~collections.abc.Mapping`
1864 Maps opaque table name to its definition.
1866 Notes
1867 -----
1868 This method should disappear in the future when opaque table
1869 definitions will be provided during `Registry` construction.
1870 """
1871 datastore_record_classes = {}
1872 for table_name, table_def in tables.items():
1873 datastore_record_classes[table_name] = table_def.record_class
1874 try:
1875 self._managers.opaque.register(table_name, table_def.table_spec)
1876 except ReadOnlyDatabaseError:
1877 # If the database is read only and we just tried and failed to
1878 # create a table, it means someone is trying to create a
1879 # read-only butler client for an empty repo. That should be
1880 # okay, as long as they then try to get any datasets before
1881 # some other client creates the table. Chances are they're
1882 # just validating configuration.
1883 pass
1884 self._datastore_record_classes = datastore_record_classes
1886 def preload_cache(self, *, load_dimension_record_cache: bool) -> None:
1887 """Immediately load caches that are used for common operations.
1889 Parameters
1890 ----------
1891 load_dimension_record_cache : `bool`
1892 If True, preload the dimension record cache. When this cache is
1893 preloaded, subsequent external changes to governor dimension
1894 records will not be visible to this Butler.
1895 """
1896 self._managers.datasets.preload_cache()
1898 if load_dimension_record_cache:
1899 self.dimension_record_cache.preload_cache()
1901 @property
1902 def obsCoreTableManager(self) -> ObsCoreTableManager | None:
1903 """The ObsCore manager instance for this registry
1904 (`~.interfaces.ObsCoreTableManager`
1905 or `None`).
1907 ObsCore manager may not be implemented for all registry backend, or
1908 may not be enabled for many repositories.
1909 """
1910 return self._managers.obscore
1912 storageClasses: StorageClassFactory
1913 """All storage classes known to the registry (`StorageClassFactory`).
1914 """
1916 _defaults: RegistryDefaults
1917 """Default collections used for registry queries (`RegistryDefaults`)."""