Coverage for python/lsst/daf/butler/registry/managers.py: 32%
168 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 01:27 -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__ = (
33 "RegistryManagerInstances",
34 "RegistryManagerTypes",
35)
37import dataclasses
38import logging
39from collections.abc import Iterator, Mapping
40from contextlib import contextmanager
41from typing import Any, Generic, TypeVar
43from lsst.utils import doImportType
45from .._column_type_info import ColumnTypeInfo
46from .._config import Config
47from ..dimensions import DimensionConfig, DimensionUniverse
48from ._caching_context import CachingContext
49from ._config import RegistryConfig
50from .interfaces import (
51 ButlerAttributeManager,
52 CollectionManager,
53 Database,
54 DatasetRecordStorageManager,
55 DatastoreRegistryBridgeManager,
56 DimensionRecordStorageManager,
57 ObsCoreTableManager,
58 OpaqueTableStorageManager,
59 StaticTablesContext,
60 VersionedExtension,
61 VersionTuple,
62)
63from .versions import ButlerVersionsManager
65_Attributes = TypeVar("_Attributes")
66_Dimensions = TypeVar("_Dimensions")
67_Collections = TypeVar("_Collections")
68_Datasets = TypeVar("_Datasets")
69_Opaque = TypeVar("_Opaque")
70_Datastores = TypeVar("_Datastores")
71_ObsCore = TypeVar("_ObsCore")
74_LOG = logging.getLogger(__name__)
76# key for dimensions configuration in attributes table
77_DIMENSIONS_ATTR = "config:dimensions.json"
79# key for obscore configuration in attributes table
80_OBSCORE_ATTR = "config:obscore.json"
83@dataclasses.dataclass(frozen=True, eq=False)
84class _GenericRegistryManagers(
85 Generic[_Attributes, _Dimensions, _Collections, _Datasets, _Opaque, _Datastores, _ObsCore]
86):
87 """Base struct used to pass around the manager instances or types that back
88 a `Registry`.
90 This class should only be used via its non-generic subclasses,
91 `RegistryManagerInstances` and `RegistryManagerTypes`.
92 """
94 attributes: _Attributes
95 """Manager for flat key-value pairs, including versions.
96 """
98 dimensions: _Dimensions
99 """Manager for dimensions.
100 """
102 collections: _Collections
103 """Manager for collections.
104 """
106 datasets: _Datasets
107 """Manager for datasets, dataset types, and collection summaries.
108 """
110 opaque: _Opaque
111 """Manager for opaque (to the Registry) tables.
112 """
114 datastores: _Datastores
115 """Manager for the interface between `Registry` and `Datastore`.
116 """
118 obscore: _ObsCore | None
119 """Manager for `ObsCore` table(s).
120 """
123@dataclasses.dataclass(frozen=True, eq=False)
124class RegistryManagerTypes(
125 _GenericRegistryManagers[
126 type[ButlerAttributeManager],
127 type[DimensionRecordStorageManager],
128 type[CollectionManager],
129 type[DatasetRecordStorageManager],
130 type[OpaqueTableStorageManager],
131 type[DatastoreRegistryBridgeManager],
132 type[ObsCoreTableManager],
133 ]
134):
135 """A struct used to pass around the types of the manager objects that back
136 a `Registry`.
137 """
139 @classmethod
140 def fromConfig(cls, config: RegistryConfig) -> RegistryManagerTypes:
141 """Construct by extracting class names from configuration and importing
142 them.
144 Parameters
145 ----------
146 config : `RegistryConfig`
147 Configuration object with a "managers" section that contains all
148 fully-qualified class names for all manager types.
150 Returns
151 -------
152 types : `RegistryManagerTypes`
153 A new struct containing type objects.
154 """
155 # We only check for manager names defined in class attributes.
156 # TODO: Maybe we need to check keys for unknown names/typos?
157 managers = {field.name for field in dataclasses.fields(cls)} - {"manager_configs", "schema_versions"}
158 # Values of "config" sub-key, if any, indexed by manager name.
159 configs: dict[str, Mapping] = {}
160 schema_versions: dict[str, VersionTuple] = {}
161 manager_types: dict[str, type] = {}
162 for manager in managers:
163 manager_config = config["managers"].get(manager)
164 if isinstance(manager_config, Config):
165 # Expect "cls" and optional "config" and "schema_version"
166 # sub-keys.
167 manager_config_dict = manager_config.toDict()
168 try:
169 class_name = manager_config_dict.pop("cls")
170 except KeyError:
171 raise KeyError(f"'cls' key is not defined in {manager!r} manager configuration") from None
172 if (mgr_config := manager_config_dict.pop("config", None)) is not None:
173 configs[manager] = mgr_config
174 if (mgr_version := manager_config_dict.pop("schema_version", None)) is not None:
175 # Note that we do not check versions that come from config
176 # for compatibility, they may be overriden later by
177 # versions from registry.
178 schema_versions[manager] = VersionTuple.fromString(mgr_version)
179 if manager_config_dict:
180 raise ValueError(
181 f"{manager!r} manager configuration has unexpected keys: {set(manager_config_dict)}"
182 )
183 elif isinstance(manager_config, str):
184 class_name = manager_config
185 elif manager_config is None:
186 # Some managers may be optional.
187 continue
188 else:
189 raise KeyError(f"Unexpected type of {manager!r} manager configuration: {manager_config!r}")
190 manager_types[manager] = doImportType(class_name)
192 # obscore need special care because it's the only manager which can be
193 # None, and we cannot define default value for it.
194 if "obscore" in manager_types:
195 return cls(**manager_types, manager_configs=configs, schema_versions=schema_versions)
196 else:
197 return cls(
198 **manager_types, obscore=None, manager_configs=configs, schema_versions=schema_versions
199 )
201 def makeRepo(self, database: Database, dimensionConfig: DimensionConfig) -> RegistryManagerInstances:
202 """Create all persistent `Registry` state for a new, empty data
203 repository, and return a new struct containing manager instances.
205 Parameters
206 ----------
207 database : `Database`
208 Object that represents a connection to the SQL database that will
209 back the data repository. Must point to an empty namespace, or at
210 least one with no tables or other entities whose names might clash
211 with those used by butler.
212 dimensionConfig : `DimensionConfig`
213 Configuration that defines a `DimensionUniverse`, to be written
214 into the data repository and used to define aspects of the schema.
216 Returns
217 -------
218 instances : `RegistryManagerInstances`
219 Struct containing instances of the types contained by ``self``,
220 pointing to the new repository and backed by ``database``.
221 """
222 # If schema versions were specified in the config, check that they are
223 # compatible with their managers.
224 managers = self.as_dict()
225 for manager_type, schema_version in self.schema_versions.items():
226 manager_class = managers[manager_type]
227 manager_class.checkNewSchemaVersion(schema_version)
229 universe = DimensionUniverse(dimensionConfig)
230 with database.declareStaticTables(create=True) as context:
231 instances = RegistryManagerInstances.initialize(database, context, types=self, universe=universe)
233 # store managers and their versions in attributes table
234 versions = ButlerVersionsManager(instances.attributes)
235 versions.storeManagersConfig(instances.as_dict())
237 # dump universe config as json into attributes (faster than YAML)
238 json = dimensionConfig.dump(format="json")
239 if json is not None:
240 instances.attributes.set(_DIMENSIONS_ATTR, json)
241 else:
242 raise RuntimeError("Unexpectedly failed to serialize DimensionConfig to JSON")
243 if instances.obscore is not None:
244 json = instances.obscore.config_json()
245 instances.attributes.set(_OBSCORE_ATTR, json)
246 return instances
248 def loadRepo(self, database: Database) -> RegistryManagerInstances:
249 """Construct manager instances that point to an existing data
250 repository.
252 Parameters
253 ----------
254 database : `Database`
255 Object that represents a connection to the SQL database that backs
256 the data repository. Must point to a namespace that already holds
257 all tables and other persistent entities used by butler.
259 Returns
260 -------
261 instances : `RegistryManagerInstances`
262 Struct containing instances of the types contained by ``self``,
263 pointing to the new repository and backed by ``database``.
264 """
265 # Create attributes manager only first, so we can use it to load the
266 # embedded dimensions configuration. Note that we do not check this
267 # manager version before initializing it, it is supposed to be
268 # completely backward- and forward-compatible.
269 with database.declareStaticTables(create=False) as context:
270 attributes = self.attributes.initialize(database, context)
272 # Verify that configured classes are compatible with the ones stored
273 # in registry.
274 versions = ButlerVersionsManager(attributes)
275 versions.checkManagersConfig(self.as_dict())
277 # Read schema versions from registry and validate them.
278 self.schema_versions.update(versions.managerVersions())
279 for manager_type, manager_class in self.as_dict().items():
280 schema_version = self.schema_versions.get(manager_type)
281 if schema_version is not None:
282 manager_class.checkCompatibility(schema_version, database.isWriteable())
284 # get serialized as a string from database
285 dimensionsString = attributes.get(_DIMENSIONS_ATTR)
286 if dimensionsString is not None:
287 dimensionConfig = DimensionConfig(Config.fromString(dimensionsString, format="json"))
288 else:
289 raise LookupError(f"Registry attribute {_DIMENSIONS_ATTR} is missing from database")
290 universe = DimensionUniverse(dimensionConfig)
291 if self.obscore is not None:
292 # Get ObsCore configuration from attributes table, this silently
293 # overrides whatever may come from config file. Idea is that we do
294 # not want to carry around the whole thing, and butler config will
295 # have empty obscore configuration after initialization. When
296 # configuration is missing from attributes table, the obscore table
297 # does not exist, and we do not instantiate obscore manager.
298 obscoreString = attributes.get(_OBSCORE_ATTR)
299 if obscoreString is not None:
300 self.manager_configs["obscore"] = Config.fromString(obscoreString, format="json")
302 with database.declareStaticTables(create=False) as context:
303 instances = RegistryManagerInstances.initialize(database, context, types=self, universe=universe)
305 # Load content from database that we try to keep in-memory.
306 instances.refresh()
307 return instances
309 def as_dict(self) -> Mapping[str, type[VersionedExtension]]:
310 """Return contained managers as a dictionary with manager type name as
311 a key.
313 Returns
314 -------
315 extensions : `~collections.abc.Mapping` [`str`, `VersionedExtension`]
316 Maps manager type name (e.g. "datasets") to its corresponding
317 manager class. Only existing managers are returned.
318 """
319 extras = {"manager_configs", "schema_versions"}
320 managers = {f.name: getattr(self, f.name) for f in dataclasses.fields(self) if f.name not in extras}
321 return {key: value for key, value in managers.items() if value is not None}
323 manager_configs: dict[str, Mapping] = dataclasses.field(default_factory=dict)
324 """Per-manager configuration options passed to their initialize methods.
325 """
327 schema_versions: dict[str, VersionTuple] = dataclasses.field(default_factory=dict)
328 """Per-manager schema versions defined by configuration, optional."""
331@dataclasses.dataclass(frozen=True, eq=False)
332class RegistryManagerInstances(
333 _GenericRegistryManagers[
334 ButlerAttributeManager,
335 DimensionRecordStorageManager,
336 CollectionManager,
337 DatasetRecordStorageManager,
338 OpaqueTableStorageManager,
339 DatastoreRegistryBridgeManager,
340 ObsCoreTableManager,
341 ]
342):
343 """A struct used to pass around the manager instances that back a
344 `Registry`.
345 """
347 column_types: ColumnTypeInfo
348 """Information about column types that can differ between data repositories
349 and registry instances, including the dimension universe.
350 """
352 caching_context: CachingContext
353 """Object containing caches for for various information generated by
354 managers.
355 """
357 @contextmanager
358 def caching_context_manager(self) -> Iterator[None]:
359 """Context manager that enables caching.
361 Calls to this method may be nested and the returned context managers
362 may even be closed out of order, with only the context manager entered
363 and the last context manager exited having any effect.
364 """
365 with (
366 self.caching_context.enable_collection_record_cache(),
367 self.caching_context.enable_collection_summary_cache(),
368 ):
369 yield
371 @classmethod
372 def initialize(
373 cls,
374 database: Database,
375 context: StaticTablesContext,
376 *,
377 types: RegistryManagerTypes,
378 universe: DimensionUniverse,
379 caching_context: CachingContext | None = None,
380 ) -> RegistryManagerInstances:
381 """Construct manager instances from their types and an existing
382 database connection.
384 Parameters
385 ----------
386 database : `Database`
387 Object that represents a connection to the SQL database that backs
388 the data repository.
389 context : `StaticTablesContext`
390 Object used to create tables in ``database``.
391 types : `RegistryManagerTypes`
392 Struct containing type objects for the manager instances to
393 construct.
394 universe : `DimensionUniverse`
395 Object that describes all dimensions in this data repository.
396 caching_context : `CachingContext` or `None`, optional
397 Caching context to use.
399 Returns
400 -------
401 instances : `RegistryManagerInstances`
402 Struct containing manager instances.
403 """
404 if caching_context is None:
405 caching_context = CachingContext()
406 dummy_table = ddl.TableSpec(fields=())
407 kwargs: dict[str, Any] = {}
408 schema_versions = types.schema_versions
409 kwargs["attributes"] = types.attributes.initialize(
410 database, context, registry_schema_version=schema_versions.get("attributes")
411 )
412 kwargs["dimensions"] = types.dimensions.initialize(
413 database, context, universe=universe, registry_schema_version=schema_versions.get("dimensions")
414 )
415 kwargs["collections"] = types.collections.initialize(
416 database,
417 context,
418 caching_context=caching_context,
419 registry_schema_version=schema_versions.get("collections"),
420 )
421 datasets = types.datasets.initialize(
422 database,
423 context,
424 collections=kwargs["collections"],
425 dimensions=kwargs["dimensions"],
426 registry_schema_version=schema_versions.get("datasets"),
427 caching_context=caching_context,
428 )
429 kwargs["datasets"] = datasets
430 kwargs["opaque"] = types.opaque.initialize(
431 database, context, registry_schema_version=schema_versions.get("opaque")
432 )
433 kwargs["datastores"] = types.datastores.initialize(
434 database,
435 context,
436 opaque=kwargs["opaque"],
437 datasets=types.datasets,
438 universe=universe,
439 registry_schema_version=schema_versions.get("datastores"),
440 )
441 kwargs["column_types"] = ColumnTypeInfo(
442 database.getTimespanRepresentation(),
443 universe,
444 dataset_id_spec=types.datasets.addDatasetForeignKey(
445 dummy_table,
446 primaryKey=False,
447 nullable=False,
448 ),
449 run_key_spec=types.collections.addRunForeignKey(dummy_table, primaryKey=False, nullable=False),
450 ingest_date_dtype=datasets.ingest_date_dtype(),
451 )
452 if types.obscore is not None and "obscore" in types.manager_configs:
453 kwargs["obscore"] = types.obscore.initialize(
454 database,
455 context,
456 universe=universe,
457 config=types.manager_configs["obscore"],
458 datasets=types.datasets,
459 dimensions=kwargs["dimensions"],
460 registry_schema_version=schema_versions.get("obscore"),
461 column_type_info=kwargs["column_types"],
462 )
463 else:
464 kwargs["obscore"] = None
465 kwargs["caching_context"] = caching_context
466 return cls(**kwargs)
468 def clone(
469 self,
470 db: Database,
471 ) -> RegistryManagerInstances:
472 """Make an independent copy of the manager instances with a new
473 `Database` instance.
475 Parameters
476 ----------
477 db : `Database`
478 New `Database` object to use when instantiating managers.
480 Returns
481 -------
482 instances : `RegistryManagerInstances`
483 New manager instances with the same configuration as this instance,
484 but bound to a new Database object.
485 """
486 caching_context = CachingContext()
487 dimensions = self.dimensions.clone(db)
488 collections = self.collections.clone(db, caching_context)
489 opaque = self.opaque.clone(db)
490 datasets = self.datasets.clone(
491 db=db, collections=collections, dimensions=dimensions, caching_context=caching_context
492 )
493 obscore = None
494 if self.obscore is not None:
495 obscore = self.obscore.clone(db=db, dimensions=dimensions)
496 return RegistryManagerInstances(
497 attributes=self.attributes.clone(db),
498 dimensions=dimensions,
499 collections=collections,
500 datasets=datasets,
501 opaque=opaque,
502 datastores=self.datastores.clone(db=db, opaque=opaque),
503 obscore=obscore,
504 column_types=self.column_types,
505 caching_context=caching_context,
506 )
508 def as_dict(self) -> Mapping[str, VersionedExtension]:
509 """Return contained managers as a dictionary with manager type name as
510 a key.
512 Returns
513 -------
514 extensions : `~collections.abc.Mapping` [`str`, `VersionedExtension`]
515 Maps manager type name (e.g. "datasets") to its corresponding
516 manager instance. Only existing managers are returned.
517 """
518 instances = {
519 f.name: getattr(self, f.name)
520 for f in dataclasses.fields(self)
521 if f.name not in ("column_types", "caching_context")
522 }
523 return {key: value for key, value in instances.items() if value is not None}
525 def refresh(self) -> None:
526 """Refresh all in-memory state by querying the database or clearing
527 caches.
528 """
529 self.collections.refresh()
530 self.datasets.refresh()