Coverage for tests/test_butler.py: 14%
1700 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-20 01:50 -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/>.
28"""Tests for Butler."""
30from __future__ import annotations
32import json
33import logging
34import os
35import pathlib
36import pickle
37import posixpath
38import random
39import re
40import shutil
41import string
42import tempfile
43import unittest
44import unittest.mock
45import uuid
46from collections.abc import Callable, Mapping
47from typing import TYPE_CHECKING, Any, cast
49try:
50 import boto3
51 import botocore
53 from lsst.resources.s3utils import clean_test_environment_for_s3
55 try:
56 from moto import mock_aws # v5
57 except ImportError:
58 from moto import mock_s3 as mock_aws
59except ImportError:
60 boto3 = None
62 def mock_aws(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-untyped-def]
63 """No-op decorator in case moto mock_aws can not be imported."""
64 return None
67import astropy.time
68from sqlalchemy.exc import IntegrityError
70from lsst.daf.butler import (
71 Butler,
72 ButlerConfig,
73 ButlerMetrics,
74 ButlerRepoIndex,
75 CollectionCycleError,
76 CollectionType,
77 Config,
78 DataCoordinate,
79 DatasetExistence,
80 DatasetNotFoundError,
81 DatasetProvenance,
82 DatasetRef,
83 DatasetType,
84 FileDataset,
85 NoDefaultCollectionError,
86 StorageClassFactory,
87 ValidationError,
88 script,
89)
90from lsst.daf.butler._rubin.file_datasets import transfer_datasets_to_datastore
91from lsst.daf.butler.datastore import NullDatastore
92from lsst.daf.butler.datastore.file_templates import FileTemplate, FileTemplateValidationError
93from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex
94from lsst.daf.butler.datastores.fileDatastore import FileDatastore
95from lsst.daf.butler.direct_butler import DirectButler
96from lsst.daf.butler.registry import (
97 CollectionError,
98 CollectionTypeError,
99 ConflictingDefinitionError,
100 DataIdValueError,
101 DatasetTypeExpressionError,
102 MissingCollectionError,
103 OrphanedRecordError,
104)
105from lsst.daf.butler.registry.sql_registry import SqlRegistry
106from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG
107from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter
108from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db
109from lsst.daf.butler.tests.server_available import butler_server_import_error, butler_server_is_available
110from lsst.daf.butler.tests.utils import (
111 MetricTestRepo,
112 TestCaseMixin,
113 makeTestTempDir,
114 removeTestTempDir,
115 safeTestTempDir,
116)
117from lsst.resources import ResourcePath
118from lsst.resources.http import HttpResourcePath
119from lsst.utils import doImportType
120from lsst.utils.introspection import get_full_type_name
122if butler_server_is_available:
123 from lsst.daf.butler.tests.server import create_test_server
126if TYPE_CHECKING:
127 import types
129 from lsst.daf.butler import DimensionGroup, Registry, StorageClass
131TESTDIR = os.path.abspath(os.path.dirname(__file__))
134def clean_environment() -> None:
135 """Remove external environment variables that affect the tests."""
136 for k in ("DAF_BUTLER_REPOSITORY_INDEX",):
137 os.environ.pop(k, None)
140def makeExampleMetrics() -> MetricsExample:
141 """Return example dataset suitable for tests."""
142 return MetricsExample(
143 {"AM1": 5.2, "AM2": 30.6},
144 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
145 [563, 234, 456.7, 752, 8, 9, 27],
146 )
149class TransactionTestError(Exception):
150 """Specific error for testing transactions, to prevent misdiagnosing
151 that might otherwise occur when a standard exception is used.
152 """
154 pass
157class ButlerConfigTests(unittest.TestCase):
158 """Simple tests for ButlerConfig that are not tested in any other test
159 cases.
160 """
162 def testSearchPath(self) -> None:
163 configFile = os.path.join(TESTDIR, "config", "basic", "butler.yaml")
164 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
165 config1 = ButlerConfig(configFile)
166 self.assertNotIn("testConfigs", "\n".join(cm.output))
168 overrideDirectory = os.path.join(TESTDIR, "config", "testConfigs")
169 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
170 config2 = ButlerConfig(configFile, searchPaths=[overrideDirectory])
171 self.assertIn("testConfigs", "\n".join(cm.output))
173 key = ("datastore", "records", "table")
174 self.assertNotEqual(config1[key], config2[key])
175 self.assertEqual(config2[key], "override_record")
178class ButlerPutGetTests(TestCaseMixin):
179 """Helper method for running a suite of put/get tests from different
180 butler configurations.
181 """
183 root: str
184 default_run = "ingésτ😺"
185 storageClassFactory: StorageClassFactory
186 configFile: str | None
187 tmpConfigFile: str
189 @staticmethod
190 def addDatasetType(
191 datasetTypeName: str, dimensions: DimensionGroup, storageClass: StorageClass | str, registry: Registry
192 ) -> DatasetType:
193 """Create a DatasetType and register it"""
194 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
195 registry.registerDatasetType(datasetType)
196 return datasetType
198 @classmethod
199 def setUpClass(cls) -> None:
200 cls.storageClassFactory = StorageClassFactory()
201 if cls.configFile is not None:
202 cls.storageClassFactory.addFromConfig(cls.configFile)
204 def assertGetComponents(
205 self,
206 butler: Butler,
207 datasetRef: DatasetRef,
208 components: tuple[str, ...],
209 reference: Any,
210 collections: Any = None,
211 ) -> None:
212 datasetType = datasetRef.datasetType
213 dataId = datasetRef.dataId
214 deferred = butler.getDeferred(datasetRef)
216 for component in components:
217 compTypeName = datasetType.componentTypeName(component)
218 result = butler.get(compTypeName, dataId, collections=collections)
219 self.assertEqual(result, getattr(reference, component))
220 result_deferred = deferred.get(component=component)
221 self.assertEqual(result_deferred, result)
223 def tearDown(self) -> None:
224 if self.root is not None:
225 removeTestTempDir(self.root)
227 def create_empty_butler(
228 self, run: str | None = None, writeable: bool | None = None, metrics: ButlerMetrics | None = None
229 ):
230 """Create a Butler for the test repository, without inserting test
231 data.
232 """
233 butler = Butler.from_config(self.tmpConfigFile, run=run, writeable=writeable, metrics=metrics)
234 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
235 return butler
237 def create_butler(
238 self,
239 run: str,
240 storageClass: StorageClass | str,
241 datasetTypeName: str,
242 metrics: ButlerMetrics | None = None,
243 ) -> tuple[Butler, DatasetType]:
244 """Create a Butler for the test repository and insert some test data
245 into it.
246 """
247 butler = self.create_empty_butler(run=run, metrics=metrics)
249 collections = set(butler.collections.query("*"))
250 self.assertEqual(collections, {run})
251 # Create and register a DatasetType
252 dimensions = butler.dimensions.conform(["instrument", "visit"])
254 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
256 # Add needed Dimensions
257 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
258 butler.registry.insertDimensionData(
259 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
260 )
261 butler.registry.insertDimensionData(
262 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
263 )
264 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101})
265 visit_start = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
266 visit_end = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
267 butler.registry.insertDimensionData(
268 "visit",
269 {
270 "instrument": "DummyCamComp",
271 "id": 423,
272 "name": "fourtwentythree",
273 "physical_filter": "d-r",
274 "datetime_begin": visit_start,
275 "datetime_end": visit_end,
276 "day_obs": 20200101,
277 },
278 )
280 # Add more visits for some later tests
281 for visit_id in (424, 425):
282 butler.registry.insertDimensionData(
283 "visit",
284 {
285 "instrument": "DummyCamComp",
286 "id": visit_id,
287 "name": f"fourtwentyfour_{visit_id}",
288 "physical_filter": "d-r",
289 "day_obs": 20200101,
290 },
291 )
292 return butler, datasetType
294 def runPutGetTest(self, storageClass: StorageClass, datasetTypeName: str) -> Butler:
295 # New datasets will be added to run and tag, but we will only look in
296 # tag when looking up datasets.
297 run = self.default_run
298 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
299 assert butler.run is not None
301 # Create and store a dataset
302 metric = makeExampleMetrics()
303 dataId = butler.registry.expandDataId({"instrument": "DummyCamComp", "visit": 423})
305 # Dataset should not exist if we haven't added it
306 with self.assertRaises(DatasetNotFoundError):
307 butler.get(datasetTypeName, dataId)
309 # Put and remove the dataset once as a DatasetRef, once as a dataId,
310 # and once with a DatasetType
312 # Keep track of any collections we add and do not clean up
313 expected_collections = {run}
315 counter = 0
316 ref = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run="put_run_1")
317 args = tuple[DatasetRef] | tuple[str | DatasetType, DataCoordinate]
318 for args in ((ref,), (datasetTypeName, dataId), (datasetType, dataId)):
319 # Since we are using subTest we can get cascading failures
320 # here with the first attempt failing and the others failing
321 # immediately because the dataset already exists. Work around
322 # this by using a distinct run collection each time
323 counter += 1
324 this_run = f"put_run_{counter}"
325 butler.collections.register(this_run)
326 expected_collections.update({this_run})
328 with self.subTest(args=repr(args)):
329 kwargs: dict[str, Any] = {}
330 if not isinstance(args[0], DatasetRef): # type: ignore
331 kwargs["run"] = this_run
332 ref = butler.put(metric, *args, **kwargs)
333 self.assertIsInstance(ref, DatasetRef)
335 # Test get of a ref.
336 metricOut = butler.get(ref)
337 self.assertEqual(metric, metricOut)
338 # Test get
339 metricOut = butler.get(ref.datasetType.name, dataId, collections=this_run)
340 self.assertEqual(metric, metricOut)
341 # Test get with a datasetRef
342 metricOut = butler.get(ref)
343 self.assertEqual(metric, metricOut)
344 # Test getDeferred with dataId
345 metricOut = butler.getDeferred(ref.datasetType.name, dataId, collections=this_run).get()
346 self.assertEqual(metric, metricOut)
347 # Test getDeferred with a ref
348 metricOut = butler.getDeferred(ref).get()
349 self.assertEqual(metric, metricOut)
351 # Check we can get components
352 if storageClass.isComposite():
353 self.assertGetComponents(
354 butler, ref, ("summary", "data", "output"), metric, collections=this_run
355 )
357 primary_uri, secondary_uris = butler.getURIs(ref)
358 n_uris = len(secondary_uris)
359 if primary_uri:
360 n_uris += 1
362 # Can the artifacts themselves be retrieved?
363 if not butler._datastore.isEphemeral:
364 # Create a temporary directory to hold the retrieved
365 # artifacts.
366 with tempfile.TemporaryDirectory(
367 prefix="butler-artifacts-", ignore_cleanup_errors=True
368 ) as artifact_root:
369 root_uri = ResourcePath(artifact_root, forceDirectory=True)
371 for preserve_path in (True, False):
372 destination = root_uri.join(f"{preserve_path}_{counter}/")
373 log = logging.getLogger("lsst.x")
374 log.debug("Using destination %s for args %s", destination, args)
375 # Use copy so that we can test that overwrite
376 # protection works (using "auto" for File URIs
377 # would use hard links and subsequent transfer
378 # would work because it knows they are the same
379 # file).
380 transferred = butler.retrieveArtifacts(
381 [ref], destination, preserve_path=preserve_path, transfer="copy"
382 )
383 self.assertGreater(len(transferred), 0)
384 artifacts = list(ResourcePath.findFileResources([destination]))
385 # Filter out the index file.
386 artifacts = [a for a in artifacts if a.basename() != ZipIndex.index_name]
387 self.assertEqual(set(transferred), set(artifacts))
389 for artifact in transferred:
390 path_in_destination = artifact.relative_to(destination)
391 self.assertIsNotNone(path_in_destination)
392 assert path_in_destination is not None
394 # When path is not preserved there should not
395 # be any path separators.
396 num_seps = path_in_destination.count("/")
397 if preserve_path:
398 self.assertGreater(num_seps, 0)
399 else:
400 self.assertEqual(num_seps, 0)
402 self.assertEqual(
403 len(artifacts),
404 n_uris,
405 "Comparing expected artifacts vs actual:"
406 f" {artifacts} vs {primary_uri} and {secondary_uris}",
407 )
409 if preserve_path:
410 # No need to run these twice
411 with self.assertRaises(ValueError):
412 butler.retrieveArtifacts([ref], destination, transfer="move")
414 with self.assertRaisesRegex(
415 ValueError, "^Destination location must refer to a directory"
416 ):
417 butler.retrieveArtifacts(
418 [ref], ResourcePath("/some/file.txt", forceDirectory=False)
419 )
421 with self.assertRaises(FileExistsError):
422 butler.retrieveArtifacts([ref], destination)
424 transferred_again = butler.retrieveArtifacts(
425 [ref], destination, preserve_path=preserve_path, overwrite=True
426 )
427 self.assertEqual(set(transferred_again), set(transferred))
429 # Now remove the dataset completely.
430 butler.pruneDatasets([ref], purge=True, unstore=True)
431 # Lookup with original args should still fail.
432 kwargs = {"collections": this_run}
433 if isinstance(args[0], DatasetRef):
434 kwargs = {} # Prevent warning from being issued.
435 self.assertFalse(butler.exists(*args, **kwargs))
436 # get() should still fail.
437 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
438 butler.get(ref)
439 # Registry shouldn't be able to find it by dataset_id anymore.
440 self.assertIsNone(butler.get_dataset(ref.id))
442 # Do explicit registry removal since we know they are
443 # empty
444 butler.collections.x_remove(this_run)
445 expected_collections.remove(this_run)
447 # Create DatasetRef for put using default run.
448 refIn = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run=butler.run)
450 # Check that getDeferred fails with standalone ref.
451 with self.assertRaises(LookupError):
452 butler.getDeferred(refIn)
454 # Put the dataset again, since the last thing we did was remove it
455 # and we want to use the default collection.
456 ref = butler.put(metric, refIn)
458 # Get with parameters
459 stop = 4
460 sliced = butler.get(ref, parameters={"slice": slice(stop)})
461 self.assertNotEqual(metric, sliced)
462 self.assertEqual(metric.summary, sliced.summary)
463 self.assertEqual(metric.output, sliced.output)
464 assert metric.data is not None # for mypy
465 self.assertEqual(metric.data[:stop], sliced.data)
466 # getDeferred with parameters
467 sliced = butler.getDeferred(ref, parameters={"slice": slice(stop)}).get()
468 self.assertNotEqual(metric, sliced)
469 self.assertEqual(metric.summary, sliced.summary)
470 self.assertEqual(metric.output, sliced.output)
471 self.assertEqual(metric.data[:stop], sliced.data)
472 # getDeferred with deferred parameters
473 sliced = butler.getDeferred(ref).get(parameters={"slice": slice(stop)})
474 self.assertNotEqual(metric, sliced)
475 self.assertEqual(metric.summary, sliced.summary)
476 self.assertEqual(metric.output, sliced.output)
477 self.assertEqual(metric.data[:stop], sliced.data)
479 if storageClass.isComposite():
480 # Check that components can be retrieved
481 metricOut = butler.get(ref.datasetType.name, dataId)
482 compNameS = ref.datasetType.componentTypeName("summary")
483 compNameD = ref.datasetType.componentTypeName("data")
484 summary = butler.get(compNameS, dataId)
485 self.assertEqual(summary, metric.summary)
486 data = butler.get(compNameD, dataId)
487 self.assertEqual(data, metric.data)
489 if "counter" in storageClass.derivedComponents:
490 count = butler.get(ref.datasetType.componentTypeName("counter"), dataId)
491 self.assertEqual(count, len(data))
493 count = butler.get(
494 ref.datasetType.componentTypeName("counter"), dataId, parameters={"slice": slice(stop)}
495 )
496 self.assertEqual(count, stop)
498 compRef = butler.find_dataset(compNameS, dataId, collections=butler.collections.defaults)
499 assert compRef is not None
500 summary = butler.get(compRef)
501 self.assertEqual(summary, metric.summary)
503 # Create a Dataset type that has the same name but is inconsistent.
504 inconsistentDatasetType = DatasetType(
505 datasetTypeName, datasetType.dimensions, self.storageClassFactory.getStorageClass("Config")
506 )
508 # Getting with a dataset type that does not match registry fails
509 with self.assertRaisesRegex(
510 ValueError,
511 "(Supplied dataset type .* inconsistent with registry)"
512 "|(The new storage class .* is not compatible with the existing storage class)",
513 ):
514 butler.get(inconsistentDatasetType, dataId)
516 # Combining a DatasetRef with a dataId should fail
517 with self.assertRaisesRegex(ValueError, "DatasetRef given, cannot use dataId as well"):
518 butler.get(ref, dataId)
519 # Getting with an explicit ref should fail if the id doesn't match.
520 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
521 butler.get(DatasetRef(ref.datasetType, ref.dataId, id=uuid.UUID(int=101), run=butler.run))
523 # Getting a dataset with unknown parameters should fail
524 with self.assertRaisesRegex(KeyError, "Parameter 'unsupported' not understood"):
525 butler.get(ref, parameters={"unsupported": True})
527 # Check we have a collection
528 collections = set(butler.collections.query("*"))
529 self.assertEqual(collections, expected_collections)
531 # Clean up to check that we can remove something that may have
532 # already had a component removed
533 butler.pruneDatasets([ref], unstore=True, purge=True)
535 # Add the same ref again, so we can check that duplicate put fails.
536 ref = butler.put(metric, datasetType, dataId)
538 # Repeat put will fail.
539 with self.assertRaisesRegex(
540 ConflictingDefinitionError, "A database constraint failure was triggered"
541 ):
542 butler.put(metric, datasetType, dataId)
544 # Remove the datastore entry.
545 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
547 # Put will still fail
548 with self.assertRaisesRegex(
549 ConflictingDefinitionError, "A database constraint failure was triggered"
550 ):
551 butler.put(metric, datasetType, dataId)
553 # Repeat the same sequence with resolved ref.
554 butler.pruneDatasets([ref], unstore=True, purge=True)
555 ref = butler.put(metric, refIn)
557 # Repeat put will fail.
558 with self.assertRaisesRegex(ConflictingDefinitionError, "Datastore already contains dataset"):
559 butler.put(metric, refIn)
561 # Remove the datastore entry.
562 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
564 # In case of resolved ref this write will succeed.
565 ref = butler.put(metric, refIn)
567 # Leave the dataset in place since some downstream tests require
568 # something to be present
570 return butler
572 def testDeferredCollectionPassing(self) -> None:
573 # Construct a butler with no run or collection, but make it writeable.
574 butler = self.create_empty_butler(writeable=True)
575 # Create and register a DatasetType
576 dimensions = butler.dimensions.conform(["instrument", "visit"])
577 datasetType = self.addDatasetType(
578 "example", dimensions, self.storageClassFactory.getStorageClass("StructuredData"), butler.registry
579 )
580 # Add needed Dimensions
581 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
582 butler.registry.insertDimensionData(
583 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
584 )
585 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
586 butler.registry.insertDimensionData(
587 "visit",
588 {
589 "instrument": "DummyCamComp",
590 "id": 423,
591 "name": "fourtwentythree",
592 "physical_filter": "d-r",
593 "day_obs": 20250101,
594 },
595 )
596 dataId = {"instrument": "DummyCamComp", "visit": 423}
597 # Create dataset.
598 metric = makeExampleMetrics()
599 # Register a new run and put dataset.
600 run = "deferred"
601 self.assertTrue(butler.collections.register(run))
602 # Second time it will be allowed but indicate no-op
603 self.assertFalse(butler.collections.register(run))
604 ref = butler.put(metric, datasetType, dataId, run=run)
605 # Putting with no run should fail with TypeError.
606 with self.assertRaises(CollectionError):
607 butler.put(metric, datasetType, dataId)
608 # Dataset should exist.
609 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
610 # We should be able to get the dataset back, but with and without
611 # a deferred dataset handle.
612 self.assertEqual(metric, butler.get(datasetType, dataId, collections=[run]))
613 self.assertEqual(metric, butler.getDeferred(datasetType, dataId, collections=[run]).get())
614 # Trying to find the dataset without any collection is an error.
615 with self.assertRaises(NoDefaultCollectionError):
616 butler.exists(datasetType, dataId)
617 with self.assertRaises(CollectionError):
618 butler.get(datasetType, dataId)
619 # Associate the dataset with a different collection.
620 butler.collections.register("tagged", type=CollectionType.TAGGED)
621 butler.registry.associate("tagged", [ref])
622 # Deleting the dataset from the new collection should make it findable
623 # in the original collection.
624 butler.pruneDatasets([ref], tags=["tagged"])
625 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
628class ButlerTests(ButlerPutGetTests):
629 """Tests for Butler."""
631 useTempRoot = True
632 validationCanFail: bool
633 fullConfigKey: str | None
634 registryStr: str | None
635 datastoreName: list[str] | None
636 datastoreStr: list[str]
637 predictionSupported = True
638 """Does getURIs support 'prediction mode'?"""
640 def setUp(self) -> None:
641 """Create a new butler root for each test."""
642 self.root = makeTestTempDir(TESTDIR)
643 Butler.makeRepo(self.root, config=Config(self.configFile))
644 self.tmpConfigFile = os.path.join(self.root, "butler.yaml")
646 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
647 """Return True if two URIs refer to the same resource.
649 Subclasses may override to handle unique requirements.
650 """
651 return uri1 == uri2
653 def testConstructor(self) -> None:
654 """Independent test of constructor."""
655 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
656 self.assertIsInstance(butler, Butler)
658 # Check that butler.yaml is added automatically.
659 if self.tmpConfigFile.endswith(end := "/butler.yaml"):
660 config_dir = self.tmpConfigFile[: -len(end)]
661 butler = Butler.from_config(config_dir, run=self.default_run)
662 self.assertIsInstance(butler, Butler)
664 # Even with a ResourcePath.
665 butler = Butler.from_config(ResourcePath(config_dir, forceDirectory=True), run=self.default_run)
666 self.assertIsInstance(butler, Butler)
668 collections = set(butler.collections.query("*"))
669 self.assertEqual(collections, {self.default_run})
671 # Check that some special characters can be included in run name.
672 special_run = "u@b.c-A"
673 butler_special = Butler.from_config(butler=butler, run=special_run)
674 collections = set(butler_special.registry.queryCollections("*@*"))
675 self.assertEqual(collections, {special_run})
677 butler2 = Butler.from_config(butler=butler, collections=["other"])
678 self.assertEqual(butler2.collections.defaults, ("other",))
679 self.assertIsNone(butler2.run)
680 self.assertEqual(type(butler._datastore), type(butler2._datastore))
681 self.assertEqual(butler._datastore.config, butler2._datastore.config)
683 # Test that we can use an environment variable to find this
684 # repository.
685 butler_index = Config()
686 butler_index["label"] = self.tmpConfigFile
687 for suffix in (".yaml", ".json"):
688 # Ensure that the content differs so that we know that
689 # we aren't reusing the cache.
690 bad_label = f"file://bucket/not_real{suffix}"
691 butler_index["bad_label"] = bad_label
692 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
693 butler_index.dumpToUri(temp_file)
694 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
695 self.assertEqual(Butler.get_known_repos(), {"label", "bad_label"})
696 uri = Butler.get_repo_uri("bad_label")
697 self.assertEqual(uri, ResourcePath(bad_label))
698 uri = Butler.get_repo_uri("label")
699 butler = Butler.from_config(uri, writeable=False)
700 self.assertIsInstance(butler, Butler)
701 butler = Butler.from_config("label", writeable=False)
702 self.assertIsInstance(butler, Butler)
703 with self.assertRaisesRegex(FileNotFoundError, "aliases:.*bad_label"):
704 Butler.from_config("not_there", writeable=False)
705 with self.assertRaisesRegex(FileNotFoundError, "resolved from alias 'bad_label'"):
706 Butler.from_config("bad_label")
707 with self.assertRaises(FileNotFoundError):
708 # Should ignore aliases.
709 Butler.from_config(ResourcePath("label", forceAbsolute=False))
710 with self.assertRaises(KeyError) as cm:
711 Butler.get_repo_uri("missing")
712 self.assertEqual(
713 Butler.get_repo_uri("missing", True), ResourcePath("missing", forceAbsolute=False)
714 )
715 self.assertIn("not known to", str(cm.exception))
716 # Should report no failure.
717 self.assertEqual(ButlerRepoIndex.get_failure_reason(), "")
718 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
719 # Now with empty configuration.
720 butler_index = Config()
721 butler_index.dumpToUri(temp_file)
722 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
723 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases)"):
724 Butler.from_config("label")
725 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
726 # Now with bad contents.
727 with open(temp_file.ospath, "w") as fh:
728 print("'", file=fh)
729 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
730 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases:.*could not be read)"):
731 Butler.from_config("label")
732 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": "file://not_found/x.yaml"}):
733 with self.assertRaises(FileNotFoundError):
734 Butler.get_repo_uri("label")
735 self.assertEqual(Butler.get_known_repos(), set())
737 with self.assertRaisesRegex(FileNotFoundError, "index file not found"):
738 Butler.from_config("label")
740 # Check that we can create Butler when the alias file is not found.
741 butler = Butler.from_config(self.tmpConfigFile, writeable=False)
742 self.assertIsInstance(butler, Butler)
743 with self.assertRaises(RuntimeError) as cm:
744 # No environment variable set.
745 Butler.get_repo_uri("label")
746 self.assertEqual(Butler.get_repo_uri("label", True), ResourcePath("label", forceAbsolute=False))
747 self.assertIn("No repository index defined", str(cm.exception))
748 with self.assertRaisesRegex(FileNotFoundError, "no known aliases.*No repository index"):
749 # No aliases registered.
750 Butler.from_config("not_there")
751 self.assertEqual(Butler.get_known_repos(), set())
753 def testDafButlerRepositories(self):
754 with unittest.mock.patch.dict(
755 os.environ,
756 {"DAF_BUTLER_REPOSITORIES": "label: 'https://someuri.com'\notherLabel: 'https://otheruri.com'\n"},
757 ):
758 self.assertEqual(str(Butler.get_repo_uri("label")), "https://someuri.com")
760 with unittest.mock.patch.dict(
761 os.environ,
762 {
763 "DAF_BUTLER_REPOSITORIES": "label: https://someuri.com",
764 "DAF_BUTLER_REPOSITORY_INDEX": "https://someuri.com",
765 },
766 ):
767 with self.assertRaisesRegex(RuntimeError, "Only one of the environment variables"):
768 Butler.get_repo_uri("label")
770 with unittest.mock.patch.dict(
771 os.environ,
772 {"DAF_BUTLER_REPOSITORIES": "invalid"},
773 ):
774 with self.assertRaisesRegex(ValueError, "Repository index not in expected format"):
775 Butler.get_repo_uri("label")
777 def testBasicPutGet(self) -> None:
778 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
779 self.runPutGetTest(storageClass, "test_metric")
781 def testCompositePutGetConcrete(self) -> None:
782 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadCompNoDisassembly")
783 butler = self.runPutGetTest(storageClass, "test_metric")
785 # Should *not* be disassembled
786 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
787 self.assertEqual(len(datasets), 1)
788 uri, components = butler.getURIs(datasets[0])
789 self.assertIsInstance(uri, ResourcePath)
790 self.assertFalse(components)
791 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
792 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
794 # Predicted dataset
795 if self.predictionSupported:
796 dataId = {"instrument": "DummyCamComp", "visit": 424}
797 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
798 self.assertFalse(components)
799 self.assertIsInstance(uri, ResourcePath)
800 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
801 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
803 def testCompositePutGetVirtual(self) -> None:
804 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp")
805 butler = self.runPutGetTest(storageClass, "test_metric_comp")
807 # Should be disassembled
808 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
809 self.assertEqual(len(datasets), 1)
810 uri, components = butler.getURIs(datasets[0])
812 if butler._datastore.isEphemeral:
813 # Never disassemble in-memory datastore
814 self.assertIsInstance(uri, ResourcePath)
815 self.assertFalse(components)
816 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
817 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
818 else:
819 self.assertIsNone(uri)
820 self.assertEqual(set(components), set(storageClass.components))
821 for compuri in components.values():
822 self.assertIsInstance(compuri, ResourcePath)
823 self.assertIn("423", str(compuri), f"Checking visit is in URI {compuri}")
824 self.assertEqual(compuri.fragment, "", f"Checking absence of fragment in {compuri}")
826 if self.predictionSupported:
827 # Predicted dataset
828 dataId = {"instrument": "DummyCamComp", "visit": 424}
829 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
831 if butler._datastore.isEphemeral:
832 # Never disassembled
833 self.assertIsInstance(uri, ResourcePath)
834 self.assertFalse(components)
835 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
836 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
837 else:
838 self.assertIsNone(uri)
839 self.assertEqual(set(components), set(storageClass.components))
840 for compuri in components.values():
841 self.assertIsInstance(compuri, ResourcePath)
842 self.assertIn("424", str(compuri), f"Checking visit is in URI {compuri}")
843 self.assertEqual(compuri.fragment, "predicted", f"Checking for fragment in {compuri}")
845 def testStorageClassOverrideGet(self) -> None:
846 """Test storage class conversion on get with override."""
847 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
848 datasetTypeName = "anything"
849 run = self.default_run
851 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
853 # Create and store a dataset.
854 metric = makeExampleMetrics()
855 dataId = {"instrument": "DummyCamComp", "visit": 423}
857 ref = butler.put(metric, datasetType, dataId)
859 # Return native type.
860 retrieved = butler.get(ref)
861 self.assertEqual(retrieved, metric)
863 # Specify an override.
864 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
865 model = butler.get(ref, storageClass=new_sc)
866 self.assertNotEqual(type(model), type(retrieved))
867 self.assertIs(type(model), new_sc.pytype)
868 self.assertEqual(retrieved, model)
870 # Defer but override later.
871 deferred = butler.getDeferred(ref)
872 model = deferred.get(storageClass=new_sc)
873 self.assertIs(type(model), new_sc.pytype)
874 self.assertEqual(retrieved, model)
876 # Defer but override up front.
877 deferred = butler.getDeferred(ref, storageClass=new_sc)
878 model = deferred.get()
879 self.assertIs(type(model), new_sc.pytype)
880 self.assertEqual(retrieved, model)
882 # Retrieve a component. Should be a tuple.
883 data = butler.get("anything.data", dataId, storageClass="StructuredDataDataTestTuple")
884 self.assertIs(type(data), tuple)
885 self.assertEqual(data, tuple(retrieved.data))
887 # Parameter on the write storage class should work regardless
888 # of read storage class.
889 data = butler.get(
890 "anything.data",
891 dataId,
892 storageClass="StructuredDataDataTestTuple",
893 parameters={"slice": slice(2, 4)},
894 )
895 self.assertEqual(len(data), 2)
897 # Try a parameter that is known to the read storage class but not
898 # the write storage class.
899 with self.assertRaises(KeyError):
900 butler.get(
901 "anything.data",
902 dataId,
903 storageClass="StructuredDataDataTestTuple",
904 parameters={"xslice": slice(2, 4)},
905 )
907 def testPytypePutCoercion(self) -> None:
908 """Test python type coercion on Butler.get and put."""
909 # Store some data with the normal example storage class.
910 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
911 datasetTypeName = "test_metric"
912 butler, _ = self.create_butler(self.default_run, storageClass, datasetTypeName)
914 dataId = {"instrument": "DummyCamComp", "visit": 423}
916 # Put a dict and this should coerce to a MetricsExample
917 test_dict = {"summary": {"a": 1}, "output": {"b": 2}}
918 metric_ref = butler.put(test_dict, datasetTypeName, dataId=dataId, visit=424)
919 test_metric = butler.get(metric_ref)
920 self.assertEqual(get_full_type_name(test_metric), "lsst.daf.butler.tests.MetricsExample")
921 self.assertEqual(test_metric.summary, test_dict["summary"])
922 self.assertEqual(test_metric.output, test_dict["output"])
924 # Check that the put still works if a DatasetType is given with
925 # a definition matching this python type.
926 registry_type = butler.get_dataset_type(datasetTypeName)
927 this_type = DatasetType(datasetTypeName, registry_type.dimensions, "StructuredDataDictJson")
928 metric2_ref = butler.put(test_dict, this_type, dataId=dataId, visit=425)
929 self.assertEqual(metric2_ref.datasetType, registry_type)
931 # The get will return the type expected by registry.
932 test_metric2 = butler.get(metric2_ref)
933 self.assertEqual(get_full_type_name(test_metric2), "lsst.daf.butler.tests.MetricsExample")
935 # Make a new DatasetRef with the compatible but different DatasetType.
936 # This should now return a dict.
937 new_ref = DatasetRef(this_type, metric2_ref.dataId, id=metric2_ref.id, run=metric2_ref.run)
938 test_dict2 = butler.get(new_ref)
939 self.assertEqual(get_full_type_name(test_dict2), "dict")
941 # Get it again with the wrong dataset type definition using get()
942 # rather than get(). This should be consistent with get()
943 # behavior and return the type of the DatasetType.
944 test_dict3 = butler.get(this_type, dataId=dataId, visit=425)
945 self.assertEqual(get_full_type_name(test_dict3), "dict")
947 def test_ingest_zip(self) -> None:
948 """Create butler, export data, delete data, import from Zip."""
949 butler, dataset_type = self.create_butler(
950 run=self.default_run, storageClass="StructuredData", datasetTypeName="metrics"
951 )
953 metric = makeExampleMetrics()
954 refs = []
955 for visit in (423, 424, 425):
956 ref = butler.put(metric, dataset_type, instrument="DummyCamComp", visit=visit)
957 refs.append(ref)
959 # Retrieve a Zip file.
960 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
961 zip = butler.retrieve_artifacts_zip(refs, destination=tmpdir)
963 # Ingest will fail.
964 with self.assertRaises(ConflictingDefinitionError):
965 butler.ingest_zip(zip)
967 # Clear out the collection.
968 butler.removeRuns([self.default_run])
969 self.assertFalse(butler.exists(refs[0]))
971 butler.ingest_zip(zip, transfer="copy")
973 # Create an entirely new local file butler in this temp directory.
974 new_butler_cfg = Butler.makeRepo(tmpdir)
975 new_butler = Butler.from_config(new_butler_cfg, writeable=True)
977 # This will fail since dimensions records are missing.
978 with self.assertRaises(ConflictingDefinitionError):
979 new_butler.ingest_zip(zip, transfer="copy")
981 # Dry run should work.
982 new_butler.ingest_zip(zip, transfer="copy", dry_run=True)
984 new_butler.ingest_zip(zip, transfer="copy", transfer_dimensions=True)
985 self.assertTrue(butler.exists(refs[0]))
987 # Check that the refs can be read again.
988 _ = [butler.get(ref) for ref in refs]
990 uri = butler.getURI(refs[2])
991 self.assertTrue(uri.exists())
993 # Delete one dataset. The Zip file should still exist and allow
994 # remaining refs to be read.
995 butler.pruneDatasets([refs[0]], purge=True, unstore=True)
996 self.assertTrue(uri.exists())
998 metric2 = butler.get(refs[1])
999 self.assertEqual(metric2, metric, msg=f"{metric2} != {metric}")
1001 butler.removeRuns([self.default_run])
1002 self.assertFalse(uri.exists())
1003 self.assertFalse(butler.exists(refs[-1]))
1005 with self.assertRaises(ValueError):
1006 butler.retrieve_artifacts_zip([], destination=".")
1008 def testIngest(self) -> None:
1009 butler = self.create_empty_butler(run=self.default_run)
1011 # Create and register a DatasetType
1012 dimensions = butler.dimensions.conform(["instrument", "visit", "detector"])
1014 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDictYaml")
1015 datasetTypeName = "metric"
1017 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1019 # Add needed Dimensions
1020 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1021 butler.registry.insertDimensionData(
1022 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1023 )
1024 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1025 for detector in (1, 2):
1026 butler.registry.insertDimensionData(
1027 "detector", {"instrument": "DummyCamComp", "id": detector, "full_name": f"detector{detector}"}
1028 )
1030 butler.registry.insertDimensionData(
1031 "visit",
1032 {
1033 "instrument": "DummyCamComp",
1034 "id": 423,
1035 "name": "fourtwentythree",
1036 "physical_filter": "d-r",
1037 "day_obs": 20250101,
1038 },
1039 {
1040 "instrument": "DummyCamComp",
1041 "id": 424,
1042 "name": "fourtwentyfour",
1043 "physical_filter": "d-r",
1044 "day_obs": 20250101,
1045 },
1046 )
1048 formatter = doImportType("lsst.daf.butler.formatters.yaml.YamlFormatter")
1049 dataRoot = os.path.join(TESTDIR, "data", "basic")
1050 datasets = []
1051 # Test one DatasetRef with a run that exists, and the other with a run
1052 # that doesn't exist, to verify that run collections are created when
1053 # required.
1054 runs = {1: self.default_run, 2: "a/new/run"}
1055 for detector in (1, 2):
1056 detector_name = f"detector_{detector}"
1057 metricFile = os.path.join(dataRoot, f"{detector_name}.yaml")
1058 dataId = butler.registry.expandDataId(
1059 {"instrument": "DummyCamComp", "visit": 423, "detector": detector}
1060 )
1061 # Create a DatasetRef for ingest
1062 refIn = DatasetRef(datasetType, dataId, run=runs[detector])
1064 datasets.append(FileDataset(path=metricFile, refs=[refIn], formatter=formatter))
1066 butler.ingest(*datasets, transfer="copy")
1068 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 423}
1069 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 423}
1071 metrics1 = butler.get(datasetTypeName, dataId1)
1072 metrics2 = butler.get(datasetTypeName, dataId2, collections="a/new/run")
1073 self.assertNotEqual(metrics1, metrics2)
1075 # Compare URIs
1076 uri1 = butler.getURI(datasetTypeName, dataId1)
1077 uri2 = butler.getURI(datasetTypeName, dataId2, collections="a/new/run")
1078 self.assertFalse(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1080 # Re-ingesting the same datasets raises an error with
1081 # skip_existing=False.
1082 with self.assertRaises(ConflictingDefinitionError):
1083 butler.ingest(*datasets, transfer="copy")
1084 # skip_existing=True makes it a no-op to re-ingest the same datasets.
1085 butler.ingest(*datasets, transfer="copy", skip_existing=True)
1087 # Now do a multi-dataset but single file ingest
1088 metricFile = os.path.join(dataRoot, "detectors.yaml")
1089 refs = []
1090 for detector in (1, 2):
1091 detector_name = f"detector_{detector}"
1092 dataId = butler.registry.expandDataId(
1093 {"instrument": "DummyCamComp", "visit": 424, "detector": detector}
1094 )
1095 # Create a DatasetRef for ingest
1096 refs.append(DatasetRef(datasetType, dataId, run=self.default_run))
1098 # Test "move" transfer to ensure that the files themselves
1099 # have disappeared following ingest.
1100 with ResourcePath.temporary_uri(suffix=".yaml") as tempFile:
1101 tempFile.transfer_from(ResourcePath(metricFile), transfer="copy")
1103 datasets = []
1104 datasets.append(FileDataset(path=tempFile, refs=refs, formatter=MultiDetectorFormatter))
1106 # For first ingest use copy.
1107 butler.ingest(*datasets, transfer="copy", record_validation_info=False)
1109 # Now try to ingest again in "execution butler" mode where
1110 # the registry entries exist but the datastore does not have
1111 # the files. We also need to strip the dimension records to ensure
1112 # that they will be re-added by the ingest.
1113 ref = datasets[0].refs[0]
1114 datasets[0].refs = [
1115 cast(
1116 DatasetRef,
1117 butler.find_dataset(ref.datasetType, data_id=ref.dataId, collections=ref.run),
1118 )
1119 for ref in datasets[0].refs
1120 ]
1121 all_refs = []
1122 for dataset in datasets:
1123 refs = []
1124 for ref in dataset.refs:
1125 # Create a dict from the dataId to drop the records.
1126 new_data_id = dict(ref.dataId.required)
1127 new_ref = butler.find_dataset(ref.datasetType, new_data_id, collections=ref.run)
1128 assert new_ref is not None
1129 self.assertFalse(new_ref.dataId.hasRecords())
1130 refs.append(new_ref)
1131 dataset.refs = refs
1132 all_refs.extend(dataset.refs)
1133 butler.pruneDatasets(all_refs, disassociate=False, unstore=True, purge=False)
1135 # Use move mode to test that the file is deleted. Also
1136 # disable recording of file size.
1137 butler.ingest(*datasets, transfer="move", record_validation_info=False)
1139 # Check that every ref now has records.
1140 for dataset in datasets:
1141 for ref in dataset.refs:
1142 self.assertTrue(ref.dataId.hasRecords())
1144 # Ensure that the file has disappeared.
1145 self.assertFalse(tempFile.exists())
1147 # Check that the datastore recorded no file size.
1148 # Not all datastores can support this.
1149 try:
1150 infos = butler._datastore.getStoredItemsInfo(datasets[0].refs[0]) # type: ignore[attr-defined]
1151 self.assertEqual(infos[0].file_size, -1)
1152 except AttributeError:
1153 pass
1155 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 424}
1156 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 424}
1158 multi1 = butler.get(datasetTypeName, dataId1)
1159 multi2 = butler.get(datasetTypeName, dataId2)
1161 self.assertEqual(multi1, metrics1)
1162 self.assertEqual(multi2, metrics2)
1164 # Compare URIs
1165 uri1 = butler.getURI(datasetTypeName, dataId1)
1166 uri2 = butler.getURI(datasetTypeName, dataId2)
1167 self.assertTrue(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1169 # Test that removing one does not break the second
1170 # This line will issue a warning log message for a ChainedDatastore
1171 # that uses an InMemoryDatastore since in-memory can not ingest
1172 # files.
1173 butler.pruneDatasets([datasets[0].refs[0]], unstore=True, disassociate=False)
1174 self.assertFalse(butler.exists(datasetTypeName, dataId1))
1175 self.assertTrue(butler.exists(datasetTypeName, dataId2))
1176 multi2b = butler.get(datasetTypeName, dataId2)
1177 self.assertEqual(multi2, multi2b)
1179 # Ensure we can ingest 0 datasets
1180 datasets = []
1181 butler.ingest(*datasets)
1183 def testPickle(self) -> None:
1184 """Test pickle support."""
1185 butler = self.create_empty_butler(run=self.default_run)
1186 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
1187 butlerOut = pickle.loads(pickle.dumps(butler))
1188 self.assertIsInstance(butlerOut, Butler)
1189 self.assertEqual(butlerOut._config, butler._config)
1190 self.assertEqual(list(butlerOut.collections.defaults), list(butler.collections.defaults))
1191 self.assertEqual(butlerOut.run, butler.run)
1193 def testGetDatasetTypes(self) -> None:
1194 butler = self.create_empty_butler(run=self.default_run)
1195 dimensions = butler.dimensions.conform(["instrument", "visit", "physical_filter"])
1196 dimensionEntries: list[tuple[str, list[Mapping[str, Any]]]] = [
1197 (
1198 "instrument",
1199 [
1200 {"instrument": "DummyCam"},
1201 {"instrument": "DummyHSC"},
1202 {"instrument": "DummyCamComp"},
1203 ],
1204 ),
1205 ("physical_filter", [{"instrument": "DummyCam", "name": "d-r", "band": "R"}]),
1206 ("day_obs", [{"instrument": "DummyCam", "id": 20250101}]),
1207 (
1208 "visit",
1209 [
1210 {
1211 "instrument": "DummyCam",
1212 "id": 42,
1213 "name": "fortytwo",
1214 "physical_filter": "d-r",
1215 "day_obs": 20250101,
1216 }
1217 ],
1218 ),
1219 ]
1220 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1221 # Add needed Dimensions
1222 for element, data in dimensionEntries:
1223 butler.registry.insertDimensionData(element, *data)
1225 # When a DatasetType is added to the registry entries are not created
1226 # for components but querying them can return the components.
1227 datasetTypeNames = {"metric", "metric2", "metric4", "metric33", "pvi", "paramtest"}
1228 components = set()
1229 for datasetTypeName in datasetTypeNames:
1230 # Create and register a DatasetType
1231 self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1233 for componentName in storageClass.components:
1234 components.add(DatasetType.nameWithComponent(datasetTypeName, componentName))
1236 fromRegistry: set[DatasetType] = set()
1237 for parent_dataset_type in butler.registry.queryDatasetTypes():
1238 fromRegistry.add(parent_dataset_type)
1239 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes())
1240 self.assertEqual({d.name for d in fromRegistry}, datasetTypeNames | components)
1242 # Query with wildcard.
1243 dataset_types = butler.registry.queryDatasetTypes("metric*")
1244 self.assertEqual(len(dataset_types), 4, f"Got: {dataset_types}")
1245 # but not regex.
1246 with self.assertRaises(DatasetTypeExpressionError):
1247 butler.registry.queryDatasetTypes(["pvi", re.compile("metric.*")])
1249 # Now that we have some dataset types registered, validate them
1250 butler.validateConfiguration(
1251 ignore=[
1252 "test_metric_comp",
1253 "metric3",
1254 "metric5",
1255 "calexp",
1256 "DummySC",
1257 "datasetType.component",
1258 "random_data",
1259 "random_data_2",
1260 ]
1261 )
1263 # Add a new datasetType that will fail template validation
1264 self.addDatasetType("test_metric_comp", dimensions, storageClass, butler.registry)
1265 if self.validationCanFail:
1266 with self.assertRaises(ValidationError):
1267 butler.validateConfiguration()
1269 # Rerun validation but with a subset of dataset type names
1270 butler.validateConfiguration(datasetTypeNames=["metric4"])
1272 # Rerun validation but ignore the bad datasetType
1273 butler.validateConfiguration(
1274 ignore=[
1275 "test_metric_comp",
1276 "metric3",
1277 "metric5",
1278 "calexp",
1279 "DummySC",
1280 "datasetType.component",
1281 "random_data",
1282 "random_data_2",
1283 ]
1284 )
1286 def testTransaction(self) -> None:
1287 butler = self.create_empty_butler(run=self.default_run)
1288 datasetTypeName = "test_metric"
1289 dimensions = butler.dimensions.conform(["instrument", "visit"])
1290 dimensionEntries: tuple[tuple[str, Mapping[str, Any]], ...] = (
1291 ("instrument", {"instrument": "DummyCam"}),
1292 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}),
1293 ("day_obs", {"instrument": "DummyCam", "id": 20250101}),
1294 (
1295 "visit",
1296 {
1297 "instrument": "DummyCam",
1298 "id": 42,
1299 "name": "fortytwo",
1300 "physical_filter": "d-r",
1301 "day_obs": 20250101,
1302 },
1303 ),
1304 )
1305 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1306 metric = makeExampleMetrics()
1307 dataId = {"instrument": "DummyCam", "visit": 42}
1308 # Create and register a DatasetType
1309 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1310 with self.assertRaises(TransactionTestError):
1311 with butler.transaction():
1312 # Add needed Dimensions
1313 for args in dimensionEntries:
1314 butler.registry.insertDimensionData(*args)
1315 # Store a dataset
1316 ref = butler.put(metric, datasetTypeName, dataId)
1317 self.assertIsInstance(ref, DatasetRef)
1318 # Test get of a ref.
1319 metricOut = butler.get(ref)
1320 self.assertEqual(metric, metricOut)
1321 # Test get
1322 metricOut = butler.get(datasetTypeName, dataId)
1323 self.assertEqual(metric, metricOut)
1324 # Check we can get components
1325 self.assertGetComponents(butler, ref, ("summary", "data", "output"), metric)
1326 raise TransactionTestError("This should roll back the entire transaction")
1327 with self.assertRaises(DataIdValueError, msg=f"Check can't expand DataId {dataId}"):
1328 butler.registry.expandDataId(dataId)
1329 # Should raise LookupError for missing data ID value
1330 with self.assertRaises(LookupError, msg=f"Check can't get by {datasetTypeName} and {dataId}"):
1331 butler.get(datasetTypeName, dataId)
1332 # Also check explicitly if Dataset entry is missing
1333 self.assertIsNone(butler.find_dataset(datasetType, dataId, collections=butler.collections.defaults))
1334 # Direct retrieval should not find the file in the Datastore
1335 with self.assertRaises(FileNotFoundError, msg=f"Check {ref} can't be retrieved directly"):
1336 butler.get(ref)
1338 def testMakeRepo(self) -> None:
1339 """Test that we can write butler configuration to a new repository via
1340 the Butler.makeRepo interface and then instantiate a butler from the
1341 repo root.
1342 """
1343 # Do not run the test if we know this datastore configuration does
1344 # not support a file system root
1345 if self.fullConfigKey is None:
1346 return
1348 # create two separate directories
1349 root1 = tempfile.mkdtemp(dir=self.root)
1350 root2 = tempfile.mkdtemp(dir=self.root)
1352 butlerConfig = Butler.makeRepo(root1, config=Config(self.configFile))
1353 limited = Config(self.configFile)
1354 butler1 = Butler.from_config(butlerConfig)
1355 assert isinstance(butler1, DirectButler), "Expect DirectButler in configuration"
1356 butlerConfig = Butler.makeRepo(root2, standalone=True, config=Config(self.configFile))
1357 full = Config(self.tmpConfigFile)
1358 butler2 = Butler.from_config(butlerConfig)
1359 assert isinstance(butler2, DirectButler), "Expect DirectButler in configuration"
1360 # Butlers should have the same configuration regardless of whether
1361 # defaults were expanded.
1362 self.assertEqual(butler1._config, butler2._config)
1363 # Config files loaded directly should not be the same.
1364 self.assertNotEqual(limited, full)
1365 # Make sure "limited" doesn't have a few keys we know it should be
1366 # inheriting from defaults.
1367 self.assertIn(self.fullConfigKey, full)
1368 self.assertNotIn(self.fullConfigKey, limited)
1370 # Collections don't appear until something is put in them
1371 collections1 = set(butler1.registry.queryCollections())
1372 self.assertEqual(collections1, set())
1373 self.assertEqual(set(butler2.registry.queryCollections()), collections1)
1375 # Check that a config with no associated file name will not
1376 # work properly with relocatable Butler repo
1377 butlerConfig.configFile = None
1378 with self.assertRaises(ValueError):
1379 Butler.from_config(butlerConfig)
1381 with self.assertRaises(FileExistsError):
1382 Butler.makeRepo(self.root, standalone=True, config=Config(self.configFile), overwrite=False)
1384 def testStringification(self) -> None:
1385 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
1386 butlerStr = str(butler)
1388 if self.datastoreStr is not None:
1389 for testStr in self.datastoreStr:
1390 self.assertIn(testStr, butlerStr)
1391 if self.registryStr is not None:
1392 self.assertIn(self.registryStr, butlerStr)
1394 datastoreName = butler._datastore.name
1395 if self.datastoreName is not None:
1396 for testStr in self.datastoreName:
1397 self.assertIn(testStr, datastoreName)
1399 def testButlerRewriteDataId(self) -> None:
1400 """Test that dataIds can be rewritten based on dimension records."""
1401 butler = self.create_empty_butler(run=self.default_run)
1403 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
1404 datasetTypeName = "random_data"
1406 # Create dimension records.
1407 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1408 butler.registry.insertDimensionData(
1409 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1410 )
1411 butler.registry.insertDimensionData(
1412 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
1413 )
1415 dimensions = butler.dimensions.conform(["instrument", "exposure"])
1416 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
1417 butler.registry.registerDatasetType(datasetType)
1419 n_exposures = 5
1420 dayobs = 20210530
1422 # Create records for multiple day_obs but same seq_num to test that
1423 # we are constraining gets properly when day_obs/seq_num is used
1424 # for an exposure. Second day is year in future but is not used.
1425 for day_obs in (dayobs, dayobs + 1_00_00):
1426 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": day_obs})
1428 for i in range(n_exposures):
1429 group_name = f"group_{day_obs}_{i}"
1430 butler.registry.insertDimensionData(
1431 "group", {"instrument": "DummyCamComp", "name": group_name}
1432 )
1433 butler.registry.insertDimensionData(
1434 "exposure",
1435 {
1436 "instrument": "DummyCamComp",
1437 "id": day_obs + i,
1438 "obs_id": f"exp_{day_obs}_{i}",
1439 "seq_num": i,
1440 "day_obs": day_obs,
1441 "physical_filter": "d-r",
1442 "group": group_name,
1443 },
1444 )
1446 # Write some data.
1447 for i in range(n_exposures):
1448 metric = {"something": i, "other": "metric", "list": [2 * x for x in range(i)]}
1450 # Use the seq_num for the put to test rewriting.
1451 dataId = {"seq_num": i, "day_obs": dayobs, "instrument": "DummyCamComp", "physical_filter": "d-r"}
1452 ref = butler.put(metric, datasetTypeName, dataId=dataId)
1454 # Check that the exposure is correct in the dataId
1455 self.assertEqual(ref.dataId["exposure"], dayobs + i)
1457 # and check that we can get the dataset back with the same dataId
1458 new_metric = butler.get(datasetTypeName, dataId=dataId)
1459 self.assertEqual(new_metric, metric)
1461 # Check that we can find the datasets using the day_obs or the
1462 # exposure.day_obs.
1463 datasets_1 = list(
1464 butler.registry.queryDatasets(
1465 datasetType,
1466 collections=self.default_run,
1467 where="day_obs = :dayObs AND instrument = :instr",
1468 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1469 )
1470 )
1471 datasets_2 = list(
1472 butler.registry.queryDatasets(
1473 datasetType,
1474 collections=self.default_run,
1475 where="exposure.day_obs = :dayObs AND instrument = :instr",
1476 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1477 )
1478 )
1479 self.assertEqual(datasets_1, datasets_2)
1481 def testGetDatasetCollectionCaching(self):
1482 # Prior to DM-41117, there was a bug where get_dataset would throw
1483 # MissingCollectionError if you tried to fetch a dataset that was added
1484 # after the collection cache was last updated.
1485 reader_butler, datasetType = self.create_butler(self.default_run, "int", "datasettypename")
1486 writer_butler = self.create_empty_butler(writeable=True, run="new_run")
1487 dataId = {"instrument": "DummyCamComp", "visit": 423}
1488 put_ref = writer_butler.put(123, datasetType, dataId)
1489 get_ref = reader_butler.get_dataset(put_ref.id)
1490 self.assertEqual(get_ref.id, put_ref.id)
1491 # Also works when looking up via a hexadecimal string instead of a UUID
1492 # instance.
1493 hex_ref = reader_butler.get_dataset(put_ref.id.hex)
1494 self.assertEqual(hex_ref.id, put_ref.id)
1496 def testCollectionChainRedefine(self):
1497 butler = self._setup_to_test_collection_chain()
1499 butler.collections.redefine_chain("chain", "a")
1500 self._check_chain(butler, ["a"])
1502 # Duplicates are removed from the list of children
1503 butler.collections.redefine_chain("chain", ["c", "b", "c"])
1504 self._check_chain(butler, ["c", "b"])
1506 # Empty list clears the chain
1507 butler.collections.redefine_chain("chain", [])
1508 self._check_chain(butler, [])
1510 self._test_common_chain_functionality(butler, butler.collections.redefine_chain)
1512 def testCollectionChainPrepend(self):
1513 butler = self._setup_to_test_collection_chain()
1515 # Duplicates are removed from the list of children
1516 butler.collections.prepend_chain("chain", ["c", "b", "c"])
1517 self._check_chain(butler, ["c", "b"])
1519 # Prepend goes on the front of existing chain
1520 butler.collections.prepend_chain("chain", ["a"])
1521 self._check_chain(butler, ["a", "c", "b"])
1523 # Empty prepend does nothing
1524 butler.collections.prepend_chain("chain", [])
1525 self._check_chain(butler, ["a", "c", "b"])
1527 # Prepending children that already exist in the chain removes them from
1528 # their current position.
1529 butler.collections.prepend_chain("chain", ["d", "b", "c"])
1530 self._check_chain(butler, ["d", "b", "c", "a"])
1532 self._test_common_chain_functionality(butler, butler.collections.prepend_chain)
1534 def testCollectionChainExtend(self):
1535 butler = self._setup_to_test_collection_chain()
1537 # Duplicates are removed from the list of children
1538 butler.collections.extend_chain("chain", ["c", "b", "c"])
1539 self._check_chain(butler, ["c", "b"])
1541 # Extend goes on the end of existing chain
1542 butler.collections.extend_chain("chain", ["a"])
1543 self._check_chain(butler, ["c", "b", "a"])
1545 # Empty extend does nothing
1546 butler.collections.extend_chain("chain", [])
1547 self._check_chain(butler, ["c", "b", "a"])
1549 # Extending children that already exist in the chain removes them from
1550 # their current position.
1551 butler.collections.extend_chain("chain", ["d", "b", "c"])
1552 self._check_chain(butler, ["a", "d", "b", "c"])
1554 self._test_common_chain_functionality(butler, butler.collections.extend_chain)
1556 def testCollectionChainRemove(self) -> None:
1557 butler = self._setup_to_test_collection_chain()
1559 butler.collections.redefine_chain("chain", ["a", "b", "c", "d"])
1561 butler.collections.remove_from_chain("chain", "c")
1562 self._check_chain(butler, ["a", "b", "d"])
1564 # Duplicates are allowed in the list of children
1565 butler.collections.remove_from_chain("chain", ["b", "b", "a"])
1566 self._check_chain(butler, ["d"])
1568 # Empty remove does nothing
1569 butler.collections.remove_from_chain("chain", [])
1570 self._check_chain(butler, ["d"])
1572 # Removing children that aren't in the chain does nothing
1573 butler.collections.remove_from_chain("chain", ["a", "chain"])
1574 self._check_chain(butler, ["d"])
1576 self._test_common_chain_functionality(
1577 butler, butler.collections.remove_from_chain, skip_cycle_check=True
1578 )
1580 def _setup_to_test_collection_chain(self) -> Butler:
1581 butler = self.create_empty_butler(writeable=True)
1583 butler.collections.register("chain", CollectionType.CHAINED)
1585 runs = ["a", "b", "c", "d"]
1586 for run in runs:
1587 butler.collections.register(run)
1589 butler.collections.register("staticchain", CollectionType.CHAINED)
1590 butler.collections.redefine_chain("staticchain", ["a", "b"])
1592 return butler
1594 def _check_chain(self, butler: Butler, expected: list[str]) -> None:
1595 children = butler.collections.get_info("chain").children
1596 self.assertEqual(expected, list(children))
1598 def _test_common_chain_functionality(
1599 self, butler, func: Callable[[str, str | list[str]], Any], *, skip_cycle_check=False
1600 ) -> None:
1601 # Missing parent collection
1602 with self.assertRaises(MissingCollectionError):
1603 func("doesnotexist", [])
1604 # Missing child collection
1605 with self.assertRaises(MissingCollectionError):
1606 func("chain", ["doesnotexist"])
1607 # Forbid operations on non-chained collections
1608 with self.assertRaises(CollectionTypeError):
1609 func("d", ["a"])
1611 # Prevent collection cycles
1612 if not skip_cycle_check:
1613 butler.collections.register("chain2", CollectionType.CHAINED)
1614 func("chain2", "chain")
1615 with self.assertRaises(CollectionCycleError):
1616 func("chain", "chain2")
1618 # Make sure none of the earlier operations interfered with unrelated
1619 # chains.
1620 self.assertEqual(["a", "b"], list(butler.collections.get_info("staticchain").children))
1622 with butler._caching_context():
1623 with self.assertRaisesRegex(RuntimeError, "Chained collection modification not permitted"):
1624 func("chain", "a")
1627class FileDatastoreButlerTests(ButlerTests):
1628 """Common tests and specialization of ButlerTests for butlers backed
1629 by datastores that inherit from FileDatastore.
1630 """
1632 trustModeSupported = True
1634 def checkFileExists(self, root: str | ResourcePath, relpath: str | ResourcePath) -> bool:
1635 """Check if file exists at a given path (relative to root).
1637 Test testPutTemplates verifies actual physical existance of the files
1638 in the requested location.
1639 """
1640 uri = ResourcePath(root, forceDirectory=True)
1641 return uri.join(relpath).exists()
1643 def testPutTemplates(self) -> None:
1644 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1645 butler = self.create_empty_butler(run=self.default_run)
1647 # Add needed Dimensions
1648 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1649 butler.registry.insertDimensionData(
1650 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1651 )
1652 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1653 butler.registry.insertDimensionData(
1654 "visit",
1655 {
1656 "instrument": "DummyCamComp",
1657 "id": 423,
1658 "name": "v423",
1659 "physical_filter": "d-r",
1660 "day_obs": 20250101,
1661 },
1662 )
1663 butler.registry.insertDimensionData(
1664 "visit",
1665 {
1666 "instrument": "DummyCamComp",
1667 "id": 425,
1668 "name": "v425",
1669 "physical_filter": "d-r",
1670 "day_obs": 20250101,
1671 },
1672 )
1674 # Create and store a dataset
1675 metric = makeExampleMetrics()
1677 # Create two almost-identical DatasetTypes (both will use default
1678 # template)
1679 dimensions = butler.dimensions.conform(["instrument", "visit"])
1680 butler.registry.registerDatasetType(DatasetType("metric1", dimensions, storageClass))
1681 butler.registry.registerDatasetType(DatasetType("metric2", dimensions, storageClass))
1682 butler.registry.registerDatasetType(DatasetType("metric3", dimensions, storageClass))
1684 dataId1 = {"instrument": "DummyCamComp", "visit": 423}
1685 dataId2 = {"instrument": "DummyCamComp", "visit": 423, "physical_filter": "d-r"}
1687 # Put with exactly the data ID keys needed
1688 ref = butler.put(metric, "metric1", dataId1)
1689 uri = butler.getURI(ref)
1690 self.assertTrue(uri.exists())
1691 self.assertTrue(
1692 uri.unquoted_path.endswith(f"{self.default_run}/metric1/??#?/d-r/DummyCamComp_423.pickle")
1693 )
1695 # Check the template based on dimensions
1696 if hasattr(butler._datastore, "templates"):
1697 butler._datastore.templates.validateTemplates([ref])
1699 # Put with extra data ID keys (physical_filter is an optional
1700 # dependency); should not change template (at least the way we're
1701 # defining them to behave now; the important thing is that they
1702 # must be consistent).
1703 ref = butler.put(metric, "metric2", dataId2)
1704 uri = butler.getURI(ref)
1705 self.assertTrue(uri.exists())
1706 self.assertTrue(
1707 uri.unquoted_path.endswith(f"{self.default_run}/metric2/d-r/DummyCamComp_v423.pickle")
1708 )
1710 # Check the template based on dimensions
1711 if hasattr(butler._datastore, "templates"):
1712 butler._datastore.templates.validateTemplates([ref])
1714 # Use a template that has a typo in dimension record metadata.
1715 # Easier to test with a butler that has a ref with records attached.
1716 template = FileTemplate("a/{visit.name}/{id}_{visit.namex:?}.fits")
1717 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1718 path = template.format(ref)
1719 self.assertEqual(path, f"a/v423/{ref.id}_fits")
1721 template = FileTemplate("a/{visit.name}/{id}_{visit.namex}.fits")
1722 with self.assertRaises(KeyError):
1723 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1724 template.format(ref)
1726 # Now use a file template that will not result in unique filenames
1727 with self.assertRaises(FileTemplateValidationError):
1728 butler.put(metric, "metric3", dataId1)
1730 def testImportExport(self) -> None:
1731 # Run put/get tests just to create and populate a repo.
1732 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1733 self.runImportExportTest(storageClass)
1735 @unittest.expectedFailure
1736 def testImportExportVirtualComposite(self) -> None:
1737 # Run put/get tests just to create and populate a repo.
1738 storageClass = self.storageClassFactory.getStorageClass("StructuredComposite")
1739 self.runImportExportTest(storageClass)
1741 def runImportExportTest(self, storageClass: StorageClass) -> None:
1742 """Test exporting and importing.
1744 This test does an export to a temp directory and an import back
1745 into a new temp directory repo. It does not assume a posix datastore.
1746 """
1747 exportButler = self.runPutGetTest(storageClass, "test_metric")
1749 # Test that we must have a file extension.
1750 with self.assertRaises(ValueError):
1751 with exportButler.export(filename="dump", directory=".") as export:
1752 pass
1754 # Test that unknown format is not allowed.
1755 with self.assertRaises(ValueError):
1756 with exportButler.export(filename="dump.fits", directory=".") as export:
1757 pass
1759 # Test that the repo actually has at least one dataset.
1760 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
1761 self.assertGreater(len(datasets), 0)
1762 # Add a DimensionRecord that's unused by those datasets.
1763 skymapRecord = {"name": "example_skymap", "hash": (50).to_bytes(8, byteorder="little")}
1764 exportButler.registry.insertDimensionData("skymap", skymapRecord)
1765 # Export and then import datasets.
1766 with safeTestTempDir(TESTDIR) as exportDir:
1767 exportFile = os.path.join(exportDir, "exports.yaml")
1768 with exportButler.export(filename=exportFile, directory=exportDir, transfer="auto") as export:
1769 export.saveDatasets(datasets)
1770 # Export the same datasets again. This should quietly do
1771 # nothing because of internal deduplication, and it shouldn't
1772 # complain about being asked to export the "htm7" elements even
1773 # though there aren't any in these datasets or in the database.
1774 export.saveDatasets(datasets, elements=["htm7"])
1775 # Save one of the data IDs again; this should be harmless
1776 # because of internal deduplication.
1777 export.saveDataIds([datasets[0].dataId])
1778 # Save some dimension records directly.
1779 export.saveDimensionData("skymap", [skymapRecord])
1780 self.assertTrue(os.path.exists(exportFile))
1781 with safeTestTempDir(TESTDIR) as importDir:
1782 # We always want this to be a local posix butler
1783 Butler.makeRepo(importDir, config=Config(os.path.join(TESTDIR, "config/basic/butler.yaml")))
1784 # Calling script.butlerImport tests the implementation of the
1785 # butler command line interface "import" subcommand. Functions
1786 # in the script folder are generally considered protected and
1787 # should not be used as public api.
1788 with open(exportFile) as f:
1789 script.butlerImport(
1790 importDir,
1791 export_file=f,
1792 directory=exportDir,
1793 transfer="auto",
1794 skip_dimensions=None,
1795 )
1796 importButler = Butler.from_config(importDir, run=self.default_run)
1797 for ref in datasets:
1798 with self.subTest(ref=repr(ref)):
1799 # Test for existence by passing in the DatasetType and
1800 # data ID separately, to avoid lookup by dataset_id.
1801 self.assertTrue(importButler.exists(ref.datasetType, ref.dataId))
1802 self.assertEqual(
1803 list(importButler.registry.queryDimensionRecords("skymap")),
1804 [importButler.dimensions["skymap"].RecordClass(**skymapRecord)],
1805 )
1807 def testRemoveRuns(self) -> None:
1808 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1809 butler = self.create_empty_butler(writeable=True)
1810 # Load registry data with dimensions to hang datasets off of.
1811 butler.import_(filename=ResourcePath("resource://lsst.daf.butler/tests/registry_data/base.yaml"))
1812 # Add some RUN-type collection.
1813 run1 = "run1"
1814 butler.collections.register(run1)
1815 run2 = "run2"
1816 butler.collections.register(run2)
1817 # put a dataset in each
1818 metric = makeExampleMetrics()
1819 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1820 datasetType = self.addDatasetType(
1821 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1822 )
1823 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1824 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1825 uri1 = butler.getURI(ref1)
1826 uri2 = butler.getURI(ref2)
1828 # Put one of the runs in a chain.
1829 butler.collections.register("Chain", CollectionType.CHAINED)
1830 butler.collections.extend_chain("Chain", run1)
1832 with self.assertRaises(OrphanedRecordError):
1833 butler.registry.removeDatasetType(datasetType.name)
1835 # Remove a non-run.
1836 with self.assertRaises(TypeError):
1837 butler.removeRuns(["Chain"])
1839 # Remove without unlinking from chain should fail.
1840 with self.assertRaises(IntegrityError):
1841 butler.removeRuns([run1])
1843 # Remove from both runs. No longer use unstore parameter since it
1844 # always purges.
1845 butler.removeRuns([run1, run2], unlink_from_chains=True)
1847 # Should be nothing in registry for either one, and datastore should
1848 # not think either exists.
1849 with self.assertRaises(MissingCollectionError):
1850 butler.collections.get_info(run1)
1851 with self.assertRaises(MissingCollectionError):
1852 butler.collections.get_info(run1)
1853 self.assertFalse(butler.stored(ref1))
1854 self.assertFalse(butler.stored(ref2))
1855 # We always unstore so both URIs should be gone.
1856 self.assertFalse(uri1.exists())
1857 self.assertFalse(uri2.exists())
1859 # Now that the collections have been pruned we can remove the
1860 # dataset type
1861 butler.registry.removeDatasetType(datasetType.name)
1863 with self.assertLogs("lsst.daf.butler.registry", "INFO") as cm:
1864 butler.registry.removeDatasetType(("test*", "test*"))
1865 self.assertIn("not defined", "\n".join(cm.output))
1867 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
1868 """Simulate an external actor removing a file outside of Butler's
1869 knowledge.
1871 Subclasses may override to handle more complicated datastore
1872 configurations.
1873 """
1874 uri = butler.getURI(ref)
1875 uri.remove()
1876 datastore = cast(FileDatastore, butler._datastore)
1877 datastore.cacheManager.remove_from_cache(ref)
1879 def testPruneDatasets(self) -> None:
1880 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1881 butler = self.create_empty_butler(writeable=True)
1882 # Load registry data with dimensions to hang datasets off of.
1883 butler.import_(filename=_get_test_data_path("base.yaml"))
1884 # Add some RUN-type collections.
1885 run1 = "run1"
1886 butler.collections.register(run1)
1887 run2 = "run2"
1888 butler.collections.register(run2)
1889 # put some datasets. ref1 and ref2 have the same data ID, and are in
1890 # different runs. ref3 has a different data ID.
1891 metric = makeExampleMetrics()
1892 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1893 datasetType = self.addDatasetType(
1894 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1895 )
1896 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1897 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1898 ref3 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, run=run1)
1900 many_stored = butler.stored_many([ref1, ref2, ref3])
1901 for ref, stored in many_stored.items():
1902 self.assertTrue(stored, f"Ref {ref} should be stored")
1904 many_exists = butler._exists_many([ref1, ref2, ref3])
1905 for ref, exists in many_exists.items():
1906 self.assertTrue(exists, f"Checking ref {ref} exists.")
1907 self.assertEqual(exists, DatasetExistence.VERIFIED, f"Ref {ref} should be stored")
1909 # Simple prune.
1910 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
1911 self.assertFalse(butler.exists(ref1.datasetType, ref1.dataId, collections=run1))
1913 many_stored = butler.stored_many([ref1, ref2, ref3])
1914 for ref, stored in many_stored.items():
1915 self.assertFalse(stored, f"Ref {ref} should not be stored")
1917 many_exists = butler._exists_many([ref1, ref2, ref3])
1918 for ref, exists in many_exists.items():
1919 self.assertEqual(exists, DatasetExistence.UNRECOGNIZED, f"Ref {ref} should not be stored")
1921 # Put data back.
1922 ref1_new = butler.put(metric, ref1)
1923 self.assertEqual(ref1_new, ref1) # Reuses original ID.
1924 ref2 = butler.put(metric, ref2)
1926 many_stored = butler.stored_many([ref1, ref2, ref3])
1927 self.assertTrue(many_stored[ref1])
1928 self.assertTrue(many_stored[ref2])
1929 self.assertFalse(many_stored[ref3])
1931 ref3 = butler.put(metric, ref3)
1933 many_exists = butler._exists_many([ref1, ref2, ref3])
1934 for ref, exists in many_exists.items():
1935 self.assertTrue(exists, f"Ref {ref} should not be stored")
1937 # Clear out the datasets from registry and start again.
1938 refs = [ref1, ref2, ref3]
1939 butler.pruneDatasets(refs, purge=True, unstore=True)
1940 for ref in refs:
1941 butler.put(metric, ref)
1943 # Confirm we can retrieve deferred.
1944 dref1 = butler.getDeferred(ref1) # known and exists
1945 metric1 = dref1.get()
1946 self.assertEqual(metric1, metric)
1948 # Test different forms of file availability.
1949 # Need to be in a state where:
1950 # - one ref just has registry record.
1951 # - one ref has a missing file but a datastore record.
1952 # - one ref has a missing datastore record but file is there.
1953 # - one ref does not exist anywhere.
1954 # Do not need to test a ref that has everything since that is tested
1955 # above.
1956 ref0 = DatasetRef(
1957 datasetType,
1958 DataCoordinate.standardize(
1959 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, universe=butler.dimensions
1960 ),
1961 run=run1,
1962 )
1964 # Delete from datastore and retain in Registry.
1965 butler.pruneDatasets([ref1], purge=False, unstore=True, disassociate=False)
1967 # File has been removed.
1968 self.remove_dataset_out_of_band(butler, ref2)
1970 # Datastore has lost track.
1971 butler._datastore.forget([ref3])
1973 # First test with a standard butler.
1974 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
1975 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
1976 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
1977 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
1978 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED)
1980 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=False)
1981 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
1982 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
1983 self.assertEqual(exists_many[ref2], DatasetExistence.KNOWN)
1984 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
1985 self.assertTrue(exists_many[ref2])
1987 # Check that per-ref query gives the same answer as many query.
1988 for ref, exists in exists_many.items():
1989 self.assertEqual(butler.exists(ref, full_check=False), exists)
1991 # Get deferred checks for existence before it allows it to be
1992 # retrieved.
1993 with self.assertRaises(LookupError):
1994 butler.getDeferred(ref3) # not known, file exists
1995 dref2 = butler.getDeferred(ref2) # known but file missing
1996 with self.assertRaises(FileNotFoundError):
1997 dref2.get()
1999 # Test again with a trusting butler.
2000 if self.trustModeSupported:
2001 butler._datastore.trustGetRequest = True
2002 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
2003 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2004 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2005 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2006 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ARTIFACT)
2008 # When trusting we can get a deferred dataset handle that is not
2009 # known but does exist.
2010 dref3 = butler.getDeferred(ref3)
2011 metric3 = dref3.get()
2012 self.assertEqual(metric3, metric)
2014 # Check that per-ref query gives the same answer as many query.
2015 for ref, exists in exists_many.items():
2016 self.assertEqual(butler.exists(ref, full_check=True), exists)
2018 # Create a ref that surprisingly has the UUID of an existing ref
2019 # but is not the same.
2020 ref_bad = DatasetRef(datasetType, dataId=ref3.dataId, run=ref3.run, id=ref2.id)
2021 with self.assertRaises(ValueError):
2022 butler.exists(ref_bad)
2024 # Create a ref that has a compatible storage class.
2025 ref_compat = ref2.overrideStorageClass("StructuredDataDict")
2026 exists = butler.exists(ref_compat)
2027 self.assertEqual(exists, exists_many[ref2])
2029 # Remove everything and start from scratch.
2030 butler._datastore.trustGetRequest = False
2031 butler.pruneDatasets(refs, purge=True, unstore=True)
2032 for ref in refs:
2033 butler.put(metric, ref)
2035 # These tests mess directly with the trash table and can leave the
2036 # datastore in an odd state. Do them at the end.
2037 # Check that in normal mode, deleting the record will lead to
2038 # trash not touching the file.
2039 uri1 = butler.getURI(ref1)
2040 butler._datastore.bridge.moveToTrash(
2041 [ref1], transaction=None
2042 ) # Update the dataset_location table
2043 butler._datastore.forget([ref1])
2044 butler._datastore.trash(ref1)
2045 butler._datastore.emptyTrash()
2046 self.assertTrue(uri1.exists())
2047 uri1.remove() # Clean it up.
2049 # Simulate execution butler setup by deleting the datastore
2050 # record but keeping the file around and trusting.
2051 butler._datastore.trustGetRequest = True
2052 uris = butler.get_many_uris([ref2, ref3])
2053 uri2 = uris[ref2].primaryURI
2054 uri3 = uris[ref3].primaryURI
2055 self.assertTrue(uri2.exists())
2056 self.assertTrue(uri3.exists())
2058 # Remove the datastore record.
2059 butler._datastore.bridge.moveToTrash(
2060 [ref2], transaction=None
2061 ) # Update the dataset_location table
2062 butler._datastore.forget([ref2])
2063 self.assertTrue(uri2.exists())
2064 butler._datastore.trash([ref2, ref3])
2065 # Immediate removal for ref2 file
2066 self.assertFalse(uri2.exists())
2067 # But ref3 has to wait for the empty.
2068 self.assertTrue(uri3.exists())
2069 butler._datastore.emptyTrash()
2070 self.assertFalse(uri3.exists())
2072 # Clear out the datasets from registry.
2073 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
2075 def test_butler_metrics(self):
2076 """Test that metrics are collected."""
2077 run = "test_run"
2078 metrics = ButlerMetrics()
2079 butler, datasetType = self.create_butler(
2080 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics
2081 )
2082 data = MetricsExampleModel(
2083 summary={"AM1": 5.2, "AM2": 30.6},
2084 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2085 data=[563, 234, 456.7, 752, 8, 9, 27],
2086 )
2088 data_ref = butler.put(data, datasetType, visit=424, instrument="DummyCamComp")
2089 butler.get(data_ref)
2090 butler.get(data_ref)
2091 self.assertEqual(metrics.n_get, 2)
2092 self.assertGreater(metrics.time_in_get, 0.0)
2093 self.assertEqual(metrics.n_put, 1)
2094 self.assertGreater(metrics.time_in_put, 0.0)
2096 deferred = butler.getDeferred(data_ref)
2097 deferred.get()
2098 self.assertEqual(metrics.n_get, 3)
2100 with butler.record_metrics() as new:
2101 butler.put(data, datasetType, visit=425, instrument="DummyCamComp")
2102 butler.get(data_ref)
2103 self.assertEqual(new.n_get, 1)
2104 self.assertEqual(new.n_put, 1)
2107class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2108 """PosixDatastore specialization of a butler"""
2110 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2111 fullConfigKey: str | None = ".datastore.formatters"
2112 validationCanFail = True
2113 datastoreStr = ["/tmp"]
2114 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2115 registryStr = "/gen3.sqlite3"
2117 def testPathConstructor(self) -> None:
2118 """Independent test of constructor using PathLike."""
2119 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
2120 self.assertIsInstance(butler, Butler)
2122 # And again with a Path object with the butler yaml
2123 path = pathlib.Path(self.tmpConfigFile)
2124 butler = Butler.from_config(path, writeable=False)
2125 self.assertIsInstance(butler, Butler)
2127 # And again with a Path object without the butler yaml
2128 # (making sure we skip it if the tmp config doesn't end
2129 # in butler.yaml -- which is the case for a subclass)
2130 if self.tmpConfigFile.endswith("butler.yaml"):
2131 path = pathlib.Path(os.path.dirname(self.tmpConfigFile))
2132 butler = Butler.from_config(path, writeable=False)
2133 self.assertIsInstance(butler, Butler)
2135 def testExportTransferCopy(self) -> None:
2136 """Test local export using all transfer modes"""
2137 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2138 exportButler = self.runPutGetTest(storageClass, "test_metric")
2139 # Test that the repo actually has at least one dataset.
2140 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
2141 self.assertGreater(len(datasets), 0)
2142 uris = [exportButler.getURI(d) for d in datasets]
2143 assert isinstance(exportButler._datastore, FileDatastore)
2144 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]]
2146 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris]
2148 for path in pathsInStore:
2149 # Assume local file system
2150 assert path is not None
2151 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}")
2153 for transfer in ("copy", "link", "symlink", "relsymlink"):
2154 with safeTestTempDir(TESTDIR) as exportDir:
2155 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export:
2156 export.saveDatasets(datasets)
2157 for path in pathsInStore:
2158 assert path is not None
2159 self.assertTrue(
2160 self.checkFileExists(exportDir, path),
2161 f"Check that mode {transfer} exported files",
2162 )
2164 def testPytypeCoercion(self) -> None:
2165 """Test python type coercion on Butler.get and put."""
2166 # Store some data with the normal example storage class.
2167 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2168 datasetTypeName = "test_metric"
2169 butler = self.runPutGetTest(storageClass, datasetTypeName)
2171 dataId = {"instrument": "DummyCamComp", "visit": 423}
2172 metric = butler.get(datasetTypeName, dataId=dataId)
2173 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
2175 datasetType_ori = butler.get_dataset_type(datasetTypeName)
2176 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
2178 # Now need to hack the registry dataset type definition.
2179 # There is no API for this.
2180 assert isinstance(butler._registry, SqlRegistry)
2181 manager = butler._registry._managers.datasets
2182 assert hasattr(manager, "_db") and hasattr(manager, "_static")
2183 manager._db.update(
2184 manager._static.dataset_type,
2185 {"name": datasetTypeName},
2186 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
2187 )
2189 # Force reset of dataset type cache
2190 butler.registry.refresh()
2192 datasetType_new = butler.get_dataset_type(datasetTypeName)
2193 self.assertEqual(datasetType_new.name, datasetType_ori.name)
2194 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
2196 metric_model = butler.get(datasetTypeName, dataId=dataId)
2197 self.assertNotEqual(type(metric_model), type(metric))
2198 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
2200 # Put the model and read it back to show that everything now
2201 # works as normal.
2202 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
2203 metric_model_new = butler.get(metric_ref)
2204 self.assertEqual(metric_model_new, metric_model)
2206 # Hack the storage class again to something that will fail on the
2207 # get with no conversion class.
2208 manager._db.update(
2209 manager._static.dataset_type,
2210 {"name": datasetTypeName},
2211 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
2212 )
2213 butler.registry.refresh()
2215 with self.assertRaises(ValueError):
2216 butler.get(datasetTypeName, dataId=dataId)
2218 def test_provenance(self):
2219 """Test that provenance is attached on put."""
2220 run = "test_run"
2221 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
2222 metric = MetricsExampleModel(
2223 summary={"AM1": 5.2, "AM2": 30.6},
2224 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2225 data=[563, 234, 456.7, 752, 8, 9, 27],
2226 )
2227 # Provenance can be attached to the object being put. Whether
2228 # it is or not is dependent on the formatter. For this test we
2229 # copy on adding provenance to ensure they differ.
2230 self.assertIsNone(metric.dataset_id)
2231 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
2232 self.assertIsNone(metric.dataset_id)
2233 metric_2 = butler.get(metric_ref)
2234 self.assertEqual(metric_2.data, metric.data)
2235 self.assertEqual(metric_2.dataset_id, metric_ref.id)
2236 self.assertIsNone(metric_2.provenance)
2238 # Put with provenance.
2239 prov = DatasetProvenance(quantum_id=uuid.uuid4())
2240 prov.add_input(metric_ref)
2241 prov.add_extra_provenance(metric_ref.id, {"answer": 42})
2242 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
2243 metric_3 = butler.get(metric_ref2)
2244 self.assertEqual(metric_3.provenance, prov)
2246 # Check that we can extract provenance from dict form.
2247 prov_dict = prov.to_flat_dict(metric_ref2)
2248 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2249 self.assertEqual(ref_from_prov, metric_ref2)
2250 # Direct __eq__ of the provenance does not work because one side
2251 # includes dimension records.
2252 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2253 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id)
2254 self.assertEqual(prov_from_prov.extras, prov.extras)
2256 # Force a bad ID into the dict.
2257 prov_dict["id"] = uuid.uuid4()
2258 with self.assertRaises(ValueError):
2259 DatasetProvenance.from_flat_dict(prov_dict, butler)
2260 del prov_dict["id"]
2261 prov_dict["input 0 id"] = uuid.uuid4()
2262 with self.assertRaises(ValueError):
2263 DatasetProvenance.from_flat_dict(prov_dict, butler)
2265 # Check that simple types can be reconstructed with non-standard
2266 # separators.
2267 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True)
2268 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2269 self.assertEqual(ref_from_prov, metric_ref2)
2270 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2272 with self.assertRaises(ValueError):
2273 DatasetProvenance.from_flat_dict({"unknown": 42}, butler)
2275 def test_specialized_file_datasets_functions(self):
2276 """Test a workflow used in Prompt Processing where we export datasets
2277 from one repository and write them in-place to the datastore of
2278 another, without immediately inserting registry entries for the
2279 datasets.
2280 """
2281 repo = MetricTestRepo.create_from_butler(
2282 self.create_empty_butler(writeable=True),
2283 self.tmpConfigFile,
2284 "StructuredCompositeReadCompNoDisassembly",
2285 )
2286 source_butler = repo.butler
2288 # Test writing outputs to a FileDatastore.
2289 with tempfile.TemporaryDirectory() as tempdir:
2290 target_repo_config = Butler.makeRepo(tempdir)
2291 refs = [repo.ref1, repo.ref2]
2292 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2293 self.assertEqual(len(datasets), 2)
2294 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2295 for dataset in datasets:
2296 path = ResourcePath(dataset.path, forceAbsolute=False)
2297 # Paths should be relative paths to the target datastore.
2298 self.assertFalse(path.isabs())
2299 # Files should have been copied into the target datastore
2300 self.assertTrue(ResourcePath(tempdir).join(path).exists())
2302 # Make sure the target Butler can ingest the datasets.
2303 target_butler = Butler(target_repo_config, writeable=True)
2304 target_butler.transfer_dimension_records_from(source_butler, refs)
2305 target_butler.ingest(*datasets, transfer=None)
2306 self.assertIsNotNone(target_butler.get(repo.ref1))
2307 self.assertIsNotNone(target_butler.get(repo.ref2))
2309 # Giving an empty list of files is a no-op.
2310 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), [])
2311 self.assertEqual(len(no_datasets), 0)
2313 # Test writing outputs to a ChainedDatastore.
2314 with tempfile.TemporaryDirectory() as tempdir:
2315 # Set up a second dataset type, so we can split the files across
2316 # multiple datastore roots.
2317 dt1 = repo.datasetType
2318 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass)
2319 source_butler.registry.registerDatasetType(dt2)
2320 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2)
2321 config = Config.fromString(
2322 f"""
2323 datastore:
2324 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore
2325 datastore_constraints:
2326 - constraints:
2327 accept:
2328 - {dt1.name}
2329 - constraints:
2330 accept:
2331 - {dt2.name}
2332 datastores:
2333 - datastore:
2334 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2335 root: <butlerRoot>/FileDatastore_0
2336 - datastore:
2337 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2338 root: <butlerRoot>/FileDatastore_1
2339 """
2340 )
2341 target_repo_config = Butler.makeRepo(tempdir, config)
2342 refs = [repo.ref1, repo.ref2, other_ref]
2343 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2344 self.assertEqual(len(datasets), 3)
2345 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2346 for dataset in datasets:
2347 path = ResourcePath(dataset.path, forceAbsolute=False)
2348 # Paths should be relative paths to the target datastore.
2349 self.assertFalse(path.isabs())
2350 # Files should have been split up between the two datastores
2351 # in the chain.
2352 datastore_root = ResourcePath(tempdir)
2353 if dataset.refs[0].datasetType.name == dt1.name:
2354 datastore_root = datastore_root.join("FileDatastore_0")
2355 else:
2356 datastore_root = datastore_root.join("FileDatastore_1")
2357 self.assertTrue(datastore_root.join(path).exists())
2359 # Make sure the target Butler can ingest the datasets.
2360 target_butler = Butler(target_repo_config, writeable=True)
2361 target_butler.transfer_dimension_records_from(source_butler, refs)
2362 target_butler.ingest(*datasets, transfer=None)
2363 self.assertIsNotNone(target_butler.get(repo.ref1))
2364 self.assertIsNotNone(target_butler.get(repo.ref2))
2365 self.assertIsNotNone(target_butler.get(other_ref))
2368class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2369 """PosixDatastore specialization of a butler using Postgres"""
2371 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2372 fullConfigKey = ".datastore.formatters"
2373 validationCanFail = True
2374 datastoreStr = ["/tmp"]
2375 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2376 registryStr = "PostgreSQL@test"
2378 @classmethod
2379 def setUpClass(cls) -> None:
2380 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
2381 super().setUpClass()
2383 def setUp(self) -> None:
2384 # Need to add a registry section to the config.
2385 self._temp_config = False
2386 config = Config(self.configFile)
2387 self.postgresql.patch_butler_config(config)
2388 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
2389 config.dump(fh)
2390 self.configFile = fh.name
2391 self._temp_config = True
2392 super().setUp()
2394 def tearDown(self) -> None:
2395 if self._temp_config and os.path.exists(self.configFile):
2396 os.remove(self.configFile)
2397 super().tearDown()
2399 def testMakeRepo(self) -> None:
2400 # The base class test assumes that it's using sqlite and assumes
2401 # the config file is acceptable to sqlite.
2402 raise unittest.SkipTest("Postgres config is not compatible with this test.")
2405class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase):
2406 """Test that Butler with a Postgres registry still works after cloning."""
2408 def create_butler(
2409 self,
2410 run: str,
2411 storageClass: StorageClass | str,
2412 datasetTypeName: str,
2413 metrics: ButlerMetrics | None = None,
2414 ) -> tuple[DirectButler, DatasetType]:
2415 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2416 return butler.clone(run=run, metrics=metrics), datasetType
2419class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
2420 """InMemoryDatastore specialization of a butler"""
2422 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")
2423 fullConfigKey = None
2424 useTempRoot = False
2425 validationCanFail = False
2426 datastoreStr = ["datastore='InMemory"]
2427 datastoreName = ["InMemoryDatastore@"]
2428 registryStr = "/gen3.sqlite3"
2430 def testIngest(self) -> None:
2431 pass
2433 def test_ingest_zip(self) -> None:
2434 pass
2437class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase):
2438 """Test that a Butler with a Sqlite registry still works after cloning."""
2440 def create_butler(
2441 self,
2442 run: str,
2443 storageClass: StorageClass | str,
2444 datasetTypeName: str,
2445 metrics: ButlerMetrics | None = None,
2446 ) -> tuple[DirectButler, DatasetType]:
2447 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2448 return butler.clone(run=run), datasetType
2451class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2452 """PosixDatastore specialization"""
2454 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
2455 fullConfigKey = ".datastore.datastores.1.formatters"
2456 validationCanFail = True
2457 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"]
2458 datastoreName = [
2459 "InMemoryDatastore@",
2460 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1",
2461 "SecondDatastore",
2462 ]
2463 registryStr = "/gen3.sqlite3"
2465 def testPruneDatasets(self) -> None:
2466 # This test relies on manipulating files out-of-band, which is
2467 # impossible for this configuration because of the InMemoryDatastore in
2468 # the ChainedDatastore.
2469 pass
2472class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase):
2473 """Test that a yaml file in one location can refer to a root in another."""
2475 datastoreStr = ["dir1"]
2476 # Disable the makeRepo test since we are deliberately not using
2477 # butler.yaml as the config name.
2478 fullConfigKey = None
2480 def setUp(self) -> None:
2481 self.root = makeTestTempDir(TESTDIR)
2483 # Make a new repository in one place
2484 self.dir1 = os.path.join(self.root, "dir1")
2485 Butler.makeRepo(self.dir1, config=Config(self.configFile))
2487 # Move the yaml file to a different place and add a "root"
2488 self.dir2 = os.path.join(self.root, "dir2")
2489 os.makedirs(self.dir2, exist_ok=True)
2490 configFile1 = os.path.join(self.dir1, "butler.yaml")
2491 config = Config(configFile1)
2492 config["root"] = self.dir1
2493 configFile2 = os.path.join(self.dir2, "butler2.yaml")
2494 config.dumpToUri(configFile2)
2495 os.remove(configFile1)
2496 self.tmpConfigFile = configFile2
2498 def testFileLocations(self) -> None:
2499 self.assertNotEqual(self.dir1, self.dir2)
2500 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml")))
2501 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml")))
2502 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3")))
2505class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase):
2506 """Test that a config file created by makeRepo outside of repo works."""
2508 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2510 def setUp(self) -> None:
2511 self.root = makeTestTempDir(TESTDIR)
2512 self.root2 = makeTestTempDir(TESTDIR)
2514 self.tmpConfigFile = os.path.join(self.root2, "different.yaml")
2515 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2517 def tearDown(self) -> None:
2518 if os.path.exists(self.root2):
2519 shutil.rmtree(self.root2, ignore_errors=True)
2520 super().tearDown()
2522 def testConfigExistence(self) -> None:
2523 c = Config(self.tmpConfigFile)
2524 uri_config = ResourcePath(c["root"])
2525 uri_expected = ResourcePath(self.root, forceDirectory=True)
2526 self.assertEqual(uri_config.geturl(), uri_expected.geturl())
2527 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path")
2529 def testPutGet(self) -> None:
2530 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2531 self.runPutGetTest(storageClass, "test_metric")
2534class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase):
2535 """Test that a config file created by makeRepo outside of repo works."""
2537 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2539 def setUp(self) -> None:
2540 self.root = makeTestTempDir(TESTDIR)
2541 self.root2 = makeTestTempDir(TESTDIR)
2543 self.tmpConfigFile = self.root2
2544 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2546 def testConfigExistence(self) -> None:
2547 # Append the yaml file else Config constructor does not know the file
2548 # type.
2549 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml")
2550 super().testConfigExistence()
2553class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase):
2554 """Test that a config file created by makeRepo outside of repo works."""
2556 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2558 def setUp(self) -> None:
2559 self.root = makeTestTempDir(TESTDIR)
2560 self.root2 = makeTestTempDir(TESTDIR)
2562 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl()
2563 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2566@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!")
2567class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2568 """S3Datastore specialization of a butler; an S3 storage Datastore +
2569 a local in-memory SqlRegistry.
2570 """
2572 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")
2573 fullConfigKey = None
2574 validationCanFail = True
2576 bucketName = "anybucketname"
2577 """Name of the Bucket that will be used in the tests. The name is read from
2578 the config file used with the tests during set-up.
2579 """
2581 root = "butlerRoot/"
2582 """Root repository directory expected to be used in case useTempRoot=False.
2583 Otherwise the root is set to a 20 characters long randomly generated string
2584 during set-up.
2585 """
2587 datastoreStr = [f"datastore={root}"]
2588 """Contains all expected root locations in a format expected to be
2589 returned by Butler stringification.
2590 """
2592 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]
2593 """The expected format of the S3 Datastore string."""
2595 registryStr = "/gen3.sqlite3"
2596 """Expected format of the Registry string."""
2598 mock_aws = mock_aws()
2599 """The mocked s3 interface from moto."""
2601 def genRoot(self) -> str:
2602 """Return a random string of len 20 to serve as a root
2603 name for the temporary bucket repo.
2605 This is equivalent to tempfile.mkdtemp as this is what self.root
2606 becomes when useTempRoot is True.
2607 """
2608 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20))
2609 return rndstr + "/"
2611 def setUp(self) -> None:
2612 config = Config(self.configFile)
2613 uri = ResourcePath(config[".datastore.datastore.root"])
2614 self.bucketName = uri.netloc
2616 # Enable S3 mocking of tests.
2617 self.enterContext(clean_test_environment_for_s3())
2618 self.mock_aws.start()
2620 if self.useTempRoot:
2621 self.root = self.genRoot()
2622 rooturi = f"s3://{self.bucketName}/{self.root}"
2623 config.update({"datastore": {"datastore": {"root": rooturi}}})
2625 # need local folder to store registry database
2626 self.reg_dir = makeTestTempDir(TESTDIR)
2627 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"
2629 # MOTO needs to know that we expect Bucket bucketname to exist
2630 # (this used to be the class attribute bucketName)
2631 s3 = boto3.resource("s3")
2632 s3.create_bucket(Bucket=self.bucketName)
2634 self.datastoreStr = [f"datastore='{rooturi}'"]
2635 self.datastoreName = [f"FileDatastore@{rooturi}"]
2636 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)
2637 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")
2639 def tearDown(self) -> None:
2640 s3 = boto3.resource("s3")
2641 bucket = s3.Bucket(self.bucketName)
2642 try:
2643 bucket.objects.all().delete()
2644 except botocore.exceptions.ClientError as e:
2645 if e.response["Error"]["Code"] == "404":
2646 # the key was not reachable - pass
2647 pass
2648 else:
2649 raise
2651 bucket = s3.Bucket(self.bucketName)
2652 bucket.delete()
2654 # Stop the S3 mock.
2655 self.mock_aws.stop()
2657 if self.reg_dir is not None and os.path.exists(self.reg_dir):
2658 shutil.rmtree(self.reg_dir, ignore_errors=True)
2660 if self.useTempRoot and os.path.exists(self.root):
2661 shutil.rmtree(self.root, ignore_errors=True)
2663 super().tearDown()
2666class DatastoreTransfers(TestCaseMixin):
2667 """Base test setup for data transfers between butlers. The concrete tests
2668 for specific configurations are in other classes, below.
2669 """
2671 storageClassFactory: StorageClassFactory
2673 @classmethod
2674 def setUpClass(cls) -> None:
2675 cls.storageClassFactory = StorageClassFactory()
2677 def setUp(self) -> None:
2678 self.root = makeTestTempDir(TESTDIR)
2679 self.config = Config(self.configFile)
2681 # Some tests cause convertors to be replaced so ensure
2682 # the storage class factory is reset each time.
2683 self.storageClassFactory.reset()
2684 self.storageClassFactory.addFromConfig(self.configFile)
2686 def tearDown(self) -> None:
2687 removeTestTempDir(self.root)
2689 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler:
2690 if manager is None:
2691 manager = (
2692 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
2693 )
2694 config = Config(config_file if config_file is not None else self.configFile)
2695 config["registry", "managers", "datasets"] = manager
2696 return Butler.from_config(
2697 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True
2698 )
2700 def assertButlerTransfers(
2701 self,
2702 purge: bool = False,
2703 storageClassName: str = "StructuredData",
2704 storageClassNameTarget: str | None = None,
2705 ) -> None:
2706 """Test that a run can be transferred to another butler."""
2707 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
2708 if storageClassNameTarget is not None:
2709 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget)
2710 else:
2711 storageClassTarget = storageClass
2713 datasetTypeName = "random_data"
2715 # Test will create 3 collections and we will want to transfer
2716 # two of those three.
2717 runs = ["run1", "run2", "other"]
2719 # Also want to use two different dataset types to ensure that
2720 # grouping works.
2721 datasetTypeNames = ["random_data", "random_data_2"]
2723 # Create the run collections in the source butler.
2724 for run in runs:
2725 self.source_butler.collections.register(run)
2727 # Create dimensions in source butler.
2728 n_exposures = 30
2729 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
2730 self.source_butler.registry.insertDimensionData(
2731 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
2732 )
2733 self.source_butler.registry.insertDimensionData(
2734 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
2735 )
2736 self.source_butler.registry.insertDimensionData(
2737 "day_obs",
2738 {
2739 "instrument": "DummyCamComp",
2740 "id": 20250101,
2741 },
2742 )
2744 for i in range(n_exposures):
2745 self.source_butler.registry.insertDimensionData(
2746 "group", {"instrument": "DummyCamComp", "name": f"group{i}"}
2747 )
2748 self.source_butler.registry.insertDimensionData(
2749 "exposure",
2750 {
2751 "instrument": "DummyCamComp",
2752 "id": i,
2753 "obs_id": f"exp{i}",
2754 "physical_filter": "d-r",
2755 "group": f"group{i}",
2756 "day_obs": 20250101,
2757 },
2758 )
2760 # Create dataset types in the source butler.
2761 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"])
2762 for datasetTypeName in datasetTypeNames:
2763 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
2764 self.source_butler.registry.registerDatasetType(datasetType)
2766 # Write a dataset to an unrelated run -- this will ensure that
2767 # we are rewriting integer dataset ids in the target if necessary.
2768 # Will not be relevant for UUID.
2769 run = "distraction"
2770 butler = Butler.from_config(butler=self.source_butler, run=run)
2771 butler.put(
2772 makeExampleMetrics(),
2773 datasetTypeName,
2774 exposure=1,
2775 instrument="DummyCamComp",
2776 physical_filter="d-r",
2777 )
2779 # Write some example metrics to the source
2780 butler = Butler.from_config(butler=self.source_butler)
2782 # Set of DatasetRefs that should be in the list of refs to transfer
2783 # but which will not be transferred.
2784 deleted: set[DatasetRef] = set()
2786 n_expected = 20 # Number of datasets expected to be transferred
2787 source_refs = []
2788 for i in range(n_exposures):
2789 # Put a third of datasets into each collection, only retain
2790 # two thirds.
2791 index = i % 3
2792 run = runs[index]
2793 datasetTypeName = datasetTypeNames[i % 2]
2795 metric = MetricsExample(
2796 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)]
2797 )
2798 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"}
2799 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run)
2801 # Remove the datastore record using low-level API, but only
2802 # for a specific index.
2803 if purge and index == 1:
2804 # For one of these delete the file as well.
2805 # This allows the "missing" code to filter the
2806 # file out.
2807 # Access the individual datastores.
2808 datastores = []
2809 if hasattr(butler._datastore, "datastores"):
2810 datastores.extend(butler._datastore.datastores)
2811 else:
2812 datastores.append(butler._datastore)
2814 if not deleted:
2815 # For a chained datastore we need to remove
2816 # files in each chain.
2817 for datastore in datastores:
2818 # The file might not be known to the datastore
2819 # if constraints are used.
2820 try:
2821 primary, uris = datastore.getURIs(ref)
2822 except FileNotFoundError:
2823 continue
2824 if primary and primary.scheme != "mem":
2825 primary.remove()
2826 for uri in uris.values():
2827 if uri.scheme != "mem":
2828 uri.remove()
2829 n_expected -= 1
2830 deleted.add(ref)
2832 # Remove the datastore record.
2833 for datastore in datastores:
2834 if hasattr(datastore, "removeStoredItemInfo"):
2835 datastore.removeStoredItemInfo(ref)
2837 if index < 2:
2838 source_refs.append(ref)
2839 if ref not in deleted:
2840 new_metric = butler.get(ref)
2841 self.assertEqual(new_metric, metric)
2843 # Create some bad dataset types to ensure we check for inconsistent
2844 # definitions.
2845 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList")
2846 for datasetTypeName in datasetTypeNames:
2847 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass)
2848 self.target_butler.registry.registerDatasetType(datasetType)
2849 with self.assertRaises(ConflictingDefinitionError) as cm:
2850 self.target_butler.transfer_from(self.source_butler, source_refs)
2851 self.assertIn("dataset type differs", str(cm.exception))
2853 # And remove the bad definitions.
2854 for datasetTypeName in datasetTypeNames:
2855 self.target_butler.registry.removeDatasetType(datasetTypeName)
2857 # Transfer without creating dataset types should fail.
2858 with self.assertRaises(KeyError):
2859 self.target_butler.transfer_from(self.source_butler, source_refs)
2861 # Transfer without creating dimensions should fail.
2862 with self.assertRaises(ConflictingDefinitionError) as cm:
2863 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True)
2864 self.assertIn("dimension", str(cm.exception))
2866 # The dry run test requires dataset types to exist. If we have
2867 # been given distinct storage classes for the target we have
2868 # to redefine at least one of the dataset types in the target butler.
2869 if storageClass != storageClassTarget:
2870 self.target_butler.registry.removeDatasetType(datasetTypeNames[0])
2871 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget)
2872 self.target_butler.registry.registerDatasetType(datasetType)
2874 # The failed transfer above leaves registry in an inconsistent
2875 # state because the run is created but then rolled back without
2876 # the collection cache being cleared. For now force a refresh.
2877 # Can remove with DM-35498.
2878 self.target_butler.registry.refresh()
2880 # Do a dry run -- this should not have any effect on the target butler.
2881 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True)
2883 # Transfer the records for one ref to test the alternative API.
2884 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2885 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]])
2886 self.assertIn("number of records transferred: 1", ";".join(log_cm.output))
2888 # Now transfer them to the second butler, including dimensions.
2889 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2890 transferred = self.target_butler.transfer_from(
2891 self.source_butler,
2892 source_refs,
2893 register_dataset_types=True,
2894 transfer_dimensions=True,
2895 )
2896 self.assertEqual(len(transferred), n_expected)
2897 log_output = ";".join(log_cm.output)
2899 # A ChainedDatastore will use the in-memory datastore for mexists
2900 # so we can not rely on the mexists log message.
2901 self.assertIn("Number of datastore records found in source", log_output)
2902 self.assertIn("Creating output run", log_output)
2904 # Do the transfer twice to ensure that it will do nothing extra.
2905 # Only do this if purge=True because it does not work for int
2906 # dataset_id.
2907 if purge:
2908 # This should not need to register dataset types.
2909 transferred = self.target_butler.transfer_from(self.source_butler, source_refs)
2910 self.assertEqual(len(transferred), n_expected)
2912 with self.assertRaises((TypeError, AttributeError)):
2913 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore
2915 with self.assertRaises(ValueError):
2916 self.target_butler._datastore.transfer_from(
2917 self.source_butler._datastore, source_refs, transfer="split"
2918 )
2920 # Now try to get the same refs from the new butler.
2921 for ref in source_refs:
2922 if ref not in deleted:
2923 new_metric = self.target_butler.get(ref)
2924 old_metric = self.source_butler.get(ref)
2925 self.assertEqual(new_metric, old_metric)
2927 # Try again without implicit storage class conversion
2928 # triggered by using the source ref. This will do conversion
2929 # since the formatter will be returning the source python type.
2930 target_ref = self.target_butler.get_dataset(ref.id)
2931 if target_ref.datasetType.storageClass != ref.datasetType.storageClass:
2932 new_metric = self.target_butler.get(target_ref)
2933 self.assertNotEqual(type(new_metric), type(old_metric))
2935 # Remove the dataset from the target and put it again
2936 # as if it was the right type all along for this butler.
2937 self.target_butler.pruneDatasets(
2938 [target_ref], unstore=True, purge=True, disassociate=True
2939 )
2940 self.target_butler.put(new_metric, target_ref)
2941 new_new_metric = self.target_butler.get(target_ref)
2942 new_old_metric = self.target_butler.get(
2943 target_ref, storageClass=ref.datasetType.storageClass
2944 )
2945 self.assertEqual(new_new_metric, new_metric)
2946 self.assertEqual(new_old_metric, old_metric)
2948 # Now prune run2 collection and create instead a CHAINED collection.
2949 # This should block the transfer.
2950 self.target_butler.removeRuns(["run2"])
2951 self.target_butler.collections.register("run2", CollectionType.CHAINED)
2952 with self.assertRaises(CollectionTypeError):
2953 # Re-importing the run1 datasets can be problematic if they
2954 # use integer IDs so filter those out.
2955 to_transfer = [ref for ref in source_refs if ref.run == "run2"]
2956 self.target_butler.transfer_from(self.source_butler, to_transfer)
2959class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
2960 """Test data transfers between butlers.
2962 Test for different managers. UUID to UUID and integer to integer are
2963 tested. UUID to integer is not supported since we do not currently
2964 want to allow that. Integer to UUID is supported with the caveat
2965 that UUID4 will be generated and this will be incorrect for raw
2966 dataset types. The test ignores that.
2967 """
2969 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2971 def create_butlers(
2972 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None
2973 ) -> None:
2974 self.source_butler = self.create_butler(manager1, "1", config_file=source_config)
2975 self.target_butler = self.create_butler(manager2, "2")
2977 def testTransferUuidToUuid(self) -> None:
2978 self.create_butlers()
2979 self.assertButlerTransfers()
2981 def testTransferFromChainedUuidToUuid(self) -> None:
2982 """Force the source butler to be a ChainedDatastore."""
2983 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml"))
2984 self.assertButlerTransfers()
2986 def testTransferFromIncompatibleUuidToUuid(self) -> None:
2987 """Force the source butler to be a incompatible datastore."""
2988 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml"))
2989 with self.assertRaises(NotImplementedError):
2990 self.assertButlerTransfers()
2992 def testTransferFromIncompatibleChainUuidToUuid(self) -> None:
2993 """Force the source butler to be a incompatible datastore."""
2994 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml"))
2995 with self.assertRaises(TypeError):
2996 self.assertButlerTransfers()
2998 def testTransferFromFileUuidToUuid(self) -> None:
2999 """Force the source butler to be a FileDatastore."""
3000 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml"))
3001 self.assertButlerTransfers()
3003 def testTransferMissing(self) -> None:
3004 """Test transfers where datastore records are missing.
3006 This is how execution butler works.
3007 """
3008 self.create_butlers()
3010 # Configure the source butler to allow trust.
3011 self.source_butler._datastore._set_trust_mode(True)
3013 self.assertButlerTransfers(purge=True)
3015 def testTransferMissingDisassembly(self) -> None:
3016 """Test transfers where datastore records are missing.
3018 This is how execution butler works.
3019 """
3020 self.create_butlers()
3022 # Configure the source butler to allow trust.
3023 self.source_butler._datastore._set_trust_mode(True)
3025 # Test disassembly.
3026 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite")
3028 def testTransferDifferingStorageClasses(self) -> None:
3029 """Test transfers when the source butler dataset type has a different
3030 but compatible storage class.
3031 """
3032 self.create_butlers()
3034 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion")
3036 def testTransferDifferingStorageClassesDisassembly(self) -> None:
3037 """Test transfers when the source butler dataset type has a different
3038 but compatible storage class and where the source butler has
3039 disassembled.
3040 """
3041 self.create_butlers()
3043 self.assertButlerTransfers(
3044 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion"
3045 )
3047 def testAbsoluteURITransferDirect(self) -> None:
3048 """Test transfer using an absolute URI."""
3049 self._absolute_transfer("auto")
3051 def testAbsoluteURITransferCopy(self) -> None:
3052 """Test transfer using an absolute URI."""
3053 self._absolute_transfer("copy")
3055 def _absolute_transfer(self, transfer: str) -> None:
3056 self.create_butlers()
3058 storageClassName = "StructuredData"
3059 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
3060 datasetTypeName = "random_data"
3061 run = "run1"
3062 self.source_butler.collections.register(run)
3064 dimensions = self.source_butler.dimensions.conform(())
3065 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3066 self.source_butler.registry.registerDatasetType(datasetType)
3068 metrics = makeExampleMetrics()
3069 with ResourcePath.temporary_uri(suffix=".json") as temp:
3070 dataId = DataCoordinate.make_empty(self.source_butler.dimensions)
3071 source_refs = [DatasetRef(datasetType, dataId, run=run)]
3072 temp.write(json.dumps(metrics.exportAsDict()).encode())
3073 dataset = FileDataset(path=temp, refs=source_refs)
3074 self.source_butler.ingest(dataset, transfer="direct")
3076 self.target_butler.transfer_from(
3077 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer
3078 )
3080 uri = self.target_butler.getURI(dataset.refs[0])
3081 if transfer == "auto":
3082 self.assertEqual(uri, temp)
3083 else:
3084 self.assertNotEqual(uri, temp)
3086 def test_shared_dimension_group(self):
3087 """Test internal logic that divides dataset types by dimension group
3088 when doing registry updates.
3089 """
3090 self.create_butlers()
3091 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
3092 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
3094 source_butler = self.source_butler
3095 target_butler = self.target_butler
3097 # Create a dataset type with the same dimensions as the 'bias' dataset
3098 # type from base.yaml
3099 dataset_type = DatasetType(
3100 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
3101 )
3102 source_butler.registry.registerDatasetType(dataset_type)
3103 # This has the same data ID as one of the bias datasets in
3104 # datasets.yaml.
3105 test_ref = source_butler.registry.insertDatasets(
3106 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
3107 )[0]
3109 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
3110 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
3111 refs = [test_ref, *biases, *flats]
3113 # Test setup will be even more convoluted if we want the datastore to
3114 # actually transfer files. For testing the dimension group behavior,
3115 # we really only care about the registry.
3116 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
3117 mock.return_value = (set(refs), set())
3118 target_butler.transfer_from(
3119 source_butler,
3120 refs,
3121 transfer=None,
3122 register_dataset_types=True,
3123 skip_missing=False,
3124 transfer_dimensions=True,
3125 )
3127 transferred_test_ref = target_butler.find_dataset(
3128 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3129 )
3130 self.assertEqual(transferred_test_ref.id, test_ref.id)
3132 transferred_bias = target_butler.find_dataset(
3133 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3134 )
3135 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
3137 transferred_flat = target_butler.find_dataset(
3138 "flat",
3139 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
3140 collections="imported_r",
3141 )
3142 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
3145class ChainedDatastoreTransfers(PosixDatastoreTransfers):
3146 """Test transfers using a chained datastore."""
3148 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
3151@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3152class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3153 """Test ``transfer_from`` involving Butler server."""
3155 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3157 def test_transfers_from_remote_to_direct(self) -> None:
3158 from lsst.daf.butler.remote_butler._remote_file_transfer_source import (
3159 mock_file_transfer_uris_for_unit_test,
3160 )
3162 self.target_butler = self.create_butler(None, "2")
3163 with create_test_server(TESTDIR) as server:
3164 self.source_butler = server.hybrid_butler
3166 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath:
3167 # The Butler server returns HTTP URIs with a domain name that
3168 # is not resolvable because there is no actual HTTP server
3169 # involved in these tests. Strip this first layer of
3170 # indirection, and return the target of the redirect instead.
3171 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers)
3172 return ResourcePath(str(response.next_request.url))
3174 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url):
3175 self.assertButlerTransfers()
3178class NullDatastoreTestCase(unittest.TestCase):
3179 """Test that we can fall back to a null datastore."""
3181 # Need a good config to create the repo.
3182 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3183 storageClassFactory: StorageClassFactory
3185 @classmethod
3186 def setUpClass(cls) -> None:
3187 cls.storageClassFactory = StorageClassFactory()
3188 cls.storageClassFactory.addFromConfig(cls.configFile)
3190 def setUp(self) -> None:
3191 """Create a new butler root for each test."""
3192 self.root = makeTestTempDir(TESTDIR)
3193 Butler.makeRepo(self.root, config=Config(self.configFile))
3195 def tearDown(self) -> None:
3196 removeTestTempDir(self.root)
3198 def test_fallback(self) -> None:
3199 # Read the butler config and mess with the datastore section.
3200 config_path = os.path.join(self.root, "butler.yaml")
3201 bad_config = Config(config_path)
3202 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore"
3203 bad_config.dumpToUri(config_path)
3205 with self.assertRaises(RuntimeError):
3206 Butler(self.root, without_datastore=False)
3208 with self.assertRaises(RuntimeError):
3209 Butler.from_config(self.root, without_datastore=False)
3211 butler = Butler.from_config(self.root, writeable=True, without_datastore=True)
3212 self.assertIsInstance(butler._datastore, NullDatastore)
3214 # Check that registry is working.
3215 butler.collections.register("MYRUN")
3216 collections = butler.collections.query("*")
3217 self.assertIn("MYRUN", set(collections))
3219 # Create a ref.
3220 dimensions = butler.dimensions.conform([])
3221 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
3222 datasetTypeName = "metric"
3223 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3224 butler.registry.registerDatasetType(datasetType)
3225 ref = DatasetRef(datasetType, {}, run="MYRUN")
3227 # Check that datastore will complain.
3228 with self.assertRaises(FileNotFoundError):
3229 butler.get(ref)
3230 with self.assertRaises(FileNotFoundError):
3231 butler.getURI(ref)
3234@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3235class ButlerServerTests(FileDatastoreButlerTests):
3236 """Test RemoteButler and Butler server."""
3238 configFile = None
3239 predictionSupported = False
3240 trustModeSupported = False
3242 postgres: TemporaryPostgresInstance | None
3244 def setUp(self):
3245 self.server_instance = self.enterContext(create_test_server(TESTDIR))
3247 def tearDown(self):
3248 pass
3250 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
3251 # S3 pre-signed URLs may end up with differing expiration times in the
3252 # query parameters, so ignore query parameters when comparing.
3253 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path
3255 def create_empty_butler(
3256 self, run: str | None = None, writeable: bool | None = None, metrics: ButlerMetrics | None = None
3257 ) -> Butler:
3258 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics)
3260 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
3261 # Can't delete a file via S3 signed URLs, so we need to reach in
3262 # through DirectButler to delete the dataset.
3263 uri = self.server_instance.direct_butler.getURI(ref)
3264 uri.remove()
3266 def testConstructor(self):
3267 # RemoteButler constructor is tested in test_server.py and
3268 # test_remote_butler.py.
3269 pass
3271 def testDafButlerRepositories(self):
3272 # Loading of RemoteButler via repository index is tested in
3273 # test_server.py.
3274 pass
3276 def testGetDatasetTypes(self) -> None:
3277 # This is mostly a test of validateConfiguration, which is for
3278 # validating Datastore configuration and thus isn't relevant to
3279 # RemoteButler.
3280 pass
3282 def testMakeRepo(self) -> None:
3283 # Only applies to DirectButler.
3284 pass
3286 # Pickling not yet implemented for RemoteButler/HybridButler.
3287 @unittest.expectedFailure
3288 def testPickle(self) -> None:
3289 return super().testPickle()
3291 def testStringification(self) -> None:
3292 self.assertEqual(
3293 str(self.server_instance.remote_butler),
3294 "RemoteButler(https://test.example/api/butler/repo/testrepo/)",
3295 )
3297 def testTransaction(self) -> None:
3298 # Transactions will never be supported for RemoteButler.
3299 pass
3301 def testPutTemplates(self) -> None:
3302 # The Butler server instance is configured with different file naming
3303 # templates than this test is expecting.
3304 pass
3307@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3308class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase):
3309 """Tests for RemoteButler's registry shim, with a SQLite DB backing the
3310 server.
3311 """
3313 postgres = None
3316@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3317class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase):
3318 """Tests for RemoteButler's registry shim, with a Postgres DB backing the
3319 server.
3320 """
3322 @classmethod
3323 def setUpClass(cls):
3324 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
3325 super().setUpClass()
3328def setup_module(module: types.ModuleType) -> None:
3329 """Set up the module for pytest."""
3330 clean_environment()
3333def _get_test_data_path(filename: str) -> ResourcePath:
3334 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}")
3337if __name__ == "__main__":
3338 clean_environment()
3339 unittest.main()