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