Coverage for tests / test_datastore.py: 13%

1200 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-12 01:12 -0800

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 

31import os 

32import pickle 

33import shutil 

34import tempfile 

35import time 

36import unittest 

37import unittest.mock 

38import uuid 

39from collections.abc import Callable, Iterator 

40from typing import Any, cast 

41 

42import yaml 

43 

44import lsst.utils.tests 

45from lsst.daf.butler import ( 

46 Config, 

47 DataCoordinate, 

48 DatasetIdGenEnum, 

49 DatasetRef, 

50 DatasetType, 

51 DatasetTypeNotSupportedError, 

52 Datastore, 

53 DimensionUniverse, 

54 FileDataset, 

55 StorageClass, 

56 StorageClassFactory, 

57) 

58from lsst.daf.butler.datastore import DatasetRefURIs, DatastoreConfig, DatastoreValidationError, NullDatastore 

59from lsst.daf.butler.datastore.cache_manager import ( 

60 DatastoreCacheManager, 

61 DatastoreCacheManagerConfig, 

62 DatastoreDisabledCacheManager, 

63) 

64from lsst.daf.butler.datastore.stored_file_info import StoredFileInfo, make_datastore_path_relative 

65from lsst.daf.butler.formatters.yaml import YamlFormatter 

66from lsst.daf.butler.tests import ( 

67 BadNoWriteFormatter, 

68 BadWriteFormatter, 

69 DatasetTestHelper, 

70 DatastoreTestHelper, 

71 DummyRegistry, 

72 MetricsExample, 

73 MetricsExampleDataclass, 

74 MetricsExampleModel, 

75) 

76from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel 

77from lsst.daf.butler.tests.utils import TestCaseMixin 

78from lsst.resources import ResourcePath 

79from lsst.utils import doImport 

80 

81TESTDIR = os.path.dirname(__file__) 

82 

83 

84def makeExampleMetrics(use_none: bool = False) -> MetricsExample: 

85 """Make example dataset that can be stored in butler.""" 

86 if use_none: 

87 array = None 

88 else: 

89 array = [563, 234, 456.7, 105, 2054, -1045] 

90 return MetricsExample( 

91 {"AM1": 5.2, "AM2": 30.6}, 

92 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

93 array, 

94 ) 

95 

96 

97class TransactionTestError(Exception): 

98 """Specific error for transactions, to prevent misdiagnosing 

99 that might otherwise occur when a standard exception is used. 

100 """ 

101 

102 pass 

103 

104 

105class DatastoreTestsBase(DatasetTestHelper, DatastoreTestHelper, TestCaseMixin): 

106 """Support routines for datastore testing""" 

107 

108 root: str | None = None 

109 universe: DimensionUniverse 

110 storageClassFactory: StorageClassFactory 

111 

112 @classmethod 

113 def setUpClass(cls) -> None: 

114 # Storage Classes are fixed for all datastores in these tests 

115 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml") 

116 cls.storageClassFactory = StorageClassFactory() 

117 cls.storageClassFactory.addFromConfig(scConfigFile) 

118 

119 # Read the Datastore config so we can get the class 

120 # information (since we should not assume the constructor 

121 # name here, but rely on the configuration file itself) 

122 datastoreConfig = DatastoreConfig(cls.configFile) 

123 cls.datastoreType = cast(type[Datastore], doImport(datastoreConfig["cls"])) 

124 cls.universe = DimensionUniverse() 

125 

126 def setUp(self) -> None: 

127 self.setUpDatastoreTests(DummyRegistry, DatastoreConfig) 

128 

129 def tearDown(self) -> None: 

130 if self.root is not None and os.path.exists(self.root): 

131 shutil.rmtree(self.root, ignore_errors=True) 

132 

133 

134class DatastoreTests(DatastoreTestsBase): 

135 """Some basic tests of a simple datastore.""" 

136 

137 hasUnsupportedPut = True 

138 rootKeys: tuple[str, ...] | None = None 

139 isEphemeral: bool = False 

140 validationCanFail: bool = False 

141 

142 def testConfigRoot(self) -> None: 

143 full = DatastoreConfig(self.configFile) 

144 config = DatastoreConfig(self.configFile, mergeDefaults=False) 

145 newroot = "/random/location" 

146 self.datastoreType.setConfigRoot(newroot, config, full) 

147 if self.rootKeys: 

148 for k in self.rootKeys: 

149 self.assertIn(newroot, config[k]) 

150 

151 def testConstructor(self) -> None: 

152 datastore = self.makeDatastore() 

153 self.assertIsNotNone(datastore) 

154 self.assertIs(datastore.isEphemeral, self.isEphemeral) 

155 

156 def testConfigurationValidation(self) -> None: 

157 datastore = self.makeDatastore() 

158 sc = self.storageClassFactory.getStorageClass("ThingOne") 

159 datastore.validateConfiguration([sc]) 

160 

161 sc2 = self.storageClassFactory.getStorageClass("ThingTwo") 

162 if self.validationCanFail: 

163 with self.assertRaises(DatastoreValidationError): 

164 datastore.validateConfiguration([sc2], logFailures=True) 

165 

166 dimensions = self.universe.conform(("visit", "physical_filter")) 

167 dataId = { 

168 "instrument": "dummy", 

169 "visit": 52, 

170 "physical_filter": "V", 

171 "band": "v", 

172 "day_obs": 20250101, 

173 } 

174 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

175 datastore.validateConfiguration([ref]) 

176 

177 def testParameterValidation(self) -> None: 

178 """Check that parameters are validated""" 

179 sc = self.storageClassFactory.getStorageClass("ThingOne") 

180 dimensions = self.universe.conform(("visit", "physical_filter")) 

181 dataId = { 

182 "instrument": "dummy", 

183 "visit": 52, 

184 "physical_filter": "V", 

185 "band": "v", 

186 "day_obs": 20250101, 

187 } 

188 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

189 datastore = self.makeDatastore() 

190 data = {1: 2, 3: 4} 

191 datastore.put(data, ref) 

192 newdata = datastore.get(ref) 

193 self.assertEqual(data, newdata) 

194 with self.assertRaises(KeyError): 

195 newdata = datastore.get(ref, parameters={"missing": 5}) 

196 

197 def testBasicPutGet(self) -> None: 

198 metrics = makeExampleMetrics() 

199 datastore = self.makeDatastore() 

200 

201 # Create multiple storage classes for testing different formulations 

202 storageClasses = [ 

203 self.storageClassFactory.getStorageClass(sc) 

204 for sc in ("StructuredData", "StructuredDataJson", "StructuredDataPickle") 

205 ] 

206 

207 dimensions = self.universe.conform(("visit", "physical_filter")) 

208 dataId = { 

209 "instrument": "dummy", 

210 "visit": 52, 

211 "physical_filter": "V", 

212 "band": "v", 

213 "day_obs": 20250101, 

214 } 

215 dataId2 = { 

216 "instrument": "dummy", 

217 "visit": 53, 

218 "physical_filter": "V", 

219 "band": "v", 

220 "day_obs": 20250101, 

221 } 

222 

223 for sc in storageClasses: 

224 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

225 ref2 = self.makeDatasetRef("metric", dimensions, sc, dataId2) 

226 

227 # Make sure that using getManyURIs without predicting before the 

228 # dataset has been put raises. 

229 with self.assertRaises(FileNotFoundError): 

230 datastore.getManyURIs([ref], predict=False) 

231 

232 # Make sure that using getManyURIs with predicting before the 

233 # dataset has been put predicts the URI. 

234 uris = datastore.getManyURIs([ref, ref2], predict=True) 

235 self.assertIn("52", uris[ref].primaryURI.geturl()) 

236 self.assertIn("#predicted", uris[ref].primaryURI.geturl()) 

237 self.assertIn("53", uris[ref2].primaryURI.geturl()) 

238 self.assertIn("#predicted", uris[ref2].primaryURI.geturl()) 

239 

240 datastore.put(metrics, ref) 

241 

242 # Does it exist? 

243 self.assertTrue(datastore.exists(ref)) 

244 self.assertTrue(datastore.knows(ref)) 

245 multi = datastore.knows_these([ref]) 

246 self.assertTrue(multi[ref]) 

247 multi = datastore.mexists([ref, ref2]) 

248 self.assertTrue(multi[ref]) 

249 self.assertFalse(multi[ref2]) 

250 

251 # Get 

252 metricsOut = datastore.get(ref, parameters=None) 

253 self.assertEqual(metrics, metricsOut) 

254 

255 uri = datastore.getURI(ref) 

256 self.assertEqual(uri.scheme, self.uriScheme) 

257 

258 uris = datastore.getManyURIs([ref]) 

259 self.assertEqual(len(uris), 1) 

260 ref, uri = uris.popitem() 

261 self.assertTrue(uri.primaryURI.exists()) 

262 self.assertFalse(uri.componentURIs) 

263 

264 # Get a component -- we need to construct new refs for them 

265 # with derived storage classes but with parent ID 

266 for comp in ("data", "output"): 

267 compRef = ref.makeComponentRef(comp) 

268 output = datastore.get(compRef) 

269 self.assertEqual(output, getattr(metricsOut, comp)) 

270 

271 uri = datastore.getURI(compRef) 

272 self.assertEqual(uri.scheme, self.uriScheme) 

273 

274 uris = datastore.getManyURIs([compRef]) 

275 self.assertEqual(len(uris), 1) 

276 

277 storageClass = sc 

278 

279 # Check that we can put a metric with None in a component and 

280 # get it back as None 

281 metricsNone = makeExampleMetrics(use_none=True) 

282 dataIdNone = { 

283 "instrument": "dummy", 

284 "visit": 54, 

285 "physical_filter": "V", 

286 "band": "v", 

287 "day_obs": 20250101, 

288 } 

289 refNone = self.makeDatasetRef("metric", dimensions, sc, dataIdNone) 

290 datastore.put(metricsNone, refNone) 

291 

292 comp = "data" 

293 for comp in ("data", "output"): 

294 compRef = refNone.makeComponentRef(comp) 

295 output = datastore.get(compRef) 

296 self.assertEqual(output, getattr(metricsNone, comp)) 

297 

298 # Check that a put fails if the dataset type is not supported 

299 if self.hasUnsupportedPut: 

300 sc = StorageClass("UnsupportedSC", pytype=type(metrics)) 

301 ref = self.makeDatasetRef("unsupportedType", dimensions, sc, dataId) 

302 with self.assertRaises(DatasetTypeNotSupportedError): 

303 datastore.put(metrics, ref) 

304 

305 # These should raise 

306 ref = self.makeDatasetRef("metrics", dimensions, storageClass, dataId) 

307 with self.assertRaises(FileNotFoundError): 

308 # non-existing file 

309 datastore.get(ref) 

310 

311 # Get a URI from it 

312 uri = datastore.getURI(ref, predict=True) 

313 self.assertEqual(uri.scheme, self.uriScheme) 

314 

315 with self.assertRaises(FileNotFoundError): 

316 datastore.getURI(ref) 

317 

318 def testTrustGetRequest(self) -> None: 

319 """Check that we can get datasets that registry knows nothing about.""" 

320 datastore = self.makeDatastore() 

321 

322 # Skip test if the attribute is not defined 

323 if not hasattr(datastore, "trustGetRequest"): 

324 return 

325 

326 metrics = makeExampleMetrics() 

327 

328 i = 0 

329 for sc_name in ("StructuredDataNoComponents", "StructuredData", "StructuredComposite"): 

330 i += 1 

331 datasetTypeName = f"test_metric{i}" # Different dataset type name each time. 

332 

333 if sc_name == "StructuredComposite": 

334 disassembled = True 

335 else: 

336 disassembled = False 

337 

338 # Start datastore in default configuration of using registry 

339 datastore.trustGetRequest = False 

340 

341 # Create multiple storage classes for testing with or without 

342 # disassembly 

343 sc = self.storageClassFactory.getStorageClass(sc_name) 

344 dimensions = self.universe.conform(("visit", "physical_filter")) 

345 

346 dataId = { 

347 "instrument": "dummy", 

348 "visit": 52 + i, 

349 "physical_filter": "V", 

350 "band": "v", 

351 "day_obs": 20250101, 

352 } 

353 

354 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId) 

355 datastore.put(metrics, ref) 

356 

357 # Does it exist? 

358 self.assertTrue(datastore.exists(ref)) 

359 self.assertTrue(datastore.knows(ref)) 

360 multi = datastore.knows_these([ref]) 

361 self.assertTrue(multi[ref]) 

362 multi = datastore.mexists([ref]) 

363 self.assertTrue(multi[ref]) 

364 

365 # Get 

366 metricsOut = datastore.get(ref) 

367 self.assertEqual(metrics, metricsOut) 

368 

369 # Get the URI(s) 

370 primaryURI, componentURIs = datastore.getURIs(ref) 

371 if disassembled: 

372 self.assertIsNone(primaryURI) 

373 self.assertEqual(len(componentURIs), 3) 

374 else: 

375 self.assertIn(datasetTypeName, primaryURI.path) 

376 self.assertFalse(componentURIs) 

377 

378 # Delete registry entry so now we are trusting 

379 datastore.removeStoredItemInfo(ref) 

380 

381 # Now stop trusting and check that things break 

382 datastore.trustGetRequest = False 

383 

384 # Does it exist? 

385 self.assertFalse(datastore.exists(ref)) 

386 self.assertFalse(datastore.knows(ref)) 

387 multi = datastore.knows_these([ref]) 

388 self.assertFalse(multi[ref]) 

389 multi = datastore.mexists([ref]) 

390 self.assertFalse(multi[ref]) 

391 

392 with self.assertRaises(FileNotFoundError): 

393 datastore.get(ref) 

394 

395 if sc_name != "StructuredDataNoComponents": 

396 with self.assertRaises(FileNotFoundError): 

397 datastore.get(ref.makeComponentRef("data")) 

398 

399 # URI should fail unless we ask for prediction 

400 with self.assertRaises(FileNotFoundError): 

401 datastore.getURIs(ref) 

402 

403 predicted_primary, predicted_disassembled = datastore.getURIs(ref, predict=True) 

404 if disassembled: 

405 self.assertIsNone(predicted_primary) 

406 self.assertEqual(len(predicted_disassembled), 3) 

407 for uri in predicted_disassembled.values(): 

408 self.assertEqual(uri.fragment, "predicted") 

409 self.assertIn(datasetTypeName, uri.path) 

410 else: 

411 self.assertIn(datasetTypeName, predicted_primary.path) 

412 self.assertFalse(predicted_disassembled) 

413 self.assertEqual(predicted_primary.fragment, "predicted") 

414 

415 # Now enable registry-free trusting mode 

416 datastore.trustGetRequest = True 

417 

418 # Try again to get it 

419 metricsOut = datastore.get(ref) 

420 self.assertEqual(metricsOut, metrics) 

421 

422 # Does it exist? 

423 self.assertTrue(datastore.exists(ref)) 

424 

425 # Get a component 

426 if sc_name != "StructuredDataNoComponents": 

427 comp = "data" 

428 compRef = ref.makeComponentRef(comp) 

429 output = datastore.get(compRef) 

430 self.assertEqual(output, getattr(metrics, comp)) 

431 

432 # Get the URI -- if we trust this should work even without 

433 # enabling prediction. 

434 primaryURI2, componentURIs2 = datastore.getURIs(ref) 

435 self.assertEqual(primaryURI2, primaryURI) 

436 self.assertEqual(componentURIs2, componentURIs) 

437 

438 # Check for compatible storage class. 

439 if sc_name in ("StructuredDataNoComponents", "StructuredData"): 

440 # Make new dataset ref with compatible storage class. 

441 ref_comp = ref.overrideStorageClass("StructuredDataDictJson") 

442 

443 # Without `set_retrieve_dataset_type_method` it will fail to 

444 # find correct file. 

445 self.assertFalse(datastore.exists(ref_comp)) 

446 with self.assertRaises(FileNotFoundError): 

447 datastore.get(ref_comp) 

448 with self.assertRaises(FileNotFoundError): 

449 datastore.get(ref, storageClass="StructuredDataDictJson") 

450 

451 # Need a special method to generate stored dataset type. 

452 def _stored_dataset_type(name: str, ref: DatasetRef = ref) -> DatasetType: 

453 if name == ref.datasetType.name: 

454 return ref.datasetType 

455 raise ValueError(f"Unexpected dataset type name {ref.datasetType.name}") 

456 

457 datastore.set_retrieve_dataset_type_method(_stored_dataset_type) 

458 

459 # Storage class override with original dataset ref. 

460 metrics_as_dict = datastore.get(ref, storageClass="StructuredDataDictJson") 

461 self.assertIsInstance(metrics_as_dict, dict) 

462 

463 # get() should return a dict now. 

464 metrics_as_dict = datastore.get(ref_comp) 

465 self.assertIsInstance(metrics_as_dict, dict) 

466 

467 # exists() should work as well. 

468 self.assertTrue(datastore.exists(ref_comp)) 

469 

470 datastore.set_retrieve_dataset_type_method(None) 

471 

472 def testDisassembly(self) -> None: 

473 """Test disassembly within datastore.""" 

474 metrics = makeExampleMetrics() 

475 if self.isEphemeral: 

476 # in-memory datastore does not disassemble 

477 return 

478 

479 # Create multiple storage classes for testing different formulations 

480 # of composites. One of these will not disassemble to provide 

481 # a reference. 

482 storageClasses = [ 

483 self.storageClassFactory.getStorageClass(sc) 

484 for sc in ( 

485 "StructuredComposite", 

486 "StructuredCompositeTestA", 

487 "StructuredCompositeTestB", 

488 "StructuredCompositeReadComp", 

489 "StructuredData", # No disassembly 

490 "StructuredCompositeReadCompNoDisassembly", 

491 ) 

492 ] 

493 

494 # Create the test datastore 

495 datastore = self.makeDatastore() 

496 

497 # Dummy dataId 

498 dimensions = self.universe.conform(("visit", "physical_filter")) 

499 dataId = {"instrument": "dummy", "visit": 428, "physical_filter": "R"} 

500 

501 for i, sc in enumerate(storageClasses): 

502 with self.subTest(storageClass=sc.name): 

503 # Create a different dataset type each time round 

504 # so that a test failure in this subtest does not trigger 

505 # a cascade of tests because of file clashes 

506 ref = self.makeDatasetRef(f"metric_comp_{i}", dimensions, sc, dataId) 

507 

508 disassembled = sc.name not in {"StructuredData", "StructuredCompositeReadCompNoDisassembly"} 

509 

510 datastore.put(metrics, ref) 

511 

512 baseURI, compURIs = datastore.getURIs(ref) 

513 if disassembled: 

514 self.assertIsNone(baseURI) 

515 self.assertEqual(set(compURIs), {"data", "output", "summary"}) 

516 else: 

517 self.assertIsNotNone(baseURI) 

518 self.assertEqual(compURIs, {}) 

519 

520 metrics_get = datastore.get(ref) 

521 self.assertEqual(metrics_get, metrics) 

522 

523 # Retrieve the composite with read parameter 

524 stop = 4 

525 metrics_get = datastore.get(ref, parameters={"slice": slice(stop)}) 

526 self.assertEqual(metrics_get.summary, metrics.summary) 

527 self.assertEqual(metrics_get.output, metrics.output) 

528 self.assertEqual(metrics_get.data, metrics.data[:stop]) 

529 

530 # Retrieve a component 

531 data = datastore.get(ref.makeComponentRef("data")) 

532 self.assertEqual(data, metrics.data) 

533 

534 # On supported storage classes attempt to access a read 

535 # only component 

536 if "ReadComp" in sc.name: 

537 cRef = ref.makeComponentRef("counter") 

538 counter = datastore.get(cRef) 

539 self.assertEqual(counter, len(metrics.data)) 

540 

541 counter = datastore.get(cRef, parameters={"slice": slice(stop)}) 

542 self.assertEqual(counter, stop) 

543 

544 datastore.remove(ref) 

545 

546 def prepDeleteTest(self, n_refs: int = 1) -> tuple[Datastore, tuple[DatasetRef, ...]]: 

547 metrics = makeExampleMetrics() 

548 datastore = self.makeDatastore() 

549 # Put 

550 dimensions = self.universe.conform(("visit", "physical_filter")) 

551 sc = self.storageClassFactory.getStorageClass("StructuredData") 

552 refs = [] 

553 for i in range(n_refs): 

554 dataId = { 

555 "instrument": "dummy", 

556 "visit": 638 + i, 

557 "physical_filter": "U", 

558 "band": "u", 

559 "day_obs": 20250101, 

560 } 

561 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

562 datastore.put(metrics, ref) 

563 

564 # Does it exist? 

565 self.assertTrue(datastore.exists(ref)) 

566 

567 # Get 

568 metricsOut = datastore.get(ref) 

569 self.assertEqual(metrics, metricsOut) 

570 refs.append(ref) 

571 

572 return datastore, *refs 

573 

574 def testRemove(self) -> None: 

575 datastore, ref = self.prepDeleteTest() 

576 

577 # Remove 

578 datastore.remove(ref) 

579 

580 # Does it exist? 

581 self.assertFalse(datastore.exists(ref)) 

582 

583 # Do we now get a predicted URI? 

584 uri = datastore.getURI(ref, predict=True) 

585 self.assertEqual(uri.fragment, "predicted") 

586 

587 # Get should now fail 

588 with self.assertRaises(FileNotFoundError): 

589 datastore.get(ref) 

590 # Can only delete once 

591 with self.assertRaises(FileNotFoundError): 

592 datastore.remove(ref) 

593 

594 def testForget(self) -> None: 

595 datastore, ref = self.prepDeleteTest() 

596 

597 # Remove 

598 datastore.forget([ref]) 

599 

600 # Does it exist (as far as we know)? 

601 self.assertFalse(datastore.exists(ref)) 

602 

603 # Do we now get a predicted URI? 

604 uri = datastore.getURI(ref, predict=True) 

605 self.assertEqual(uri.fragment, "predicted") 

606 

607 # Get should now fail 

608 with self.assertRaises(FileNotFoundError): 

609 datastore.get(ref) 

610 

611 # Forgetting again is a silent no-op 

612 datastore.forget([ref]) 

613 

614 # Predicted URI should still point to the file. 

615 self.assertTrue(uri.exists()) 

616 

617 def testTransfer(self) -> None: 

618 metrics = makeExampleMetrics() 

619 

620 dimensions = self.universe.conform(("visit", "physical_filter")) 

621 dataId = { 

622 "instrument": "dummy", 

623 "visit": 2048, 

624 "physical_filter": "Uprime", 

625 "band": "u", 

626 "day_obs": 20250101, 

627 } 

628 

629 sc = self.storageClassFactory.getStorageClass("StructuredData") 

630 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

631 

632 inputDatastore = self.makeDatastore("test_input_datastore") 

633 outputDatastore = self.makeDatastore("test_output_datastore") 

634 

635 inputDatastore.put(metrics, ref) 

636 outputDatastore.transfer(inputDatastore, ref) 

637 

638 metricsOut = outputDatastore.get(ref) 

639 self.assertEqual(metrics, metricsOut) 

640 

641 def testBasicTransaction(self) -> None: 

642 datastore = self.makeDatastore() 

643 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

644 dimensions = self.universe.conform(("visit", "physical_filter")) 

645 nDatasets = 6 

646 dataIds = [ 

647 {"instrument": "dummy", "visit": i, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

648 for i in range(nDatasets) 

649 ] 

650 data = [ 

651 ( 

652 self.makeDatasetRef("metric", dimensions, storageClass, dataId), 

653 makeExampleMetrics(), 

654 ) 

655 for dataId in dataIds 

656 ] 

657 succeed = data[: nDatasets // 2] 

658 fail = data[nDatasets // 2 :] 

659 # All datasets added in this transaction should continue to exist 

660 with datastore.transaction(): 

661 for ref, metrics in succeed: 

662 datastore.put(metrics, ref) 

663 # Whereas datasets added in this transaction should not 

664 with self.assertRaises(TransactionTestError): 

665 with datastore.transaction(): 

666 for ref, metrics in fail: 

667 datastore.put(metrics, ref) 

668 raise TransactionTestError("This should propagate out of the context manager") 

669 # Check for datasets that should exist 

670 for ref, metrics in succeed: 

671 # Does it exist? 

672 self.assertTrue(datastore.exists(ref)) 

673 # Get 

674 metricsOut = datastore.get(ref, parameters=None) 

675 self.assertEqual(metrics, metricsOut) 

676 # URI 

677 uri = datastore.getURI(ref) 

678 self.assertEqual(uri.scheme, self.uriScheme) 

679 # Check for datasets that should not exist 

680 for ref, _ in fail: 

681 # These should raise 

682 with self.assertRaises(FileNotFoundError): 

683 # non-existing file 

684 datastore.get(ref) 

685 with self.assertRaises(FileNotFoundError): 

686 datastore.getURI(ref) 

687 

688 def testNestedTransaction(self) -> None: 

689 datastore = self.makeDatastore() 

690 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

691 dimensions = self.universe.conform(("visit", "physical_filter")) 

692 metrics = makeExampleMetrics() 

693 

694 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

695 refBefore = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

696 datastore.put(metrics, refBefore) 

697 with self.assertRaises(TransactionTestError): 

698 with datastore.transaction(): 

699 dataId = { 

700 "instrument": "dummy", 

701 "visit": 1, 

702 "physical_filter": "V", 

703 "band": "v", 

704 "day_obs": 20250101, 

705 } 

706 refOuter = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

707 datastore.put(metrics, refOuter) 

708 with datastore.transaction(): 

709 dataId = { 

710 "instrument": "dummy", 

711 "visit": 2, 

712 "physical_filter": "V", 

713 "band": "v", 

714 "day_obs": 20250101, 

715 } 

716 refInner = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

717 datastore.put(metrics, refInner) 

718 # All datasets should exist 

719 for ref in (refBefore, refOuter, refInner): 

720 metricsOut = datastore.get(ref, parameters=None) 

721 self.assertEqual(metrics, metricsOut) 

722 raise TransactionTestError("This should roll back the transaction") 

723 # Dataset(s) inserted before the transaction should still exist 

724 metricsOut = datastore.get(refBefore, parameters=None) 

725 self.assertEqual(metrics, metricsOut) 

726 # But all datasets inserted during the (rolled back) transaction 

727 # should be gone 

728 with self.assertRaises(FileNotFoundError): 

729 datastore.get(refOuter) 

730 with self.assertRaises(FileNotFoundError): 

731 datastore.get(refInner) 

732 

733 def _prepareIngestTest(self) -> tuple[MetricsExample, DatasetRef]: 

734 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

735 dimensions = self.universe.conform(("visit", "physical_filter")) 

736 metrics = makeExampleMetrics() 

737 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

738 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

739 return metrics, ref 

740 

741 def runIngestTest(self, func: Callable[[MetricsExample, str, DatasetRef], None]) -> None: 

742 metrics, ref = self._prepareIngestTest() 

743 # The file will be deleted after the test. 

744 # For symlink tests this leads to a situation where the datastore 

745 # points to a file that does not exist. This will make os.path.exist 

746 # return False but then the new symlink will fail with 

747 # FileExistsError later in the code so the test still passes. 

748 with _temp_yaml_file(metrics._asdict()) as path: 

749 func(metrics, path, ref) 

750 

751 def testIngestNoTransfer(self) -> None: 

752 """Test ingesting existing files with no transfer.""" 

753 for mode in (None, "auto"): 

754 # Some datastores have auto but can't do in place transfer 

755 if mode == "auto" and "auto" in self.ingestTransferModes and not self.canIngestNoTransferAuto: 

756 continue 

757 

758 with self.subTest(mode=mode): 

759 datastore = self.makeDatastore() 

760 

761 def succeed( 

762 obj: MetricsExample, 

763 path: str, 

764 ref: DatasetRef, 

765 mode: str | None = mode, 

766 datastore: Datastore = datastore, 

767 ) -> None: 

768 """Ingest a file already in the datastore root.""" 

769 # first move it into the root, and adjust the path 

770 # accordingly. 

771 # In the case of a ChainedDatastore, we have multiple 

772 # roots, all of which will accept the file, so we 

773 # have to copy it into all the roots. 

774 relative_path = None 

775 for root in datastore.roots.values(): 

776 if root is not None: 

777 copied_path = shutil.copy(path, root.ospath) 

778 relative_path = os.path.relpath(copied_path, start=root.ospath) 

779 assert relative_path is not None, ( 

780 "Running a FileDatastore test on a Datastore instance without any roots" 

781 ) 

782 datastore.ingest(FileDataset(path=relative_path, refs=ref), transfer=mode) 

783 self.assertEqual(obj, datastore.get(ref)) 

784 

785 def failInputDoesNotExist( 

786 obj: MetricsExample, 

787 path: str, 

788 ref: DatasetRef, 

789 mode: str | None = mode, 

790 datastore: Datastore = datastore, 

791 ) -> None: 

792 """Can't ingest files if we're given a bad path.""" 

793 with self.assertRaises(FileNotFoundError): 

794 datastore.ingest( 

795 FileDataset(path="this-file-does-not-exist.yaml", refs=ref), transfer=mode 

796 ) 

797 self.assertFalse(datastore.exists(ref)) 

798 

799 def failOutsideRoot( 

800 obj: MetricsExample, 

801 path: str, 

802 ref: DatasetRef, 

803 mode: str | None = mode, 

804 datastore: Datastore = datastore, 

805 ) -> None: 

806 """Can't ingest files outside of datastore root unless 

807 auto. 

808 """ 

809 if mode == "auto": 

810 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

811 self.assertTrue(datastore.exists(ref)) 

812 else: 

813 with self.assertRaises(RuntimeError): 

814 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

815 self.assertFalse(datastore.exists(ref)) 

816 

817 def failNotImplemented( 

818 obj: MetricsExample, 

819 path: str, 

820 ref: DatasetRef, 

821 mode: str | None = mode, 

822 datastore: Datastore = datastore, 

823 ) -> None: 

824 with self.assertRaises(NotImplementedError): 

825 datastore.ingest(FileDataset(path=path, refs=ref), transfer=mode) 

826 

827 if mode in self.ingestTransferModes: 

828 self.runIngestTest(failOutsideRoot) 

829 self.runIngestTest(failInputDoesNotExist) 

830 self.runIngestTest(succeed) 

831 else: 

832 self.runIngestTest(failNotImplemented) 

833 

834 def testIngestTransfer(self) -> None: 

835 """Test ingesting existing files after transferring them.""" 

836 for mode in ("copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto"): 

837 with self.subTest(mode=mode): 

838 datastore = self.makeDatastore(mode) 

839 

840 def succeed( 

841 obj: MetricsExample, 

842 path: str, 

843 ref: DatasetRef, 

844 mode: str | None = mode, 

845 datastore: Datastore = datastore, 

846 ) -> None: 

847 """Ingest a file by transferring it to the template 

848 location. 

849 """ 

850 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

851 self.assertEqual(obj, datastore.get(ref)) 

852 file_exists = os.path.exists(path) 

853 if mode == "move": 

854 self.assertFalse(file_exists) 

855 else: 

856 self.assertTrue(file_exists) 

857 

858 def failInputDoesNotExist( 

859 obj: MetricsExample, 

860 path: str, 

861 ref: DatasetRef, 

862 mode: str | None = mode, 

863 datastore: Datastore = datastore, 

864 ) -> None: 

865 """Can't ingest files if we're given a bad path.""" 

866 with self.assertRaises(FileNotFoundError): 

867 # Ensure the file does not look like it is in 

868 # datastore for auto mode 

869 datastore.ingest( 

870 FileDataset(path="../this-file-does-not-exist.yaml", refs=ref), transfer=mode 

871 ) 

872 self.assertFalse(datastore.exists(ref), f"Checking not in datastore using mode {mode}") 

873 

874 def failNotImplemented( 

875 obj: MetricsExample, 

876 path: str, 

877 ref: DatasetRef, 

878 mode: str | None = mode, 

879 datastore: Datastore = datastore, 

880 ) -> None: 

881 with self.assertRaises(NotImplementedError): 

882 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode) 

883 

884 if mode in self.ingestTransferModes: 

885 self.runIngestTest(failInputDoesNotExist) 

886 self.runIngestTest(succeed) 

887 else: 

888 self.runIngestTest(failNotImplemented) 

889 

890 def testIngestSymlinkOfSymlink(self) -> None: 

891 """Special test for symlink to a symlink ingest""" 

892 metrics, ref = self._prepareIngestTest() 

893 # The aim of this test is to create a dataset on disk, then 

894 # create a symlink to it and finally ingest the symlink such that 

895 # the symlink in the datastore points to the original dataset. 

896 for mode in ("symlink", "relsymlink"): 

897 if mode not in self.ingestTransferModes: 

898 continue 

899 

900 print(f"Trying mode {mode}") 

901 with _temp_yaml_file(metrics._asdict()) as realpath: 

902 with tempfile.TemporaryDirectory() as tmpdir: 

903 sympath = os.path.join(tmpdir, "symlink.yaml") 

904 os.symlink(os.path.realpath(realpath), sympath) 

905 

906 datastore = self.makeDatastore() 

907 datastore.ingest(FileDataset(path=os.path.abspath(sympath), refs=ref), transfer=mode) 

908 

909 uri = datastore.getURI(ref) 

910 self.assertTrue(uri.isLocal, f"Check {uri.scheme}") 

911 self.assertTrue(os.path.islink(uri.ospath), f"Check {uri} is a symlink") 

912 

913 linkTarget = os.readlink(uri.ospath) 

914 if mode == "relsymlink": 

915 self.assertFalse(os.path.isabs(linkTarget)) 

916 else: 

917 self.assertTrue(os.path.samefile(linkTarget, realpath)) 

918 

919 # Check that we can get the dataset back regardless of mode 

920 metric2 = datastore.get(ref) 

921 self.assertEqual(metric2, metrics) 

922 

923 # Cleanup the file for next time round loop 

924 # since it will get the same file name in store 

925 datastore.remove(ref) 

926 

927 def _populate_export_datastore(self, name: str) -> tuple[Datastore, list[DatasetRef]]: 

928 datastore = self.makeDatastore(name) 

929 

930 # For now only the FileDatastore can be used for this test. 

931 # ChainedDatastore that only includes InMemoryDatastores have to be 

932 # skipped as well. 

933 for name in datastore.names: 

934 if not name.startswith("InMemoryDatastore"): 

935 break 

936 else: 

937 raise unittest.SkipTest("in-memory datastore does not support record export/import") 

938 

939 metrics = makeExampleMetrics() 

940 dimensions = self.universe.conform(("visit", "physical_filter")) 

941 sc = self.storageClassFactory.getStorageClass("StructuredData") 

942 

943 refs = [] 

944 for visit in (2048, 2049, 2050): 

945 dataId = { 

946 "instrument": "dummy", 

947 "visit": visit, 

948 "physical_filter": "Uprime", 

949 "band": "u", 

950 "day_obs": 20250101, 

951 } 

952 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

953 datastore.put(metrics, ref) 

954 refs.append(ref) 

955 return datastore, refs 

956 

957 def testExportImportRecords(self) -> None: 

958 """Test for export_records and import_records methods.""" 

959 datastore, refs = self._populate_export_datastore("test_datastore") 

960 for exported_refs in (refs, refs[1:]): 

961 n_refs = len(exported_refs) 

962 records = datastore.export_records(exported_refs) 

963 self.assertGreater(len(records), 0) 

964 self.assertTrue(set(records.keys()) <= set(datastore.names)) 

965 # In a ChainedDatastore each FileDatastore will have a complete set 

966 for datastore_name in records: 

967 record_data = records[datastore_name] 

968 self.assertEqual(len(record_data.records), n_refs) 

969 

970 # Check that subsetting works, include non-existing dataset ID. 

971 dataset_ids = {exported_refs[0].id, uuid.uuid4()} 

972 subset = record_data.subset(dataset_ids) 

973 assert subset is not None 

974 self.assertEqual(len(subset.records), 1) 

975 subset = record_data.subset({uuid.uuid4()}) 

976 self.assertIsNone(subset) 

977 

978 # Use the same datastore name to import relative path. 

979 datastore2 = self.makeDatastore("test_datastore") 

980 

981 records = datastore.export_records(refs[1:]) 

982 datastore2.import_records(records) 

983 

984 with self.assertRaises(FileNotFoundError): 

985 data = datastore2.get(refs[0]) 

986 data = datastore2.get(refs[1]) 

987 self.assertIsNotNone(data) 

988 data = datastore2.get(refs[2]) 

989 self.assertIsNotNone(data) 

990 

991 def testExportPredictedRecords(self): 

992 if self.isEphemeral: 

993 raise unittest.SkipTest("in-memory datastore does not support record export/import") 

994 sc = self.storageClassFactory.getStorageClass("ThingOne") 

995 dimensions = self.universe.conform(("visit", "physical_filter")) 

996 dataId = { 

997 "instrument": "dummy", 

998 "visit": 52, 

999 "physical_filter": "V", 

1000 "band": "v", 

1001 "day_obs": 20250101, 

1002 } 

1003 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

1004 

1005 datastore = self.makeDatastore("test_datastore") 

1006 names = {n for n in datastore.names if not n.startswith("InMemory")} 

1007 records = datastore.export_predicted_records([ref]) 

1008 

1009 # Expect predicted records from all datastores. 

1010 self.assertEqual(set(records.keys()), names) 

1011 

1012 for record_data in records.values(): 

1013 self.assertEqual(len(record_data.records), 1) 

1014 

1015 def testExport(self) -> None: 

1016 datastore, refs = self._populate_export_datastore("test_datastore") 

1017 

1018 datasets = list(datastore.export(refs)) 

1019 self.assertEqual(len(datasets), 3) 

1020 

1021 for transfer in (None, "auto"): 

1022 # Both will default to None 

1023 datasets = list(datastore.export(refs, transfer=transfer)) 

1024 self.assertEqual(len(datasets), 3) 

1025 

1026 with self.assertRaises(TypeError): 

1027 list(datastore.export(refs, transfer="copy")) 

1028 

1029 with self.assertRaises(TypeError): 

1030 list(datastore.export(refs, directory="exportDir", transfer="move")) 

1031 

1032 # Create a new ref that is not known to the datastore and try to 

1033 # export it. 

1034 sc = self.storageClassFactory.getStorageClass("ThingOne") 

1035 dimensions = self.universe.conform(("visit", "physical_filter")) 

1036 dataId = { 

1037 "instrument": "dummy", 

1038 "visit": 52, 

1039 "physical_filter": "V", 

1040 "band": "v", 

1041 "day_obs": 20250101, 

1042 } 

1043 ref = self.makeDatasetRef("metric", dimensions, sc, dataId) 

1044 with self.assertRaises(FileNotFoundError): 

1045 list(datastore.export(refs + [ref], transfer=None)) 

1046 

1047 def test_pydantic_dict_storage_class_conversions(self) -> None: 

1048 """Test converting a dataset stored as a pydantic model into a dict on 

1049 read. 

1050 """ 

1051 datastore = self.makeDatastore() 

1052 store_as_model = self.makeDatasetRef( 

1053 "store_as_model", 

1054 dimensions=self.universe.empty, 

1055 storageClass="DictConvertibleModel", 

1056 dataId=DataCoordinate.make_empty(self.universe), 

1057 ) 

1058 content = {"a": "one", "b": "two"} 

1059 model = DictConvertibleModel.from_dict(content, extra="original content") 

1060 datastore.put(model, store_as_model) 

1061 retrieved_model = datastore.get(store_as_model) 

1062 self.assertEqual(retrieved_model, model) 

1063 loaded = datastore.get(store_as_model.overrideStorageClass("NativeDictForConvertibleModel")) 

1064 self.assertEqual(type(loaded), dict) 

1065 self.assertEqual(loaded, content) 

1066 

1067 def test_simple_class_put_get(self) -> None: 

1068 """Test that we can put and get a simple class with dict() 

1069 constructor. 

1070 """ 

1071 datastore = self.makeDatastore() 

1072 data = MetricsExample(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1073 self._assert_different_puts(datastore, "MetricsExample", data) 

1074 

1075 def test_dataclass_put_get(self) -> None: 

1076 """Test that we can put and get a simple dataclass.""" 

1077 datastore = self.makeDatastore() 

1078 data = MetricsExampleDataclass(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1079 self._assert_different_puts(datastore, "MetricsExampleDataclass", data) 

1080 

1081 def test_pydantic_put_get(self) -> None: 

1082 """Test that we can put and get a simple Pydantic model.""" 

1083 datastore = self.makeDatastore() 

1084 data = MetricsExampleModel(summary={"a": 1}, data=[1, 2, 3], output={"b": 2}) 

1085 self._assert_different_puts(datastore, "MetricsExampleModel", data) 

1086 

1087 def test_tuple_put_get(self) -> None: 

1088 """Test that we can put and get a tuple.""" 

1089 datastore = self.makeDatastore() 

1090 data = ("a", "b", 1) 

1091 self._assert_different_puts(datastore, "TupleExample", data) 

1092 

1093 def _assert_different_puts(self, datastore: Datastore, storageClass_root: str, data: Any) -> None: 

1094 refs = { 

1095 x: self.makeDatasetRef( 

1096 f"stora_as_{x}", 

1097 dimensions=self.universe.empty, 

1098 storageClass=f"{storageClass_root}{x}", 

1099 dataId=DataCoordinate.make_empty(self.universe), 

1100 ) 

1101 for x in ["A", "B"] 

1102 } 

1103 

1104 for ref in refs.values(): 

1105 datastore.put(data, ref) 

1106 

1107 self.assertEqual(datastore.get(refs["A"]), datastore.get(refs["B"])) 

1108 

1109 

1110class PosixDatastoreTestCase(DatastoreTests, unittest.TestCase): 

1111 """PosixDatastore specialization""" 

1112 

1113 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1114 uriScheme = "file" 

1115 canIngestNoTransferAuto = True 

1116 ingestTransferModes = (None, "copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto") 

1117 isEphemeral = False 

1118 rootKeys = ("root",) 

1119 validationCanFail = True 

1120 

1121 def setUp(self) -> None: 

1122 # The call to os.path.realpath is necessary because Mac temporary files 

1123 # can end up in either /private/var/folders or /var/folders, which 

1124 # refer to the same location but don't appear to. 

1125 # This matters for "relsymlink" transfer mode, because it needs to be 

1126 # able to read the file through a relative symlink, but some of the 

1127 # intermediate directories are not traversable if you try to get from a 

1128 # tempfile in /var/folders to one in /private/var/folders via a 

1129 # relative path. 

1130 self.root = os.path.realpath(self.enterContext(tempfile.TemporaryDirectory())) 

1131 super().setUp() 

1132 

1133 def testAtomicWrite(self) -> None: 

1134 """Test that we write to a temporary and then rename""" 

1135 datastore = self.makeDatastore() 

1136 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1137 dimensions = self.universe.conform(("visit", "physical_filter")) 

1138 metrics = makeExampleMetrics() 

1139 

1140 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

1141 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1142 

1143 with self.assertLogs("lsst.resources", "DEBUG") as cm: 

1144 datastore.put(metrics, ref) 

1145 move_logs = [ll for ll in cm.output if "transfer=" in ll] 

1146 self.assertIn("transfer=move", move_logs[0]) 

1147 

1148 # And the transfer should be file to file. 

1149 self.assertEqual(move_logs[0].count("file://"), 2) 

1150 

1151 def testCanNotDeterminePutFormatterLocation(self) -> None: 

1152 """Verify that the expected exception is raised if the FileDatastore 

1153 can not determine the put formatter location. 

1154 """ 

1155 _ = makeExampleMetrics() 

1156 datastore = self.makeDatastore() 

1157 

1158 # Create multiple storage classes for testing different formulations 

1159 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1160 

1161 sccomp = StorageClass("Dummy") 

1162 compositeStorageClass = StorageClass( 

1163 "StructuredComposite", components={"dummy": sccomp, "dummy2": sccomp} 

1164 ) 

1165 

1166 dimensions = self.universe.conform(("visit", "physical_filter")) 

1167 dataId = { 

1168 "instrument": "dummy", 

1169 "visit": 52, 

1170 "physical_filter": "V", 

1171 "band": "v", 

1172 "day_obs": 20250101, 

1173 } 

1174 

1175 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1176 compRef = self.makeDatasetRef("metric", dimensions, compositeStorageClass, dataId) 

1177 

1178 def raiser(ref: DatasetRef) -> None: 

1179 raise DatasetTypeNotSupportedError() 

1180 

1181 with unittest.mock.patch.object( 

1182 lsst.daf.butler.datastores.fileDatastore.FileDatastore, 

1183 "_determine_put_formatter_location", 

1184 side_effect=raiser, 

1185 ): 

1186 # verify the non-composite ref execution path: 

1187 with self.assertRaises(DatasetTypeNotSupportedError): 

1188 datastore.getURIs(ref, predict=True) 

1189 

1190 # verify the composite-ref execution path: 

1191 with self.assertRaises(DatasetTypeNotSupportedError): 

1192 datastore.getURIs(compRef, predict=True) 

1193 

1194 def test_roots(self): 

1195 datastore = self.makeDatastore() 

1196 

1197 self.assertEqual(set(datastore.names), set(datastore.roots.keys())) 

1198 for root in datastore.roots.values(): 

1199 if root is not None: 

1200 self.assertTrue(root.exists()) 

1201 

1202 def test_prepare_get_for_external_client(self): 

1203 datastore = self.makeDatastore() 

1204 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1205 dimensions = self.universe.conform(("visit", "physical_filter")) 

1206 dataId = {"instrument": "dummy", "visit": 52, "physical_filter": "V", "band": "v"} 

1207 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1208 # Most of the coverage for this function is in test_server.py, 

1209 # because it requires a file backend that supports URL signing. 

1210 self.assertIsNone(datastore.prepare_get_for_external_client(ref)) 

1211 

1212 

1213class PosixDatastoreNoChecksumsTestCase(PosixDatastoreTestCase): 

1214 """Posix datastore tests but with checksums disabled.""" 

1215 

1216 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreNoChecksums.yaml") 

1217 

1218 def testChecksum(self) -> None: 

1219 """Ensure that checksums have not been calculated.""" 

1220 datastore = self.makeDatastore() 

1221 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1222 dimensions = self.universe.conform(("visit", "physical_filter")) 

1223 metrics = makeExampleMetrics() 

1224 

1225 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101} 

1226 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1227 

1228 # Configuration should have disabled checksum calculation 

1229 datastore.put(metrics, ref) 

1230 infos = datastore.getStoredItemsInfo(ref) 

1231 self.assertIsNone(infos[0].checksum) 

1232 

1233 # Remove put back but with checksums enabled explicitly 

1234 datastore.remove(ref) 

1235 datastore.useChecksum = True 

1236 datastore.put(metrics, ref) 

1237 

1238 infos = datastore.getStoredItemsInfo(ref) 

1239 self.assertIsNotNone(infos[0].checksum) 

1240 

1241 def test_repeat_ingest(self): 

1242 """Test that repeatedly ingesting the same file in direct mode 

1243 is allowed. 

1244 

1245 Test can only run with FileDatastore since that is the only one 

1246 supporting "direct" ingest. 

1247 """ 

1248 metrics, v4ref = self._prepareIngestTest() 

1249 datastore = self.makeDatastore() 

1250 v5ref = DatasetRef( 

1251 v4ref.datasetType, v4ref.dataId, v4ref.run, id_generation_mode=DatasetIdGenEnum.DATAID_TYPE_RUN 

1252 ) 

1253 

1254 with _temp_yaml_file(metrics._asdict()) as path: 

1255 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct") 

1256 

1257 # This will fail because the ref is using UUIDv4. 

1258 with self.assertRaises(RuntimeError): 

1259 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct") 

1260 

1261 # UUIDv5 can be repeatedly ingested in direct mode. 

1262 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct") 

1263 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct") 

1264 

1265 with self.assertRaises(RuntimeError): 

1266 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="copy") 

1267 

1268 

1269class TrashDatastoreTestCase(PosixDatastoreTestCase): 

1270 """Restrict trash test to FileDatastore.""" 

1271 

1272 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1273 

1274 def testTrash(self) -> None: 

1275 datastore, *refs = self.prepDeleteTest(n_refs=10) 

1276 

1277 # Trash one of them. 

1278 ref = refs.pop() 

1279 uri = datastore.getURI(ref) 

1280 datastore.trash(ref) 

1281 self.assertTrue(uri.exists(), uri) # Not deleted yet 

1282 datastore.emptyTrash() 

1283 self.assertFalse(uri.exists(), uri) 

1284 

1285 # Trash it again should be fine. 

1286 datastore.trash(ref) 

1287 

1288 # Trash multiple items at once. 

1289 subset = [refs.pop(), refs.pop()] 

1290 datastore.trash(subset) 

1291 datastore.emptyTrash() 

1292 

1293 # Remove a record and trash should do nothing. 

1294 # This is execution butler scenario. 

1295 ref = refs.pop() 

1296 uri = datastore.getURI(ref) 

1297 datastore._table.delete(["dataset_id"], {"dataset_id": ref.id}) 

1298 self.assertTrue(uri.exists()) 

1299 datastore.trash(ref) 

1300 datastore.emptyTrash() 

1301 self.assertTrue(uri.exists()) 

1302 

1303 # Switch on trust and it should delete the file. 

1304 datastore.trustGetRequest = True 

1305 datastore.trash([ref]) 

1306 self.assertFalse(uri.exists()) 

1307 

1308 # Remove multiples at once in trust mode. 

1309 subset = [refs.pop() for i in range(3)] 

1310 datastore.trash(subset) 

1311 datastore.trash(refs.pop()) # Check that a single ref can trash 

1312 

1313 def test_empty_trash(self) -> None: 

1314 """Test parameters and return value for empty trash.""" 

1315 datastore, *refs = self.prepDeleteTest(n_refs=10) 

1316 

1317 # Trash one of them. 

1318 ref = refs.pop() 

1319 uri = datastore.getURI(ref) 

1320 datastore.trash(ref) 

1321 self.assertTrue(uri.exists(), uri) # Not deleted yet 

1322 

1323 # Empty trash but with a list of refs that does not include the 

1324 # one in the trash table. 

1325 removed = datastore.emptyTrash(refs=refs) 

1326 self.assertEqual(len(removed), 0) 

1327 self.assertTrue(uri.exists(), uri) 

1328 

1329 # Empty the entire trash but in dry_run mode. 

1330 removed = datastore.emptyTrash(dry_run=True) 

1331 self.assertEqual(len(removed), 1) 

1332 self.assertEqual(removed.pop(), uri) 

1333 self.assertTrue(uri.exists(), uri) 

1334 

1335 # Empty the trash specifying the actual ref. 

1336 removed = datastore.emptyTrash(refs=[ref]) 

1337 self.assertEqual(len(removed), 1) 

1338 self.assertEqual(removed.pop(), uri) 

1339 self.assertFalse(uri.exists(), uri) 

1340 

1341 # Trash everything and empty. 

1342 datastore.trash(refs) 

1343 removed = datastore.emptyTrash(dry_run=True) 

1344 for u in removed: 

1345 self.assertTrue(u.exists()) 

1346 removed = datastore.emptyTrash() 

1347 for u in removed: 

1348 self.assertFalse(u.exists()) 

1349 

1350 

1351class CleanupPosixDatastoreTestCase(DatastoreTestsBase, unittest.TestCase): 

1352 """Test datastore cleans up on failure.""" 

1353 

1354 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

1355 

1356 def setUp(self) -> None: 

1357 # Override the working directory before calling the base class 

1358 self.root = tempfile.mkdtemp() 

1359 super().setUp() 

1360 

1361 def testCleanup(self) -> None: 

1362 """Test that a failed formatter write does cleanup a partial file.""" 

1363 metrics = makeExampleMetrics() 

1364 datastore = self.makeDatastore() 

1365 

1366 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1367 

1368 dimensions = self.universe.conform(("visit", "physical_filter")) 

1369 dataId = { 

1370 "instrument": "dummy", 

1371 "visit": 52, 

1372 "physical_filter": "V", 

1373 "band": "v", 

1374 "day_obs": 20250101, 

1375 } 

1376 

1377 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId) 

1378 

1379 # Determine where the file will end up (we assume Formatters use 

1380 # the same file extension) 

1381 expectedUri = datastore.getURI(ref, predict=True) 

1382 self.assertEqual(expectedUri.fragment, "predicted") 

1383 

1384 self.assertEqual(expectedUri.getExtension(), ".yaml", f"Is there a file extension in {expectedUri}") 

1385 

1386 # Try formatter that fails and formatter that fails and leaves 

1387 # a file behind 

1388 for formatter in (BadWriteFormatter, BadNoWriteFormatter): 

1389 with self.subTest(formatter=str(formatter)): 

1390 # Monkey patch the formatter 

1391 datastore.formatterFactory.registerFormatter(ref.datasetType, formatter, overwrite=True) 

1392 

1393 # Try to put the dataset, it should fail 

1394 with self.assertRaises(RuntimeError): 

1395 datastore.put(metrics, ref) 

1396 

1397 # Check that there is no file on disk 

1398 self.assertFalse(expectedUri.exists(), f"Check for existence of {expectedUri}") 

1399 

1400 # Check that there is a directory 

1401 dir = expectedUri.dirname() 

1402 self.assertTrue(dir.exists(), f"Check for existence of directory {dir}") 

1403 

1404 # Force YamlFormatter and check that this time a file is written 

1405 datastore.formatterFactory.registerFormatter(ref.datasetType, YamlFormatter, overwrite=True) 

1406 datastore.put(metrics, ref) 

1407 self.assertTrue(expectedUri.exists(), f"Check for existence of {expectedUri}") 

1408 datastore.remove(ref) 

1409 self.assertFalse(expectedUri.exists(), f"Check for existence of now removed {expectedUri}") 

1410 

1411 

1412class InMemoryDatastoreTestCase(DatastoreTests, unittest.TestCase): 

1413 """PosixDatastore specialization""" 

1414 

1415 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastore.yaml") 

1416 uriScheme = "mem" 

1417 hasUnsupportedPut = False 

1418 ingestTransferModes = () 

1419 isEphemeral = True 

1420 rootKeys = None 

1421 validationCanFail = False 

1422 

1423 

1424class ChainedDatastoreTestCase(PosixDatastoreTestCase): 

1425 """ChainedDatastore specialization using a POSIXDatastore""" 

1426 

1427 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore.yaml") 

1428 hasUnsupportedPut = False 

1429 canIngestNoTransferAuto = False 

1430 ingestTransferModes = (None, "copy", "move", "hardlink", "symlink", "relsymlink", "link", "auto") 

1431 isEphemeral = False 

1432 rootKeys = (".datastores.1.root", ".datastores.2.root") 

1433 validationCanFail = True 

1434 

1435 

1436class ChainedDatastoreMemoryTestCase(InMemoryDatastoreTestCase): 

1437 """ChainedDatastore specialization using all InMemoryDatastore""" 

1438 

1439 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2.yaml") 

1440 validationCanFail = False 

1441 isEphemeral = True 

1442 

1443 

1444class DatastoreConstraintsTests(DatastoreTestsBase): 

1445 """Basic tests of constraints model of Datastores.""" 

1446 

1447 def testConstraints(self) -> None: 

1448 """Test constraints model. Assumes that each test class has the 

1449 same constraints. 

1450 """ 

1451 metrics = makeExampleMetrics() 

1452 datastore = self.makeDatastore() 

1453 

1454 sc1 = self.storageClassFactory.getStorageClass("StructuredData") 

1455 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson") 

1456 dimensions = self.universe.conform(("visit", "physical_filter", "instrument")) 

1457 dataId = { 

1458 "visit": 52, 

1459 "physical_filter": "V", 

1460 "band": "v", 

1461 "instrument": "DummyCamComp", 

1462 "day_obs": 20250101, 

1463 } 

1464 

1465 # Write empty file suitable for ingest check (JSON and YAML variants) 

1466 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml") 

1467 testfile_j = tempfile.NamedTemporaryFile(suffix=".json") 

1468 for datasetTypeName, sc, accepted in ( 

1469 ("metric", sc1, True), 

1470 ("metric5", sc1, False), 

1471 ("metric33", sc1, True), 

1472 ("metric5", sc2, True), 

1473 ): 

1474 # Choose different temp file depending on StorageClass 

1475 testfile = testfile_j if sc.name.endswith("Json") else testfile_y 

1476 

1477 with self.subTest(datasetTypeName=datasetTypeName, storageClass=sc.name, file=testfile.name): 

1478 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId) 

1479 if accepted: 

1480 datastore.put(metrics, ref) 

1481 self.assertTrue(datastore.exists(ref)) 

1482 datastore.remove(ref) 

1483 

1484 # Try ingest 

1485 if self.canIngest: 

1486 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1487 self.assertTrue(datastore.exists(ref)) 

1488 datastore.remove(ref) 

1489 else: 

1490 with self.assertRaises(DatasetTypeNotSupportedError): 

1491 datastore.put(metrics, ref) 

1492 self.assertFalse(datastore.exists(ref)) 

1493 

1494 # Again with ingest 

1495 if self.canIngest: 

1496 with self.assertRaises(DatasetTypeNotSupportedError): 

1497 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1498 self.assertFalse(datastore.exists(ref)) 

1499 

1500 

1501class PosixDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase): 

1502 """PosixDatastore specialization""" 

1503 

1504 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreP.yaml") 

1505 canIngest = True 

1506 

1507 def setUp(self) -> None: 

1508 # Override the working directory before calling the base class 

1509 self.root = tempfile.mkdtemp() 

1510 super().setUp() 

1511 

1512 

1513class InMemoryDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase): 

1514 """InMemoryDatastore specialization.""" 

1515 

1516 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastoreP.yaml") 

1517 canIngest = False 

1518 

1519 

1520class ChainedDatastoreConstraintsNativeTestCase(PosixDatastoreConstraintsTestCase): 

1521 """ChainedDatastore specialization using a POSIXDatastore and constraints 

1522 at the ChainedDatstore. 

1523 """ 

1524 

1525 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePa.yaml") 

1526 

1527 

1528class ChainedDatastoreConstraintsTestCase(PosixDatastoreConstraintsTestCase): 

1529 """ChainedDatastore specialization using a POSIXDatastore.""" 

1530 

1531 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastoreP.yaml") 

1532 

1533 

1534class ChainedDatastoreMemoryConstraintsTestCase(InMemoryDatastoreConstraintsTestCase): 

1535 """ChainedDatastore specialization using all InMemoryDatastore.""" 

1536 

1537 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2P.yaml") 

1538 canIngest = False 

1539 

1540 

1541class ChainedDatastorePerStoreConstraintsTests(DatastoreTestsBase, unittest.TestCase): 

1542 """Test that a chained datastore can control constraints per-datastore 

1543 even if child datastore would accept. 

1544 """ 

1545 

1546 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePb.yaml") 

1547 

1548 def setUp(self) -> None: 

1549 # Override the working directory before calling the base class 

1550 self.root = tempfile.mkdtemp() 

1551 super().setUp() 

1552 

1553 def testConstraints(self) -> None: 

1554 """Test chained datastore constraints model.""" 

1555 metrics = makeExampleMetrics() 

1556 datastore = self.makeDatastore() 

1557 

1558 sc1 = self.storageClassFactory.getStorageClass("StructuredData") 

1559 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson") 

1560 dimensions = self.universe.conform(("visit", "physical_filter", "instrument")) 

1561 dataId1 = { 

1562 "visit": 52, 

1563 "physical_filter": "V", 

1564 "band": "v", 

1565 "instrument": "DummyCamComp", 

1566 "day_obs": 20250101, 

1567 } 

1568 dataId2 = {"visit": 52, "physical_filter": "V", "band": "v", "instrument": "HSC", "day_obs": 20250101} 

1569 

1570 # Write empty file suitable for ingest check (JSON and YAML variants) 

1571 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml") 

1572 testfile_j = tempfile.NamedTemporaryFile(suffix=".json") 

1573 

1574 for typeName, dataId, sc, accept, ingest in ( 

1575 ("metric", dataId1, sc1, (False, True, False), True), 

1576 ("metric5", dataId1, sc1, (False, False, False), False), 

1577 ("metric5", dataId2, sc1, (True, False, False), False), 

1578 ("metric33", dataId2, sc2, (True, True, False), True), 

1579 ("metric5", dataId1, sc2, (False, True, False), True), 

1580 ): 

1581 # Choose different temp file depending on StorageClass 

1582 testfile = testfile_j if sc.name.endswith("Json") else testfile_y 

1583 

1584 with self.subTest(datasetTypeName=typeName, dataId=dataId, sc=sc.name): 

1585 ref = self.makeDatasetRef(typeName, dimensions, sc, dataId) 

1586 if any(accept): 

1587 datastore.put(metrics, ref) 

1588 self.assertTrue(datastore.exists(ref)) 

1589 

1590 # Check each datastore inside the chained datastore 

1591 for childDatastore, expected in zip(datastore.datastores, accept, strict=True): 

1592 self.assertEqual( 

1593 childDatastore.exists(ref), 

1594 expected, 

1595 f"Testing presence of {ref} in datastore {childDatastore.name}", 

1596 ) 

1597 

1598 datastore.remove(ref) 

1599 

1600 # Check that ingest works 

1601 if ingest: 

1602 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1603 self.assertTrue(datastore.exists(ref)) 

1604 

1605 # Check each datastore inside the chained datastore 

1606 for childDatastore, expected in zip(datastore.datastores, accept, strict=True): 

1607 # Ephemeral datastores means InMemory at the moment 

1608 # and that does not accept ingest of files. 

1609 if childDatastore.isEphemeral: 

1610 expected = False 

1611 self.assertEqual( 

1612 childDatastore.exists(ref), 

1613 expected, 

1614 f"Testing presence of ingested {ref} in datastore {childDatastore.name}", 

1615 ) 

1616 

1617 datastore.remove(ref) 

1618 else: 

1619 with self.assertRaises(DatasetTypeNotSupportedError): 

1620 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1621 

1622 else: 

1623 with self.assertRaises(DatasetTypeNotSupportedError): 

1624 datastore.put(metrics, ref) 

1625 self.assertFalse(datastore.exists(ref)) 

1626 

1627 # Again with ingest 

1628 with self.assertRaises(DatasetTypeNotSupportedError): 

1629 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link") 

1630 self.assertFalse(datastore.exists(ref)) 

1631 

1632 

1633@unittest.mock.patch.dict(os.environ, {}, clear=True) 

1634class DatastoreCacheTestCase(DatasetTestHelper, unittest.TestCase): 

1635 """Tests for datastore caching infrastructure.""" 

1636 

1637 @classmethod 

1638 def setUpClass(cls) -> None: 

1639 cls.storageClassFactory = StorageClassFactory() 

1640 cls.universe = DimensionUniverse() 

1641 

1642 # Ensure that we load the test storage class definitions. 

1643 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml") 

1644 cls.storageClassFactory.addFromConfig(scConfigFile) 

1645 

1646 def setUp(self) -> None: 

1647 self.id = 0 

1648 

1649 # Create a root that we can use for caching tests. 

1650 self.root = tempfile.mkdtemp() 

1651 

1652 # Create some test dataset refs and associated test files 

1653 sc = self.storageClassFactory.getStorageClass("StructuredDataDict") 

1654 dimensions = self.universe.conform(("visit", "physical_filter")) 

1655 dataId = { 

1656 "instrument": "dummy", 

1657 "visit": 52, 

1658 "physical_filter": "V", 

1659 "band": "v", 

1660 "day_obs": 20250101, 

1661 } 

1662 

1663 # Create list of refs and list of temporary files 

1664 n_datasets = 10 

1665 self.refs = [self.makeDatasetRef(f"metric{n}", dimensions, sc, dataId) for n in range(n_datasets)] 

1666 

1667 root_uri = ResourcePath(self.root, forceDirectory=True) 

1668 self.files = [root_uri.join(f"file{n}.txt") for n in range(n_datasets)] 

1669 

1670 # Create test files. 

1671 for uri in self.files: 

1672 uri.write(b"0123456789") 

1673 

1674 # Create some composite refs with component files. 

1675 sc = self.storageClassFactory.getStorageClass("StructuredData") 

1676 self.composite_refs = [self.makeDatasetRef(f"composite{n}", dimensions, sc, dataId) for n in range(3)] 

1677 self.comp_files = [] 

1678 self.comp_refs = [] 

1679 for n, ref in enumerate(self.composite_refs): 

1680 component_refs = [] 

1681 component_files = [] 

1682 for component in sc.components: 

1683 component_ref = ref.makeComponentRef(component) 

1684 file = root_uri.join(f"composite_file-{n}-{component}.txt") 

1685 component_refs.append(component_ref) 

1686 component_files.append(file) 

1687 file.write(b"9876543210") 

1688 

1689 self.comp_files.append(component_files) 

1690 self.comp_refs.append(component_refs) 

1691 

1692 def tearDown(self) -> None: 

1693 if self.root is not None and os.path.exists(self.root): 

1694 shutil.rmtree(self.root, ignore_errors=True) 

1695 

1696 def _make_cache_manager(self, config_str: str) -> DatastoreCacheManager: 

1697 config = Config.fromYaml(config_str) 

1698 return DatastoreCacheManager(DatastoreCacheManagerConfig(config), universe=self.universe) 

1699 

1700 def testNoCacheDir(self) -> None: 

1701 config_str = """ 

1702cached: 

1703 root: null 

1704 cacheable: 

1705 metric0: true 

1706 """ 

1707 cache_manager = self._make_cache_manager(config_str) 

1708 

1709 # Look inside to check we don't have a cache directory 

1710 self.assertIsNone(cache_manager._cache_directory) 

1711 

1712 self.assertCache(cache_manager) 

1713 

1714 # Test that the cache directory is marked temporary 

1715 self.assertTrue(cache_manager.cache_directory.isTemporary) 

1716 

1717 def testNoCacheDirReversed(self) -> None: 

1718 """Use default caching status and metric1 to false""" 

1719 config_str = """ 

1720cached: 

1721 root: null 

1722 default: true 

1723 cacheable: 

1724 metric1: false 

1725 """ 

1726 cache_manager = self._make_cache_manager(config_str) 

1727 

1728 self.assertCache(cache_manager) 

1729 

1730 def testEnvvarCacheDir(self) -> None: 

1731 config_str = f""" 

1732cached: 

1733 root: '{self.root}' 

1734 cacheable: 

1735 metric0: true 

1736 """ 

1737 

1738 root = ResourcePath(self.root, forceDirectory=True) 

1739 env_dir = root.join("somewhere", forceDirectory=True) 

1740 elsewhere = root.join("elsewhere", forceDirectory=True) 

1741 

1742 # Environment variable should override the config value. 

1743 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath}): 

1744 cache_manager = self._make_cache_manager(config_str) 

1745 self.assertEqual(cache_manager.cache_directory, env_dir) 

1746 

1747 # This environment variable should not override the config value. 

1748 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}): 

1749 cache_manager = self._make_cache_manager(config_str) 

1750 self.assertEqual(cache_manager.cache_directory, root) 

1751 

1752 # No default setting. 

1753 config_str = """ 

1754cached: 

1755 root: null 

1756 default: true 

1757 cacheable: 

1758 metric1: false 

1759 """ 

1760 cache_manager = self._make_cache_manager(config_str) 

1761 

1762 # This environment variable should override the config value. 

1763 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}): 

1764 cache_manager = self._make_cache_manager(config_str) 

1765 self.assertEqual(cache_manager.cache_directory, env_dir) 

1766 

1767 # If both environment variables are set the main (not IF_UNSET) 

1768 # variable should win. 

1769 with unittest.mock.patch.dict( 

1770 os.environ, 

1771 { 

1772 "DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath, 

1773 "DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": elsewhere.ospath, 

1774 }, 

1775 ): 

1776 cache_manager = self._make_cache_manager(config_str) 

1777 self.assertEqual(cache_manager.cache_directory, env_dir) 

1778 

1779 # Use the API to set the environment variable, making sure that the 

1780 # variable is reset on exit. 

1781 with unittest.mock.patch.dict( 

1782 os.environ, 

1783 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""}, 

1784 ): 

1785 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

1786 self.assertTrue(defined) 

1787 cache_manager = self._make_cache_manager(config_str) 

1788 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True)) 

1789 

1790 # Now create the cache manager ahead of time and set the fallback 

1791 # later. 

1792 cache_manager = self._make_cache_manager(config_str) 

1793 self.assertIsNone(cache_manager._cache_directory) 

1794 with unittest.mock.patch.dict( 

1795 os.environ, 

1796 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""}, 

1797 ): 

1798 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

1799 self.assertTrue(defined) 

1800 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True)) 

1801 

1802 def testExplicitCacheDir(self) -> None: 

1803 config_str = f""" 

1804cached: 

1805 root: '{self.root}' 

1806 cacheable: 

1807 metric0: true 

1808 """ 

1809 cache_manager = self._make_cache_manager(config_str) 

1810 

1811 # Look inside to check we do have a cache directory. 

1812 self.assertEqual(cache_manager.cache_directory, ResourcePath(self.root, forceDirectory=True)) 

1813 

1814 self.assertCache(cache_manager) 

1815 

1816 # Test that the cache directory is not marked temporary 

1817 self.assertFalse(cache_manager.cache_directory.isTemporary) 

1818 

1819 def testUnexpectedFilesInCacheDir(self) -> None: 

1820 """Test for regression of a bug where extraneous files in a cache 

1821 directory would cause all cache lookups to raise an exception. 

1822 """ 

1823 config_str = f""" 

1824cached: 

1825 root: '{self.root}' 

1826 cacheable: 

1827 metric0: true 

1828 """ 

1829 

1830 for filename in ["unexpected.txt", "unexpected", "un_expected", "un_expected.txt"]: 

1831 unexpected_file = os.path.join(self.root, filename) 

1832 with open(unexpected_file, "w") as fh: 

1833 fh.write("test") 

1834 

1835 cache_manager = self._make_cache_manager(config_str) 

1836 cache_manager.scan_cache() 

1837 self.assertCache(cache_manager) 

1838 

1839 def assertCache(self, cache_manager: DatastoreCacheManager) -> None: 

1840 self.assertTrue(cache_manager.should_be_cached(self.refs[0])) 

1841 self.assertFalse(cache_manager.should_be_cached(self.refs[1])) 

1842 

1843 uri = cache_manager.move_to_cache(self.files[0], self.refs[0]) 

1844 self.assertIsInstance(uri, ResourcePath) 

1845 self.assertIsNone(cache_manager.move_to_cache(self.files[1], self.refs[1])) 

1846 

1847 # Check presence in cache using ref and then using file extension. 

1848 self.assertFalse(cache_manager.known_to_cache(self.refs[1])) 

1849 self.assertTrue(cache_manager.known_to_cache(self.refs[0])) 

1850 self.assertFalse(cache_manager.known_to_cache(self.refs[1], self.files[1].getExtension())) 

1851 self.assertTrue(cache_manager.known_to_cache(self.refs[0], self.files[0].getExtension())) 

1852 

1853 # Cached file should no longer exist but uncached file should be 

1854 # unaffected. 

1855 self.assertFalse(self.files[0].exists()) 

1856 self.assertTrue(self.files[1].exists()) 

1857 

1858 # Should find this file and it should be within the cache directory. 

1859 with cache_manager.find_in_cache(self.refs[0], ".txt") as found: 

1860 self.assertTrue(found.exists()) 

1861 self.assertIsNotNone(found.relative_to(cache_manager.cache_directory)) 

1862 

1863 # Should not be able to find these in cache 

1864 with cache_manager.find_in_cache(self.refs[0], ".fits") as found: 

1865 self.assertIsNone(found) 

1866 with cache_manager.find_in_cache(self.refs[1], ".fits") as found: 

1867 self.assertIsNone(found) 

1868 

1869 def testNoCache(self) -> None: 

1870 cache_manager = DatastoreDisabledCacheManager("", universe=self.universe) 

1871 for uri, ref in zip(self.files, self.refs, strict=True): 

1872 self.assertFalse(cache_manager.should_be_cached(ref)) 

1873 self.assertIsNone(cache_manager.move_to_cache(uri, ref)) 

1874 self.assertFalse(cache_manager.known_to_cache(ref)) 

1875 with cache_manager.find_in_cache(ref, ".txt") as found: 

1876 self.assertIsNone(found, msg=f"{cache_manager}") 

1877 

1878 def _expiration_config(self, mode: str, threshold: int) -> str: 

1879 return f""" 

1880cached: 

1881 default: true 

1882 expiry: 

1883 mode: {mode} 

1884 threshold: {threshold} 

1885 cacheable: 

1886 unused: true 

1887 """ 

1888 

1889 def testCacheExpiryFiles(self) -> None: 

1890 threshold = 2 # Keep at least 2 files. 

1891 mode = "files" 

1892 config_str = self._expiration_config(mode, threshold) 

1893 

1894 cache_manager = self._make_cache_manager(config_str) 

1895 

1896 # Check that an empty cache returns unknown for arbitrary ref 

1897 self.assertFalse(cache_manager.known_to_cache(self.refs[0])) 

1898 

1899 # Should end with datasets: 2, 3, 4 

1900 self.assertExpiration(cache_manager, 5, threshold + 1) 

1901 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1902 

1903 # Check that we will not expire a file that is actively in use. 

1904 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1905 self.assertIsNotNone(found) 

1906 

1907 # Trigger cache expiration that should remove the file 

1908 # we just retrieved. Should now have: 3, 4, 5 

1909 cached = cache_manager.move_to_cache(self.files[5], self.refs[5]) 

1910 self.assertIsNotNone(cached) 

1911 

1912 # Cache should still report the standard file count. 

1913 self.assertEqual(cache_manager.file_count, threshold + 1) 

1914 

1915 # Add additional entry to cache. 

1916 # Should now have 4, 5, 6 

1917 cached = cache_manager.move_to_cache(self.files[6], self.refs[6]) 

1918 self.assertIsNotNone(cached) 

1919 

1920 # Is the file still there? 

1921 self.assertTrue(found.exists()) 

1922 

1923 # Can we read it? 

1924 data = found.read() 

1925 self.assertGreater(len(data), 0) 

1926 

1927 # Outside context the file should no longer exist. 

1928 self.assertFalse(found.exists()) 

1929 

1930 # File count should not have changed. 

1931 self.assertEqual(cache_manager.file_count, threshold + 1) 

1932 

1933 # Dataset 2 was in the exempt directory but because hardlinks 

1934 # are used it was deleted from the main cache during cache expiry 

1935 # above and so should no longer be found. 

1936 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1937 self.assertIsNone(found) 

1938 

1939 # And the one stored after it is also gone. 

1940 with cache_manager.find_in_cache(self.refs[3], ".txt") as found: 

1941 self.assertIsNone(found) 

1942 

1943 # But dataset 4 is present. 

1944 with cache_manager.find_in_cache(self.refs[4], ".txt") as found: 

1945 self.assertIsNotNone(found) 

1946 

1947 # Adding a new dataset to the cache should now delete it. 

1948 cache_manager.move_to_cache(self.files[7], self.refs[7]) 

1949 

1950 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

1951 self.assertIsNone(found) 

1952 

1953 def testCacheExpiryDatasets(self) -> None: 

1954 threshold = 2 # Keep 2 datasets. 

1955 mode = "datasets" 

1956 config_str = self._expiration_config(mode, threshold) 

1957 

1958 cache_manager = self._make_cache_manager(config_str) 

1959 self.assertExpiration(cache_manager, 5, threshold + 1) 

1960 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1961 

1962 def testCacheExpiryDatasetsFromDisabled(self) -> None: 

1963 threshold = 2 

1964 mode = "datasets" 

1965 with unittest.mock.patch.dict( 

1966 os.environ, 

1967 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"}, 

1968 ): 

1969 cache_manager = DatastoreCacheManager.create_disabled(universe=DimensionUniverse()) 

1970 self.assertExpiration(cache_manager, 5, threshold + 1) 

1971 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1972 

1973 def testExpirationModeOverride(self) -> None: 

1974 threshold = 2 # Keep 2 datasets. 

1975 mode = "datasets" 

1976 config_str = self._expiration_config(mode, threshold) 

1977 

1978 mode = "size" 

1979 threshold = 55 

1980 with unittest.mock.patch.dict( 

1981 os.environ, 

1982 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"}, 

1983 ): 

1984 cache_manager = self._make_cache_manager(config_str) 

1985 self.assertExpiration(cache_manager, 10, 6) 

1986 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

1987 

1988 # Check we get a warning with unrecognized form. 

1989 with unittest.mock.patch.dict( 

1990 os.environ, 

1991 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something"}, 

1992 ): 

1993 with self.assertLogs(level="WARNING") as cm: 

1994 self._make_cache_manager(config_str) 

1995 self.assertIn("Unrecognized form (something)", cm.output[0]) 

1996 

1997 with unittest.mock.patch.dict( 

1998 os.environ, 

1999 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something=5"}, 

2000 ): 

2001 with self.assertRaises(ValueError) as cm: 

2002 self._make_cache_manager(config_str) 

2003 self.assertIn("Unrecognized value", str(cm.exception)) 

2004 

2005 def testMissingThreshold(self) -> None: 

2006 threshold = "" 

2007 mode = "datasets" 

2008 config_str = self._expiration_config(mode, threshold) 

2009 

2010 with self.assertRaises(ValueError) as cm: 

2011 self._make_cache_manager(config_str) 

2012 self.assertIn("Cache expiration threshold", str(cm.exception)) 

2013 

2014 def testCacheExpiryDatasetsComposite(self) -> None: 

2015 threshold = 2 # Keep 2 datasets. 

2016 mode = "datasets" 

2017 config_str = self._expiration_config(mode, threshold) 

2018 

2019 cache_manager = self._make_cache_manager(config_str) 

2020 

2021 n_datasets = 3 

2022 for i in range(n_datasets): 

2023 for component_file, component_ref in zip(self.comp_files[i], self.comp_refs[i], strict=True): 

2024 cached = cache_manager.move_to_cache(component_file, component_ref) 

2025 self.assertIsNotNone(cached) 

2026 self.assertTrue(cache_manager.known_to_cache(component_ref)) 

2027 self.assertTrue(cache_manager.known_to_cache(component_ref.makeCompositeRef())) 

2028 self.assertTrue(cache_manager.known_to_cache(component_ref, component_file.getExtension())) 

2029 

2030 self.assertEqual(cache_manager.file_count, 6) # 2 datasets each of 3 files 

2031 

2032 # Write two new non-composite and the number of files should drop. 

2033 self.assertExpiration(cache_manager, 2, 5) 

2034 

2035 def testCacheExpirySize(self) -> None: 

2036 threshold = 55 # Each file is 10 bytes 

2037 mode = "size" 

2038 config_str = self._expiration_config(mode, threshold) 

2039 

2040 cache_manager = self._make_cache_manager(config_str) 

2041 self.assertExpiration(cache_manager, 10, 6) 

2042 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

2043 

2044 def testDisabledCache(self) -> None: 

2045 # Configure an active cache but disable via environment. 

2046 threshold = 2 

2047 mode = "datasets" 

2048 config_str = self._expiration_config(mode, threshold) 

2049 

2050 with unittest.mock.patch.dict( 

2051 os.environ, 

2052 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "disabled"}, 

2053 ): 

2054 env_cache_manager = self._make_cache_manager(config_str) 

2055 

2056 # Configure to be disabled 

2057 threshold = 0 

2058 mode = "disabled" 

2059 config_str = self._expiration_config(mode, threshold) 

2060 cfg_cache_manager = self._make_cache_manager(config_str) 

2061 

2062 for cache_manager in ( 

2063 cfg_cache_manager, 

2064 env_cache_manager, 

2065 DatastoreCacheManager.create_disabled(universe=DimensionUniverse()), 

2066 ): 

2067 for uri, ref in zip(self.files, self.refs, strict=True): 

2068 self.assertFalse(cache_manager.should_be_cached(ref)) 

2069 self.assertIsNone(cache_manager.move_to_cache(uri, ref)) 

2070 self.assertFalse(cache_manager.known_to_cache(ref)) 

2071 with cache_manager.find_in_cache(ref, ".txt") as found: 

2072 self.assertIsNone(found, msg=f"{cache_manager}") 

2073 self.assertIn("disabled", str(cache_manager)) 

2074 

2075 def assertExpiration( 

2076 self, cache_manager: DatastoreCacheManager, n_datasets: int, n_retained: int 

2077 ) -> None: 

2078 """Insert the datasets and then check the number retained.""" 

2079 for i in range(n_datasets): 

2080 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2081 self.assertIsNotNone(cached) 

2082 

2083 self.assertEqual(cache_manager.file_count, n_retained) 

2084 

2085 # The oldest file should not be in the cache any more. 

2086 for i in range(n_datasets): 

2087 with cache_manager.find_in_cache(self.refs[i], ".txt") as found: 

2088 if i >= n_datasets - n_retained: 

2089 self.assertIsInstance(found, ResourcePath) 

2090 else: 

2091 self.assertIsNone(found) 

2092 

2093 def testCacheExpiryAge(self) -> None: 

2094 threshold = 1 # Expire older than 2 seconds 

2095 mode = "age" 

2096 config_str = self._expiration_config(mode, threshold) 

2097 

2098 cache_manager = self._make_cache_manager(config_str) 

2099 self.assertIn(f"{mode}={threshold}", str(cache_manager)) 

2100 

2101 # Insert 3 files, then sleep, then insert more. 

2102 for i in range(2): 

2103 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2104 self.assertIsNotNone(cached) 

2105 time.sleep(2.0) 

2106 for j in range(4): 

2107 i = 2 + j # Continue the counting 

2108 cached = cache_manager.move_to_cache(self.files[i], self.refs[i]) 

2109 self.assertIsNotNone(cached) 

2110 

2111 # Only the files written after the sleep should exist. 

2112 self.assertEqual(cache_manager.file_count, 4) 

2113 with cache_manager.find_in_cache(self.refs[1], ".txt") as found: 

2114 self.assertIsNone(found) 

2115 with cache_manager.find_in_cache(self.refs[2], ".txt") as found: 

2116 self.assertIsInstance(found, ResourcePath) 

2117 

2118 

2119class NullDatastoreTestCase(DatasetTestHelper, unittest.TestCase): 

2120 """Test the null datastore.""" 

2121 

2122 storageClassFactory = StorageClassFactory() 

2123 

2124 def test_basics(self) -> None: 

2125 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

2126 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2127 

2128 null = NullDatastore(None, None) 

2129 

2130 self.assertFalse(null.exists(ref)) 

2131 self.assertFalse(null.knows(ref)) 

2132 knows = null.knows_these([ref]) 

2133 self.assertFalse(knows[ref]) 

2134 null.validateConfiguration(ref) 

2135 

2136 with self.assertRaises(FileNotFoundError): 

2137 null.get(ref) 

2138 with self.assertRaises(NotImplementedError): 

2139 null.put("", ref) 

2140 with self.assertRaises(FileNotFoundError): 

2141 null.getURI(ref) 

2142 with self.assertRaises(FileNotFoundError): 

2143 null.getURIs(ref) 

2144 with self.assertRaises(FileNotFoundError): 

2145 null.getManyURIs([ref]) 

2146 with self.assertRaises(NotImplementedError): 

2147 null.getLookupKeys() 

2148 with self.assertRaises(NotImplementedError): 

2149 null.import_records({}) 

2150 with self.assertRaises(NotImplementedError): 

2151 null.export_records([]) 

2152 with self.assertRaises(NotImplementedError): 

2153 null.export_predicted_records([]) 

2154 with self.assertRaises(NotImplementedError): 

2155 null.export([ref]) 

2156 with self.assertRaises(NotImplementedError): 

2157 null.transfer(null, ref) 

2158 with self.assertRaises(NotImplementedError): 

2159 null.emptyTrash() 

2160 with self.assertRaises(NotImplementedError): 

2161 null.trash(ref) 

2162 with self.assertRaises(NotImplementedError): 

2163 null.forget([ref]) 

2164 with self.assertRaises(NotImplementedError): 

2165 null.remove(ref) 

2166 with self.assertRaises(NotImplementedError): 

2167 null.retrieveArtifacts([ref], ResourcePath(".")) 

2168 with self.assertRaises(NotImplementedError): 

2169 null.transfer_from(null, [ref]) 

2170 with self.assertRaises(NotImplementedError): 

2171 null.ingest() 

2172 

2173 

2174class DatasetRefURIsTestCase(unittest.TestCase): 

2175 """Tests for DatasetRefURIs.""" 

2176 

2177 def testSequenceAccess(self) -> None: 

2178 """Verify that DatasetRefURIs can be treated like a two-item tuple.""" 

2179 uris = DatasetRefURIs() 

2180 

2181 self.assertEqual(len(uris), 2) 

2182 self.assertEqual(uris[0], None) 

2183 self.assertEqual(uris[1], {}) 

2184 

2185 primaryURI = ResourcePath("1/2/3") 

2186 componentURI = ResourcePath("a/b/c") 

2187 

2188 # affirm that DatasetRefURIs does not support MutableSequence functions 

2189 with self.assertRaises(TypeError): 

2190 uris[0] = primaryURI 

2191 with self.assertRaises(TypeError): 

2192 uris[1] = {"foo": componentURI} 

2193 

2194 # but DatasetRefURIs can be set by property name: 

2195 uris.primaryURI = primaryURI 

2196 uris.componentURIs = {"foo": componentURI} 

2197 self.assertEqual(uris.primaryURI, primaryURI) 

2198 self.assertEqual(uris[0], primaryURI) 

2199 

2200 primary, components = uris 

2201 self.assertEqual(primary, primaryURI) 

2202 self.assertEqual(components, {"foo": componentURI}) 

2203 

2204 def testRepr(self) -> None: 

2205 """Verify __repr__ output.""" 

2206 uris = DatasetRefURIs(ResourcePath("/1/2/3"), {"comp": ResourcePath("/a/b/c")}) 

2207 self.assertEqual( 

2208 repr(uris), 

2209 'DatasetRefURIs(ResourcePath("file:///1/2/3"), {\'comp\': ResourcePath("file:///a/b/c")})', 

2210 ) 

2211 

2212 

2213class StoredFileInfoTestCase(DatasetTestHelper, unittest.TestCase): 

2214 """Test the StoredFileInfo class.""" 

2215 

2216 storageClassFactory = StorageClassFactory() 

2217 

2218 def test_StoredFileInfo(self) -> None: 

2219 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

2220 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2221 

2222 record = dict( 

2223 storage_class="StructuredDataDict", 

2224 formatter="lsst.daf.butler.Formatter", 

2225 path="a/b/c.txt", 

2226 component="component", 

2227 checksum=None, 

2228 file_size=5, 

2229 ) 

2230 info = StoredFileInfo.from_record(record) 

2231 

2232 self.assertEqual(info.to_record(), record) 

2233 

2234 ref2 = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {}) 

2235 rebased = info.rebase(ref2) 

2236 self.assertEqual(rebased.rebase(ref), info) 

2237 

2238 with self.assertRaises(TypeError): 

2239 rebased.update(formatter=42) 

2240 

2241 with self.assertRaises(ValueError): 

2242 rebased.update(something=42, new="42") 

2243 

2244 # Check that pickle works on StoredFileInfo. 

2245 pickled_info = pickle.dumps(info) 

2246 unpickled_info = pickle.loads(pickled_info) 

2247 self.assertEqual(unpickled_info, info) 

2248 

2249 def test_make_datastore_path_relative(self): 

2250 self.assertEqual(make_datastore_path_relative("a/relative/path"), "a/relative/path") 

2251 self.assertEqual(make_datastore_path_relative("path/with#fragment"), "path/with#fragment") 

2252 self.assertEqual(make_datastore_path_relative("http://server.com/some/path"), "some/path") 

2253 self.assertEqual(make_datastore_path_relative("http://server.com/some/path#frag"), "some/path#frag") 

2254 

2255 

2256@contextlib.contextmanager 

2257def _temp_yaml_file(data: Any) -> Iterator[str]: 

2258 fh = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") 

2259 try: 

2260 yaml.dump(data, stream=fh) 

2261 fh.flush() 

2262 yield fh.name 

2263 finally: 

2264 # Some tests delete the file 

2265 with contextlib.suppress(FileNotFoundError): 

2266 fh.close() 

2267 

2268 

2269if __name__ == "__main__": 

2270 unittest.main()