Coverage for python/lsst/daf/butler/registry/dimensions/spatial.py : 86%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
21from __future__ import annotations
23__all__ = ["SpatialDimensionRecordStorage"]
25from typing import List, Optional
27import sqlalchemy
29from ...core import DimensionElement, DimensionRecord, Timespan, ddl
30from ...core.utils import NamedKeyDict, NamedValueSet
31from ...core.dimensions.schema import makeElementTableSpec, REGION_FIELD_SPEC, addDimensionForeignKey
32from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext
33from ..queries import QueryBuilder
34from .table import TableDimensionRecordStorage
37_OVERLAP_TABLE_NAME_PATTERN = "{0}_{1}_overlap"
40def _makeOverlapTableSpec(a: DimensionElement, b: DimensionElement) -> ddl.TableSpec:
41 """Create a specification for a table that represents a many-to-many
42 relationship between two `DimensionElement` tables.
44 Parameters
45 ----------
46 a : `DimensionElement`
47 First element in the relationship.
48 b : `DimensionElement`
49 Second element in the relationship.
51 Returns
52 -------
53 spec : `TableSpec`
54 Database-agnostic specification for a table.
55 """
56 tableSpec = ddl.TableSpec(
57 fields=NamedValueSet(),
58 unique=set(),
59 foreignKeys=[],
60 )
61 for dimension in a.required:
62 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
63 for dimension in b.required:
64 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
65 return tableSpec
68class SpatialDimensionRecordStorage(TableDimensionRecordStorage):
69 """A record storage implementation for spatial dimension elements that uses
70 a regular database table.
72 Parameters
73 ----------
74 db : `Database`
75 Interface to the database engine and namespace that will hold these
76 dimension records.
77 element : `DimensionElement`
78 The element whose records this storage will manage.
79 table : `sqlalchemy.schema.Table`
80 The logical table for the element.
81 commonSkyPixOverlapTable : `sqlalchemy.schema.Table`, optional
82 The logical table for the overlap table with the dimension universe's
83 common skypix dimension.
84 """
85 def __init__(self, db: Database, element: DimensionElement, *, table: sqlalchemy.schema.Table,
86 commonSkyPixOverlapTable: sqlalchemy.schema.Table):
87 super().__init__(db, element, table=table)
88 self._commonSkyPixOverlapTable = commonSkyPixOverlapTable
89 assert element.spatial is not None
91 @classmethod
92 def initialize(cls, db: Database, element: DimensionElement, *,
93 context: Optional[StaticTablesContext] = None) -> DimensionRecordStorage:
94 # Docstring inherited from DimensionRecordStorage.
95 if context is not None: 95 ↛ 98line 95 didn't jump to line 98, because the condition on line 95 was never false
96 method = context.addTable
97 else:
98 method = db.ensureTableExists
99 return cls(
100 db,
101 element,
102 table=method(element.name, makeElementTableSpec(element)),
103 commonSkyPixOverlapTable=method(
104 _OVERLAP_TABLE_NAME_PATTERN.format(element.name, element.universe.commonSkyPix.name),
105 _makeOverlapTableSpec(element, element.universe.commonSkyPix)
106 )
107 )
109 def join(
110 self,
111 builder: QueryBuilder, *,
112 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
113 timespans: Optional[NamedKeyDict[DimensionElement, Timespan[sqlalchemy.sql.ColumnElement]]] = None,
114 ):
115 # Docstring inherited from DimensionRecordStorage.
116 if regions is not None:
117 dimensions = NamedValueSet(self.element.required)
118 dimensions.add(self.element.universe.universe.commonSkyPix)
119 builder.joinTable(self._commonSkyPixOverlapTable, dimensions)
120 regions[self.element] = self._table.columns[REGION_FIELD_SPEC.name]
121 return super().join(builder, regions=None, timespans=timespans)
123 def _computeCommonSkyPixRows(self, *records: DimensionRecord) -> List[dict]:
124 commonSkyPixRows = []
125 commonSkyPix = self.element.universe.commonSkyPix
126 for record in records:
127 if record.region is None:
128 # TODO: should we warn about this case?
129 continue
130 base = record.dataId.byName()
131 for begin, end in commonSkyPix.pixelization.envelope(record.region):
132 for skypix in range(begin, end):
133 row = base.copy()
134 row[commonSkyPix.name] = skypix
135 commonSkyPixRows.append(row)
136 return commonSkyPixRows
138 def insert(self, *records: DimensionRecord):
139 # Docstring inherited from DimensionRecordStorage.insert.
140 commonSkyPixRows = self._computeCommonSkyPixRows(*records)
141 with self._db.transaction():
142 super().insert(*records)
143 if commonSkyPixRows:
144 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
146 def sync(self, record: DimensionRecord) -> bool:
147 # Docstring inherited from DimensionRecordStorage.sync.
148 inserted = super().sync(record)
149 if inserted:
150 try:
151 commonSkyPixRows = self._computeCommonSkyPixRows(record)
152 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
153 except Exception as err:
154 # EEK. We've just failed to insert the overlap table rows
155 # after succesfully inserting the main dimension element table
156 # row, which means the database is now in a slightly
157 # inconsistent state.
158 # Note that we can't use transactions to solve this, because
159 # Database.sync needs to begin and commit its own transation;
160 # see also DM-24355.
161 raise RuntimeError(
162 f"Failed to add overlap records for {self.element} after "
163 f"successfully inserting the main row. This means the "
164 f"database is in an inconsistent state; please manually "
165 f"remove the row corresponding to data ID "
166 f"{record.dataId.byName()}."
167 ) from err
168 return inserted