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