Coverage for python/lsst/daf/butler/tests/_dummyRegistry.py: 84%

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27from __future__ import annotations 

28 

29__all__ = ("DummyRegistry",) 

30 

31from collections.abc import Iterable, Iterator 

32from typing import Any 

33 

34from lsst.daf.butler import DimensionUniverse, ddl 

35from lsst.daf.butler.registry import ConflictingDefinitionError 

36from lsst.daf.butler.registry.bridge.ephemeral import EphemeralDatastoreRegistryBridge 

37from lsst.daf.butler.registry.interfaces import ( 

38 Database, 

39 DatabaseInsertMode, 

40 DatasetIdRef, 

41 DatasetRecordStorageManager, 

42 DatastoreRegistryBridge, 

43 DatastoreRegistryBridgeManager, 

44 OpaqueTableStorage, 

45 OpaqueTableStorageManager, 

46 StaticTablesContext, 

47 VersionTuple, 

48) 

49 

50from ..datastore import DatastoreTransaction 

51 

52 

53class DummyOpaqueTableStorage(OpaqueTableStorage): 

54 def __init__(self, name: str, spec: ddl.TableSpec) -> None: 

55 super().__init__(name=name) 

56 self._rows: list[dict] = [] 

57 self._spec = spec 

58 

59 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

60 # Docstring inherited from OpaqueTableStorage. 

61 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.INSERT) 

62 

63 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

64 # Docstring inherited from OpaqueTableStorage. 

65 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.REPLACE) 

66 

67 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

68 # Docstring inherited from OpaqueTableStorage. 

69 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.ENSURE) 

70 

71 def _insert( 

72 self, 

73 *data: dict, 

74 transaction: DatastoreTransaction | None = None, 

75 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT, 

76 ) -> None: 

77 uniqueConstraints = list(self._spec.unique) 

78 uniqueConstraints.append(tuple(field.name for field in self._spec.fields if field.primaryKey)) 

79 for d in data: 

80 skipping = False 

81 for constraint in uniqueConstraints: 

82 matching = list(self.fetch(**{k: d[k] for k in constraint})) 

83 if len(matching) != 0: 

84 match insert_mode: 

85 case DatabaseInsertMode.INSERT: 85 ↛ 86line 85 didn't jump to line 86 because the pattern on line 85 never matched

86 raise ConflictingDefinitionError( 

87 f"Table {self.name} already contains matching record(s): {matching}" 

88 ) 

89 case DatabaseInsertMode.ENSURE: 89 ↛ 91line 89 didn't jump to line 91 because the pattern on line 89 never matched

90 # Row already exists. Skip. 

91 skipping = True 

92 case DatabaseInsertMode.REPLACE: 92 ↛ 96line 92 didn't jump to line 96 because the pattern on line 92 always matched

93 # Should try to put these rows back on transaction 

94 # rollback... 

95 self.delete([], *matching) 

96 case _: 

97 raise ValueError(f"Unrecognized insert mode: {insert_mode}.") 

98 

99 if skipping: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true

100 continue 

101 self._rows.append(d) 

102 if transaction is not None: 

103 transaction.registerUndo("insert", self.delete, [], d) 

104 

105 def fetch(self, **where: Any) -> Iterator[dict]: 

106 # Docstring inherited from OpaqueTableStorage. 

107 where = where.copy() # May need to modify it. 

108 

109 # Can support an IN operator if given list. 

110 wherein = {} 

111 for k in list(where): 

112 if isinstance(where[k], tuple | list | set): 

113 wherein[k] = set(where[k]) 

114 del where[k] 

115 

116 for d in self._rows: 

117 if all(d[k] == v for k, v in where.items()): 

118 if wherein: 

119 match = True 

120 for k, v in wherein.items(): 

121 if d[k] not in v: 

122 match = False 

123 break 

124 if match: 

125 yield d 

126 else: 

127 yield d 

128 

129 def fetch_batches(self, **where: Any) -> Iterator[list[dict]]: 

130 # Docstring inherited from OpaqueTableStorage. 

131 yield list(self.fetch(**where)) 

132 

133 def delete(self, columns: Iterable[str], *rows: dict) -> None: 

134 # Docstring inherited from OpaqueTableStorage. 

135 kept_rows = [] 

136 for table_row in self._rows: 

137 for where_row in rows: 

138 if all(table_row[k] == v for k, v in where_row.items()): 

139 break 

140 else: 

141 kept_rows.append(table_row) 

142 self._rows = kept_rows 

143 

144 

145class DummyOpaqueTableStorageManager(OpaqueTableStorageManager): 

146 def __init__(self, registry_schema_version: VersionTuple | None = None) -> None: 

147 super().__init__(registry_schema_version=registry_schema_version) 

148 self._storages: dict[str, DummyOpaqueTableStorage] = {} 

149 

150 def clone(self, db: Database) -> OpaqueTableStorageManager: 

151 return self 

152 

153 @classmethod 

154 def initialize( 

155 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None 

156 ) -> OpaqueTableStorageManager: 

157 # Docstring inherited from OpaqueTableStorageManager. 

158 # Not used, but needed to satisfy ABC requirement. 

159 return cls(registry_schema_version=registry_schema_version) 

160 

161 def get(self, name: str) -> OpaqueTableStorage | None: 

162 # Docstring inherited from OpaqueTableStorageManager. 

163 return self._storages.get(name) 

164 

165 def register(self, name: str, spec: ddl.TableSpec) -> OpaqueTableStorage: 

166 # Docstring inherited from OpaqueTableStorageManager. 

167 return self._storages.setdefault(name, DummyOpaqueTableStorage(name, spec)) 

168 

169 @classmethod 

170 def currentVersions(cls) -> list[VersionTuple]: 

171 # Docstring inherited from VersionedExtension. 

172 return [] 

173 

174 

175class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager): 

176 def __init__( 

177 self, 

178 opaque: OpaqueTableStorageManager, 

179 universe: DimensionUniverse, 

180 registry_schema_version: VersionTuple | None = None, 

181 ): 

182 super().__init__( 

183 opaque=opaque, 

184 universe=universe, 

185 registry_schema_version=registry_schema_version, 

186 ) 

187 self._bridges: dict[str, EphemeralDatastoreRegistryBridge] = {} 

188 

189 def clone(self, *, db: Database, opaque: OpaqueTableStorageManager) -> DatastoreRegistryBridgeManager: 

190 return DummyDatastoreRegistryBridgeManager( 

191 opaque=opaque, 

192 universe=self.universe, 

193 registry_schema_version=self._registry_schema_version, 

194 ) 

195 

196 @classmethod 

197 def initialize( 

198 cls, 

199 db: Database, 

200 context: StaticTablesContext, 

201 *, 

202 opaque: OpaqueTableStorageManager, 

203 datasets: type[DatasetRecordStorageManager], 

204 universe: DimensionUniverse, 

205 registry_schema_version: VersionTuple | None = None, 

206 ) -> DatastoreRegistryBridgeManager: 

207 # Docstring inherited from DatastoreRegistryBridgeManager 

208 # Not used, but needed to satisfy ABC requirement. 

209 return cls( 

210 opaque=opaque, 

211 universe=universe, 

212 registry_schema_version=registry_schema_version, 

213 ) 

214 

215 def refresh(self) -> None: 

216 # Docstring inherited from DatastoreRegistryBridgeManager 

217 pass 

218 

219 def register(self, name: str, *, ephemeral: bool = False) -> DatastoreRegistryBridge: 

220 # Docstring inherited from DatastoreRegistryBridgeManager 

221 return self._bridges.setdefault(name, EphemeralDatastoreRegistryBridge(name)) 

222 

223 def findDatastores(self, ref: DatasetIdRef) -> Iterable[str]: 

224 # Docstring inherited from DatastoreRegistryBridgeManager 

225 for name, bridge in self._bridges.items(): 

226 if ref in bridge: 

227 yield name 

228 

229 @classmethod 

230 def currentVersions(cls) -> list[VersionTuple]: 

231 # Docstring inherited from VersionedExtension. 

232 return [] 

233 

234 

235class DummyRegistry: 

236 """Dummy Registry, for Datastore test purposes.""" 

237 

238 def __init__(self) -> None: 

239 self._opaque = DummyOpaqueTableStorageManager() 

240 self.dimensions = DimensionUniverse() 

241 self._datastoreBridges = DummyDatastoreRegistryBridgeManager(self._opaque, self.dimensions, None) 

242 

243 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager: 

244 return self._datastoreBridges