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

90 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-14 19:21 +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 DatasetIdRef, 

34 DatasetRecordStorageManager, 

35 DatastoreRegistryBridge, 

36 DatastoreRegistryBridgeManager, 

37 OpaqueTableStorage, 

38 OpaqueTableStorageManager, 

39 StaticTablesContext, 

40 VersionTuple, 

41) 

42 

43from ..core.datastore import DatastoreTransaction 

44 

45 

46class DummyOpaqueTableStorage(OpaqueTableStorage): 

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

48 super().__init__(name=name) 

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

50 self._spec = spec 

51 

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

53 # Docstring inherited from OpaqueTableStorage. 

54 uniqueConstraints = list(self._spec.unique) 

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

56 for d in data: 

57 for constraint in uniqueConstraints: 

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

59 if len(matching) != 0: 

60 raise RuntimeError( 

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

62 ) 

63 self._rows.append(d) 

64 if transaction is not None: 

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

66 

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

68 # Docstring inherited from OpaqueTableStorage. 

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

70 

71 # Can support an IN operator if given list. 

72 wherein = {} 

73 for k in list(where): 

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

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

76 del where[k] 

77 

78 for d in self._rows: 

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

80 if wherein: 

81 match = True 

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

83 if d[k] not in v: 

84 match = False 

85 break 

86 if match: 

87 yield d 

88 else: 

89 yield d 

90 

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

92 # Docstring inherited from OpaqueTableStorage. 

93 kept_rows = [] 

94 for table_row in self._rows: 

95 for where_row in rows: 

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

97 break 

98 else: 

99 kept_rows.append(table_row) 

100 self._rows = kept_rows 

101 

102 

103class DummyOpaqueTableStorageManager(OpaqueTableStorageManager): 

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

105 super().__init__(registry_schema_version=registry_schema_version) 

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

107 

108 @classmethod 

109 def initialize( 

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

111 ) -> OpaqueTableStorageManager: 

112 # Docstring inherited from OpaqueTableStorageManager. 

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

114 return cls(registry_schema_version=registry_schema_version) 

115 

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

117 # Docstring inherited from OpaqueTableStorageManager. 

118 return self._storages.get(name) 

119 

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

121 # Docstring inherited from OpaqueTableStorageManager. 

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

123 

124 @classmethod 

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

126 # Docstring inherited from VersionedExtension. 

127 return [] 

128 

129 

130class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager): 

131 def __init__( 

132 self, 

133 opaque: OpaqueTableStorageManager, 

134 universe: DimensionUniverse, 

135 datasetIdColumnType: type, 

136 registry_schema_version: VersionTuple | None = None, 

137 ): 

138 super().__init__( 

139 opaque=opaque, 

140 universe=universe, 

141 datasetIdColumnType=datasetIdColumnType, 

142 registry_schema_version=registry_schema_version, 

143 ) 

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

145 

146 @classmethod 

147 def initialize( 

148 cls, 

149 db: Database, 

150 context: StaticTablesContext, 

151 *, 

152 opaque: OpaqueTableStorageManager, 

153 datasets: type[DatasetRecordStorageManager], 

154 universe: DimensionUniverse, 

155 registry_schema_version: VersionTuple | None = None, 

156 ) -> DatastoreRegistryBridgeManager: 

157 # Docstring inherited from DatastoreRegistryBridgeManager 

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

159 return cls( 

160 opaque=opaque, 

161 universe=universe, 

162 datasetIdColumnType=datasets.getIdColumnType(), 

163 registry_schema_version=registry_schema_version, 

164 ) 

165 

166 def refresh(self) -> None: 

167 # Docstring inherited from DatastoreRegistryBridgeManager 

168 pass 

169 

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

171 # Docstring inherited from DatastoreRegistryBridgeManager 

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

173 

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

175 # Docstring inherited from DatastoreRegistryBridgeManager 

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

177 if ref in bridge: 

178 yield name 

179 

180 @classmethod 

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

182 # Docstring inherited from VersionedExtension. 

183 return [] 

184 

185 

186class DummyRegistry: 

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

188 

189 def __init__(self) -> None: 

190 self._opaque = DummyOpaqueTableStorageManager() 

191 self.dimensions = DimensionUniverse() 

192 self._datastoreBridges = DummyDatastoreRegistryBridgeManager( 

193 self._opaque, self.dimensions, sqlalchemy.BigInteger 

194 ) 

195 

196 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager: 

197 return self._datastoreBridges