Coverage for python/lsst/daf/butler/tests/hybrid_butler_registry.py: 0%

117 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-25 08:37 +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 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/>. 

27 

28from __future__ import annotations 

29 

30import contextlib 

31from collections.abc import Iterable, Iterator, Mapping, Sequence 

32from typing import Any 

33 

34from .._collection_type import CollectionType 

35from .._dataset_association import DatasetAssociation 

36from .._dataset_ref import DatasetId, DatasetIdGenEnum, DatasetRef 

37from .._dataset_type import DatasetType 

38from .._storage_class import StorageClassFactory 

39from .._timespan import Timespan 

40from ..dimensions import ( 

41 DataCoordinate, 

42 DataId, 

43 DimensionElement, 

44 DimensionGroup, 

45 DimensionRecord, 

46 DimensionUniverse, 

47) 

48from ..registry import CollectionArgType, CollectionSummary, Registry, RegistryDefaults 

49from ..registry.queries import ( 

50 DataCoordinateQueryResults, 

51 DatasetQueryResults, 

52 DimensionRecordQueryResults, 

53) 

54from ..registry.sql_registry import SqlRegistry 

55 

56 

57class HybridButlerRegistry(Registry): 

58 """A `Registry` that delegates methods to internal `RemoteButlerRegistry` 

59 and `SqlRegistry` instances. Intended to allow testing of `RemoteButler` 

60 before its implementation is complete, by delegating unsupported methods to 

61 the direct `SqlRegistry`. 

62 

63 Parameters 

64 ---------- 

65 direct : `SqlRegistry` 

66 DirectButler SqlRegistry used to provide methods not yet implemented by 

67 RemoteButlerRegistry. 

68 

69 remote : `Registry` 

70 The RemoteButler Registry implementation we are intending to test. 

71 """ 

72 

73 def __init__(self, direct: SqlRegistry, remote: Registry): 

74 self._direct = direct 

75 self._remote = remote 

76 

77 def isWriteable(self) -> bool: 

78 return self._remote.isWriteable() 

79 

80 @property 

81 def dimensions(self) -> DimensionUniverse: 

82 return self._remote.dimensions 

83 

84 @property 

85 def defaults(self) -> RegistryDefaults: 

86 return self._remote.defaults 

87 

88 @defaults.setter 

89 def defaults(self, value: RegistryDefaults) -> None: 

90 # Make a copy before assigning the value. 

91 # When assigned, it will have finish() called on it -- we don't want to 

92 # intermingle the results of that between Remote and Direct, because 

93 # that could let the Remote side cheat. 

94 copy = RegistryDefaults(value.collections, value.run, value._infer, **value._kwargs) 

95 self._remote.defaults = value 

96 self._direct.defaults = copy 

97 

98 def refresh(self) -> None: 

99 self._direct.refresh() 

100 

101 def refresh_collection_summaries(self) -> None: 

102 self._direct.refresh_collection_summaries() 

103 

104 @contextlib.contextmanager 

105 def caching_context(self) -> Iterator[None]: 

106 with self._direct.caching_context(): 

107 with self._remote.caching_context(): 

108 yield 

109 

110 @contextlib.contextmanager 

111 def transaction(self, *, savepoint: bool = False) -> Iterator[None]: 

112 # RemoteButler doesn't support transactions, and if the direct registry 

113 # enters one its changes are invisible to the remote side. 

114 raise NotImplementedError() 

115 

116 def registerCollection( 

117 self, name: str, type: CollectionType = CollectionType.TAGGED, doc: str | None = None 

118 ) -> bool: 

119 return self._direct.registerCollection(name, type, doc) 

120 

121 def getCollectionType(self, name: str) -> CollectionType: 

122 return self._remote.getCollectionType(name) 

123 

124 def registerRun(self, name: str, doc: str | None = None) -> bool: 

125 return self._direct.registerRun(name, doc) 

126 

127 def removeCollection(self, name: str) -> None: 

128 return self._direct.removeCollection(name) 

129 

130 def getCollectionChain(self, parent: str) -> Sequence[str]: 

131 return self._remote.getCollectionChain(parent) 

132 

133 def setCollectionChain(self, parent: str, children: Any, *, flatten: bool = False) -> None: 

134 return self._direct.setCollectionChain(parent, children, flatten=flatten) 

135 

136 def getCollectionParentChains(self, collection: str) -> set[str]: 

137 return self._remote.getCollectionParentChains(collection) 

138 

139 def getCollectionDocumentation(self, collection: str) -> str | None: 

140 return self._remote.getCollectionDocumentation(collection) 

141 

142 def setCollectionDocumentation(self, collection: str, doc: str | None) -> None: 

143 return self._direct.setCollectionDocumentation(collection, doc) 

144 

145 def getCollectionSummary(self, collection: str) -> CollectionSummary: 

146 return self._remote.getCollectionSummary(collection) 

147 

148 def registerDatasetType(self, datasetType: DatasetType) -> bool: 

149 # We need to make sure that dataset type universe is the same as 

150 # direct registry universe. Only the remote universe is converted; 

151 # anything else is passed through so that the direct registry can 

152 # apply its strict universe check. 

153 if datasetType.dimensions.universe is self._remote.dimensions: 

154 datasetType = datasetType.conform_to(self._direct.dimensions) 

155 return self._direct.registerDatasetType(datasetType) 

156 

157 def removeDatasetType(self, name: str | tuple[str, ...]) -> None: 

158 return self._direct.removeDatasetType(name) 

159 

160 def getDatasetType(self, name: str) -> DatasetType: 

161 return self._remote.getDatasetType(name) 

162 

163 def supportsIdGenerationMode(self, mode: DatasetIdGenEnum) -> bool: 

164 return self._direct.supportsIdGenerationMode(mode) 

165 

166 def findDataset( 

167 self, 

168 datasetType: DatasetType | str, 

169 dataId: DataId | None = None, 

170 *, 

171 collections: CollectionArgType | None = None, 

172 timespan: Timespan | None = None, 

173 datastore_records: bool = False, 

174 **kwargs: Any, 

175 ) -> DatasetRef | None: 

176 return self._remote.findDataset( 

177 datasetType, 

178 dataId, 

179 collections=collections, 

180 timespan=timespan, 

181 datastore_records=datastore_records, 

182 **kwargs, 

183 ) 

184 

185 def insertDatasets( 

186 self, 

187 datasetType: DatasetType | str, 

188 dataIds: Iterable[DataId], 

189 run: str | None = None, 

190 expand: bool = True, 

191 idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE, 

192 ) -> list[DatasetRef]: 

193 return self._direct.insertDatasets(datasetType, dataIds, run, expand, idGenerationMode) 

194 

195 def _importDatasets( 

196 self, datasets: Iterable[DatasetRef], expand: bool = True, assume_new: bool = False 

197 ) -> list[DatasetRef]: 

198 return self._direct._importDatasets(datasets, expand, assume_new) 

199 

200 def getDataset(self, id: DatasetId) -> DatasetRef | None: 

201 return self._remote.getDataset(id) 

202 

203 def _fetch_run_dataset_ids(self, run: str) -> list[DatasetId]: 

204 return self._direct._fetch_run_dataset_ids(run) 

205 

206 def removeDatasets(self, refs: Iterable[DatasetRef]) -> None: 

207 return self._direct.removeDatasets(refs) 

208 

209 def associate(self, collection: str, refs: Iterable[DatasetRef]) -> None: 

210 return self._direct.associate(collection, refs) 

211 

212 def disassociate(self, collection: str, refs: Iterable[DatasetRef]) -> None: 

213 return self._direct.disassociate(collection, refs) 

214 

215 def certify(self, collection: str, refs: Iterable[DatasetRef], timespan: Timespan) -> None: 

216 return self._direct.certify(collection, refs, timespan) 

217 

218 def decertify( 

219 self, 

220 collection: str, 

221 datasetType: str | DatasetType, 

222 timespan: Timespan, 

223 *, 

224 dataIds: Iterable[DataId] | None = None, 

225 ) -> None: 

226 return self._direct.decertify(collection, datasetType, timespan, dataIds=dataIds) 

227 

228 def getDatasetLocations(self, ref: DatasetRef) -> Iterable[str]: 

229 return self._direct.getDatasetLocations(ref) 

230 

231 def expandDataId( 

232 self, 

233 dataId: DataId | None = None, 

234 *, 

235 dimensions: Iterable[str] | DimensionGroup | None = None, 

236 records: Mapping[str, DimensionRecord | None] | None = None, 

237 withDefaults: bool = True, 

238 **kwargs: Any, 

239 ) -> DataCoordinate: 

240 return self._remote.expandDataId( 

241 dataId, dimensions=dimensions, records=records, withDefaults=withDefaults, **kwargs 

242 ) 

243 

244 def insertDimensionData( 

245 self, 

246 element: DimensionElement | str, 

247 *data: Mapping[str, Any] | DimensionRecord, 

248 conform: bool = True, 

249 replace: bool = False, 

250 skip_existing: bool = False, 

251 ) -> None: 

252 return self._direct.insertDimensionData( 

253 element, *data, conform=conform, replace=replace, skip_existing=skip_existing 

254 ) 

255 

256 def syncDimensionData( 

257 self, 

258 element: DimensionElement | str, 

259 row: Mapping[str, Any] | DimensionRecord, 

260 conform: bool = True, 

261 update: bool = False, 

262 ) -> bool | dict[str, Any]: 

263 return self._direct.syncDimensionData(element, row, conform, update) 

264 

265 def queryDatasetTypes( 

266 self, 

267 expression: Any = ..., 

268 *, 

269 missing: list[str] | None = None, 

270 ) -> Iterable[DatasetType]: 

271 return self._remote.queryDatasetTypes(expression, missing=missing) 

272 

273 def queryCollections( 

274 self, 

275 expression: Any = ..., 

276 datasetType: DatasetType | None = None, 

277 collectionTypes: Iterable[CollectionType] | CollectionType = CollectionType.all(), 

278 flattenChains: bool = False, 

279 includeChains: bool | None = None, 

280 ) -> Sequence[str]: 

281 return self._remote.queryCollections( 

282 expression, datasetType, collectionTypes, flattenChains, includeChains 

283 ) 

284 

285 def queryDatasets( 

286 self, 

287 datasetType: Any, 

288 *, 

289 collections: CollectionArgType | None = None, 

290 dimensions: Iterable[str] | None = None, 

291 dataId: DataId | None = None, 

292 where: str = "", 

293 findFirst: bool = False, 

294 bind: Mapping[str, Any] | None = None, 

295 check: bool = True, 

296 **kwargs: Any, 

297 ) -> DatasetQueryResults: 

298 return self._remote.queryDatasets( 

299 datasetType, 

300 collections=collections, 

301 dimensions=dimensions, 

302 dataId=dataId, 

303 where=where, 

304 findFirst=findFirst, 

305 bind=bind, 

306 check=check, 

307 **kwargs, 

308 ) 

309 

310 def queryDataIds( 

311 self, 

312 dimensions: DimensionGroup | Iterable[str] | str, 

313 *, 

314 dataId: DataId | None = None, 

315 datasets: Any = None, 

316 collections: CollectionArgType | None = None, 

317 where: str = "", 

318 bind: Mapping[str, Any] | None = None, 

319 check: bool = True, 

320 **kwargs: Any, 

321 ) -> DataCoordinateQueryResults: 

322 return self._remote.queryDataIds( 

323 dimensions, 

324 dataId=dataId, 

325 datasets=datasets, 

326 collections=collections, 

327 where=where, 

328 bind=bind, 

329 check=check, 

330 **kwargs, 

331 ) 

332 

333 def queryDimensionRecords( 

334 self, 

335 element: DimensionElement | str, 

336 *, 

337 dataId: DataId | None = None, 

338 datasets: Any = None, 

339 collections: CollectionArgType | None = None, 

340 where: str = "", 

341 bind: Mapping[str, Any] | None = None, 

342 check: bool = True, 

343 **kwargs: Any, 

344 ) -> DimensionRecordQueryResults: 

345 return self._remote.queryDimensionRecords( 

346 element, 

347 dataId=dataId, 

348 datasets=datasets, 

349 collections=collections, 

350 where=where, 

351 bind=bind, 

352 check=check, 

353 **kwargs, 

354 ) 

355 

356 def queryDatasetAssociations( 

357 self, 

358 datasetType: str | DatasetType, 

359 collections: CollectionArgType | None = ..., 

360 *, 

361 collectionTypes: Iterable[CollectionType] = CollectionType.all(), 

362 flattenChains: bool = False, 

363 ) -> Iterator[DatasetAssociation]: 

364 return self._remote.queryDatasetAssociations( 

365 datasetType, collections, collectionTypes=collectionTypes, flattenChains=flattenChains 

366 ) 

367 

368 @property 

369 def storageClasses(self) -> StorageClassFactory: 

370 return self._remote.storageClasses 

371 

372 @storageClasses.setter 

373 def storageClasses(self, value: StorageClassFactory) -> None: 

374 raise NotImplementedError()