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