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

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 ):
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 if record.region is None:
136 # TODO: should we warn about this case?
137 continue
138 base = record.dataId.byName()
139 for begin, end in commonSkyPix.pixelization.envelope(record.region):
140 for skypix in range(begin, end):
141 row = base.copy()
142 row[commonSkyPix.name] = skypix
143 commonSkyPixRows.append(row)
144 return commonSkyPixRows
146 def insert(self, *records: DimensionRecord):
147 # Docstring inherited from DimensionRecordStorage.insert.
148 commonSkyPixRows = self._computeCommonSkyPixRows(*records)
149 with self._db.transaction():
150 super().insert(*records)
151 if commonSkyPixRows:
152 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
154 def sync(self, record: DimensionRecord) -> bool:
155 # Docstring inherited from DimensionRecordStorage.sync.
156 inserted = super().sync(record)
157 if inserted:
158 try:
159 commonSkyPixRows = self._computeCommonSkyPixRows(record)
160 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
161 except Exception as err:
162 # EEK. We've just failed to insert the overlap table rows
163 # after succesfully inserting the main dimension element table
164 # row, which means the database is now in a slightly
165 # inconsistent state.
166 # Note that we can't use transactions to solve this, because
167 # Database.sync needs to begin and commit its own transation;
168 # see also DM-24355.
169 raise RuntimeError(
170 f"Failed to add overlap records for {self.element} after "
171 f"successfully inserting the main row. This means the "
172 f"database is in an inconsistent state; please manually "
173 f"remove the row corresponding to data ID "
174 f"{record.dataId.byName()}."
175 ) from err
176 return inserted