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

107 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-25 15:14 +0000

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__ = ("DummyRegistry",) 

24 

25from collections.abc import Iterable, Iterator 

26from typing import Any 

27 

28import sqlalchemy 

29from lsst.daf.butler import DimensionUniverse, ddl 

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

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

32 Database, 

33 DatabaseInsertMode, 

34 DatasetIdRef, 

35 DatasetRecordStorageManager, 

36 DatastoreRegistryBridge, 

37 DatastoreRegistryBridgeManager, 

38 OpaqueTableStorage, 

39 OpaqueTableStorageManager, 

40 StaticTablesContext, 

41 VersionTuple, 

42) 

43 

44from ..core.datastore import DatastoreTransaction 

45 

46 

47class DummyOpaqueTableStorage(OpaqueTableStorage): 

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

49 super().__init__(name=name) 

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

51 self._spec = spec 

52 

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

54 # Docstring inherited from OpaqueTableStorage. 

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

56 

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

58 # Docstring inherited from OpaqueTableStorage. 

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

60 

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

62 # Docstring inherited from OpaqueTableStorage. 

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

64 

65 def _insert( 

66 self, 

67 *data: dict, 

68 transaction: DatastoreTransaction | None = None, 

69 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT, 

70 ) -> None: 

71 uniqueConstraints = list(self._spec.unique) 

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

73 for d in data: 

74 skipping = False 

75 for constraint in uniqueConstraints: 

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

77 if len(matching) != 0: 

78 match insert_mode: 

79 case DatabaseInsertMode.INSERT: 

80 raise RuntimeError( 

81 f"Unique constraint {constraint} violation in external table {self.name}." 

82 ) 

83 case DatabaseInsertMode.ENSURE: 

84 # Row already exists. Skip. 

85 skipping = True 

86 case DatabaseInsertMode.REPLACE: 

87 # Should try to put these rows back on transaction 

88 # rollback... 

89 self.delete([], *matching) 

90 case _: 

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

92 

93 if skipping: 

94 continue 

95 self._rows.append(d) 

96 if transaction is not None: 

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

98 

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

100 # Docstring inherited from OpaqueTableStorage. 

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

102 

103 # Can support an IN operator if given list. 

104 wherein = {} 

105 for k in list(where): 

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

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

108 del where[k] 

109 

110 for d in self._rows: 

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

112 if wherein: 

113 match = True 

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

115 if d[k] not in v: 

116 match = False 

117 break 

118 if match: 

119 yield d 

120 else: 

121 yield d 

122 

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

124 # Docstring inherited from OpaqueTableStorage. 

125 kept_rows = [] 

126 for table_row in self._rows: 

127 for where_row in rows: 

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

129 break 

130 else: 

131 kept_rows.append(table_row) 

132 self._rows = kept_rows 

133 

134 

135class DummyOpaqueTableStorageManager(OpaqueTableStorageManager): 

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

137 super().__init__(registry_schema_version=registry_schema_version) 

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

139 

140 @classmethod 

141 def initialize( 

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

143 ) -> OpaqueTableStorageManager: 

144 # Docstring inherited from OpaqueTableStorageManager. 

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

146 return cls(registry_schema_version=registry_schema_version) 

147 

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

149 # Docstring inherited from OpaqueTableStorageManager. 

150 return self._storages.get(name) 

151 

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

153 # Docstring inherited from OpaqueTableStorageManager. 

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

155 

156 @classmethod 

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

158 # Docstring inherited from VersionedExtension. 

159 return [] 

160 

161 

162class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager): 

163 def __init__( 

164 self, 

165 opaque: OpaqueTableStorageManager, 

166 universe: DimensionUniverse, 

167 datasetIdColumnType: type, 

168 registry_schema_version: VersionTuple | None = None, 

169 ): 

170 super().__init__( 

171 opaque=opaque, 

172 universe=universe, 

173 datasetIdColumnType=datasetIdColumnType, 

174 registry_schema_version=registry_schema_version, 

175 ) 

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

177 

178 @classmethod 

179 def initialize( 

180 cls, 

181 db: Database, 

182 context: StaticTablesContext, 

183 *, 

184 opaque: OpaqueTableStorageManager, 

185 datasets: type[DatasetRecordStorageManager], 

186 universe: DimensionUniverse, 

187 registry_schema_version: VersionTuple | None = None, 

188 ) -> DatastoreRegistryBridgeManager: 

189 # Docstring inherited from DatastoreRegistryBridgeManager 

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

191 return cls( 

192 opaque=opaque, 

193 universe=universe, 

194 datasetIdColumnType=datasets.getIdColumnType(), 

195 registry_schema_version=registry_schema_version, 

196 ) 

197 

198 def refresh(self) -> None: 

199 # Docstring inherited from DatastoreRegistryBridgeManager 

200 pass 

201 

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

203 # Docstring inherited from DatastoreRegistryBridgeManager 

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

205 

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

207 # Docstring inherited from DatastoreRegistryBridgeManager 

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

209 if ref in bridge: 

210 yield name 

211 

212 @classmethod 

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

214 # Docstring inherited from VersionedExtension. 

215 return [] 

216 

217 

218class DummyRegistry: 

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

220 

221 def __init__(self) -> None: 

222 self._opaque = DummyOpaqueTableStorageManager() 

223 self.dimensions = DimensionUniverse() 

224 self._datastoreBridges = DummyDatastoreRegistryBridgeManager( 

225 self._opaque, self.dimensions, sqlalchemy.BigInteger 

226 ) 

227 

228 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager: 

229 return self._datastoreBridges