Hide keyboard shortcuts

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 

22 

23__all__ = ["SpatialDimensionRecordStorage"] 

24 

25from typing import List, Optional 

26 

27import sqlalchemy 

28 

29from ...core import ( 

30 addDimensionForeignKey, 

31 DatabaseTimespanRepresentation, 

32 ddl, 

33 DimensionElement, 

34 DimensionRecord, 

35 NamedKeyDict, 

36 NamedValueSet, 

37 REGION_FIELD_SPEC, 

38) 

39from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext 

40from ..queries import QueryBuilder 

41from .table import TableDimensionRecordStorage 

42 

43 

44_OVERLAP_TABLE_NAME_PATTERN = "{0}_{1}_overlap" 

45 

46 

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. 

50 

51 Parameters 

52 ---------- 

53 a : `DimensionElement` 

54 First element in the relationship. 

55 b : `DimensionElement` 

56 Second element in the relationship. 

57 

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 

73 

74 

75class SpatialDimensionRecordStorage(TableDimensionRecordStorage): 

76 """A record storage implementation for spatial dimension elements that uses 

77 a regular database table. 

78 

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 

97 

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( 

110 element.name, 

111 element.RecordClass.fields.makeTableSpec(tsRepr=db.getTimespanRepresentation()) 

112 ), 

113 commonSkyPixOverlapTable=method( 

114 _OVERLAP_TABLE_NAME_PATTERN.format(element.name, element.universe.commonSkyPix.name), 

115 _makeOverlapTableSpec(element, element.universe.commonSkyPix) 

116 ) 

117 ) 

118 

119 def join( 

120 self, 

121 builder: QueryBuilder, *, 

122 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None, 

123 timespans: Optional[NamedKeyDict[DimensionElement, DatabaseTimespanRepresentation]] = None, 

124 ) -> None: 

125 # Docstring inherited from DimensionRecordStorage. 

126 if regions is not None: 

127 dimensions = NamedValueSet(self.element.required) 

128 dimensions.add(self.element.universe.universe.commonSkyPix) 

129 builder.joinTable(self._commonSkyPixOverlapTable, dimensions) 

130 regions[self.element] = self._table.columns[REGION_FIELD_SPEC.name] 

131 return super().join(builder, regions=None, timespans=timespans) 

132 

133 def _computeCommonSkyPixRows(self, *records: DimensionRecord) -> List[dict]: 

134 commonSkyPixRows = [] 

135 commonSkyPix = self.element.universe.commonSkyPix 

136 for record in records: 

137 # MyPy can't tell that some DimensionRecords have regions, because 

138 # they're dynamically-created types. 

139 region = record.region # type: ignore 

140 if region is None: 

141 # TODO: should we warn about this case? 

142 continue 

143 base = record.dataId.byName() 

144 for begin, end in commonSkyPix.pixelization.envelope(region): 

145 for skypix in range(begin, end): 

146 row = base.copy() 

147 row[commonSkyPix.name] = skypix 

148 commonSkyPixRows.append(row) 

149 return commonSkyPixRows 

150 

151 def insert(self, *records: DimensionRecord) -> None: 

152 # Docstring inherited from DimensionRecordStorage.insert. 

153 commonSkyPixRows = self._computeCommonSkyPixRows(*records) 

154 with self._db.transaction(): 

155 super().insert(*records) 

156 if commonSkyPixRows: 

157 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows) 

158 

159 def sync(self, record: DimensionRecord) -> bool: 

160 # Docstring inherited from DimensionRecordStorage.sync. 

161 inserted = super().sync(record) 

162 if inserted: 

163 try: 

164 commonSkyPixRows = self._computeCommonSkyPixRows(record) 

165 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows) 

166 except Exception as err: 

167 # EEK. We've just failed to insert the overlap table rows 

168 # after succesfully inserting the main dimension element table 

169 # row, which means the database is now in a slightly 

170 # inconsistent state. 

171 # Note that we can't use transactions to solve this, because 

172 # Database.sync needs to begin and commit its own transation; 

173 # see also DM-24355. 

174 raise RuntimeError( 

175 f"Failed to add overlap records for {self.element} after " 

176 f"successfully inserting the main row. This means the " 

177 f"database is in an inconsistent state; please manually " 

178 f"remove the row corresponding to data ID " 

179 f"{record.dataId.byName()}." 

180 ) from err 

181 return inserted