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 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 

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(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 ) 

115 

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) 

129 

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 

144 

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) 

152 

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