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