Coverage for python/lsst/daf/butler/registry/obscore/_manager.py: 17%
179 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-06 01:33 -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
30__all__ = ["ObsCoreLiveTableManager"]
32import re
33import warnings
34from collections import defaultdict
35from collections.abc import Collection, Iterable, Iterator, Mapping
36from contextlib import AbstractContextManager, contextmanager
37from typing import TYPE_CHECKING, Any
39import sqlalchemy
41from lsst.daf.butler import Config, DataCoordinate, DatasetRef, DimensionUniverse
42from lsst.sphgeom import Region
43from lsst.utils.introspection import find_outside_stacklevel
44from lsst.utils.iteration import chunk_iterable
46from ..._column_type_info import ColumnTypeInfo
47from ...queries import Query, QueryFactoryFunction
48from ..interfaces import ObsCoreTableManager, VersionTuple
49from ._config import ConfigCollectionType, ObsCoreManagerConfig
50from ._records import DerivedRegionFactory, Record, RecordFactory
51from ._schema import ObsCoreSchema
52from ._spatial import RegionTypeError, RegionTypeWarning, SpatialObsCorePlugin
54if TYPE_CHECKING:
55 from ..interfaces import (
56 CollectionRecord,
57 Database,
58 DatasetRecordStorageManager,
59 DimensionRecordStorageManager,
60 StaticTablesContext,
61 )
63_VERSION = VersionTuple(0, 0, 1)
66class _ExposureRegionFactory(DerivedRegionFactory):
67 """Find exposure region from a matching visit dimensions records.
69 Parameters
70 ----------
71 dimensions : `DimensionRecordStorageManager`
72 The dimension records storage manager.
73 query_func : `QueryFactoryFunction`
74 Function returning a context manager that sets up a `Query` object
75 for querying the registry. (That is, a function equivalent to
76 ``Butler.query()``).
77 """
79 def __init__(
80 self,
81 dimensions: DimensionRecordStorageManager,
82 query_func: QueryFactoryFunction,
83 ):
84 self.dimensions = dimensions
85 self.universe = dimensions.universe
86 self.exposure_dimensions = self.universe["exposure"].minimal_group
87 self.exposure_detector_dimensions = self.universe.conform(["exposure", "detector"])
88 self._query_func = query_func
90 def derived_region(self, dataId: DataCoordinate) -> Region | None:
91 # Docstring is inherited from a base class.
93 # Make sure the dimension universe contains a table that can be used
94 # to find visits from exposures.
95 if "visit_definition" not in self.universe.elements.names:
96 return None
97 # Choose the table we will use to look up the visit region:
98 # either visit+detector regions or visit regions.
99 if "detector" in dataId.dimensions:
100 if "visit_detector_region" not in self.universe:
101 return None
102 constraint_data_id = dataId.subset(self.exposure_detector_dimensions)
103 region_dimension = "visit_detector_region"
104 else:
105 if "visit" not in self.universe:
106 return None
107 constraint_data_id = dataId.subset(self.exposure_dimensions)
108 region_dimension = "visit"
110 with self._query_func() as query:
111 result = list(
112 query.dimension_records(region_dimension)
113 # Constrain the relation to match the given exposure and (if
114 # present) detector IDs.
115 .where(constraint_data_id)
116 # If we get more than one result (because the exposure belongs
117 # to multiple visits), just pick an arbitrary one.
118 .limit(1)
119 )
120 if len(result) > 0:
121 return result[0].region
122 else:
123 return None
126class ObsCoreLiveTableManager(ObsCoreTableManager):
127 """A manager class for ObsCore table, implements methods for updating the
128 records in that table.
130 Parameters
131 ----------
132 db : `Database`
133 The active database.
134 table : `sqlalchemy.schema.Table`
135 The ObsCore table.
136 schema : `ObsCoreSchema`
137 The relevant schema.
138 universe : `DimensionUniverse`
139 The dimension universe.
140 config : `ObsCoreManagerConfig`
141 The config controlling the manager.
142 dimensions : `DimensionRecordStorageManager`
143 The storage manager for the dimension records.
144 spatial_plugins : `~collections.abc.Collection` of `SpatialObsCorePlugin`
145 Spatial plugins.
146 registry_schema_version : `VersionTuple` or `None`, optional
147 Version of registry schema.
148 column_type_info : `ColumnTypeInfo`
149 Information about column types that can differ between data
150 repositories and registry instances.
151 """
153 def __init__(
154 self,
155 *,
156 db: Database,
157 table: sqlalchemy.schema.Table,
158 schema: ObsCoreSchema,
159 universe: DimensionUniverse,
160 config: ObsCoreManagerConfig,
161 dimensions: DimensionRecordStorageManager,
162 spatial_plugins: Collection[SpatialObsCorePlugin],
163 registry_schema_version: VersionTuple | None = None,
164 column_type_info: ColumnTypeInfo,
165 ):
166 super().__init__(registry_schema_version=registry_schema_version)
167 self.db = db
168 self.table = table
169 self.schema = schema
170 self.universe = universe
171 self.config = config
172 self.spatial_plugins = spatial_plugins
173 self._column_type_info = column_type_info
174 self._query_func: QueryFactoryFunction | None = None
175 exposure_region_factory = _ExposureRegionFactory(dimensions, self._get_query_object)
176 self.record_factory = RecordFactory.get_record_type_from_universe(universe)(
177 config, schema, universe, spatial_plugins, exposure_region_factory
178 )
179 self.tagged_collection: str | None = None
180 self.run_patterns: list[re.Pattern] = []
181 if config.collection_type is ConfigCollectionType.TAGGED:
182 assert config.collections is not None and len(config.collections) == 1, (
183 "Exactly one collection name required for tagged type."
184 )
185 self.tagged_collection = config.collections[0]
186 elif config.collection_type is ConfigCollectionType.RUN:
187 if config.collections:
188 for coll in config.collections:
189 try:
190 self.run_patterns.append(re.compile(coll))
191 except re.error as exc:
192 raise ValueError(f"Failed to compile regex: {coll!r}") from exc
193 else:
194 raise ValueError(f"Unexpected value of collection_type: {config.collection_type}")
196 def clone(self, *, db: Database, dimensions: DimensionRecordStorageManager) -> ObsCoreLiveTableManager:
197 manager = ObsCoreLiveTableManager(
198 db=db,
199 table=self.table,
200 schema=self.schema,
201 universe=self.universe,
202 config=self.config,
203 dimensions=dimensions,
204 # Current spatial plugins are safe to share without cloning -- they
205 # are immutable and do not use their Database object outside of
206 # 'initialize'.
207 spatial_plugins=self.spatial_plugins,
208 registry_schema_version=self._registry_schema_version,
209 column_type_info=self._column_type_info,
210 )
211 if self._query_func is not None:
212 manager.set_query_function(self._query_func)
213 return manager
215 @classmethod
216 def initialize(
217 cls,
218 db: Database,
219 context: StaticTablesContext,
220 *,
221 universe: DimensionUniverse,
222 config: Mapping,
223 datasets: type[DatasetRecordStorageManager],
224 dimensions: DimensionRecordStorageManager,
225 registry_schema_version: VersionTuple | None = None,
226 column_type_info: ColumnTypeInfo,
227 ) -> ObsCoreTableManager:
228 # Docstring inherited from base class.
229 config_data = Config(config)
230 obscore_config = ObsCoreManagerConfig.model_validate(config_data)
232 # Instantiate all spatial plugins.
233 spatial_plugins = SpatialObsCorePlugin.load_plugins(obscore_config.spatial_plugins, db)
235 schema = ObsCoreSchema(config=obscore_config, spatial_plugins=spatial_plugins, datasets=datasets)
237 # Generate table specification for main obscore table.
238 table_spec = schema.table_spec
239 for plugin in spatial_plugins:
240 plugin.extend_table_spec(table_spec)
241 table = context.addTable(obscore_config.table_name, schema.table_spec)
243 return ObsCoreLiveTableManager(
244 db=db,
245 table=table,
246 schema=schema,
247 universe=universe,
248 config=obscore_config,
249 dimensions=dimensions,
250 spatial_plugins=spatial_plugins,
251 registry_schema_version=registry_schema_version,
252 column_type_info=column_type_info,
253 )
255 def _get_query_object(self) -> AbstractContextManager[Query]:
256 if self._query_func is None:
257 raise AssertionError("set_query_function should have been called prior to this.")
259 return self._query_func()
261 def set_query_function(self, query_func: QueryFactoryFunction) -> None:
262 self._query_func = query_func
264 def config_json(self) -> str:
265 """Dump configuration in JSON format.
267 Returns
268 -------
269 json : `str`
270 Configuration serialized in JSON format.
271 """
272 return self.config.model_dump_json()
274 @classmethod
275 def currentVersions(cls) -> list[VersionTuple]:
276 # Docstring inherited from base class.
277 return [_VERSION]
279 def add_datasets(self, refs: Iterable[DatasetRef]) -> int:
280 # Docstring inherited from base class.
282 # Only makes sense for RUN collection types
283 if self.config.collection_type is not ConfigCollectionType.RUN:
284 return 0
286 obscore_refs: Iterable[DatasetRef]
287 if self.run_patterns:
288 # Check each dataset run against configured run list. We want to
289 # reduce number of calls to _check_dataset_run, which may be
290 # expensive. Normally references are grouped by run, if there are
291 # multiple input references, they should have the same run.
292 # Instead of just checking that, we group them by run again.
293 refs_by_run: dict[str, list[DatasetRef]] = defaultdict(list)
294 for ref in refs:
295 # Record factory will filter dataset types, but to reduce
296 # collection checks we also pre-filter it here.
297 if ref.datasetType.name not in self.config.dataset_types:
298 continue
300 assert ref.run is not None, "Run cannot be None"
301 refs_by_run[ref.run].append(ref)
303 good_refs: list[DatasetRef] = []
304 for run, run_refs in refs_by_run.items():
305 if not self._check_dataset_run(run):
306 continue
307 good_refs.extend(run_refs)
308 obscore_refs = good_refs
310 else:
311 # Take all refs, no collection check.
312 obscore_refs = refs
314 return self._populate(obscore_refs)
316 def associate(self, refs: Iterable[DatasetRef], collection: CollectionRecord) -> int:
317 # Docstring inherited from base class.
319 # Only works when collection type is TAGGED
320 if self.tagged_collection is None:
321 return 0
323 if collection.name == self.tagged_collection:
324 return self._populate(refs)
325 else:
326 return 0
328 def disassociate(self, refs: Iterable[DatasetRef], collection: CollectionRecord) -> int:
329 # Docstring inherited from base class.
331 # Only works when collection type is TAGGED
332 if self.tagged_collection is None:
333 return 0
335 count = 0
336 if collection.name == self.tagged_collection:
337 # Sorting may improve performance
338 dataset_ids = sorted(ref.id for ref in refs)
339 if dataset_ids:
340 fk_field = self.schema.dataset_fk
341 assert fk_field is not None, "Cannot be None by construction"
342 # There may be too many of them, do it in chunks.
343 for ids in chunk_iterable(dataset_ids):
344 where = self.table.columns[fk_field.name].in_(ids)
345 count += self.db.deleteWhere(self.table, where)
346 return count
348 def _populate(self, refs: Iterable[DatasetRef]) -> int:
349 """Populate obscore table with the data from given datasets."""
350 records: list[Record] = []
351 for ref in refs:
352 record = self.record_factory(ref)
353 if record is not None:
354 records.append(record)
356 if records:
357 # Ignore potential conflicts with existing datasets.
358 return self.db.ensure(self.table, *records, primary_key_only=True)
359 else:
360 return 0
362 def _check_dataset_run(self, run: str) -> bool:
363 """Check that specified run collection matches know patterns."""
364 if not self.run_patterns:
365 # Empty list means take anything.
366 return True
368 # Try each pattern in turn.
369 return any(pattern.fullmatch(run) for pattern in self.run_patterns)
371 def update_exposure_regions(self, instrument: str, region_data: Iterable[tuple[int, int, Region]]) -> int:
372 # Docstring inherited from base class.
373 instrument_column = self.schema.dimension_column("instrument")
374 exposure_column = self.schema.dimension_column("exposure")
375 detector_column = self.schema.dimension_column("detector")
376 if instrument_column is None or exposure_column is None or detector_column is None:
377 # Not all needed columns are in the table.
378 return 0
380 update_rows: list[Record] = []
381 for exposure, detector, region in region_data:
382 try:
383 record = self.record_factory.make_spatial_records(region)
384 except RegionTypeError as exc:
385 warnings.warn(
386 f"Failed to convert region for exposure={exposure} detector={detector}: {exc}",
387 category=RegionTypeWarning,
388 stacklevel=find_outside_stacklevel("lsst.daf.butler"),
389 )
390 continue
392 record.update(
393 {
394 "instrument_column": instrument,
395 "exposure_column": exposure,
396 "detector_column": detector,
397 }
398 )
399 update_rows.append(record)
401 where_dict: dict[str, str] = {
402 instrument_column: "instrument_column",
403 exposure_column: "exposure_column",
404 detector_column: "detector_column",
405 }
407 count = self.db.update(self.table, where_dict, *update_rows)
408 return count
410 @contextmanager
411 def query(
412 self, columns: Iterable[str | sqlalchemy.sql.expression.ColumnElement] | None = None, /, **kwargs: Any
413 ) -> Iterator[sqlalchemy.engine.CursorResult]:
414 # Docstring inherited from base class.
415 if columns is not None:
416 column_elements: list[sqlalchemy.sql.ColumnElement] = []
417 for column in columns:
418 if isinstance(column, str):
419 column_elements.append(self.table.columns[column])
420 else:
421 column_elements.append(column)
422 query = sqlalchemy.sql.select(*column_elements).select_from(self.table)
423 else:
424 query = self.table.select()
426 if kwargs:
427 query = query.where(
428 sqlalchemy.sql.expression.and_(
429 *[self.table.columns[column] == value for column, value in kwargs.items()]
430 )
431 )
432 with self.db.query(query) as result:
433 yield result