Coverage for python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py: 92%
583 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-29 02:00 -0700
1from __future__ import annotations
3__all__ = ("ByDimensionsDatasetRecordStorageManagerUUID",)
5import dataclasses
6import datetime
7import logging
8from collections import defaultdict
9from collections.abc import Iterable, Mapping, Sequence, Set
10from typing import TYPE_CHECKING, Any, ClassVar
12import astropy.time
13import pydantic
14import sqlalchemy
16from lsst.utils.iteration import chunk_iterable
18from .... import ddl
19from ...._collection_type import CollectionType
20from ...._dataset_ref import DatasetId, DatasetIdFactory, DatasetIdGenEnum, DatasetRef
21from ...._dataset_type import DatasetType, get_dataset_type_name, validate_dataset_type_name
22from ...._exceptions import CollectionTypeError, DatasetTypeExpressionError, MissingDatasetTypeError
23from ...._exceptions_legacy import DatasetTypeError
24from ...._timespan import Timespan
25from ....dimensions import DataCoordinate, DimensionGroup, DimensionUniverse
26from ....direct_query_driver import SqlJoinsBuilder, SqlSelectBuilder # new query system, server+direct only
27from ....queries import QueryFactoryFunction
28from ....queries import tree as qt # new query system, both clients + server
29from ..._caching_context import CachingContext
30from ..._collection_summary import CollectionSummary
31from ..._exceptions import ConflictingDefinitionError, OrphanedRecordError
32from ...interfaces import DatasetRecordStorageManager, RunRecord, VersionTuple
33from ...wildcards import DatasetTypeWildcard
34from ._dataset_type_cache import DatasetTypeCache
35from .summaries import CollectionSummaryManager
36from .tables import (
37 DynamicTables,
38 StaticDatasetTableSpecTuple,
39 StaticDatasetTablesTuple,
40 addDatasetForeignKey,
41 makeStaticTableSpecs,
42 makeTagTableSpec,
43)
45if TYPE_CHECKING:
46 from ...interfaces import (
47 CollectionManager,
48 CollectionRecord,
49 Database,
50 DimensionRecordStorageManager,
51 StaticTablesContext,
52 )
55# This has to be updated on every schema change
56# TODO: 1.0.0 can be removed once all repos were migrated to 2.0.0.
57_VERSION_UUID = VersionTuple(1, 0, 0)
58# Starting with 2.0.0 the `ingest_date` column type uses nanoseconds instead
59# of TIMESTAMP. The code supports both 1.0.0 and 2.0.0 for the duration of
60# client migration period.
61_VERSION_UUID_NS = VersionTuple(2, 0, 0)
63_LOG = logging.getLogger(__name__)
66@dataclasses.dataclass
67class _DatasetTypeRecord:
68 """Contents of a single dataset type record."""
70 dataset_type: DatasetType
71 dataset_type_id: int
72 dimensions_key: int
73 tag_table_name: str
74 calib_table_name: str | None
76 def make_dynamic_tables(self) -> DynamicTables:
77 return DynamicTables(
78 self.dataset_type.dimensions, self.dimensions_key, self.tag_table_name, self.calib_table_name
79 )
81 def update_dynamic_tables(self, current: DynamicTables) -> DynamicTables:
82 assert self.dimensions_key == current.dimensions_key
83 assert self.tag_table_name == current.tags_name
84 if self.calib_table_name is not None:
85 if current.calibs_name is not None: 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true
86 assert self.calib_table_name == current.calibs_name
87 else:
88 # Some previously-cached dataset type had the same dimensions
89 # but was not a calibration.
90 current = current.copy(calibs_name=self.calib_table_name)
91 # If some previously-cached dataset type was a calibration but this
92 # one isn't, we don't want to forget the calibs table.
93 return current
96@dataclasses.dataclass
97class _DatasetRecordStorage:
98 """Information cached about a dataset type.
100 This combines information cached with different keys - the dataset type
101 and its ID are cached by name, while the tables are cached by the dataset
102 types dimensions (and hence shared with other dataset types that have the
103 same dimensions).
104 """
106 dataset_type: DatasetType
107 dataset_type_id: int
108 dynamic_tables: DynamicTables
111class _DatasetTypeOverride(pydantic.BaseModel):
112 """A configuration struct that holds storage class overrides and/or a new
113 name for a dataset type.
114 """
116 storageClass: str | None = None
117 rename: str | None = None
120class ByDimensionsDatasetRecordStorageManagerUUID(DatasetRecordStorageManager):
121 """A manager class for datasets that uses one dataset-collection table for
122 each group of dataset types that share the same dimensions.
124 In addition to the table organization, this class makes a number of
125 other design choices that would have been cumbersome (to say the least) to
126 try to pack into its name:
128 - It uses a private surrogate integer autoincrement field to identify
129 dataset types, instead of using the name as the primary and foreign key
130 directly.
132 - It aggressively loads all DatasetTypes into memory instead of fetching
133 them from the database only when needed or attempting more clever forms
134 of caching.
136 Alternative implementations that make different choices for these while
137 keeping the same general table organization might be reasonable as well.
139 Parameters
140 ----------
141 db : `Database`
142 Interface to the underlying database engine and namespace.
143 collections : `CollectionManager`
144 Manager object for the collections in this `Registry`.
145 dimensions : `DimensionRecordStorageManager`
146 Manager object for the dimensions in this `Registry`.
147 static : `StaticDatasetTablesTuple`
148 Named tuple of `sqlalchemy.schema.Table` instances for all static
149 tables used by this class.
150 summaries : `CollectionSummaryManager`
151 Structure containing tables that summarize the contents of collections.
152 renames
153 Mapping from a dataset type name in the database to an client-level
154 override name.
155 reversed_renames
156 Mapping from the override name for a dataset type to its true (DB)
157 name.
158 storage_class_overrides
159 Mapping from the true (DB) name of a dataset type to an override
160 storage class this client should use instead of the one in the
161 database.
162 registry_schema_version : `VersionTuple` or `None`, optional
163 Version of registry schema.
164 _cache : `None`, optional
165 For internal use only.
166 """
168 def __init__(
169 self,
170 *,
171 db: Database,
172 collections: CollectionManager,
173 dimensions: DimensionRecordStorageManager,
174 static: StaticDatasetTablesTuple,
175 summaries: CollectionSummaryManager,
176 renames: Mapping[str, str],
177 reversed_renames: Mapping[str, str],
178 storage_class_overrides: Mapping[str, str],
179 registry_schema_version: VersionTuple | None = None,
180 _cache: DatasetTypeCache | None = None,
181 ):
182 super().__init__(registry_schema_version=registry_schema_version)
183 self._db = db
184 self._collections = collections
185 self._dimensions = dimensions
186 self._static = static
187 self._summaries = summaries
188 self._cache = _cache if _cache is not None else DatasetTypeCache()
189 self._use_astropy_ingest_date = self.ingest_date_dtype() is ddl.AstropyTimeNsecTai
190 self._run_key_column = collections.getRunForeignKeyName()
191 self._renames = renames
192 self._reversed_renames = reversed_renames
193 self._storage_class_overrides = storage_class_overrides
195 _versions: ClassVar[list[VersionTuple]] = [_VERSION_UUID, _VERSION_UUID_NS]
197 _id_maker: ClassVar[DatasetIdFactory] = DatasetIdFactory()
198 """Factory for dataset IDs. In the future this factory may be shared with
199 other classes (e.g. Registry).
200 """
202 @classmethod
203 def initialize(
204 cls,
205 db: Database,
206 context: StaticTablesContext,
207 *,
208 collections: CollectionManager,
209 dimensions: DimensionRecordStorageManager,
210 caching_context: CachingContext,
211 config: Mapping,
212 registry_schema_version: VersionTuple | None = None,
213 ) -> DatasetRecordStorageManager:
214 # Docstring inherited from DatasetRecordStorageManager.
215 specs = cls.makeStaticTableSpecs(
216 type(collections), universe=dimensions.universe, schema_version=registry_schema_version
217 )
218 static: StaticDatasetTablesTuple = context.addTableTuple(specs) # type: ignore
219 dataset_type_overrides = {
220 dataset_type_name: _DatasetTypeOverride.model_validate(override_config)
221 for dataset_type_name, override_config in config.get("overrides", {}).items()
222 }
223 renames = {
224 db_name: override.rename
225 for db_name, override in dataset_type_overrides.items()
226 if override.rename is not None
227 }
228 reversed_renames = {v: k for k, v in renames.items()}
229 storage_class_overrides = {
230 db_name: override.storageClass
231 for db_name, override in dataset_type_overrides.items()
232 if override.storageClass is not None
233 }
234 summaries = CollectionSummaryManager.initialize(
235 db,
236 context,
237 collections=collections,
238 dimensions=dimensions,
239 dataset_type_table=static.dataset_type,
240 caching_context=caching_context,
241 reversed_renames=reversed_renames,
242 )
243 return cls(
244 db=db,
245 collections=collections,
246 dimensions=dimensions,
247 static=static,
248 summaries=summaries,
249 registry_schema_version=registry_schema_version,
250 renames=renames,
251 reversed_renames=reversed_renames,
252 storage_class_overrides=storage_class_overrides,
253 )
255 @classmethod
256 def currentVersions(cls) -> list[VersionTuple]:
257 # Docstring inherited from VersionedExtension.
258 return cls._versions
260 @classmethod
261 def makeStaticTableSpecs(
262 cls,
263 collections: type[CollectionManager],
264 universe: DimensionUniverse,
265 schema_version: VersionTuple | None,
266 ) -> StaticDatasetTableSpecTuple:
267 """Construct all static tables used by the classes in this package.
269 Static tables are those that are present in all Registries and do not
270 depend on what DatasetTypes have been registered.
272 Parameters
273 ----------
274 collections : `CollectionManager`
275 Manager object for the collections in this `Registry`.
276 universe : `DimensionUniverse`
277 Universe graph containing all dimensions known to this `Registry`.
278 schema_version : `VersionTuple` or `None`
279 Version of the schema that should be created, if `None` then
280 default schema should be used.
282 Returns
283 -------
284 specs : `StaticDatasetTablesTuple`
285 A named tuple containing `ddl.TableSpec` instances.
286 """
287 schema_version = cls.clsNewSchemaVersion(schema_version)
288 assert schema_version is not None, "New schema version cannot be None"
289 return makeStaticTableSpecs(
290 collections,
291 universe=universe,
292 schema_version=schema_version,
293 )
295 @classmethod
296 def addDatasetForeignKey(
297 cls,
298 tableSpec: ddl.TableSpec,
299 *,
300 name: str = "dataset",
301 constraint: bool = True,
302 onDelete: str | None = None,
303 **kwargs: Any,
304 ) -> ddl.FieldSpec:
305 # Docstring inherited from DatasetRecordStorageManager.
306 return addDatasetForeignKey(tableSpec, name=name, onDelete=onDelete, constraint=constraint, **kwargs)
308 @classmethod
309 def _newDefaultSchemaVersion(cls) -> VersionTuple:
310 # Docstring inherited from VersionedExtension.
311 return _VERSION_UUID_NS
313 def clone(
314 self,
315 *,
316 db: Database,
317 collections: CollectionManager,
318 dimensions: DimensionRecordStorageManager,
319 caching_context: CachingContext,
320 ) -> ByDimensionsDatasetRecordStorageManagerUUID:
321 return ByDimensionsDatasetRecordStorageManagerUUID(
322 db=db,
323 collections=collections,
324 dimensions=dimensions,
325 static=self._static,
326 summaries=self._summaries.clone(db=db, collections=collections, caching_context=caching_context),
327 renames=self._renames,
328 reversed_renames=self._reversed_renames,
329 storage_class_overrides=self._storage_class_overrides,
330 registry_schema_version=self._registry_schema_version,
331 # See notes on DatasetTypeCache.clone() about cache behavior after
332 # cloning.
333 _cache=self._cache.clone(),
334 )
336 def refresh(self) -> None:
337 # Docstring inherited from DatasetRecordStorageManager.
338 self._cache.clear()
340 def remove_dataset_type(self, name: str) -> None:
341 # Docstring inherited from DatasetRecordStorageManager.
342 db_name = self._reversed_renames.get(name, name)
343 compositeName, componentName = DatasetType.splitDatasetTypeName(name)
344 if componentName is not None: 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true
345 raise ValueError(f"Cannot delete a dataset type of a component of a composite (given {name})")
347 # Delete the row
348 try:
349 self._db.delete(self._static.dataset_type, ["name"], {"name": db_name})
350 except sqlalchemy.exc.IntegrityError as e:
351 msg = name if db_name == name else f"{db_name} (renamed to {name})"
352 raise OrphanedRecordError(
353 f"Dataset type {msg} can not be removed."
354 " It is associated with datasets that must be removed first."
355 ) from e
357 # Now refresh everything -- removal is rare enough that this does
358 # not need to be fast.
359 self.refresh()
361 def get_dataset_type(self, name: str) -> DatasetType:
362 # Docstring inherited from DatasetRecordStorageManager.
363 return self._find_storage(name).dataset_type
365 def register_dataset_type(self, dataset_type: DatasetType) -> bool:
366 # Docstring inherited from DatasetRecordStorageManager.
367 #
368 # This is one of the places where we populate the dataset type cache.
369 # See the comment in _fetch_dataset_types for how these are related and
370 # invariants they must maintain.
371 #
372 if dataset_type.isComponent(): 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 raise ValueError(
374 f"Component dataset types can not be stored in registry. Rejecting {dataset_type.name}"
375 )
377 # If database universe and dimension group universe are different it
378 # can cause unexpected effects.
379 if dataset_type.dimensions.universe is not self._dimensions.universe:
380 raise ValueError(
381 "Incompatible dimension universe versions - "
382 f"database universe: {self._dimensions.universe}, "
383 f"dataset type universe: {dataset_type.dimensions.universe}."
384 )
386 record = self._fetch_dataset_type_record(dataset_type.name)
387 if record is None:
388 if (dynamic_tables := self._cache.get_by_dimensions(dataset_type.dimensions)) is None:
389 dimensions_key = self._dimensions.save_dimension_group(dataset_type.dimensions)
390 dynamic_tables = DynamicTables.from_dimensions_key(
391 dataset_type.dimensions, dimensions_key, dataset_type.isCalibration()
392 )
393 dynamic_tables.create(self._db, type(self._collections), self._cache.tables)
394 elif dataset_type.isCalibration() and dynamic_tables.calibs_name is None:
395 dynamic_tables = dynamic_tables.add_calibs(
396 self._db, type(self._collections), self._cache.tables
397 )
398 db_name = self._reversed_renames.get(dataset_type.name, dataset_type.name)
399 if db_name == dataset_type.name and db_name in self._renames:
400 # We were asked to fetch a dataset type that has been renamed
401 # a config override, with no other existing dataset type
402 # to replace it. It should appear missing, but we still can't
403 # register it.
404 raise ConflictingDefinitionError(
405 f"A dataset type with the name {db_name!r} has been renamed to , "
406 f"{self._renames[db_name]!r} in this client; it can only be registered "
407 "from a client without any rename configuration."
408 )
409 if db_name in self._storage_class_overrides:
410 raise ConflictingDefinitionError(
411 f"Dataset type {db_name!r} has had its storage class overridden to "
412 f"{dataset_type.storageClass_name!r} via butler configuration, and cannot "
413 "be registered via this client since the database-level storage class is unknown."
414 )
415 row, inserted = self._db.sync(
416 self._static.dataset_type,
417 keys={"name": db_name},
418 compared={
419 "dimensions_key": dynamic_tables.dimensions_key,
420 # Force the storage class to be loaded to ensure it
421 # exists and there is no typo in the name.
422 "storage_class": dataset_type.storageClass.name,
423 },
424 extra={
425 "tag_association_table": dynamic_tables.tags_name,
426 "calibration_association_table": (
427 dynamic_tables.calibs_name if dataset_type.isCalibration() else None
428 ),
429 },
430 returning=["id", "tag_association_table"],
431 )
432 # Make sure that cache is updated
433 if row is not None: 433 ↛ 443line 433 didn't jump to line 443 because the condition on line 433 was always true
434 self._cache.add(dataset_type, row["id"])
435 self._cache.add_by_dimensions(dataset_type.dimensions, dynamic_tables)
436 else:
437 if dataset_type != record.dataset_type:
438 raise ConflictingDefinitionError(
439 f"Given dataset type {dataset_type} is inconsistent "
440 f"with database definition {record.dataset_type}."
441 )
442 inserted = False
443 return bool(inserted)
445 def resolve_wildcard(
446 self,
447 expression: Any,
448 missing: list[str] | None = None,
449 explicit_only: bool = False,
450 ) -> list[DatasetType]:
451 wildcard = DatasetTypeWildcard.from_expression(expression)
452 result: list[DatasetType] = []
453 for name, dataset_type in wildcard.values.items():
454 parent_name, component_name = DatasetType.splitDatasetTypeName(name)
455 if component_name is not None:
456 # Distinguish a real component name from a string that is not
457 # a valid dataset type name at all, such as "...", which is
458 # sometimes passed in place of the "..." wildcard.
459 validate_dataset_type_name(name)
460 raise DatasetTypeExpressionError(
461 "Component dataset types are not supported in Registry methods; use DatasetRef or "
462 "DatasetType methods to obtain components from parents instead."
463 )
464 try:
465 resolved_dataset_type = self.get_dataset_type(parent_name)
466 except MissingDatasetTypeError:
467 if missing is not None:
468 missing.append(name)
469 else:
470 if dataset_type is not None:
471 if dataset_type.is_compatible_with(resolved_dataset_type): 471 ↛ 476line 471 didn't jump to line 476 because the condition on line 471 was always true
472 # Prefer the given dataset type to enable storage class
473 # conversions.
474 resolved_dataset_type = dataset_type
475 else:
476 raise DatasetTypeExpressionError(
477 f"Dataset type definition in query expression {dataset_type} is "
478 f"not compatible with the registered type {resolved_dataset_type}."
479 )
480 result.append(resolved_dataset_type)
481 if wildcard.patterns is ...:
482 if explicit_only: 482 ↛ 483line 482 didn't jump to line 483 because the condition on line 482 was never true
483 raise TypeError(
484 "Universal wildcard '...' is not permitted for dataset types in this context."
485 )
486 for datasetType in self._fetch_dataset_types():
487 result.append(datasetType)
488 elif wildcard.patterns:
489 if explicit_only: 489 ↛ 490line 489 didn't jump to line 490 because the condition on line 489 was never true
490 raise DatasetTypeExpressionError(
491 "Dataset type wildcard expressions are not supported in this context."
492 )
493 dataset_types = self._fetch_dataset_types()
494 for datasetType in dataset_types:
495 if any(p.fullmatch(datasetType.name) for p in wildcard.patterns):
496 result.append(datasetType)
498 return result
500 def get_dataset_refs(self, ids: list[DatasetId]) -> list[DatasetRef]:
501 dataset_type_map: dict[DatasetId, DatasetType] = {}
502 for batch in chunk_iterable(set(ids), 50000):
503 # Look up the dataset types corresponding to the given Dataset IDs.
504 id_col = self._static.dataset.columns["id"]
505 sql = sqlalchemy.sql.select(
506 id_col,
507 self._static.dataset.columns["dataset_type_id"],
508 ).where(id_col.in_(batch))
509 with self._db.query(sql) as sql_result:
510 dataset_rows = sql_result.mappings().all()
511 for row in dataset_rows:
512 dataset_type_map[row["id"]] = self._get_dataset_type_by_id(row["dataset_type_id"])
514 # Group the given dataset IDs by the DimensionGroup of their dataset
515 # types -- there is a separate tags table for each DimensionGroup.
516 dimension_groups = defaultdict[DimensionGroup, set[DatasetId]](set)
517 for id, dataset_type in dataset_type_map.items():
518 dimension_groups[dataset_type.dimensions].add(id)
520 output_refs: list[DatasetRef] = []
521 for dimension_group, datasets in dimension_groups.items():
522 # Query the tags table for each dimension group to look up the
523 # data IDs corresponding to the UUIDs found from the dataset table.
524 dynamic_tables = self._get_dynamic_tables(dimension_group)
525 tags_table = self._get_tags_table(dynamic_tables)
526 for batch in chunk_iterable(datasets, 50000):
527 tags_sql = tags_table.select().where(tags_table.columns["dataset_id"].in_(batch))
528 # Join in the collection table to fetch the run name.
529 collection_column = tags_table.columns[self._collections.getCollectionForeignKeyName()]
530 joined_collections = self._collections.join_collections_sql(collection_column, tags_sql)
531 tags_sql = joined_collections.joined_sql
532 run_name_column = joined_collections.name_column
533 tags_sql = tags_sql.add_columns(run_name_column)
534 # Tags table includes run collections and tagged
535 # collections.
536 # In theory the data ID for a given dataset should be the
537 # same in both, but nothing actually guarantees this.
538 # So skip any tagged collections, using the run collection
539 # as the definitive definition.
540 tags_sql = tags_sql.where(joined_collections.type_column == int(CollectionType.RUN))
542 with self._db.query(tags_sql) as sql_result:
543 data_id_rows = sql_result.mappings().all()
545 assert run_name_column.key is not None
546 for data_id_row in data_id_rows:
547 id = data_id_row["dataset_id"]
548 dataset_type = dataset_type_map[id]
549 run_name = data_id_row[run_name_column.key]
550 data_id = DataCoordinate.from_required_values(
551 dimension_group,
552 tuple(data_id_row[dimension] for dimension in dimension_group.required),
553 )
554 ref = DatasetRef(
555 datasetType=dataset_type,
556 dataId=data_id,
557 id=id,
558 run=run_name,
559 )
560 output_refs.append(ref)
562 return output_refs
564 def _fetch_dataset_type_record(self, name: str) -> _DatasetTypeRecord | None:
565 """Retrieve all dataset types defined in database.
567 Returns
568 -------
569 dataset_type
570 Information from a single database record.
571 """
572 db_name = self._reversed_renames.get(name, name)
573 if db_name == name and db_name in self._renames:
574 # We were asked to fetch a dataset type that has been renamed via
575 # a config override, with no other existing dataset type renamed
576 # to replace it. It should appear missing.
577 return None
578 c = self._static.dataset_type.columns
579 stmt = self._static.dataset_type.select().where(c.name == db_name)
580 with self._db.query(stmt) as sql_result:
581 row = sql_result.mappings().one_or_none()
582 if row is None:
583 return None
584 else:
585 return self._record_from_row(row)
587 def _record_from_row(self, row: Mapping) -> _DatasetTypeRecord:
588 name = self._renames.get(row["name"], row["name"])
589 dimensions = self._dimensions.load_dimension_group(row["dimensions_key"])
590 calibTableName = row["calibration_association_table"]
591 storage_class_name = self._storage_class_overrides.get(row["name"], row["storage_class"])
592 datasetType = DatasetType(
593 name, dimensions, storage_class_name, isCalibration=(calibTableName is not None)
594 )
595 return _DatasetTypeRecord(
596 dataset_type=datasetType,
597 dataset_type_id=row["id"],
598 dimensions_key=row["dimensions_key"],
599 tag_table_name=row["tag_association_table"],
600 calib_table_name=calibTableName,
601 )
603 def _dataset_type_from_row(self, row: Mapping) -> DatasetType:
604 return self._record_from_row(row).dataset_type
606 def preload_cache(self) -> None:
607 self._fetch_dataset_types()
609 def _fetch_dataset_types(self, force_refresh: bool = False) -> list[DatasetType]:
610 """Fetch list of all defined dataset types."""
611 # This is one of two places we populate the dataset type cache:
612 #
613 # - This method handles almost all requests for dataset types that
614 # should already exist. It always marks the cache as "full" in both
615 # dataset type names and dimensions.
616 #
617 # - register_dataset_type handles the case where the dataset type might
618 # not existing yet. Since it can only add a single dataset type, it
619 # never changes whether the cache is full.
620 #
621 # In both cases, we require that the per-dimensions data be cached
622 # whenever a dataset type is added to the cache by name, to reduce the
623 # number of possible states the cache can be in and minimize the number
624 # of queries.
625 if self._cache.full and not force_refresh:
626 return [dataset_type for dataset_type, _ in self._cache.items()]
627 with self._db.query(self._static.dataset_type.select()) as sql_result:
628 sql_rows = sql_result.mappings().fetchall()
629 records = [self._record_from_row(row) for row in sql_rows]
630 # Cache everything and specify that cache is complete.
631 cache_data: list[tuple[DatasetType, int]] = []
632 cache_dimensions_data: dict[DimensionGroup, DynamicTables] = {}
633 for record in records:
634 cache_data.append((record.dataset_type, record.dataset_type_id))
635 if (dynamic_tables := cache_dimensions_data.get(record.dataset_type.dimensions)) is None:
636 tables = record.make_dynamic_tables()
637 else:
638 tables = record.update_dynamic_tables(dynamic_tables)
639 cache_dimensions_data[record.dataset_type.dimensions] = tables
640 self._cache.set(
641 cache_data, full=True, dimensions_data=cache_dimensions_data.items(), dimensions_full=True
642 )
643 return [record.dataset_type for record in records]
645 def _get_dataset_type_by_id(self, id: int) -> DatasetType:
646 dt = self._cache.get_by_id(id)
647 if dt is None:
648 # Since the ID is not a concept exposed to the public API, it
649 # had to have come from a dataset table, and our cache must be
650 # empty or out of date.
651 self._fetch_dataset_types(force_refresh=True)
652 dt = self._cache.get_by_id(id)
653 if dt is None: 653 ↛ 654line 653 didn't jump to line 654 because the condition on line 653 was never true
654 raise RuntimeError(f"Failed to look up dataset type with ID {id}")
655 return dt
657 def _find_storage(self, name: str) -> _DatasetRecordStorage:
658 """Find a dataset type and the extra information needed to work with
659 it, utilizing and populating the cache as needed.
660 """
661 dataset_type, dataset_type_id = self._cache.get(name)
662 if dataset_type is not None:
663 tables = self._get_dynamic_tables(dataset_type.dimensions)
664 assert dataset_type_id is not None, "Dataset type cache population is incomplete."
665 return _DatasetRecordStorage(
666 dataset_type=dataset_type, dataset_type_id=dataset_type_id, dynamic_tables=tables
667 )
668 else:
669 # On the first cache miss populate the cache with complete list
670 # of dataset types (if it was not done yet).
671 if not self._cache.full:
672 self._fetch_dataset_types()
673 # Try again
674 dataset_type, dataset_type_id = self._cache.get(name)
675 if dataset_type is not None:
676 tables = self._get_dynamic_tables(dataset_type.dimensions)
677 assert dataset_type_id is not None, "Dataset type cache population is incomplete."
678 return _DatasetRecordStorage(
679 dataset_type=dataset_type, dataset_type_id=dataset_type_id, dynamic_tables=tables
680 )
681 record = self._fetch_dataset_type_record(name)
682 if record is not None:
683 self._cache.add(record.dataset_type, record.dataset_type_id)
684 tables = record.make_dynamic_tables()
685 self._cache.add_by_dimensions(record.dataset_type.dimensions, tables)
686 return _DatasetRecordStorage(record.dataset_type, record.dataset_type_id, tables)
687 raise MissingDatasetTypeError(f"Dataset type {name!r} does not exist.")
689 def _get_dynamic_tables(self, dimensions: DimensionGroup) -> DynamicTables:
690 tables = self._cache.get_by_dimensions(dimensions)
691 assert tables is not None, (
692 "_fetch_dataset_types is supposed to guarantee that the tables cache is populated."
693 )
694 return tables
696 def getCollectionSummary(self, collection: CollectionRecord) -> CollectionSummary:
697 # Docstring inherited from DatasetRecordStorageManager.
698 summaries = self._summaries.fetch_summaries([collection], None, self._dataset_type_from_row)
699 return summaries[collection.key]
701 def fetch_summaries(
702 self,
703 collections: Iterable[CollectionRecord],
704 dataset_types: Iterable[DatasetType] | Iterable[str] | None = None,
705 ) -> Mapping[Any, CollectionSummary]:
706 # Docstring inherited from DatasetRecordStorageManager.
707 dataset_type_names: Iterable[str] | None = None
708 if dataset_types is not None:
709 dataset_type_names = set(get_dataset_type_name(dt) for dt in dataset_types)
710 return self._summaries.fetch_summaries(collections, dataset_type_names, self._dataset_type_from_row)
712 def fetch_run_dataset_ids(self, run: RunRecord) -> list[DatasetId]:
713 # Docstring inherited.
714 sql = sqlalchemy.select(self._static.dataset.c.id).where(
715 self._static.dataset.c[self._run_key_column] == run.key
716 )
717 with self._db.query(sql) as result:
718 return list(result.scalars())
720 def ingest_date_dtype(self) -> type:
721 """Return type of the ``ingest_date`` column."""
722 schema_version = self.newSchemaVersion()
723 if schema_version is not None and schema_version.major > 1: 723 ↛ 726line 723 didn't jump to line 726 because the condition on line 723 was always true
724 return ddl.AstropyTimeNsecTai
725 else:
726 return sqlalchemy.TIMESTAMP
728 def insert(
729 self,
730 dataset_type_name: str,
731 run: RunRecord,
732 data_ids: Iterable[DataCoordinate],
733 id_generation_mode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE,
734 ) -> list[DatasetRef]:
735 # Docstring inherited from DatasetRecordStorageManager.
736 if (storage := self._find_storage(dataset_type_name)) is None: 736 ↛ 737line 736 didn't jump to line 737 because the condition on line 736 was never true
737 raise MissingDatasetTypeError(f"Dataset type {dataset_type_name!r} has not been registered.")
738 # Current timestamp, type depends on schema version. Use microsecond
739 # precision for astropy time to keep things consistent with
740 # TIMESTAMP(6) SQL type.
741 timestamp: datetime.datetime | astropy.time.Time
742 if self._use_astropy_ingest_date: 742 ↛ 747line 742 didn't jump to line 747 because the condition on line 742 was always true
743 # Astropy `now()` precision should be the same as `now()` which
744 # should mean microsecond.
745 timestamp = astropy.time.Time.now()
746 else:
747 timestamp = datetime.datetime.now(datetime.UTC)
749 # Iterate over data IDs, transforming a possibly-single-pass iterable
750 # into a list.
751 data_id_list: list[DataCoordinate] = []
752 rows = []
753 summary = CollectionSummary()
754 for dataId in summary.add_data_ids_generator(storage.dataset_type, data_ids):
755 data_id_list.append(dataId)
756 rows.append(
757 {
758 "id": self._id_maker.makeDatasetId(
759 run.name, storage.dataset_type, dataId, id_generation_mode
760 ),
761 "dataset_type_id": storage.dataset_type_id,
762 self._run_key_column: run.key,
763 "ingest_date": timestamp,
764 }
765 )
766 if not rows: 766 ↛ 769line 766 didn't jump to line 769 because the condition on line 766 was never true
767 # Just in case an empty collection is provided we want to avoid
768 # adding dataset type to summary tables.
769 return []
771 with self._db.transaction():
772 # Insert into the static dataset table.
773 self._db.insert(self._static.dataset, *rows)
774 # Update the summary tables for this collection in case this is the
775 # first time this dataset type or these governor values will be
776 # inserted there.
777 self._summaries.update(run, [storage.dataset_type_id], summary)
778 # Combine the generated dataset_id values and data ID fields to
779 # form rows to be inserted into the tags table.
780 protoTagsRow = {
781 "dataset_type_id": storage.dataset_type_id,
782 self._collections.getCollectionForeignKeyName(): run.key,
783 }
784 tagsRows = [
785 dict(protoTagsRow, dataset_id=row["id"], **dataId.required)
786 for dataId, row in zip(data_id_list, rows, strict=True)
787 ]
788 # Insert those rows into the tags table.
789 self._db.insert(self._get_tags_table(storage.dynamic_tables), *tagsRows)
791 return [
792 DatasetRef(
793 datasetType=storage.dataset_type,
794 dataId=dataId,
795 id=row["id"],
796 run=run.name,
797 )
798 for dataId, row in zip(data_id_list, rows, strict=True)
799 ]
801 def import_(self, run: RunRecord, refs: list[DatasetRef], assume_new: bool = False) -> None:
802 # Docstring inherited from DatasetRecordStorageManager.
803 if not refs: 803 ↛ 806line 803 didn't jump to line 806 because the condition on line 803 was never true
804 # Just in case an empty mapping is provided we want to avoid
805 # adding dataset type to summary tables.
806 return
807 assert all(ref.run == run.name for ref in refs), (
808 "Run names in refs must match the run we are inserting into"
809 )
811 dataset_types = {ref.datasetType for ref in refs}
812 dimensions = _ensure_dimension_groups_match(dataset_types)
813 dataset_type_storage: dict[str, _DatasetRecordStorage] = {}
814 for dt in dataset_types:
815 if (storage := self._find_storage(dt.name)) is None: 815 ↛ 816line 815 didn't jump to line 816 because the condition on line 815 was never true
816 raise MissingDatasetTypeError(f"Dataset type {dt.name!r} has not been registered.")
817 dataset_type_storage[dt.name] = storage
819 dynamic_tables = self._get_dynamic_tables(dimensions)
820 tags_table = self._get_tags_table(dynamic_tables)
821 # Current timestamp, type depends on schema version.
822 if self._use_astropy_ingest_date: 822 ↛ 827line 822 didn't jump to line 827 because the condition on line 822 was always true
823 # Astropy `now()` precision should be the same as `now()` which
824 # should mean microsecond.
825 timestamp = sqlalchemy.sql.literal(astropy.time.Time.now(), type_=ddl.AstropyTimeNsecTai)
826 else:
827 timestamp = sqlalchemy.sql.literal(datetime.datetime.now(datetime.UTC))
828 collection_fkey_name = self._collections.getCollectionForeignKeyName()
829 tags_rows = [
830 {
831 "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
832 collection_fkey_name: run.key,
833 "dataset_id": ref.id,
834 **ref.dataId.required,
835 }
836 for ref in refs
837 ]
838 if assume_new:
839 self._import_new(run, refs, dataset_type_storage, tags_table, tags_rows, timestamp)
840 else:
841 self._import_guarded(
842 run, refs, dimensions, dataset_type_storage, tags_table, tags_rows, timestamp
843 )
845 def _import_guarded(
846 self,
847 run: RunRecord,
848 refs: list[DatasetRef],
849 dimensions: DimensionGroup,
850 dataset_type_storage: dict[str, _DatasetRecordStorage],
851 tags_table: sqlalchemy.Table,
852 tags_rows: list[dict[str, object]],
853 timestamp: sqlalchemy.BindParameter[astropy.time.Time | datetime.datetime],
854 ) -> None:
855 # We'll insert all new rows into a temporary table
856 table_spec = makeTagTableSpec(dimensions, type(self._collections), constraints=False)
857 collection_fkey_name = self._collections.getCollectionForeignKeyName()
858 with self._db.transaction(for_temp_tables=True), self._db.temporary_table(table_spec) as tmp_tags:
859 # store all incoming data in a temporary table
860 self._db.insert(tmp_tags, *tags_rows)
861 # There are some checks that we want to make for consistency
862 # of the new datasets with existing ones.
863 self._validate_import(dimensions, tags_table, tmp_tags, run)
864 # Before we merge temporary table into dataset/tags we need to
865 # drop datasets which are already there (and do not conflict).
866 self._db.deleteWhere(
867 tmp_tags,
868 tmp_tags.columns.dataset_id.in_(sqlalchemy.sql.select(self._static.dataset.columns.id)),
869 )
870 # Copy it into dataset table, need to re-label some columns.
871 self._db.insert(
872 self._static.dataset,
873 select=sqlalchemy.sql.select(
874 tmp_tags.columns.dataset_id.label("id"),
875 tmp_tags.columns.dataset_type_id,
876 tmp_tags.columns[collection_fkey_name].label(self._run_key_column),
877 timestamp.label("ingest_date"),
878 ),
879 )
880 self._update_summaries(run, refs, dataset_type_storage)
881 # Copy from temp table into tags table.
882 self._db.insert(tags_table, select=tmp_tags.select())
884 def _update_summaries(
885 self, run: RunRecord, refs: list[DatasetRef], dataset_type_storage: dict[str, _DatasetRecordStorage]
886 ) -> None:
887 summary = CollectionSummary()
888 summary.add_datasets(refs)
889 self._summaries.update(
890 run, [storage.dataset_type_id for storage in dataset_type_storage.values()], summary
891 )
893 def _validate_import(
894 self,
895 dimensions: DimensionGroup,
896 tags: sqlalchemy.schema.Table,
897 tmp_tags: sqlalchemy.schema.Table,
898 run: RunRecord,
899 ) -> None:
900 """Validate imported refs against existing datasets.
902 Parameters
903 ----------
904 dimensions : `DimensionGroup`
905 Dimensions to validate.
906 tags : `sqlalchemy.schema.Table`
907 ???
908 tmp_tags : `sqlalchemy.schema.Table`
909 Temporary table with new datasets and the same schema as tags
910 table.
911 run : `RunRecord`
912 The record object describing the `~CollectionType.RUN` collection.
914 Raises
915 ------
916 ConflictingDefinitionError
917 Raise if new datasets conflict with existing ones.
918 """
919 dataset = self._static.dataset
920 collection_fkey_name = self._collections.getCollectionForeignKeyName()
922 # Check that existing datasets have the same dataset type and
923 # run.
924 query = (
925 sqlalchemy.sql.select(
926 dataset.columns.id.label("dataset_id"),
927 dataset.columns.dataset_type_id.label("dataset_type_id"),
928 tmp_tags.columns.dataset_type_id.label("new_dataset_type_id"),
929 dataset.columns[self._run_key_column].label("run"),
930 tmp_tags.columns[collection_fkey_name].label("new_run"),
931 )
932 .select_from(dataset.join(tmp_tags, dataset.columns.id == tmp_tags.columns.dataset_id))
933 .where(
934 sqlalchemy.sql.or_(
935 dataset.columns.dataset_type_id != tmp_tags.columns.dataset_type_id,
936 dataset.columns[self._run_key_column] != tmp_tags.columns[collection_fkey_name],
937 )
938 )
939 .limit(1)
940 )
941 with self._db.query(query) as result:
942 # Only include the first one in the exception message
943 if (row := result.first()) is not None:
944 existing_run = self._collections[row.run].name
945 new_run = self._collections[row.new_run].name
946 if row.dataset_type_id == row.new_dataset_type_id:
947 raise ConflictingDefinitionError(
948 f"Current run {existing_run!r} and new run {new_run!r} do not agree for "
949 f"dataset {row.dataset_id}."
950 )
951 else:
952 raise ConflictingDefinitionError(
953 f"Dataset {row.dataset_id} was provided with type ID {row.new_dataset_type_id} "
954 f"in run {new_run!r}, but was already defined with type ID "
955 f"{row.dataset_type_id} in run {run!r}."
956 )
958 # Check that matching dataset in tags table has the same DataId.
959 query = (
960 sqlalchemy.sql.select(
961 tags.columns.dataset_id,
962 tags.columns.dataset_type_id.label("type_id"),
963 tmp_tags.columns.dataset_type_id.label("new_type_id"),
964 *[tags.columns[dim] for dim in dimensions.required],
965 *[tmp_tags.columns[dim].label(f"new_{dim}") for dim in dimensions.required],
966 )
967 .select_from(tags.join(tmp_tags, tags.columns.dataset_id == tmp_tags.columns.dataset_id))
968 .where(
969 sqlalchemy.sql.or_(
970 tags.columns.dataset_type_id != tmp_tags.columns.dataset_type_id,
971 *[tags.columns[dim] != tmp_tags.columns[dim] for dim in dimensions.required],
972 )
973 )
974 .limit(1)
975 )
977 with self._db.query(query) as result:
978 if (row := result.first()) is not None:
979 # Only include the first one in the exception message
980 raise ConflictingDefinitionError(
981 f"Existing dataset type or dataId do not match new dataset: {row._asdict()}"
982 )
984 # Check that matching run+dataId have the same dataset ID.
985 query = (
986 sqlalchemy.sql.select(
987 *[tags.columns[dim] for dim in dimensions.required],
988 tags.columns.dataset_id,
989 tmp_tags.columns.dataset_id.label("new_dataset_id"),
990 tmp_tags.columns.dataset_type_id.label("new_dataset_type_id"),
991 tags.columns[collection_fkey_name],
992 tmp_tags.columns[collection_fkey_name].label(f"new_{collection_fkey_name}"),
993 )
994 .select_from(
995 tags.join(
996 tmp_tags,
997 sqlalchemy.sql.and_(
998 tags.columns.dataset_type_id == tmp_tags.columns.dataset_type_id,
999 tags.columns[collection_fkey_name] == tmp_tags.columns[collection_fkey_name],
1000 *[tags.columns[dim] == tmp_tags.columns[dim] for dim in dimensions.required],
1001 ),
1002 )
1003 )
1004 .where(tags.columns.dataset_id != tmp_tags.columns.dataset_id)
1005 .limit(1)
1006 )
1007 with self._db.query(query) as result:
1008 # only include the first one in the exception message
1009 if (row := result.first()) is not None:
1010 data_id = {dim: getattr(row, dim) for dim in dimensions.required}
1011 existing_collection = self._collections[getattr(row, collection_fkey_name)].name
1012 new_collection = self._collections[getattr(row, f"new_{collection_fkey_name}")].name
1013 raise ConflictingDefinitionError(
1014 f"Dataset with type ID {row.new_dataset_type_id} and data ID {data_id} "
1015 f"has ID {row.dataset_id} in existing collection {existing_collection!r} "
1016 f"but ID {row.new_dataset_id} in new collection {new_collection!r}."
1017 )
1019 def _import_new(
1020 self,
1021 run: RunRecord,
1022 refs: list[DatasetRef],
1023 dataset_type_storage: dict[str, _DatasetRecordStorage],
1024 tags_table: sqlalchemy.Table,
1025 tags_rows: list[dict[str, object]],
1026 timestamp: sqlalchemy.BindParameter[astropy.time.Time | datetime.datetime],
1027 ) -> None:
1028 static_rows = [
1029 {
1030 "id": ref.id,
1031 "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
1032 self._run_key_column: run.key,
1033 "ingest_date": timestamp.value,
1034 }
1035 for ref in refs
1036 ]
1037 with self._db.transaction():
1038 self._db.insert(self._static.dataset, *static_rows)
1039 self._update_summaries(run, refs, dataset_type_storage)
1040 self._db.insert(tags_table, *tags_rows)
1042 def delete(self, datasets: Iterable[DatasetId | DatasetRef]) -> None:
1043 # Docstring inherited from DatasetRecordStorageManager.
1044 # Only delete from common dataset table; ON DELETE foreign key clauses
1045 # will handle the rest.
1046 self._db.delete(
1047 self._static.dataset,
1048 ["id"],
1049 *[{"id": getattr(dataset, "id", dataset)} for dataset in datasets],
1050 )
1052 def associate(
1053 self, dataset_type: DatasetType, collection: CollectionRecord, datasets: Iterable[DatasetRef]
1054 ) -> None:
1055 # Docstring inherited from DatasetRecordStorageManager.
1056 if (storage := self._find_storage(dataset_type.name)) is None: 1056 ↛ 1057line 1056 didn't jump to line 1057 because the condition on line 1056 was never true
1057 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1058 if collection.type is not CollectionType.TAGGED:
1059 raise CollectionTypeError(
1060 f"Cannot associate into collection '{collection.name}' "
1061 f"of type {collection.type.name}; must be TAGGED."
1062 )
1063 proto_row = {
1064 self._collections.getCollectionForeignKeyName(): collection.key,
1065 "dataset_type_id": storage.dataset_type_id,
1066 }
1067 rows = []
1068 summary = CollectionSummary()
1069 for dataset in summary.add_datasets_generator(datasets):
1070 rows.append(dict(proto_row, dataset_id=dataset.id, **dataset.dataId.required))
1071 if rows:
1072 # Update the summary tables for this collection in case this is the
1073 # first time this dataset type or these governor values will be
1074 # inserted there.
1075 self._summaries.update(collection, [storage.dataset_type_id], summary)
1076 # Update the tag table itself.
1077 self._db.replace(self._get_tags_table(storage.dynamic_tables), *rows)
1079 def disassociate(
1080 self, dataset_type: DatasetType, collection: CollectionRecord, datasets: Iterable[DatasetRef]
1081 ) -> None:
1082 # Docstring inherited from DatasetRecordStorageManager.
1083 if (storage := self._find_storage(dataset_type.name)) is None: 1083 ↛ 1084line 1083 didn't jump to line 1084 because the condition on line 1083 was never true
1084 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1085 if collection.type is not CollectionType.TAGGED: 1085 ↛ 1086line 1085 didn't jump to line 1086 because the condition on line 1085 was never true
1086 raise CollectionTypeError(
1087 f"Cannot disassociate from collection '{collection.name}' "
1088 f"of type {collection.type.name}; must be TAGGED."
1089 )
1090 rows = [
1091 {
1092 "dataset_id": dataset.id,
1093 self._collections.getCollectionForeignKeyName(): collection.key,
1094 }
1095 for dataset in datasets
1096 ]
1097 self._db.delete(
1098 self._get_tags_table(storage.dynamic_tables),
1099 ["dataset_id", self._collections.getCollectionForeignKeyName()],
1100 *rows,
1101 )
1103 def certify(
1104 self,
1105 dataset_type: DatasetType,
1106 collection: CollectionRecord,
1107 datasets: Iterable[DatasetRef],
1108 timespan: Timespan,
1109 query_func: QueryFactoryFunction,
1110 ) -> None:
1111 # Docstring inherited from DatasetRecordStorageManager.
1112 if (storage := self._find_storage(dataset_type.name)) is None: 1112 ↛ 1113line 1112 didn't jump to line 1113 because the condition on line 1112 was never true
1113 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1114 if not dataset_type.isCalibration(): 1114 ↛ 1115line 1114 didn't jump to line 1115 because the condition on line 1114 was never true
1115 raise DatasetTypeError(
1116 f"Cannot certify datasets of type {dataset_type.name!r}, for which "
1117 "DatasetType.isCalibration() is False."
1118 )
1119 if collection.type is not CollectionType.CALIBRATION: 1119 ↛ 1120line 1119 didn't jump to line 1120 because the condition on line 1119 was never true
1120 raise CollectionTypeError(
1121 f"Cannot certify into collection '{collection.name}' "
1122 f"of type {collection.type.name}; must be CALIBRATION."
1123 )
1124 TimespanReprClass = self._db.getTimespanRepresentation()
1125 proto_row = {
1126 self._collections.getCollectionForeignKeyName(): collection.key,
1127 "dataset_type_id": storage.dataset_type_id,
1128 }
1129 rows = []
1130 data_ids: set[DataCoordinate] | None = (
1131 set() if not TimespanReprClass.hasExclusionConstraint() else None
1132 )
1133 summary = CollectionSummary()
1134 for dataset in summary.add_datasets_generator(datasets):
1135 row = dict(proto_row, dataset_id=dataset.id, **dataset.dataId.required)
1136 TimespanReprClass.update(timespan, result=row)
1137 rows.append(row)
1138 if data_ids is not None: 1138 ↛ 1134line 1138 didn't jump to line 1134 because the condition on line 1138 was always true
1139 data_ids.add(dataset.dataId)
1140 if not rows: 1140 ↛ 1143line 1140 didn't jump to line 1143 because the condition on line 1140 was never true
1141 # Just in case an empty dataset collection is provided we want to
1142 # avoid adding dataset type to summary tables.
1143 return
1144 # Update the summary tables for this collection in case this is the
1145 # first time this dataset type or these governor values will be
1146 # inserted there.
1147 self._summaries.update(collection, [storage.dataset_type_id], summary)
1148 # Update the association table itself.
1149 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1150 if TimespanReprClass.hasExclusionConstraint(): 1150 ↛ 1153line 1150 didn't jump to line 1153 because the condition on line 1150 was never true
1151 # Rely on database constraint to enforce invariants; we just
1152 # reraise the exception for consistency across DB engines.
1153 try:
1154 self._db.insert(calibs_table, *rows)
1155 except sqlalchemy.exc.IntegrityError as err:
1156 raise ConflictingDefinitionError(
1157 f"Validity range conflict certifying datasets of type {dataset_type.name!r} "
1158 f"into {collection.name!r} for range {timespan}."
1159 ) from err
1160 else:
1161 # Have to implement exclusion constraint ourselves.
1162 # Acquire a table lock to ensure there are no concurrent writes
1163 # could invalidate our checking before we finish the inserts. We
1164 # use a SAVEPOINT in case there is an outer transaction that a
1165 # failure here should not roll back.
1166 with self._db.transaction(
1167 lock=[calibs_table],
1168 savepoint=True,
1169 # join_data_coordinates sometimes requires a temp table
1170 for_temp_tables=True,
1171 ):
1172 # Query for any rows that would overlap this one.
1173 with query_func() as query:
1174 if data_ids is not None: 1174 ↛ 1176line 1174 didn't jump to line 1176 because the condition on line 1174 was always true
1175 query = query.join_data_coordinates(data_ids)
1176 timespan_column = query.expression_factory[dataset_type.name].timespan
1177 result = query.datasets(dataset_type, collection.name, find_first=False).where(
1178 timespan_column.overlaps(timespan)
1179 )
1180 conflicting = result.count()
1181 if conflicting > 0:
1182 raise ConflictingDefinitionError(
1183 f"{conflicting} validity range conflicts certifying datasets of type "
1184 f"{dataset_type.name} into {collection.name} for range "
1185 f"[{timespan.begin}, {timespan.end})."
1186 )
1187 # Proceed with the insert.
1188 self._db.insert(calibs_table, *rows)
1190 def decertify(
1191 self,
1192 dataset_type: DatasetType,
1193 collection: CollectionRecord,
1194 timespan: Timespan,
1195 *,
1196 data_ids: Iterable[DataCoordinate] | None = None,
1197 query_func: QueryFactoryFunction,
1198 ) -> None:
1199 # Docstring inherited from DatasetRecordStorageManager.
1200 if (storage := self._find_storage(dataset_type.name)) is None: 1200 ↛ 1201line 1200 didn't jump to line 1201 because the condition on line 1200 was never true
1201 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1202 if not dataset_type.isCalibration(): 1202 ↛ 1203line 1202 didn't jump to line 1203 because the condition on line 1202 was never true
1203 raise DatasetTypeError(
1204 f"Cannot certify datasets of type {dataset_type.name!r}, for which "
1205 "DatasetType.isCalibration() is False."
1206 )
1207 if collection.type is not CollectionType.CALIBRATION: 1207 ↛ 1208line 1207 didn't jump to line 1208 because the condition on line 1207 was never true
1208 raise CollectionTypeError(
1209 f"Cannot decertify from collection '{collection.name}' "
1210 f"of type {collection.type.name}; must be CALIBRATION."
1211 )
1212 TimespanReprClass = self._db.getTimespanRepresentation()
1213 data_id_set: set[DataCoordinate] | None
1214 if data_ids is not None:
1215 data_id_set = set(data_ids)
1216 else:
1217 data_id_set = None
1219 # Set up collections to populate with the rows we'll want to modify.
1220 # The insert rows will have the same values for collection and
1221 # dataset type.
1222 proto_insert_row = {
1223 self._collections.getCollectionForeignKeyName(): collection.key,
1224 "dataset_type_id": storage.dataset_type_id,
1225 }
1226 rows_to_delete = []
1227 rows_to_insert = []
1228 # Acquire a table lock to ensure there are no concurrent writes
1229 # between the SELECT and the DELETE and INSERT queries based on it.
1230 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1231 with self._db.transaction(lock=[calibs_table], savepoint=True):
1232 # Find rows overlapping our inputs.
1233 with query_func() as query:
1234 query = query.join_dataset_search(dataset_type, [collection.name])
1235 if data_id_set is not None:
1236 query = query.join_data_coordinates(data_id_set)
1237 timespan_column = query.expression_factory[dataset_type.name].timespan
1238 query = query.where(timespan_column.overlaps(timespan))
1239 result = query.general(
1240 dataset_type.dimensions,
1241 dataset_fields={dataset_type.name: {"dataset_id", "timespan"}},
1242 find_first=False,
1243 )._with_added_dataset_field(dataset_type.name, "calib_pkey")
1245 calib_pkey_key = f"{dataset_type.name}.calib_pkey"
1246 dataset_id_key = f"{dataset_type.name}.dataset_id"
1247 timespan_key = f"{dataset_type.name}.timespan"
1248 for row in result.iter_tuples():
1249 rows_to_delete.append({"id": row.raw_row[calib_pkey_key]})
1250 # Construct the insert row(s) by copying the prototype row,
1251 # then adding the dimension column values, then adding
1252 # what's left of the timespan from that row after we
1253 # subtract the given timespan.
1254 new_insert_row = proto_insert_row.copy()
1255 new_insert_row["dataset_id"] = row.raw_row[dataset_id_key]
1256 for name, value in row.data_id.required.items():
1257 new_insert_row[name] = value
1258 row_timespan = row.raw_row[timespan_key]
1259 assert row_timespan is not None, "Field should have a NOT NULL constraint."
1260 for diff_timespan in row_timespan.difference(timespan):
1261 rows_to_insert.append(
1262 TimespanReprClass.update(diff_timespan, result=new_insert_row.copy())
1263 )
1264 # Run the DELETE and INSERT queries.
1265 self._db.delete(calibs_table, ["id"], *rows_to_delete)
1266 self._db.insert(calibs_table, *rows_to_insert)
1268 def make_joins_builder(
1269 self,
1270 dataset_type: DatasetType,
1271 collections: Sequence[CollectionRecord],
1272 fields: Set[qt.AnyDatasetFieldName],
1273 is_union: bool = False,
1274 ) -> SqlJoinsBuilder:
1275 if (storage := self._find_storage(dataset_type.name)) is None: 1275 ↛ 1276line 1275 didn't jump to line 1276 because the condition on line 1275 was never true
1276 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1277 # This method largely mimics `make_relation`, but it uses the new query
1278 # system primitives instead of the old one. In terms of the SQL
1279 # queries it builds, there are two more main differences:
1280 #
1281 # - Collection and run columns are now string names rather than IDs.
1282 # This insulates the query result-processing code from collection
1283 # caching and the collection manager subclass details.
1284 #
1285 # - The subquery always has unique rows, which is achieved by using
1286 # SELECT DISTINCT when necessary.
1287 #
1288 collection_types = {collection.type for collection in collections}
1289 assert CollectionType.CHAINED not in collection_types, "CHAINED collections must be flattened."
1290 #
1291 # There are two kinds of table in play here:
1292 #
1293 # - the static dataset table (with the dataset ID, dataset type ID,
1294 # run ID/name, and ingest date);
1295 #
1296 # - the dynamic tags/calibs table (with the dataset ID, dataset type
1297 # type ID, collection ID/name, data ID, and possibly validity
1298 # range).
1299 #
1300 # That means that we might want to return a query against either table
1301 # or a JOIN of both, depending on which quantities the caller wants.
1302 # But the data ID is always included, which means we'll always include
1303 # the tags/calibs table and join in the static dataset table only if we
1304 # need things from it that we can't get from the tags/calibs table.
1305 #
1306 # Note that it's important that we include a WHERE constraint on both
1307 # tables for any column (e.g. dataset_type_id) that is in both when
1308 # it's given explicitly; not doing can prevent the query planner from
1309 # using very important indexes. At present, we don't include those
1310 # redundant columns in the JOIN ON expression, however, because the
1311 # FOREIGN KEY (and its index) are defined only on dataset_id.
1312 columns = qt.ColumnSet(dataset_type.dimensions)
1313 columns.drop_implied_dimension_keys()
1314 fields_key: str | qt.AnyDatasetType = qt.ANY_DATASET if is_union else dataset_type.name
1315 columns.dataset_fields[fields_key].update(fields)
1316 tags_builder: SqlSelectBuilder | None = None
1317 if collection_types != {CollectionType.CALIBRATION}:
1318 # We'll need a subquery for the tags table if any of the given
1319 # collections are not a CALIBRATION collection. This intentionally
1320 # also fires when the list of collections is empty as a way to
1321 # create a dummy subquery that we know will fail.
1322 # We give the table an alias because it might appear multiple times
1323 # in the same query, for different dataset types.
1324 tags_table = self._get_tags_table(storage.dynamic_tables).alias(
1325 f"{dataset_type.name}_tags{'_union' if is_union else ''}"
1326 )
1327 tags_builder = self._finish_query_builder(
1328 storage,
1329 SqlJoinsBuilder(db=self._db, from_clause=tags_table).to_select_builder(columns),
1330 [record for record in collections if record.type is not CollectionType.CALIBRATION],
1331 fields,
1332 fields_key,
1333 )
1334 if "timespan" in fields:
1335 tags_builder.joins.timespans[fields_key] = self._db.getTimespanRepresentation().fromLiteral(
1336 Timespan(None, None)
1337 )
1338 assert "calib_pkey" not in fields, (
1339 "Calibration primary key for internal use only on calibration collections."
1340 )
1341 calibs_builder: SqlSelectBuilder | None = None
1342 if CollectionType.CALIBRATION in collection_types:
1343 # If at least one collection is a CALIBRATION collection, we'll
1344 # need a subquery for the calibs table, and could include the
1345 # timespan as a result or constraint.
1346 calibs_table = self._get_calibs_table(storage.dynamic_tables).alias(
1347 f"{dataset_type.name}_calibs{'_union' if is_union else ''}"
1348 )
1349 calibs_builder = self._finish_query_builder(
1350 storage,
1351 SqlJoinsBuilder(db=self._db, from_clause=calibs_table).to_select_builder(columns),
1352 [record for record in collections if record.type is CollectionType.CALIBRATION],
1353 fields,
1354 fields_key,
1355 )
1356 if "timespan" in fields:
1357 calibs_builder.joins.timespans[fields_key] = (
1358 self._db.getTimespanRepresentation().from_columns(calibs_table.columns)
1359 )
1360 if "calib_pkey" in fields:
1361 calibs_builder.joins.fields[fields_key]["calib_pkey"] = calibs_table.columns["id"]
1363 # In calibration collections, we need timespan as well as data ID
1364 # to ensure unique rows.
1365 calibs_builder.distinct = calibs_builder.distinct and "timespan" not in fields
1366 if tags_builder is not None:
1367 if calibs_builder is not None:
1368 # Need a UNION subquery.
1369 return tags_builder.union_subquery([calibs_builder])
1370 else:
1371 return tags_builder.into_joins_builder(postprocessing=None)
1372 elif calibs_builder is not None:
1373 return calibs_builder.into_joins_builder(postprocessing=None)
1374 else:
1375 raise AssertionError("Branch should be unreachable.")
1377 def _finish_query_builder(
1378 self,
1379 storage: _DatasetRecordStorage,
1380 sql_projection: SqlSelectBuilder,
1381 collections: Sequence[CollectionRecord],
1382 fields: Set[qt.AnyDatasetFieldName],
1383 fields_key: str | qt.AnyDatasetType,
1384 ) -> SqlSelectBuilder:
1385 # This method plays the same role as _finish_single_relation in the new
1386 # query system. It is called exactly one or two times by
1387 # make_sql_builder, just as _finish_single_relation is called exactly
1388 # one or two times by make_relation. See make_sql_builder comments for
1389 # what's different.
1390 assert sql_projection.joins.from_clause is not None
1391 run_collections_only = all(record.type is CollectionType.RUN for record in collections)
1392 sql_projection.joins.where(
1393 sql_projection.joins.from_clause.c.dataset_type_id == storage.dataset_type_id
1394 )
1395 dataset_id_col = sql_projection.joins.from_clause.c.dataset_id
1396 collection_col = sql_projection.joins.from_clause.c[self._collections.getCollectionForeignKeyName()]
1397 fields_provided = sql_projection.joins.fields[fields_key]
1398 # We always constrain and optionally retrieve the collection(s) via the
1399 # tags/calibs table.
1400 if "collection_key" in fields:
1401 sql_projection.joins.fields[fields_key]["collection_key"] = collection_col
1402 if len(collections) == 1:
1403 only_collection_record = collections[0]
1404 sql_projection.joins.where(collection_col == only_collection_record.key)
1405 if "collection" in fields:
1406 fields_provided["collection"] = sqlalchemy.literal(only_collection_record.name).cast(
1407 # This cast is necessary to ensure that Postgres knows the
1408 # type of this column if it is used in an aggregate
1409 # function.
1410 sqlalchemy.String
1411 )
1413 elif not collections:
1414 sql_projection.joins.where(sqlalchemy.literal(False))
1415 if "collection" in fields:
1416 fields_provided["collection"] = sqlalchemy.literal("NO COLLECTIONS")
1417 else:
1418 sql_projection.joins.where(collection_col.in_([collection.key for collection in collections]))
1419 if "collection" in fields:
1420 # Avoid a join to the collection table to get the name by using
1421 # a CASE statement. The SQL will be a bit more verbose but
1422 # more efficient.
1423 fields_provided["collection"] = _create_case_expression_for_collections(
1424 collections, collection_col
1425 )
1426 # Add more column definitions, starting with the data ID.
1427 sql_projection.joins.extract_dimensions(storage.dataset_type.dimensions.required)
1428 # We can always get the dataset_id from the tags/calibs table, even if
1429 # could also get it from the 'static' dataset table.
1430 if "dataset_id" in fields:
1431 fields_provided["dataset_id"] = dataset_id_col
1433 # It's possible we now have everything we need, from just the
1434 # tags/calibs table. The things we might need to get from the static
1435 # dataset table are the run key and the ingest date.
1436 need_static_table = False
1437 need_collection_table = False
1438 # Ingest date can only come from the static table.
1439 if "ingest_date" in fields:
1440 fields_provided["ingest_date"] = self._static.dataset.c.ingest_date
1441 need_static_table = True
1442 if "run" in fields:
1443 if len(collections) == 1 and run_collections_only:
1444 # If we are searching exactly one RUN collection, we
1445 # know that if we find the dataset in that collection,
1446 # then that's the datasets's run; we don't need to
1447 # query for it.
1448 #
1449 fields_provided["run"] = sqlalchemy.literal(only_collection_record.name).cast(
1450 # This cast is necessary to ensure that Postgres knows the
1451 # type of this column if it is used in an aggregate
1452 # function.
1453 sqlalchemy.String
1454 )
1455 elif run_collections_only:
1456 # Once again we can avoid joining to the collection table by
1457 # adding a CASE statement.
1458 fields_provided["run"] = _create_case_expression_for_collections(
1459 collections, self._static.dataset.c[self._run_key_column]
1460 )
1461 need_static_table = True
1462 else:
1463 # Here we can't avoid a join to the collection table, because
1464 # we might find a dataset via something other than its RUN
1465 # collection.
1466 #
1467 # We have to defer adding the join until after we have joined
1468 # in the static dataset table, because the ON clause involves
1469 # the run collection from the static dataset table. Postgres
1470 # cares about the join ordering (though SQLite does not.)
1471 need_collection_table = True
1472 need_static_table = True
1473 if need_static_table:
1474 # If we need the static table, join it in via dataset_id. We don't
1475 # use SqlJoinsBuilder.join because we're joining on dataset ID, not
1476 # dimensions.
1477 sql_projection.joins.from_clause = sql_projection.joins.from_clause.join(
1478 self._static.dataset, onclause=(dataset_id_col == self._static.dataset.c.id)
1479 )
1480 # Also constrain dataset_type_id in static table in case that helps
1481 # generate a better plan. We could also include this in the JOIN ON
1482 # clause, but my guess is that that's a good idea IFF it's in the
1483 # foreign key, and right now it isn't.
1484 sql_projection.joins.where(self._static.dataset.c.dataset_type_id == storage.dataset_type_id)
1485 if need_collection_table:
1486 # Join the collection table to look up the RUN collection name
1487 # associated with the dataset.
1488 (
1489 fields_provided["run"],
1490 sql_projection.joins.from_clause,
1491 ) = self._collections.lookup_name_sql(
1492 self._static.dataset.c[self._run_key_column],
1493 sql_projection.joins.from_clause,
1494 )
1496 sql_projection.distinct = (
1497 # If there are multiple collections, this subquery might have
1498 # non-unique rows.
1499 len(collections) > 1 and not fields
1500 )
1501 return sql_projection
1503 def refresh_collection_summaries(self, dataset_type: DatasetType) -> None:
1504 # Docstring inherited.
1505 if (storage := self._find_storage(dataset_type.name)) is None: 1505 ↛ 1506line 1505 didn't jump to line 1506 because the condition on line 1505 was never true
1506 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1507 with self._db.transaction():
1508 # The main issue here is consistency in the presence of concurrent
1509 # updates (using default READ COMMITTED isolation). Regular clients
1510 # only add to summary tables, and we want to avoid deleting what
1511 # other concurrent transactions may add while we are in this
1512 # transaction. This ordering of operations should guarantee it:
1513 # - read collections for this dataset type from summary tables,
1514 # - read collections for this dataset type from dataset tables
1515 # (both tags and calibs),
1516 # - whatever is in the first set but not in the second can be
1517 # dropped from summary tables.
1518 summary_collection_ids = set(self._summaries.get_collection_ids(storage.dataset_type_id))
1520 # Query datasets tables for associated collections.
1521 column_name = self._collections.getCollectionForeignKeyName()
1522 tags_table = self._get_tags_table(storage.dynamic_tables)
1523 query: sqlalchemy.sql.expression.SelectBase = (
1524 sqlalchemy.select(tags_table.columns[column_name])
1525 .where(tags_table.columns.dataset_type_id == storage.dataset_type_id)
1526 .distinct()
1527 )
1528 if dataset_type.isCalibration():
1529 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1530 query2 = (
1531 sqlalchemy.select(calibs_table.columns[column_name])
1532 .where(calibs_table.columns.dataset_type_id == storage.dataset_type_id)
1533 .distinct()
1534 )
1535 query = sqlalchemy.sql.expression.union(query, query2)
1537 with self._db.query(query) as result:
1538 collection_ids = set(result.scalars())
1540 collections_to_delete = summary_collection_ids - collection_ids
1541 self._summaries.delete_collections(storage.dataset_type_id, collections_to_delete)
1543 def _get_tags_table(self, table: DynamicTables) -> sqlalchemy.Table:
1544 return table.tags(self._db, type(self._collections), self._cache.tables)
1546 def _get_calibs_table(self, table: DynamicTables) -> sqlalchemy.Table:
1547 return table.calibs(self._db, type(self._collections), self._cache.tables)
1550def _create_case_expression_for_collections(
1551 collections: Iterable[CollectionRecord], id_column: sqlalchemy.ColumnElement
1552) -> sqlalchemy.ColumnElement:
1553 """Return a SQLAlchemy Case expression that converts collection IDs to
1554 collection names for the given set of collections.
1556 Parameters
1557 ----------
1558 collections : `~collections.abc.Iterable` [ `CollectionRecord` ]
1559 List of collections to include in conversion table. This should be an
1560 exhaustive list of collections that could appear in `id_column`.
1561 id_column : `sqlalchemy.ColumnElement`
1562 The column containing the collection ID that we want to convert to a
1563 collection name.
1564 """
1565 mapping = {record.key: record.name for record in collections}
1566 if not mapping:
1567 # SQLAlchemy does not correctly handle an empty mapping in case() -- it
1568 # crashes when trying to compile the expression with an
1569 # "AttributeError('NoneType' object has no attribute 'dialect_impl')"
1570 # when trying to access the 'type' property of the Case object. If you
1571 # explicitly specify a type via type_coerce it instead generates
1572 # invalid SQL syntax.
1573 #
1574 # We can end up with empty mappings here in certain "doomed query" edge
1575 # cases, e.g. we start with a list of valid collections but they are
1576 # all filtered out by higher-level code on the basis of collection
1577 # summaries.
1578 return sqlalchemy.cast(sqlalchemy.null(), sqlalchemy.String)
1580 return sqlalchemy.case(mapping, value=id_column)
1583def _ensure_dimension_groups_match(dataset_types: Iterable[DatasetType]) -> DimensionGroup:
1584 dimensions = set(dt.dimensions for dt in dataset_types)
1585 assert len(dimensions) > 0, "At least one dataset type is required"
1586 if len(dimensions) != 1: 1586 ↛ 1587line 1586 didn't jump to line 1587 because the condition on line 1586 was never true
1587 raise DatasetTypeError(
1588 "Dataset types have more than one dimension group.\n"
1589 f"Dataset types: {dataset_types}\n"
1590 f"Dimension groups: {dimensions}"
1591 )
1592 return dimensions.pop()