Coverage for python/lsst/daf/butler/registry/dimensions/governor.py: 93%

80 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-26 02:04 -0800

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__ = ["BasicGovernorDimensionRecordStorage"] 

24 

25from collections.abc import Callable, Mapping 

26from typing import Any, cast 

27 

28import sqlalchemy 

29from lsst.daf.relation import Relation 

30 

31from ...core import DataCoordinate, DimensionRecord, GovernorDimension 

32from .. import queries 

33from ..interfaces import Database, GovernorDimensionRecordStorage, StaticTablesContext 

34 

35 

36class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage): 

37 """A record storage implementation for `GovernorDimension` that 

38 aggressively fetches and caches all values from the database. 

39 

40 Parameters 

41 ---------- 

42 db : `Database` 

43 Interface to the database engine and namespace that will hold these 

44 dimension records. 

45 dimension : `GovernorDimension` 

46 The dimension whose records this storage will manage. 

47 table : `sqlalchemy.schema.Table` 

48 The logical table for the dimension. 

49 """ 

50 

51 def __init__( 

52 self, 

53 db: Database, 

54 dimension: GovernorDimension, 

55 table: sqlalchemy.schema.Table, 

56 ): 

57 self._db = db 

58 self._dimension = dimension 

59 self._table = table 

60 # We need to allow the cache to be None so we have some recourse when 

61 # it is cleared as part of transaction rollback - we can't run 

62 # queries to repopulate them at that point, so we need to defer it 

63 # until next use. 

64 self._cache: dict[DataCoordinate, DimensionRecord] | None = None 

65 self._callbacks: list[Callable[[DimensionRecord], None]] = [] 

66 

67 @classmethod 

68 def initialize( 

69 cls, 

70 db: Database, 

71 element: GovernorDimension, 

72 *, 

73 context: StaticTablesContext | None = None, 

74 config: Mapping[str, Any], 

75 ) -> GovernorDimensionRecordStorage: 

76 # Docstring inherited from GovernorDimensionRecordStorage. 

77 spec = element.RecordClass.fields.makeTableSpec( 

78 TimespanReprClass=db.getTimespanRepresentation(), 

79 ) 

80 if context is not None: 80 ↛ 83line 80 didn't jump to line 83, because the condition on line 80 was never false

81 table = context.addTable(element.name, spec) 

82 else: 

83 table = db.ensureTableExists(element.name, spec) 

84 return cls(db, element, table) 

85 

86 @property 

87 def element(self) -> GovernorDimension: 

88 # Docstring inherited from DimensionRecordStorage.element. 

89 return self._dimension 

90 

91 @property 

92 def table(self) -> sqlalchemy.schema.Table: 

93 return self._table 

94 

95 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None: 

96 # Docstring inherited from GovernorDimensionRecordStorage. 

97 self._callbacks.append(callback) 

98 

99 def clearCaches(self) -> None: 

100 # Docstring inherited from DimensionRecordStorage.clearCaches. 

101 self._cache = None 

102 

103 def make_relation(self, context: queries.SqlQueryContext, _sized: bool = True) -> Relation: 

104 # Docstring inherited. 

105 payload = self._build_sql_payload(self._table, context.column_types) 

106 if _sized: 

107 cache = self.get_record_cache(context) 

108 return context.sql_engine.make_leaf( 

109 payload.columns_available.keys(), 

110 name=self.element.name, 

111 payload=payload, 

112 min_rows=len(cache) if _sized else 0, 

113 max_rows=len(cache) if _sized else None, 

114 ) 

115 

116 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None: 

117 # Docstring inherited from DimensionRecordStorage.insert. 

118 elementRows = [record.toDict() for record in records] 

119 with self._db.transaction(): 

120 if replace: 120 ↛ 121line 120 didn't jump to line 121, because the condition on line 120 was never true

121 self._db.replace(self._table, *elementRows) 

122 elif skip_existing: 

123 self._db.ensure(self._table, *elementRows, primary_key_only=True) 

124 else: 

125 self._db.insert(self._table, *elementRows) 

126 for record in records: 

127 # We really shouldn't ever get into a situation where the 

128 # record here differs from the one in the DB, but the last 

129 # thing we want is to make it harder to debug by making the 

130 # cache different from the DB. 

131 if self._cache is not None: 

132 # We really shouldn't ever get into a situation where the 

133 # record here differs from the one in the DB, but the last 

134 # thing we want is to make it harder to debug by making the 

135 # cache different from the DB. 

136 if skip_existing: 

137 self._cache.setdefault(record.dataId, record) 

138 else: 

139 self._cache[record.dataId] = record 

140 for callback in self._callbacks: 

141 callback(record) 

142 

143 def sync(self, record: DimensionRecord, update: bool = False) -> bool | dict[str, Any]: 

144 # Docstring inherited from DimensionRecordStorage.sync. 

145 compared = record.toDict() 

146 keys = {} 

147 for name in record.fields.required.names: 

148 keys[name] = compared.pop(name) 

149 with self._db.transaction(): 

150 _, inserted_or_updated = self._db.sync( 

151 self._table, 

152 keys=keys, 

153 compared=compared, 

154 update=update, 

155 ) 

156 if inserted_or_updated: 156 ↛ 161line 156 didn't jump to line 161, because the condition on line 156 was never false

157 if self._cache is not None: 157 ↛ 158line 157 didn't jump to line 158, because the condition on line 157 was never true

158 self._cache[record.dataId] = record 

159 for callback in self._callbacks: 

160 callback(record) 

161 return inserted_or_updated 

162 

163 def fetch_one(self, data_id: DataCoordinate, context: queries.SqlQueryContext) -> DimensionRecord | None: 

164 # Docstring inherited. 

165 cache = self.get_record_cache(context) 

166 return cache.get(data_id) 

167 

168 def get_record_cache(self, context: queries.SqlQueryContext) -> Mapping[DataCoordinate, DimensionRecord]: 

169 # Docstring inherited. 

170 if self._cache is None: 

171 reader = queries.DimensionRecordReader(self.element) 

172 cache = {} 

173 for row in context.fetch_iterable(self.make_relation(context, _sized=False)): 

174 record = reader.read(row) 

175 cache[record.dataId] = record 

176 self._cache = cache 

177 return cast(Mapping[DataCoordinate, DimensionRecord], self._cache) 

178 

179 def digestTables(self) -> list[sqlalchemy.schema.Table]: 

180 # Docstring inherited from DimensionRecordStorage.digestTables. 

181 return [self._table]