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

90 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-10-26 15:15 +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 

43 

44class DummyOpaqueTableStorage(OpaqueTableStorage): 

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

46 super().__init__(name=name) 

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

48 self._spec = spec 

49 

50 def insert(self, *data: dict) -> None: 

51 # Docstring inherited from OpaqueTableStorage. 

52 uniqueConstraints = list(self._spec.unique) 

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

54 for d in data: 

55 for constraint in uniqueConstraints: 

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

57 if len(matching) != 0: 

58 raise RuntimeError( 

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

60 ) 

61 self._rows.append(d) 

62 

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

64 # Docstring inherited from OpaqueTableStorage. 

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

66 

67 # Can support an IN operator if given list. 

68 wherein = {} 

69 for k in list(where): 

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

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

72 del where[k] 

73 

74 for d in self._rows: 

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

76 if wherein: 

77 match = True 

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

79 if d[k] not in v: 

80 match = False 

81 break 

82 if match: 

83 yield d 

84 else: 

85 yield d 

86 

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

88 # Docstring inherited from OpaqueTableStorage. 

89 kept_rows = [] 

90 for table_row in self._rows: 

91 for where_row in rows: 

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

93 break 

94 else: 

95 kept_rows.append(table_row) 

96 self._rows = kept_rows 

97 

98 

99class DummyOpaqueTableStorageManager(OpaqueTableStorageManager): 

100 def __init__(self) -> None: 

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

102 

103 @classmethod 

104 def initialize(cls, db: Database, context: StaticTablesContext) -> OpaqueTableStorageManager: 

105 # Docstring inherited from OpaqueTableStorageManager. 

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

107 return cls() 

108 

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

110 # Docstring inherited from OpaqueTableStorageManager. 

111 return self._storages.get(name) 

112 

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

114 # Docstring inherited from OpaqueTableStorageManager. 

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

116 

117 @classmethod 

118 def currentVersion(cls) -> VersionTuple | None: 

119 # Docstring inherited from VersionedExtension. 

120 return None 

121 

122 def schemaDigest(self) -> str | None: 

123 # Docstring inherited from VersionedExtension. 

124 return None 

125 

126 

127class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager): 

128 def __init__( 

129 self, opaque: OpaqueTableStorageManager, universe: DimensionUniverse, datasetIdColumnType: type 

130 ): 

131 super().__init__(opaque=opaque, universe=universe, datasetIdColumnType=datasetIdColumnType) 

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

133 

134 @classmethod 

135 def initialize( 

136 cls, 

137 db: Database, 

138 context: StaticTablesContext, 

139 *, 

140 opaque: OpaqueTableStorageManager, 

141 datasets: type[DatasetRecordStorageManager], 

142 universe: DimensionUniverse, 

143 ) -> DatastoreRegistryBridgeManager: 

144 # Docstring inherited from DatastoreRegistryBridgeManager 

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

146 return cls(opaque=opaque, universe=universe, datasetIdColumnType=datasets.getIdColumnType()) 

147 

148 def refresh(self) -> None: 

149 # Docstring inherited from DatastoreRegistryBridgeManager 

150 pass 

151 

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

153 # Docstring inherited from DatastoreRegistryBridgeManager 

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

155 

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

157 # Docstring inherited from DatastoreRegistryBridgeManager 

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

159 if ref in bridge: 

160 yield name 

161 

162 @classmethod 

163 def currentVersion(cls) -> VersionTuple | None: 

164 # Docstring inherited from VersionedExtension. 

165 return None 

166 

167 def schemaDigest(self) -> str | None: 

168 # Docstring inherited from VersionedExtension. 

169 return None 

170 

171 

172class DummyRegistry: 

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

174 

175 def __init__(self) -> None: 

176 self._opaque = DummyOpaqueTableStorageManager() 

177 self.dimensions = DimensionUniverse() 

178 self._datastoreBridges = DummyDatastoreRegistryBridgeManager( 

179 self._opaque, self.dimensions, sqlalchemy.BigInteger 

180 ) 

181 

182 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager: 

183 return self._datastoreBridges