Coverage for python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py: 92%
583 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 18:11 +0000
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(self._db.make_in_array_constraint(id_col, 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(
528 self._db.make_in_array_constraint(tags_table.columns["dataset_id"], batch)
529 )
530 # Join in the collection table to fetch the run name.
531 collection_column = tags_table.columns[self._collections.getCollectionForeignKeyName()]
532 joined_collections = self._collections.join_collections_sql(collection_column, tags_sql)
533 tags_sql = joined_collections.joined_sql
534 run_name_column = joined_collections.name_column
535 tags_sql = tags_sql.add_columns(run_name_column)
536 # Tags table includes run collections and tagged
537 # collections.
538 # In theory the data ID for a given dataset should be the
539 # same in both, but nothing actually guarantees this.
540 # So skip any tagged collections, using the run collection
541 # as the definitive definition.
542 tags_sql = tags_sql.where(joined_collections.type_column == int(CollectionType.RUN))
544 with self._db.query(tags_sql) as sql_result:
545 data_id_rows = sql_result.mappings().all()
547 assert run_name_column.key is not None
548 for data_id_row in data_id_rows:
549 id = data_id_row["dataset_id"]
550 dataset_type = dataset_type_map[id]
551 run_name = data_id_row[run_name_column.key]
552 data_id = DataCoordinate.from_required_values(
553 dimension_group,
554 tuple(data_id_row[dimension] for dimension in dimension_group.required),
555 )
556 ref = DatasetRef(
557 datasetType=dataset_type,
558 dataId=data_id,
559 id=id,
560 run=run_name,
561 )
562 output_refs.append(ref)
564 return output_refs
566 def _fetch_dataset_type_record(self, name: str) -> _DatasetTypeRecord | None:
567 """Retrieve all dataset types defined in database.
569 Returns
570 -------
571 dataset_type
572 Information from a single database record.
573 """
574 db_name = self._reversed_renames.get(name, name)
575 if db_name == name and db_name in self._renames:
576 # We were asked to fetch a dataset type that has been renamed via
577 # a config override, with no other existing dataset type renamed
578 # to replace it. It should appear missing.
579 return None
580 c = self._static.dataset_type.columns
581 stmt = self._static.dataset_type.select().where(c.name == db_name)
582 with self._db.query(stmt) as sql_result:
583 row = sql_result.mappings().one_or_none()
584 if row is None:
585 return None
586 else:
587 return self._record_from_row(row)
589 def _record_from_row(self, row: Mapping) -> _DatasetTypeRecord:
590 name = self._renames.get(row["name"], row["name"])
591 dimensions = self._dimensions.load_dimension_group(row["dimensions_key"])
592 calibTableName = row["calibration_association_table"]
593 storage_class_name = self._storage_class_overrides.get(row["name"], row["storage_class"])
594 datasetType = DatasetType(
595 name, dimensions, storage_class_name, isCalibration=(calibTableName is not None)
596 )
597 return _DatasetTypeRecord(
598 dataset_type=datasetType,
599 dataset_type_id=row["id"],
600 dimensions_key=row["dimensions_key"],
601 tag_table_name=row["tag_association_table"],
602 calib_table_name=calibTableName,
603 )
605 def _dataset_type_from_row(self, row: Mapping) -> DatasetType:
606 return self._record_from_row(row).dataset_type
608 def preload_cache(self) -> None:
609 self._fetch_dataset_types()
611 def _fetch_dataset_types(self, force_refresh: bool = False) -> list[DatasetType]:
612 """Fetch list of all defined dataset types."""
613 # This is one of two places we populate the dataset type cache:
614 #
615 # - This method handles almost all requests for dataset types that
616 # should already exist. It always marks the cache as "full" in both
617 # dataset type names and dimensions.
618 #
619 # - register_dataset_type handles the case where the dataset type might
620 # not existing yet. Since it can only add a single dataset type, it
621 # never changes whether the cache is full.
622 #
623 # In both cases, we require that the per-dimensions data be cached
624 # whenever a dataset type is added to the cache by name, to reduce the
625 # number of possible states the cache can be in and minimize the number
626 # of queries.
627 if self._cache.full and not force_refresh:
628 return [dataset_type for dataset_type, _ in self._cache.items()]
629 with self._db.query(self._static.dataset_type.select()) as sql_result:
630 sql_rows = sql_result.mappings().fetchall()
631 records = [self._record_from_row(row) for row in sql_rows]
632 # Cache everything and specify that cache is complete.
633 cache_data: list[tuple[DatasetType, int]] = []
634 cache_dimensions_data: dict[DimensionGroup, DynamicTables] = {}
635 for record in records:
636 cache_data.append((record.dataset_type, record.dataset_type_id))
637 if (dynamic_tables := cache_dimensions_data.get(record.dataset_type.dimensions)) is None:
638 tables = record.make_dynamic_tables()
639 else:
640 tables = record.update_dynamic_tables(dynamic_tables)
641 cache_dimensions_data[record.dataset_type.dimensions] = tables
642 self._cache.set(
643 cache_data, full=True, dimensions_data=cache_dimensions_data.items(), dimensions_full=True
644 )
645 return [record.dataset_type for record in records]
647 def _get_dataset_type_by_id(self, id: int) -> DatasetType:
648 dt = self._cache.get_by_id(id)
649 if dt is None:
650 # Since the ID is not a concept exposed to the public API, it
651 # had to have come from a dataset table, and our cache must be
652 # empty or out of date.
653 self._fetch_dataset_types(force_refresh=True)
654 dt = self._cache.get_by_id(id)
655 if dt is None: 655 ↛ 656line 655 didn't jump to line 656 because the condition on line 655 was never true
656 raise RuntimeError(f"Failed to look up dataset type with ID {id}")
657 return dt
659 def _find_storage(self, name: str) -> _DatasetRecordStorage:
660 """Find a dataset type and the extra information needed to work with
661 it, utilizing and populating the cache as needed.
662 """
663 dataset_type, dataset_type_id = self._cache.get(name)
664 if dataset_type is not None:
665 tables = self._get_dynamic_tables(dataset_type.dimensions)
666 assert dataset_type_id is not None, "Dataset type cache population is incomplete."
667 return _DatasetRecordStorage(
668 dataset_type=dataset_type, dataset_type_id=dataset_type_id, dynamic_tables=tables
669 )
670 else:
671 # On the first cache miss populate the cache with complete list
672 # of dataset types (if it was not done yet).
673 if not self._cache.full:
674 self._fetch_dataset_types()
675 # Try again
676 dataset_type, dataset_type_id = self._cache.get(name)
677 if dataset_type is not None:
678 tables = self._get_dynamic_tables(dataset_type.dimensions)
679 assert dataset_type_id is not None, "Dataset type cache population is incomplete."
680 return _DatasetRecordStorage(
681 dataset_type=dataset_type, dataset_type_id=dataset_type_id, dynamic_tables=tables
682 )
683 record = self._fetch_dataset_type_record(name)
684 if record is not None:
685 self._cache.add(record.dataset_type, record.dataset_type_id)
686 tables = record.make_dynamic_tables()
687 self._cache.add_by_dimensions(record.dataset_type.dimensions, tables)
688 return _DatasetRecordStorage(record.dataset_type, record.dataset_type_id, tables)
689 raise MissingDatasetTypeError(f"Dataset type {name!r} does not exist.")
691 def _get_dynamic_tables(self, dimensions: DimensionGroup) -> DynamicTables:
692 tables = self._cache.get_by_dimensions(dimensions)
693 assert tables is not None, (
694 "_fetch_dataset_types is supposed to guarantee that the tables cache is populated."
695 )
696 return tables
698 def getCollectionSummary(self, collection: CollectionRecord) -> CollectionSummary:
699 # Docstring inherited from DatasetRecordStorageManager.
700 summaries = self._summaries.fetch_summaries([collection], None, self._dataset_type_from_row)
701 return summaries[collection.key]
703 def fetch_summaries(
704 self,
705 collections: Iterable[CollectionRecord],
706 dataset_types: Iterable[DatasetType] | Iterable[str] | None = None,
707 ) -> Mapping[Any, CollectionSummary]:
708 # Docstring inherited from DatasetRecordStorageManager.
709 dataset_type_names: Iterable[str] | None = None
710 if dataset_types is not None:
711 dataset_type_names = set(get_dataset_type_name(dt) for dt in dataset_types)
712 return self._summaries.fetch_summaries(collections, dataset_type_names, self._dataset_type_from_row)
714 def fetch_run_dataset_ids(self, run: RunRecord) -> list[DatasetId]:
715 # Docstring inherited.
716 sql = sqlalchemy.select(self._static.dataset.c.id).where(
717 self._static.dataset.c[self._run_key_column] == run.key
718 )
719 with self._db.query(sql) as result:
720 return list(result.scalars())
722 def ingest_date_dtype(self) -> type:
723 """Return type of the ``ingest_date`` column."""
724 schema_version = self.newSchemaVersion()
725 if schema_version is not None and schema_version.major > 1: 725 ↛ 728line 725 didn't jump to line 728 because the condition on line 725 was always true
726 return ddl.AstropyTimeNsecTai
727 else:
728 return sqlalchemy.TIMESTAMP
730 def insert(
731 self,
732 dataset_type_name: str,
733 run: RunRecord,
734 data_ids: Iterable[DataCoordinate],
735 id_generation_mode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE,
736 ) -> list[DatasetRef]:
737 # Docstring inherited from DatasetRecordStorageManager.
738 if (storage := self._find_storage(dataset_type_name)) is None: 738 ↛ 739line 738 didn't jump to line 739 because the condition on line 738 was never true
739 raise MissingDatasetTypeError(f"Dataset type {dataset_type_name!r} has not been registered.")
740 # Current timestamp, type depends on schema version. Use microsecond
741 # precision for astropy time to keep things consistent with
742 # TIMESTAMP(6) SQL type.
743 timestamp: datetime.datetime | astropy.time.Time
744 if self._use_astropy_ingest_date: 744 ↛ 749line 744 didn't jump to line 749 because the condition on line 744 was always true
745 # Astropy `now()` precision should be the same as `now()` which
746 # should mean microsecond.
747 timestamp = astropy.time.Time.now()
748 else:
749 timestamp = datetime.datetime.now(datetime.UTC)
751 # Iterate over data IDs, transforming a possibly-single-pass iterable
752 # into a list.
753 data_id_list: list[DataCoordinate] = []
754 rows = []
755 summary = CollectionSummary()
756 for dataId in summary.add_data_ids_generator(storage.dataset_type, data_ids):
757 data_id_list.append(dataId)
758 rows.append(
759 {
760 "id": self._id_maker.makeDatasetId(
761 run.name, storage.dataset_type, dataId, id_generation_mode
762 ),
763 "dataset_type_id": storage.dataset_type_id,
764 self._run_key_column: run.key,
765 "ingest_date": timestamp,
766 }
767 )
768 if not rows: 768 ↛ 771line 768 didn't jump to line 771 because the condition on line 768 was never true
769 # Just in case an empty collection is provided we want to avoid
770 # adding dataset type to summary tables.
771 return []
773 with self._db.transaction():
774 # Insert into the static dataset table.
775 self._db.insert(self._static.dataset, *rows)
776 # Update the summary tables for this collection in case this is the
777 # first time this dataset type or these governor values will be
778 # inserted there.
779 self._summaries.update(run, [storage.dataset_type_id], summary)
780 # Combine the generated dataset_id values and data ID fields to
781 # form rows to be inserted into the tags table.
782 protoTagsRow = {
783 "dataset_type_id": storage.dataset_type_id,
784 self._collections.getCollectionForeignKeyName(): run.key,
785 }
786 tagsRows = [
787 dict(protoTagsRow, dataset_id=row["id"], **dataId.required)
788 for dataId, row in zip(data_id_list, rows, strict=True)
789 ]
790 # Insert those rows into the tags table.
791 self._db.insert(self._get_tags_table(storage.dynamic_tables), *tagsRows)
793 return [
794 DatasetRef(
795 datasetType=storage.dataset_type,
796 dataId=dataId,
797 id=row["id"],
798 run=run.name,
799 )
800 for dataId, row in zip(data_id_list, rows, strict=True)
801 ]
803 def import_(self, run: RunRecord, refs: list[DatasetRef], assume_new: bool = False) -> None:
804 # Docstring inherited from DatasetRecordStorageManager.
805 if not refs: 805 ↛ 808line 805 didn't jump to line 808 because the condition on line 805 was never true
806 # Just in case an empty mapping is provided we want to avoid
807 # adding dataset type to summary tables.
808 return
809 assert all(ref.run == run.name for ref in refs), (
810 "Run names in refs must match the run we are inserting into"
811 )
813 dataset_types = {ref.datasetType for ref in refs}
814 dimensions = _ensure_dimension_groups_match(dataset_types)
815 dataset_type_storage: dict[str, _DatasetRecordStorage] = {}
816 for dt in dataset_types:
817 if (storage := self._find_storage(dt.name)) is None: 817 ↛ 818line 817 didn't jump to line 818 because the condition on line 817 was never true
818 raise MissingDatasetTypeError(f"Dataset type {dt.name!r} has not been registered.")
819 dataset_type_storage[dt.name] = storage
821 dynamic_tables = self._get_dynamic_tables(dimensions)
822 tags_table = self._get_tags_table(dynamic_tables)
823 # Current timestamp, type depends on schema version.
824 if self._use_astropy_ingest_date: 824 ↛ 829line 824 didn't jump to line 829 because the condition on line 824 was always true
825 # Astropy `now()` precision should be the same as `now()` which
826 # should mean microsecond.
827 timestamp = sqlalchemy.sql.literal(astropy.time.Time.now(), type_=ddl.AstropyTimeNsecTai)
828 else:
829 timestamp = sqlalchemy.sql.literal(datetime.datetime.now(datetime.UTC))
830 collection_fkey_name = self._collections.getCollectionForeignKeyName()
831 tags_rows = [
832 {
833 "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
834 collection_fkey_name: run.key,
835 "dataset_id": ref.id,
836 **ref.dataId.required,
837 }
838 for ref in refs
839 ]
840 if assume_new:
841 self._import_new(run, refs, dataset_type_storage, tags_table, tags_rows, timestamp)
842 else:
843 self._import_guarded(
844 run, refs, dimensions, dataset_type_storage, tags_table, tags_rows, timestamp
845 )
847 def _import_guarded(
848 self,
849 run: RunRecord,
850 refs: list[DatasetRef],
851 dimensions: DimensionGroup,
852 dataset_type_storage: dict[str, _DatasetRecordStorage],
853 tags_table: sqlalchemy.Table,
854 tags_rows: list[dict[str, object]],
855 timestamp: sqlalchemy.BindParameter[astropy.time.Time | datetime.datetime],
856 ) -> None:
857 # We'll insert all new rows into a temporary table
858 table_spec = makeTagTableSpec(dimensions, type(self._collections), constraints=False)
859 collection_fkey_name = self._collections.getCollectionForeignKeyName()
860 with self._db.transaction(for_temp_tables=True), self._db.temporary_table(table_spec) as tmp_tags:
861 # store all incoming data in a temporary table
862 self._db.insert(tmp_tags, *tags_rows)
863 # There are some checks that we want to make for consistency
864 # of the new datasets with existing ones.
865 self._validate_import(dimensions, tags_table, tmp_tags, run)
866 # Before we merge temporary table into dataset/tags we need to
867 # drop datasets which are already there (and do not conflict).
868 self._db.deleteWhere(
869 tmp_tags,
870 tmp_tags.columns.dataset_id.in_(sqlalchemy.sql.select(self._static.dataset.columns.id)),
871 )
872 # Copy it into dataset table, need to re-label some columns.
873 self._db.insert(
874 self._static.dataset,
875 select=sqlalchemy.sql.select(
876 tmp_tags.columns.dataset_id.label("id"),
877 tmp_tags.columns.dataset_type_id,
878 tmp_tags.columns[collection_fkey_name].label(self._run_key_column),
879 timestamp.label("ingest_date"),
880 ),
881 )
882 self._update_summaries(run, refs, dataset_type_storage)
883 # Copy from temp table into tags table.
884 self._db.insert(tags_table, select=tmp_tags.select())
886 def _update_summaries(
887 self, run: RunRecord, refs: list[DatasetRef], dataset_type_storage: dict[str, _DatasetRecordStorage]
888 ) -> None:
889 summary = CollectionSummary()
890 summary.add_datasets(refs)
891 self._summaries.update(
892 run, [storage.dataset_type_id for storage in dataset_type_storage.values()], summary
893 )
895 def _validate_import(
896 self,
897 dimensions: DimensionGroup,
898 tags: sqlalchemy.schema.Table,
899 tmp_tags: sqlalchemy.schema.Table,
900 run: RunRecord,
901 ) -> None:
902 """Validate imported refs against existing datasets.
904 Parameters
905 ----------
906 dimensions : `DimensionGroup`
907 Dimensions to validate.
908 tags : `sqlalchemy.schema.Table`
909 ???
910 tmp_tags : `sqlalchemy.schema.Table`
911 Temporary table with new datasets and the same schema as tags
912 table.
913 run : `RunRecord`
914 The record object describing the `~CollectionType.RUN` collection.
916 Raises
917 ------
918 ConflictingDefinitionError
919 Raise if new datasets conflict with existing ones.
920 """
921 dataset = self._static.dataset
922 collection_fkey_name = self._collections.getCollectionForeignKeyName()
924 # Check that existing datasets have the same dataset type and
925 # run.
926 query = (
927 sqlalchemy.sql.select(
928 dataset.columns.id.label("dataset_id"),
929 dataset.columns.dataset_type_id.label("dataset_type_id"),
930 tmp_tags.columns.dataset_type_id.label("new_dataset_type_id"),
931 dataset.columns[self._run_key_column].label("run"),
932 tmp_tags.columns[collection_fkey_name].label("new_run"),
933 )
934 .select_from(dataset.join(tmp_tags, dataset.columns.id == tmp_tags.columns.dataset_id))
935 .where(
936 sqlalchemy.sql.or_(
937 dataset.columns.dataset_type_id != tmp_tags.columns.dataset_type_id,
938 dataset.columns[self._run_key_column] != tmp_tags.columns[collection_fkey_name],
939 )
940 )
941 .limit(1)
942 )
943 with self._db.query(query) as result:
944 # Only include the first one in the exception message
945 if (row := result.first()) is not None:
946 existing_run = self._collections[row.run].name
947 new_run = self._collections[row.new_run].name
948 if row.dataset_type_id == row.new_dataset_type_id:
949 raise ConflictingDefinitionError(
950 f"Current run {existing_run!r} and new run {new_run!r} do not agree for "
951 f"dataset {row.dataset_id}."
952 )
953 else:
954 raise ConflictingDefinitionError(
955 f"Dataset {row.dataset_id} was provided with type ID {row.new_dataset_type_id} "
956 f"in run {new_run!r}, but was already defined with type ID "
957 f"{row.dataset_type_id} in run {run!r}."
958 )
960 # Check that matching dataset in tags table has the same DataId.
961 query = (
962 sqlalchemy.sql.select(
963 tags.columns.dataset_id,
964 tags.columns.dataset_type_id.label("type_id"),
965 tmp_tags.columns.dataset_type_id.label("new_type_id"),
966 *[tags.columns[dim] for dim in dimensions.required],
967 *[tmp_tags.columns[dim].label(f"new_{dim}") for dim in dimensions.required],
968 )
969 .select_from(tags.join(tmp_tags, tags.columns.dataset_id == tmp_tags.columns.dataset_id))
970 .where(
971 sqlalchemy.sql.or_(
972 tags.columns.dataset_type_id != tmp_tags.columns.dataset_type_id,
973 *[tags.columns[dim] != tmp_tags.columns[dim] for dim in dimensions.required],
974 )
975 )
976 .limit(1)
977 )
979 with self._db.query(query) as result:
980 if (row := result.first()) is not None:
981 # Only include the first one in the exception message
982 raise ConflictingDefinitionError(
983 f"Existing dataset type or dataId do not match new dataset: {row._asdict()}"
984 )
986 # Check that matching run+dataId have the same dataset ID.
987 query = (
988 sqlalchemy.sql.select(
989 *[tags.columns[dim] for dim in dimensions.required],
990 tags.columns.dataset_id,
991 tmp_tags.columns.dataset_id.label("new_dataset_id"),
992 tmp_tags.columns.dataset_type_id.label("new_dataset_type_id"),
993 tags.columns[collection_fkey_name],
994 tmp_tags.columns[collection_fkey_name].label(f"new_{collection_fkey_name}"),
995 )
996 .select_from(
997 tags.join(
998 tmp_tags,
999 sqlalchemy.sql.and_(
1000 tags.columns.dataset_type_id == tmp_tags.columns.dataset_type_id,
1001 tags.columns[collection_fkey_name] == tmp_tags.columns[collection_fkey_name],
1002 *[tags.columns[dim] == tmp_tags.columns[dim] for dim in dimensions.required],
1003 ),
1004 )
1005 )
1006 .where(tags.columns.dataset_id != tmp_tags.columns.dataset_id)
1007 .limit(1)
1008 )
1009 with self._db.query(query) as result:
1010 # only include the first one in the exception message
1011 if (row := result.first()) is not None:
1012 data_id = {dim: getattr(row, dim) for dim in dimensions.required}
1013 existing_collection = self._collections[getattr(row, collection_fkey_name)].name
1014 new_collection = self._collections[getattr(row, f"new_{collection_fkey_name}")].name
1015 raise ConflictingDefinitionError(
1016 f"Dataset with type ID {row.new_dataset_type_id} and data ID {data_id} "
1017 f"has ID {row.dataset_id} in existing collection {existing_collection!r} "
1018 f"but ID {row.new_dataset_id} in new collection {new_collection!r}."
1019 )
1021 def _import_new(
1022 self,
1023 run: RunRecord,
1024 refs: list[DatasetRef],
1025 dataset_type_storage: dict[str, _DatasetRecordStorage],
1026 tags_table: sqlalchemy.Table,
1027 tags_rows: list[dict[str, object]],
1028 timestamp: sqlalchemy.BindParameter[astropy.time.Time | datetime.datetime],
1029 ) -> None:
1030 static_rows = [
1031 {
1032 "id": ref.id,
1033 "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
1034 self._run_key_column: run.key,
1035 "ingest_date": timestamp.value,
1036 }
1037 for ref in refs
1038 ]
1039 with self._db.transaction():
1040 self._db.insert(self._static.dataset, *static_rows)
1041 self._update_summaries(run, refs, dataset_type_storage)
1042 self._db.insert(tags_table, *tags_rows)
1044 def delete(self, datasets: Iterable[DatasetId | DatasetRef]) -> None:
1045 # Docstring inherited from DatasetRecordStorageManager.
1046 # Only delete from common dataset table; ON DELETE foreign key clauses
1047 # will handle the rest.
1048 self._db.delete(
1049 self._static.dataset,
1050 ["id"],
1051 *[{"id": getattr(dataset, "id", dataset)} for dataset in datasets],
1052 )
1054 def associate(
1055 self, dataset_type: DatasetType, collection: CollectionRecord, datasets: Iterable[DatasetRef]
1056 ) -> None:
1057 # Docstring inherited from DatasetRecordStorageManager.
1058 if (storage := self._find_storage(dataset_type.name)) is None: 1058 ↛ 1059line 1058 didn't jump to line 1059 because the condition on line 1058 was never true
1059 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1060 if collection.type is not CollectionType.TAGGED:
1061 raise CollectionTypeError(
1062 f"Cannot associate into collection '{collection.name}' "
1063 f"of type {collection.type.name}; must be TAGGED."
1064 )
1065 proto_row = {
1066 self._collections.getCollectionForeignKeyName(): collection.key,
1067 "dataset_type_id": storage.dataset_type_id,
1068 }
1069 rows = []
1070 summary = CollectionSummary()
1071 for dataset in summary.add_datasets_generator(datasets):
1072 rows.append(dict(proto_row, dataset_id=dataset.id, **dataset.dataId.required))
1073 if rows:
1074 # Update the summary tables for this collection in case this is the
1075 # first time this dataset type or these governor values will be
1076 # inserted there.
1077 self._summaries.update(collection, [storage.dataset_type_id], summary)
1078 # Update the tag table itself.
1079 self._db.replace(self._get_tags_table(storage.dynamic_tables), *rows)
1081 def disassociate(
1082 self, dataset_type: DatasetType, collection: CollectionRecord, datasets: Iterable[DatasetRef]
1083 ) -> None:
1084 # Docstring inherited from DatasetRecordStorageManager.
1085 if (storage := self._find_storage(dataset_type.name)) is None: 1085 ↛ 1086line 1085 didn't jump to line 1086 because the condition on line 1085 was never true
1086 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1087 if collection.type is not CollectionType.TAGGED: 1087 ↛ 1088line 1087 didn't jump to line 1088 because the condition on line 1087 was never true
1088 raise CollectionTypeError(
1089 f"Cannot disassociate from collection '{collection.name}' "
1090 f"of type {collection.type.name}; must be TAGGED."
1091 )
1092 rows = [
1093 {
1094 "dataset_id": dataset.id,
1095 self._collections.getCollectionForeignKeyName(): collection.key,
1096 }
1097 for dataset in datasets
1098 ]
1099 self._db.delete(
1100 self._get_tags_table(storage.dynamic_tables),
1101 ["dataset_id", self._collections.getCollectionForeignKeyName()],
1102 *rows,
1103 )
1105 def certify(
1106 self,
1107 dataset_type: DatasetType,
1108 collection: CollectionRecord,
1109 datasets: Iterable[DatasetRef],
1110 timespan: Timespan,
1111 query_func: QueryFactoryFunction,
1112 ) -> None:
1113 # Docstring inherited from DatasetRecordStorageManager.
1114 if (storage := self._find_storage(dataset_type.name)) is None: 1114 ↛ 1115line 1114 didn't jump to line 1115 because the condition on line 1114 was never true
1115 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1116 if not dataset_type.isCalibration(): 1116 ↛ 1117line 1116 didn't jump to line 1117 because the condition on line 1116 was never true
1117 raise DatasetTypeError(
1118 f"Cannot certify datasets of type {dataset_type.name!r}, for which "
1119 "DatasetType.isCalibration() is False."
1120 )
1121 if collection.type is not CollectionType.CALIBRATION: 1121 ↛ 1122line 1121 didn't jump to line 1122 because the condition on line 1121 was never true
1122 raise CollectionTypeError(
1123 f"Cannot certify into collection '{collection.name}' "
1124 f"of type {collection.type.name}; must be CALIBRATION."
1125 )
1126 TimespanReprClass = self._db.getTimespanRepresentation()
1127 proto_row = {
1128 self._collections.getCollectionForeignKeyName(): collection.key,
1129 "dataset_type_id": storage.dataset_type_id,
1130 }
1131 rows = []
1132 data_ids: set[DataCoordinate] | None = (
1133 set() if not TimespanReprClass.hasExclusionConstraint() else None
1134 )
1135 summary = CollectionSummary()
1136 for dataset in summary.add_datasets_generator(datasets):
1137 row = dict(proto_row, dataset_id=dataset.id, **dataset.dataId.required)
1138 TimespanReprClass.update(timespan, result=row)
1139 rows.append(row)
1140 if data_ids is not None: 1140 ↛ 1136line 1140 didn't jump to line 1136 because the condition on line 1140 was always true
1141 data_ids.add(dataset.dataId)
1142 if not rows: 1142 ↛ 1145line 1142 didn't jump to line 1145 because the condition on line 1142 was never true
1143 # Just in case an empty dataset collection is provided we want to
1144 # avoid adding dataset type to summary tables.
1145 return
1146 # Update the summary tables for this collection in case this is the
1147 # first time this dataset type or these governor values will be
1148 # inserted there.
1149 self._summaries.update(collection, [storage.dataset_type_id], summary)
1150 # Update the association table itself.
1151 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1152 if TimespanReprClass.hasExclusionConstraint(): 1152 ↛ 1155line 1152 didn't jump to line 1155 because the condition on line 1152 was never true
1153 # Rely on database constraint to enforce invariants; we just
1154 # reraise the exception for consistency across DB engines.
1155 try:
1156 self._db.insert(calibs_table, *rows)
1157 except sqlalchemy.exc.IntegrityError as err:
1158 raise ConflictingDefinitionError(
1159 f"Validity range conflict certifying datasets of type {dataset_type.name!r} "
1160 f"into {collection.name!r} for range {timespan}."
1161 ) from err
1162 else:
1163 # Have to implement exclusion constraint ourselves.
1164 # Acquire a table lock to ensure there are no concurrent writes
1165 # could invalidate our checking before we finish the inserts. We
1166 # use a SAVEPOINT in case there is an outer transaction that a
1167 # failure here should not roll back.
1168 with self._db.transaction(
1169 lock=[calibs_table],
1170 savepoint=True,
1171 # join_data_coordinates sometimes requires a temp table
1172 for_temp_tables=True,
1173 ):
1174 # Query for any rows that would overlap this one.
1175 with query_func() as query:
1176 if data_ids is not None: 1176 ↛ 1178line 1176 didn't jump to line 1178 because the condition on line 1176 was always true
1177 query = query.join_data_coordinates(data_ids)
1178 timespan_column = query.expression_factory[dataset_type.name].timespan
1179 result = query.datasets(dataset_type, collection.name, find_first=False).where(
1180 timespan_column.overlaps(timespan)
1181 )
1182 conflicting = result.count()
1183 if conflicting > 0:
1184 raise ConflictingDefinitionError(
1185 f"{conflicting} validity range conflicts certifying datasets of type "
1186 f"{dataset_type.name} into {collection.name} for range "
1187 f"[{timespan.begin}, {timespan.end})."
1188 )
1189 # Proceed with the insert.
1190 self._db.insert(calibs_table, *rows)
1192 def decertify(
1193 self,
1194 dataset_type: DatasetType,
1195 collection: CollectionRecord,
1196 timespan: Timespan,
1197 *,
1198 data_ids: Iterable[DataCoordinate] | None = None,
1199 query_func: QueryFactoryFunction,
1200 ) -> None:
1201 # Docstring inherited from DatasetRecordStorageManager.
1202 if (storage := self._find_storage(dataset_type.name)) is None: 1202 ↛ 1203line 1202 didn't jump to line 1203 because the condition on line 1202 was never true
1203 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1204 if not dataset_type.isCalibration(): 1204 ↛ 1205line 1204 didn't jump to line 1205 because the condition on line 1204 was never true
1205 raise DatasetTypeError(
1206 f"Cannot certify datasets of type {dataset_type.name!r}, for which "
1207 "DatasetType.isCalibration() is False."
1208 )
1209 if collection.type is not CollectionType.CALIBRATION: 1209 ↛ 1210line 1209 didn't jump to line 1210 because the condition on line 1209 was never true
1210 raise CollectionTypeError(
1211 f"Cannot decertify from collection '{collection.name}' "
1212 f"of type {collection.type.name}; must be CALIBRATION."
1213 )
1214 TimespanReprClass = self._db.getTimespanRepresentation()
1215 data_id_set: set[DataCoordinate] | None
1216 if data_ids is not None:
1217 data_id_set = set(data_ids)
1218 else:
1219 data_id_set = None
1221 # Set up collections to populate with the rows we'll want to modify.
1222 # The insert rows will have the same values for collection and
1223 # dataset type.
1224 proto_insert_row = {
1225 self._collections.getCollectionForeignKeyName(): collection.key,
1226 "dataset_type_id": storage.dataset_type_id,
1227 }
1228 rows_to_delete = []
1229 rows_to_insert = []
1230 # Acquire a table lock to ensure there are no concurrent writes
1231 # between the SELECT and the DELETE and INSERT queries based on it.
1232 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1233 with self._db.transaction(lock=[calibs_table], savepoint=True):
1234 # Find rows overlapping our inputs.
1235 with query_func() as query:
1236 query = query.join_dataset_search(dataset_type, [collection.name])
1237 if data_id_set is not None:
1238 query = query.join_data_coordinates(data_id_set)
1239 timespan_column = query.expression_factory[dataset_type.name].timespan
1240 query = query.where(timespan_column.overlaps(timespan))
1241 result = query.general(
1242 dataset_type.dimensions,
1243 dataset_fields={dataset_type.name: {"dataset_id", "timespan"}},
1244 find_first=False,
1245 )._with_added_dataset_field(dataset_type.name, "calib_pkey")
1247 calib_pkey_key = f"{dataset_type.name}.calib_pkey"
1248 dataset_id_key = f"{dataset_type.name}.dataset_id"
1249 timespan_key = f"{dataset_type.name}.timespan"
1250 for row in result.iter_tuples():
1251 rows_to_delete.append({"id": row.raw_row[calib_pkey_key]})
1252 # Construct the insert row(s) by copying the prototype row,
1253 # then adding the dimension column values, then adding
1254 # what's left of the timespan from that row after we
1255 # subtract the given timespan.
1256 new_insert_row = proto_insert_row.copy()
1257 new_insert_row["dataset_id"] = row.raw_row[dataset_id_key]
1258 for name, value in row.data_id.required.items():
1259 new_insert_row[name] = value
1260 row_timespan = row.raw_row[timespan_key]
1261 assert row_timespan is not None, "Field should have a NOT NULL constraint."
1262 for diff_timespan in row_timespan.difference(timespan):
1263 rows_to_insert.append(
1264 TimespanReprClass.update(diff_timespan, result=new_insert_row.copy())
1265 )
1266 # Run the DELETE and INSERT queries.
1267 self._db.delete(calibs_table, ["id"], *rows_to_delete)
1268 self._db.insert(calibs_table, *rows_to_insert)
1270 def make_joins_builder(
1271 self,
1272 dataset_type: DatasetType,
1273 collections: Sequence[CollectionRecord],
1274 fields: Set[qt.AnyDatasetFieldName],
1275 is_union: bool = False,
1276 ) -> SqlJoinsBuilder:
1277 if (storage := self._find_storage(dataset_type.name)) is None: 1277 ↛ 1278line 1277 didn't jump to line 1278 because the condition on line 1277 was never true
1278 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1279 # This method largely mimics `make_relation`, but it uses the new query
1280 # system primitives instead of the old one. In terms of the SQL
1281 # queries it builds, there are two more main differences:
1282 #
1283 # - Collection and run columns are now string names rather than IDs.
1284 # This insulates the query result-processing code from collection
1285 # caching and the collection manager subclass details.
1286 #
1287 # - The subquery always has unique rows, which is achieved by using
1288 # SELECT DISTINCT when necessary.
1289 #
1290 collection_types = {collection.type for collection in collections}
1291 assert CollectionType.CHAINED not in collection_types, "CHAINED collections must be flattened."
1292 #
1293 # There are two kinds of table in play here:
1294 #
1295 # - the static dataset table (with the dataset ID, dataset type ID,
1296 # run ID/name, and ingest date);
1297 #
1298 # - the dynamic tags/calibs table (with the dataset ID, dataset type
1299 # type ID, collection ID/name, data ID, and possibly validity
1300 # range).
1301 #
1302 # That means that we might want to return a query against either table
1303 # or a JOIN of both, depending on which quantities the caller wants.
1304 # But the data ID is always included, which means we'll always include
1305 # the tags/calibs table and join in the static dataset table only if we
1306 # need things from it that we can't get from the tags/calibs table.
1307 #
1308 # Note that it's important that we include a WHERE constraint on both
1309 # tables for any column (e.g. dataset_type_id) that is in both when
1310 # it's given explicitly; not doing can prevent the query planner from
1311 # using very important indexes. At present, we don't include those
1312 # redundant columns in the JOIN ON expression, however, because the
1313 # FOREIGN KEY (and its index) are defined only on dataset_id.
1314 columns = qt.ColumnSet(dataset_type.dimensions)
1315 columns.drop_implied_dimension_keys()
1316 fields_key: str | qt.AnyDatasetType = qt.ANY_DATASET if is_union else dataset_type.name
1317 columns.dataset_fields[fields_key].update(fields)
1318 tags_builder: SqlSelectBuilder | None = None
1319 if collection_types != {CollectionType.CALIBRATION}:
1320 # We'll need a subquery for the tags table if any of the given
1321 # collections are not a CALIBRATION collection. This intentionally
1322 # also fires when the list of collections is empty as a way to
1323 # create a dummy subquery that we know will fail.
1324 # We give the table an alias because it might appear multiple times
1325 # in the same query, for different dataset types.
1326 tags_table = self._get_tags_table(storage.dynamic_tables).alias(
1327 f"{dataset_type.name}_tags{'_union' if is_union else ''}"
1328 )
1329 tags_builder = self._finish_query_builder(
1330 storage,
1331 SqlJoinsBuilder(db=self._db, from_clause=tags_table).to_select_builder(columns),
1332 [record for record in collections if record.type is not CollectionType.CALIBRATION],
1333 fields,
1334 fields_key,
1335 )
1336 if "timespan" in fields:
1337 tags_builder.joins.timespans[fields_key] = self._db.getTimespanRepresentation().fromLiteral(
1338 Timespan(None, None)
1339 )
1340 assert "calib_pkey" not in fields, (
1341 "Calibration primary key for internal use only on calibration collections."
1342 )
1343 calibs_builder: SqlSelectBuilder | None = None
1344 if CollectionType.CALIBRATION in collection_types:
1345 # If at least one collection is a CALIBRATION collection, we'll
1346 # need a subquery for the calibs table, and could include the
1347 # timespan as a result or constraint.
1348 calibs_table = self._get_calibs_table(storage.dynamic_tables).alias(
1349 f"{dataset_type.name}_calibs{'_union' if is_union else ''}"
1350 )
1351 calibs_builder = self._finish_query_builder(
1352 storage,
1353 SqlJoinsBuilder(db=self._db, from_clause=calibs_table).to_select_builder(columns),
1354 [record for record in collections if record.type is CollectionType.CALIBRATION],
1355 fields,
1356 fields_key,
1357 )
1358 if "timespan" in fields:
1359 calibs_builder.joins.timespans[fields_key] = (
1360 self._db.getTimespanRepresentation().from_columns(calibs_table.columns)
1361 )
1362 if "calib_pkey" in fields:
1363 calibs_builder.joins.fields[fields_key]["calib_pkey"] = calibs_table.columns["id"]
1365 # In calibration collections, we need timespan as well as data ID
1366 # to ensure unique rows.
1367 calibs_builder.distinct = calibs_builder.distinct and "timespan" not in fields
1368 if tags_builder is not None:
1369 if calibs_builder is not None:
1370 # Need a UNION subquery.
1371 return tags_builder.union_subquery([calibs_builder])
1372 else:
1373 return tags_builder.into_joins_builder(postprocessing=None)
1374 elif calibs_builder is not None:
1375 return calibs_builder.into_joins_builder(postprocessing=None)
1376 else:
1377 raise AssertionError("Branch should be unreachable.")
1379 def _finish_query_builder(
1380 self,
1381 storage: _DatasetRecordStorage,
1382 sql_projection: SqlSelectBuilder,
1383 collections: Sequence[CollectionRecord],
1384 fields: Set[qt.AnyDatasetFieldName],
1385 fields_key: str | qt.AnyDatasetType,
1386 ) -> SqlSelectBuilder:
1387 # This method plays the same role as _finish_single_relation in the new
1388 # query system. It is called exactly one or two times by
1389 # make_sql_builder, just as _finish_single_relation is called exactly
1390 # one or two times by make_relation. See make_sql_builder comments for
1391 # what's different.
1392 assert sql_projection.joins.from_clause is not None
1393 run_collections_only = all(record.type is CollectionType.RUN for record in collections)
1394 sql_projection.joins.where(
1395 sql_projection.joins.from_clause.c.dataset_type_id == storage.dataset_type_id
1396 )
1397 dataset_id_col = sql_projection.joins.from_clause.c.dataset_id
1398 collection_col = sql_projection.joins.from_clause.c[self._collections.getCollectionForeignKeyName()]
1399 fields_provided = sql_projection.joins.fields[fields_key]
1400 # We always constrain and optionally retrieve the collection(s) via the
1401 # tags/calibs table.
1402 if "collection_key" in fields:
1403 sql_projection.joins.fields[fields_key]["collection_key"] = collection_col
1404 if len(collections) == 1:
1405 only_collection_record = collections[0]
1406 sql_projection.joins.where(collection_col == only_collection_record.key)
1407 if "collection" in fields:
1408 fields_provided["collection"] = sqlalchemy.literal(only_collection_record.name).cast(
1409 # This cast is necessary to ensure that Postgres knows the
1410 # type of this column if it is used in an aggregate
1411 # function.
1412 sqlalchemy.String
1413 )
1415 elif not collections:
1416 sql_projection.joins.where(sqlalchemy.literal(False))
1417 if "collection" in fields:
1418 fields_provided["collection"] = sqlalchemy.literal("NO COLLECTIONS")
1419 else:
1420 sql_projection.joins.where(
1421 self._db.make_in_array_constraint(
1422 collection_col, [collection.key for collection in collections]
1423 )
1424 )
1425 if "collection" in fields:
1426 # Avoid a join to the collection table to get the name by using
1427 # a CASE statement. The SQL will be a bit more verbose but
1428 # more efficient.
1429 fields_provided["collection"] = _create_case_expression_for_collections(
1430 collections, collection_col
1431 )
1432 # Add more column definitions, starting with the data ID.
1433 sql_projection.joins.extract_dimensions(storage.dataset_type.dimensions.required)
1434 # We can always get the dataset_id from the tags/calibs table, even if
1435 # could also get it from the 'static' dataset table.
1436 if "dataset_id" in fields:
1437 fields_provided["dataset_id"] = dataset_id_col
1439 # It's possible we now have everything we need, from just the
1440 # tags/calibs table. The things we might need to get from the static
1441 # dataset table are the run key and the ingest date.
1442 need_static_table = False
1443 need_collection_table = False
1444 # Ingest date can only come from the static table.
1445 if "ingest_date" in fields:
1446 fields_provided["ingest_date"] = self._static.dataset.c.ingest_date
1447 need_static_table = True
1448 if "run" in fields:
1449 if len(collections) == 1 and run_collections_only:
1450 # If we are searching exactly one RUN collection, we
1451 # know that if we find the dataset in that collection,
1452 # then that's the datasets's run; we don't need to
1453 # query for it.
1454 #
1455 fields_provided["run"] = sqlalchemy.literal(only_collection_record.name).cast(
1456 # This cast is necessary to ensure that Postgres knows the
1457 # type of this column if it is used in an aggregate
1458 # function.
1459 sqlalchemy.String
1460 )
1461 elif run_collections_only:
1462 # Once again we can avoid joining to the collection table by
1463 # adding a CASE statement.
1464 fields_provided["run"] = _create_case_expression_for_collections(
1465 collections, self._static.dataset.c[self._run_key_column]
1466 )
1467 need_static_table = True
1468 else:
1469 # Here we can't avoid a join to the collection table, because
1470 # we might find a dataset via something other than its RUN
1471 # collection.
1472 #
1473 # We have to defer adding the join until after we have joined
1474 # in the static dataset table, because the ON clause involves
1475 # the run collection from the static dataset table. Postgres
1476 # cares about the join ordering (though SQLite does not.)
1477 need_collection_table = True
1478 need_static_table = True
1479 if need_static_table:
1480 # If we need the static table, join it in via dataset_id. We don't
1481 # use SqlJoinsBuilder.join because we're joining on dataset ID, not
1482 # dimensions.
1483 sql_projection.joins.from_clause = sql_projection.joins.from_clause.join(
1484 self._static.dataset, onclause=(dataset_id_col == self._static.dataset.c.id)
1485 )
1486 # Also constrain dataset_type_id in static table in case that helps
1487 # generate a better plan. We could also include this in the JOIN ON
1488 # clause, but my guess is that that's a good idea IFF it's in the
1489 # foreign key, and right now it isn't.
1490 sql_projection.joins.where(self._static.dataset.c.dataset_type_id == storage.dataset_type_id)
1491 if need_collection_table:
1492 # Join the collection table to look up the RUN collection name
1493 # associated with the dataset.
1494 (
1495 fields_provided["run"],
1496 sql_projection.joins.from_clause,
1497 ) = self._collections.lookup_name_sql(
1498 self._static.dataset.c[self._run_key_column],
1499 sql_projection.joins.from_clause,
1500 )
1502 sql_projection.distinct = (
1503 # If there are multiple collections, this subquery might have
1504 # non-unique rows.
1505 len(collections) > 1 and not fields
1506 )
1507 return sql_projection
1509 def refresh_collection_summaries(self, dataset_type: DatasetType) -> None:
1510 # Docstring inherited.
1511 if (storage := self._find_storage(dataset_type.name)) is None: 1511 ↛ 1512line 1511 didn't jump to line 1512 because the condition on line 1511 was never true
1512 raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
1513 with self._db.transaction():
1514 # The main issue here is consistency in the presence of concurrent
1515 # updates (using default READ COMMITTED isolation). Regular clients
1516 # only add to summary tables, and we want to avoid deleting what
1517 # other concurrent transactions may add while we are in this
1518 # transaction. This ordering of operations should guarantee it:
1519 # - read collections for this dataset type from summary tables,
1520 # - read collections for this dataset type from dataset tables
1521 # (both tags and calibs),
1522 # - whatever is in the first set but not in the second can be
1523 # dropped from summary tables.
1524 summary_collection_ids = set(self._summaries.get_collection_ids(storage.dataset_type_id))
1526 # Query datasets tables for associated collections.
1527 column_name = self._collections.getCollectionForeignKeyName()
1528 tags_table = self._get_tags_table(storage.dynamic_tables)
1529 query: sqlalchemy.sql.expression.SelectBase = (
1530 sqlalchemy.select(tags_table.columns[column_name])
1531 .where(tags_table.columns.dataset_type_id == storage.dataset_type_id)
1532 .distinct()
1533 )
1534 if dataset_type.isCalibration():
1535 calibs_table = self._get_calibs_table(storage.dynamic_tables)
1536 query2 = (
1537 sqlalchemy.select(calibs_table.columns[column_name])
1538 .where(calibs_table.columns.dataset_type_id == storage.dataset_type_id)
1539 .distinct()
1540 )
1541 query = sqlalchemy.sql.expression.union(query, query2)
1543 with self._db.query(query) as result:
1544 collection_ids = set(result.scalars())
1546 collections_to_delete = summary_collection_ids - collection_ids
1547 self._summaries.delete_collections(storage.dataset_type_id, collections_to_delete)
1549 def _get_tags_table(self, table: DynamicTables) -> sqlalchemy.Table:
1550 return table.tags(self._db, type(self._collections), self._cache.tables)
1552 def _get_calibs_table(self, table: DynamicTables) -> sqlalchemy.Table:
1553 return table.calibs(self._db, type(self._collections), self._cache.tables)
1556def _create_case_expression_for_collections(
1557 collections: Iterable[CollectionRecord], id_column: sqlalchemy.ColumnElement
1558) -> sqlalchemy.ColumnElement:
1559 """Return a SQLAlchemy Case expression that converts collection IDs to
1560 collection names for the given set of collections.
1562 Parameters
1563 ----------
1564 collections : `~collections.abc.Iterable` [ `CollectionRecord` ]
1565 List of collections to include in conversion table. This should be an
1566 exhaustive list of collections that could appear in `id_column`.
1567 id_column : `sqlalchemy.ColumnElement`
1568 The column containing the collection ID that we want to convert to a
1569 collection name.
1570 """
1571 mapping = {record.key: record.name for record in collections}
1572 if not mapping:
1573 # SQLAlchemy does not correctly handle an empty mapping in case() -- it
1574 # crashes when trying to compile the expression with an
1575 # "AttributeError('NoneType' object has no attribute 'dialect_impl')"
1576 # when trying to access the 'type' property of the Case object. If you
1577 # explicitly specify a type via type_coerce it instead generates
1578 # invalid SQL syntax.
1579 #
1580 # We can end up with empty mappings here in certain "doomed query" edge
1581 # cases, e.g. we start with a list of valid collections but they are
1582 # all filtered out by higher-level code on the basis of collection
1583 # summaries.
1584 return sqlalchemy.cast(sqlalchemy.null(), sqlalchemy.String)
1586 return sqlalchemy.case(mapping, value=id_column)
1589def _ensure_dimension_groups_match(dataset_types: Iterable[DatasetType]) -> DimensionGroup:
1590 dimensions = set(dt.dimensions for dt in dataset_types)
1591 assert len(dimensions) > 0, "At least one dataset type is required"
1592 if len(dimensions) != 1: 1592 ↛ 1593line 1592 didn't jump to line 1593 because the condition on line 1592 was never true
1593 raise DatasetTypeError(
1594 "Dataset types have more than one dimension group.\n"
1595 f"Dataset types: {dataset_types}\n"
1596 f"Dimension groups: {dimensions}"
1597 )
1598 return dimensions.pop()