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

36 NamedValueSet, 

37 REGION_FIELD_SPEC, 

38 Timespan, 

39) 

40from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext 

41from ..queries import QueryBuilder 

42from .table import TableDimensionRecordStorage 

43 

44 

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

46 

47 

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. 

51 

52 Parameters 

53 ---------- 

54 a : `DimensionElement` 

55 First element in the relationship. 

56 b : `DimensionElement` 

57 Second element in the relationship. 

58 

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 

74 

75 

76class SpatialDimensionRecordStorage(TableDimensionRecordStorage): 

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

78 a regular database table. 

79 

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 

98 

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 ) 

116 

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) 

130 

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 

145 

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) 

153 

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