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 (
30 addDimensionForeignKey,
31 ddl,
32 DimensionElement,
33 DimensionRecord,
34 makeDimensionElementTableSpec,
35 NamedKeyDict,
36 NamedValueSet,
37 REGION_FIELD_SPEC,
38 Timespan,
39)
40from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext
41from ..queries import QueryBuilder
42from .table import TableDimensionRecordStorage
45_OVERLAP_TABLE_NAME_PATTERN = "{0}_{1}_overlap"
48def _makeOverlapTableSpec(a: DimensionElement, b: DimensionElement) -> ddl.TableSpec:
49 """Create a specification for a table that represents a many-to-many
50 relationship between two `DimensionElement` tables.
52 Parameters
53 ----------
54 a : `DimensionElement`
55 First element in the relationship.
56 b : `DimensionElement`
57 Second element in the relationship.
59 Returns
60 -------
61 spec : `TableSpec`
62 Database-agnostic specification for a table.
63 """
64 tableSpec = ddl.TableSpec(
65 fields=NamedValueSet(),
66 unique=set(),
67 foreignKeys=[],
68 )
69 for dimension in a.required:
70 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
71 for dimension in b.required:
72 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
73 return tableSpec
76class SpatialDimensionRecordStorage(TableDimensionRecordStorage):
77 """A record storage implementation for spatial dimension elements that uses
78 a regular database table.
80 Parameters
81 ----------
82 db : `Database`
83 Interface to the database engine and namespace that will hold these
84 dimension records.
85 element : `DimensionElement`
86 The element whose records this storage will manage.
87 table : `sqlalchemy.schema.Table`
88 The logical table for the element.
89 commonSkyPixOverlapTable : `sqlalchemy.schema.Table`, optional
90 The logical table for the overlap table with the dimension universe's
91 common skypix dimension.
92 """
93 def __init__(self, db: Database, element: DimensionElement, *, table: sqlalchemy.schema.Table,
94 commonSkyPixOverlapTable: sqlalchemy.schema.Table):
95 super().__init__(db, element, table=table)
96 self._commonSkyPixOverlapTable = commonSkyPixOverlapTable
97 assert element.spatial is not None
99 @classmethod
100 def initialize(cls, db: Database, element: DimensionElement, *,
101 context: Optional[StaticTablesContext] = None) -> DimensionRecordStorage:
102 # Docstring inherited from DimensionRecordStorage.
103 if context is not None: 103 ↛ 106line 103 didn't jump to line 106, because the condition on line 103 was never false
104 method = context.addTable
105 else:
106 method = db.ensureTableExists
107 return cls(
108 db,
109 element,
110 table=method(element.name, makeDimensionElementTableSpec(element)),
111 commonSkyPixOverlapTable=method(
112 _OVERLAP_TABLE_NAME_PATTERN.format(element.name, element.universe.commonSkyPix.name),
113 _makeOverlapTableSpec(element, element.universe.commonSkyPix)
114 )
115 )
117 def join(
118 self,
119 builder: QueryBuilder, *,
120 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
121 timespans: Optional[NamedKeyDict[DimensionElement, Timespan[sqlalchemy.sql.ColumnElement]]] = None,
122 ) -> None:
123 # Docstring inherited from DimensionRecordStorage.
124 if regions is not None:
125 dimensions = NamedValueSet(self.element.required)
126 dimensions.add(self.element.universe.universe.commonSkyPix)
127 builder.joinTable(self._commonSkyPixOverlapTable, dimensions)
128 regions[self.element] = self._table.columns[REGION_FIELD_SPEC.name]
129 return super().join(builder, regions=None, timespans=timespans)
131 def _computeCommonSkyPixRows(self, *records: DimensionRecord) -> List[dict]:
132 commonSkyPixRows = []
133 commonSkyPix = self.element.universe.commonSkyPix
134 for record in records:
135 # MyPy can't tell that some DimensionRecords have regions, because
136 # they're dynamically-created types.
137 region = record.region # type: ignore
138 if region is None:
139 # TODO: should we warn about this case?
140 continue
141 base = record.dataId.byName()
142 for begin, end in commonSkyPix.pixelization.envelope(region):
143 for skypix in range(begin, end):
144 row = base.copy()
145 row[commonSkyPix.name] = skypix
146 commonSkyPixRows.append(row)
147 return commonSkyPixRows
149 def insert(self, *records: DimensionRecord) -> None:
150 # Docstring inherited from DimensionRecordStorage.insert.
151 commonSkyPixRows = self._computeCommonSkyPixRows(*records)
152 with self._db.transaction():
153 super().insert(*records)
154 if commonSkyPixRows:
155 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
157 def sync(self, record: DimensionRecord) -> bool:
158 # Docstring inherited from DimensionRecordStorage.sync.
159 inserted = super().sync(record)
160 if inserted:
161 try:
162 commonSkyPixRows = self._computeCommonSkyPixRows(record)
163 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
164 except Exception as err:
165 # EEK. We've just failed to insert the overlap table rows
166 # after succesfully inserting the main dimension element table
167 # row, which means the database is now in a slightly
168 # inconsistent state.
169 # Note that we can't use transactions to solve this, because
170 # Database.sync needs to begin and commit its own transation;
171 # see also DM-24355.
172 raise RuntimeError(
173 f"Failed to add overlap records for {self.element} after "
174 f"successfully inserting the main row. This means the "
175 f"database is in an inconsistent state; please manually "
176 f"remove the row corresponding to data ID "
177 f"{record.dataId.byName()}."
178 ) from err
179 return inserted