Coverage for python/lsst/daf/butler/registry/datasets/byDimensions/tables.py: 89%
114 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-07 01:48 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "StaticDatasetTablesTuple",
32 "addDatasetForeignKey",
33 "makeCalibTableName",
34 "makeCalibTableSpec",
35 "makeStaticTableSpecs",
36 "makeTagTableName",
37 "makeTagTableSpec",
38)
40from collections import namedtuple
41from typing import Any, TypeAlias
43import sqlalchemy
45from lsst.utils.classes import immutable
47from .... import ddl
48from ...._utilities.thread_safe_cache import ThreadSafeCache
49from ....dimensions import DimensionGroup, DimensionUniverse, GovernorDimension, addDimensionForeignKey
50from ....timespan_database_representation import TimespanDatabaseRepresentation
51from ...interfaces import CollectionManager, Database, VersionTuple
53DATASET_TYPE_NAME_LENGTH = 128
56class MissingDatabaseTableError(RuntimeError):
57 """Exception raised when a table is not found in a database."""
60StaticDatasetTablesTuple = namedtuple(
61 "StaticDatasetTablesTuple",
62 [
63 "dataset_type",
64 "dataset",
65 ],
66)
69def addDatasetForeignKey(
70 tableSpec: ddl.TableSpec,
71 *,
72 name: str = "dataset",
73 onDelete: str | None = None,
74 constraint: bool = True,
75 **kwargs: Any,
76) -> ddl.FieldSpec:
77 """Add a foreign key column for datasets and (optionally) a constraint to
78 a table.
80 This is an internal interface for the ``byDimensions`` package; external
81 code should use `DatasetRecordStorageManager.addDatasetForeignKey` instead.
83 Parameters
84 ----------
85 tableSpec : `ddl.TableSpec`
86 Specification for the table that should reference the dataset
87 table. Will be modified in place.
88 name : `str`, optional
89 A name to use for the prefix of the new field; the full name is
90 ``{name}_id``.
91 onDelete : `str`, optional
92 One of "CASCADE" or "SET NULL", indicating what should happen to
93 the referencing row if the collection row is deleted. `None`
94 indicates that this should be an integrity error.
95 constraint : `bool`, optional
96 If `False` (`True` is default), add a field that can be joined to
97 the dataset primary key, but do not add a foreign key constraint.
98 **kwargs
99 Additional keyword arguments are forwarded to the `ddl.FieldSpec`
100 constructor (only the ``name`` and ``dtype`` arguments are
101 otherwise provided).
103 Returns
104 -------
105 idSpec : `ddl.FieldSpec`
106 Specification for the ID field.
107 """
108 idFieldSpec = ddl.FieldSpec(f"{name}_id", dtype=ddl.GUID, **kwargs)
109 tableSpec.fields.add(idFieldSpec)
110 if constraint:
111 tableSpec.foreignKeys.append(
112 ddl.ForeignKeySpec("dataset", source=(idFieldSpec.name,), target=("id",), onDelete=onDelete)
113 )
114 return idFieldSpec
117def makeStaticTableSpecs(
118 collections: type[CollectionManager],
119 universe: DimensionUniverse,
120 schema_version: VersionTuple,
121) -> StaticDatasetTablesTuple:
122 """Construct all static tables used by the classes in this package.
124 Static tables are those that are present in all Registries and do not
125 depend on what DatasetTypes have been registered.
127 Parameters
128 ----------
129 collections : `CollectionManager`
130 Manager object for the collections in this `Registry`.
131 universe : `DimensionUniverse`
132 Universe graph containing all dimensions known to this `Registry`.
133 schema_version : `VersionTuple`
134 The version of this schema.
136 Returns
137 -------
138 specs : `StaticDatasetTablesTuple`
139 A named tuple containing `ddl.TableSpec` instances.
140 """
141 ingest_date_type: type
142 ingest_date_default: Any = None
143 if schema_version.major > 1: 143 ↛ 146line 143 didn't jump to line 146 because the condition on line 143 was always true
144 ingest_date_type = ddl.AstropyTimeNsecTai
145 else:
146 ingest_date_type = sqlalchemy.TIMESTAMP
147 # New code provides explicit values for ingest_data, but we keep
148 # default just to be consistent with the existing schema.
149 ingest_date_default = sqlalchemy.sql.func.now()
151 specs = StaticDatasetTablesTuple(
152 dataset_type=ddl.TableSpec(
153 fields=[
154 ddl.FieldSpec(
155 name="id",
156 dtype=sqlalchemy.BigInteger,
157 autoincrement=True,
158 primaryKey=True,
159 doc=(
160 "Autoincrement ID that uniquely identifies a dataset "
161 "type in other tables. Python code outside the "
162 "`Registry` class should never interact with this; "
163 "its existence is considered an implementation detail."
164 ),
165 ),
166 ddl.FieldSpec(
167 name="name",
168 dtype=sqlalchemy.String,
169 length=DATASET_TYPE_NAME_LENGTH,
170 nullable=False,
171 doc="String name that uniquely identifies a dataset type.",
172 ),
173 ddl.FieldSpec(
174 name="storage_class",
175 dtype=sqlalchemy.String,
176 length=64,
177 nullable=False,
178 doc=(
179 "Name of the storage class associated with all "
180 "datasets of this type. Storage classes are "
181 "generally associated with a Python class, and are "
182 "enumerated in butler configuration."
183 ),
184 ),
185 ddl.FieldSpec(
186 name="dimensions_key",
187 dtype=sqlalchemy.BigInteger,
188 nullable=False,
189 doc="Unique key for the set of dimensions that identifies datasets of this type.",
190 ),
191 ddl.FieldSpec(
192 name="tag_association_table",
193 dtype=sqlalchemy.String,
194 length=128,
195 nullable=False,
196 doc=(
197 "Name of the table that holds associations between "
198 "datasets of this type and most types of collections."
199 ),
200 ),
201 ddl.FieldSpec(
202 name="calibration_association_table",
203 dtype=sqlalchemy.String,
204 length=128,
205 nullable=True,
206 doc=(
207 "Name of the table that holds associations between "
208 "datasets of this type and CALIBRATION collections. "
209 "NULL values indicate dataset types with "
210 "isCalibration=False."
211 ),
212 ),
213 ],
214 unique=[("name",)],
215 ),
216 dataset=ddl.TableSpec(
217 fields=[
218 ddl.FieldSpec(
219 name="id",
220 dtype=ddl.GUID,
221 primaryKey=True,
222 doc="A unique field used as the primary key for dataset.",
223 ),
224 ddl.FieldSpec(
225 name="dataset_type_id",
226 dtype=sqlalchemy.BigInteger,
227 nullable=False,
228 doc="Reference to the associated entry in the dataset_type table.",
229 ),
230 ddl.FieldSpec(
231 name="ingest_date",
232 dtype=ingest_date_type,
233 default=ingest_date_default,
234 nullable=False,
235 doc="Time of dataset ingestion.",
236 ),
237 # Foreign key field/constraint to run added below.
238 ],
239 foreignKeys=[
240 ddl.ForeignKeySpec("dataset_type", source=("dataset_type_id",), target=("id",)),
241 ],
242 ),
243 )
244 # Add foreign key fields programmatically.
245 collections.addRunForeignKey(specs.dataset, onDelete="CASCADE", nullable=False)
246 return specs
249def makeTagTableName(dimensionsKey: int) -> str:
250 """Construct the name for a dynamic (DatasetType-dependent) tag table used
251 by the classes in this package.
253 Parameters
254 ----------
255 dimensionsKey : `int`
256 Integer key used to save ``datasetType.dimensions`` to the database.
258 Returns
259 -------
260 name : `str`
261 Name for the table.
262 """
263 return f"dataset_tags_{dimensionsKey:08d}"
266def makeCalibTableName(dimensionsKey: int) -> str:
267 """Construct the name for a dynamic (DatasetType-dependent) tag + validity
268 range table used by the classes in this package.
270 Parameters
271 ----------
272 dimensionsKey : `int`
273 Integer key used to save ``datasetType.dimensions`` to the database.
275 Returns
276 -------
277 name : `str`
278 Name for the table.
279 """
280 return f"dataset_calibs_{dimensionsKey:08d}"
283def makeTagTableSpec(
284 dimensions: DimensionGroup, collections: type[CollectionManager], *, constraints: bool = True
285) -> ddl.TableSpec:
286 """Construct the specification for a dynamic (DatasetType-dependent) tag
287 table used by the classes in this package.
289 Parameters
290 ----------
291 dimensions : `DimensionGroup`
292 Dimensions of the dataset type.
293 collections : `type` [ `CollectionManager` ]
294 `CollectionManager` subclass that can be used to construct foreign keys
295 to the run and/or collection tables.
296 constraints : `bool`, optional
297 If `False` (`True` is default), do not define foreign key constraints.
299 Returns
300 -------
301 spec : `ddl.TableSpec`
302 Specification for the table.
303 """
304 tableSpec = ddl.TableSpec(
305 fields=[
306 # Foreign key fields to dataset, collection, and usually dimension
307 # tables added below.
308 # The dataset_type_id field here would be redundant with the one
309 # in the main monolithic dataset table, but we need it here for an
310 # important unique constraint.
311 ddl.FieldSpec("dataset_type_id", dtype=sqlalchemy.BigInteger, nullable=False),
312 ]
313 )
314 if constraints:
315 tableSpec.foreignKeys.append(
316 ddl.ForeignKeySpec("dataset_type", source=("dataset_type_id",), target=("id",))
317 )
318 # We'll also have a unique constraint on dataset type, collection, and data
319 # ID. We only include the required part of the data ID, as that's
320 # sufficient and saves us from worrying about nulls in the constraint.
321 constraint = ["dataset_type_id"]
322 # Add foreign key fields to dataset table (part of the primary key)
323 addDatasetForeignKey(tableSpec, primaryKey=True, onDelete="CASCADE", constraint=constraints)
324 # Add foreign key fields to collection table (part of the primary key and
325 # the data ID unique constraint).
326 collectionFieldSpec = collections.addCollectionForeignKey(
327 tableSpec, primaryKey=True, onDelete="CASCADE", constraint=constraints
328 )
329 constraint.append(collectionFieldSpec.name)
330 # Add foreign key constraint to the collection_summary_dataset_type table.
331 if constraints:
332 tableSpec.foreignKeys.append(
333 ddl.ForeignKeySpec(
334 "collection_summary_dataset_type",
335 source=(collectionFieldSpec.name, "dataset_type_id"),
336 target=(collectionFieldSpec.name, "dataset_type_id"),
337 )
338 )
339 for dimension_name in dimensions.required:
340 dimension = dimensions.universe.dimensions[dimension_name]
341 fieldSpec = addDimensionForeignKey(
342 tableSpec, dimension=dimension, nullable=False, primaryKey=False, constraint=constraints
343 )
344 constraint.append(fieldSpec.name)
345 # If this is a governor dimension, add a foreign key constraint to the
346 # collection_summary_<dimension> table.
347 if isinstance(dimension, GovernorDimension) and constraints:
348 tableSpec.foreignKeys.append(
349 ddl.ForeignKeySpec(
350 f"collection_summary_{dimension.name}",
351 source=(collectionFieldSpec.name, fieldSpec.name),
352 target=(collectionFieldSpec.name, fieldSpec.name),
353 )
354 )
355 # Actually add the unique constraint.
356 tableSpec.unique.add(tuple(constraint))
357 return tableSpec
360def makeCalibTableSpec(
361 dimensions: DimensionGroup,
362 collections: type[CollectionManager],
363 TimespanReprClass: type[TimespanDatabaseRepresentation],
364) -> ddl.TableSpec:
365 """Construct the specification for a dynamic (DatasetType-dependent) tag +
366 validity range table used by the classes in this package.
368 Parameters
369 ----------
370 dimensions : `DimensionGroup`
371 Dimensions of the dataset type.
372 collections : `type` [ `CollectionManager` ]
373 `CollectionManager` subclass that can be used to construct foreign keys
374 to the run and/or collection tables.
375 TimespanReprClass : `type` of `TimespanDatabaseRepresentation`
376 The Python type to use to represent a timespan.
378 Returns
379 -------
380 spec : `ddl.TableSpec`
381 Specification for the table.
382 """
383 tableSpec = ddl.TableSpec(
384 fields=[
385 # This table has no natural primary key, compound or otherwise, so
386 # we add an autoincrement key. We may use this field a bit
387 # internally, but its presence is an implementation detail and it
388 # shouldn't appear as a foreign key in any other tables.
389 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, autoincrement=True, primaryKey=True),
390 # Foreign key fields to dataset, collection, and usually dimension
391 # tables added below. The dataset_type_id field here is redundant
392 # with the one in the main monolithic dataset table, but this bit
393 # of denormalization lets us define what should be a much more
394 # useful index.
395 ddl.FieldSpec("dataset_type_id", dtype=sqlalchemy.BigInteger, nullable=False),
396 ],
397 foreignKeys=[
398 ddl.ForeignKeySpec("dataset_type", source=("dataset_type_id",), target=("id",)),
399 ],
400 )
401 # Record fields that should go in the temporal lookup index/constraint,
402 # starting with the dataset type.
403 index: list[str | type[TimespanDatabaseRepresentation]] = ["dataset_type_id"]
404 # Add foreign key fields to dataset table (not part of the temporal
405 # lookup/constraint).
406 addDatasetForeignKey(tableSpec, nullable=False, onDelete="CASCADE")
407 # Add foreign key fields to collection table (part of the temporal lookup
408 # index/constraint).
409 collectionFieldSpec = collections.addCollectionForeignKey(tableSpec, nullable=False, onDelete="CASCADE")
410 index.append(collectionFieldSpec.name)
411 # Add foreign key constraint to the collection_summary_dataset_type table.
412 tableSpec.foreignKeys.append(
413 ddl.ForeignKeySpec(
414 "collection_summary_dataset_type",
415 source=(collectionFieldSpec.name, "dataset_type_id"),
416 target=(collectionFieldSpec.name, "dataset_type_id"),
417 )
418 )
419 # Add dimension fields (part of the temporal lookup index.constraint).
420 for dimension_name in dimensions.required:
421 dimension = dimensions.universe.dimensions[dimension_name]
422 fieldSpec = addDimensionForeignKey(tableSpec, dimension=dimension, nullable=False, primaryKey=False)
423 index.append(fieldSpec.name)
424 # If this is a governor dimension, add a foreign key constraint to the
425 # collection_summary_<dimension> table.
426 if isinstance(dimension, GovernorDimension):
427 tableSpec.foreignKeys.append(
428 ddl.ForeignKeySpec(
429 f"collection_summary_{dimension.name}",
430 source=(collectionFieldSpec.name, fieldSpec.name),
431 target=(collectionFieldSpec.name, fieldSpec.name),
432 )
433 )
434 # Add validity-range field(s) (part of the temporal lookup
435 # index/constraint).
436 tsFieldSpecs = TimespanReprClass.makeFieldSpecs(nullable=False)
437 for fieldSpec in tsFieldSpecs:
438 tableSpec.fields.add(fieldSpec)
439 if TimespanReprClass.hasExclusionConstraint(): 439 ↛ 444line 439 didn't jump to line 444 because the condition on line 439 was never true
440 # This database's timespan representation can define a database-level
441 # constraint that prevents overlapping validity ranges for entries with
442 # the same DatasetType, collection, and data ID.
443 # This also creates an index.
444 index.append(TimespanReprClass)
445 tableSpec.exclusion.add(tuple(index))
446 else:
447 # No database-level constraint possible. We'll have to simulate that
448 # in our DatasetRecordStorage.certify() implementation, and just create
449 # a regular index here in the hope that helps with lookups.
450 index.extend(fieldSpec.name for fieldSpec in tsFieldSpecs)
451 tableSpec.indexes.add(ddl.IndexSpec(*index)) # type: ignore
452 return tableSpec
455TableCache: TypeAlias = ThreadSafeCache[str, sqlalchemy.Table]
458@immutable
459class DynamicTables:
460 """A struct that holds the "dynamic" tables common to dataset types that
461 share the same dimensions.
463 Objects of this class may be shared between multiple threads, so it must be
464 immutable to prevent concurrency issues.
466 Parameters
467 ----------
468 dimensions : `DimensionGroup`
469 Dimensions of the dataset types that use these tables.
470 dimensions_key : `int`
471 Integer key used to persist this dimension group in the database and
472 name the associated tables.
473 tags_name : `str`
474 Name of the "tags" table that associates datasets with data IDs in
475 RUN and TAGGED collections.
476 calibs_name : `str` or `None`
477 Name of the "calibs" table that associates datasets with data IDs and
478 timespans in CALIBRATION collections. This is `None` if none of the
479 dataset types (or at least none of those seen by this client) are
480 calibrations.
481 """
483 def __init__(
484 self, dimensions: DimensionGroup, dimensions_key: int, tags_name: str, calibs_name: str | None
485 ):
486 self._dimensions = dimensions
487 self.dimensions_key = dimensions_key
488 self.tags_name = tags_name
489 self.calibs_name = calibs_name
491 def copy(self, calibs_name: str) -> DynamicTables:
492 return DynamicTables(self._dimensions, self.dimensions_key, self.tags_name, calibs_name)
494 @classmethod
495 def from_dimensions_key(
496 cls, dimensions: DimensionGroup, dimensions_key: int, is_calibration: bool
497 ) -> DynamicTables:
498 """Construct with table names generated from the dimension key.
500 Parameters
501 ----------
502 dimensions : `DimensionGroup`
503 Dimensions of the dataset types that use these tables.
504 dimensions_key : `int`
505 Integer key used to persist this dimension group in the database
506 and name the associated tables.
507 is_calibration : `bool`
508 Whether any of the dataset types that use these tables are
509 calibrations.
511 Returns
512 -------
513 dynamic_tables : `DynamicTables`
514 Struct that holds tables for a group of dataset types.
515 """
516 return cls(
517 dimensions,
518 dimensions_key=dimensions_key,
519 tags_name=makeTagTableName(dimensions_key),
520 calibs_name=makeCalibTableName(dimensions_key) if is_calibration else None,
521 )
523 def create(self, db: Database, collections: type[CollectionManager], cache: TableCache) -> None:
524 """Create the tables if they don't already exist.
526 Parameters
527 ----------
528 db : `Database`
529 Database interface.
530 collections : `type` [ `CollectionManager` ]
531 Manager class for collections; used to create foreign key columns
532 for collections.
533 cache : `DynamicTablesCache`
534 Cache used to store sqlalchemy Table objects.
535 """
536 if cache.get(self.tags_name) is None:
537 cache.set_or_get(
538 self.tags_name,
539 db.ensureTableExists(
540 self.tags_name,
541 makeTagTableSpec(self._dimensions, collections),
542 ),
543 )
545 if self.calibs_name is not None and cache.get(self.calibs_name) is None:
546 cache.set_or_get(
547 self.calibs_name,
548 db.ensureTableExists(
549 self.calibs_name,
550 makeCalibTableSpec(self._dimensions, collections, db.getTimespanRepresentation()),
551 ),
552 )
554 def add_calibs(
555 self, db: Database, collections: type[CollectionManager], cache: TableCache
556 ) -> DynamicTables:
557 """Create a calibs table for a dataset type whose dimensions already
558 have a tags table.
560 Parameters
561 ----------
562 db : `Database`
563 Database interface.
564 collections : `type` [ `CollectionManager` ]
565 Manager class for collections; used to create foreign key columns
566 for collections.
567 cache : `DynamicTablesCache`
568 Cache used to store sqlalchemy Table objects.
569 """
570 calibs_name = makeCalibTableName(self.dimensions_key)
571 cache.set_or_get(
572 calibs_name,
573 db.ensureTableExists(
574 calibs_name,
575 makeCalibTableSpec(self._dimensions, collections, db.getTimespanRepresentation()),
576 ),
577 )
579 return self.copy(calibs_name=calibs_name)
581 def tags(self, db: Database, collections: type[CollectionManager], cache: TableCache) -> sqlalchemy.Table:
582 """Return the "tags" table that associates datasets with data IDs in
583 TAGGED and RUN collections.
585 This method caches its result the first time it is called (and assumes
586 the arguments it is given never change).
588 Parameters
589 ----------
590 db : `Database`
591 Database interface.
592 collections : `type` [ `CollectionManager` ]
593 Manager class for collections; used to create foreign key columns
594 for collections.
595 cache : `DynamicTablesCache`
596 Cache used to store sqlalchemy Table objects.
598 Returns
599 -------
600 table : `sqlalchemy.Table`
601 SQLAlchemy table object.
602 """
603 table = cache.get(self.tags_name)
604 if table is not None:
605 return table
607 spec = makeTagTableSpec(self._dimensions, collections)
608 table = db.getExistingTable(self.tags_name, spec)
609 if table is None: 609 ↛ 610line 609 didn't jump to line 610 because the condition on line 609 was never true
610 raise MissingDatabaseTableError(f"Table {self.tags_name!r} is missing from database schema.")
611 return cache.set_or_get(self.tags_name, table)
613 def calibs(
614 self, db: Database, collections: type[CollectionManager], cache: TableCache
615 ) -> sqlalchemy.Table:
616 """Return the "calibs" table that associates datasets with data IDs and
617 timespans in CALIBRATION collections.
619 This method caches its result the first time it is called (and assumes
620 the arguments it is given never change). It may only be called if the
621 dataset type is calibration.
623 Parameters
624 ----------
625 db : `Database`
626 Database interface.
627 collections : `type` [ `CollectionManager` ]
628 Manager class for collections; used to create foreign key columns
629 for collections.
630 cache : `DynamicTablesCache`
631 Cache used to store sqlalchemy Table objects.
633 Returns
634 -------
635 table : `sqlalchemy.Table`
636 SQLAlchemy table object.
637 """
638 assert self.calibs_name is not None, (
639 "Dataset type should be checked to be calibration by calling code."
640 )
641 table = cache.get(self.calibs_name)
642 if table is not None: 642 ↛ 645line 642 didn't jump to line 645 because the condition on line 642 was always true
643 return table
645 spec = makeCalibTableSpec(self._dimensions, collections, db.getTimespanRepresentation())
646 table = db.getExistingTable(self.calibs_name, spec)
647 if table is None:
648 raise MissingDatabaseTableError(f"Table {self.calibs_name!r} is missing from database schema.")
649 return cache.set_or_get(self.calibs_name, table)