Coverage for tests/test_datastore.py: 13%

1# This file is part of daf_butler.

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.

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

28from __future__ import annotations

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

42import yaml

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.record_data import DatastoreRecordData, SerializedDatastoreRecordData

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

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

67from lsst.daf.butler.tests import (

68 BadNoWriteFormatter,

69 BadWriteFormatter,

70 DatasetTestHelper,

71 DatastoreTestHelper,

72 DummyRegistry,

73 MetricsExample,

74 MetricsExampleDataclass,

75 MetricsExampleModel,

76)

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

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

79from lsst.resources import ResourcePath

80from lsst.utils import doImport

81from lsst.utils.introspection import get_full_type_name

83TESTDIR = os.path.dirname(__file__)

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

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

88 if use_none:

89 array = None

90 else:

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

92 return MetricsExample(

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

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

95 array,

96 )

99class TransactionTestError(Exception):

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

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

102 """

103

104 pass

105

106

107class DatastoreTestsBase(DatasetTestHelper, DatastoreTestHelper, TestCaseMixin):

108 """Support routines for datastore testing"""

109

110 root: str | None = None

111 universe: DimensionUniverse

112 storageClassFactory: StorageClassFactory

113

114 @classmethod

115 def setUpClass(cls) -> None:

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

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

118 cls.storageClassFactory = StorageClassFactory()

119 cls.storageClassFactory.addFromConfig(scConfigFile)

120

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

122 # information (since we should not assume the constructor

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

124 datastoreConfig = DatastoreConfig(cls.configFile)

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

126 cls.universe = DimensionUniverse()

127

128 def setUp(self) -> None:

129 self.setUpDatastoreTests(DummyRegistry, DatastoreConfig)

130

131 def tearDown(self) -> None:

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

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

134

135

136class DatastoreTests(DatastoreTestsBase):

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

138

139 hasUnsupportedPut = True

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

141 isEphemeral: bool = False

142 validationCanFail: bool = False

143

144 def testConfigRoot(self) -> None:

145 full = DatastoreConfig(self.configFile)

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

147 newroot = "/random/location"

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

149 if self.rootKeys:

150 for k in self.rootKeys:

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

152

153 def testConstructor(self) -> None:

154 datastore = self.makeDatastore()

155 self.assertIsNotNone(datastore)

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

157

158 def testConfigurationValidation(self) -> None:

159 datastore = self.makeDatastore()

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

161 datastore.validateConfiguration([sc])

162

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

164 if self.validationCanFail:

165 with self.assertRaises(DatastoreValidationError):

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

167

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

169 dataId = {

170 "instrument": "dummy",

171 "visit": 52,

172 "physical_filter": "V",

173 "band": "v",

174 "day_obs": 20250101,

175 }

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

177 datastore.validateConfiguration([ref])

178

179 def testParameterValidation(self) -> None:

180 """Check that parameters are validated"""

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

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

183 dataId = {

184 "instrument": "dummy",

185 "visit": 52,

186 "physical_filter": "V",

187 "band": "v",

188 "day_obs": 20250101,

189 }

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

191 datastore = self.makeDatastore()

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

193 datastore.put(data, ref)

194 newdata = datastore.get(ref)

195 self.assertEqual(data, newdata)

196 with self.assertRaises(KeyError):

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

198

199 def testBasicPutGet(self) -> None:

200 metrics = makeExampleMetrics()

201 datastore = self.makeDatastore()

202

203 # Create multiple storage classes for testing different formulations

204 storageClasses = [

205 self.storageClassFactory.getStorageClass(sc)

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

207 ]

208

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

210 dataId = {

211 "instrument": "dummy",

212 "visit": 52,

213 "physical_filter": "V",

214 "band": "v",

215 "day_obs": 20250101,

216 }

217 dataId2 = {

218 "instrument": "dummy",

219 "visit": 53,

220 "physical_filter": "V",

221 "band": "v",

222 "day_obs": 20250101,

223 }

224

225 for sc in storageClasses:

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

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

228

229 # Make sure that using getManyURIs without predicting before the

230 # dataset has been put raises.

231 with self.assertRaises(FileNotFoundError):

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

233

234 # Make sure that using getManyURIs with predicting before the

235 # dataset has been put predicts the URI.

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

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

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

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

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

241

242 datastore.put(metrics, ref)

243

244 # Does it exist?

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

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

247 multi = datastore.knows_these([ref])

248 self.assertTrue(multi[ref])

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

250 self.assertTrue(multi[ref])

251 self.assertFalse(multi[ref2])

252

253 # Get

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

255 self.assertEqual(metrics, metricsOut)

256

257 uri = datastore.getURI(ref)

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

259

260 uris = datastore.getManyURIs([ref])

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

262 ref, uri = uris.popitem()

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

264 self.assertFalse(uri.componentURIs)

265

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

267 # with derived storage classes but with parent ID

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

269 compRef = ref.makeComponentRef(comp)

270 output = datastore.get(compRef)

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

272

273 uri = datastore.getURI(compRef)

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

275

276 uris = datastore.getManyURIs([compRef])

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

278

279 storageClass = sc

280

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

282 # get it back as None

283 metricsNone = makeExampleMetrics(use_none=True)

284 dataIdNone = {

285 "instrument": "dummy",

286 "visit": 54,

287 "physical_filter": "V",

288 "band": "v",

289 "day_obs": 20250101,

290 }

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

292 datastore.put(metricsNone, refNone)

293

294 comp = "data"

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

296 compRef = refNone.makeComponentRef(comp)

297 output = datastore.get(compRef)

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

299

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

301 if self.hasUnsupportedPut:

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

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

304 with self.assertRaises(DatasetTypeNotSupportedError):

305 datastore.put(metrics, ref)

306

307 # These should raise

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

309 with self.assertRaises(FileNotFoundError):

310 # non-existing file

311 datastore.get(ref)

312

313 # Get a URI from it

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

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

316

317 with self.assertRaises(FileNotFoundError):

318 datastore.getURI(ref)

319

320 def testTrustGetRequest(self) -> None:

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

322 datastore = self.makeDatastore()

323

324 # Skip test if the attribute is not defined

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

326 return

327

328 metrics = makeExampleMetrics()

329

330 i = 0

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

332 i += 1

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

334

335 if sc_name == "StructuredComposite":

336 disassembled = True

337 else:

338 disassembled = False

339

340 # Start datastore in default configuration of using registry

341 datastore.trustGetRequest = False

342

343 # Create multiple storage classes for testing with or without

344 # disassembly

345 sc = self.storageClassFactory.getStorageClass(sc_name)

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

347

348 dataId = {

349 "instrument": "dummy",

350 "visit": 52 + i,

351 "physical_filter": "V",

352 "band": "v",

353 "day_obs": 20250101,

354 }

355

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

357 datastore.put(metrics, ref)

358

359 # Does it exist?

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

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

362 multi = datastore.knows_these([ref])

363 self.assertTrue(multi[ref])

364 multi = datastore.mexists([ref])

365 self.assertTrue(multi[ref])

366

367 # Get

368 metricsOut = datastore.get(ref)

369 self.assertEqual(metrics, metricsOut)

370

371 # Get the URI(s)

372 primaryURI, componentURIs = datastore.getURIs(ref)

373 if disassembled:

374 self.assertIsNone(primaryURI)

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

376 else:

377 self.assertIn(datasetTypeName, primaryURI.path)

378 self.assertFalse(componentURIs)

379

380 # Delete registry entry so now we are trusting

381 datastore.removeStoredItemInfo(ref)

382

383 # Now stop trusting and check that things break

384 datastore.trustGetRequest = False

385

386 # Does it exist?

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

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

389 multi = datastore.knows_these([ref])

390 self.assertFalse(multi[ref])

391 multi = datastore.mexists([ref])

392 self.assertFalse(multi[ref])

393

394 with self.assertRaises(FileNotFoundError):

395 datastore.get(ref)

396

397 if sc_name != "StructuredDataNoComponents":

398 with self.assertRaises(FileNotFoundError):

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

400

401 # URI should fail unless we ask for prediction

402 with self.assertRaises(FileNotFoundError):

403 datastore.getURIs(ref)

404

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

406 if disassembled:

407 self.assertIsNone(predicted_primary)

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

409 for uri in predicted_disassembled.values():

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

411 self.assertIn(datasetTypeName, uri.path)

412 else:

413 self.assertIn(datasetTypeName, predicted_primary.path)

414 self.assertFalse(predicted_disassembled)

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

416

417 # Now enable registry-free trusting mode

418 datastore.trustGetRequest = True

419

420 # Try again to get it

421 metricsOut = datastore.get(ref)

422 self.assertEqual(metricsOut, metrics)

423

424 # Does it exist?

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

426

427 # Get a component

428 if sc_name != "StructuredDataNoComponents":

429 comp = "data"

430 compRef = ref.makeComponentRef(comp)

431 output = datastore.get(compRef)

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

433

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

435 # enabling prediction.

436 primaryURI2, componentURIs2 = datastore.getURIs(ref)

437 self.assertEqual(primaryURI2, primaryURI)

438 self.assertEqual(componentURIs2, componentURIs)

439

440 # Check for compatible storage class.

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

442 # Make new dataset ref with compatible storage class.

443 ref_comp = ref.overrideStorageClass("StructuredDataDictJson")

444

445 # Without `set_retrieve_dataset_type_method` it will fail to

446 # find correct file.

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

448 with self.assertRaises(FileNotFoundError):

449 datastore.get(ref_comp)

450 with self.assertRaises(FileNotFoundError):

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

452

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

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

455 if name == ref.datasetType.name:

456 return ref.datasetType

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

458

459 datastore.set_retrieve_dataset_type_method(_stored_dataset_type)

460

461 # Storage class override with original dataset ref.

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

463 self.assertIsInstance(metrics_as_dict, dict)

464

465 # get() should return a dict now.

466 metrics_as_dict = datastore.get(ref_comp)

467 self.assertIsInstance(metrics_as_dict, dict)

468

469 # exists() should work as well.

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

471

472 datastore.set_retrieve_dataset_type_method(None)

473

474 def testDisassembly(self) -> None:

475 """Test disassembly within datastore."""

476 metrics = makeExampleMetrics()

477 if self.isEphemeral:

478 # in-memory datastore does not disassemble

479 return

480

481 # Create multiple storage classes for testing different formulations

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

483 # a reference.

484 storageClasses = [

485 self.storageClassFactory.getStorageClass(sc)

486 for sc in (

487 "StructuredComposite",

488 "StructuredCompositeTestA",

489 "StructuredCompositeTestB",

490 "StructuredCompositeReadComp",

491 "StructuredData", # No disassembly

492 "StructuredCompositeReadCompNoDisassembly",

493 )

494 ]

495

496 # Create the test datastore

497 datastore = self.makeDatastore()

498

499 # Dummy dataId

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

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

502

503 for i, sc in enumerate(storageClasses):

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

505 # Create a different dataset type each time round

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

507 # a cascade of tests because of file clashes

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

509

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

511

512 datastore.put(metrics, ref)

513

514 baseURI, compURIs = datastore.getURIs(ref)

515 if disassembled:

516 self.assertIsNone(baseURI)

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

518 else:

519 self.assertIsNotNone(baseURI)

520 self.assertEqual(compURIs, {})

521

522 metrics_get = datastore.get(ref)

523 self.assertEqual(metrics_get, metrics)

524

525 # Retrieve the composite with read parameter

526 stop = 4

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

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

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

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

531

532 # Retrieve a component

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

534 self.assertEqual(data, metrics.data)

535

536 # On supported storage classes attempt to access a read

537 # only component

538 if "ReadComp" in sc.name:

539 cRef = ref.makeComponentRef("counter")

540 counter = datastore.get(cRef)

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

542

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

544 self.assertEqual(counter, stop)

545

546 datastore.remove(ref)

547

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

549 metrics = makeExampleMetrics()

550 datastore = self.makeDatastore()

551 # Put

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

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

554 refs = []

555 for i in range(n_refs):

556 dataId = {

557 "instrument": "dummy",

558 "visit": 638 + i,

559 "physical_filter": "U",

560 "band": "u",

561 "day_obs": 20250101,

562 }

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

564 datastore.put(metrics, ref)

565

566 # Does it exist?

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

568

569 # Get

570 metricsOut = datastore.get(ref)

571 self.assertEqual(metrics, metricsOut)

572 refs.append(ref)

573

574 return datastore, *refs

575

576 def testRemove(self) -> None:

577 datastore, ref = self.prepDeleteTest()

578

579 # Remove

580 datastore.remove(ref)

581

582 # Does it exist?

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

584

585 # Do we now get a predicted URI?

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

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

588

589 # Get should now fail

590 with self.assertRaises(FileNotFoundError):

591 datastore.get(ref)

592 # Can only delete once

593 with self.assertRaises(FileNotFoundError):

594 datastore.remove(ref)

595

596 def testForget(self) -> None:

597 datastore, ref = self.prepDeleteTest()

598

599 # Remove

600 datastore.forget([ref])

601

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

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

604

605 # Do we now get a predicted URI?

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

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

608

609 # Get should now fail

610 with self.assertRaises(FileNotFoundError):

611 datastore.get(ref)

612

613 # Forgetting again is a silent no-op

614 datastore.forget([ref])

615

616 # Predicted URI should still point to the file.

617 self.assertTrue(uri.exists())

618

619 def testTransfer(self) -> None:

620 metrics = makeExampleMetrics()

621

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

623 dataId = {

624 "instrument": "dummy",

625 "visit": 2048,

626 "physical_filter": "Uprime",

627 "band": "u",

628 "day_obs": 20250101,

629 }

630

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

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

633

634 inputDatastore = self.makeDatastore("test_input_datastore")

635 outputDatastore = self.makeDatastore("test_output_datastore")

636

637 inputDatastore.put(metrics, ref)

638 outputDatastore.transfer(inputDatastore, ref)

639

640 metricsOut = outputDatastore.get(ref)

641 self.assertEqual(metrics, metricsOut)

642

643 def testBasicTransaction(self) -> None:

644 datastore = self.makeDatastore()

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

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

647 nDatasets = 6

648 dataIds = [

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

650 for i in range(nDatasets)

651 ]

652 data = [

653 (

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

655 makeExampleMetrics(),

656 )

657 for dataId in dataIds

658 ]

659 succeed = data[: nDatasets // 2]

660 fail = data[nDatasets // 2 :]

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

662 with datastore.transaction():

663 for ref, metrics in succeed:

664 datastore.put(metrics, ref)

665 # Whereas datasets added in this transaction should not

666 with self.assertRaises(TransactionTestError):

667 with datastore.transaction():

668 for ref, metrics in fail:

669 datastore.put(metrics, ref)

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

671 # Check for datasets that should exist

672 for ref, metrics in succeed:

673 # Does it exist?

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

675 # Get

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

677 self.assertEqual(metrics, metricsOut)

678 # URI

679 uri = datastore.getURI(ref)

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

681 # Check for datasets that should not exist

682 for ref, _ in fail:

683 # These should raise

684 with self.assertRaises(FileNotFoundError):

685 # non-existing file

686 datastore.get(ref)

687 with self.assertRaises(FileNotFoundError):

688 datastore.getURI(ref)

689

690 def testNestedTransaction(self) -> None:

691 datastore = self.makeDatastore()

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

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

694 metrics = makeExampleMetrics()

695

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

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

698 datastore.put(metrics, refBefore)

699 with self.assertRaises(TransactionTestError):

700 with datastore.transaction():

701 dataId = {

702 "instrument": "dummy",

703 "visit": 1,

704 "physical_filter": "V",

705 "band": "v",

706 "day_obs": 20250101,

707 }

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

709 datastore.put(metrics, refOuter)

710 with datastore.transaction():

711 dataId = {

712 "instrument": "dummy",

713 "visit": 2,

714 "physical_filter": "V",

715 "band": "v",

716 "day_obs": 20250101,

717 }

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

719 datastore.put(metrics, refInner)

720 # All datasets should exist

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

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

723 self.assertEqual(metrics, metricsOut)

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

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

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

727 self.assertEqual(metrics, metricsOut)

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

729 # should be gone

730 with self.assertRaises(FileNotFoundError):

731 datastore.get(refOuter)

732 with self.assertRaises(FileNotFoundError):

733 datastore.get(refInner)

734

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

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

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

738 metrics = makeExampleMetrics()

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

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

741 return metrics, ref

742

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

744 metrics, ref = self._prepareIngestTest()

745 # The file will be deleted after the test.

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

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

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

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

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

751 func(metrics, path, ref)

752

753 def testIngestNoTransfer(self) -> None:

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

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

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

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

758 continue

759

760 with self.subTest(mode=mode):

761 datastore = self.makeDatastore()

762

763 def succeed(

764 obj: MetricsExample,

765 path: str,

766 ref: DatasetRef,

767 mode: str | None = mode,

768 datastore: Datastore = datastore,

769 ) -> None:

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

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

772 # accordingly.

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

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

775 # have to copy it into all the roots.

776 relative_path = None

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

778 if root is not None:

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

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

781 assert relative_path is not None, (

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

783 )

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

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

786

787 def failInputDoesNotExist(

788 obj: MetricsExample,

789 path: str,

790 ref: DatasetRef,

791 mode: str | None = mode,

792 datastore: Datastore = datastore,

793 ) -> None:

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

795 with self.assertRaises(FileNotFoundError):

796 datastore.ingest(

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

798 )

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

800

801 def failOutsideRoot(

802 obj: MetricsExample,

803 path: str,

804 ref: DatasetRef,

805 mode: str | None = mode,

806 datastore: Datastore = datastore,

807 ) -> None:

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

809 auto.

810 """

811 if mode == "auto":

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

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

814 else:

815 with self.assertRaises(RuntimeError):

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

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

818

819 def failNotImplemented(

820 obj: MetricsExample,

821 path: str,

822 ref: DatasetRef,

823 mode: str | None = mode,

824 datastore: Datastore = datastore,

825 ) -> None:

826 with self.assertRaises(NotImplementedError):

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

828

829 if mode in self.ingestTransferModes:

830 self.runIngestTest(failOutsideRoot)

831 self.runIngestTest(failInputDoesNotExist)

832 self.runIngestTest(succeed)

833 else:

834 self.runIngestTest(failNotImplemented)

835

836 def testIngestTransfer(self) -> None:

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

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

839 with self.subTest(mode=mode):

840 datastore = self.makeDatastore(mode)

841

842 def succeed(

843 obj: MetricsExample,

844 path: str,

845 ref: DatasetRef,

846 mode: str | None = mode,

847 datastore: Datastore = datastore,

848 ) -> None:

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

850 location.

851 """

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

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

854 file_exists = os.path.exists(path)

855 if mode == "move":

856 self.assertFalse(file_exists)

857 else:

858 self.assertTrue(file_exists)

859

860 def failInputDoesNotExist(

861 obj: MetricsExample,

862 path: str,

863 ref: DatasetRef,

864 mode: str | None = mode,

865 datastore: Datastore = datastore,

866 ) -> None:

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

868 with self.assertRaises(FileNotFoundError):

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

870 # datastore for auto mode

871 datastore.ingest(

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

873 )

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

875

876 def failNotImplemented(

877 obj: MetricsExample,

878 path: str,

879 ref: DatasetRef,

880 mode: str | None = mode,

881 datastore: Datastore = datastore,

882 ) -> None:

883 with self.assertRaises(NotImplementedError):

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

885

886 if mode in self.ingestTransferModes:

887 self.runIngestTest(failInputDoesNotExist)

888 self.runIngestTest(succeed)

889 else:

890 self.runIngestTest(failNotImplemented)

891

892 def testIngestSymlinkOfSymlink(self) -> None:

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

894 metrics, ref = self._prepareIngestTest()

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

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

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

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

899 if mode not in self.ingestTransferModes:

900 continue

901

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

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

904 with tempfile.TemporaryDirectory() as tmpdir:

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

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

907

908 datastore = self.makeDatastore()

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

910

911 uri = datastore.getURI(ref)

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

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

914

915 linkTarget = os.readlink(uri.ospath)

916 if mode == "relsymlink":

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

918 else:

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

920

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

922 metric2 = datastore.get(ref)

923 self.assertEqual(metric2, metrics)

924

925 # Cleanup the file for next time round loop

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

927 datastore.remove(ref)

928

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

930 datastore = self.makeDatastore(name)

931

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

933 # ChainedDatastore that only includes InMemoryDatastores have to be

934 # skipped as well.

935 for name in datastore.names:

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

937 break

938 else:

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

940

941 metrics = makeExampleMetrics()

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

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

944

945 refs = []

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

947 dataId = {

948 "instrument": "dummy",

949 "visit": visit,

950 "physical_filter": "Uprime",

951 "band": "u",

952 "day_obs": 20250101,

953 }

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

955 datastore.put(metrics, ref)

956 refs.append(ref)

957 return datastore, refs

958

959 def testExportImportRecords(self) -> None:

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

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

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

963 n_refs = len(exported_refs)

964 records = datastore.export_records(exported_refs)

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

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

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

968 for datastore_name in records:

969 record_data = records[datastore_name]

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

971

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

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

974 subset = record_data.subset(dataset_ids)

975 assert subset is not None

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

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

978 self.assertIsNone(subset)

979

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

981 datastore2 = self.makeDatastore("test_datastore")

982

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

984 datastore2.import_records(records)

985

986 with self.assertRaises(FileNotFoundError):

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

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

989 self.assertIsNotNone(data)

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

991 self.assertIsNotNone(data)

992

993 def testExportPredictedRecords(self):

994 if self.isEphemeral:

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

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

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

998 dataId = {

999 "instrument": "dummy",

1000 "visit": 52,

1001 "physical_filter": "V",

1002 "band": "v",

1003 "day_obs": 20250101,

1004 }

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

1006

1007 datastore = self.makeDatastore("test_datastore")

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

1009 records = datastore.export_predicted_records([ref])

1010

1011 # Expect predicted records from all datastores.

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

1013

1014 for record_data in records.values():

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

1016

1017 def testExport(self) -> None:

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

1019

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

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

1022

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

1024 # Both will default to None

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

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

1027

1028 with self.assertRaises(TypeError):

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

1030

1031 with self.assertRaises(TypeError):

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

1033

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

1035 # export it.

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

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

1038 dataId = {

1039 "instrument": "dummy",

1040 "visit": 52,

1041 "physical_filter": "V",

1042 "band": "v",

1043 "day_obs": 20250101,

1044 }

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

1046 with self.assertRaises(FileNotFoundError):

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

1048

1049 def test_pydantic_dict_storage_class_conversions(self) -> None:

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

1051 read.

1052 """

1053 datastore = self.makeDatastore()

1054 store_as_model = self.makeDatasetRef(

1055 "store_as_model",

1056 dimensions=self.universe.empty,

1057 storageClass="DictConvertibleModel",

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

1059 )

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

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

1062 datastore.put(model, store_as_model)

1063 retrieved_model = datastore.get(store_as_model)

1064 self.assertEqual(retrieved_model, model)

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

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

1067 self.assertEqual(loaded, content)

1068

1069 def test_simple_class_put_get(self) -> None:

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

1071 constructor.

1072 """

1073 datastore = self.makeDatastore()

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

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

1076

1077 def test_dataclass_put_get(self) -> None:

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

1079 datastore = self.makeDatastore()

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

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

1082

1083 def test_pydantic_put_get(self) -> None:

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

1085 datastore = self.makeDatastore()

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

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

1088

1089 def test_tuple_put_get(self) -> None:

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

1091 datastore = self.makeDatastore()

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

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

1094

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

1096 refs = {

1097 x: self.makeDatasetRef(

1098 f"stora_as_{x}",

1099 dimensions=self.universe.empty,

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

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

1102 )

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

1104 }

1105

1106 for ref in refs.values():

1107 datastore.put(data, ref)

1108

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

1110

1111

1112class PosixDatastoreTestCase(DatastoreTests, unittest.TestCase):

1113 """PosixDatastore specialization"""

1114

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

1116 uriScheme = "file"

1117 canIngestNoTransferAuto = True

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

1119 isEphemeral = False

1120 rootKeys = ("root",)

1121 validationCanFail = True

1122

1123 def setUp(self) -> None:

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

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

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

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

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

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

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

1131 # relative path.

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

1133 super().setUp()

1134

1135 def testAtomicWrite(self) -> None:

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

1137 datastore = self.makeDatastore()

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

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

1140 metrics = makeExampleMetrics()

1141

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

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

1144

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

1146 datastore.put(metrics, ref)

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

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

1149

1150 # And the transfer should be file to file.

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

1152

1153 def testCanNotDeterminePutFormatterLocation(self) -> None:

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

1155 can not determine the put formatter location.

1156 """

1157 _ = makeExampleMetrics()

1158 datastore = self.makeDatastore()

1159

1160 # Create multiple storage classes for testing different formulations

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

1162

1163 sccomp = StorageClass("Dummy")

1164 compositeStorageClass = StorageClass(

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

1166 )

1167

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

1169 dataId = {

1170 "instrument": "dummy",

1171 "visit": 52,

1172 "physical_filter": "V",

1173 "band": "v",

1174 "day_obs": 20250101,

1175 }

1176

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

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

1179

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

1181 raise DatasetTypeNotSupportedError()

1182

1183 with unittest.mock.patch.object(

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

1185 "_determine_put_formatter_location",

1186 side_effect=raiser,

1187 ):

1188 # verify the non-composite ref execution path:

1189 with self.assertRaises(DatasetTypeNotSupportedError):

1190 datastore.getURIs(ref, predict=True)

1191

1192 # verify the composite-ref execution path:

1193 with self.assertRaises(DatasetTypeNotSupportedError):

1194 datastore.getURIs(compRef, predict=True)

1195

1196 def test_roots(self):

1197 datastore = self.makeDatastore()

1198

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

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

1201 if root is not None:

1202 self.assertTrue(root.exists())

1203

1204 def test_prepare_get_for_external_client(self):

1205 datastore = self.makeDatastore()

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

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

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

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

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

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

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

1213

1214

1215class PosixDatastoreNoChecksumsTestCase(PosixDatastoreTestCase):

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

1217

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

1219

1220 def testChecksum(self) -> None:

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

1222 datastore = self.makeDatastore()

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

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

1225 metrics = makeExampleMetrics()

1226

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

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

1229

1230 # Configuration should have disabled checksum calculation

1231 datastore.put(metrics, ref)

1232 infos = datastore.getStoredItemsInfo(ref)

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

1234

1235 # Remove put back but with checksums enabled explicitly

1236 datastore.remove(ref)

1237 datastore.useChecksum = True

1238 datastore.put(metrics, ref)

1239

1240 infos = datastore.getStoredItemsInfo(ref)

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

1242

1243 def test_repeat_ingest(self):

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

1245 is allowed.

1246

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

1248 supporting "direct" ingest.

1249 """

1250 metrics, v4ref = self._prepareIngestTest()

1251 datastore = self.makeDatastore()

1252 v5ref = DatasetRef(

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

1254 )

1255

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

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

1258

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

1260 with self.assertRaises(RuntimeError):

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

1262

1263 # UUIDv5 can be repeatedly ingested in direct mode.

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

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

1266

1267 with self.assertRaises(RuntimeError):

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

1269

1270

1271class TrashDatastoreTestCase(PosixDatastoreTestCase):

1272 """Restrict trash test to FileDatastore."""

1273

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

1275

1276 def testTrash(self) -> None:

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

1278

1279 # Trash one of them.

1280 ref = refs.pop()

1281 uri = datastore.getURI(ref)

1282 datastore.trash(ref)

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

1284 datastore.emptyTrash()

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

1286

1287 # Trash it again should be fine.

1288 datastore.trash(ref)

1289

1290 # Trash multiple items at once.

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

1292 datastore.trash(subset)

1293 datastore.emptyTrash()

1294

1295 # Remove a record and trash should do nothing.

1296 # This is execution butler scenario.

1297 ref = refs.pop()

1298 uri = datastore.getURI(ref)

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

1300 self.assertTrue(uri.exists())

1301 datastore.trash(ref)

1302 datastore.emptyTrash()

1303 self.assertTrue(uri.exists())

1304

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

1306 datastore.trustGetRequest = True

1307 datastore.trash([ref])

1308 self.assertFalse(uri.exists())

1309

1310 # Remove multiples at once in trust mode.

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

1312 datastore.trash(subset)

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

1314

1315 def test_empty_trash(self) -> None:

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

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

1318

1319 # Trash one of them.

1320 ref = refs.pop()

1321 uri = datastore.getURI(ref)

1322 datastore.trash(ref)

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

1324

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

1326 # one in the trash table.

1327 removed = datastore.emptyTrash(refs=refs)

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

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

1330

1331 # Empty the entire trash but in dry_run mode.

1332 removed = datastore.emptyTrash(dry_run=True)

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

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

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

1336

1337 # Empty the trash specifying the actual ref.

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

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

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

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

1342

1343 # Trash everything and empty.

1344 datastore.trash(refs)

1345 removed = datastore.emptyTrash(dry_run=True)

1346 for u in removed:

1347 self.assertTrue(u.exists())

1348 removed = datastore.emptyTrash()

1349 for u in removed:

1350 self.assertFalse(u.exists())

1351

1352

1353class CleanupPosixDatastoreTestCase(DatastoreTestsBase, unittest.TestCase):

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

1355

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

1357

1358 def setUp(self) -> None:

1359 # Override the working directory before calling the base class

1360 self.root = tempfile.mkdtemp()

1361 super().setUp()

1362

1363 def testCleanup(self) -> None:

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

1365 metrics = makeExampleMetrics()

1366 datastore = self.makeDatastore()

1367

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

1369

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

1371 dataId = {

1372 "instrument": "dummy",

1373 "visit": 52,

1374 "physical_filter": "V",

1375 "band": "v",

1376 "day_obs": 20250101,

1377 }

1378

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

1380

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

1382 # the same file extension)

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

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

1385

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

1387

1388 # Try formatter that fails and formatter that fails and leaves

1389 # a file behind

1390 for formatter in (BadWriteFormatter, BadNoWriteFormatter):

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

1392 # Monkey patch the formatter

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

1394

1395 # Try to put the dataset, it should fail

1396 with self.assertRaises(RuntimeError):

1397 datastore.put(metrics, ref)

1398

1399 # Check that there is no file on disk

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

1401

1402 # Check that there is a directory

1403 dir = expectedUri.dirname()

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

1405

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

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

1408 datastore.put(metrics, ref)

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

1410 datastore.remove(ref)

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

1412

1413

1414class InMemoryDatastoreTestCase(DatastoreTests, unittest.TestCase):

1415 """PosixDatastore specialization"""

1416

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

1418 uriScheme = "mem"

1419 hasUnsupportedPut = False

1420 ingestTransferModes = ()

1421 isEphemeral = True

1422 rootKeys = None

1423 validationCanFail = False

1424

1425

1426class ChainedDatastoreTestCase(PosixDatastoreTestCase):

1427 """ChainedDatastore specialization using a POSIXDatastore"""

1428

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

1430 hasUnsupportedPut = False

1431 canIngestNoTransferAuto = False

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

1433 isEphemeral = False

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

1435 validationCanFail = True

1436

1437

1438class ChainedDatastoreMemoryTestCase(InMemoryDatastoreTestCase):

1439 """ChainedDatastore specialization using all InMemoryDatastore"""

1440

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

1442 validationCanFail = False

1443 isEphemeral = True

1444

1445

1446class DatastoreConstraintsTests(DatastoreTestsBase):

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

1448

1449 def testConstraints(self) -> None:

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

1451 same constraints.

1452 """

1453 metrics = makeExampleMetrics()

1454 datastore = self.makeDatastore()

1455

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

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

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

1459 dataId = {

1460 "visit": 52,

1461 "physical_filter": "V",

1462 "band": "v",

1463 "instrument": "DummyCamComp",

1464 "day_obs": 20250101,

1465 }

1466

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

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

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

1470 for datasetTypeName, sc, accepted in (

1471 ("metric", sc1, True),

1472 ("metric5", sc1, False),

1473 ("metric33", sc1, True),

1474 ("metric5", sc2, True),

1475 ):

1476 # Choose different temp file depending on StorageClass

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

1478

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

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

1481 if accepted:

1482 datastore.put(metrics, ref)

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

1484 datastore.remove(ref)

1485

1486 # Try ingest

1487 if self.canIngest:

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

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

1490 datastore.remove(ref)

1491 else:

1492 with self.assertRaises(DatasetTypeNotSupportedError):

1493 datastore.put(metrics, ref)

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

1495

1496 # Again with ingest

1497 if self.canIngest:

1498 with self.assertRaises(DatasetTypeNotSupportedError):

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

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

1501

1502

1503class PosixDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):

1504 """PosixDatastore specialization"""

1505

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

1507 canIngest = True

1508

1509 def setUp(self) -> None:

1510 # Override the working directory before calling the base class

1511 self.root = tempfile.mkdtemp()

1512 super().setUp()

1513

1514

1515class InMemoryDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):

1516 """InMemoryDatastore specialization."""

1517

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

1519 canIngest = False

1520

1521

1522class ChainedDatastoreConstraintsNativeTestCase(PosixDatastoreConstraintsTestCase):

1523 """ChainedDatastore specialization using a POSIXDatastore and constraints

1524 at the ChainedDatstore.

1525 """

1526

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

1528

1529

1530class ChainedDatastoreConstraintsTestCase(PosixDatastoreConstraintsTestCase):

1531 """ChainedDatastore specialization using a POSIXDatastore."""

1532

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

1534

1535

1536class ChainedDatastoreMemoryConstraintsTestCase(InMemoryDatastoreConstraintsTestCase):

1537 """ChainedDatastore specialization using all InMemoryDatastore."""

1538

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

1540 canIngest = False

1541

1542

1543class ChainedDatastorePerStoreConstraintsTests(DatastoreTestsBase, unittest.TestCase):

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

1545 even if child datastore would accept.

1546 """

1547

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

1549

1550 def setUp(self) -> None:

1551 # Override the working directory before calling the base class

1552 self.root = tempfile.mkdtemp()

1553 super().setUp()

1554

1555 def testConstraints(self) -> None:

1556 """Test chained datastore constraints model."""

1557 metrics = makeExampleMetrics()

1558 datastore = self.makeDatastore()

1559

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

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

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

1563 dataId1 = {

1564 "visit": 52,

1565 "physical_filter": "V",

1566 "band": "v",

1567 "instrument": "DummyCamComp",

1568 "day_obs": 20250101,

1569 }

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

1571

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

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

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

1575

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

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

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

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

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

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

1582 ):

1583 # Choose different temp file depending on StorageClass

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

1585

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

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

1588 if any(accept):

1589 datastore.put(metrics, ref)

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

1591

1592 # Check each datastore inside the chained datastore

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

1594 self.assertEqual(

1595 childDatastore.exists(ref),

1596 expected,

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

1598 )

1599

1600 datastore.remove(ref)

1601

1602 # Check that ingest works

1603 if ingest:

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

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

1606

1607 # Check each datastore inside the chained datastore

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

1609 # Ephemeral datastores means InMemory at the moment

1610 # and that does not accept ingest of files.

1611 if childDatastore.isEphemeral:

1612 expected = False

1613 self.assertEqual(

1614 childDatastore.exists(ref),

1615 expected,

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

1617 )

1618

1619 datastore.remove(ref)

1620 else:

1621 with self.assertRaises(DatasetTypeNotSupportedError):

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

1623

1624 else:

1625 with self.assertRaises(DatasetTypeNotSupportedError):

1626 datastore.put(metrics, ref)

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

1628

1629 # Again with ingest

1630 with self.assertRaises(DatasetTypeNotSupportedError):

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

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

1633

1634

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

1636class DatastoreCacheTestCase(DatasetTestHelper, unittest.TestCase):

1637 """Tests for datastore caching infrastructure."""

1638

1639 @classmethod

1640 def setUpClass(cls) -> None:

1641 cls.storageClassFactory = StorageClassFactory()

1642 cls.universe = DimensionUniverse()

1643

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

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

1646 cls.storageClassFactory.addFromConfig(scConfigFile)

1647

1648 def setUp(self) -> None:

1649 self.id = 0

1650

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

1652 self.root = tempfile.mkdtemp()

1653

1654 # Create some test dataset refs and associated test files

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

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

1657 dataId = {

1658 "instrument": "dummy",

1659 "visit": 52,

1660 "physical_filter": "V",

1661 "band": "v",

1662 "day_obs": 20250101,

1663 }

1664

1665 # Create list of refs and list of temporary files

1666 n_datasets = 10

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

1668

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

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

1671

1672 # Create test files.

1673 for uri in self.files:

1674 uri.write(b"0123456789")

1675

1676 # Create some composite refs with component files.

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

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

1679 self.comp_files = []

1680 self.comp_refs = []

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

1682 component_refs = []

1683 component_files = []

1684 for component in sc.components:

1685 component_ref = ref.makeComponentRef(component)

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

1687 component_refs.append(component_ref)

1688 component_files.append(file)

1689 file.write(b"9876543210")

1690

1691 self.comp_files.append(component_files)

1692 self.comp_refs.append(component_refs)

1693

1694 def tearDown(self) -> None:

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

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

1697

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

1699 config = Config.fromYaml(config_str)

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

1701

1702 def testNoCacheDir(self) -> None:

1703 config_str = """

1704cached:

1705 root: null

1706 cacheable:

1707 metric0: true

1708 """

1709 cache_manager = self._make_cache_manager(config_str)

1710

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

1712 self.assertIsNone(cache_manager._cache_directory)

1713

1714 self.assertCache(cache_manager)

1715

1716 # Test that the cache directory is marked temporary

1717 self.assertTrue(cache_manager.cache_directory.isTemporary)

1718

1719 def testNoCacheDirReversed(self) -> None:

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

1721 config_str = """

1722cached:

1723 root: null

1724 default: true

1725 cacheable:

1726 metric1: false

1727 """

1728 cache_manager = self._make_cache_manager(config_str)

1729

1730 self.assertCache(cache_manager)

1731

1732 def testEnvvarCacheDir(self) -> None:

1733 config_str = f"""

1734cached:

1735 root: '{self.root}'

1736 cacheable:

1737 metric0: true

1738 """

1739

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

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

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

1743

1744 # Environment variable should override the config value.

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

1746 cache_manager = self._make_cache_manager(config_str)

1747 self.assertEqual(cache_manager.cache_directory, env_dir)

1748

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

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

1751 cache_manager = self._make_cache_manager(config_str)

1752 self.assertEqual(cache_manager.cache_directory, root)

1753

1754 # No default setting.

1755 config_str = """

1756cached:

1757 root: null

1758 default: true

1759 cacheable:

1760 metric1: false

1761 """

1762 cache_manager = self._make_cache_manager(config_str)

1763

1764 # This environment variable should override the config value.

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

1766 cache_manager = self._make_cache_manager(config_str)

1767 self.assertEqual(cache_manager.cache_directory, env_dir)

1768

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

1770 # variable should win.

1771 with unittest.mock.patch.dict(

1772 os.environ,

1773 {

1774 "DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath,

1775 "DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": elsewhere.ospath,

1776 },

1777 ):

1778 cache_manager = self._make_cache_manager(config_str)

1779 self.assertEqual(cache_manager.cache_directory, env_dir)

1780

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

1782 # variable is reset on exit.

1783 with unittest.mock.patch.dict(

1784 os.environ,

1785 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""},

1786 ):

1787 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()

1788 self.assertTrue(defined)

1789 cache_manager = self._make_cache_manager(config_str)

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

1791

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

1793 # later.

1794 cache_manager = self._make_cache_manager(config_str)

1795 self.assertIsNone(cache_manager._cache_directory)

1796 with unittest.mock.patch.dict(

1797 os.environ,

1798 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""},

1799 ):

1800 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()

1801 self.assertTrue(defined)

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

1803

1804 def testExplicitCacheDir(self) -> None:

1805 config_str = f"""

1806cached:

1807 root: '{self.root}'

1808 cacheable:

1809 metric0: true

1810 """

1811 cache_manager = self._make_cache_manager(config_str)

1812

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

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

1815

1816 self.assertCache(cache_manager)

1817

1818 # Test that the cache directory is not marked temporary

1819 self.assertFalse(cache_manager.cache_directory.isTemporary)

1820

1821 def testUnexpectedFilesInCacheDir(self) -> None:

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

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

1824 """

1825 config_str = f"""

1826cached:

1827 root: '{self.root}'

1828 cacheable:

1829 metric0: true

1830 """

1831

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

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

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

1835 fh.write("test")

1836

1837 cache_manager = self._make_cache_manager(config_str)

1838 cache_manager.scan_cache()

1839 self.assertCache(cache_manager)

1840

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

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

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

1844

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

1846 self.assertIsInstance(uri, ResourcePath)

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

1848

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

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

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

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

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

1854

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

1856 # unaffected.

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

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

1859

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

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

1862 self.assertTrue(found.exists())

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

1864

1865 # Should not be able to find these in cache

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

1867 self.assertIsNone(found)

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

1869 self.assertIsNone(found)

1870

1871 def testNoCache(self) -> None:

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

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

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

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

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

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

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

1879

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

1881 return f"""

1882cached:

1883 default: true

1884 expiry:

1885 mode: {mode}

1886 threshold: {threshold}

1887 cacheable:

1888 unused: true

1889 """

1890

1891 def testCacheExpiryFiles(self) -> None:

1892 threshold = 2 # Keep at least 2 files.

1893 mode = "files"

1894 config_str = self._expiration_config(mode, threshold)

1895

1896 cache_manager = self._make_cache_manager(config_str)

1897

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

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

1900

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

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

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

1904

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

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

1907 self.assertIsNotNone(found)

1908

1909 # Trigger cache expiration that should remove the file

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

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

1912 self.assertIsNotNone(cached)

1913

1914 # Cache should still report the standard file count.

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

1916

1917 # Add additional entry to cache.

1918 # Should now have 4, 5, 6

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

1920 self.assertIsNotNone(cached)

1921

1922 # Is the file still there?

1923 self.assertTrue(found.exists())

1924

1925 # Can we read it?

1926 data = found.read()

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

1928

1929 # Outside context the file should no longer exist.

1930 self.assertFalse(found.exists())

1931

1932 # File count should not have changed.

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

1934

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

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

1937 # above and so should no longer be found.

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

1939 self.assertIsNone(found)

1940

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

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

1943 self.assertIsNone(found)

1944

1945 # But dataset 4 is present.

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

1947 self.assertIsNotNone(found)

1948

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

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

1951

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

1953 self.assertIsNone(found)

1954

1955 def testCacheExpiryDatasets(self) -> None:

1956 threshold = 2 # Keep 2 datasets.

1957 mode = "datasets"

1958 config_str = self._expiration_config(mode, threshold)

1959

1960 cache_manager = self._make_cache_manager(config_str)

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

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

1963

1964 def testCacheExpiryDatasetsFromDisabled(self) -> None:

1965 threshold = 2

1966 mode = "datasets"

1967 with unittest.mock.patch.dict(

1968 os.environ,

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

1970 ):

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

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

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

1974

1975 def testExpirationModeOverride(self) -> None:

1976 threshold = 2 # Keep 2 datasets.

1977 mode = "datasets"

1978 config_str = self._expiration_config(mode, threshold)

1979

1980 mode = "size"

1981 threshold = 55

1982 with unittest.mock.patch.dict(

1983 os.environ,

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

1985 ):

1986 cache_manager = self._make_cache_manager(config_str)

1987 self.assertExpiration(cache_manager, 10, 6)

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

1989

1990 # Check we get a warning with unrecognized form.

1991 with unittest.mock.patch.dict(

1992 os.environ,

1993 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something"},

1994 ):

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

1996 self._make_cache_manager(config_str)

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

1998

1999 with unittest.mock.patch.dict(

2000 os.environ,

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

2002 ):

2003 with self.assertRaises(ValueError) as cm:

2004 self._make_cache_manager(config_str)

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

2006

2007 def testMissingThreshold(self) -> None:

2008 threshold = ""

2009 mode = "datasets"

2010 config_str = self._expiration_config(mode, threshold)

2011

2012 with self.assertRaises(ValueError) as cm:

2013 self._make_cache_manager(config_str)

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

2015

2016 def testCacheExpiryDatasetsComposite(self) -> None:

2017 threshold = 2 # Keep 2 datasets.

2018 mode = "datasets"

2019 config_str = self._expiration_config(mode, threshold)

2020

2021 cache_manager = self._make_cache_manager(config_str)

2022

2023 n_datasets = 3

2024 for i in range(n_datasets):

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

2026 cached = cache_manager.move_to_cache(component_file, component_ref)

2027 self.assertIsNotNone(cached)

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

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

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

2031

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

2033

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

2035 self.assertExpiration(cache_manager, 2, 5)

2036

2037 def testCacheExpirySize(self) -> None:

2038 threshold = 55 # Each file is 10 bytes

2039 mode = "size"

2040 config_str = self._expiration_config(mode, threshold)

2041

2042 cache_manager = self._make_cache_manager(config_str)

2043 self.assertExpiration(cache_manager, 10, 6)

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

2045

2046 def testDisabledCache(self) -> None:

2047 # Configure an active cache but disable via environment.

2048 threshold = 2

2049 mode = "datasets"

2050 config_str = self._expiration_config(mode, threshold)

2051

2052 with unittest.mock.patch.dict(

2053 os.environ,

2054 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "disabled"},

2055 ):

2056 env_cache_manager = self._make_cache_manager(config_str)

2057

2058 # Configure to be disabled

2059 threshold = 0

2060 mode = "disabled"

2061 config_str = self._expiration_config(mode, threshold)

2062 cfg_cache_manager = self._make_cache_manager(config_str)

2063

2064 for cache_manager in (

2065 cfg_cache_manager,

2066 env_cache_manager,

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

2068 ):

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

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

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

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

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

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

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

2076

2077 def assertExpiration(

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

2079 ) -> None:

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

2081 for i in range(n_datasets):

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

2083 self.assertIsNotNone(cached)

2084

2085 self.assertEqual(cache_manager.file_count, n_retained)

2086

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

2088 for i in range(n_datasets):

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

2090 if i >= n_datasets - n_retained:

2091 self.assertIsInstance(found, ResourcePath)

2092 else:

2093 self.assertIsNone(found)

2094

2095 def testCacheExpiryAge(self) -> None:

2096 threshold = 1 # Expire older than 2 seconds

2097 mode = "age"

2098 config_str = self._expiration_config(mode, threshold)

2099

2100 cache_manager = self._make_cache_manager(config_str)

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

2102

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

2104 for i in range(2):

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

2106 self.assertIsNotNone(cached)

2107 time.sleep(2.0)

2108 for j in range(4):

2109 i = 2 + j # Continue the counting

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

2111 self.assertIsNotNone(cached)

2112

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

2114 self.assertEqual(cache_manager.file_count, 4)

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

2116 self.assertIsNone(found)

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

2118 self.assertIsInstance(found, ResourcePath)

2119

2120

2121class NullDatastoreTestCase(DatasetTestHelper, unittest.TestCase):

2122 """Test the null datastore."""

2123

2124 storageClassFactory = StorageClassFactory()

2125

2126 def test_basics(self) -> None:

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

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

2129

2130 null = NullDatastore(None, None)

2131

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

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

2134 knows = null.knows_these([ref])

2135 self.assertFalse(knows[ref])

2136 null.validateConfiguration(ref)

2137

2138 with self.assertRaises(FileNotFoundError):

2139 null.get(ref)

2140 with self.assertRaises(NotImplementedError):

2141 null.put("", ref)

2142 with self.assertRaises(FileNotFoundError):

2143 null.getURI(ref)

2144 with self.assertRaises(FileNotFoundError):

2145 null.getURIs(ref)

2146 with self.assertRaises(FileNotFoundError):

2147 null.getManyURIs([ref])

2148 with self.assertRaises(NotImplementedError):

2149 null.getLookupKeys()

2150 with self.assertRaises(NotImplementedError):

2151 null.import_records({})

2152 with self.assertRaises(NotImplementedError):

2153 null.export_records([])

2154 with self.assertRaises(NotImplementedError):

2155 null.export_predicted_records([])

2156 with self.assertRaises(NotImplementedError):

2157 null.export([ref])

2158 with self.assertRaises(NotImplementedError):

2159 null.transfer(null, ref)

2160 with self.assertRaises(NotImplementedError):

2161 null.emptyTrash()

2162 with self.assertRaises(NotImplementedError):

2163 null.trash(ref)

2164 with self.assertRaises(NotImplementedError):

2165 null.forget([ref])

2166 with self.assertRaises(NotImplementedError):

2167 null.remove(ref)

2168 with self.assertRaises(NotImplementedError):

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

2170 with self.assertRaises(NotImplementedError):

2171 null.transfer_from(null, [ref])

2172 with self.assertRaises(NotImplementedError):

2173 null.ingest()

2174

2175

2176class DatasetRefURIsTestCase(unittest.TestCase):

2177 """Tests for DatasetRefURIs."""

2178

2179 def testSequenceAccess(self) -> None:

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

2181 uris = DatasetRefURIs()

2182

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

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

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

2186

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

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

2189

2190 # affirm that DatasetRefURIs does not support MutableSequence functions

2191 with self.assertRaises(TypeError):

2192 uris[0] = primaryURI

2193 with self.assertRaises(TypeError):

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

2195

2196 # but DatasetRefURIs can be set by property name:

2197 uris.primaryURI = primaryURI

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

2199 self.assertEqual(uris.primaryURI, primaryURI)

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

2201

2202 primary, components = uris

2203 self.assertEqual(primary, primaryURI)

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

2205

2206 def testRepr(self) -> None:

2207 """Verify __repr__ output."""

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

2209 self.assertEqual(

2210 repr(uris),

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

2212 )

2213

2214

2215class StoredFileInfoTestCase(DatasetTestHelper, unittest.TestCase):

2216 """Test the StoredFileInfo class."""

2217

2218 storageClassFactory = StorageClassFactory()

2219

2220 def test_StoredFileInfo(self) -> None:

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

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

2223

2224 record = dict(

2225 storage_class="StructuredDataDict",

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

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

2228 component="component",

2229 checksum=None,

2230 file_size=5,

2231 )

2232 info = StoredFileInfo.from_record(record)

2233

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

2235

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

2237 rebased = info.rebase(ref2)

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

2239

2240 with self.assertRaises(TypeError):

2241 rebased.update(formatter=42)

2242

2243 with self.assertRaises(ValueError):

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

2245

2246 # Check that pickle works on StoredFileInfo.

2247 pickled_info = pickle.dumps(info)

2248 unpickled_info = pickle.loads(pickled_info)

2249 self.assertEqual(unpickled_info, info)

2250

2251 def test_make_datastore_path_relative(self):

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

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

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

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

2256

2257 def test_datastore_record_data_json_types(self):

2258 """Test that we don't round-trip checksums to UUIDs when deserializing

2259 datastore record data.

2260 """

2261 test_json = """

2262 {

2263 "dataset_ids": [

2264 "74478304-abf1-4a9c-9eb2-926090a84446"

2265 ],

2266 "records": {

2267 "lsst.daf.butler.datastore.stored_file_info.StoredFileInfo": {

2268 "74478304abf14a9c9eb2926090a84446": {

2269 "file_datastore_records": [

2270 {

2271 "formatter": "lsst.daf.butler.formatters.yaml.YamlFormatter",

2272 "path": "gain_factors/base-2025-158/gain_factors_spx_base-2025-158.yaml",

2273 "storage_class": "GainFactors",

2274 "component": "__NULL_STRING__",

2275 "checksum": "cab515f6-ab67-0484-393f-aaa525dd526f",

2276 "file_size": 5412

2277 }

2278 ]

2279 }

2280 }

2281 }

2282 }

2283 """

2284 id_str = "74478304abf14a9c9eb2926090a84446"

2285 s = SerializedDatastoreRecordData.model_validate_json(test_json)

2286 self.assertIsInstance(

2287 s.records[get_full_type_name(StoredFileInfo)][id_str]["file_datastore_records"][0]["checksum"],

2288 str,

2289 )

2290 id = uuid.UUID(id_str)

2291 d = DatastoreRecordData.from_simple(s)

2292 self.assertIsInstance(d.records[id]["file_datastore_records"][0], StoredFileInfo)

2293 self.assertIsInstance(d.records[id]["file_datastore_records"][0].checksum, str)

2294

2295

2296@contextlib.contextmanager

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

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

2299 try:

2300 yaml.dump(data, stream=fh)

2301 fh.flush()

2302 yield fh.name

2303 finally:

2304 # Some tests delete the file

2305 with contextlib.suppress(FileNotFoundError):

2306 fh.close()

2307

2308

2309if __name__ == "__main__":

2310 unittest.main()

Coverage for tests / test_datastore.py: 13%

1211 statements