Coverage for tests / test_datastore.py: 13%
1211 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:11 -0800
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:11 -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/>.
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 """
104 pass
107class DatastoreTestsBase(DatasetTestHelper, DatastoreTestHelper, TestCaseMixin):
108 """Support routines for datastore testing"""
110 root: str | None = None
111 universe: DimensionUniverse
112 storageClassFactory: StorageClassFactory
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)
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()
128 def setUp(self) -> None:
129 self.setUpDatastoreTests(DummyRegistry, DatastoreConfig)
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)
136class DatastoreTests(DatastoreTestsBase):
137 """Some basic tests of a simple datastore."""
139 hasUnsupportedPut = True
140 rootKeys: tuple[str, ...] | None = None
141 isEphemeral: bool = False
142 validationCanFail: bool = False
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])
153 def testConstructor(self) -> None:
154 datastore = self.makeDatastore()
155 self.assertIsNotNone(datastore)
156 self.assertIs(datastore.isEphemeral, self.isEphemeral)
158 def testConfigurationValidation(self) -> None:
159 datastore = self.makeDatastore()
160 sc = self.storageClassFactory.getStorageClass("ThingOne")
161 datastore.validateConfiguration([sc])
163 sc2 = self.storageClassFactory.getStorageClass("ThingTwo")
164 if self.validationCanFail:
165 with self.assertRaises(DatastoreValidationError):
166 datastore.validateConfiguration([sc2], logFailures=True)
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])
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})
199 def testBasicPutGet(self) -> None:
200 metrics = makeExampleMetrics()
201 datastore = self.makeDatastore()
203 # Create multiple storage classes for testing different formulations
204 storageClasses = [
205 self.storageClassFactory.getStorageClass(sc)
206 for sc in ("StructuredData", "StructuredDataJson", "StructuredDataPickle")
207 ]
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 }
225 for sc in storageClasses:
226 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
227 ref2 = self.makeDatasetRef("metric", dimensions, sc, dataId2)
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)
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())
242 datastore.put(metrics, ref)
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])
253 # Get
254 metricsOut = datastore.get(ref, parameters=None)
255 self.assertEqual(metrics, metricsOut)
257 uri = datastore.getURI(ref)
258 self.assertEqual(uri.scheme, self.uriScheme)
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)
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))
273 uri = datastore.getURI(compRef)
274 self.assertEqual(uri.scheme, self.uriScheme)
276 uris = datastore.getManyURIs([compRef])
277 self.assertEqual(len(uris), 1)
279 storageClass = sc
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)
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))
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)
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)
313 # Get a URI from it
314 uri = datastore.getURI(ref, predict=True)
315 self.assertEqual(uri.scheme, self.uriScheme)
317 with self.assertRaises(FileNotFoundError):
318 datastore.getURI(ref)
320 def testTrustGetRequest(self) -> None:
321 """Check that we can get datasets that registry knows nothing about."""
322 datastore = self.makeDatastore()
324 # Skip test if the attribute is not defined
325 if not hasattr(datastore, "trustGetRequest"):
326 return
328 metrics = makeExampleMetrics()
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.
335 if sc_name == "StructuredComposite":
336 disassembled = True
337 else:
338 disassembled = False
340 # Start datastore in default configuration of using registry
341 datastore.trustGetRequest = False
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"))
348 dataId = {
349 "instrument": "dummy",
350 "visit": 52 + i,
351 "physical_filter": "V",
352 "band": "v",
353 "day_obs": 20250101,
354 }
356 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId)
357 datastore.put(metrics, ref)
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])
367 # Get
368 metricsOut = datastore.get(ref)
369 self.assertEqual(metrics, metricsOut)
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)
380 # Delete registry entry so now we are trusting
381 datastore.removeStoredItemInfo(ref)
383 # Now stop trusting and check that things break
384 datastore.trustGetRequest = False
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])
394 with self.assertRaises(FileNotFoundError):
395 datastore.get(ref)
397 if sc_name != "StructuredDataNoComponents":
398 with self.assertRaises(FileNotFoundError):
399 datastore.get(ref.makeComponentRef("data"))
401 # URI should fail unless we ask for prediction
402 with self.assertRaises(FileNotFoundError):
403 datastore.getURIs(ref)
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")
417 # Now enable registry-free trusting mode
418 datastore.trustGetRequest = True
420 # Try again to get it
421 metricsOut = datastore.get(ref)
422 self.assertEqual(metricsOut, metrics)
424 # Does it exist?
425 self.assertTrue(datastore.exists(ref))
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))
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)
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")
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")
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}")
459 datastore.set_retrieve_dataset_type_method(_stored_dataset_type)
461 # Storage class override with original dataset ref.
462 metrics_as_dict = datastore.get(ref, storageClass="StructuredDataDictJson")
463 self.assertIsInstance(metrics_as_dict, dict)
465 # get() should return a dict now.
466 metrics_as_dict = datastore.get(ref_comp)
467 self.assertIsInstance(metrics_as_dict, dict)
469 # exists() should work as well.
470 self.assertTrue(datastore.exists(ref_comp))
472 datastore.set_retrieve_dataset_type_method(None)
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
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 ]
496 # Create the test datastore
497 datastore = self.makeDatastore()
499 # Dummy dataId
500 dimensions = self.universe.conform(("visit", "physical_filter"))
501 dataId = {"instrument": "dummy", "visit": 428, "physical_filter": "R"}
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)
510 disassembled = sc.name not in {"StructuredData", "StructuredCompositeReadCompNoDisassembly"}
512 datastore.put(metrics, ref)
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, {})
522 metrics_get = datastore.get(ref)
523 self.assertEqual(metrics_get, metrics)
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])
532 # Retrieve a component
533 data = datastore.get(ref.makeComponentRef("data"))
534 self.assertEqual(data, metrics.data)
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))
543 counter = datastore.get(cRef, parameters={"slice": slice(stop)})
544 self.assertEqual(counter, stop)
546 datastore.remove(ref)
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)
566 # Does it exist?
567 self.assertTrue(datastore.exists(ref))
569 # Get
570 metricsOut = datastore.get(ref)
571 self.assertEqual(metrics, metricsOut)
572 refs.append(ref)
574 return datastore, *refs
576 def testRemove(self) -> None:
577 datastore, ref = self.prepDeleteTest()
579 # Remove
580 datastore.remove(ref)
582 # Does it exist?
583 self.assertFalse(datastore.exists(ref))
585 # Do we now get a predicted URI?
586 uri = datastore.getURI(ref, predict=True)
587 self.assertEqual(uri.fragment, "predicted")
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)
596 def testForget(self) -> None:
597 datastore, ref = self.prepDeleteTest()
599 # Remove
600 datastore.forget([ref])
602 # Does it exist (as far as we know)?
603 self.assertFalse(datastore.exists(ref))
605 # Do we now get a predicted URI?
606 uri = datastore.getURI(ref, predict=True)
607 self.assertEqual(uri.fragment, "predicted")
609 # Get should now fail
610 with self.assertRaises(FileNotFoundError):
611 datastore.get(ref)
613 # Forgetting again is a silent no-op
614 datastore.forget([ref])
616 # Predicted URI should still point to the file.
617 self.assertTrue(uri.exists())
619 def testTransfer(self) -> None:
620 metrics = makeExampleMetrics()
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 }
631 sc = self.storageClassFactory.getStorageClass("StructuredData")
632 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
634 inputDatastore = self.makeDatastore("test_input_datastore")
635 outputDatastore = self.makeDatastore("test_output_datastore")
637 inputDatastore.put(metrics, ref)
638 outputDatastore.transfer(inputDatastore, ref)
640 metricsOut = outputDatastore.get(ref)
641 self.assertEqual(metrics, metricsOut)
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)
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()
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)
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
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)
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
760 with self.subTest(mode=mode):
761 datastore = self.makeDatastore()
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))
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))
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))
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)
829 if mode in self.ingestTransferModes:
830 self.runIngestTest(failOutsideRoot)
831 self.runIngestTest(failInputDoesNotExist)
832 self.runIngestTest(succeed)
833 else:
834 self.runIngestTest(failNotImplemented)
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)
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)
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}")
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)
886 if mode in self.ingestTransferModes:
887 self.runIngestTest(failInputDoesNotExist)
888 self.runIngestTest(succeed)
889 else:
890 self.runIngestTest(failNotImplemented)
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
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)
908 datastore = self.makeDatastore()
909 datastore.ingest(FileDataset(path=os.path.abspath(sympath), refs=ref), transfer=mode)
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")
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))
921 # Check that we can get the dataset back regardless of mode
922 metric2 = datastore.get(ref)
923 self.assertEqual(metric2, metrics)
925 # Cleanup the file for next time round loop
926 # since it will get the same file name in store
927 datastore.remove(ref)
929 def _populate_export_datastore(self, name: str) -> tuple[Datastore, list[DatasetRef]]:
930 datastore = self.makeDatastore(name)
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")
941 metrics = makeExampleMetrics()
942 dimensions = self.universe.conform(("visit", "physical_filter"))
943 sc = self.storageClassFactory.getStorageClass("StructuredData")
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
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)
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)
980 # Use the same datastore name to import relative path.
981 datastore2 = self.makeDatastore("test_datastore")
983 records = datastore.export_records(refs[1:])
984 datastore2.import_records(records)
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)
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)
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])
1011 # Expect predicted records from all datastores.
1012 self.assertEqual(set(records.keys()), names)
1014 for record_data in records.values():
1015 self.assertEqual(len(record_data.records), 1)
1017 def testExport(self) -> None:
1018 datastore, refs = self._populate_export_datastore("test_datastore")
1020 datasets = list(datastore.export(refs))
1021 self.assertEqual(len(datasets), 3)
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)
1028 with self.assertRaises(TypeError):
1029 list(datastore.export(refs, transfer="copy"))
1031 with self.assertRaises(TypeError):
1032 list(datastore.export(refs, directory="exportDir", transfer="move"))
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))
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)
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)
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)
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)
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)
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 }
1106 for ref in refs.values():
1107 datastore.put(data, ref)
1109 self.assertEqual(datastore.get(refs["A"]), datastore.get(refs["B"]))
1112class PosixDatastoreTestCase(DatastoreTests, unittest.TestCase):
1113 """PosixDatastore specialization"""
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
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()
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()
1142 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
1143 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
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])
1150 # And the transfer should be file to file.
1151 self.assertEqual(move_logs[0].count("file://"), 2)
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()
1160 # Create multiple storage classes for testing different formulations
1161 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1163 sccomp = StorageClass("Dummy")
1164 compositeStorageClass = StorageClass(
1165 "StructuredComposite", components={"dummy": sccomp, "dummy2": sccomp}
1166 )
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 }
1177 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1178 compRef = self.makeDatasetRef("metric", dimensions, compositeStorageClass, dataId)
1180 def raiser(ref: DatasetRef) -> None:
1181 raise DatasetTypeNotSupportedError()
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)
1192 # verify the composite-ref execution path:
1193 with self.assertRaises(DatasetTypeNotSupportedError):
1194 datastore.getURIs(compRef, predict=True)
1196 def test_roots(self):
1197 datastore = self.makeDatastore()
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())
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))
1215class PosixDatastoreNoChecksumsTestCase(PosixDatastoreTestCase):
1216 """Posix datastore tests but with checksums disabled."""
1218 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreNoChecksums.yaml")
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()
1227 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
1228 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1230 # Configuration should have disabled checksum calculation
1231 datastore.put(metrics, ref)
1232 infos = datastore.getStoredItemsInfo(ref)
1233 self.assertIsNone(infos[0].checksum)
1235 # Remove put back but with checksums enabled explicitly
1236 datastore.remove(ref)
1237 datastore.useChecksum = True
1238 datastore.put(metrics, ref)
1240 infos = datastore.getStoredItemsInfo(ref)
1241 self.assertIsNotNone(infos[0].checksum)
1243 def test_repeat_ingest(self):
1244 """Test that repeatedly ingesting the same file in direct mode
1245 is allowed.
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 )
1256 with _temp_yaml_file(metrics._asdict()) as path:
1257 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct")
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")
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")
1267 with self.assertRaises(RuntimeError):
1268 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="copy")
1271class TrashDatastoreTestCase(PosixDatastoreTestCase):
1272 """Restrict trash test to FileDatastore."""
1274 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
1276 def testTrash(self) -> None:
1277 datastore, *refs = self.prepDeleteTest(n_refs=10)
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)
1287 # Trash it again should be fine.
1288 datastore.trash(ref)
1290 # Trash multiple items at once.
1291 subset = [refs.pop(), refs.pop()]
1292 datastore.trash(subset)
1293 datastore.emptyTrash()
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())
1305 # Switch on trust and it should delete the file.
1306 datastore.trustGetRequest = True
1307 datastore.trash([ref])
1308 self.assertFalse(uri.exists())
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
1315 def test_empty_trash(self) -> None:
1316 """Test parameters and return value for empty trash."""
1317 datastore, *refs = self.prepDeleteTest(n_refs=10)
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
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)
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)
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)
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())
1353class CleanupPosixDatastoreTestCase(DatastoreTestsBase, unittest.TestCase):
1354 """Test datastore cleans up on failure."""
1356 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
1358 def setUp(self) -> None:
1359 # Override the working directory before calling the base class
1360 self.root = tempfile.mkdtemp()
1361 super().setUp()
1363 def testCleanup(self) -> None:
1364 """Test that a failed formatter write does cleanup a partial file."""
1365 metrics = makeExampleMetrics()
1366 datastore = self.makeDatastore()
1368 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
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 }
1379 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
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")
1386 self.assertEqual(expectedUri.getExtension(), ".yaml", f"Is there a file extension in {expectedUri}")
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)
1395 # Try to put the dataset, it should fail
1396 with self.assertRaises(RuntimeError):
1397 datastore.put(metrics, ref)
1399 # Check that there is no file on disk
1400 self.assertFalse(expectedUri.exists(), f"Check for existence of {expectedUri}")
1402 # Check that there is a directory
1403 dir = expectedUri.dirname()
1404 self.assertTrue(dir.exists(), f"Check for existence of directory {dir}")
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}")
1414class InMemoryDatastoreTestCase(DatastoreTests, unittest.TestCase):
1415 """PosixDatastore specialization"""
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
1426class ChainedDatastoreTestCase(PosixDatastoreTestCase):
1427 """ChainedDatastore specialization using a POSIXDatastore"""
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
1438class ChainedDatastoreMemoryTestCase(InMemoryDatastoreTestCase):
1439 """ChainedDatastore specialization using all InMemoryDatastore"""
1441 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2.yaml")
1442 validationCanFail = False
1443 isEphemeral = True
1446class DatastoreConstraintsTests(DatastoreTestsBase):
1447 """Basic tests of constraints model of Datastores."""
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()
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 }
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
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)
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))
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))
1503class PosixDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):
1504 """PosixDatastore specialization"""
1506 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreP.yaml")
1507 canIngest = True
1509 def setUp(self) -> None:
1510 # Override the working directory before calling the base class
1511 self.root = tempfile.mkdtemp()
1512 super().setUp()
1515class InMemoryDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):
1516 """InMemoryDatastore specialization."""
1518 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastoreP.yaml")
1519 canIngest = False
1522class ChainedDatastoreConstraintsNativeTestCase(PosixDatastoreConstraintsTestCase):
1523 """ChainedDatastore specialization using a POSIXDatastore and constraints
1524 at the ChainedDatstore.
1525 """
1527 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePa.yaml")
1530class ChainedDatastoreConstraintsTestCase(PosixDatastoreConstraintsTestCase):
1531 """ChainedDatastore specialization using a POSIXDatastore."""
1533 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastoreP.yaml")
1536class ChainedDatastoreMemoryConstraintsTestCase(InMemoryDatastoreConstraintsTestCase):
1537 """ChainedDatastore specialization using all InMemoryDatastore."""
1539 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2P.yaml")
1540 canIngest = False
1543class ChainedDatastorePerStoreConstraintsTests(DatastoreTestsBase, unittest.TestCase):
1544 """Test that a chained datastore can control constraints per-datastore
1545 even if child datastore would accept.
1546 """
1548 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePb.yaml")
1550 def setUp(self) -> None:
1551 # Override the working directory before calling the base class
1552 self.root = tempfile.mkdtemp()
1553 super().setUp()
1555 def testConstraints(self) -> None:
1556 """Test chained datastore constraints model."""
1557 metrics = makeExampleMetrics()
1558 datastore = self.makeDatastore()
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}
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")
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
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))
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 )
1600 datastore.remove(ref)
1602 # Check that ingest works
1603 if ingest:
1604 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1605 self.assertTrue(datastore.exists(ref))
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 )
1619 datastore.remove(ref)
1620 else:
1621 with self.assertRaises(DatasetTypeNotSupportedError):
1622 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1624 else:
1625 with self.assertRaises(DatasetTypeNotSupportedError):
1626 datastore.put(metrics, ref)
1627 self.assertFalse(datastore.exists(ref))
1629 # Again with ingest
1630 with self.assertRaises(DatasetTypeNotSupportedError):
1631 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1632 self.assertFalse(datastore.exists(ref))
1635@unittest.mock.patch.dict(os.environ, {}, clear=True)
1636class DatastoreCacheTestCase(DatasetTestHelper, unittest.TestCase):
1637 """Tests for datastore caching infrastructure."""
1639 @classmethod
1640 def setUpClass(cls) -> None:
1641 cls.storageClassFactory = StorageClassFactory()
1642 cls.universe = DimensionUniverse()
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)
1648 def setUp(self) -> None:
1649 self.id = 0
1651 # Create a root that we can use for caching tests.
1652 self.root = tempfile.mkdtemp()
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 }
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)]
1669 root_uri = ResourcePath(self.root, forceDirectory=True)
1670 self.files = [root_uri.join(f"file{n}.txt") for n in range(n_datasets)]
1672 # Create test files.
1673 for uri in self.files:
1674 uri.write(b"0123456789")
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")
1691 self.comp_files.append(component_files)
1692 self.comp_refs.append(component_refs)
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)
1698 def _make_cache_manager(self, config_str: str) -> DatastoreCacheManager:
1699 config = Config.fromYaml(config_str)
1700 return DatastoreCacheManager(DatastoreCacheManagerConfig(config), universe=self.universe)
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)
1711 # Look inside to check we don't have a cache directory
1712 self.assertIsNone(cache_manager._cache_directory)
1714 self.assertCache(cache_manager)
1716 # Test that the cache directory is marked temporary
1717 self.assertTrue(cache_manager.cache_directory.isTemporary)
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)
1730 self.assertCache(cache_manager)
1732 def testEnvvarCacheDir(self) -> None:
1733 config_str = f"""
1734cached:
1735 root: '{self.root}'
1736 cacheable:
1737 metric0: true
1738 """
1740 root = ResourcePath(self.root, forceDirectory=True)
1741 env_dir = root.join("somewhere", forceDirectory=True)
1742 elsewhere = root.join("elsewhere", forceDirectory=True)
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)
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)
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)
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)
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)
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))
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))
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)
1813 # Look inside to check we do have a cache directory.
1814 self.assertEqual(cache_manager.cache_directory, ResourcePath(self.root, forceDirectory=True))
1816 self.assertCache(cache_manager)
1818 # Test that the cache directory is not marked temporary
1819 self.assertFalse(cache_manager.cache_directory.isTemporary)
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 """
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")
1837 cache_manager = self._make_cache_manager(config_str)
1838 cache_manager.scan_cache()
1839 self.assertCache(cache_manager)
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]))
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]))
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()))
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())
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))
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)
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}")
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 """
1891 def testCacheExpiryFiles(self) -> None:
1892 threshold = 2 # Keep at least 2 files.
1893 mode = "files"
1894 config_str = self._expiration_config(mode, threshold)
1896 cache_manager = self._make_cache_manager(config_str)
1898 # Check that an empty cache returns unknown for arbitrary ref
1899 self.assertFalse(cache_manager.known_to_cache(self.refs[0]))
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))
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)
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)
1914 # Cache should still report the standard file count.
1915 self.assertEqual(cache_manager.file_count, threshold + 1)
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)
1922 # Is the file still there?
1923 self.assertTrue(found.exists())
1925 # Can we read it?
1926 data = found.read()
1927 self.assertGreater(len(data), 0)
1929 # Outside context the file should no longer exist.
1930 self.assertFalse(found.exists())
1932 # File count should not have changed.
1933 self.assertEqual(cache_manager.file_count, threshold + 1)
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)
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)
1945 # But dataset 4 is present.
1946 with cache_manager.find_in_cache(self.refs[4], ".txt") as found:
1947 self.assertIsNotNone(found)
1949 # Adding a new dataset to the cache should now delete it.
1950 cache_manager.move_to_cache(self.files[7], self.refs[7])
1952 with cache_manager.find_in_cache(self.refs[2], ".txt") as found:
1953 self.assertIsNone(found)
1955 def testCacheExpiryDatasets(self) -> None:
1956 threshold = 2 # Keep 2 datasets.
1957 mode = "datasets"
1958 config_str = self._expiration_config(mode, threshold)
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))
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))
1975 def testExpirationModeOverride(self) -> None:
1976 threshold = 2 # Keep 2 datasets.
1977 mode = "datasets"
1978 config_str = self._expiration_config(mode, threshold)
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))
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])
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))
2007 def testMissingThreshold(self) -> None:
2008 threshold = ""
2009 mode = "datasets"
2010 config_str = self._expiration_config(mode, threshold)
2012 with self.assertRaises(ValueError) as cm:
2013 self._make_cache_manager(config_str)
2014 self.assertIn("Cache expiration threshold", str(cm.exception))
2016 def testCacheExpiryDatasetsComposite(self) -> None:
2017 threshold = 2 # Keep 2 datasets.
2018 mode = "datasets"
2019 config_str = self._expiration_config(mode, threshold)
2021 cache_manager = self._make_cache_manager(config_str)
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()))
2032 self.assertEqual(cache_manager.file_count, 6) # 2 datasets each of 3 files
2034 # Write two new non-composite and the number of files should drop.
2035 self.assertExpiration(cache_manager, 2, 5)
2037 def testCacheExpirySize(self) -> None:
2038 threshold = 55 # Each file is 10 bytes
2039 mode = "size"
2040 config_str = self._expiration_config(mode, threshold)
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))
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)
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)
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)
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))
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)
2085 self.assertEqual(cache_manager.file_count, n_retained)
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)
2095 def testCacheExpiryAge(self) -> None:
2096 threshold = 1 # Expire older than 2 seconds
2097 mode = "age"
2098 config_str = self._expiration_config(mode, threshold)
2100 cache_manager = self._make_cache_manager(config_str)
2101 self.assertIn(f"{mode}={threshold}", str(cache_manager))
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)
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)
2121class NullDatastoreTestCase(DatasetTestHelper, unittest.TestCase):
2122 """Test the null datastore."""
2124 storageClassFactory = StorageClassFactory()
2126 def test_basics(self) -> None:
2127 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
2128 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
2130 null = NullDatastore(None, None)
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)
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()
2176class DatasetRefURIsTestCase(unittest.TestCase):
2177 """Tests for DatasetRefURIs."""
2179 def testSequenceAccess(self) -> None:
2180 """Verify that DatasetRefURIs can be treated like a two-item tuple."""
2181 uris = DatasetRefURIs()
2183 self.assertEqual(len(uris), 2)
2184 self.assertEqual(uris[0], None)
2185 self.assertEqual(uris[1], {})
2187 primaryURI = ResourcePath("1/2/3")
2188 componentURI = ResourcePath("a/b/c")
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}
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)
2202 primary, components = uris
2203 self.assertEqual(primary, primaryURI)
2204 self.assertEqual(components, {"foo": componentURI})
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 )
2215class StoredFileInfoTestCase(DatasetTestHelper, unittest.TestCase):
2216 """Test the StoredFileInfo class."""
2218 storageClassFactory = StorageClassFactory()
2220 def test_StoredFileInfo(self) -> None:
2221 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
2222 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
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)
2234 self.assertEqual(info.to_record(), record)
2236 ref2 = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
2237 rebased = info.rebase(ref2)
2238 self.assertEqual(rebased.rebase(ref), info)
2240 with self.assertRaises(TypeError):
2241 rebased.update(formatter=42)
2243 with self.assertRaises(ValueError):
2244 rebased.update(something=42, new="42")
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)
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")
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)
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()
2309if __name__ == "__main__":
2310 unittest.main()