Coverage for tests / test_datastore.py: 13%
1200 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:12 -0800
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:12 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
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.stored_file_info import StoredFileInfo, make_datastore_path_relative
65from lsst.daf.butler.formatters.yaml import YamlFormatter
66from lsst.daf.butler.tests import (
67 BadNoWriteFormatter,
68 BadWriteFormatter,
69 DatasetTestHelper,
70 DatastoreTestHelper,
71 DummyRegistry,
72 MetricsExample,
73 MetricsExampleDataclass,
74 MetricsExampleModel,
75)
76from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel
77from lsst.daf.butler.tests.utils import TestCaseMixin
78from lsst.resources import ResourcePath
79from lsst.utils import doImport
81TESTDIR = os.path.dirname(__file__)
84def makeExampleMetrics(use_none: bool = False) -> MetricsExample:
85 """Make example dataset that can be stored in butler."""
86 if use_none:
87 array = None
88 else:
89 array = [563, 234, 456.7, 105, 2054, -1045]
90 return MetricsExample(
91 {"AM1": 5.2, "AM2": 30.6},
92 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
93 array,
94 )
97class TransactionTestError(Exception):
98 """Specific error for transactions, to prevent misdiagnosing
99 that might otherwise occur when a standard exception is used.
100 """
102 pass
105class DatastoreTestsBase(DatasetTestHelper, DatastoreTestHelper, TestCaseMixin):
106 """Support routines for datastore testing"""
108 root: str | None = None
109 universe: DimensionUniverse
110 storageClassFactory: StorageClassFactory
112 @classmethod
113 def setUpClass(cls) -> None:
114 # Storage Classes are fixed for all datastores in these tests
115 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml")
116 cls.storageClassFactory = StorageClassFactory()
117 cls.storageClassFactory.addFromConfig(scConfigFile)
119 # Read the Datastore config so we can get the class
120 # information (since we should not assume the constructor
121 # name here, but rely on the configuration file itself)
122 datastoreConfig = DatastoreConfig(cls.configFile)
123 cls.datastoreType = cast(type[Datastore], doImport(datastoreConfig["cls"]))
124 cls.universe = DimensionUniverse()
126 def setUp(self) -> None:
127 self.setUpDatastoreTests(DummyRegistry, DatastoreConfig)
129 def tearDown(self) -> None:
130 if self.root is not None and os.path.exists(self.root):
131 shutil.rmtree(self.root, ignore_errors=True)
134class DatastoreTests(DatastoreTestsBase):
135 """Some basic tests of a simple datastore."""
137 hasUnsupportedPut = True
138 rootKeys: tuple[str, ...] | None = None
139 isEphemeral: bool = False
140 validationCanFail: bool = False
142 def testConfigRoot(self) -> None:
143 full = DatastoreConfig(self.configFile)
144 config = DatastoreConfig(self.configFile, mergeDefaults=False)
145 newroot = "/random/location"
146 self.datastoreType.setConfigRoot(newroot, config, full)
147 if self.rootKeys:
148 for k in self.rootKeys:
149 self.assertIn(newroot, config[k])
151 def testConstructor(self) -> None:
152 datastore = self.makeDatastore()
153 self.assertIsNotNone(datastore)
154 self.assertIs(datastore.isEphemeral, self.isEphemeral)
156 def testConfigurationValidation(self) -> None:
157 datastore = self.makeDatastore()
158 sc = self.storageClassFactory.getStorageClass("ThingOne")
159 datastore.validateConfiguration([sc])
161 sc2 = self.storageClassFactory.getStorageClass("ThingTwo")
162 if self.validationCanFail:
163 with self.assertRaises(DatastoreValidationError):
164 datastore.validateConfiguration([sc2], logFailures=True)
166 dimensions = self.universe.conform(("visit", "physical_filter"))
167 dataId = {
168 "instrument": "dummy",
169 "visit": 52,
170 "physical_filter": "V",
171 "band": "v",
172 "day_obs": 20250101,
173 }
174 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
175 datastore.validateConfiguration([ref])
177 def testParameterValidation(self) -> None:
178 """Check that parameters are validated"""
179 sc = self.storageClassFactory.getStorageClass("ThingOne")
180 dimensions = self.universe.conform(("visit", "physical_filter"))
181 dataId = {
182 "instrument": "dummy",
183 "visit": 52,
184 "physical_filter": "V",
185 "band": "v",
186 "day_obs": 20250101,
187 }
188 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
189 datastore = self.makeDatastore()
190 data = {1: 2, 3: 4}
191 datastore.put(data, ref)
192 newdata = datastore.get(ref)
193 self.assertEqual(data, newdata)
194 with self.assertRaises(KeyError):
195 newdata = datastore.get(ref, parameters={"missing": 5})
197 def testBasicPutGet(self) -> None:
198 metrics = makeExampleMetrics()
199 datastore = self.makeDatastore()
201 # Create multiple storage classes for testing different formulations
202 storageClasses = [
203 self.storageClassFactory.getStorageClass(sc)
204 for sc in ("StructuredData", "StructuredDataJson", "StructuredDataPickle")
205 ]
207 dimensions = self.universe.conform(("visit", "physical_filter"))
208 dataId = {
209 "instrument": "dummy",
210 "visit": 52,
211 "physical_filter": "V",
212 "band": "v",
213 "day_obs": 20250101,
214 }
215 dataId2 = {
216 "instrument": "dummy",
217 "visit": 53,
218 "physical_filter": "V",
219 "band": "v",
220 "day_obs": 20250101,
221 }
223 for sc in storageClasses:
224 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
225 ref2 = self.makeDatasetRef("metric", dimensions, sc, dataId2)
227 # Make sure that using getManyURIs without predicting before the
228 # dataset has been put raises.
229 with self.assertRaises(FileNotFoundError):
230 datastore.getManyURIs([ref], predict=False)
232 # Make sure that using getManyURIs with predicting before the
233 # dataset has been put predicts the URI.
234 uris = datastore.getManyURIs([ref, ref2], predict=True)
235 self.assertIn("52", uris[ref].primaryURI.geturl())
236 self.assertIn("#predicted", uris[ref].primaryURI.geturl())
237 self.assertIn("53", uris[ref2].primaryURI.geturl())
238 self.assertIn("#predicted", uris[ref2].primaryURI.geturl())
240 datastore.put(metrics, ref)
242 # Does it exist?
243 self.assertTrue(datastore.exists(ref))
244 self.assertTrue(datastore.knows(ref))
245 multi = datastore.knows_these([ref])
246 self.assertTrue(multi[ref])
247 multi = datastore.mexists([ref, ref2])
248 self.assertTrue(multi[ref])
249 self.assertFalse(multi[ref2])
251 # Get
252 metricsOut = datastore.get(ref, parameters=None)
253 self.assertEqual(metrics, metricsOut)
255 uri = datastore.getURI(ref)
256 self.assertEqual(uri.scheme, self.uriScheme)
258 uris = datastore.getManyURIs([ref])
259 self.assertEqual(len(uris), 1)
260 ref, uri = uris.popitem()
261 self.assertTrue(uri.primaryURI.exists())
262 self.assertFalse(uri.componentURIs)
264 # Get a component -- we need to construct new refs for them
265 # with derived storage classes but with parent ID
266 for comp in ("data", "output"):
267 compRef = ref.makeComponentRef(comp)
268 output = datastore.get(compRef)
269 self.assertEqual(output, getattr(metricsOut, comp))
271 uri = datastore.getURI(compRef)
272 self.assertEqual(uri.scheme, self.uriScheme)
274 uris = datastore.getManyURIs([compRef])
275 self.assertEqual(len(uris), 1)
277 storageClass = sc
279 # Check that we can put a metric with None in a component and
280 # get it back as None
281 metricsNone = makeExampleMetrics(use_none=True)
282 dataIdNone = {
283 "instrument": "dummy",
284 "visit": 54,
285 "physical_filter": "V",
286 "band": "v",
287 "day_obs": 20250101,
288 }
289 refNone = self.makeDatasetRef("metric", dimensions, sc, dataIdNone)
290 datastore.put(metricsNone, refNone)
292 comp = "data"
293 for comp in ("data", "output"):
294 compRef = refNone.makeComponentRef(comp)
295 output = datastore.get(compRef)
296 self.assertEqual(output, getattr(metricsNone, comp))
298 # Check that a put fails if the dataset type is not supported
299 if self.hasUnsupportedPut:
300 sc = StorageClass("UnsupportedSC", pytype=type(metrics))
301 ref = self.makeDatasetRef("unsupportedType", dimensions, sc, dataId)
302 with self.assertRaises(DatasetTypeNotSupportedError):
303 datastore.put(metrics, ref)
305 # These should raise
306 ref = self.makeDatasetRef("metrics", dimensions, storageClass, dataId)
307 with self.assertRaises(FileNotFoundError):
308 # non-existing file
309 datastore.get(ref)
311 # Get a URI from it
312 uri = datastore.getURI(ref, predict=True)
313 self.assertEqual(uri.scheme, self.uriScheme)
315 with self.assertRaises(FileNotFoundError):
316 datastore.getURI(ref)
318 def testTrustGetRequest(self) -> None:
319 """Check that we can get datasets that registry knows nothing about."""
320 datastore = self.makeDatastore()
322 # Skip test if the attribute is not defined
323 if not hasattr(datastore, "trustGetRequest"):
324 return
326 metrics = makeExampleMetrics()
328 i = 0
329 for sc_name in ("StructuredDataNoComponents", "StructuredData", "StructuredComposite"):
330 i += 1
331 datasetTypeName = f"test_metric{i}" # Different dataset type name each time.
333 if sc_name == "StructuredComposite":
334 disassembled = True
335 else:
336 disassembled = False
338 # Start datastore in default configuration of using registry
339 datastore.trustGetRequest = False
341 # Create multiple storage classes for testing with or without
342 # disassembly
343 sc = self.storageClassFactory.getStorageClass(sc_name)
344 dimensions = self.universe.conform(("visit", "physical_filter"))
346 dataId = {
347 "instrument": "dummy",
348 "visit": 52 + i,
349 "physical_filter": "V",
350 "band": "v",
351 "day_obs": 20250101,
352 }
354 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId)
355 datastore.put(metrics, ref)
357 # Does it exist?
358 self.assertTrue(datastore.exists(ref))
359 self.assertTrue(datastore.knows(ref))
360 multi = datastore.knows_these([ref])
361 self.assertTrue(multi[ref])
362 multi = datastore.mexists([ref])
363 self.assertTrue(multi[ref])
365 # Get
366 metricsOut = datastore.get(ref)
367 self.assertEqual(metrics, metricsOut)
369 # Get the URI(s)
370 primaryURI, componentURIs = datastore.getURIs(ref)
371 if disassembled:
372 self.assertIsNone(primaryURI)
373 self.assertEqual(len(componentURIs), 3)
374 else:
375 self.assertIn(datasetTypeName, primaryURI.path)
376 self.assertFalse(componentURIs)
378 # Delete registry entry so now we are trusting
379 datastore.removeStoredItemInfo(ref)
381 # Now stop trusting and check that things break
382 datastore.trustGetRequest = False
384 # Does it exist?
385 self.assertFalse(datastore.exists(ref))
386 self.assertFalse(datastore.knows(ref))
387 multi = datastore.knows_these([ref])
388 self.assertFalse(multi[ref])
389 multi = datastore.mexists([ref])
390 self.assertFalse(multi[ref])
392 with self.assertRaises(FileNotFoundError):
393 datastore.get(ref)
395 if sc_name != "StructuredDataNoComponents":
396 with self.assertRaises(FileNotFoundError):
397 datastore.get(ref.makeComponentRef("data"))
399 # URI should fail unless we ask for prediction
400 with self.assertRaises(FileNotFoundError):
401 datastore.getURIs(ref)
403 predicted_primary, predicted_disassembled = datastore.getURIs(ref, predict=True)
404 if disassembled:
405 self.assertIsNone(predicted_primary)
406 self.assertEqual(len(predicted_disassembled), 3)
407 for uri in predicted_disassembled.values():
408 self.assertEqual(uri.fragment, "predicted")
409 self.assertIn(datasetTypeName, uri.path)
410 else:
411 self.assertIn(datasetTypeName, predicted_primary.path)
412 self.assertFalse(predicted_disassembled)
413 self.assertEqual(predicted_primary.fragment, "predicted")
415 # Now enable registry-free trusting mode
416 datastore.trustGetRequest = True
418 # Try again to get it
419 metricsOut = datastore.get(ref)
420 self.assertEqual(metricsOut, metrics)
422 # Does it exist?
423 self.assertTrue(datastore.exists(ref))
425 # Get a component
426 if sc_name != "StructuredDataNoComponents":
427 comp = "data"
428 compRef = ref.makeComponentRef(comp)
429 output = datastore.get(compRef)
430 self.assertEqual(output, getattr(metrics, comp))
432 # Get the URI -- if we trust this should work even without
433 # enabling prediction.
434 primaryURI2, componentURIs2 = datastore.getURIs(ref)
435 self.assertEqual(primaryURI2, primaryURI)
436 self.assertEqual(componentURIs2, componentURIs)
438 # Check for compatible storage class.
439 if sc_name in ("StructuredDataNoComponents", "StructuredData"):
440 # Make new dataset ref with compatible storage class.
441 ref_comp = ref.overrideStorageClass("StructuredDataDictJson")
443 # Without `set_retrieve_dataset_type_method` it will fail to
444 # find correct file.
445 self.assertFalse(datastore.exists(ref_comp))
446 with self.assertRaises(FileNotFoundError):
447 datastore.get(ref_comp)
448 with self.assertRaises(FileNotFoundError):
449 datastore.get(ref, storageClass="StructuredDataDictJson")
451 # Need a special method to generate stored dataset type.
452 def _stored_dataset_type(name: str, ref: DatasetRef = ref) -> DatasetType:
453 if name == ref.datasetType.name:
454 return ref.datasetType
455 raise ValueError(f"Unexpected dataset type name {ref.datasetType.name}")
457 datastore.set_retrieve_dataset_type_method(_stored_dataset_type)
459 # Storage class override with original dataset ref.
460 metrics_as_dict = datastore.get(ref, storageClass="StructuredDataDictJson")
461 self.assertIsInstance(metrics_as_dict, dict)
463 # get() should return a dict now.
464 metrics_as_dict = datastore.get(ref_comp)
465 self.assertIsInstance(metrics_as_dict, dict)
467 # exists() should work as well.
468 self.assertTrue(datastore.exists(ref_comp))
470 datastore.set_retrieve_dataset_type_method(None)
472 def testDisassembly(self) -> None:
473 """Test disassembly within datastore."""
474 metrics = makeExampleMetrics()
475 if self.isEphemeral:
476 # in-memory datastore does not disassemble
477 return
479 # Create multiple storage classes for testing different formulations
480 # of composites. One of these will not disassemble to provide
481 # a reference.
482 storageClasses = [
483 self.storageClassFactory.getStorageClass(sc)
484 for sc in (
485 "StructuredComposite",
486 "StructuredCompositeTestA",
487 "StructuredCompositeTestB",
488 "StructuredCompositeReadComp",
489 "StructuredData", # No disassembly
490 "StructuredCompositeReadCompNoDisassembly",
491 )
492 ]
494 # Create the test datastore
495 datastore = self.makeDatastore()
497 # Dummy dataId
498 dimensions = self.universe.conform(("visit", "physical_filter"))
499 dataId = {"instrument": "dummy", "visit": 428, "physical_filter": "R"}
501 for i, sc in enumerate(storageClasses):
502 with self.subTest(storageClass=sc.name):
503 # Create a different dataset type each time round
504 # so that a test failure in this subtest does not trigger
505 # a cascade of tests because of file clashes
506 ref = self.makeDatasetRef(f"metric_comp_{i}", dimensions, sc, dataId)
508 disassembled = sc.name not in {"StructuredData", "StructuredCompositeReadCompNoDisassembly"}
510 datastore.put(metrics, ref)
512 baseURI, compURIs = datastore.getURIs(ref)
513 if disassembled:
514 self.assertIsNone(baseURI)
515 self.assertEqual(set(compURIs), {"data", "output", "summary"})
516 else:
517 self.assertIsNotNone(baseURI)
518 self.assertEqual(compURIs, {})
520 metrics_get = datastore.get(ref)
521 self.assertEqual(metrics_get, metrics)
523 # Retrieve the composite with read parameter
524 stop = 4
525 metrics_get = datastore.get(ref, parameters={"slice": slice(stop)})
526 self.assertEqual(metrics_get.summary, metrics.summary)
527 self.assertEqual(metrics_get.output, metrics.output)
528 self.assertEqual(metrics_get.data, metrics.data[:stop])
530 # Retrieve a component
531 data = datastore.get(ref.makeComponentRef("data"))
532 self.assertEqual(data, metrics.data)
534 # On supported storage classes attempt to access a read
535 # only component
536 if "ReadComp" in sc.name:
537 cRef = ref.makeComponentRef("counter")
538 counter = datastore.get(cRef)
539 self.assertEqual(counter, len(metrics.data))
541 counter = datastore.get(cRef, parameters={"slice": slice(stop)})
542 self.assertEqual(counter, stop)
544 datastore.remove(ref)
546 def prepDeleteTest(self, n_refs: int = 1) -> tuple[Datastore, tuple[DatasetRef, ...]]:
547 metrics = makeExampleMetrics()
548 datastore = self.makeDatastore()
549 # Put
550 dimensions = self.universe.conform(("visit", "physical_filter"))
551 sc = self.storageClassFactory.getStorageClass("StructuredData")
552 refs = []
553 for i in range(n_refs):
554 dataId = {
555 "instrument": "dummy",
556 "visit": 638 + i,
557 "physical_filter": "U",
558 "band": "u",
559 "day_obs": 20250101,
560 }
561 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
562 datastore.put(metrics, ref)
564 # Does it exist?
565 self.assertTrue(datastore.exists(ref))
567 # Get
568 metricsOut = datastore.get(ref)
569 self.assertEqual(metrics, metricsOut)
570 refs.append(ref)
572 return datastore, *refs
574 def testRemove(self) -> None:
575 datastore, ref = self.prepDeleteTest()
577 # Remove
578 datastore.remove(ref)
580 # Does it exist?
581 self.assertFalse(datastore.exists(ref))
583 # Do we now get a predicted URI?
584 uri = datastore.getURI(ref, predict=True)
585 self.assertEqual(uri.fragment, "predicted")
587 # Get should now fail
588 with self.assertRaises(FileNotFoundError):
589 datastore.get(ref)
590 # Can only delete once
591 with self.assertRaises(FileNotFoundError):
592 datastore.remove(ref)
594 def testForget(self) -> None:
595 datastore, ref = self.prepDeleteTest()
597 # Remove
598 datastore.forget([ref])
600 # Does it exist (as far as we know)?
601 self.assertFalse(datastore.exists(ref))
603 # Do we now get a predicted URI?
604 uri = datastore.getURI(ref, predict=True)
605 self.assertEqual(uri.fragment, "predicted")
607 # Get should now fail
608 with self.assertRaises(FileNotFoundError):
609 datastore.get(ref)
611 # Forgetting again is a silent no-op
612 datastore.forget([ref])
614 # Predicted URI should still point to the file.
615 self.assertTrue(uri.exists())
617 def testTransfer(self) -> None:
618 metrics = makeExampleMetrics()
620 dimensions = self.universe.conform(("visit", "physical_filter"))
621 dataId = {
622 "instrument": "dummy",
623 "visit": 2048,
624 "physical_filter": "Uprime",
625 "band": "u",
626 "day_obs": 20250101,
627 }
629 sc = self.storageClassFactory.getStorageClass("StructuredData")
630 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
632 inputDatastore = self.makeDatastore("test_input_datastore")
633 outputDatastore = self.makeDatastore("test_output_datastore")
635 inputDatastore.put(metrics, ref)
636 outputDatastore.transfer(inputDatastore, ref)
638 metricsOut = outputDatastore.get(ref)
639 self.assertEqual(metrics, metricsOut)
641 def testBasicTransaction(self) -> None:
642 datastore = self.makeDatastore()
643 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
644 dimensions = self.universe.conform(("visit", "physical_filter"))
645 nDatasets = 6
646 dataIds = [
647 {"instrument": "dummy", "visit": i, "physical_filter": "V", "band": "v", "day_obs": 20250101}
648 for i in range(nDatasets)
649 ]
650 data = [
651 (
652 self.makeDatasetRef("metric", dimensions, storageClass, dataId),
653 makeExampleMetrics(),
654 )
655 for dataId in dataIds
656 ]
657 succeed = data[: nDatasets // 2]
658 fail = data[nDatasets // 2 :]
659 # All datasets added in this transaction should continue to exist
660 with datastore.transaction():
661 for ref, metrics in succeed:
662 datastore.put(metrics, ref)
663 # Whereas datasets added in this transaction should not
664 with self.assertRaises(TransactionTestError):
665 with datastore.transaction():
666 for ref, metrics in fail:
667 datastore.put(metrics, ref)
668 raise TransactionTestError("This should propagate out of the context manager")
669 # Check for datasets that should exist
670 for ref, metrics in succeed:
671 # Does it exist?
672 self.assertTrue(datastore.exists(ref))
673 # Get
674 metricsOut = datastore.get(ref, parameters=None)
675 self.assertEqual(metrics, metricsOut)
676 # URI
677 uri = datastore.getURI(ref)
678 self.assertEqual(uri.scheme, self.uriScheme)
679 # Check for datasets that should not exist
680 for ref, _ in fail:
681 # These should raise
682 with self.assertRaises(FileNotFoundError):
683 # non-existing file
684 datastore.get(ref)
685 with self.assertRaises(FileNotFoundError):
686 datastore.getURI(ref)
688 def testNestedTransaction(self) -> None:
689 datastore = self.makeDatastore()
690 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
691 dimensions = self.universe.conform(("visit", "physical_filter"))
692 metrics = makeExampleMetrics()
694 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
695 refBefore = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
696 datastore.put(metrics, refBefore)
697 with self.assertRaises(TransactionTestError):
698 with datastore.transaction():
699 dataId = {
700 "instrument": "dummy",
701 "visit": 1,
702 "physical_filter": "V",
703 "band": "v",
704 "day_obs": 20250101,
705 }
706 refOuter = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
707 datastore.put(metrics, refOuter)
708 with datastore.transaction():
709 dataId = {
710 "instrument": "dummy",
711 "visit": 2,
712 "physical_filter": "V",
713 "band": "v",
714 "day_obs": 20250101,
715 }
716 refInner = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
717 datastore.put(metrics, refInner)
718 # All datasets should exist
719 for ref in (refBefore, refOuter, refInner):
720 metricsOut = datastore.get(ref, parameters=None)
721 self.assertEqual(metrics, metricsOut)
722 raise TransactionTestError("This should roll back the transaction")
723 # Dataset(s) inserted before the transaction should still exist
724 metricsOut = datastore.get(refBefore, parameters=None)
725 self.assertEqual(metrics, metricsOut)
726 # But all datasets inserted during the (rolled back) transaction
727 # should be gone
728 with self.assertRaises(FileNotFoundError):
729 datastore.get(refOuter)
730 with self.assertRaises(FileNotFoundError):
731 datastore.get(refInner)
733 def _prepareIngestTest(self) -> tuple[MetricsExample, DatasetRef]:
734 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
735 dimensions = self.universe.conform(("visit", "physical_filter"))
736 metrics = makeExampleMetrics()
737 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
738 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
739 return metrics, ref
741 def runIngestTest(self, func: Callable[[MetricsExample, str, DatasetRef], None]) -> None:
742 metrics, ref = self._prepareIngestTest()
743 # The file will be deleted after the test.
744 # For symlink tests this leads to a situation where the datastore
745 # points to a file that does not exist. This will make os.path.exist
746 # return False but then the new symlink will fail with
747 # FileExistsError later in the code so the test still passes.
748 with _temp_yaml_file(metrics._asdict()) as path:
749 func(metrics, path, ref)
751 def testIngestNoTransfer(self) -> None:
752 """Test ingesting existing files with no transfer."""
753 for mode in (None, "auto"):
754 # Some datastores have auto but can't do in place transfer
755 if mode == "auto" and "auto" in self.ingestTransferModes and not self.canIngestNoTransferAuto:
756 continue
758 with self.subTest(mode=mode):
759 datastore = self.makeDatastore()
761 def succeed(
762 obj: MetricsExample,
763 path: str,
764 ref: DatasetRef,
765 mode: str | None = mode,
766 datastore: Datastore = datastore,
767 ) -> None:
768 """Ingest a file already in the datastore root."""
769 # first move it into the root, and adjust the path
770 # accordingly.
771 # In the case of a ChainedDatastore, we have multiple
772 # roots, all of which will accept the file, so we
773 # have to copy it into all the roots.
774 relative_path = None
775 for root in datastore.roots.values():
776 if root is not None:
777 copied_path = shutil.copy(path, root.ospath)
778 relative_path = os.path.relpath(copied_path, start=root.ospath)
779 assert relative_path is not None, (
780 "Running a FileDatastore test on a Datastore instance without any roots"
781 )
782 datastore.ingest(FileDataset(path=relative_path, refs=ref), transfer=mode)
783 self.assertEqual(obj, datastore.get(ref))
785 def failInputDoesNotExist(
786 obj: MetricsExample,
787 path: str,
788 ref: DatasetRef,
789 mode: str | None = mode,
790 datastore: Datastore = datastore,
791 ) -> None:
792 """Can't ingest files if we're given a bad path."""
793 with self.assertRaises(FileNotFoundError):
794 datastore.ingest(
795 FileDataset(path="this-file-does-not-exist.yaml", refs=ref), transfer=mode
796 )
797 self.assertFalse(datastore.exists(ref))
799 def failOutsideRoot(
800 obj: MetricsExample,
801 path: str,
802 ref: DatasetRef,
803 mode: str | None = mode,
804 datastore: Datastore = datastore,
805 ) -> None:
806 """Can't ingest files outside of datastore root unless
807 auto.
808 """
809 if mode == "auto":
810 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode)
811 self.assertTrue(datastore.exists(ref))
812 else:
813 with self.assertRaises(RuntimeError):
814 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode)
815 self.assertFalse(datastore.exists(ref))
817 def failNotImplemented(
818 obj: MetricsExample,
819 path: str,
820 ref: DatasetRef,
821 mode: str | None = mode,
822 datastore: Datastore = datastore,
823 ) -> None:
824 with self.assertRaises(NotImplementedError):
825 datastore.ingest(FileDataset(path=path, refs=ref), transfer=mode)
827 if mode in self.ingestTransferModes:
828 self.runIngestTest(failOutsideRoot)
829 self.runIngestTest(failInputDoesNotExist)
830 self.runIngestTest(succeed)
831 else:
832 self.runIngestTest(failNotImplemented)
834 def testIngestTransfer(self) -> None:
835 """Test ingesting existing files after transferring them."""
836 for mode in ("copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto"):
837 with self.subTest(mode=mode):
838 datastore = self.makeDatastore(mode)
840 def succeed(
841 obj: MetricsExample,
842 path: str,
843 ref: DatasetRef,
844 mode: str | None = mode,
845 datastore: Datastore = datastore,
846 ) -> None:
847 """Ingest a file by transferring it to the template
848 location.
849 """
850 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode)
851 self.assertEqual(obj, datastore.get(ref))
852 file_exists = os.path.exists(path)
853 if mode == "move":
854 self.assertFalse(file_exists)
855 else:
856 self.assertTrue(file_exists)
858 def failInputDoesNotExist(
859 obj: MetricsExample,
860 path: str,
861 ref: DatasetRef,
862 mode: str | None = mode,
863 datastore: Datastore = datastore,
864 ) -> None:
865 """Can't ingest files if we're given a bad path."""
866 with self.assertRaises(FileNotFoundError):
867 # Ensure the file does not look like it is in
868 # datastore for auto mode
869 datastore.ingest(
870 FileDataset(path="../this-file-does-not-exist.yaml", refs=ref), transfer=mode
871 )
872 self.assertFalse(datastore.exists(ref), f"Checking not in datastore using mode {mode}")
874 def failNotImplemented(
875 obj: MetricsExample,
876 path: str,
877 ref: DatasetRef,
878 mode: str | None = mode,
879 datastore: Datastore = datastore,
880 ) -> None:
881 with self.assertRaises(NotImplementedError):
882 datastore.ingest(FileDataset(path=os.path.abspath(path), refs=ref), transfer=mode)
884 if mode in self.ingestTransferModes:
885 self.runIngestTest(failInputDoesNotExist)
886 self.runIngestTest(succeed)
887 else:
888 self.runIngestTest(failNotImplemented)
890 def testIngestSymlinkOfSymlink(self) -> None:
891 """Special test for symlink to a symlink ingest"""
892 metrics, ref = self._prepareIngestTest()
893 # The aim of this test is to create a dataset on disk, then
894 # create a symlink to it and finally ingest the symlink such that
895 # the symlink in the datastore points to the original dataset.
896 for mode in ("symlink", "relsymlink"):
897 if mode not in self.ingestTransferModes:
898 continue
900 print(f"Trying mode {mode}")
901 with _temp_yaml_file(metrics._asdict()) as realpath:
902 with tempfile.TemporaryDirectory() as tmpdir:
903 sympath = os.path.join(tmpdir, "symlink.yaml")
904 os.symlink(os.path.realpath(realpath), sympath)
906 datastore = self.makeDatastore()
907 datastore.ingest(FileDataset(path=os.path.abspath(sympath), refs=ref), transfer=mode)
909 uri = datastore.getURI(ref)
910 self.assertTrue(uri.isLocal, f"Check {uri.scheme}")
911 self.assertTrue(os.path.islink(uri.ospath), f"Check {uri} is a symlink")
913 linkTarget = os.readlink(uri.ospath)
914 if mode == "relsymlink":
915 self.assertFalse(os.path.isabs(linkTarget))
916 else:
917 self.assertTrue(os.path.samefile(linkTarget, realpath))
919 # Check that we can get the dataset back regardless of mode
920 metric2 = datastore.get(ref)
921 self.assertEqual(metric2, metrics)
923 # Cleanup the file for next time round loop
924 # since it will get the same file name in store
925 datastore.remove(ref)
927 def _populate_export_datastore(self, name: str) -> tuple[Datastore, list[DatasetRef]]:
928 datastore = self.makeDatastore(name)
930 # For now only the FileDatastore can be used for this test.
931 # ChainedDatastore that only includes InMemoryDatastores have to be
932 # skipped as well.
933 for name in datastore.names:
934 if not name.startswith("InMemoryDatastore"):
935 break
936 else:
937 raise unittest.SkipTest("in-memory datastore does not support record export/import")
939 metrics = makeExampleMetrics()
940 dimensions = self.universe.conform(("visit", "physical_filter"))
941 sc = self.storageClassFactory.getStorageClass("StructuredData")
943 refs = []
944 for visit in (2048, 2049, 2050):
945 dataId = {
946 "instrument": "dummy",
947 "visit": visit,
948 "physical_filter": "Uprime",
949 "band": "u",
950 "day_obs": 20250101,
951 }
952 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
953 datastore.put(metrics, ref)
954 refs.append(ref)
955 return datastore, refs
957 def testExportImportRecords(self) -> None:
958 """Test for export_records and import_records methods."""
959 datastore, refs = self._populate_export_datastore("test_datastore")
960 for exported_refs in (refs, refs[1:]):
961 n_refs = len(exported_refs)
962 records = datastore.export_records(exported_refs)
963 self.assertGreater(len(records), 0)
964 self.assertTrue(set(records.keys()) <= set(datastore.names))
965 # In a ChainedDatastore each FileDatastore will have a complete set
966 for datastore_name in records:
967 record_data = records[datastore_name]
968 self.assertEqual(len(record_data.records), n_refs)
970 # Check that subsetting works, include non-existing dataset ID.
971 dataset_ids = {exported_refs[0].id, uuid.uuid4()}
972 subset = record_data.subset(dataset_ids)
973 assert subset is not None
974 self.assertEqual(len(subset.records), 1)
975 subset = record_data.subset({uuid.uuid4()})
976 self.assertIsNone(subset)
978 # Use the same datastore name to import relative path.
979 datastore2 = self.makeDatastore("test_datastore")
981 records = datastore.export_records(refs[1:])
982 datastore2.import_records(records)
984 with self.assertRaises(FileNotFoundError):
985 data = datastore2.get(refs[0])
986 data = datastore2.get(refs[1])
987 self.assertIsNotNone(data)
988 data = datastore2.get(refs[2])
989 self.assertIsNotNone(data)
991 def testExportPredictedRecords(self):
992 if self.isEphemeral:
993 raise unittest.SkipTest("in-memory datastore does not support record export/import")
994 sc = self.storageClassFactory.getStorageClass("ThingOne")
995 dimensions = self.universe.conform(("visit", "physical_filter"))
996 dataId = {
997 "instrument": "dummy",
998 "visit": 52,
999 "physical_filter": "V",
1000 "band": "v",
1001 "day_obs": 20250101,
1002 }
1003 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
1005 datastore = self.makeDatastore("test_datastore")
1006 names = {n for n in datastore.names if not n.startswith("InMemory")}
1007 records = datastore.export_predicted_records([ref])
1009 # Expect predicted records from all datastores.
1010 self.assertEqual(set(records.keys()), names)
1012 for record_data in records.values():
1013 self.assertEqual(len(record_data.records), 1)
1015 def testExport(self) -> None:
1016 datastore, refs = self._populate_export_datastore("test_datastore")
1018 datasets = list(datastore.export(refs))
1019 self.assertEqual(len(datasets), 3)
1021 for transfer in (None, "auto"):
1022 # Both will default to None
1023 datasets = list(datastore.export(refs, transfer=transfer))
1024 self.assertEqual(len(datasets), 3)
1026 with self.assertRaises(TypeError):
1027 list(datastore.export(refs, transfer="copy"))
1029 with self.assertRaises(TypeError):
1030 list(datastore.export(refs, directory="exportDir", transfer="move"))
1032 # Create a new ref that is not known to the datastore and try to
1033 # export it.
1034 sc = self.storageClassFactory.getStorageClass("ThingOne")
1035 dimensions = self.universe.conform(("visit", "physical_filter"))
1036 dataId = {
1037 "instrument": "dummy",
1038 "visit": 52,
1039 "physical_filter": "V",
1040 "band": "v",
1041 "day_obs": 20250101,
1042 }
1043 ref = self.makeDatasetRef("metric", dimensions, sc, dataId)
1044 with self.assertRaises(FileNotFoundError):
1045 list(datastore.export(refs + [ref], transfer=None))
1047 def test_pydantic_dict_storage_class_conversions(self) -> None:
1048 """Test converting a dataset stored as a pydantic model into a dict on
1049 read.
1050 """
1051 datastore = self.makeDatastore()
1052 store_as_model = self.makeDatasetRef(
1053 "store_as_model",
1054 dimensions=self.universe.empty,
1055 storageClass="DictConvertibleModel",
1056 dataId=DataCoordinate.make_empty(self.universe),
1057 )
1058 content = {"a": "one", "b": "two"}
1059 model = DictConvertibleModel.from_dict(content, extra="original content")
1060 datastore.put(model, store_as_model)
1061 retrieved_model = datastore.get(store_as_model)
1062 self.assertEqual(retrieved_model, model)
1063 loaded = datastore.get(store_as_model.overrideStorageClass("NativeDictForConvertibleModel"))
1064 self.assertEqual(type(loaded), dict)
1065 self.assertEqual(loaded, content)
1067 def test_simple_class_put_get(self) -> None:
1068 """Test that we can put and get a simple class with dict()
1069 constructor.
1070 """
1071 datastore = self.makeDatastore()
1072 data = MetricsExample(summary={"a": 1}, data=[1, 2, 3], output={"b": 2})
1073 self._assert_different_puts(datastore, "MetricsExample", data)
1075 def test_dataclass_put_get(self) -> None:
1076 """Test that we can put and get a simple dataclass."""
1077 datastore = self.makeDatastore()
1078 data = MetricsExampleDataclass(summary={"a": 1}, data=[1, 2, 3], output={"b": 2})
1079 self._assert_different_puts(datastore, "MetricsExampleDataclass", data)
1081 def test_pydantic_put_get(self) -> None:
1082 """Test that we can put and get a simple Pydantic model."""
1083 datastore = self.makeDatastore()
1084 data = MetricsExampleModel(summary={"a": 1}, data=[1, 2, 3], output={"b": 2})
1085 self._assert_different_puts(datastore, "MetricsExampleModel", data)
1087 def test_tuple_put_get(self) -> None:
1088 """Test that we can put and get a tuple."""
1089 datastore = self.makeDatastore()
1090 data = ("a", "b", 1)
1091 self._assert_different_puts(datastore, "TupleExample", data)
1093 def _assert_different_puts(self, datastore: Datastore, storageClass_root: str, data: Any) -> None:
1094 refs = {
1095 x: self.makeDatasetRef(
1096 f"stora_as_{x}",
1097 dimensions=self.universe.empty,
1098 storageClass=f"{storageClass_root}{x}",
1099 dataId=DataCoordinate.make_empty(self.universe),
1100 )
1101 for x in ["A", "B"]
1102 }
1104 for ref in refs.values():
1105 datastore.put(data, ref)
1107 self.assertEqual(datastore.get(refs["A"]), datastore.get(refs["B"]))
1110class PosixDatastoreTestCase(DatastoreTests, unittest.TestCase):
1111 """PosixDatastore specialization"""
1113 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
1114 uriScheme = "file"
1115 canIngestNoTransferAuto = True
1116 ingestTransferModes = (None, "copy", "move", "link", "hardlink", "symlink", "relsymlink", "auto")
1117 isEphemeral = False
1118 rootKeys = ("root",)
1119 validationCanFail = True
1121 def setUp(self) -> None:
1122 # The call to os.path.realpath is necessary because Mac temporary files
1123 # can end up in either /private/var/folders or /var/folders, which
1124 # refer to the same location but don't appear to.
1125 # This matters for "relsymlink" transfer mode, because it needs to be
1126 # able to read the file through a relative symlink, but some of the
1127 # intermediate directories are not traversable if you try to get from a
1128 # tempfile in /var/folders to one in /private/var/folders via a
1129 # relative path.
1130 self.root = os.path.realpath(self.enterContext(tempfile.TemporaryDirectory()))
1131 super().setUp()
1133 def testAtomicWrite(self) -> None:
1134 """Test that we write to a temporary and then rename"""
1135 datastore = self.makeDatastore()
1136 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1137 dimensions = self.universe.conform(("visit", "physical_filter"))
1138 metrics = makeExampleMetrics()
1140 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
1141 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1143 with self.assertLogs("lsst.resources", "DEBUG") as cm:
1144 datastore.put(metrics, ref)
1145 move_logs = [ll for ll in cm.output if "transfer=" in ll]
1146 self.assertIn("transfer=move", move_logs[0])
1148 # And the transfer should be file to file.
1149 self.assertEqual(move_logs[0].count("file://"), 2)
1151 def testCanNotDeterminePutFormatterLocation(self) -> None:
1152 """Verify that the expected exception is raised if the FileDatastore
1153 can not determine the put formatter location.
1154 """
1155 _ = makeExampleMetrics()
1156 datastore = self.makeDatastore()
1158 # Create multiple storage classes for testing different formulations
1159 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1161 sccomp = StorageClass("Dummy")
1162 compositeStorageClass = StorageClass(
1163 "StructuredComposite", components={"dummy": sccomp, "dummy2": sccomp}
1164 )
1166 dimensions = self.universe.conform(("visit", "physical_filter"))
1167 dataId = {
1168 "instrument": "dummy",
1169 "visit": 52,
1170 "physical_filter": "V",
1171 "band": "v",
1172 "day_obs": 20250101,
1173 }
1175 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1176 compRef = self.makeDatasetRef("metric", dimensions, compositeStorageClass, dataId)
1178 def raiser(ref: DatasetRef) -> None:
1179 raise DatasetTypeNotSupportedError()
1181 with unittest.mock.patch.object(
1182 lsst.daf.butler.datastores.fileDatastore.FileDatastore,
1183 "_determine_put_formatter_location",
1184 side_effect=raiser,
1185 ):
1186 # verify the non-composite ref execution path:
1187 with self.assertRaises(DatasetTypeNotSupportedError):
1188 datastore.getURIs(ref, predict=True)
1190 # verify the composite-ref execution path:
1191 with self.assertRaises(DatasetTypeNotSupportedError):
1192 datastore.getURIs(compRef, predict=True)
1194 def test_roots(self):
1195 datastore = self.makeDatastore()
1197 self.assertEqual(set(datastore.names), set(datastore.roots.keys()))
1198 for root in datastore.roots.values():
1199 if root is not None:
1200 self.assertTrue(root.exists())
1202 def test_prepare_get_for_external_client(self):
1203 datastore = self.makeDatastore()
1204 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1205 dimensions = self.universe.conform(("visit", "physical_filter"))
1206 dataId = {"instrument": "dummy", "visit": 52, "physical_filter": "V", "band": "v"}
1207 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1208 # Most of the coverage for this function is in test_server.py,
1209 # because it requires a file backend that supports URL signing.
1210 self.assertIsNone(datastore.prepare_get_for_external_client(ref))
1213class PosixDatastoreNoChecksumsTestCase(PosixDatastoreTestCase):
1214 """Posix datastore tests but with checksums disabled."""
1216 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreNoChecksums.yaml")
1218 def testChecksum(self) -> None:
1219 """Ensure that checksums have not been calculated."""
1220 datastore = self.makeDatastore()
1221 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1222 dimensions = self.universe.conform(("visit", "physical_filter"))
1223 metrics = makeExampleMetrics()
1225 dataId = {"instrument": "dummy", "visit": 0, "physical_filter": "V", "band": "v", "day_obs": 20250101}
1226 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1228 # Configuration should have disabled checksum calculation
1229 datastore.put(metrics, ref)
1230 infos = datastore.getStoredItemsInfo(ref)
1231 self.assertIsNone(infos[0].checksum)
1233 # Remove put back but with checksums enabled explicitly
1234 datastore.remove(ref)
1235 datastore.useChecksum = True
1236 datastore.put(metrics, ref)
1238 infos = datastore.getStoredItemsInfo(ref)
1239 self.assertIsNotNone(infos[0].checksum)
1241 def test_repeat_ingest(self):
1242 """Test that repeatedly ingesting the same file in direct mode
1243 is allowed.
1245 Test can only run with FileDatastore since that is the only one
1246 supporting "direct" ingest.
1247 """
1248 metrics, v4ref = self._prepareIngestTest()
1249 datastore = self.makeDatastore()
1250 v5ref = DatasetRef(
1251 v4ref.datasetType, v4ref.dataId, v4ref.run, id_generation_mode=DatasetIdGenEnum.DATAID_TYPE_RUN
1252 )
1254 with _temp_yaml_file(metrics._asdict()) as path:
1255 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct")
1257 # This will fail because the ref is using UUIDv4.
1258 with self.assertRaises(RuntimeError):
1259 datastore.ingest(FileDataset(path=path, refs=v4ref), transfer="direct")
1261 # UUIDv5 can be repeatedly ingested in direct mode.
1262 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct")
1263 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="direct")
1265 with self.assertRaises(RuntimeError):
1266 datastore.ingest(FileDataset(path=path, refs=v5ref), transfer="copy")
1269class TrashDatastoreTestCase(PosixDatastoreTestCase):
1270 """Restrict trash test to FileDatastore."""
1272 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
1274 def testTrash(self) -> None:
1275 datastore, *refs = self.prepDeleteTest(n_refs=10)
1277 # Trash one of them.
1278 ref = refs.pop()
1279 uri = datastore.getURI(ref)
1280 datastore.trash(ref)
1281 self.assertTrue(uri.exists(), uri) # Not deleted yet
1282 datastore.emptyTrash()
1283 self.assertFalse(uri.exists(), uri)
1285 # Trash it again should be fine.
1286 datastore.trash(ref)
1288 # Trash multiple items at once.
1289 subset = [refs.pop(), refs.pop()]
1290 datastore.trash(subset)
1291 datastore.emptyTrash()
1293 # Remove a record and trash should do nothing.
1294 # This is execution butler scenario.
1295 ref = refs.pop()
1296 uri = datastore.getURI(ref)
1297 datastore._table.delete(["dataset_id"], {"dataset_id": ref.id})
1298 self.assertTrue(uri.exists())
1299 datastore.trash(ref)
1300 datastore.emptyTrash()
1301 self.assertTrue(uri.exists())
1303 # Switch on trust and it should delete the file.
1304 datastore.trustGetRequest = True
1305 datastore.trash([ref])
1306 self.assertFalse(uri.exists())
1308 # Remove multiples at once in trust mode.
1309 subset = [refs.pop() for i in range(3)]
1310 datastore.trash(subset)
1311 datastore.trash(refs.pop()) # Check that a single ref can trash
1313 def test_empty_trash(self) -> None:
1314 """Test parameters and return value for empty trash."""
1315 datastore, *refs = self.prepDeleteTest(n_refs=10)
1317 # Trash one of them.
1318 ref = refs.pop()
1319 uri = datastore.getURI(ref)
1320 datastore.trash(ref)
1321 self.assertTrue(uri.exists(), uri) # Not deleted yet
1323 # Empty trash but with a list of refs that does not include the
1324 # one in the trash table.
1325 removed = datastore.emptyTrash(refs=refs)
1326 self.assertEqual(len(removed), 0)
1327 self.assertTrue(uri.exists(), uri)
1329 # Empty the entire trash but in dry_run mode.
1330 removed = datastore.emptyTrash(dry_run=True)
1331 self.assertEqual(len(removed), 1)
1332 self.assertEqual(removed.pop(), uri)
1333 self.assertTrue(uri.exists(), uri)
1335 # Empty the trash specifying the actual ref.
1336 removed = datastore.emptyTrash(refs=[ref])
1337 self.assertEqual(len(removed), 1)
1338 self.assertEqual(removed.pop(), uri)
1339 self.assertFalse(uri.exists(), uri)
1341 # Trash everything and empty.
1342 datastore.trash(refs)
1343 removed = datastore.emptyTrash(dry_run=True)
1344 for u in removed:
1345 self.assertTrue(u.exists())
1346 removed = datastore.emptyTrash()
1347 for u in removed:
1348 self.assertFalse(u.exists())
1351class CleanupPosixDatastoreTestCase(DatastoreTestsBase, unittest.TestCase):
1352 """Test datastore cleans up on failure."""
1354 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
1356 def setUp(self) -> None:
1357 # Override the working directory before calling the base class
1358 self.root = tempfile.mkdtemp()
1359 super().setUp()
1361 def testCleanup(self) -> None:
1362 """Test that a failed formatter write does cleanup a partial file."""
1363 metrics = makeExampleMetrics()
1364 datastore = self.makeDatastore()
1366 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1368 dimensions = self.universe.conform(("visit", "physical_filter"))
1369 dataId = {
1370 "instrument": "dummy",
1371 "visit": 52,
1372 "physical_filter": "V",
1373 "band": "v",
1374 "day_obs": 20250101,
1375 }
1377 ref = self.makeDatasetRef("metric", dimensions, storageClass, dataId)
1379 # Determine where the file will end up (we assume Formatters use
1380 # the same file extension)
1381 expectedUri = datastore.getURI(ref, predict=True)
1382 self.assertEqual(expectedUri.fragment, "predicted")
1384 self.assertEqual(expectedUri.getExtension(), ".yaml", f"Is there a file extension in {expectedUri}")
1386 # Try formatter that fails and formatter that fails and leaves
1387 # a file behind
1388 for formatter in (BadWriteFormatter, BadNoWriteFormatter):
1389 with self.subTest(formatter=str(formatter)):
1390 # Monkey patch the formatter
1391 datastore.formatterFactory.registerFormatter(ref.datasetType, formatter, overwrite=True)
1393 # Try to put the dataset, it should fail
1394 with self.assertRaises(RuntimeError):
1395 datastore.put(metrics, ref)
1397 # Check that there is no file on disk
1398 self.assertFalse(expectedUri.exists(), f"Check for existence of {expectedUri}")
1400 # Check that there is a directory
1401 dir = expectedUri.dirname()
1402 self.assertTrue(dir.exists(), f"Check for existence of directory {dir}")
1404 # Force YamlFormatter and check that this time a file is written
1405 datastore.formatterFactory.registerFormatter(ref.datasetType, YamlFormatter, overwrite=True)
1406 datastore.put(metrics, ref)
1407 self.assertTrue(expectedUri.exists(), f"Check for existence of {expectedUri}")
1408 datastore.remove(ref)
1409 self.assertFalse(expectedUri.exists(), f"Check for existence of now removed {expectedUri}")
1412class InMemoryDatastoreTestCase(DatastoreTests, unittest.TestCase):
1413 """PosixDatastore specialization"""
1415 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastore.yaml")
1416 uriScheme = "mem"
1417 hasUnsupportedPut = False
1418 ingestTransferModes = ()
1419 isEphemeral = True
1420 rootKeys = None
1421 validationCanFail = False
1424class ChainedDatastoreTestCase(PosixDatastoreTestCase):
1425 """ChainedDatastore specialization using a POSIXDatastore"""
1427 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore.yaml")
1428 hasUnsupportedPut = False
1429 canIngestNoTransferAuto = False
1430 ingestTransferModes = (None, "copy", "move", "hardlink", "symlink", "relsymlink", "link", "auto")
1431 isEphemeral = False
1432 rootKeys = (".datastores.1.root", ".datastores.2.root")
1433 validationCanFail = True
1436class ChainedDatastoreMemoryTestCase(InMemoryDatastoreTestCase):
1437 """ChainedDatastore specialization using all InMemoryDatastore"""
1439 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2.yaml")
1440 validationCanFail = False
1441 isEphemeral = True
1444class DatastoreConstraintsTests(DatastoreTestsBase):
1445 """Basic tests of constraints model of Datastores."""
1447 def testConstraints(self) -> None:
1448 """Test constraints model. Assumes that each test class has the
1449 same constraints.
1450 """
1451 metrics = makeExampleMetrics()
1452 datastore = self.makeDatastore()
1454 sc1 = self.storageClassFactory.getStorageClass("StructuredData")
1455 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson")
1456 dimensions = self.universe.conform(("visit", "physical_filter", "instrument"))
1457 dataId = {
1458 "visit": 52,
1459 "physical_filter": "V",
1460 "band": "v",
1461 "instrument": "DummyCamComp",
1462 "day_obs": 20250101,
1463 }
1465 # Write empty file suitable for ingest check (JSON and YAML variants)
1466 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml")
1467 testfile_j = tempfile.NamedTemporaryFile(suffix=".json")
1468 for datasetTypeName, sc, accepted in (
1469 ("metric", sc1, True),
1470 ("metric5", sc1, False),
1471 ("metric33", sc1, True),
1472 ("metric5", sc2, True),
1473 ):
1474 # Choose different temp file depending on StorageClass
1475 testfile = testfile_j if sc.name.endswith("Json") else testfile_y
1477 with self.subTest(datasetTypeName=datasetTypeName, storageClass=sc.name, file=testfile.name):
1478 ref = self.makeDatasetRef(datasetTypeName, dimensions, sc, dataId)
1479 if accepted:
1480 datastore.put(metrics, ref)
1481 self.assertTrue(datastore.exists(ref))
1482 datastore.remove(ref)
1484 # Try ingest
1485 if self.canIngest:
1486 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1487 self.assertTrue(datastore.exists(ref))
1488 datastore.remove(ref)
1489 else:
1490 with self.assertRaises(DatasetTypeNotSupportedError):
1491 datastore.put(metrics, ref)
1492 self.assertFalse(datastore.exists(ref))
1494 # Again with ingest
1495 if self.canIngest:
1496 with self.assertRaises(DatasetTypeNotSupportedError):
1497 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1498 self.assertFalse(datastore.exists(ref))
1501class PosixDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):
1502 """PosixDatastore specialization"""
1504 configFile = os.path.join(TESTDIR, "config/basic/posixDatastoreP.yaml")
1505 canIngest = True
1507 def setUp(self) -> None:
1508 # Override the working directory before calling the base class
1509 self.root = tempfile.mkdtemp()
1510 super().setUp()
1513class InMemoryDatastoreConstraintsTestCase(DatastoreConstraintsTests, unittest.TestCase):
1514 """InMemoryDatastore specialization."""
1516 configFile = os.path.join(TESTDIR, "config/basic/inMemoryDatastoreP.yaml")
1517 canIngest = False
1520class ChainedDatastoreConstraintsNativeTestCase(PosixDatastoreConstraintsTestCase):
1521 """ChainedDatastore specialization using a POSIXDatastore and constraints
1522 at the ChainedDatstore.
1523 """
1525 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePa.yaml")
1528class ChainedDatastoreConstraintsTestCase(PosixDatastoreConstraintsTestCase):
1529 """ChainedDatastore specialization using a POSIXDatastore."""
1531 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastoreP.yaml")
1534class ChainedDatastoreMemoryConstraintsTestCase(InMemoryDatastoreConstraintsTestCase):
1535 """ChainedDatastore specialization using all InMemoryDatastore."""
1537 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastore2P.yaml")
1538 canIngest = False
1541class ChainedDatastorePerStoreConstraintsTests(DatastoreTestsBase, unittest.TestCase):
1542 """Test that a chained datastore can control constraints per-datastore
1543 even if child datastore would accept.
1544 """
1546 configFile = os.path.join(TESTDIR, "config/basic/chainedDatastorePb.yaml")
1548 def setUp(self) -> None:
1549 # Override the working directory before calling the base class
1550 self.root = tempfile.mkdtemp()
1551 super().setUp()
1553 def testConstraints(self) -> None:
1554 """Test chained datastore constraints model."""
1555 metrics = makeExampleMetrics()
1556 datastore = self.makeDatastore()
1558 sc1 = self.storageClassFactory.getStorageClass("StructuredData")
1559 sc2 = self.storageClassFactory.getStorageClass("StructuredDataJson")
1560 dimensions = self.universe.conform(("visit", "physical_filter", "instrument"))
1561 dataId1 = {
1562 "visit": 52,
1563 "physical_filter": "V",
1564 "band": "v",
1565 "instrument": "DummyCamComp",
1566 "day_obs": 20250101,
1567 }
1568 dataId2 = {"visit": 52, "physical_filter": "V", "band": "v", "instrument": "HSC", "day_obs": 20250101}
1570 # Write empty file suitable for ingest check (JSON and YAML variants)
1571 testfile_y = tempfile.NamedTemporaryFile(suffix=".yaml")
1572 testfile_j = tempfile.NamedTemporaryFile(suffix=".json")
1574 for typeName, dataId, sc, accept, ingest in (
1575 ("metric", dataId1, sc1, (False, True, False), True),
1576 ("metric5", dataId1, sc1, (False, False, False), False),
1577 ("metric5", dataId2, sc1, (True, False, False), False),
1578 ("metric33", dataId2, sc2, (True, True, False), True),
1579 ("metric5", dataId1, sc2, (False, True, False), True),
1580 ):
1581 # Choose different temp file depending on StorageClass
1582 testfile = testfile_j if sc.name.endswith("Json") else testfile_y
1584 with self.subTest(datasetTypeName=typeName, dataId=dataId, sc=sc.name):
1585 ref = self.makeDatasetRef(typeName, dimensions, sc, dataId)
1586 if any(accept):
1587 datastore.put(metrics, ref)
1588 self.assertTrue(datastore.exists(ref))
1590 # Check each datastore inside the chained datastore
1591 for childDatastore, expected in zip(datastore.datastores, accept, strict=True):
1592 self.assertEqual(
1593 childDatastore.exists(ref),
1594 expected,
1595 f"Testing presence of {ref} in datastore {childDatastore.name}",
1596 )
1598 datastore.remove(ref)
1600 # Check that ingest works
1601 if ingest:
1602 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1603 self.assertTrue(datastore.exists(ref))
1605 # Check each datastore inside the chained datastore
1606 for childDatastore, expected in zip(datastore.datastores, accept, strict=True):
1607 # Ephemeral datastores means InMemory at the moment
1608 # and that does not accept ingest of files.
1609 if childDatastore.isEphemeral:
1610 expected = False
1611 self.assertEqual(
1612 childDatastore.exists(ref),
1613 expected,
1614 f"Testing presence of ingested {ref} in datastore {childDatastore.name}",
1615 )
1617 datastore.remove(ref)
1618 else:
1619 with self.assertRaises(DatasetTypeNotSupportedError):
1620 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1622 else:
1623 with self.assertRaises(DatasetTypeNotSupportedError):
1624 datastore.put(metrics, ref)
1625 self.assertFalse(datastore.exists(ref))
1627 # Again with ingest
1628 with self.assertRaises(DatasetTypeNotSupportedError):
1629 datastore.ingest(FileDataset(testfile.name, [ref]), transfer="link")
1630 self.assertFalse(datastore.exists(ref))
1633@unittest.mock.patch.dict(os.environ, {}, clear=True)
1634class DatastoreCacheTestCase(DatasetTestHelper, unittest.TestCase):
1635 """Tests for datastore caching infrastructure."""
1637 @classmethod
1638 def setUpClass(cls) -> None:
1639 cls.storageClassFactory = StorageClassFactory()
1640 cls.universe = DimensionUniverse()
1642 # Ensure that we load the test storage class definitions.
1643 scConfigFile = os.path.join(TESTDIR, "config/basic/storageClasses.yaml")
1644 cls.storageClassFactory.addFromConfig(scConfigFile)
1646 def setUp(self) -> None:
1647 self.id = 0
1649 # Create a root that we can use for caching tests.
1650 self.root = tempfile.mkdtemp()
1652 # Create some test dataset refs and associated test files
1653 sc = self.storageClassFactory.getStorageClass("StructuredDataDict")
1654 dimensions = self.universe.conform(("visit", "physical_filter"))
1655 dataId = {
1656 "instrument": "dummy",
1657 "visit": 52,
1658 "physical_filter": "V",
1659 "band": "v",
1660 "day_obs": 20250101,
1661 }
1663 # Create list of refs and list of temporary files
1664 n_datasets = 10
1665 self.refs = [self.makeDatasetRef(f"metric{n}", dimensions, sc, dataId) for n in range(n_datasets)]
1667 root_uri = ResourcePath(self.root, forceDirectory=True)
1668 self.files = [root_uri.join(f"file{n}.txt") for n in range(n_datasets)]
1670 # Create test files.
1671 for uri in self.files:
1672 uri.write(b"0123456789")
1674 # Create some composite refs with component files.
1675 sc = self.storageClassFactory.getStorageClass("StructuredData")
1676 self.composite_refs = [self.makeDatasetRef(f"composite{n}", dimensions, sc, dataId) for n in range(3)]
1677 self.comp_files = []
1678 self.comp_refs = []
1679 for n, ref in enumerate(self.composite_refs):
1680 component_refs = []
1681 component_files = []
1682 for component in sc.components:
1683 component_ref = ref.makeComponentRef(component)
1684 file = root_uri.join(f"composite_file-{n}-{component}.txt")
1685 component_refs.append(component_ref)
1686 component_files.append(file)
1687 file.write(b"9876543210")
1689 self.comp_files.append(component_files)
1690 self.comp_refs.append(component_refs)
1692 def tearDown(self) -> None:
1693 if self.root is not None and os.path.exists(self.root):
1694 shutil.rmtree(self.root, ignore_errors=True)
1696 def _make_cache_manager(self, config_str: str) -> DatastoreCacheManager:
1697 config = Config.fromYaml(config_str)
1698 return DatastoreCacheManager(DatastoreCacheManagerConfig(config), universe=self.universe)
1700 def testNoCacheDir(self) -> None:
1701 config_str = """
1702cached:
1703 root: null
1704 cacheable:
1705 metric0: true
1706 """
1707 cache_manager = self._make_cache_manager(config_str)
1709 # Look inside to check we don't have a cache directory
1710 self.assertIsNone(cache_manager._cache_directory)
1712 self.assertCache(cache_manager)
1714 # Test that the cache directory is marked temporary
1715 self.assertTrue(cache_manager.cache_directory.isTemporary)
1717 def testNoCacheDirReversed(self) -> None:
1718 """Use default caching status and metric1 to false"""
1719 config_str = """
1720cached:
1721 root: null
1722 default: true
1723 cacheable:
1724 metric1: false
1725 """
1726 cache_manager = self._make_cache_manager(config_str)
1728 self.assertCache(cache_manager)
1730 def testEnvvarCacheDir(self) -> None:
1731 config_str = f"""
1732cached:
1733 root: '{self.root}'
1734 cacheable:
1735 metric0: true
1736 """
1738 root = ResourcePath(self.root, forceDirectory=True)
1739 env_dir = root.join("somewhere", forceDirectory=True)
1740 elsewhere = root.join("elsewhere", forceDirectory=True)
1742 # Environment variable should override the config value.
1743 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath}):
1744 cache_manager = self._make_cache_manager(config_str)
1745 self.assertEqual(cache_manager.cache_directory, env_dir)
1747 # This environment variable should not override the config value.
1748 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}):
1749 cache_manager = self._make_cache_manager(config_str)
1750 self.assertEqual(cache_manager.cache_directory, root)
1752 # No default setting.
1753 config_str = """
1754cached:
1755 root: null
1756 default: true
1757 cacheable:
1758 metric1: false
1759 """
1760 cache_manager = self._make_cache_manager(config_str)
1762 # This environment variable should override the config value.
1763 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": env_dir.ospath}):
1764 cache_manager = self._make_cache_manager(config_str)
1765 self.assertEqual(cache_manager.cache_directory, env_dir)
1767 # If both environment variables are set the main (not IF_UNSET)
1768 # variable should win.
1769 with unittest.mock.patch.dict(
1770 os.environ,
1771 {
1772 "DAF_BUTLER_CACHE_DIRECTORY": env_dir.ospath,
1773 "DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": elsewhere.ospath,
1774 },
1775 ):
1776 cache_manager = self._make_cache_manager(config_str)
1777 self.assertEqual(cache_manager.cache_directory, env_dir)
1779 # Use the API to set the environment variable, making sure that the
1780 # variable is reset on exit.
1781 with unittest.mock.patch.dict(
1782 os.environ,
1783 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""},
1784 ):
1785 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()
1786 self.assertTrue(defined)
1787 cache_manager = self._make_cache_manager(config_str)
1788 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True))
1790 # Now create the cache manager ahead of time and set the fallback
1791 # later.
1792 cache_manager = self._make_cache_manager(config_str)
1793 self.assertIsNone(cache_manager._cache_directory)
1794 with unittest.mock.patch.dict(
1795 os.environ,
1796 {"DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET": ""},
1797 ):
1798 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()
1799 self.assertTrue(defined)
1800 self.assertEqual(cache_manager.cache_directory, ResourcePath(cache_dir, forceDirectory=True))
1802 def testExplicitCacheDir(self) -> None:
1803 config_str = f"""
1804cached:
1805 root: '{self.root}'
1806 cacheable:
1807 metric0: true
1808 """
1809 cache_manager = self._make_cache_manager(config_str)
1811 # Look inside to check we do have a cache directory.
1812 self.assertEqual(cache_manager.cache_directory, ResourcePath(self.root, forceDirectory=True))
1814 self.assertCache(cache_manager)
1816 # Test that the cache directory is not marked temporary
1817 self.assertFalse(cache_manager.cache_directory.isTemporary)
1819 def testUnexpectedFilesInCacheDir(self) -> None:
1820 """Test for regression of a bug where extraneous files in a cache
1821 directory would cause all cache lookups to raise an exception.
1822 """
1823 config_str = f"""
1824cached:
1825 root: '{self.root}'
1826 cacheable:
1827 metric0: true
1828 """
1830 for filename in ["unexpected.txt", "unexpected", "un_expected", "un_expected.txt"]:
1831 unexpected_file = os.path.join(self.root, filename)
1832 with open(unexpected_file, "w") as fh:
1833 fh.write("test")
1835 cache_manager = self._make_cache_manager(config_str)
1836 cache_manager.scan_cache()
1837 self.assertCache(cache_manager)
1839 def assertCache(self, cache_manager: DatastoreCacheManager) -> None:
1840 self.assertTrue(cache_manager.should_be_cached(self.refs[0]))
1841 self.assertFalse(cache_manager.should_be_cached(self.refs[1]))
1843 uri = cache_manager.move_to_cache(self.files[0], self.refs[0])
1844 self.assertIsInstance(uri, ResourcePath)
1845 self.assertIsNone(cache_manager.move_to_cache(self.files[1], self.refs[1]))
1847 # Check presence in cache using ref and then using file extension.
1848 self.assertFalse(cache_manager.known_to_cache(self.refs[1]))
1849 self.assertTrue(cache_manager.known_to_cache(self.refs[0]))
1850 self.assertFalse(cache_manager.known_to_cache(self.refs[1], self.files[1].getExtension()))
1851 self.assertTrue(cache_manager.known_to_cache(self.refs[0], self.files[0].getExtension()))
1853 # Cached file should no longer exist but uncached file should be
1854 # unaffected.
1855 self.assertFalse(self.files[0].exists())
1856 self.assertTrue(self.files[1].exists())
1858 # Should find this file and it should be within the cache directory.
1859 with cache_manager.find_in_cache(self.refs[0], ".txt") as found:
1860 self.assertTrue(found.exists())
1861 self.assertIsNotNone(found.relative_to(cache_manager.cache_directory))
1863 # Should not be able to find these in cache
1864 with cache_manager.find_in_cache(self.refs[0], ".fits") as found:
1865 self.assertIsNone(found)
1866 with cache_manager.find_in_cache(self.refs[1], ".fits") as found:
1867 self.assertIsNone(found)
1869 def testNoCache(self) -> None:
1870 cache_manager = DatastoreDisabledCacheManager("", universe=self.universe)
1871 for uri, ref in zip(self.files, self.refs, strict=True):
1872 self.assertFalse(cache_manager.should_be_cached(ref))
1873 self.assertIsNone(cache_manager.move_to_cache(uri, ref))
1874 self.assertFalse(cache_manager.known_to_cache(ref))
1875 with cache_manager.find_in_cache(ref, ".txt") as found:
1876 self.assertIsNone(found, msg=f"{cache_manager}")
1878 def _expiration_config(self, mode: str, threshold: int) -> str:
1879 return f"""
1880cached:
1881 default: true
1882 expiry:
1883 mode: {mode}
1884 threshold: {threshold}
1885 cacheable:
1886 unused: true
1887 """
1889 def testCacheExpiryFiles(self) -> None:
1890 threshold = 2 # Keep at least 2 files.
1891 mode = "files"
1892 config_str = self._expiration_config(mode, threshold)
1894 cache_manager = self._make_cache_manager(config_str)
1896 # Check that an empty cache returns unknown for arbitrary ref
1897 self.assertFalse(cache_manager.known_to_cache(self.refs[0]))
1899 # Should end with datasets: 2, 3, 4
1900 self.assertExpiration(cache_manager, 5, threshold + 1)
1901 self.assertIn(f"{mode}={threshold}", str(cache_manager))
1903 # Check that we will not expire a file that is actively in use.
1904 with cache_manager.find_in_cache(self.refs[2], ".txt") as found:
1905 self.assertIsNotNone(found)
1907 # Trigger cache expiration that should remove the file
1908 # we just retrieved. Should now have: 3, 4, 5
1909 cached = cache_manager.move_to_cache(self.files[5], self.refs[5])
1910 self.assertIsNotNone(cached)
1912 # Cache should still report the standard file count.
1913 self.assertEqual(cache_manager.file_count, threshold + 1)
1915 # Add additional entry to cache.
1916 # Should now have 4, 5, 6
1917 cached = cache_manager.move_to_cache(self.files[6], self.refs[6])
1918 self.assertIsNotNone(cached)
1920 # Is the file still there?
1921 self.assertTrue(found.exists())
1923 # Can we read it?
1924 data = found.read()
1925 self.assertGreater(len(data), 0)
1927 # Outside context the file should no longer exist.
1928 self.assertFalse(found.exists())
1930 # File count should not have changed.
1931 self.assertEqual(cache_manager.file_count, threshold + 1)
1933 # Dataset 2 was in the exempt directory but because hardlinks
1934 # are used it was deleted from the main cache during cache expiry
1935 # above and so should no longer be found.
1936 with cache_manager.find_in_cache(self.refs[2], ".txt") as found:
1937 self.assertIsNone(found)
1939 # And the one stored after it is also gone.
1940 with cache_manager.find_in_cache(self.refs[3], ".txt") as found:
1941 self.assertIsNone(found)
1943 # But dataset 4 is present.
1944 with cache_manager.find_in_cache(self.refs[4], ".txt") as found:
1945 self.assertIsNotNone(found)
1947 # Adding a new dataset to the cache should now delete it.
1948 cache_manager.move_to_cache(self.files[7], self.refs[7])
1950 with cache_manager.find_in_cache(self.refs[2], ".txt") as found:
1951 self.assertIsNone(found)
1953 def testCacheExpiryDatasets(self) -> None:
1954 threshold = 2 # Keep 2 datasets.
1955 mode = "datasets"
1956 config_str = self._expiration_config(mode, threshold)
1958 cache_manager = self._make_cache_manager(config_str)
1959 self.assertExpiration(cache_manager, 5, threshold + 1)
1960 self.assertIn(f"{mode}={threshold}", str(cache_manager))
1962 def testCacheExpiryDatasetsFromDisabled(self) -> None:
1963 threshold = 2
1964 mode = "datasets"
1965 with unittest.mock.patch.dict(
1966 os.environ,
1967 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"},
1968 ):
1969 cache_manager = DatastoreCacheManager.create_disabled(universe=DimensionUniverse())
1970 self.assertExpiration(cache_manager, 5, threshold + 1)
1971 self.assertIn(f"{mode}={threshold}", str(cache_manager))
1973 def testExpirationModeOverride(self) -> None:
1974 threshold = 2 # Keep 2 datasets.
1975 mode = "datasets"
1976 config_str = self._expiration_config(mode, threshold)
1978 mode = "size"
1979 threshold = 55
1980 with unittest.mock.patch.dict(
1981 os.environ,
1982 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": f"{mode}={threshold}"},
1983 ):
1984 cache_manager = self._make_cache_manager(config_str)
1985 self.assertExpiration(cache_manager, 10, 6)
1986 self.assertIn(f"{mode}={threshold}", str(cache_manager))
1988 # Check we get a warning with unrecognized form.
1989 with unittest.mock.patch.dict(
1990 os.environ,
1991 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something"},
1992 ):
1993 with self.assertLogs(level="WARNING") as cm:
1994 self._make_cache_manager(config_str)
1995 self.assertIn("Unrecognized form (something)", cm.output[0])
1997 with unittest.mock.patch.dict(
1998 os.environ,
1999 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "something=5"},
2000 ):
2001 with self.assertRaises(ValueError) as cm:
2002 self._make_cache_manager(config_str)
2003 self.assertIn("Unrecognized value", str(cm.exception))
2005 def testMissingThreshold(self) -> None:
2006 threshold = ""
2007 mode = "datasets"
2008 config_str = self._expiration_config(mode, threshold)
2010 with self.assertRaises(ValueError) as cm:
2011 self._make_cache_manager(config_str)
2012 self.assertIn("Cache expiration threshold", str(cm.exception))
2014 def testCacheExpiryDatasetsComposite(self) -> None:
2015 threshold = 2 # Keep 2 datasets.
2016 mode = "datasets"
2017 config_str = self._expiration_config(mode, threshold)
2019 cache_manager = self._make_cache_manager(config_str)
2021 n_datasets = 3
2022 for i in range(n_datasets):
2023 for component_file, component_ref in zip(self.comp_files[i], self.comp_refs[i], strict=True):
2024 cached = cache_manager.move_to_cache(component_file, component_ref)
2025 self.assertIsNotNone(cached)
2026 self.assertTrue(cache_manager.known_to_cache(component_ref))
2027 self.assertTrue(cache_manager.known_to_cache(component_ref.makeCompositeRef()))
2028 self.assertTrue(cache_manager.known_to_cache(component_ref, component_file.getExtension()))
2030 self.assertEqual(cache_manager.file_count, 6) # 2 datasets each of 3 files
2032 # Write two new non-composite and the number of files should drop.
2033 self.assertExpiration(cache_manager, 2, 5)
2035 def testCacheExpirySize(self) -> None:
2036 threshold = 55 # Each file is 10 bytes
2037 mode = "size"
2038 config_str = self._expiration_config(mode, threshold)
2040 cache_manager = self._make_cache_manager(config_str)
2041 self.assertExpiration(cache_manager, 10, 6)
2042 self.assertIn(f"{mode}={threshold}", str(cache_manager))
2044 def testDisabledCache(self) -> None:
2045 # Configure an active cache but disable via environment.
2046 threshold = 2
2047 mode = "datasets"
2048 config_str = self._expiration_config(mode, threshold)
2050 with unittest.mock.patch.dict(
2051 os.environ,
2052 {"DAF_BUTLER_CACHE_EXPIRATION_MODE": "disabled"},
2053 ):
2054 env_cache_manager = self._make_cache_manager(config_str)
2056 # Configure to be disabled
2057 threshold = 0
2058 mode = "disabled"
2059 config_str = self._expiration_config(mode, threshold)
2060 cfg_cache_manager = self._make_cache_manager(config_str)
2062 for cache_manager in (
2063 cfg_cache_manager,
2064 env_cache_manager,
2065 DatastoreCacheManager.create_disabled(universe=DimensionUniverse()),
2066 ):
2067 for uri, ref in zip(self.files, self.refs, strict=True):
2068 self.assertFalse(cache_manager.should_be_cached(ref))
2069 self.assertIsNone(cache_manager.move_to_cache(uri, ref))
2070 self.assertFalse(cache_manager.known_to_cache(ref))
2071 with cache_manager.find_in_cache(ref, ".txt") as found:
2072 self.assertIsNone(found, msg=f"{cache_manager}")
2073 self.assertIn("disabled", str(cache_manager))
2075 def assertExpiration(
2076 self, cache_manager: DatastoreCacheManager, n_datasets: int, n_retained: int
2077 ) -> None:
2078 """Insert the datasets and then check the number retained."""
2079 for i in range(n_datasets):
2080 cached = cache_manager.move_to_cache(self.files[i], self.refs[i])
2081 self.assertIsNotNone(cached)
2083 self.assertEqual(cache_manager.file_count, n_retained)
2085 # The oldest file should not be in the cache any more.
2086 for i in range(n_datasets):
2087 with cache_manager.find_in_cache(self.refs[i], ".txt") as found:
2088 if i >= n_datasets - n_retained:
2089 self.assertIsInstance(found, ResourcePath)
2090 else:
2091 self.assertIsNone(found)
2093 def testCacheExpiryAge(self) -> None:
2094 threshold = 1 # Expire older than 2 seconds
2095 mode = "age"
2096 config_str = self._expiration_config(mode, threshold)
2098 cache_manager = self._make_cache_manager(config_str)
2099 self.assertIn(f"{mode}={threshold}", str(cache_manager))
2101 # Insert 3 files, then sleep, then insert more.
2102 for i in range(2):
2103 cached = cache_manager.move_to_cache(self.files[i], self.refs[i])
2104 self.assertIsNotNone(cached)
2105 time.sleep(2.0)
2106 for j in range(4):
2107 i = 2 + j # Continue the counting
2108 cached = cache_manager.move_to_cache(self.files[i], self.refs[i])
2109 self.assertIsNotNone(cached)
2111 # Only the files written after the sleep should exist.
2112 self.assertEqual(cache_manager.file_count, 4)
2113 with cache_manager.find_in_cache(self.refs[1], ".txt") as found:
2114 self.assertIsNone(found)
2115 with cache_manager.find_in_cache(self.refs[2], ".txt") as found:
2116 self.assertIsInstance(found, ResourcePath)
2119class NullDatastoreTestCase(DatasetTestHelper, unittest.TestCase):
2120 """Test the null datastore."""
2122 storageClassFactory = StorageClassFactory()
2124 def test_basics(self) -> None:
2125 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
2126 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
2128 null = NullDatastore(None, None)
2130 self.assertFalse(null.exists(ref))
2131 self.assertFalse(null.knows(ref))
2132 knows = null.knows_these([ref])
2133 self.assertFalse(knows[ref])
2134 null.validateConfiguration(ref)
2136 with self.assertRaises(FileNotFoundError):
2137 null.get(ref)
2138 with self.assertRaises(NotImplementedError):
2139 null.put("", ref)
2140 with self.assertRaises(FileNotFoundError):
2141 null.getURI(ref)
2142 with self.assertRaises(FileNotFoundError):
2143 null.getURIs(ref)
2144 with self.assertRaises(FileNotFoundError):
2145 null.getManyURIs([ref])
2146 with self.assertRaises(NotImplementedError):
2147 null.getLookupKeys()
2148 with self.assertRaises(NotImplementedError):
2149 null.import_records({})
2150 with self.assertRaises(NotImplementedError):
2151 null.export_records([])
2152 with self.assertRaises(NotImplementedError):
2153 null.export_predicted_records([])
2154 with self.assertRaises(NotImplementedError):
2155 null.export([ref])
2156 with self.assertRaises(NotImplementedError):
2157 null.transfer(null, ref)
2158 with self.assertRaises(NotImplementedError):
2159 null.emptyTrash()
2160 with self.assertRaises(NotImplementedError):
2161 null.trash(ref)
2162 with self.assertRaises(NotImplementedError):
2163 null.forget([ref])
2164 with self.assertRaises(NotImplementedError):
2165 null.remove(ref)
2166 with self.assertRaises(NotImplementedError):
2167 null.retrieveArtifacts([ref], ResourcePath("."))
2168 with self.assertRaises(NotImplementedError):
2169 null.transfer_from(null, [ref])
2170 with self.assertRaises(NotImplementedError):
2171 null.ingest()
2174class DatasetRefURIsTestCase(unittest.TestCase):
2175 """Tests for DatasetRefURIs."""
2177 def testSequenceAccess(self) -> None:
2178 """Verify that DatasetRefURIs can be treated like a two-item tuple."""
2179 uris = DatasetRefURIs()
2181 self.assertEqual(len(uris), 2)
2182 self.assertEqual(uris[0], None)
2183 self.assertEqual(uris[1], {})
2185 primaryURI = ResourcePath("1/2/3")
2186 componentURI = ResourcePath("a/b/c")
2188 # affirm that DatasetRefURIs does not support MutableSequence functions
2189 with self.assertRaises(TypeError):
2190 uris[0] = primaryURI
2191 with self.assertRaises(TypeError):
2192 uris[1] = {"foo": componentURI}
2194 # but DatasetRefURIs can be set by property name:
2195 uris.primaryURI = primaryURI
2196 uris.componentURIs = {"foo": componentURI}
2197 self.assertEqual(uris.primaryURI, primaryURI)
2198 self.assertEqual(uris[0], primaryURI)
2200 primary, components = uris
2201 self.assertEqual(primary, primaryURI)
2202 self.assertEqual(components, {"foo": componentURI})
2204 def testRepr(self) -> None:
2205 """Verify __repr__ output."""
2206 uris = DatasetRefURIs(ResourcePath("/1/2/3"), {"comp": ResourcePath("/a/b/c")})
2207 self.assertEqual(
2208 repr(uris),
2209 'DatasetRefURIs(ResourcePath("file:///1/2/3"), {\'comp\': ResourcePath("file:///a/b/c")})',
2210 )
2213class StoredFileInfoTestCase(DatasetTestHelper, unittest.TestCase):
2214 """Test the StoredFileInfo class."""
2216 storageClassFactory = StorageClassFactory()
2218 def test_StoredFileInfo(self) -> None:
2219 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
2220 ref = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
2222 record = dict(
2223 storage_class="StructuredDataDict",
2224 formatter="lsst.daf.butler.Formatter",
2225 path="a/b/c.txt",
2226 component="component",
2227 checksum=None,
2228 file_size=5,
2229 )
2230 info = StoredFileInfo.from_record(record)
2232 self.assertEqual(info.to_record(), record)
2234 ref2 = self.makeDatasetRef("metric", DimensionUniverse().empty, storageClass, {})
2235 rebased = info.rebase(ref2)
2236 self.assertEqual(rebased.rebase(ref), info)
2238 with self.assertRaises(TypeError):
2239 rebased.update(formatter=42)
2241 with self.assertRaises(ValueError):
2242 rebased.update(something=42, new="42")
2244 # Check that pickle works on StoredFileInfo.
2245 pickled_info = pickle.dumps(info)
2246 unpickled_info = pickle.loads(pickled_info)
2247 self.assertEqual(unpickled_info, info)
2249 def test_make_datastore_path_relative(self):
2250 self.assertEqual(make_datastore_path_relative("a/relative/path"), "a/relative/path")
2251 self.assertEqual(make_datastore_path_relative("path/with#fragment"), "path/with#fragment")
2252 self.assertEqual(make_datastore_path_relative("http://server.com/some/path"), "some/path")
2253 self.assertEqual(make_datastore_path_relative("http://server.com/some/path#frag"), "some/path#frag")
2256@contextlib.contextmanager
2257def _temp_yaml_file(data: Any) -> Iterator[str]:
2258 fh = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml")
2259 try:
2260 yaml.dump(data, stream=fh)
2261 fh.flush()
2262 yield fh.name
2263 finally:
2264 # Some tests delete the file
2265 with contextlib.suppress(FileNotFoundError):
2266 fh.close()
2269if __name__ == "__main__":
2270 unittest.main()