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