Coverage for tests/test_butler.py: 14%
1699 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-10 01:29 -0700
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
67try:
68 from lsst.daf.butler.tests.server import create_test_server
69except ImportError:
70 create_test_server = None
72import astropy.time
73from sqlalchemy.exc import IntegrityError
75from lsst.daf.butler import (
76 Butler,
77 ButlerConfig,
78 ButlerMetrics,
79 ButlerRepoIndex,
80 CollectionCycleError,
81 CollectionType,
82 Config,
83 DataCoordinate,
84 DatasetExistence,
85 DatasetNotFoundError,
86 DatasetProvenance,
87 DatasetRef,
88 DatasetType,
89 FileDataset,
90 NoDefaultCollectionError,
91 StorageClassFactory,
92 ValidationError,
93 script,
94)
95from lsst.daf.butler._rubin.file_datasets import transfer_datasets_to_datastore
96from lsst.daf.butler.datastore import NullDatastore
97from lsst.daf.butler.datastore.file_templates import FileTemplate, FileTemplateValidationError
98from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex
99from lsst.daf.butler.datastores.fileDatastore import FileDatastore
100from lsst.daf.butler.direct_butler import DirectButler
101from lsst.daf.butler.registry import (
102 CollectionError,
103 CollectionTypeError,
104 ConflictingDefinitionError,
105 DataIdValueError,
106 DatasetTypeExpressionError,
107 MissingCollectionError,
108 OrphanedRecordError,
109)
110from lsst.daf.butler.registry.sql_registry import SqlRegistry
111from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG
112from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter
113from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db
114from lsst.daf.butler.tests.utils import (
115 MetricTestRepo,
116 TestCaseMixin,
117 makeTestTempDir,
118 removeTestTempDir,
119 safeTestTempDir,
120)
121from lsst.resources import ResourcePath
122from lsst.resources.http import HttpResourcePath
123from lsst.utils import doImportType
124from lsst.utils.introspection import get_full_type_name
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=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)
1492 def testCollectionChainRedefine(self):
1493 butler = self._setup_to_test_collection_chain()
1495 butler.collections.redefine_chain("chain", "a")
1496 self._check_chain(butler, ["a"])
1498 # Duplicates are removed from the list of children
1499 butler.collections.redefine_chain("chain", ["c", "b", "c"])
1500 self._check_chain(butler, ["c", "b"])
1502 # Empty list clears the chain
1503 butler.collections.redefine_chain("chain", [])
1504 self._check_chain(butler, [])
1506 self._test_common_chain_functionality(butler, butler.collections.redefine_chain)
1508 def testCollectionChainPrepend(self):
1509 butler = self._setup_to_test_collection_chain()
1511 # Duplicates are removed from the list of children
1512 butler.collections.prepend_chain("chain", ["c", "b", "c"])
1513 self._check_chain(butler, ["c", "b"])
1515 # Prepend goes on the front of existing chain
1516 butler.collections.prepend_chain("chain", ["a"])
1517 self._check_chain(butler, ["a", "c", "b"])
1519 # Empty prepend does nothing
1520 butler.collections.prepend_chain("chain", [])
1521 self._check_chain(butler, ["a", "c", "b"])
1523 # Prepending children that already exist in the chain removes them from
1524 # their current position.
1525 butler.collections.prepend_chain("chain", ["d", "b", "c"])
1526 self._check_chain(butler, ["d", "b", "c", "a"])
1528 self._test_common_chain_functionality(butler, butler.collections.prepend_chain)
1530 def testCollectionChainExtend(self):
1531 butler = self._setup_to_test_collection_chain()
1533 # Duplicates are removed from the list of children
1534 butler.collections.extend_chain("chain", ["c", "b", "c"])
1535 self._check_chain(butler, ["c", "b"])
1537 # Extend goes on the end of existing chain
1538 butler.collections.extend_chain("chain", ["a"])
1539 self._check_chain(butler, ["c", "b", "a"])
1541 # Empty extend does nothing
1542 butler.collections.extend_chain("chain", [])
1543 self._check_chain(butler, ["c", "b", "a"])
1545 # Extending children that already exist in the chain removes them from
1546 # their current position.
1547 butler.collections.extend_chain("chain", ["d", "b", "c"])
1548 self._check_chain(butler, ["a", "d", "b", "c"])
1550 self._test_common_chain_functionality(butler, butler.collections.extend_chain)
1552 def testCollectionChainRemove(self) -> None:
1553 butler = self._setup_to_test_collection_chain()
1555 butler.collections.redefine_chain("chain", ["a", "b", "c", "d"])
1557 butler.collections.remove_from_chain("chain", "c")
1558 self._check_chain(butler, ["a", "b", "d"])
1560 # Duplicates are allowed in the list of children
1561 butler.collections.remove_from_chain("chain", ["b", "b", "a"])
1562 self._check_chain(butler, ["d"])
1564 # Empty remove does nothing
1565 butler.collections.remove_from_chain("chain", [])
1566 self._check_chain(butler, ["d"])
1568 # Removing children that aren't in the chain does nothing
1569 butler.collections.remove_from_chain("chain", ["a", "chain"])
1570 self._check_chain(butler, ["d"])
1572 self._test_common_chain_functionality(
1573 butler, butler.collections.remove_from_chain, skip_cycle_check=True
1574 )
1576 def _setup_to_test_collection_chain(self) -> Butler:
1577 butler = self.create_empty_butler(writeable=True)
1579 butler.collections.register("chain", CollectionType.CHAINED)
1581 runs = ["a", "b", "c", "d"]
1582 for run in runs:
1583 butler.collections.register(run)
1585 butler.collections.register("staticchain", CollectionType.CHAINED)
1586 butler.collections.redefine_chain("staticchain", ["a", "b"])
1588 return butler
1590 def _check_chain(self, butler: Butler, expected: list[str]) -> None:
1591 children = butler.collections.get_info("chain").children
1592 self.assertEqual(expected, list(children))
1594 def _test_common_chain_functionality(
1595 self, butler, func: Callable[[str, str | list[str]], Any], *, skip_cycle_check=False
1596 ) -> None:
1597 # Missing parent collection
1598 with self.assertRaises(MissingCollectionError):
1599 func("doesnotexist", [])
1600 # Missing child collection
1601 with self.assertRaises(MissingCollectionError):
1602 func("chain", ["doesnotexist"])
1603 # Forbid operations on non-chained collections
1604 with self.assertRaises(CollectionTypeError):
1605 func("d", ["a"])
1607 # Prevent collection cycles
1608 if not skip_cycle_check:
1609 butler.collections.register("chain2", CollectionType.CHAINED)
1610 func("chain2", "chain")
1611 with self.assertRaises(CollectionCycleError):
1612 func("chain", "chain2")
1614 # Make sure none of the earlier operations interfered with unrelated
1615 # chains.
1616 self.assertEqual(["a", "b"], list(butler.collections.get_info("staticchain").children))
1618 with butler._caching_context():
1619 with self.assertRaisesRegex(RuntimeError, "Chained collection modification not permitted"):
1620 func("chain", "a")
1623class FileDatastoreButlerTests(ButlerTests):
1624 """Common tests and specialization of ButlerTests for butlers backed
1625 by datastores that inherit from FileDatastore.
1626 """
1628 trustModeSupported = True
1630 def checkFileExists(self, root: str | ResourcePath, relpath: str | ResourcePath) -> bool:
1631 """Check if file exists at a given path (relative to root).
1633 Test testPutTemplates verifies actual physical existance of the files
1634 in the requested location.
1635 """
1636 uri = ResourcePath(root, forceDirectory=True)
1637 return uri.join(relpath).exists()
1639 def testPutTemplates(self) -> None:
1640 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1641 butler = self.create_empty_butler(run=self.default_run)
1643 # Add needed Dimensions
1644 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1645 butler.registry.insertDimensionData(
1646 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1647 )
1648 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1649 butler.registry.insertDimensionData(
1650 "visit",
1651 {
1652 "instrument": "DummyCamComp",
1653 "id": 423,
1654 "name": "v423",
1655 "physical_filter": "d-r",
1656 "day_obs": 20250101,
1657 },
1658 )
1659 butler.registry.insertDimensionData(
1660 "visit",
1661 {
1662 "instrument": "DummyCamComp",
1663 "id": 425,
1664 "name": "v425",
1665 "physical_filter": "d-r",
1666 "day_obs": 20250101,
1667 },
1668 )
1670 # Create and store a dataset
1671 metric = makeExampleMetrics()
1673 # Create two almost-identical DatasetTypes (both will use default
1674 # template)
1675 dimensions = butler.dimensions.conform(["instrument", "visit"])
1676 butler.registry.registerDatasetType(DatasetType("metric1", dimensions, storageClass))
1677 butler.registry.registerDatasetType(DatasetType("metric2", dimensions, storageClass))
1678 butler.registry.registerDatasetType(DatasetType("metric3", dimensions, storageClass))
1680 dataId1 = {"instrument": "DummyCamComp", "visit": 423}
1681 dataId2 = {"instrument": "DummyCamComp", "visit": 423, "physical_filter": "d-r"}
1683 # Put with exactly the data ID keys needed
1684 ref = butler.put(metric, "metric1", dataId1)
1685 uri = butler.getURI(ref)
1686 self.assertTrue(uri.exists())
1687 self.assertTrue(
1688 uri.unquoted_path.endswith(f"{self.default_run}/metric1/??#?/d-r/DummyCamComp_423.pickle")
1689 )
1691 # Check the template based on dimensions
1692 if hasattr(butler._datastore, "templates"):
1693 butler._datastore.templates.validateTemplates([ref])
1695 # Put with extra data ID keys (physical_filter is an optional
1696 # dependency); should not change template (at least the way we're
1697 # defining them to behave now; the important thing is that they
1698 # must be consistent).
1699 ref = butler.put(metric, "metric2", dataId2)
1700 uri = butler.getURI(ref)
1701 self.assertTrue(uri.exists())
1702 self.assertTrue(
1703 uri.unquoted_path.endswith(f"{self.default_run}/metric2/d-r/DummyCamComp_v423.pickle")
1704 )
1706 # Check the template based on dimensions
1707 if hasattr(butler._datastore, "templates"):
1708 butler._datastore.templates.validateTemplates([ref])
1710 # Use a template that has a typo in dimension record metadata.
1711 # Easier to test with a butler that has a ref with records attached.
1712 template = FileTemplate("a/{visit.name}/{id}_{visit.namex:?}.fits")
1713 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1714 path = template.format(ref)
1715 self.assertEqual(path, f"a/v423/{ref.id}_fits")
1717 template = FileTemplate("a/{visit.name}/{id}_{visit.namex}.fits")
1718 with self.assertRaises(KeyError):
1719 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1720 template.format(ref)
1722 # Now use a file template that will not result in unique filenames
1723 with self.assertRaises(FileTemplateValidationError):
1724 butler.put(metric, "metric3", dataId1)
1726 def testImportExport(self) -> None:
1727 # Run put/get tests just to create and populate a repo.
1728 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1729 self.runImportExportTest(storageClass)
1731 @unittest.expectedFailure
1732 def testImportExportVirtualComposite(self) -> None:
1733 # Run put/get tests just to create and populate a repo.
1734 storageClass = self.storageClassFactory.getStorageClass("StructuredComposite")
1735 self.runImportExportTest(storageClass)
1737 def runImportExportTest(self, storageClass: StorageClass) -> None:
1738 """Test exporting and importing.
1740 This test does an export to a temp directory and an import back
1741 into a new temp directory repo. It does not assume a posix datastore.
1742 """
1743 exportButler = self.runPutGetTest(storageClass, "test_metric")
1745 # Test that we must have a file extension.
1746 with self.assertRaises(ValueError):
1747 with exportButler.export(filename="dump", directory=".") as export:
1748 pass
1750 # Test that unknown format is not allowed.
1751 with self.assertRaises(ValueError):
1752 with exportButler.export(filename="dump.fits", directory=".") as export:
1753 pass
1755 # Test that the repo actually has at least one dataset.
1756 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
1757 self.assertGreater(len(datasets), 0)
1758 # Add a DimensionRecord that's unused by those datasets.
1759 skymapRecord = {"name": "example_skymap", "hash": (50).to_bytes(8, byteorder="little")}
1760 exportButler.registry.insertDimensionData("skymap", skymapRecord)
1761 # Export and then import datasets.
1762 with safeTestTempDir(TESTDIR) as exportDir:
1763 exportFile = os.path.join(exportDir, "exports.yaml")
1764 with exportButler.export(filename=exportFile, directory=exportDir, transfer="auto") as export:
1765 export.saveDatasets(datasets)
1766 # Export the same datasets again. This should quietly do
1767 # nothing because of internal deduplication, and it shouldn't
1768 # complain about being asked to export the "htm7" elements even
1769 # though there aren't any in these datasets or in the database.
1770 export.saveDatasets(datasets, elements=["htm7"])
1771 # Save one of the data IDs again; this should be harmless
1772 # because of internal deduplication.
1773 export.saveDataIds([datasets[0].dataId])
1774 # Save some dimension records directly.
1775 export.saveDimensionData("skymap", [skymapRecord])
1776 self.assertTrue(os.path.exists(exportFile))
1777 with safeTestTempDir(TESTDIR) as importDir:
1778 # We always want this to be a local posix butler
1779 Butler.makeRepo(importDir, config=Config(os.path.join(TESTDIR, "config/basic/butler.yaml")))
1780 # Calling script.butlerImport tests the implementation of the
1781 # butler command line interface "import" subcommand. Functions
1782 # in the script folder are generally considered protected and
1783 # should not be used as public api.
1784 with open(exportFile) as f:
1785 script.butlerImport(
1786 importDir,
1787 export_file=f,
1788 directory=exportDir,
1789 transfer="auto",
1790 skip_dimensions=None,
1791 )
1792 importButler = Butler.from_config(importDir, run=self.default_run)
1793 for ref in datasets:
1794 with self.subTest(ref=ref):
1795 # Test for existence by passing in the DatasetType and
1796 # data ID separately, to avoid lookup by dataset_id.
1797 self.assertTrue(importButler.exists(ref.datasetType, ref.dataId))
1798 self.assertEqual(
1799 list(importButler.registry.queryDimensionRecords("skymap")),
1800 [importButler.dimensions["skymap"].RecordClass(**skymapRecord)],
1801 )
1803 def testRemoveRuns(self) -> None:
1804 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1805 butler = self.create_empty_butler(writeable=True)
1806 # Load registry data with dimensions to hang datasets off of.
1807 butler.import_(filename=ResourcePath("resource://lsst.daf.butler/tests/registry_data/base.yaml"))
1808 # Add some RUN-type collection.
1809 run1 = "run1"
1810 butler.collections.register(run1)
1811 run2 = "run2"
1812 butler.collections.register(run2)
1813 # put a dataset in each
1814 metric = makeExampleMetrics()
1815 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1816 datasetType = self.addDatasetType(
1817 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1818 )
1819 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1820 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1821 uri1 = butler.getURI(ref1)
1822 uri2 = butler.getURI(ref2)
1824 # Put one of the runs in a chain.
1825 butler.collections.register("Chain", CollectionType.CHAINED)
1826 butler.collections.extend_chain("Chain", run1)
1828 with self.assertRaises(OrphanedRecordError):
1829 butler.registry.removeDatasetType(datasetType.name)
1831 # Remove a non-run.
1832 with self.assertRaises(TypeError):
1833 butler.removeRuns(["Chain"])
1835 # Remove without unlinking from chain should fail.
1836 with self.assertRaises(IntegrityError):
1837 butler.removeRuns([run1])
1839 # Remove from both runs. No longer use unstore parameter since it
1840 # always purges.
1841 butler.removeRuns([run1, run2], unlink_from_chains=True)
1843 # Should be nothing in registry for either one, and datastore should
1844 # not think either exists.
1845 with self.assertRaises(MissingCollectionError):
1846 butler.collections.get_info(run1)
1847 with self.assertRaises(MissingCollectionError):
1848 butler.collections.get_info(run1)
1849 self.assertFalse(butler.stored(ref1))
1850 self.assertFalse(butler.stored(ref2))
1851 # We always unstore so both URIs should be gone.
1852 self.assertFalse(uri1.exists())
1853 self.assertFalse(uri2.exists())
1855 # Now that the collections have been pruned we can remove the
1856 # dataset type
1857 butler.registry.removeDatasetType(datasetType.name)
1859 with self.assertLogs("lsst.daf.butler.registry", "INFO") as cm:
1860 butler.registry.removeDatasetType(("test*", "test*"))
1861 self.assertIn("not defined", "\n".join(cm.output))
1863 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
1864 """Simulate an external actor removing a file outside of Butler's
1865 knowledge.
1867 Subclasses may override to handle more complicated datastore
1868 configurations.
1869 """
1870 uri = butler.getURI(ref)
1871 uri.remove()
1872 datastore = cast(FileDatastore, butler._datastore)
1873 datastore.cacheManager.remove_from_cache(ref)
1875 def testPruneDatasets(self) -> None:
1876 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1877 butler = self.create_empty_butler(writeable=True)
1878 # Load registry data with dimensions to hang datasets off of.
1879 butler.import_(filename=_get_test_data_path("base.yaml"))
1880 # Add some RUN-type collections.
1881 run1 = "run1"
1882 butler.collections.register(run1)
1883 run2 = "run2"
1884 butler.collections.register(run2)
1885 # put some datasets. ref1 and ref2 have the same data ID, and are in
1886 # different runs. ref3 has a different data ID.
1887 metric = makeExampleMetrics()
1888 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1889 datasetType = self.addDatasetType(
1890 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1891 )
1892 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1893 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1894 ref3 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, run=run1)
1896 many_stored = butler.stored_many([ref1, ref2, ref3])
1897 for ref, stored in many_stored.items():
1898 self.assertTrue(stored, f"Ref {ref} should be stored")
1900 many_exists = butler._exists_many([ref1, ref2, ref3])
1901 for ref, exists in many_exists.items():
1902 self.assertTrue(exists, f"Checking ref {ref} exists.")
1903 self.assertEqual(exists, DatasetExistence.VERIFIED, f"Ref {ref} should be stored")
1905 # Simple prune.
1906 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
1907 self.assertFalse(butler.exists(ref1.datasetType, ref1.dataId, collections=run1))
1909 many_stored = butler.stored_many([ref1, ref2, ref3])
1910 for ref, stored in many_stored.items():
1911 self.assertFalse(stored, f"Ref {ref} should not be stored")
1913 many_exists = butler._exists_many([ref1, ref2, ref3])
1914 for ref, exists in many_exists.items():
1915 self.assertEqual(exists, DatasetExistence.UNRECOGNIZED, f"Ref {ref} should not be stored")
1917 # Put data back.
1918 ref1_new = butler.put(metric, ref1)
1919 self.assertEqual(ref1_new, ref1) # Reuses original ID.
1920 ref2 = butler.put(metric, ref2)
1922 many_stored = butler.stored_many([ref1, ref2, ref3])
1923 self.assertTrue(many_stored[ref1])
1924 self.assertTrue(many_stored[ref2])
1925 self.assertFalse(many_stored[ref3])
1927 ref3 = butler.put(metric, ref3)
1929 many_exists = butler._exists_many([ref1, ref2, ref3])
1930 for ref, exists in many_exists.items():
1931 self.assertTrue(exists, f"Ref {ref} should not be stored")
1933 # Clear out the datasets from registry and start again.
1934 refs = [ref1, ref2, ref3]
1935 butler.pruneDatasets(refs, purge=True, unstore=True)
1936 for ref in refs:
1937 butler.put(metric, ref)
1939 # Confirm we can retrieve deferred.
1940 dref1 = butler.getDeferred(ref1) # known and exists
1941 metric1 = dref1.get()
1942 self.assertEqual(metric1, metric)
1944 # Test different forms of file availability.
1945 # Need to be in a state where:
1946 # - one ref just has registry record.
1947 # - one ref has a missing file but a datastore record.
1948 # - one ref has a missing datastore record but file is there.
1949 # - one ref does not exist anywhere.
1950 # Do not need to test a ref that has everything since that is tested
1951 # above.
1952 ref0 = DatasetRef(
1953 datasetType,
1954 DataCoordinate.standardize(
1955 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, universe=butler.dimensions
1956 ),
1957 run=run1,
1958 )
1960 # Delete from datastore and retain in Registry.
1961 butler.pruneDatasets([ref1], purge=False, unstore=True, disassociate=False)
1963 # File has been removed.
1964 self.remove_dataset_out_of_band(butler, ref2)
1966 # Datastore has lost track.
1967 butler._datastore.forget([ref3])
1969 # First test with a standard butler.
1970 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
1971 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
1972 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
1973 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
1974 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED)
1976 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=False)
1977 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
1978 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
1979 self.assertEqual(exists_many[ref2], DatasetExistence.KNOWN)
1980 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
1981 self.assertTrue(exists_many[ref2])
1983 # Check that per-ref query gives the same answer as many query.
1984 for ref, exists in exists_many.items():
1985 self.assertEqual(butler.exists(ref, full_check=False), exists)
1987 # Get deferred checks for existence before it allows it to be
1988 # retrieved.
1989 with self.assertRaises(LookupError):
1990 butler.getDeferred(ref3) # not known, file exists
1991 dref2 = butler.getDeferred(ref2) # known but file missing
1992 with self.assertRaises(FileNotFoundError):
1993 dref2.get()
1995 # Test again with a trusting butler.
1996 if self.trustModeSupported:
1997 butler._datastore.trustGetRequest = True
1998 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
1999 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2000 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2001 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2002 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ARTIFACT)
2004 # When trusting we can get a deferred dataset handle that is not
2005 # known but does exist.
2006 dref3 = butler.getDeferred(ref3)
2007 metric3 = dref3.get()
2008 self.assertEqual(metric3, metric)
2010 # Check that per-ref query gives the same answer as many query.
2011 for ref, exists in exists_many.items():
2012 self.assertEqual(butler.exists(ref, full_check=True), exists)
2014 # Create a ref that surprisingly has the UUID of an existing ref
2015 # but is not the same.
2016 ref_bad = DatasetRef(datasetType, dataId=ref3.dataId, run=ref3.run, id=ref2.id)
2017 with self.assertRaises(ValueError):
2018 butler.exists(ref_bad)
2020 # Create a ref that has a compatible storage class.
2021 ref_compat = ref2.overrideStorageClass("StructuredDataDict")
2022 exists = butler.exists(ref_compat)
2023 self.assertEqual(exists, exists_many[ref2])
2025 # Remove everything and start from scratch.
2026 butler._datastore.trustGetRequest = False
2027 butler.pruneDatasets(refs, purge=True, unstore=True)
2028 for ref in refs:
2029 butler.put(metric, ref)
2031 # These tests mess directly with the trash table and can leave the
2032 # datastore in an odd state. Do them at the end.
2033 # Check that in normal mode, deleting the record will lead to
2034 # trash not touching the file.
2035 uri1 = butler.getURI(ref1)
2036 butler._datastore.bridge.moveToTrash(
2037 [ref1], transaction=None
2038 ) # Update the dataset_location table
2039 butler._datastore.forget([ref1])
2040 butler._datastore.trash(ref1)
2041 butler._datastore.emptyTrash()
2042 self.assertTrue(uri1.exists())
2043 uri1.remove() # Clean it up.
2045 # Simulate execution butler setup by deleting the datastore
2046 # record but keeping the file around and trusting.
2047 butler._datastore.trustGetRequest = True
2048 uris = butler.get_many_uris([ref2, ref3])
2049 uri2 = uris[ref2].primaryURI
2050 uri3 = uris[ref3].primaryURI
2051 self.assertTrue(uri2.exists())
2052 self.assertTrue(uri3.exists())
2054 # Remove the datastore record.
2055 butler._datastore.bridge.moveToTrash(
2056 [ref2], transaction=None
2057 ) # Update the dataset_location table
2058 butler._datastore.forget([ref2])
2059 self.assertTrue(uri2.exists())
2060 butler._datastore.trash([ref2, ref3])
2061 # Immediate removal for ref2 file
2062 self.assertFalse(uri2.exists())
2063 # But ref3 has to wait for the empty.
2064 self.assertTrue(uri3.exists())
2065 butler._datastore.emptyTrash()
2066 self.assertFalse(uri3.exists())
2068 # Clear out the datasets from registry.
2069 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
2071 def test_butler_metrics(self):
2072 """Test that metrics are collected."""
2073 run = "test_run"
2074 metrics = ButlerMetrics()
2075 butler, datasetType = self.create_butler(
2076 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics
2077 )
2078 data = MetricsExampleModel(
2079 summary={"AM1": 5.2, "AM2": 30.6},
2080 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2081 data=[563, 234, 456.7, 752, 8, 9, 27],
2082 )
2084 data_ref = butler.put(data, datasetType, visit=424, instrument="DummyCamComp")
2085 butler.get(data_ref)
2086 butler.get(data_ref)
2087 self.assertEqual(metrics.n_get, 2)
2088 self.assertGreater(metrics.time_in_get, 0.0)
2089 self.assertEqual(metrics.n_put, 1)
2090 self.assertGreater(metrics.time_in_put, 0.0)
2092 deferred = butler.getDeferred(data_ref)
2093 deferred.get()
2094 self.assertEqual(metrics.n_get, 3)
2096 with butler.record_metrics() as new:
2097 butler.put(data, datasetType, visit=425, instrument="DummyCamComp")
2098 butler.get(data_ref)
2099 self.assertEqual(new.n_get, 1)
2100 self.assertEqual(new.n_put, 1)
2103class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2104 """PosixDatastore specialization of a butler"""
2106 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2107 fullConfigKey: str | None = ".datastore.formatters"
2108 validationCanFail = True
2109 datastoreStr = ["/tmp"]
2110 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2111 registryStr = "/gen3.sqlite3"
2113 def testPathConstructor(self) -> None:
2114 """Independent test of constructor using PathLike."""
2115 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
2116 self.assertIsInstance(butler, Butler)
2118 # And again with a Path object with the butler yaml
2119 path = pathlib.Path(self.tmpConfigFile)
2120 butler = Butler.from_config(path, writeable=False)
2121 self.assertIsInstance(butler, Butler)
2123 # And again with a Path object without the butler yaml
2124 # (making sure we skip it if the tmp config doesn't end
2125 # in butler.yaml -- which is the case for a subclass)
2126 if self.tmpConfigFile.endswith("butler.yaml"):
2127 path = pathlib.Path(os.path.dirname(self.tmpConfigFile))
2128 butler = Butler.from_config(path, writeable=False)
2129 self.assertIsInstance(butler, Butler)
2131 def testExportTransferCopy(self) -> None:
2132 """Test local export using all transfer modes"""
2133 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2134 exportButler = self.runPutGetTest(storageClass, "test_metric")
2135 # Test that the repo actually has at least one dataset.
2136 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
2137 self.assertGreater(len(datasets), 0)
2138 uris = [exportButler.getURI(d) for d in datasets]
2139 assert isinstance(exportButler._datastore, FileDatastore)
2140 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]]
2142 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris]
2144 for path in pathsInStore:
2145 # Assume local file system
2146 assert path is not None
2147 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}")
2149 for transfer in ("copy", "link", "symlink", "relsymlink"):
2150 with safeTestTempDir(TESTDIR) as exportDir:
2151 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export:
2152 export.saveDatasets(datasets)
2153 for path in pathsInStore:
2154 assert path is not None
2155 self.assertTrue(
2156 self.checkFileExists(exportDir, path),
2157 f"Check that mode {transfer} exported files",
2158 )
2160 def testPytypeCoercion(self) -> None:
2161 """Test python type coercion on Butler.get and put."""
2162 # Store some data with the normal example storage class.
2163 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2164 datasetTypeName = "test_metric"
2165 butler = self.runPutGetTest(storageClass, datasetTypeName)
2167 dataId = {"instrument": "DummyCamComp", "visit": 423}
2168 metric = butler.get(datasetTypeName, dataId=dataId)
2169 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
2171 datasetType_ori = butler.get_dataset_type(datasetTypeName)
2172 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
2174 # Now need to hack the registry dataset type definition.
2175 # There is no API for this.
2176 assert isinstance(butler._registry, SqlRegistry)
2177 manager = butler._registry._managers.datasets
2178 assert hasattr(manager, "_db") and hasattr(manager, "_static")
2179 manager._db.update(
2180 manager._static.dataset_type,
2181 {"name": datasetTypeName},
2182 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
2183 )
2185 # Force reset of dataset type cache
2186 butler.registry.refresh()
2188 datasetType_new = butler.get_dataset_type(datasetTypeName)
2189 self.assertEqual(datasetType_new.name, datasetType_ori.name)
2190 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
2192 metric_model = butler.get(datasetTypeName, dataId=dataId)
2193 self.assertNotEqual(type(metric_model), type(metric))
2194 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
2196 # Put the model and read it back to show that everything now
2197 # works as normal.
2198 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
2199 metric_model_new = butler.get(metric_ref)
2200 self.assertEqual(metric_model_new, metric_model)
2202 # Hack the storage class again to something that will fail on the
2203 # get with no conversion class.
2204 manager._db.update(
2205 manager._static.dataset_type,
2206 {"name": datasetTypeName},
2207 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
2208 )
2209 butler.registry.refresh()
2211 with self.assertRaises(ValueError):
2212 butler.get(datasetTypeName, dataId=dataId)
2214 def test_provenance(self):
2215 """Test that provenance is attached on put."""
2216 run = "test_run"
2217 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
2218 metric = MetricsExampleModel(
2219 summary={"AM1": 5.2, "AM2": 30.6},
2220 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2221 data=[563, 234, 456.7, 752, 8, 9, 27],
2222 )
2223 # Provenance can be attached to the object being put. Whether
2224 # it is or not is dependent on the formatter. For this test we
2225 # copy on adding provenance to ensure they differ.
2226 self.assertIsNone(metric.dataset_id)
2227 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
2228 self.assertIsNone(metric.dataset_id)
2229 metric_2 = butler.get(metric_ref)
2230 self.assertEqual(metric_2.data, metric.data)
2231 self.assertEqual(metric_2.dataset_id, metric_ref.id)
2232 self.assertIsNone(metric_2.provenance)
2234 # Put with provenance.
2235 prov = DatasetProvenance(quantum_id=uuid.uuid4())
2236 prov.add_input(metric_ref)
2237 prov.add_extra_provenance(metric_ref.id, {"answer": 42})
2238 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
2239 metric_3 = butler.get(metric_ref2)
2240 self.assertEqual(metric_3.provenance, prov)
2242 # Check that we can extract provenance from dict form.
2243 prov_dict = prov.to_flat_dict(metric_ref2)
2244 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2245 self.assertEqual(ref_from_prov, metric_ref2)
2246 # Direct __eq__ of the provenance does not work because one side
2247 # includes dimension records.
2248 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2249 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id)
2250 self.assertEqual(prov_from_prov.extras, prov.extras)
2252 # Force a bad ID into the dict.
2253 prov_dict["id"] = uuid.uuid4()
2254 with self.assertRaises(ValueError):
2255 DatasetProvenance.from_flat_dict(prov_dict, butler)
2256 del prov_dict["id"]
2257 prov_dict["input 0 id"] = uuid.uuid4()
2258 with self.assertRaises(ValueError):
2259 DatasetProvenance.from_flat_dict(prov_dict, butler)
2261 # Check that simple types can be reconstructed with non-standard
2262 # separators.
2263 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True)
2264 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2265 self.assertEqual(ref_from_prov, metric_ref2)
2266 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2268 with self.assertRaises(ValueError):
2269 DatasetProvenance.from_flat_dict({"unknown": 42}, butler)
2271 def test_specialized_file_datasets_functions(self):
2272 """Test a workflow used in Prompt Processing where we export datasets
2273 from one repository and write them in-place to the datastore of
2274 another, without immediately inserting registry entries for the
2275 datasets.
2276 """
2277 repo = MetricTestRepo.create_from_butler(
2278 self.create_empty_butler(writeable=True),
2279 self.tmpConfigFile,
2280 "StructuredCompositeReadCompNoDisassembly",
2281 )
2282 source_butler = repo.butler
2284 # Test writing outputs to a FileDatastore.
2285 with tempfile.TemporaryDirectory() as tempdir:
2286 target_repo_config = Butler.makeRepo(tempdir)
2287 refs = [repo.ref1, repo.ref2]
2288 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2289 self.assertEqual(len(datasets), 2)
2290 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2291 for dataset in datasets:
2292 path = ResourcePath(dataset.path, forceAbsolute=False)
2293 # Paths should be relative paths to the target datastore.
2294 self.assertFalse(path.isabs())
2295 # Files should have been copied into the target datastore
2296 self.assertTrue(ResourcePath(tempdir).join(path).exists())
2298 # Make sure the target Butler can ingest the datasets.
2299 target_butler = Butler(target_repo_config, writeable=True)
2300 target_butler.transfer_dimension_records_from(source_butler, refs)
2301 target_butler.ingest(*datasets, transfer=None)
2302 self.assertIsNotNone(target_butler.get(repo.ref1))
2303 self.assertIsNotNone(target_butler.get(repo.ref2))
2305 # Giving an empty list of files is a no-op.
2306 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), [])
2307 self.assertEqual(len(no_datasets), 0)
2309 # Test writing outputs to a ChainedDatastore.
2310 with tempfile.TemporaryDirectory() as tempdir:
2311 # Set up a second dataset type, so we can split the files across
2312 # multiple datastore roots.
2313 dt1 = repo.datasetType
2314 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass)
2315 source_butler.registry.registerDatasetType(dt2)
2316 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2)
2317 config = Config.fromString(
2318 f"""
2319 datastore:
2320 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore
2321 datastore_constraints:
2322 - constraints:
2323 accept:
2324 - {dt1.name}
2325 - constraints:
2326 accept:
2327 - {dt2.name}
2328 datastores:
2329 - datastore:
2330 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2331 root: <butlerRoot>/FileDatastore_0
2332 - datastore:
2333 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2334 root: <butlerRoot>/FileDatastore_1
2335 """
2336 )
2337 target_repo_config = Butler.makeRepo(tempdir, config)
2338 refs = [repo.ref1, repo.ref2, other_ref]
2339 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2340 self.assertEqual(len(datasets), 3)
2341 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2342 for dataset in datasets:
2343 path = ResourcePath(dataset.path, forceAbsolute=False)
2344 # Paths should be relative paths to the target datastore.
2345 self.assertFalse(path.isabs())
2346 # Files should have been split up between the two datastores
2347 # in the chain.
2348 datastore_root = ResourcePath(tempdir)
2349 if dataset.refs[0].datasetType.name == dt1.name:
2350 datastore_root = datastore_root.join("FileDatastore_0")
2351 else:
2352 datastore_root = datastore_root.join("FileDatastore_1")
2353 self.assertTrue(datastore_root.join(path).exists())
2355 # Make sure the target Butler can ingest the datasets.
2356 target_butler = Butler(target_repo_config, writeable=True)
2357 target_butler.transfer_dimension_records_from(source_butler, refs)
2358 target_butler.ingest(*datasets, transfer=None)
2359 self.assertIsNotNone(target_butler.get(repo.ref1))
2360 self.assertIsNotNone(target_butler.get(repo.ref2))
2361 self.assertIsNotNone(target_butler.get(other_ref))
2364class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2365 """PosixDatastore specialization of a butler using Postgres"""
2367 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2368 fullConfigKey = ".datastore.formatters"
2369 validationCanFail = True
2370 datastoreStr = ["/tmp"]
2371 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2372 registryStr = "PostgreSQL@test"
2374 @classmethod
2375 def setUpClass(cls) -> None:
2376 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
2377 super().setUpClass()
2379 def setUp(self) -> None:
2380 # Need to add a registry section to the config.
2381 self._temp_config = False
2382 config = Config(self.configFile)
2383 self.postgresql.patch_butler_config(config)
2384 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
2385 config.dump(fh)
2386 self.configFile = fh.name
2387 self._temp_config = True
2388 super().setUp()
2390 def tearDown(self) -> None:
2391 if self._temp_config and os.path.exists(self.configFile):
2392 os.remove(self.configFile)
2393 super().tearDown()
2395 def testMakeRepo(self) -> None:
2396 # The base class test assumes that it's using sqlite and assumes
2397 # the config file is acceptable to sqlite.
2398 raise unittest.SkipTest("Postgres config is not compatible with this test.")
2401class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase):
2402 """Test that Butler with a Postgres registry still works after cloning."""
2404 def create_butler(
2405 self,
2406 run: str,
2407 storageClass: StorageClass | str,
2408 datasetTypeName: str,
2409 metrics: ButlerMetrics | None = None,
2410 ) -> tuple[DirectButler, DatasetType]:
2411 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2412 return butler.clone(run=run, metrics=metrics), datasetType
2415class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
2416 """InMemoryDatastore specialization of a butler"""
2418 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")
2419 fullConfigKey = None
2420 useTempRoot = False
2421 validationCanFail = False
2422 datastoreStr = ["datastore='InMemory"]
2423 datastoreName = ["InMemoryDatastore@"]
2424 registryStr = "/gen3.sqlite3"
2426 def testIngest(self) -> None:
2427 pass
2429 def test_ingest_zip(self) -> None:
2430 pass
2433class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase):
2434 """Test that a Butler with a Sqlite registry still works after cloning."""
2436 def create_butler(
2437 self,
2438 run: str,
2439 storageClass: StorageClass | str,
2440 datasetTypeName: str,
2441 metrics: ButlerMetrics | None = None,
2442 ) -> tuple[DirectButler, DatasetType]:
2443 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2444 return butler.clone(run=run), datasetType
2447class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2448 """PosixDatastore specialization"""
2450 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
2451 fullConfigKey = ".datastore.datastores.1.formatters"
2452 validationCanFail = True
2453 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"]
2454 datastoreName = [
2455 "InMemoryDatastore@",
2456 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1",
2457 "SecondDatastore",
2458 ]
2459 registryStr = "/gen3.sqlite3"
2461 def testPruneDatasets(self) -> None:
2462 # This test relies on manipulating files out-of-band, which is
2463 # impossible for this configuration because of the InMemoryDatastore in
2464 # the ChainedDatastore.
2465 pass
2468class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase):
2469 """Test that a yaml file in one location can refer to a root in another."""
2471 datastoreStr = ["dir1"]
2472 # Disable the makeRepo test since we are deliberately not using
2473 # butler.yaml as the config name.
2474 fullConfigKey = None
2476 def setUp(self) -> None:
2477 self.root = makeTestTempDir(TESTDIR)
2479 # Make a new repository in one place
2480 self.dir1 = os.path.join(self.root, "dir1")
2481 Butler.makeRepo(self.dir1, config=Config(self.configFile))
2483 # Move the yaml file to a different place and add a "root"
2484 self.dir2 = os.path.join(self.root, "dir2")
2485 os.makedirs(self.dir2, exist_ok=True)
2486 configFile1 = os.path.join(self.dir1, "butler.yaml")
2487 config = Config(configFile1)
2488 config["root"] = self.dir1
2489 configFile2 = os.path.join(self.dir2, "butler2.yaml")
2490 config.dumpToUri(configFile2)
2491 os.remove(configFile1)
2492 self.tmpConfigFile = configFile2
2494 def testFileLocations(self) -> None:
2495 self.assertNotEqual(self.dir1, self.dir2)
2496 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml")))
2497 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml")))
2498 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3")))
2501class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase):
2502 """Test that a config file created by makeRepo outside of repo works."""
2504 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2506 def setUp(self) -> None:
2507 self.root = makeTestTempDir(TESTDIR)
2508 self.root2 = makeTestTempDir(TESTDIR)
2510 self.tmpConfigFile = os.path.join(self.root2, "different.yaml")
2511 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2513 def tearDown(self) -> None:
2514 if os.path.exists(self.root2):
2515 shutil.rmtree(self.root2, ignore_errors=True)
2516 super().tearDown()
2518 def testConfigExistence(self) -> None:
2519 c = Config(self.tmpConfigFile)
2520 uri_config = ResourcePath(c["root"])
2521 uri_expected = ResourcePath(self.root, forceDirectory=True)
2522 self.assertEqual(uri_config.geturl(), uri_expected.geturl())
2523 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path")
2525 def testPutGet(self) -> None:
2526 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2527 self.runPutGetTest(storageClass, "test_metric")
2530class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase):
2531 """Test that a config file created by makeRepo outside of repo works."""
2533 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2535 def setUp(self) -> None:
2536 self.root = makeTestTempDir(TESTDIR)
2537 self.root2 = makeTestTempDir(TESTDIR)
2539 self.tmpConfigFile = self.root2
2540 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2542 def testConfigExistence(self) -> None:
2543 # Append the yaml file else Config constructor does not know the file
2544 # type.
2545 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml")
2546 super().testConfigExistence()
2549class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase):
2550 """Test that a config file created by makeRepo outside of repo works."""
2552 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2554 def setUp(self) -> None:
2555 self.root = makeTestTempDir(TESTDIR)
2556 self.root2 = makeTestTempDir(TESTDIR)
2558 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl()
2559 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2562@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!")
2563class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2564 """S3Datastore specialization of a butler; an S3 storage Datastore +
2565 a local in-memory SqlRegistry.
2566 """
2568 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")
2569 fullConfigKey = None
2570 validationCanFail = True
2572 bucketName = "anybucketname"
2573 """Name of the Bucket that will be used in the tests. The name is read from
2574 the config file used with the tests during set-up.
2575 """
2577 root = "butlerRoot/"
2578 """Root repository directory expected to be used in case useTempRoot=False.
2579 Otherwise the root is set to a 20 characters long randomly generated string
2580 during set-up.
2581 """
2583 datastoreStr = [f"datastore={root}"]
2584 """Contains all expected root locations in a format expected to be
2585 returned by Butler stringification.
2586 """
2588 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]
2589 """The expected format of the S3 Datastore string."""
2591 registryStr = "/gen3.sqlite3"
2592 """Expected format of the Registry string."""
2594 mock_aws = mock_aws()
2595 """The mocked s3 interface from moto."""
2597 def genRoot(self) -> str:
2598 """Return a random string of len 20 to serve as a root
2599 name for the temporary bucket repo.
2601 This is equivalent to tempfile.mkdtemp as this is what self.root
2602 becomes when useTempRoot is True.
2603 """
2604 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20))
2605 return rndstr + "/"
2607 def setUp(self) -> None:
2608 config = Config(self.configFile)
2609 uri = ResourcePath(config[".datastore.datastore.root"])
2610 self.bucketName = uri.netloc
2612 # Enable S3 mocking of tests.
2613 self.enterContext(clean_test_environment_for_s3())
2614 self.mock_aws.start()
2616 if self.useTempRoot:
2617 self.root = self.genRoot()
2618 rooturi = f"s3://{self.bucketName}/{self.root}"
2619 config.update({"datastore": {"datastore": {"root": rooturi}}})
2621 # need local folder to store registry database
2622 self.reg_dir = makeTestTempDir(TESTDIR)
2623 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"
2625 # MOTO needs to know that we expect Bucket bucketname to exist
2626 # (this used to be the class attribute bucketName)
2627 s3 = boto3.resource("s3")
2628 s3.create_bucket(Bucket=self.bucketName)
2630 self.datastoreStr = [f"datastore='{rooturi}'"]
2631 self.datastoreName = [f"FileDatastore@{rooturi}"]
2632 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)
2633 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")
2635 def tearDown(self) -> None:
2636 s3 = boto3.resource("s3")
2637 bucket = s3.Bucket(self.bucketName)
2638 try:
2639 bucket.objects.all().delete()
2640 except botocore.exceptions.ClientError as e:
2641 if e.response["Error"]["Code"] == "404":
2642 # the key was not reachable - pass
2643 pass
2644 else:
2645 raise
2647 bucket = s3.Bucket(self.bucketName)
2648 bucket.delete()
2650 # Stop the S3 mock.
2651 self.mock_aws.stop()
2653 if self.reg_dir is not None and os.path.exists(self.reg_dir):
2654 shutil.rmtree(self.reg_dir, ignore_errors=True)
2656 if self.useTempRoot and os.path.exists(self.root):
2657 shutil.rmtree(self.root, ignore_errors=True)
2659 super().tearDown()
2662class DatastoreTransfers(TestCaseMixin):
2663 """Base test setup for data transfers between butlers. The concrete tests
2664 for specific configurations are in other classes, below.
2665 """
2667 storageClassFactory: StorageClassFactory
2669 @classmethod
2670 def setUpClass(cls) -> None:
2671 cls.storageClassFactory = StorageClassFactory()
2673 def setUp(self) -> None:
2674 self.root = makeTestTempDir(TESTDIR)
2675 self.config = Config(self.configFile)
2677 # Some tests cause convertors to be replaced so ensure
2678 # the storage class factory is reset each time.
2679 self.storageClassFactory.reset()
2680 self.storageClassFactory.addFromConfig(self.configFile)
2682 def tearDown(self) -> None:
2683 removeTestTempDir(self.root)
2685 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler:
2686 if manager is None:
2687 manager = (
2688 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
2689 )
2690 config = Config(config_file if config_file is not None else self.configFile)
2691 config["registry", "managers", "datasets"] = manager
2692 return Butler.from_config(
2693 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True
2694 )
2696 def assertButlerTransfers(
2697 self,
2698 purge: bool = False,
2699 storageClassName: str = "StructuredData",
2700 storageClassNameTarget: str | None = None,
2701 ) -> None:
2702 """Test that a run can be transferred to another butler."""
2703 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
2704 if storageClassNameTarget is not None:
2705 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget)
2706 else:
2707 storageClassTarget = storageClass
2709 datasetTypeName = "random_data"
2711 # Test will create 3 collections and we will want to transfer
2712 # two of those three.
2713 runs = ["run1", "run2", "other"]
2715 # Also want to use two different dataset types to ensure that
2716 # grouping works.
2717 datasetTypeNames = ["random_data", "random_data_2"]
2719 # Create the run collections in the source butler.
2720 for run in runs:
2721 self.source_butler.collections.register(run)
2723 # Create dimensions in source butler.
2724 n_exposures = 30
2725 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
2726 self.source_butler.registry.insertDimensionData(
2727 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
2728 )
2729 self.source_butler.registry.insertDimensionData(
2730 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
2731 )
2732 self.source_butler.registry.insertDimensionData(
2733 "day_obs",
2734 {
2735 "instrument": "DummyCamComp",
2736 "id": 20250101,
2737 },
2738 )
2740 for i in range(n_exposures):
2741 self.source_butler.registry.insertDimensionData(
2742 "group", {"instrument": "DummyCamComp", "name": f"group{i}"}
2743 )
2744 self.source_butler.registry.insertDimensionData(
2745 "exposure",
2746 {
2747 "instrument": "DummyCamComp",
2748 "id": i,
2749 "obs_id": f"exp{i}",
2750 "physical_filter": "d-r",
2751 "group": f"group{i}",
2752 "day_obs": 20250101,
2753 },
2754 )
2756 # Create dataset types in the source butler.
2757 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"])
2758 for datasetTypeName in datasetTypeNames:
2759 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
2760 self.source_butler.registry.registerDatasetType(datasetType)
2762 # Write a dataset to an unrelated run -- this will ensure that
2763 # we are rewriting integer dataset ids in the target if necessary.
2764 # Will not be relevant for UUID.
2765 run = "distraction"
2766 butler = Butler.from_config(butler=self.source_butler, run=run)
2767 butler.put(
2768 makeExampleMetrics(),
2769 datasetTypeName,
2770 exposure=1,
2771 instrument="DummyCamComp",
2772 physical_filter="d-r",
2773 )
2775 # Write some example metrics to the source
2776 butler = Butler.from_config(butler=self.source_butler)
2778 # Set of DatasetRefs that should be in the list of refs to transfer
2779 # but which will not be transferred.
2780 deleted: set[DatasetRef] = set()
2782 n_expected = 20 # Number of datasets expected to be transferred
2783 source_refs = []
2784 for i in range(n_exposures):
2785 # Put a third of datasets into each collection, only retain
2786 # two thirds.
2787 index = i % 3
2788 run = runs[index]
2789 datasetTypeName = datasetTypeNames[i % 2]
2791 metric = MetricsExample(
2792 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)]
2793 )
2794 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"}
2795 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run)
2797 # Remove the datastore record using low-level API, but only
2798 # for a specific index.
2799 if purge and index == 1:
2800 # For one of these delete the file as well.
2801 # This allows the "missing" code to filter the
2802 # file out.
2803 # Access the individual datastores.
2804 datastores = []
2805 if hasattr(butler._datastore, "datastores"):
2806 datastores.extend(butler._datastore.datastores)
2807 else:
2808 datastores.append(butler._datastore)
2810 if not deleted:
2811 # For a chained datastore we need to remove
2812 # files in each chain.
2813 for datastore in datastores:
2814 # The file might not be known to the datastore
2815 # if constraints are used.
2816 try:
2817 primary, uris = datastore.getURIs(ref)
2818 except FileNotFoundError:
2819 continue
2820 if primary and primary.scheme != "mem":
2821 primary.remove()
2822 for uri in uris.values():
2823 if uri.scheme != "mem":
2824 uri.remove()
2825 n_expected -= 1
2826 deleted.add(ref)
2828 # Remove the datastore record.
2829 for datastore in datastores:
2830 if hasattr(datastore, "removeStoredItemInfo"):
2831 datastore.removeStoredItemInfo(ref)
2833 if index < 2:
2834 source_refs.append(ref)
2835 if ref not in deleted:
2836 new_metric = butler.get(ref)
2837 self.assertEqual(new_metric, metric)
2839 # Create some bad dataset types to ensure we check for inconsistent
2840 # definitions.
2841 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList")
2842 for datasetTypeName in datasetTypeNames:
2843 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass)
2844 self.target_butler.registry.registerDatasetType(datasetType)
2845 with self.assertRaises(ConflictingDefinitionError) as cm:
2846 self.target_butler.transfer_from(self.source_butler, source_refs)
2847 self.assertIn("dataset type differs", str(cm.exception))
2849 # And remove the bad definitions.
2850 for datasetTypeName in datasetTypeNames:
2851 self.target_butler.registry.removeDatasetType(datasetTypeName)
2853 # Transfer without creating dataset types should fail.
2854 with self.assertRaises(KeyError):
2855 self.target_butler.transfer_from(self.source_butler, source_refs)
2857 # Transfer without creating dimensions should fail.
2858 with self.assertRaises(ConflictingDefinitionError) as cm:
2859 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True)
2860 self.assertIn("dimension", str(cm.exception))
2862 # The dry run test requires dataset types to exist. If we have
2863 # been given distinct storage classes for the target we have
2864 # to redefine at least one of the dataset types in the target butler.
2865 if storageClass != storageClassTarget:
2866 self.target_butler.registry.removeDatasetType(datasetTypeNames[0])
2867 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget)
2868 self.target_butler.registry.registerDatasetType(datasetType)
2870 # The failed transfer above leaves registry in an inconsistent
2871 # state because the run is created but then rolled back without
2872 # the collection cache being cleared. For now force a refresh.
2873 # Can remove with DM-35498.
2874 self.target_butler.registry.refresh()
2876 # Do a dry run -- this should not have any effect on the target butler.
2877 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True)
2879 # Transfer the records for one ref to test the alternative API.
2880 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2881 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]])
2882 self.assertIn("number of records transferred: 1", ";".join(log_cm.output))
2884 # Now transfer them to the second butler, including dimensions.
2885 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2886 transferred = self.target_butler.transfer_from(
2887 self.source_butler,
2888 source_refs,
2889 register_dataset_types=True,
2890 transfer_dimensions=True,
2891 )
2892 self.assertEqual(len(transferred), n_expected)
2893 log_output = ";".join(log_cm.output)
2895 # A ChainedDatastore will use the in-memory datastore for mexists
2896 # so we can not rely on the mexists log message.
2897 self.assertIn("Number of datastore records found in source", log_output)
2898 self.assertIn("Creating output run", log_output)
2900 # Do the transfer twice to ensure that it will do nothing extra.
2901 # Only do this if purge=True because it does not work for int
2902 # dataset_id.
2903 if purge:
2904 # This should not need to register dataset types.
2905 transferred = self.target_butler.transfer_from(self.source_butler, source_refs)
2906 self.assertEqual(len(transferred), n_expected)
2908 with self.assertRaises((TypeError, AttributeError)):
2909 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore
2911 with self.assertRaises(ValueError):
2912 self.target_butler._datastore.transfer_from(
2913 self.source_butler._datastore, source_refs, transfer="split"
2914 )
2916 # Now try to get the same refs from the new butler.
2917 for ref in source_refs:
2918 if ref not in deleted:
2919 new_metric = self.target_butler.get(ref)
2920 old_metric = self.source_butler.get(ref)
2921 self.assertEqual(new_metric, old_metric)
2923 # Try again without implicit storage class conversion
2924 # triggered by using the source ref. This will do conversion
2925 # since the formatter will be returning the source python type.
2926 target_ref = self.target_butler.get_dataset(ref.id)
2927 if target_ref.datasetType.storageClass != ref.datasetType.storageClass:
2928 new_metric = self.target_butler.get(target_ref)
2929 self.assertNotEqual(type(new_metric), type(old_metric))
2931 # Remove the dataset from the target and put it again
2932 # as if it was the right type all along for this butler.
2933 self.target_butler.pruneDatasets(
2934 [target_ref], unstore=True, purge=True, disassociate=True
2935 )
2936 self.target_butler.put(new_metric, target_ref)
2937 new_new_metric = self.target_butler.get(target_ref)
2938 new_old_metric = self.target_butler.get(
2939 target_ref, storageClass=ref.datasetType.storageClass
2940 )
2941 self.assertEqual(new_new_metric, new_metric)
2942 self.assertEqual(new_old_metric, old_metric)
2944 # Now prune run2 collection and create instead a CHAINED collection.
2945 # This should block the transfer.
2946 self.target_butler.removeRuns(["run2"])
2947 self.target_butler.collections.register("run2", CollectionType.CHAINED)
2948 with self.assertRaises(CollectionTypeError):
2949 # Re-importing the run1 datasets can be problematic if they
2950 # use integer IDs so filter those out.
2951 to_transfer = [ref for ref in source_refs if ref.run == "run2"]
2952 self.target_butler.transfer_from(self.source_butler, to_transfer)
2955class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
2956 """Test data transfers between butlers.
2958 Test for different managers. UUID to UUID and integer to integer are
2959 tested. UUID to integer is not supported since we do not currently
2960 want to allow that. Integer to UUID is supported with the caveat
2961 that UUID4 will be generated and this will be incorrect for raw
2962 dataset types. The test ignores that.
2963 """
2965 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2967 def create_butlers(
2968 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None
2969 ) -> None:
2970 self.source_butler = self.create_butler(manager1, "1", config_file=source_config)
2971 self.target_butler = self.create_butler(manager2, "2")
2973 def testTransferUuidToUuid(self) -> None:
2974 self.create_butlers()
2975 self.assertButlerTransfers()
2977 def testTransferFromChainedUuidToUuid(self) -> None:
2978 """Force the source butler to be a ChainedDatastore."""
2979 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml"))
2980 self.assertButlerTransfers()
2982 def testTransferFromIncompatibleUuidToUuid(self) -> None:
2983 """Force the source butler to be a incompatible datastore."""
2984 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml"))
2985 with self.assertRaises(NotImplementedError):
2986 self.assertButlerTransfers()
2988 def testTransferFromIncompatibleChainUuidToUuid(self) -> None:
2989 """Force the source butler to be a incompatible datastore."""
2990 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml"))
2991 with self.assertRaises(TypeError):
2992 self.assertButlerTransfers()
2994 def testTransferFromFileUuidToUuid(self) -> None:
2995 """Force the source butler to be a FileDatastore."""
2996 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml"))
2997 self.assertButlerTransfers()
2999 def testTransferMissing(self) -> None:
3000 """Test transfers where datastore records are missing.
3002 This is how execution butler works.
3003 """
3004 self.create_butlers()
3006 # Configure the source butler to allow trust.
3007 self.source_butler._datastore._set_trust_mode(True)
3009 self.assertButlerTransfers(purge=True)
3011 def testTransferMissingDisassembly(self) -> None:
3012 """Test transfers where datastore records are missing.
3014 This is how execution butler works.
3015 """
3016 self.create_butlers()
3018 # Configure the source butler to allow trust.
3019 self.source_butler._datastore._set_trust_mode(True)
3021 # Test disassembly.
3022 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite")
3024 def testTransferDifferingStorageClasses(self) -> None:
3025 """Test transfers when the source butler dataset type has a different
3026 but compatible storage class.
3027 """
3028 self.create_butlers()
3030 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion")
3032 def testTransferDifferingStorageClassesDisassembly(self) -> None:
3033 """Test transfers when the source butler dataset type has a different
3034 but compatible storage class and where the source butler has
3035 disassembled.
3036 """
3037 self.create_butlers()
3039 self.assertButlerTransfers(
3040 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion"
3041 )
3043 def testAbsoluteURITransferDirect(self) -> None:
3044 """Test transfer using an absolute URI."""
3045 self._absolute_transfer("auto")
3047 def testAbsoluteURITransferCopy(self) -> None:
3048 """Test transfer using an absolute URI."""
3049 self._absolute_transfer("copy")
3051 def _absolute_transfer(self, transfer: str) -> None:
3052 self.create_butlers()
3054 storageClassName = "StructuredData"
3055 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
3056 datasetTypeName = "random_data"
3057 run = "run1"
3058 self.source_butler.collections.register(run)
3060 dimensions = self.source_butler.dimensions.conform(())
3061 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3062 self.source_butler.registry.registerDatasetType(datasetType)
3064 metrics = makeExampleMetrics()
3065 with ResourcePath.temporary_uri(suffix=".json") as temp:
3066 dataId = DataCoordinate.make_empty(self.source_butler.dimensions)
3067 source_refs = [DatasetRef(datasetType, dataId, run=run)]
3068 temp.write(json.dumps(metrics.exportAsDict()).encode())
3069 dataset = FileDataset(path=temp, refs=source_refs)
3070 self.source_butler.ingest(dataset, transfer="direct")
3072 self.target_butler.transfer_from(
3073 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer
3074 )
3076 uri = self.target_butler.getURI(dataset.refs[0])
3077 if transfer == "auto":
3078 self.assertEqual(uri, temp)
3079 else:
3080 self.assertNotEqual(uri, temp)
3082 def test_shared_dimension_group(self):
3083 """Test internal logic that divides dataset types by dimension group
3084 when doing registry updates.
3085 """
3086 self.create_butlers()
3087 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
3088 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
3090 source_butler = self.source_butler
3091 target_butler = self.target_butler
3093 # Create a dataset type with the same dimensions as the 'bias' dataset
3094 # type from base.yaml
3095 dataset_type = DatasetType(
3096 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
3097 )
3098 source_butler.registry.registerDatasetType(dataset_type)
3099 # This has the same data ID as one of the bias datasets in
3100 # datasets.yaml.
3101 test_ref = source_butler.registry.insertDatasets(
3102 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
3103 )[0]
3105 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
3106 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
3107 refs = [test_ref, *biases, *flats]
3109 # Test setup will be even more convoluted if we want the datastore to
3110 # actually transfer files. For testing the dimension group behavior,
3111 # we really only care about the registry.
3112 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
3113 mock.return_value = (set(refs), set())
3114 target_butler.transfer_from(
3115 source_butler,
3116 refs,
3117 transfer=None,
3118 register_dataset_types=True,
3119 skip_missing=False,
3120 transfer_dimensions=True,
3121 )
3123 transferred_test_ref = target_butler.find_dataset(
3124 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3125 )
3126 self.assertEqual(transferred_test_ref.id, test_ref.id)
3128 transferred_bias = target_butler.find_dataset(
3129 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3130 )
3131 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
3133 transferred_flat = target_butler.find_dataset(
3134 "flat",
3135 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
3136 collections="imported_r",
3137 )
3138 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
3141class ChainedDatastoreTransfers(PosixDatastoreTransfers):
3142 """Test transfers using a chained datastore."""
3144 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
3147@unittest.skipIf(create_test_server is None, "Server dependencies not installed.")
3148class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3149 """Test ``transfer_from`` involving Butler server."""
3151 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3153 def test_transfers_from_remote_to_direct(self) -> None:
3154 from lsst.daf.butler.remote_butler._remote_file_transfer_source import (
3155 mock_file_transfer_uris_for_unit_test,
3156 )
3158 self.target_butler = self.create_butler(None, "2")
3159 with create_test_server(TESTDIR) as server:
3160 self.source_butler = server.hybrid_butler
3162 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath:
3163 # The Butler server returns HTTP URIs with a domain name that
3164 # is not resolvable because there is no actual HTTP server
3165 # involved in these tests. Strip this first layer of
3166 # indirection, and return the target of the redirect instead.
3167 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers)
3168 return ResourcePath(str(response.next_request.url))
3170 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url):
3171 self.assertButlerTransfers()
3174class NullDatastoreTestCase(unittest.TestCase):
3175 """Test that we can fall back to a null datastore."""
3177 # Need a good config to create the repo.
3178 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3179 storageClassFactory: StorageClassFactory
3181 @classmethod
3182 def setUpClass(cls) -> None:
3183 cls.storageClassFactory = StorageClassFactory()
3184 cls.storageClassFactory.addFromConfig(cls.configFile)
3186 def setUp(self) -> None:
3187 """Create a new butler root for each test."""
3188 self.root = makeTestTempDir(TESTDIR)
3189 Butler.makeRepo(self.root, config=Config(self.configFile))
3191 def tearDown(self) -> None:
3192 removeTestTempDir(self.root)
3194 def test_fallback(self) -> None:
3195 # Read the butler config and mess with the datastore section.
3196 config_path = os.path.join(self.root, "butler.yaml")
3197 bad_config = Config(config_path)
3198 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore"
3199 bad_config.dumpToUri(config_path)
3201 with self.assertRaises(RuntimeError):
3202 Butler(self.root, without_datastore=False)
3204 with self.assertRaises(RuntimeError):
3205 Butler.from_config(self.root, without_datastore=False)
3207 butler = Butler.from_config(self.root, writeable=True, without_datastore=True)
3208 self.assertIsInstance(butler._datastore, NullDatastore)
3210 # Check that registry is working.
3211 butler.collections.register("MYRUN")
3212 collections = butler.collections.query("*")
3213 self.assertIn("MYRUN", set(collections))
3215 # Create a ref.
3216 dimensions = butler.dimensions.conform([])
3217 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
3218 datasetTypeName = "metric"
3219 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3220 butler.registry.registerDatasetType(datasetType)
3221 ref = DatasetRef(datasetType, {}, run="MYRUN")
3223 # Check that datastore will complain.
3224 with self.assertRaises(FileNotFoundError):
3225 butler.get(ref)
3226 with self.assertRaises(FileNotFoundError):
3227 butler.getURI(ref)
3230@unittest.skipIf(create_test_server is None, "Server dependencies not installed.")
3231class ButlerServerTests(FileDatastoreButlerTests):
3232 """Test RemoteButler and Butler server."""
3234 configFile = None
3235 predictionSupported = False
3236 trustModeSupported = False
3238 postgres: TemporaryPostgresInstance | None
3240 def setUp(self):
3241 self.server_instance = self.enterContext(create_test_server(TESTDIR))
3243 def tearDown(self):
3244 pass
3246 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
3247 # S3 pre-signed URLs may end up with differing expiration times in the
3248 # query parameters, so ignore query parameters when comparing.
3249 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path
3251 def create_empty_butler(
3252 self, run: str | None = None, writeable: bool | None = None, metrics: ButlerMetrics | None = None
3253 ) -> Butler:
3254 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics)
3256 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
3257 # Can't delete a file via S3 signed URLs, so we need to reach in
3258 # through DirectButler to delete the dataset.
3259 uri = self.server_instance.direct_butler.getURI(ref)
3260 uri.remove()
3262 def testConstructor(self):
3263 # RemoteButler constructor is tested in test_server.py and
3264 # test_remote_butler.py.
3265 pass
3267 def testDafButlerRepositories(self):
3268 # Loading of RemoteButler via repository index is tested in
3269 # test_server.py.
3270 pass
3272 def testGetDatasetTypes(self) -> None:
3273 # This is mostly a test of validateConfiguration, which is for
3274 # validating Datastore configuration and thus isn't relevant to
3275 # RemoteButler.
3276 pass
3278 def testMakeRepo(self) -> None:
3279 # Only applies to DirectButler.
3280 pass
3282 # Pickling not yet implemented for RemoteButler/HybridButler.
3283 @unittest.expectedFailure
3284 def testPickle(self) -> None:
3285 return super().testPickle()
3287 def testStringification(self) -> None:
3288 self.assertEqual(
3289 str(self.server_instance.remote_butler),
3290 "RemoteButler(https://test.example/api/butler/repo/testrepo/)",
3291 )
3293 def testTransaction(self) -> None:
3294 # Transactions will never be supported for RemoteButler.
3295 pass
3297 def testPutTemplates(self) -> None:
3298 # The Butler server instance is configured with different file naming
3299 # templates than this test is expecting.
3300 pass
3303@unittest.skipIf(create_test_server is None, "Server dependencies not installed.")
3304class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase):
3305 """Tests for RemoteButler's registry shim, with a SQLite DB backing the
3306 server.
3307 """
3309 postgres = None
3312@unittest.skipIf(create_test_server is None, "Server dependencies not installed.")
3313class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase):
3314 """Tests for RemoteButler's registry shim, with a Postgres DB backing the
3315 server.
3316 """
3318 @classmethod
3319 def setUpClass(cls):
3320 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
3321 super().setUpClass()
3324def setup_module(module: types.ModuleType) -> None:
3325 """Set up the module for pytest."""
3326 clean_environment()
3329def _get_test_data_path(filename: str) -> ResourcePath:
3330 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}")
3333if __name__ == "__main__":
3334 clean_environment()
3335 unittest.main()