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

88 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-21 02:43 -0700

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/>. 

21 

22__all__ = ("DummyRegistry",) 

23 

24 

25from typing import Any, Iterable, Iterator, Optional, Type 

26 

27import sqlalchemy 

28from lsst.daf.butler import DatasetRef, DimensionUniverse, ddl 

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

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

31 Database, 

32 DatasetRecordStorageManager, 

33 DatastoreRegistryBridge, 

34 DatastoreRegistryBridgeManager, 

35 OpaqueTableStorage, 

36 OpaqueTableStorageManager, 

37 StaticTablesContext, 

38 VersionTuple, 

39) 

40 

41 

42class DummyOpaqueTableStorage(OpaqueTableStorage): 

43 def __init__(self, name: str, spec: ddl.TableSpec): 

44 super().__init__(name=name) 

45 self._rows = [] 

46 self._spec = spec 

47 

48 def insert(self, *data: dict): 

49 # Docstring inherited from OpaqueTableStorage. 

50 uniqueConstraints = list(self._spec.unique) 

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

52 for d in data: 

53 for constraint in uniqueConstraints: 

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

55 if len(matching) != 0: 

56 raise RuntimeError( 

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

58 ) 

59 self._rows.append(d) 

60 

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

62 # Docstring inherited from OpaqueTableStorage. 

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

64 

65 # Can support an IN operator if given list. 

66 wherein = {} 

67 for k in list(where): 

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

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

70 del where[k] 

71 

72 for d in self._rows: 

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

74 if wherein: 

75 match = True 

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

77 if d[k] not in v: 

78 match = False 

79 break 

80 if match: 

81 yield d 

82 else: 

83 yield d 

84 

85 def delete(self, columns: Iterable[str], *rows: dict): 

86 # Docstring inherited from OpaqueTableStorage. 

87 kept_rows = [] 

88 for table_row in self._rows: 

89 for where_row in rows: 

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

91 break 

92 else: 

93 kept_rows.append(table_row) 

94 self._rows = kept_rows 

95 

96 

97class DummyOpaqueTableStorageManager(OpaqueTableStorageManager): 

98 def __init__(self): 

99 self._storages = {} 

100 

101 @classmethod 

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

103 # Docstring inherited from OpaqueTableStorageManager. 

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

105 return cls() 

106 

107 def get(self, name: str) -> Optional[OpaqueTableStorage]: 

108 # Docstring inherited from OpaqueTableStorageManager. 

109 return self._storage.get(name) 

110 

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

112 # Docstring inherited from OpaqueTableStorageManager. 

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

114 

115 @classmethod 

116 def currentVersion(cls) -> Optional[VersionTuple]: 

117 # Docstring inherited from VersionedExtension. 

118 return None 

119 

120 def schemaDigest(self) -> Optional[str]: 

121 # Docstring inherited from VersionedExtension. 

122 return None 

123 

124 

125class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager): 

126 def __init__( 

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

128 ): 

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

130 self._bridges = {} 

131 

132 @classmethod 

133 def initialize( 

134 cls, 

135 db: Database, 

136 context: StaticTablesContext, 

137 *, 

138 opaque: OpaqueTableStorageManager, 

139 datasets: Type[DatasetRecordStorageManager], 

140 universe: DimensionUniverse, 

141 ) -> DatastoreRegistryBridgeManager: 

142 # Docstring inherited from DatastoreRegistryBridgeManager 

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

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

145 

146 def refresh(self): 

147 # Docstring inherited from DatastoreRegistryBridgeManager 

148 pass 

149 

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

151 # Docstring inherited from DatastoreRegistryBridgeManager 

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

153 

154 def findDatastores(self, ref: DatasetRef) -> Iterable[str]: 

155 # Docstring inherited from DatastoreRegistryBridgeManager 

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

157 if ref in bridge: 

158 yield name 

159 

160 @classmethod 

161 def currentVersion(cls) -> Optional[VersionTuple]: 

162 # Docstring inherited from VersionedExtension. 

163 return None 

164 

165 def schemaDigest(self) -> Optional[str]: 

166 # Docstring inherited from VersionedExtension. 

167 return None 

168 

169 

170class DummyRegistry: 

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

172 

173 def __init__(self): 

174 self._opaque = DummyOpaqueTableStorageManager() 

175 self.dimensions = DimensionUniverse() 

176 self._datastoreBridges = DummyDatastoreRegistryBridgeManager( 

177 self._opaque, self.dimensions, sqlalchemy.BigInteger 

178 ) 

179 

180 def getDatastoreBridgeManager(self): 

181 return self._datastoreBridges