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