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