Coverage for tests / test_butler.py: 13%
1786 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-13 01:20 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""Tests for Butler."""
30from __future__ import annotations
32import json
33import logging
34import os
35import pathlib
36import pickle
37import posixpath
38import random
39import re
40import shutil
41import string
42import tempfile
43import unittest
44import unittest.mock
45import uuid
46import warnings
47import weakref
48from collections.abc import Callable, Mapping
49from typing import TYPE_CHECKING, Any, cast
51try:
52 import boto3
53 import botocore
55 from lsst.resources.s3utils import clean_test_environment_for_s3
57 try:
58 from moto import mock_aws # v5
59 except ImportError:
60 from moto import mock_s3 as mock_aws
61except ImportError:
62 boto3 = None
64 def mock_aws(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-untyped-def]
65 """No-op decorator in case moto mock_aws can not be imported."""
66 return None
69import astropy.time
70from sqlalchemy.exc import IntegrityError
72from lsst.daf.butler import (
73 Butler,
74 ButlerConfig,
75 ButlerMetrics,
76 ButlerRepoIndex,
77 CollectionCycleError,
78 CollectionType,
79 Config,
80 DataCoordinate,
81 DatasetExistence,
82 DatasetNotFoundError,
83 DatasetProvenance,
84 DatasetRef,
85 DatasetType,
86 FileDataset,
87 NoDefaultCollectionError,
88 StorageClassFactory,
89 ValidationError,
90 script,
91)
92from lsst.daf.butler._rubin.file_datasets import transfer_datasets_to_datastore
93from lsst.daf.butler._rubin.temporary_for_ingest import TemporaryForIngest
94from lsst.daf.butler.datastore import NullDatastore
95from lsst.daf.butler.datastore.file_templates import FileTemplate, FileTemplateValidationError
96from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex
97from lsst.daf.butler.datastores.fileDatastore import FileDatastore
98from lsst.daf.butler.direct_butler import DirectButler
99from lsst.daf.butler.registry import (
100 CollectionError,
101 CollectionTypeError,
102 ConflictingDefinitionError,
103 DataIdValueError,
104 DatasetTypeExpressionError,
105 MissingCollectionError,
106 OrphanedRecordError,
107)
108from lsst.daf.butler.registry.sql_registry import SqlRegistry
109from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG
110from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter
111from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db
112from lsst.daf.butler.tests.server_available import butler_server_import_error, butler_server_is_available
113from lsst.daf.butler.tests.utils import (
114 MetricTestRepo,
115 TestCaseMixin,
116 makeTestTempDir,
117 removeTestTempDir,
118 safeTestTempDir,
119)
120from lsst.resources import ResourcePath
121from lsst.resources.http import HttpResourcePath
122from lsst.utils import doImportType
123from lsst.utils.introspection import get_full_type_name
125if butler_server_is_available:
126 from lsst.daf.butler.tests.server import create_test_server
129if TYPE_CHECKING:
130 import types
132 from lsst.daf.butler import DimensionGroup, Registry, StorageClass
134TESTDIR = os.path.abspath(os.path.dirname(__file__))
137def clean_environment() -> None:
138 """Remove external environment variables that affect the tests."""
139 for k in ("DAF_BUTLER_REPOSITORY_INDEX",):
140 os.environ.pop(k, None)
143def makeExampleMetrics() -> MetricsExample:
144 """Return example dataset suitable for tests."""
145 return MetricsExample(
146 {"AM1": 5.2, "AM2": 30.6},
147 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
148 [563, 234, 456.7, 752, 8, 9, 27],
149 )
152class TransactionTestError(Exception):
153 """Specific error for testing transactions, to prevent misdiagnosing
154 that might otherwise occur when a standard exception is used.
155 """
157 pass
160class ButlerConfigTests(unittest.TestCase):
161 """Simple tests for ButlerConfig that are not tested in any other test
162 cases.
163 """
165 def testSearchPath(self) -> None:
166 configFile = os.path.join(TESTDIR, "config", "basic", "butler.yaml")
167 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
168 config1 = ButlerConfig(configFile)
169 self.assertNotIn("testConfigs", "\n".join(cm.output))
171 overrideDirectory = os.path.join(TESTDIR, "config", "testConfigs")
172 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
173 config2 = ButlerConfig(configFile, searchPaths=[overrideDirectory])
174 self.assertIn("testConfigs", "\n".join(cm.output))
176 key = ("datastore", "records", "table")
177 self.assertNotEqual(config1[key], config2[key])
178 self.assertEqual(config2[key], "override_record")
181class ButlerPutGetTests(TestCaseMixin):
182 """Helper method for running a suite of put/get tests from different
183 butler configurations.
184 """
186 root: str
187 default_run = "ingésτ😺"
188 storageClassFactory: StorageClassFactory
189 configFile: str | None
190 tmpConfigFile: str
192 @staticmethod
193 def addDatasetType(
194 datasetTypeName: str, dimensions: DimensionGroup, storageClass: StorageClass | str, registry: Registry
195 ) -> DatasetType:
196 """Create a DatasetType and register it"""
197 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
198 registry.registerDatasetType(datasetType)
199 return datasetType
201 @classmethod
202 def setUpClass(cls) -> None:
203 cls.storageClassFactory = StorageClassFactory()
204 if cls.configFile is not None:
205 cls.storageClassFactory.addFromConfig(cls.configFile)
207 def assertGetComponents(
208 self,
209 butler: Butler,
210 datasetRef: DatasetRef,
211 components: tuple[str, ...],
212 reference: Any,
213 collections: Any = None,
214 ) -> None:
215 datasetType = datasetRef.datasetType
216 dataId = datasetRef.dataId
217 deferred = butler.getDeferred(datasetRef)
219 for component in components:
220 compTypeName = datasetType.componentTypeName(component)
221 result = butler.get(compTypeName, dataId, collections=collections)
222 self.assertEqual(result, getattr(reference, component))
223 result_deferred = deferred.get(component=component)
224 self.assertEqual(result_deferred, result)
226 def tearDown(self) -> None:
227 if self.root is not None:
228 removeTestTempDir(self.root)
230 def create_empty_butler(
231 self,
232 run: str | None = None,
233 writeable: bool | None = None,
234 metrics: ButlerMetrics | None = None,
235 cleanup: bool = True,
236 ):
237 """Create a Butler for the test repository, without inserting test
238 data.
239 """
240 butler = Butler.from_config(self.tmpConfigFile, run=run, writeable=writeable, metrics=metrics)
241 if cleanup:
242 self.enterContext(butler)
243 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
244 return butler
246 def create_butler(
247 self,
248 run: str,
249 storageClass: StorageClass | str,
250 datasetTypeName: str,
251 metrics: ButlerMetrics | None = None,
252 ) -> tuple[Butler, DatasetType]:
253 """Create a Butler for the test repository and insert some test data
254 into it.
255 """
256 butler = self.create_empty_butler(run=run, metrics=metrics)
258 collections = set(butler.collections.query("*"))
259 self.assertEqual(collections, {run})
260 # Create and register a DatasetType
261 dimensions = butler.dimensions.conform(["instrument", "visit"])
263 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
265 # Add needed Dimensions
266 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
267 butler.registry.insertDimensionData(
268 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
269 )
270 butler.registry.insertDimensionData(
271 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
272 )
273 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101})
274 visit_start = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
275 visit_end = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
276 butler.registry.insertDimensionData(
277 "visit",
278 {
279 "instrument": "DummyCamComp",
280 "id": 423,
281 "name": "fourtwentythree",
282 "physical_filter": "d-r",
283 "datetime_begin": visit_start,
284 "datetime_end": visit_end,
285 "day_obs": 20200101,
286 },
287 )
289 # Add more visits for some later tests
290 for visit_id in (424, 425):
291 butler.registry.insertDimensionData(
292 "visit",
293 {
294 "instrument": "DummyCamComp",
295 "id": visit_id,
296 "name": f"fourtwentyfour_{visit_id}",
297 "physical_filter": "d-r",
298 "day_obs": 20200101,
299 },
300 )
301 return butler, datasetType
303 def runPutGetTest(self, storageClass: StorageClass, datasetTypeName: str) -> Butler:
304 # New datasets will be added to run and tag, but we will only look in
305 # tag when looking up datasets.
306 run = self.default_run
307 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
308 assert butler.run is not None
310 # Create and store a dataset
311 metric = makeExampleMetrics()
312 dataId = butler.registry.expandDataId({"instrument": "DummyCamComp", "visit": 423})
314 # Dataset should not exist if we haven't added it
315 with self.assertRaises(DatasetNotFoundError):
316 butler.get(datasetTypeName, dataId)
318 # Put and remove the dataset once as a DatasetRef, once as a dataId,
319 # and once with a DatasetType
321 # Keep track of any collections we add and do not clean up
322 expected_collections = {run}
324 counter = 0
325 ref = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run="put_run_1")
326 args = tuple[DatasetRef] | tuple[str | DatasetType, DataCoordinate]
327 for args in ((ref,), (datasetTypeName, dataId), (datasetType, dataId)):
328 # Since we are using subTest we can get cascading failures
329 # here with the first attempt failing and the others failing
330 # immediately because the dataset already exists. Work around
331 # this by using a distinct run collection each time
332 counter += 1
333 this_run = f"put_run_{counter}"
334 butler.collections.register(this_run)
335 expected_collections.update({this_run})
337 with self.subTest(args=repr(args)):
338 kwargs: dict[str, Any] = {}
339 if not isinstance(args[0], DatasetRef): # type: ignore
340 kwargs["run"] = this_run
341 ref = butler.put(metric, *args, **kwargs)
342 self.assertIsInstance(ref, DatasetRef)
344 # Test get of a ref.
345 metricOut = butler.get(ref)
346 self.assertEqual(metric, metricOut)
347 # Test get
348 metricOut = butler.get(ref.datasetType.name, dataId, collections=this_run)
349 self.assertEqual(metric, metricOut)
350 # Test get with a datasetRef
351 metricOut = butler.get(ref)
352 self.assertEqual(metric, metricOut)
353 # Test getDeferred with dataId
354 metricOut = butler.getDeferred(ref.datasetType.name, dataId, collections=this_run).get()
355 self.assertEqual(metric, metricOut)
356 # Test getDeferred with a ref
357 metricOut = butler.getDeferred(ref).get()
358 self.assertEqual(metric, metricOut)
360 # Check we can get components
361 if storageClass.isComposite():
362 self.assertGetComponents(
363 butler, ref, ("summary", "data", "output"), metric, collections=this_run
364 )
366 primary_uri, secondary_uris = butler.getURIs(ref)
367 n_uris = len(secondary_uris)
368 if primary_uri:
369 n_uris += 1
371 # Can the artifacts themselves be retrieved?
372 if not butler._datastore.isEphemeral:
373 # Create a temporary directory to hold the retrieved
374 # artifacts.
375 with tempfile.TemporaryDirectory(
376 prefix="butler-artifacts-", ignore_cleanup_errors=True
377 ) as artifact_root:
378 root_uri = ResourcePath(artifact_root, forceDirectory=True)
380 for preserve_path in (True, False):
381 destination = root_uri.join(f"{preserve_path}_{counter}/")
382 log = logging.getLogger("lsst.x")
383 log.debug("Using destination %s for args %s", destination, args)
384 # Use copy so that we can test that overwrite
385 # protection works (using "auto" for File URIs
386 # would use hard links and subsequent transfer
387 # would work because it knows they are the same
388 # file).
389 transferred = butler.retrieveArtifacts(
390 [ref], destination, preserve_path=preserve_path, transfer="copy"
391 )
392 self.assertGreater(len(transferred), 0)
393 artifacts = list(ResourcePath.findFileResources([destination]))
394 # Filter out the index file.
395 artifacts = [a for a in artifacts if a.basename() != ZipIndex.index_name]
396 self.assertEqual(set(transferred), set(artifacts))
398 for artifact in transferred:
399 path_in_destination = artifact.relative_to(destination)
400 self.assertIsNotNone(path_in_destination)
401 assert path_in_destination is not None
403 # When path is not preserved there should not
404 # be any path separators.
405 num_seps = path_in_destination.count("/")
406 if preserve_path:
407 self.assertGreater(num_seps, 0)
408 else:
409 self.assertEqual(num_seps, 0)
411 self.assertEqual(
412 len(artifacts),
413 n_uris,
414 "Comparing expected artifacts vs actual:"
415 f" {artifacts} vs {primary_uri} and {secondary_uris}",
416 )
418 if preserve_path:
419 # No need to run these twice
420 with self.assertRaises(ValueError):
421 butler.retrieveArtifacts([ref], destination, transfer="move")
423 with self.assertRaisesRegex(
424 ValueError, "^Destination location must refer to a directory"
425 ):
426 butler.retrieveArtifacts(
427 [ref], ResourcePath("/some/file.txt", forceDirectory=False)
428 )
430 with self.assertRaises(FileExistsError):
431 butler.retrieveArtifacts([ref], destination)
433 transferred_again = butler.retrieveArtifacts(
434 [ref], destination, preserve_path=preserve_path, overwrite=True
435 )
436 self.assertEqual(set(transferred_again), set(transferred))
438 # Now remove the dataset completely.
439 butler.pruneDatasets([ref], purge=True, unstore=True)
440 # Lookup with original args should still fail.
441 kwargs = {"collections": this_run}
442 if isinstance(args[0], DatasetRef):
443 kwargs = {} # Prevent warning from being issued.
444 self.assertFalse(butler.exists(*args, **kwargs))
445 # get() should still fail.
446 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
447 butler.get(ref)
448 # Registry shouldn't be able to find it by dataset_id anymore.
449 self.assertIsNone(butler.get_dataset(ref.id))
451 # Do explicit registry removal since we know they are
452 # empty
453 butler.collections.x_remove(this_run)
454 expected_collections.remove(this_run)
456 # Create DatasetRef for put using default run.
457 refIn = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run=butler.run)
459 # Check that getDeferred fails with standalone ref.
460 with self.assertRaises(LookupError):
461 butler.getDeferred(refIn)
463 # Put the dataset again, since the last thing we did was remove it
464 # and we want to use the default collection.
465 ref = butler.put(metric, refIn)
467 # Get with parameters
468 stop = 4
469 sliced = butler.get(ref, parameters={"slice": slice(stop)})
470 self.assertNotEqual(metric, sliced)
471 self.assertEqual(metric.summary, sliced.summary)
472 self.assertEqual(metric.output, sliced.output)
473 assert metric.data is not None # for mypy
474 self.assertEqual(metric.data[:stop], sliced.data)
475 # getDeferred with parameters
476 sliced = butler.getDeferred(ref, parameters={"slice": slice(stop)}).get()
477 self.assertNotEqual(metric, sliced)
478 self.assertEqual(metric.summary, sliced.summary)
479 self.assertEqual(metric.output, sliced.output)
480 self.assertEqual(metric.data[:stop], sliced.data)
481 # getDeferred with deferred parameters
482 sliced = butler.getDeferred(ref).get(parameters={"slice": slice(stop)})
483 self.assertNotEqual(metric, sliced)
484 self.assertEqual(metric.summary, sliced.summary)
485 self.assertEqual(metric.output, sliced.output)
486 self.assertEqual(metric.data[:stop], sliced.data)
488 if storageClass.isComposite():
489 # Check that components can be retrieved
490 metricOut = butler.get(ref.datasetType.name, dataId)
491 compNameS = ref.datasetType.componentTypeName("summary")
492 compNameD = ref.datasetType.componentTypeName("data")
493 summary = butler.get(compNameS, dataId)
494 self.assertEqual(summary, metric.summary)
495 data = butler.get(compNameD, dataId)
496 self.assertEqual(data, metric.data)
498 if "counter" in storageClass.derivedComponents:
499 count = butler.get(ref.datasetType.componentTypeName("counter"), dataId)
500 self.assertEqual(count, len(data))
502 count = butler.get(
503 ref.datasetType.componentTypeName("counter"), dataId, parameters={"slice": slice(stop)}
504 )
505 self.assertEqual(count, stop)
507 compRef = butler.find_dataset(compNameS, dataId, collections=butler.collections.defaults)
508 assert compRef is not None
509 summary = butler.get(compRef)
510 self.assertEqual(summary, metric.summary)
512 # Create a Dataset type that has the same name but is inconsistent.
513 inconsistentDatasetType = DatasetType(
514 datasetTypeName, datasetType.dimensions, self.storageClassFactory.getStorageClass("Config")
515 )
517 # Getting with a dataset type that does not match registry fails
518 with self.assertRaisesRegex(
519 ValueError,
520 "(Supplied dataset type .* inconsistent with registry)"
521 "|(The new storage class .* is not compatible with the existing storage class)",
522 ):
523 butler.get(inconsistentDatasetType, dataId)
525 # Combining a DatasetRef with a dataId should fail
526 with self.assertRaisesRegex(ValueError, "DatasetRef given, cannot use dataId as well"):
527 butler.get(ref, dataId)
528 # Getting with an explicit ref should fail if the id doesn't match.
529 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
530 butler.get(DatasetRef(ref.datasetType, ref.dataId, id=uuid.UUID(int=101), run=butler.run))
532 # Getting a dataset with unknown parameters should fail
533 with self.assertRaisesRegex(KeyError, "Parameter 'unsupported' not understood"):
534 butler.get(ref, parameters={"unsupported": True})
536 # Check we have a collection
537 collections = set(butler.collections.query("*"))
538 self.assertEqual(collections, expected_collections)
540 # Clean up to check that we can remove something that may have
541 # already had a component removed
542 butler.pruneDatasets([ref], unstore=True, purge=True)
544 # Add the same ref again, so we can check that duplicate put fails.
545 ref = butler.put(metric, datasetType, dataId)
547 # Repeat put will fail.
548 with self.assertRaisesRegex(
549 ConflictingDefinitionError, "A database constraint failure was triggered"
550 ):
551 butler.put(metric, datasetType, dataId)
553 # Remove the datastore entry.
554 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
556 # Put will still fail
557 with self.assertRaisesRegex(
558 ConflictingDefinitionError, "A database constraint failure was triggered"
559 ):
560 butler.put(metric, datasetType, dataId)
562 # Repeat the same sequence with resolved ref.
563 butler.pruneDatasets([ref], unstore=True, purge=True)
564 ref = butler.put(metric, refIn)
566 # Repeat put will fail.
567 with self.assertRaisesRegex(ConflictingDefinitionError, "Datastore already contains dataset"):
568 butler.put(metric, refIn)
570 # Remove the datastore entry.
571 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
573 # In case of resolved ref this write will succeed.
574 ref = butler.put(metric, refIn)
576 # Leave the dataset in place since some downstream tests require
577 # something to be present
579 return butler
581 def testDeferredCollectionPassing(self) -> None:
582 # Construct a butler with no run or collection, but make it writeable.
583 butler = self.create_empty_butler(writeable=True)
584 # Create and register a DatasetType
585 dimensions = butler.dimensions.conform(["instrument", "visit"])
586 datasetType = self.addDatasetType(
587 "example", dimensions, self.storageClassFactory.getStorageClass("StructuredData"), butler.registry
588 )
589 # Add needed Dimensions
590 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
591 butler.registry.insertDimensionData(
592 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
593 )
594 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
595 butler.registry.insertDimensionData(
596 "visit",
597 {
598 "instrument": "DummyCamComp",
599 "id": 423,
600 "name": "fourtwentythree",
601 "physical_filter": "d-r",
602 "day_obs": 20250101,
603 },
604 )
605 dataId = {"instrument": "DummyCamComp", "visit": 423}
606 # Create dataset.
607 metric = makeExampleMetrics()
608 # Register a new run and put dataset.
609 run = "deferred"
610 self.assertTrue(butler.collections.register(run))
611 # Second time it will be allowed but indicate no-op
612 self.assertFalse(butler.collections.register(run))
613 ref = butler.put(metric, datasetType, dataId, run=run)
614 # Putting with no run should fail with TypeError.
615 with self.assertRaises(CollectionError):
616 butler.put(metric, datasetType, dataId)
617 # Dataset should exist.
618 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
619 # We should be able to get the dataset back, but with and without
620 # a deferred dataset handle.
621 self.assertEqual(metric, butler.get(datasetType, dataId, collections=[run]))
622 self.assertEqual(metric, butler.getDeferred(datasetType, dataId, collections=[run]).get())
623 # Trying to find the dataset without any collection is an error.
624 with self.assertRaises(NoDefaultCollectionError):
625 butler.exists(datasetType, dataId)
626 with self.assertRaises(CollectionError):
627 butler.get(datasetType, dataId)
628 # Associate the dataset with a different collection.
629 butler.collections.register("tagged", type=CollectionType.TAGGED)
630 butler.registry.associate("tagged", [ref])
631 # Deleting the dataset from the new collection should make it findable
632 # in the original collection.
633 butler.pruneDatasets([ref], tags=["tagged"])
634 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
637class ButlerTests(ButlerPutGetTests):
638 """Tests for Butler."""
640 useTempRoot = True
641 validationCanFail: bool
642 fullConfigKey: str | None
643 registryStr: str | None
644 datastoreName: list[str] | None
645 datastoreStr: list[str]
646 predictionSupported = True
647 """Does getURIs support 'prediction mode'?"""
649 def setUp(self) -> None:
650 """Create a new butler root for each test."""
651 self.root = makeTestTempDir(TESTDIR)
652 Butler.makeRepo(self.root, config=Config(self.configFile))
653 self.tmpConfigFile = os.path.join(self.root, "butler.yaml")
655 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
656 """Return True if two URIs refer to the same resource.
658 Subclasses may override to handle unique requirements.
659 """
660 return uri1 == uri2
662 def testConstructor(self) -> None:
663 """Independent test of constructor."""
664 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
665 self.enterContext(butler)
666 self.assertIsInstance(butler, Butler)
668 # Check that butler.yaml is added automatically.
669 if self.tmpConfigFile.endswith(end := "/butler.yaml"):
670 config_dir = self.tmpConfigFile[: -len(end)]
671 butler = Butler.from_config(config_dir, run=self.default_run)
672 self.enterContext(butler)
673 self.assertIsInstance(butler, Butler)
675 # Even with a ResourcePath.
676 butler = Butler.from_config(ResourcePath(config_dir, forceDirectory=True), run=self.default_run)
677 self.enterContext(butler)
678 self.assertIsInstance(butler, Butler)
680 collections = set(butler.collections.query("*"))
681 self.assertEqual(collections, {self.default_run})
683 # Check that some special characters can be included in run name.
684 special_run = "u@b.c-A"
685 butler_special = Butler.from_config(butler=butler, run=special_run)
686 self.enterContext(butler_special)
687 collections = set(butler_special.registry.queryCollections("*@*"))
688 self.assertEqual(collections, {special_run})
690 butler2 = Butler.from_config(butler=butler, collections=["other"])
691 self.enterContext(butler2)
692 self.assertEqual(butler2.collections.defaults, ("other",))
693 self.assertIsNone(butler2.run)
694 self.assertEqual(type(butler._datastore), type(butler2._datastore))
695 self.assertEqual(butler._datastore.config, butler2._datastore.config)
697 # Test that we can use an environment variable to find this
698 # repository.
699 butler_index = Config()
700 butler_index["label"] = self.tmpConfigFile
701 for suffix in (".yaml", ".json"):
702 # Ensure that the content differs so that we know that
703 # we aren't reusing the cache.
704 bad_label = f"file://bucket/not_real{suffix}"
705 butler_index["bad_label"] = bad_label
706 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
707 butler_index.dumpToUri(temp_file)
708 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
709 self.assertEqual(Butler.get_known_repos(), {"label", "bad_label"})
710 uri = Butler.get_repo_uri("bad_label")
711 self.assertEqual(uri, ResourcePath(bad_label))
712 uri = Butler.get_repo_uri("label")
713 butler = Butler.from_config(uri, writeable=False)
714 self.assertIsInstance(butler, Butler)
715 butler.close()
716 butler = Butler.from_config("label", writeable=False)
717 self.assertIsInstance(butler, Butler)
718 butler.close()
719 with self.assertRaisesRegex(FileNotFoundError, "aliases:.*bad_label"):
720 Butler.from_config("not_there", writeable=False)
721 with self.assertRaisesRegex(FileNotFoundError, "resolved from alias 'bad_label'"):
722 Butler.from_config("bad_label")
723 with self.assertRaises(FileNotFoundError):
724 # Should ignore aliases.
725 Butler.from_config(ResourcePath("label", forceAbsolute=False))
726 with self.assertRaises(KeyError) as cm:
727 Butler.get_repo_uri("missing")
728 self.assertEqual(
729 Butler.get_repo_uri("missing", True), ResourcePath("missing", forceAbsolute=False)
730 )
731 self.assertIn("not known to", str(cm.exception))
732 # Should report no failure.
733 self.assertEqual(ButlerRepoIndex.get_failure_reason(), "")
734 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
735 # Now with empty configuration.
736 butler_index = Config()
737 butler_index.dumpToUri(temp_file)
738 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
739 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases)"):
740 Butler.from_config("label")
741 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
742 # Now with bad contents.
743 with open(temp_file.ospath, "w") as fh:
744 print("'", file=fh)
745 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
746 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases:.*could not be read)"):
747 Butler.from_config("label")
748 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": "file://not_found/x.yaml"}):
749 with self.assertRaises(FileNotFoundError):
750 Butler.get_repo_uri("label")
751 self.assertEqual(Butler.get_known_repos(), set())
753 with self.assertRaisesRegex(FileNotFoundError, "index file not found"):
754 Butler.from_config("label")
756 # Check that we can create Butler when the alias file is not found.
757 butler = Butler.from_config(self.tmpConfigFile, writeable=False)
758 self.enterContext(butler)
759 self.assertIsInstance(butler, Butler)
760 with self.assertRaises(RuntimeError) as cm:
761 # No environment variable set.
762 Butler.get_repo_uri("label")
763 self.assertEqual(Butler.get_repo_uri("label", True), ResourcePath("label", forceAbsolute=False))
764 self.assertIn("No repository index defined", str(cm.exception))
765 with self.assertRaisesRegex(FileNotFoundError, "no known aliases.*No repository index"):
766 # No aliases registered.
767 Butler.from_config("not_there")
768 self.assertEqual(Butler.get_known_repos(), set())
770 def testClose(self):
771 butler = self.create_empty_butler(cleanup=False)
772 is_direct_butler = isinstance(butler, DirectButler)
773 if is_direct_butler:
774 self.assertFalse(butler._closed)
776 with butler as butler_from_context_manager:
777 self.assertIs(butler, butler_from_context_manager)
778 if is_direct_butler:
779 self.assertTrue(butler._closed)
780 with self.assertRaisesRegex(RuntimeError, "has been closed"):
781 butler.get_dataset_type("raw")
783 # Close may be called multiple times.
784 butler.close()
785 if is_direct_butler:
786 self.assertTrue(butler._closed)
788 def testGarbageCollection(self):
789 """Test that Butler does not have any circular references that prevent
790 it from being garbage collected immediately when it goes out of scope.
791 """
792 butler = self.create_empty_butler(cleanup=False)
793 is_direct_butler = isinstance(butler, DirectButler)
794 butler_ref = weakref.ref(butler)
795 if is_direct_butler:
796 registry_ref = weakref.ref(butler._registry)
797 managers_ref = weakref.ref(butler._registry._managers)
798 datastore_ref = weakref.ref(butler._datastore)
799 db_ref = weakref.ref(butler._registry._db)
800 engine_ref = weakref.ref(butler._registry._db._engine)
802 with warnings.catch_warnings():
803 # Hide warnings from unclosed database handles.
804 warnings.simplefilter("ignore", ResourceWarning)
805 del butler
806 self.assertIsNone(butler_ref(), "Butler should have been garbage collected")
807 if is_direct_butler:
808 self.assertIsNone(registry_ref(), "SqlRegistry should have been garbage collected")
809 self.assertIsNone(managers_ref(), "Registry managers should have been garbage collected")
810 self.assertIsNone(datastore_ref(), "Datastore should have been garbage collected")
811 self.assertIsNone(db_ref(), "Database should have been garbage collected")
812 # SQLAlchemy has internal reference cycles, so the Engine instance
813 # is not cleaned up promptly even if we release our reference to
814 # it. Explicitly clean it up here to avoid file handles leaking.
815 if is_direct_butler:
816 engine = engine_ref()
817 if engine is not None:
818 engine.dispose()
820 def testDafButlerRepositories(self):
821 with unittest.mock.patch.dict(
822 os.environ,
823 {"DAF_BUTLER_REPOSITORIES": "label: 'https://someuri.com'\notherLabel: 'https://otheruri.com'\n"},
824 ):
825 self.assertEqual(str(Butler.get_repo_uri("label")), "https://someuri.com")
827 with unittest.mock.patch.dict(
828 os.environ,
829 {
830 "DAF_BUTLER_REPOSITORIES": "label: https://someuri.com",
831 "DAF_BUTLER_REPOSITORY_INDEX": "https://someuri.com",
832 },
833 ):
834 with self.assertRaisesRegex(RuntimeError, "Only one of the environment variables"):
835 Butler.get_repo_uri("label")
837 with unittest.mock.patch.dict(
838 os.environ,
839 {"DAF_BUTLER_REPOSITORIES": "invalid"},
840 ):
841 with self.assertRaisesRegex(ValueError, "Repository index not in expected format"):
842 Butler.get_repo_uri("label")
844 def testBasicPutGet(self) -> None:
845 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
846 self.runPutGetTest(storageClass, "test_metric")
848 def testCompositePutGetConcrete(self) -> None:
849 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadCompNoDisassembly")
850 butler = self.runPutGetTest(storageClass, "test_metric")
852 # Should *not* be disassembled
853 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
854 self.assertEqual(len(datasets), 1)
855 uri, components = butler.getURIs(datasets[0])
856 self.assertIsInstance(uri, ResourcePath)
857 self.assertFalse(components)
858 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
859 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
861 # Predicted dataset
862 if self.predictionSupported:
863 dataId = {"instrument": "DummyCamComp", "visit": 424}
864 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
865 self.assertFalse(components)
866 self.assertIsInstance(uri, ResourcePath)
867 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
868 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
870 def testCompositePutGetVirtual(self) -> None:
871 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp")
872 butler = self.runPutGetTest(storageClass, "test_metric_comp")
874 # Should be disassembled
875 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
876 self.assertEqual(len(datasets), 1)
877 uri, components = butler.getURIs(datasets[0])
879 if butler._datastore.isEphemeral:
880 # Never disassemble in-memory datastore
881 self.assertIsInstance(uri, ResourcePath)
882 self.assertFalse(components)
883 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
884 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
885 else:
886 self.assertIsNone(uri)
887 self.assertEqual(set(components), set(storageClass.components))
888 for compuri in components.values():
889 self.assertIsInstance(compuri, ResourcePath)
890 self.assertIn("423", str(compuri), f"Checking visit is in URI {compuri}")
891 self.assertEqual(compuri.fragment, "", f"Checking absence of fragment in {compuri}")
893 if self.predictionSupported:
894 # Predicted dataset
895 dataId = {"instrument": "DummyCamComp", "visit": 424}
896 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
898 if butler._datastore.isEphemeral:
899 # Never disassembled
900 self.assertIsInstance(uri, ResourcePath)
901 self.assertFalse(components)
902 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
903 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
904 else:
905 self.assertIsNone(uri)
906 self.assertEqual(set(components), set(storageClass.components))
907 for compuri in components.values():
908 self.assertIsInstance(compuri, ResourcePath)
909 self.assertIn("424", str(compuri), f"Checking visit is in URI {compuri}")
910 self.assertEqual(compuri.fragment, "predicted", f"Checking for fragment in {compuri}")
912 def testStorageClassOverrideGet(self) -> None:
913 """Test storage class conversion on get with override."""
914 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
915 datasetTypeName = "anything"
916 run = self.default_run
918 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
920 # Create and store a dataset.
921 metric = makeExampleMetrics()
922 dataId = {"instrument": "DummyCamComp", "visit": 423}
924 ref = butler.put(metric, datasetType, dataId)
926 # Return native type.
927 retrieved = butler.get(ref)
928 self.assertEqual(retrieved, metric)
930 # Specify an override.
931 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
932 model = butler.get(ref, storageClass=new_sc)
933 self.assertNotEqual(type(model), type(retrieved))
934 self.assertIs(type(model), new_sc.pytype)
935 self.assertEqual(retrieved, model)
937 # Defer but override later.
938 deferred = butler.getDeferred(ref)
939 model = deferred.get(storageClass=new_sc)
940 self.assertIs(type(model), new_sc.pytype)
941 self.assertEqual(retrieved, model)
943 # Defer but override up front.
944 deferred = butler.getDeferred(ref, storageClass=new_sc)
945 model = deferred.get()
946 self.assertIs(type(model), new_sc.pytype)
947 self.assertEqual(retrieved, model)
949 # Retrieve a component. Should be a tuple.
950 data = butler.get("anything.data", dataId, storageClass="StructuredDataDataTestTuple")
951 self.assertIs(type(data), tuple)
952 self.assertEqual(data, tuple(retrieved.data))
954 # Parameter on the write storage class should work regardless
955 # of read storage class.
956 data = butler.get(
957 "anything.data",
958 dataId,
959 storageClass="StructuredDataDataTestTuple",
960 parameters={"slice": slice(2, 4)},
961 )
962 self.assertEqual(len(data), 2)
964 # Try a parameter that is known to the read storage class but not
965 # the write storage class.
966 with self.assertRaises(KeyError):
967 butler.get(
968 "anything.data",
969 dataId,
970 storageClass="StructuredDataDataTestTuple",
971 parameters={"xslice": slice(2, 4)},
972 )
974 def testPytypePutCoercion(self) -> None:
975 """Test python type coercion on Butler.get and put."""
976 # Store some data with the normal example storage class.
977 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
978 datasetTypeName = "test_metric"
979 butler, _ = self.create_butler(self.default_run, storageClass, datasetTypeName)
981 dataId = {"instrument": "DummyCamComp", "visit": 423}
983 # Put a dict and this should coerce to a MetricsExample
984 test_dict = {"summary": {"a": 1}, "output": {"b": 2}}
985 metric_ref = butler.put(test_dict, datasetTypeName, dataId=dataId, visit=424)
986 test_metric = butler.get(metric_ref)
987 self.assertEqual(get_full_type_name(test_metric), "lsst.daf.butler.tests.MetricsExample")
988 self.assertEqual(test_metric.summary, test_dict["summary"])
989 self.assertEqual(test_metric.output, test_dict["output"])
991 # Check that the put still works if a DatasetType is given with
992 # a definition matching this python type.
993 registry_type = butler.get_dataset_type(datasetTypeName)
994 this_type = DatasetType(datasetTypeName, registry_type.dimensions, "StructuredDataDictJson")
995 metric2_ref = butler.put(test_dict, this_type, dataId=dataId, visit=425)
996 self.assertEqual(metric2_ref.datasetType, registry_type)
998 # The get will return the type expected by registry.
999 test_metric2 = butler.get(metric2_ref)
1000 self.assertEqual(get_full_type_name(test_metric2), "lsst.daf.butler.tests.MetricsExample")
1002 # Make a new DatasetRef with the compatible but different DatasetType.
1003 # This should now return a dict.
1004 new_ref = DatasetRef(this_type, metric2_ref.dataId, id=metric2_ref.id, run=metric2_ref.run)
1005 test_dict2 = butler.get(new_ref)
1006 self.assertEqual(get_full_type_name(test_dict2), "dict")
1008 # Get it again with the wrong dataset type definition using get()
1009 # rather than get(). This should be consistent with get()
1010 # behavior and return the type of the DatasetType.
1011 test_dict3 = butler.get(this_type, dataId=dataId, visit=425)
1012 self.assertEqual(get_full_type_name(test_dict3), "dict")
1014 def test_ingest_zip(self) -> None:
1015 """Create butler, export data, delete data, import from Zip."""
1016 butler, dataset_type = self.create_butler(
1017 run=self.default_run, storageClass="StructuredData", datasetTypeName="metrics"
1018 )
1020 metric = makeExampleMetrics()
1021 refs = []
1022 for visit in (423, 424, 425):
1023 ref = butler.put(metric, dataset_type, instrument="DummyCamComp", visit=visit)
1024 refs.append(ref)
1026 # Retrieve a Zip file.
1027 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
1028 zip = butler.retrieve_artifacts_zip(refs, destination=tmpdir)
1030 # Ingest will fail.
1031 with self.assertRaises(ConflictingDefinitionError):
1032 butler.ingest_zip(zip)
1034 # Clear out the collection.
1035 butler.removeRuns([self.default_run])
1036 self.assertFalse(butler.exists(refs[0]))
1038 butler.ingest_zip(zip, transfer="copy")
1040 # Create an entirely new local file butler in this temp directory.
1041 new_butler_cfg = Butler.makeRepo(tmpdir)
1042 new_butler = Butler.from_config(new_butler_cfg, writeable=True)
1043 self.enterContext(new_butler)
1045 # This will fail since dimensions records are missing.
1046 with self.assertRaises(ConflictingDefinitionError):
1047 new_butler.ingest_zip(zip, transfer="copy")
1049 # Dry run should work.
1050 new_butler.ingest_zip(zip, transfer="copy", dry_run=True)
1052 new_butler.ingest_zip(zip, transfer="copy", transfer_dimensions=True)
1053 self.assertTrue(butler.exists(refs[0]))
1055 # Check that the refs can be read again.
1056 _ = [butler.get(ref) for ref in refs]
1058 uri = butler.getURI(refs[2])
1059 self.assertTrue(uri.exists())
1061 # Delete one dataset. The Zip file should still exist and allow
1062 # remaining refs to be read.
1063 butler.pruneDatasets([refs[0]], purge=True, unstore=True)
1064 self.assertTrue(uri.exists())
1066 metric2 = butler.get(refs[1])
1067 self.assertEqual(metric2, metric, msg=f"{metric2} != {metric}")
1069 butler.removeRuns([self.default_run])
1070 self.assertFalse(uri.exists())
1071 self.assertFalse(butler.exists(refs[-1]))
1073 with self.assertRaises(ValueError):
1074 butler.retrieve_artifacts_zip([], destination=".")
1076 def testIngest(self) -> None:
1077 butler = self.create_empty_butler(run=self.default_run)
1079 # Create and register a DatasetType
1080 dimensions = butler.dimensions.conform(["instrument", "visit", "detector"])
1082 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDictYaml")
1083 datasetTypeName = "metric"
1085 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1087 # Add needed Dimensions
1088 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1089 butler.registry.insertDimensionData(
1090 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1091 )
1092 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1093 for detector in (1, 2):
1094 butler.registry.insertDimensionData(
1095 "detector", {"instrument": "DummyCamComp", "id": detector, "full_name": f"detector{detector}"}
1096 )
1098 butler.registry.insertDimensionData(
1099 "visit",
1100 {
1101 "instrument": "DummyCamComp",
1102 "id": 423,
1103 "name": "fourtwentythree",
1104 "physical_filter": "d-r",
1105 "day_obs": 20250101,
1106 },
1107 {
1108 "instrument": "DummyCamComp",
1109 "id": 424,
1110 "name": "fourtwentyfour",
1111 "physical_filter": "d-r",
1112 "day_obs": 20250101,
1113 },
1114 )
1116 formatter = doImportType("lsst.daf.butler.formatters.yaml.YamlFormatter")
1117 dataRoot = os.path.join(TESTDIR, "data", "basic")
1118 datasets = []
1119 # Test one DatasetRef with a run that exists, and the other with a run
1120 # that doesn't exist, to verify that run collections are created when
1121 # required.
1122 runs = {1: self.default_run, 2: "a/new/run"}
1123 for detector in (1, 2):
1124 detector_name = f"detector_{detector}"
1125 metricFile = os.path.join(dataRoot, f"{detector_name}.yaml")
1126 dataId = butler.registry.expandDataId(
1127 {"instrument": "DummyCamComp", "visit": 423, "detector": detector}
1128 )
1129 # Create a DatasetRef for ingest
1130 refIn = DatasetRef(datasetType, dataId, run=runs[detector])
1132 datasets.append(FileDataset(path=metricFile, refs=[refIn], formatter=formatter))
1134 butler.ingest(*datasets, transfer="copy")
1136 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 423}
1137 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 423}
1139 metrics1 = butler.get(datasetTypeName, dataId1)
1140 metrics2 = butler.get(datasetTypeName, dataId2, collections="a/new/run")
1141 self.assertNotEqual(metrics1, metrics2)
1143 # Compare URIs
1144 uri1 = butler.getURI(datasetTypeName, dataId1)
1145 uri2 = butler.getURI(datasetTypeName, dataId2, collections="a/new/run")
1146 self.assertFalse(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1148 # Re-ingesting the same datasets raises an error with
1149 # skip_existing=False.
1150 with self.assertRaises(ConflictingDefinitionError):
1151 butler.ingest(*datasets, transfer="copy")
1152 # skip_existing=True makes it a no-op to re-ingest the same datasets.
1153 butler.ingest(*datasets, transfer="copy", skip_existing=True)
1155 # Now do a multi-dataset but single file ingest
1156 metricFile = os.path.join(dataRoot, "detectors.yaml")
1157 refs = []
1158 for detector in (1, 2):
1159 detector_name = f"detector_{detector}"
1160 dataId = butler.registry.expandDataId(
1161 {"instrument": "DummyCamComp", "visit": 424, "detector": detector}
1162 )
1163 # Create a DatasetRef for ingest
1164 refs.append(DatasetRef(datasetType, dataId, run=self.default_run))
1166 # Test "move" transfer to ensure that the files themselves
1167 # have disappeared following ingest.
1168 with ResourcePath.temporary_uri(suffix=".yaml") as tempFile:
1169 tempFile.transfer_from(ResourcePath(metricFile), transfer="copy")
1171 datasets = []
1172 datasets.append(FileDataset(path=tempFile, refs=refs, formatter=MultiDetectorFormatter))
1174 # For first ingest use copy.
1175 butler.ingest(*datasets, transfer="copy", record_validation_info=False)
1177 # Now try to ingest again in "execution butler" mode where
1178 # the registry entries exist but the datastore does not have
1179 # the files. We also need to strip the dimension records to ensure
1180 # that they will be re-added by the ingest.
1181 ref = datasets[0].refs[0]
1182 datasets[0].refs = [
1183 cast(
1184 DatasetRef,
1185 butler.find_dataset(ref.datasetType, data_id=ref.dataId, collections=ref.run),
1186 )
1187 for ref in datasets[0].refs
1188 ]
1189 all_refs = []
1190 for dataset in datasets:
1191 refs = []
1192 for ref in dataset.refs:
1193 # Create a dict from the dataId to drop the records.
1194 new_data_id = dict(ref.dataId.required)
1195 new_ref = butler.find_dataset(ref.datasetType, new_data_id, collections=ref.run)
1196 assert new_ref is not None
1197 self.assertFalse(new_ref.dataId.hasRecords())
1198 refs.append(new_ref)
1199 dataset.refs = refs
1200 all_refs.extend(dataset.refs)
1201 butler.pruneDatasets(all_refs, disassociate=False, unstore=True, purge=False)
1203 # Use move mode to test that the file is deleted. Also
1204 # disable recording of file size.
1205 butler.ingest(*datasets, transfer="move", record_validation_info=False)
1207 # Check that every ref now has records.
1208 for dataset in datasets:
1209 for ref in dataset.refs:
1210 self.assertTrue(ref.dataId.hasRecords())
1212 # Ensure that the file has disappeared.
1213 self.assertFalse(tempFile.exists())
1215 # Check that the datastore recorded no file size.
1216 # Not all datastores can support this.
1217 try:
1218 infos = butler._datastore.getStoredItemsInfo(datasets[0].refs[0]) # type: ignore[attr-defined]
1219 self.assertEqual(infos[0].file_size, -1)
1220 except AttributeError:
1221 pass
1223 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 424}
1224 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 424}
1226 multi1 = butler.get(datasetTypeName, dataId1)
1227 multi2 = butler.get(datasetTypeName, dataId2)
1229 self.assertEqual(multi1, metrics1)
1230 self.assertEqual(multi2, metrics2)
1232 # Compare URIs
1233 uri1 = butler.getURI(datasetTypeName, dataId1)
1234 uri2 = butler.getURI(datasetTypeName, dataId2)
1235 self.assertTrue(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1237 # Test that removing one does not break the second
1238 # This line will issue a warning log message for a ChainedDatastore
1239 # that uses an InMemoryDatastore since in-memory can not ingest
1240 # files.
1241 butler.pruneDatasets([datasets[0].refs[0]], unstore=True, disassociate=False)
1242 self.assertFalse(butler.exists(datasetTypeName, dataId1))
1243 self.assertTrue(butler.exists(datasetTypeName, dataId2))
1244 multi2b = butler.get(datasetTypeName, dataId2)
1245 self.assertEqual(multi2, multi2b)
1247 # Ensure we can ingest 0 datasets
1248 datasets = []
1249 butler.ingest(*datasets)
1251 def testPickle(self) -> None:
1252 """Test pickle support."""
1253 butler = self.create_empty_butler(run=self.default_run)
1254 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
1255 butlerOut = pickle.loads(pickle.dumps(butler))
1256 self.enterContext(butlerOut)
1257 self.assertIsInstance(butlerOut, Butler)
1258 self.assertEqual(butlerOut._config, butler._config)
1259 self.assertEqual(list(butlerOut.collections.defaults), list(butler.collections.defaults))
1260 self.assertEqual(butlerOut.run, butler.run)
1262 def testGetDatasetTypes(self) -> None:
1263 butler = self.create_empty_butler(run=self.default_run)
1264 dimensions = butler.dimensions.conform(["instrument", "visit", "physical_filter"])
1265 dimensionEntries: list[tuple[str, list[Mapping[str, Any]]]] = [
1266 (
1267 "instrument",
1268 [
1269 {"instrument": "DummyCam"},
1270 {"instrument": "DummyHSC"},
1271 {"instrument": "DummyCamComp"},
1272 ],
1273 ),
1274 ("physical_filter", [{"instrument": "DummyCam", "name": "d-r", "band": "R"}]),
1275 ("day_obs", [{"instrument": "DummyCam", "id": 20250101}]),
1276 (
1277 "visit",
1278 [
1279 {
1280 "instrument": "DummyCam",
1281 "id": 42,
1282 "name": "fortytwo",
1283 "physical_filter": "d-r",
1284 "day_obs": 20250101,
1285 }
1286 ],
1287 ),
1288 ]
1289 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1290 # Add needed Dimensions
1291 for element, data in dimensionEntries:
1292 butler.registry.insertDimensionData(element, *data)
1294 # When a DatasetType is added to the registry entries are not created
1295 # for components but querying them can return the components.
1296 datasetTypeNames = {"metric", "metric2", "metric4", "metric33", "pvi", "paramtest"}
1297 components = set()
1298 for datasetTypeName in datasetTypeNames:
1299 # Create and register a DatasetType
1300 self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1302 for componentName in storageClass.components:
1303 components.add(DatasetType.nameWithComponent(datasetTypeName, componentName))
1305 fromRegistry: set[DatasetType] = set()
1306 for parent_dataset_type in butler.registry.queryDatasetTypes():
1307 fromRegistry.add(parent_dataset_type)
1308 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes())
1309 self.assertEqual({d.name for d in fromRegistry}, datasetTypeNames | components)
1311 # Query with wildcard.
1312 dataset_types = butler.registry.queryDatasetTypes("metric*")
1313 self.assertEqual(len(dataset_types), 4, f"Got: {dataset_types}")
1314 # but not regex.
1315 with self.assertRaises(DatasetTypeExpressionError):
1316 butler.registry.queryDatasetTypes(["pvi", re.compile("metric.*")])
1318 # Now that we have some dataset types registered, validate them
1319 butler.validateConfiguration(
1320 ignore=[
1321 "test_metric_comp",
1322 "metric3",
1323 "metric5",
1324 "calexp",
1325 "DummySC",
1326 "datasetType.component",
1327 "random_data",
1328 "random_data_2",
1329 ]
1330 )
1332 # Add a new datasetType that will fail template validation
1333 self.addDatasetType("test_metric_comp", dimensions, storageClass, butler.registry)
1334 if self.validationCanFail:
1335 with self.assertRaises(ValidationError):
1336 butler.validateConfiguration()
1338 # Rerun validation but with a subset of dataset type names
1339 butler.validateConfiguration(datasetTypeNames=["metric4"])
1341 # Rerun validation but ignore the bad datasetType
1342 butler.validateConfiguration(
1343 ignore=[
1344 "test_metric_comp",
1345 "metric3",
1346 "metric5",
1347 "calexp",
1348 "DummySC",
1349 "datasetType.component",
1350 "random_data",
1351 "random_data_2",
1352 ]
1353 )
1355 def testTransaction(self) -> None:
1356 butler = self.create_empty_butler(run=self.default_run)
1357 datasetTypeName = "test_metric"
1358 dimensions = butler.dimensions.conform(["instrument", "visit"])
1359 dimensionEntries: tuple[tuple[str, Mapping[str, Any]], ...] = (
1360 ("instrument", {"instrument": "DummyCam"}),
1361 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}),
1362 ("day_obs", {"instrument": "DummyCam", "id": 20250101}),
1363 (
1364 "visit",
1365 {
1366 "instrument": "DummyCam",
1367 "id": 42,
1368 "name": "fortytwo",
1369 "physical_filter": "d-r",
1370 "day_obs": 20250101,
1371 },
1372 ),
1373 )
1374 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1375 metric = makeExampleMetrics()
1376 dataId = {"instrument": "DummyCam", "visit": 42}
1377 # Create and register a DatasetType
1378 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1379 with self.assertRaises(TransactionTestError):
1380 with butler.transaction():
1381 # Add needed Dimensions
1382 for args in dimensionEntries:
1383 butler.registry.insertDimensionData(*args)
1384 # Store a dataset
1385 ref = butler.put(metric, datasetTypeName, dataId)
1386 self.assertIsInstance(ref, DatasetRef)
1387 # Test get of a ref.
1388 metricOut = butler.get(ref)
1389 self.assertEqual(metric, metricOut)
1390 # Test get
1391 metricOut = butler.get(datasetTypeName, dataId)
1392 self.assertEqual(metric, metricOut)
1393 # Check we can get components
1394 self.assertGetComponents(butler, ref, ("summary", "data", "output"), metric)
1395 raise TransactionTestError("This should roll back the entire transaction")
1396 with self.assertRaises(DataIdValueError, msg=f"Check can't expand DataId {dataId}"):
1397 butler.registry.expandDataId(dataId)
1398 # Should raise LookupError for missing data ID value
1399 with self.assertRaises(LookupError, msg=f"Check can't get by {datasetTypeName} and {dataId}"):
1400 butler.get(datasetTypeName, dataId)
1401 # Also check explicitly if Dataset entry is missing
1402 self.assertIsNone(butler.find_dataset(datasetType, dataId, collections=butler.collections.defaults))
1403 # Direct retrieval should not find the file in the Datastore
1404 with self.assertRaises(FileNotFoundError, msg=f"Check {ref} can't be retrieved directly"):
1405 butler.get(ref)
1407 def testMakeRepo(self) -> None:
1408 """Test that we can write butler configuration to a new repository via
1409 the Butler.makeRepo interface and then instantiate a butler from the
1410 repo root.
1411 """
1412 # Do not run the test if we know this datastore configuration does
1413 # not support a file system root
1414 if self.fullConfigKey is None:
1415 return
1417 # create two separate directories
1418 root1 = tempfile.mkdtemp(dir=self.root)
1419 root2 = tempfile.mkdtemp(dir=self.root)
1421 butlerConfig = Butler.makeRepo(root1, config=Config(self.configFile))
1422 limited = Config(self.configFile)
1423 butler1 = Butler.from_config(butlerConfig)
1424 self.enterContext(butler1)
1425 assert isinstance(butler1, DirectButler), "Expect DirectButler in configuration"
1426 butlerConfig = Butler.makeRepo(root2, standalone=True, config=Config(self.configFile))
1427 full = Config(self.tmpConfigFile)
1428 butler2 = Butler.from_config(butlerConfig)
1429 self.enterContext(butler2)
1430 assert isinstance(butler2, DirectButler), "Expect DirectButler in configuration"
1431 # Butlers should have the same configuration regardless of whether
1432 # defaults were expanded.
1433 self.assertEqual(butler1._config, butler2._config)
1434 # Config files loaded directly should not be the same.
1435 self.assertNotEqual(limited, full)
1436 # Make sure "limited" doesn't have a few keys we know it should be
1437 # inheriting from defaults.
1438 self.assertIn(self.fullConfigKey, full)
1439 self.assertNotIn(self.fullConfigKey, limited)
1441 # Collections don't appear until something is put in them
1442 collections1 = set(butler1.registry.queryCollections())
1443 self.assertEqual(collections1, set())
1444 self.assertEqual(set(butler2.registry.queryCollections()), collections1)
1446 # Check that a config with no associated file name will not
1447 # work properly with relocatable Butler repo
1448 butlerConfig.configFile = None
1449 with self.assertRaises(ValueError):
1450 Butler.from_config(butlerConfig)
1452 with self.assertRaises(FileExistsError):
1453 Butler.makeRepo(self.root, standalone=True, config=Config(self.configFile), overwrite=False)
1455 def testStringification(self) -> None:
1456 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
1457 self.enterContext(butler)
1458 butlerStr = str(butler)
1460 if self.datastoreStr is not None:
1461 for testStr in self.datastoreStr:
1462 self.assertIn(testStr, butlerStr)
1463 if self.registryStr is not None:
1464 self.assertIn(self.registryStr, butlerStr)
1466 datastoreName = butler._datastore.name
1467 if self.datastoreName is not None:
1468 for testStr in self.datastoreName:
1469 self.assertIn(testStr, datastoreName)
1471 def testButlerRewriteDataId(self) -> None:
1472 """Test that dataIds can be rewritten based on dimension records."""
1473 butler = self.create_empty_butler(run=self.default_run)
1475 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
1476 datasetTypeName = "random_data"
1478 # Create dimension records.
1479 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1480 butler.registry.insertDimensionData(
1481 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1482 )
1483 butler.registry.insertDimensionData(
1484 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
1485 )
1487 dimensions = butler.dimensions.conform(["instrument", "exposure"])
1488 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
1489 butler.registry.registerDatasetType(datasetType)
1491 n_exposures = 5
1492 dayobs = 20210530
1494 # Create records for multiple day_obs but same seq_num to test that
1495 # we are constraining gets properly when day_obs/seq_num is used
1496 # for an exposure. Second day is year in future but is not used.
1497 for day_obs in (dayobs, dayobs + 1_00_00):
1498 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": day_obs})
1500 for i in range(n_exposures):
1501 group_name = f"group_{day_obs}_{i}"
1502 butler.registry.insertDimensionData(
1503 "group", {"instrument": "DummyCamComp", "name": group_name}
1504 )
1505 butler.registry.insertDimensionData(
1506 "exposure",
1507 {
1508 "instrument": "DummyCamComp",
1509 "id": day_obs + i,
1510 "obs_id": f"exp_{day_obs}_{i}",
1511 "seq_num": i,
1512 "day_obs": day_obs,
1513 "physical_filter": "d-r",
1514 "group": group_name,
1515 },
1516 )
1518 # Write some data.
1519 for i in range(n_exposures):
1520 metric = {"something": i, "other": "metric", "list": [2 * x for x in range(i)]}
1522 # Use the seq_num for the put to test rewriting.
1523 dataId = {"seq_num": i, "day_obs": dayobs, "instrument": "DummyCamComp", "physical_filter": "d-r"}
1524 ref = butler.put(metric, datasetTypeName, dataId=dataId)
1526 # Check that the exposure is correct in the dataId
1527 self.assertEqual(ref.dataId["exposure"], dayobs + i)
1529 # and check that we can get the dataset back with the same dataId
1530 new_metric = butler.get(datasetTypeName, dataId=dataId)
1531 self.assertEqual(new_metric, metric)
1533 # Check that we can find the datasets using the day_obs or the
1534 # exposure.day_obs.
1535 datasets_1 = list(
1536 butler.registry.queryDatasets(
1537 datasetType,
1538 collections=self.default_run,
1539 where="day_obs = :dayObs AND instrument = :instr",
1540 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1541 )
1542 )
1543 datasets_2 = list(
1544 butler.registry.queryDatasets(
1545 datasetType,
1546 collections=self.default_run,
1547 where="exposure.day_obs = :dayObs AND instrument = :instr",
1548 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1549 )
1550 )
1551 self.assertEqual(datasets_1, datasets_2)
1553 def testGetDatasetCollectionCaching(self):
1554 # Prior to DM-41117, there was a bug where get_dataset would throw
1555 # MissingCollectionError if you tried to fetch a dataset that was added
1556 # after the collection cache was last updated.
1557 reader_butler, datasetType = self.create_butler(self.default_run, "int", "datasettypename")
1558 writer_butler = self.create_empty_butler(writeable=True, run="new_run")
1559 dataId = {"instrument": "DummyCamComp", "visit": 423}
1560 put_ref = writer_butler.put(123, datasetType, dataId)
1561 get_ref = reader_butler.get_dataset(put_ref.id)
1562 self.assertEqual(get_ref.id, put_ref.id)
1563 # Also works when looking up via a hexadecimal string instead of a UUID
1564 # instance.
1565 hex_ref = reader_butler.get_dataset(put_ref.id.hex)
1566 self.assertEqual(hex_ref.id, put_ref.id)
1568 def testCollectionChainRedefine(self):
1569 butler = self._setup_to_test_collection_chain()
1571 butler.collections.redefine_chain("chain", "a")
1572 self._check_chain(butler, ["a"])
1574 # Duplicates are removed from the list of children
1575 butler.collections.redefine_chain("chain", ["c", "b", "c"])
1576 self._check_chain(butler, ["c", "b"])
1578 # Empty list clears the chain
1579 butler.collections.redefine_chain("chain", [])
1580 self._check_chain(butler, [])
1582 self._test_common_chain_functionality(butler, butler.collections.redefine_chain)
1584 def testCollectionChainPrepend(self):
1585 butler = self._setup_to_test_collection_chain()
1587 # Duplicates are removed from the list of children
1588 butler.collections.prepend_chain("chain", ["c", "b", "c"])
1589 self._check_chain(butler, ["c", "b"])
1591 # Prepend goes on the front of existing chain
1592 butler.collections.prepend_chain("chain", ["a"])
1593 self._check_chain(butler, ["a", "c", "b"])
1595 # Empty prepend does nothing
1596 butler.collections.prepend_chain("chain", [])
1597 self._check_chain(butler, ["a", "c", "b"])
1599 # Prepending children that already exist in the chain removes them from
1600 # their current position.
1601 butler.collections.prepend_chain("chain", ["d", "b", "c"])
1602 self._check_chain(butler, ["d", "b", "c", "a"])
1604 self._test_common_chain_functionality(butler, butler.collections.prepend_chain)
1606 def testCollectionChainExtend(self):
1607 butler = self._setup_to_test_collection_chain()
1609 # Duplicates are removed from the list of children
1610 butler.collections.extend_chain("chain", ["c", "b", "c"])
1611 self._check_chain(butler, ["c", "b"])
1613 # Extend goes on the end of existing chain
1614 butler.collections.extend_chain("chain", ["a"])
1615 self._check_chain(butler, ["c", "b", "a"])
1617 # Empty extend does nothing
1618 butler.collections.extend_chain("chain", [])
1619 self._check_chain(butler, ["c", "b", "a"])
1621 # Extending children that already exist in the chain removes them from
1622 # their current position.
1623 butler.collections.extend_chain("chain", ["d", "b", "c"])
1624 self._check_chain(butler, ["a", "d", "b", "c"])
1626 self._test_common_chain_functionality(butler, butler.collections.extend_chain)
1628 def testCollectionChainRemove(self) -> None:
1629 butler = self._setup_to_test_collection_chain()
1631 butler.collections.redefine_chain("chain", ["a", "b", "c", "d"])
1633 butler.collections.remove_from_chain("chain", "c")
1634 self._check_chain(butler, ["a", "b", "d"])
1636 # Duplicates are allowed in the list of children
1637 butler.collections.remove_from_chain("chain", ["b", "b", "a"])
1638 self._check_chain(butler, ["d"])
1640 # Empty remove does nothing
1641 butler.collections.remove_from_chain("chain", [])
1642 self._check_chain(butler, ["d"])
1644 # Removing children that aren't in the chain does nothing
1645 butler.collections.remove_from_chain("chain", ["a", "chain"])
1646 self._check_chain(butler, ["d"])
1648 self._test_common_chain_functionality(
1649 butler, butler.collections.remove_from_chain, skip_cycle_check=True
1650 )
1652 def _setup_to_test_collection_chain(self) -> Butler:
1653 butler = self.create_empty_butler(writeable=True)
1655 butler.collections.register("chain", CollectionType.CHAINED)
1657 runs = ["a", "b", "c", "d"]
1658 for run in runs:
1659 butler.collections.register(run)
1661 butler.collections.register("staticchain", CollectionType.CHAINED)
1662 butler.collections.redefine_chain("staticchain", ["a", "b"])
1664 return butler
1666 def _check_chain(self, butler: Butler, expected: list[str]) -> None:
1667 children = butler.collections.get_info("chain").children
1668 self.assertEqual(expected, list(children))
1670 def _test_common_chain_functionality(
1671 self, butler, func: Callable[[str, str | list[str]], Any], *, skip_cycle_check=False
1672 ) -> None:
1673 # Missing parent collection
1674 with self.assertRaises(MissingCollectionError):
1675 func("doesnotexist", [])
1676 # Missing child collection
1677 with self.assertRaises(MissingCollectionError):
1678 func("chain", ["doesnotexist"])
1679 # Forbid operations on non-chained collections
1680 with self.assertRaises(CollectionTypeError):
1681 func("d", ["a"])
1683 # Prevent collection cycles
1684 if not skip_cycle_check:
1685 butler.collections.register("chain2", CollectionType.CHAINED)
1686 func("chain2", "chain")
1687 with self.assertRaises(CollectionCycleError):
1688 func("chain", "chain2")
1690 # Make sure none of the earlier operations interfered with unrelated
1691 # chains.
1692 self.assertEqual(["a", "b"], list(butler.collections.get_info("staticchain").children))
1694 with butler._caching_context():
1695 with self.assertRaisesRegex(RuntimeError, "Chained collection modification not permitted"):
1696 func("chain", "a")
1699class FileDatastoreButlerTests(ButlerTests):
1700 """Common tests and specialization of ButlerTests for butlers backed
1701 by datastores that inherit from FileDatastore.
1702 """
1704 trustModeSupported = True
1706 def checkFileExists(self, root: str | ResourcePath, relpath: str | ResourcePath) -> bool:
1707 """Check if file exists at a given path (relative to root).
1709 Test testPutTemplates verifies actual physical existance of the files
1710 in the requested location.
1711 """
1712 uri = ResourcePath(root, forceDirectory=True)
1713 return uri.join(relpath).exists()
1715 def testPutTemplates(self) -> None:
1716 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1717 butler = self.create_empty_butler(run=self.default_run)
1719 # Add needed Dimensions
1720 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1721 butler.registry.insertDimensionData(
1722 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1723 )
1724 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1725 butler.registry.insertDimensionData(
1726 "visit",
1727 {
1728 "instrument": "DummyCamComp",
1729 "id": 423,
1730 "name": "v423",
1731 "physical_filter": "d-r",
1732 "day_obs": 20250101,
1733 },
1734 )
1735 butler.registry.insertDimensionData(
1736 "visit",
1737 {
1738 "instrument": "DummyCamComp",
1739 "id": 425,
1740 "name": "v425",
1741 "physical_filter": "d-r",
1742 "day_obs": 20250101,
1743 },
1744 )
1746 # Create and store a dataset
1747 metric = makeExampleMetrics()
1749 # Create two almost-identical DatasetTypes (both will use default
1750 # template)
1751 dimensions = butler.dimensions.conform(["instrument", "visit"])
1752 butler.registry.registerDatasetType(DatasetType("metric1", dimensions, storageClass))
1753 butler.registry.registerDatasetType(DatasetType("metric2", dimensions, storageClass))
1754 butler.registry.registerDatasetType(DatasetType("metric3", dimensions, storageClass))
1756 dataId1 = {"instrument": "DummyCamComp", "visit": 423}
1757 dataId2 = {"instrument": "DummyCamComp", "visit": 423, "physical_filter": "d-r"}
1759 # Put with exactly the data ID keys needed
1760 ref = butler.put(metric, "metric1", dataId1)
1761 uri = butler.getURI(ref)
1762 self.assertTrue(uri.exists())
1763 self.assertTrue(
1764 uri.unquoted_path.endswith(f"{self.default_run}/metric1/??#?/d-r/DummyCamComp_423.pickle")
1765 )
1767 # Check the template based on dimensions
1768 if hasattr(butler._datastore, "templates"):
1769 butler._datastore.templates.validateTemplates([ref])
1771 # Put with extra data ID keys (physical_filter is an optional
1772 # dependency); should not change template (at least the way we're
1773 # defining them to behave now; the important thing is that they
1774 # must be consistent).
1775 ref = butler.put(metric, "metric2", dataId2)
1776 uri = butler.getURI(ref)
1777 self.assertTrue(uri.exists())
1778 self.assertTrue(
1779 uri.unquoted_path.endswith(f"{self.default_run}/metric2/d-r/DummyCamComp_v423.pickle")
1780 )
1782 # Check the template based on dimensions
1783 if hasattr(butler._datastore, "templates"):
1784 butler._datastore.templates.validateTemplates([ref])
1786 # Use a template that has a typo in dimension record metadata.
1787 # Easier to test with a butler that has a ref with records attached.
1788 template = FileTemplate("a/{visit.name}/{id}_{visit.namex:?}.fits")
1789 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1790 path = template.format(ref)
1791 self.assertEqual(path, f"a/v423/{ref.id}_fits")
1793 template = FileTemplate("a/{visit.name}/{id}_{visit.namex}.fits")
1794 with self.assertRaises(KeyError):
1795 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1796 template.format(ref)
1798 # Now use a file template that will not result in unique filenames
1799 with self.assertRaises(FileTemplateValidationError):
1800 butler.put(metric, "metric3", dataId1)
1802 def testImportExport(self) -> None:
1803 # Run put/get tests just to create and populate a repo.
1804 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1805 self.runImportExportTest(storageClass)
1807 @unittest.expectedFailure
1808 def testImportExportVirtualComposite(self) -> None:
1809 # Run put/get tests just to create and populate a repo.
1810 storageClass = self.storageClassFactory.getStorageClass("StructuredComposite")
1811 self.runImportExportTest(storageClass)
1813 def runImportExportTest(self, storageClass: StorageClass) -> None:
1814 """Test exporting and importing.
1816 This test does an export to a temp directory and an import back
1817 into a new temp directory repo. It does not assume a posix datastore.
1818 """
1819 exportButler = self.runPutGetTest(storageClass, "test_metric")
1821 # Test that we must have a file extension.
1822 with self.assertRaises(ValueError):
1823 with exportButler.export(filename="dump", directory=".") as export:
1824 pass
1826 # Test that unknown format is not allowed.
1827 with self.assertRaises(ValueError):
1828 with exportButler.export(filename="dump.fits", directory=".") as export:
1829 pass
1831 # Test that the repo actually has at least one dataset.
1832 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
1833 self.assertGreater(len(datasets), 0)
1834 # Add a DimensionRecord that's unused by those datasets.
1835 skymapRecord = {"name": "example_skymap", "hash": (50).to_bytes(8, byteorder="little")}
1836 exportButler.registry.insertDimensionData("skymap", skymapRecord)
1837 # Export and then import datasets.
1838 with safeTestTempDir(TESTDIR) as exportDir:
1839 exportFile = os.path.join(exportDir, "exports.yaml")
1840 with exportButler.export(filename=exportFile, directory=exportDir, transfer="auto") as export:
1841 export.saveDatasets(datasets)
1842 # Export the same datasets again. This should quietly do
1843 # nothing because of internal deduplication, and it shouldn't
1844 # complain about being asked to export the "htm7" elements even
1845 # though there aren't any in these datasets or in the database.
1846 export.saveDatasets(datasets, elements=["htm7"])
1847 # Save one of the data IDs again; this should be harmless
1848 # because of internal deduplication.
1849 export.saveDataIds([datasets[0].dataId])
1850 # Save some dimension records directly.
1851 export.saveDimensionData("skymap", [skymapRecord])
1852 self.assertTrue(os.path.exists(exportFile))
1853 with safeTestTempDir(TESTDIR) as importDir:
1854 # We always want this to be a local posix butler
1855 Butler.makeRepo(importDir, config=Config(os.path.join(TESTDIR, "config/basic/butler.yaml")))
1856 # Calling script.butlerImport tests the implementation of the
1857 # butler command line interface "import" subcommand. Functions
1858 # in the script folder are generally considered protected and
1859 # should not be used as public api.
1860 with open(exportFile) as f:
1861 script.butlerImport(
1862 importDir,
1863 export_file=f,
1864 directory=exportDir,
1865 transfer="auto",
1866 skip_dimensions=None,
1867 )
1868 importButler = Butler.from_config(importDir, run=self.default_run)
1869 self.enterContext(importButler)
1870 for ref in datasets:
1871 with self.subTest(ref=repr(ref)):
1872 # Test for existence by passing in the DatasetType and
1873 # data ID separately, to avoid lookup by dataset_id.
1874 self.assertTrue(importButler.exists(ref.datasetType, ref.dataId))
1875 self.assertEqual(
1876 list(importButler.registry.queryDimensionRecords("skymap")),
1877 [importButler.dimensions["skymap"].RecordClass(**skymapRecord)],
1878 )
1880 def testRemoveRuns(self) -> None:
1881 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1882 butler = self.create_empty_butler(writeable=True)
1883 # Load registry data with dimensions to hang datasets off of.
1884 butler.import_(filename=ResourcePath("resource://lsst.daf.butler/tests/registry_data/base.yaml"))
1885 # Add some RUN-type collection.
1886 run1 = "run1"
1887 butler.collections.register(run1)
1888 run2 = "run2"
1889 butler.collections.register(run2)
1890 # put a dataset in each
1891 metric = makeExampleMetrics()
1892 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1893 datasetType = self.addDatasetType(
1894 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1895 )
1896 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1897 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1898 uri1 = butler.getURI(ref1)
1899 uri2 = butler.getURI(ref2)
1901 # Put one of the runs in a chain.
1902 butler.collections.register("Chain", CollectionType.CHAINED)
1903 butler.collections.extend_chain("Chain", run1)
1905 with self.assertRaises(OrphanedRecordError):
1906 butler.registry.removeDatasetType(datasetType.name)
1908 # Remove a non-run.
1909 with self.assertRaises(TypeError):
1910 butler.removeRuns(["Chain"])
1912 # Remove without unlinking from chain should fail.
1913 with self.assertRaises(IntegrityError):
1914 butler.removeRuns([run1])
1916 # Remove from both runs. No longer use unstore parameter since it
1917 # always purges.
1918 butler.removeRuns([run1, run2], unlink_from_chains=True)
1920 # Should be nothing in registry for either one, and datastore should
1921 # not think either exists.
1922 with self.assertRaises(MissingCollectionError):
1923 butler.collections.get_info(run1)
1924 with self.assertRaises(MissingCollectionError):
1925 butler.collections.get_info(run1)
1926 self.assertFalse(butler.stored(ref1))
1927 self.assertFalse(butler.stored(ref2))
1928 # We always unstore so both URIs should be gone.
1929 self.assertFalse(uri1.exists())
1930 self.assertFalse(uri2.exists())
1932 # Now that the collections have been pruned we can remove the
1933 # dataset type
1934 butler.registry.removeDatasetType(datasetType.name)
1936 with self.assertLogs("lsst.daf.butler.registry", "INFO") as cm:
1937 butler.registry.removeDatasetType(("test*", "test*"))
1938 self.assertIn("not defined", "\n".join(cm.output))
1940 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
1941 """Simulate an external actor removing a file outside of Butler's
1942 knowledge.
1944 Subclasses may override to handle more complicated datastore
1945 configurations.
1946 """
1947 uri = butler.getURI(ref)
1948 uri.remove()
1949 datastore = cast(FileDatastore, butler._datastore)
1950 datastore.cacheManager.remove_from_cache(ref)
1952 def testPruneDatasets(self) -> None:
1953 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1954 butler = self.create_empty_butler(writeable=True)
1955 # Load registry data with dimensions to hang datasets off of.
1956 butler.import_(filename=_get_test_data_path("base.yaml"))
1957 # Add some RUN-type collections.
1958 run1 = "run1"
1959 butler.collections.register(run1)
1960 run2 = "run2"
1961 butler.collections.register(run2)
1962 # put some datasets. ref1 and ref2 have the same data ID, and are in
1963 # different runs. ref3 has a different data ID.
1964 metric = makeExampleMetrics()
1965 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1966 datasetType = self.addDatasetType(
1967 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1968 )
1969 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1970 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1971 ref3 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, run=run1)
1973 many_stored = butler.stored_many([ref1, ref2, ref3])
1974 for ref, stored in many_stored.items():
1975 self.assertTrue(stored, f"Ref {ref} should be stored")
1977 many_exists = butler._exists_many([ref1, ref2, ref3])
1978 for ref, exists in many_exists.items():
1979 self.assertTrue(exists, f"Checking ref {ref} exists.")
1980 self.assertEqual(exists, DatasetExistence.VERIFIED, f"Ref {ref} should be stored")
1982 # Simple prune.
1983 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
1984 self.assertFalse(butler.exists(ref1.datasetType, ref1.dataId, collections=run1))
1986 many_stored = butler.stored_many([ref1, ref2, ref3])
1987 for ref, stored in many_stored.items():
1988 self.assertFalse(stored, f"Ref {ref} should not be stored")
1990 many_exists = butler._exists_many([ref1, ref2, ref3])
1991 for ref, exists in many_exists.items():
1992 self.assertEqual(exists, DatasetExistence.UNRECOGNIZED, f"Ref {ref} should not be stored")
1994 # Put data back.
1995 ref1_new = butler.put(metric, ref1)
1996 self.assertEqual(ref1_new, ref1) # Reuses original ID.
1997 ref2 = butler.put(metric, ref2)
1999 many_stored = butler.stored_many([ref1, ref2, ref3])
2000 self.assertTrue(many_stored[ref1])
2001 self.assertTrue(many_stored[ref2])
2002 self.assertFalse(many_stored[ref3])
2004 ref3 = butler.put(metric, ref3)
2006 many_exists = butler._exists_many([ref1, ref2, ref3])
2007 for ref, exists in many_exists.items():
2008 self.assertTrue(exists, f"Ref {ref} should not be stored")
2010 # Clear out the datasets from registry and start again.
2011 refs = [ref1, ref2, ref3]
2012 butler.pruneDatasets(refs, purge=True, unstore=True)
2013 for ref in refs:
2014 butler.put(metric, ref)
2016 # Confirm we can retrieve deferred.
2017 dref1 = butler.getDeferred(ref1) # known and exists
2018 metric1 = dref1.get()
2019 self.assertEqual(metric1, metric)
2021 # Test different forms of file availability.
2022 # Need to be in a state where:
2023 # - one ref just has registry record.
2024 # - one ref has a missing file but a datastore record.
2025 # - one ref has a missing datastore record but file is there.
2026 # - one ref does not exist anywhere.
2027 # Do not need to test a ref that has everything since that is tested
2028 # above.
2029 ref0 = DatasetRef(
2030 datasetType,
2031 DataCoordinate.standardize(
2032 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, universe=butler.dimensions
2033 ),
2034 run=run1,
2035 )
2037 # Delete from datastore and retain in Registry.
2038 butler.pruneDatasets([ref1], purge=False, unstore=True, disassociate=False)
2040 # File has been removed.
2041 self.remove_dataset_out_of_band(butler, ref2)
2043 # Datastore has lost track.
2044 butler._datastore.forget([ref3])
2046 # First test with a standard butler.
2047 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
2048 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2049 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2050 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2051 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED)
2053 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=False)
2054 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2055 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
2056 self.assertEqual(exists_many[ref2], DatasetExistence.KNOWN)
2057 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
2058 self.assertTrue(exists_many[ref2])
2060 # Check that per-ref query gives the same answer as many query.
2061 for ref, exists in exists_many.items():
2062 self.assertEqual(butler.exists(ref, full_check=False), exists)
2064 # Get deferred checks for existence before it allows it to be
2065 # retrieved.
2066 with self.assertRaises(LookupError):
2067 butler.getDeferred(ref3) # not known, file exists
2068 dref2 = butler.getDeferred(ref2) # known but file missing
2069 with self.assertRaises(FileNotFoundError):
2070 dref2.get()
2072 # Test again with a trusting butler.
2073 if self.trustModeSupported:
2074 butler._datastore.trustGetRequest = True
2075 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
2076 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2077 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2078 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2079 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ARTIFACT)
2081 # When trusting we can get a deferred dataset handle that is not
2082 # known but does exist.
2083 dref3 = butler.getDeferred(ref3)
2084 metric3 = dref3.get()
2085 self.assertEqual(metric3, metric)
2087 # Check that per-ref query gives the same answer as many query.
2088 for ref, exists in exists_many.items():
2089 self.assertEqual(butler.exists(ref, full_check=True), exists)
2091 # Create a ref that surprisingly has the UUID of an existing ref
2092 # but is not the same.
2093 ref_bad = DatasetRef(datasetType, dataId=ref3.dataId, run=ref3.run, id=ref2.id)
2094 with self.assertRaises(ValueError):
2095 butler.exists(ref_bad)
2097 # Create a ref that has a compatible storage class.
2098 ref_compat = ref2.overrideStorageClass("StructuredDataDict")
2099 exists = butler.exists(ref_compat)
2100 self.assertEqual(exists, exists_many[ref2])
2102 # Remove everything and start from scratch.
2103 butler._datastore.trustGetRequest = False
2104 butler.pruneDatasets(refs, purge=True, unstore=True)
2105 for ref in refs:
2106 butler.put(metric, ref)
2108 # These tests mess directly with the trash table and can leave the
2109 # datastore in an odd state. Do them at the end.
2110 # Check that in normal mode, deleting the record will lead to
2111 # trash not touching the file.
2112 uri1 = butler.getURI(ref1)
2113 butler._datastore.bridge.moveToTrash(
2114 [ref1], transaction=None
2115 ) # Update the dataset_location table
2116 butler._datastore.forget([ref1])
2117 butler._datastore.trash(ref1)
2118 butler._datastore.emptyTrash()
2119 self.assertTrue(uri1.exists())
2120 uri1.remove() # Clean it up.
2122 # Simulate execution butler setup by deleting the datastore
2123 # record but keeping the file around and trusting.
2124 butler._datastore.trustGetRequest = True
2125 uris = butler.get_many_uris([ref2, ref3])
2126 uri2 = uris[ref2].primaryURI
2127 uri3 = uris[ref3].primaryURI
2128 self.assertTrue(uri2.exists())
2129 self.assertTrue(uri3.exists())
2131 # Remove the datastore record.
2132 butler._datastore.bridge.moveToTrash(
2133 [ref2], transaction=None
2134 ) # Update the dataset_location table
2135 butler._datastore.forget([ref2])
2136 self.assertTrue(uri2.exists())
2137 butler._datastore.trash([ref2, ref3])
2138 # Immediate removal for ref2 file
2139 self.assertFalse(uri2.exists())
2140 # But ref3 has to wait for the empty.
2141 self.assertTrue(uri3.exists())
2142 butler._datastore.emptyTrash()
2143 self.assertFalse(uri3.exists())
2145 # Clear out the datasets from registry.
2146 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
2148 def test_butler_metrics(self):
2149 """Test that metrics are collected."""
2150 run = "test_run"
2151 metrics = ButlerMetrics()
2152 butler, datasetType = self.create_butler(
2153 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics
2154 )
2155 data = MetricsExampleModel(
2156 summary={"AM1": 5.2, "AM2": 30.6},
2157 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2158 data=[563, 234, 456.7, 752, 8, 9, 27],
2159 )
2161 data_ref = butler.put(data, datasetType, visit=424, instrument="DummyCamComp")
2162 butler.get(data_ref)
2163 butler.get(data_ref)
2164 self.assertEqual(metrics.n_get, 2)
2165 self.assertGreater(metrics.time_in_get, 0.0)
2166 self.assertEqual(metrics.n_put, 1)
2167 self.assertGreater(metrics.time_in_put, 0.0)
2169 deferred = butler.getDeferred(data_ref)
2170 deferred.get()
2171 self.assertEqual(metrics.n_get, 3)
2173 with butler.record_metrics() as new:
2174 data_ref_2 = butler.put(data, datasetType, visit=425, instrument="DummyCamComp")
2175 butler.get(data_ref)
2177 butler.pruneDatasets([data_ref, data_ref_2], purge=True, unstore=True)
2178 with ResourcePath.temporary_uri(suffix=".json") as tmpFile:
2179 tmpFile.write(data.model_dump_json().encode())
2180 refs = [
2181 DatasetRef(datasetType, data_ref_2.dataId, run),
2182 DatasetRef(datasetType, data_ref.dataId, run),
2183 ]
2184 datasets = [FileDataset(path=tmpFile, refs=refs)]
2185 butler.ingest(*datasets, transfer="copy")
2187 self.assertEqual(new.n_get, 1)
2188 self.assertEqual(new.n_put, 1)
2189 self.assertEqual(new.n_ingest, 2)
2192class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2193 """PosixDatastore specialization of a butler"""
2195 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2196 fullConfigKey: str | None = ".datastore.formatters"
2197 validationCanFail = True
2198 datastoreStr = ["/tmp"]
2199 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2200 registryStr = "/gen3.sqlite3"
2202 def testPathConstructor(self) -> None:
2203 """Independent test of constructor using PathLike."""
2204 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
2205 self.enterContext(butler)
2206 self.assertIsInstance(butler, Butler)
2208 # And again with a Path object with the butler yaml
2209 path = pathlib.Path(self.tmpConfigFile)
2210 butler = Butler.from_config(path, writeable=False)
2211 self.enterContext(butler)
2212 self.assertIsInstance(butler, Butler)
2214 # And again with a Path object without the butler yaml
2215 # (making sure we skip it if the tmp config doesn't end
2216 # in butler.yaml -- which is the case for a subclass)
2217 if self.tmpConfigFile.endswith("butler.yaml"):
2218 path = pathlib.Path(os.path.dirname(self.tmpConfigFile))
2219 butler = Butler.from_config(path, writeable=False)
2220 self.enterContext(butler)
2221 self.assertIsInstance(butler, Butler)
2223 def testExportTransferCopy(self) -> None:
2224 """Test local export using all transfer modes"""
2225 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2226 exportButler = self.runPutGetTest(storageClass, "test_metric")
2227 # Test that the repo actually has at least one dataset.
2228 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
2229 self.assertGreater(len(datasets), 0)
2230 uris = [exportButler.getURI(d) for d in datasets]
2231 assert isinstance(exportButler._datastore, FileDatastore)
2232 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]]
2234 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris]
2236 for path in pathsInStore:
2237 # Assume local file system
2238 assert path is not None
2239 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}")
2241 for transfer in ("copy", "link", "symlink", "relsymlink"):
2242 with safeTestTempDir(TESTDIR) as exportDir:
2243 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export:
2244 export.saveDatasets(datasets)
2245 for path in pathsInStore:
2246 assert path is not None
2247 self.assertTrue(
2248 self.checkFileExists(exportDir, path),
2249 f"Check that mode {transfer} exported files",
2250 )
2252 def testPytypeCoercion(self) -> None:
2253 """Test python type coercion on Butler.get and put."""
2254 # Store some data with the normal example storage class.
2255 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2256 datasetTypeName = "test_metric"
2257 butler = self.runPutGetTest(storageClass, datasetTypeName)
2259 dataId = {"instrument": "DummyCamComp", "visit": 423}
2260 metric = butler.get(datasetTypeName, dataId=dataId)
2261 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
2263 datasetType_ori = butler.get_dataset_type(datasetTypeName)
2264 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
2266 # Now need to hack the registry dataset type definition.
2267 # There is no API for this.
2268 assert isinstance(butler._registry, SqlRegistry)
2269 manager = butler._registry._managers.datasets
2270 assert hasattr(manager, "_db") and hasattr(manager, "_static")
2271 manager._db.update(
2272 manager._static.dataset_type,
2273 {"name": datasetTypeName},
2274 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
2275 )
2277 # Force reset of dataset type cache
2278 butler.registry.refresh()
2280 datasetType_new = butler.get_dataset_type(datasetTypeName)
2281 self.assertEqual(datasetType_new.name, datasetType_ori.name)
2282 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
2284 metric_model = butler.get(datasetTypeName, dataId=dataId)
2285 self.assertNotEqual(type(metric_model), type(metric))
2286 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
2288 # Put the model and read it back to show that everything now
2289 # works as normal.
2290 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
2291 metric_model_new = butler.get(metric_ref)
2292 self.assertEqual(metric_model_new, metric_model)
2294 # Hack the storage class again to something that will fail on the
2295 # get with no conversion class.
2296 manager._db.update(
2297 manager._static.dataset_type,
2298 {"name": datasetTypeName},
2299 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
2300 )
2301 butler.registry.refresh()
2303 with self.assertRaises(ValueError):
2304 butler.get(datasetTypeName, dataId=dataId)
2306 def test_provenance(self):
2307 """Test that provenance is attached on put."""
2308 run = "test_run"
2309 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
2310 metric = MetricsExampleModel(
2311 summary={"AM1": 5.2, "AM2": 30.6},
2312 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2313 data=[563, 234, 456.7, 752, 8, 9, 27],
2314 )
2315 # Provenance can be attached to the object being put. Whether
2316 # it is or not is dependent on the formatter. For this test we
2317 # copy on adding provenance to ensure they differ.
2318 self.assertIsNone(metric.dataset_id)
2319 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
2320 self.assertIsNone(metric.dataset_id)
2321 metric_2 = butler.get(metric_ref)
2322 self.assertEqual(metric_2.data, metric.data)
2323 self.assertEqual(metric_2.dataset_id, metric_ref.id)
2324 self.assertIsNone(metric_2.provenance)
2326 # Put with provenance.
2327 prov = DatasetProvenance(quantum_id=uuid.uuid4())
2328 prov.add_input(metric_ref)
2329 prov.add_extra_provenance(metric_ref.id, {"answer": 42})
2330 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
2331 metric_3 = butler.get(metric_ref2)
2332 self.assertEqual(metric_3.provenance, prov)
2334 # Check that we can extract provenance from dict form.
2335 prov_dict = prov.to_flat_dict(metric_ref2)
2336 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2337 self.assertEqual(ref_from_prov, metric_ref2)
2338 # Direct __eq__ of the provenance does not work because one side
2339 # includes dimension records.
2340 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2341 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id)
2342 self.assertEqual(prov_from_prov.extras, prov.extras)
2344 # Force a bad ID into the dict.
2345 prov_dict["id"] = uuid.uuid4()
2346 with self.assertRaises(ValueError):
2347 DatasetProvenance.from_flat_dict(prov_dict, butler)
2348 del prov_dict["id"]
2349 prov_dict["input 0 id"] = uuid.uuid4()
2350 with self.assertRaises(ValueError):
2351 DatasetProvenance.from_flat_dict(prov_dict, butler)
2353 # Check that simple types can be reconstructed with non-standard
2354 # separators.
2355 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True)
2356 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2357 self.assertEqual(ref_from_prov, metric_ref2)
2358 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2360 with self.assertRaises(ValueError):
2361 DatasetProvenance.from_flat_dict({"unknown": 42}, butler)
2363 def test_specialized_file_datasets_functions(self):
2364 """Test a workflow used in Prompt Processing where we export datasets
2365 from one repository and write them in-place to the datastore of
2366 another, without immediately inserting registry entries for the
2367 datasets.
2368 """
2369 repo = MetricTestRepo.create_from_butler(
2370 self.create_empty_butler(writeable=True),
2371 self.tmpConfigFile,
2372 "StructuredCompositeReadCompNoDisassembly",
2373 )
2374 source_butler = repo.butler
2376 # Test writing outputs to a FileDatastore.
2377 with tempfile.TemporaryDirectory() as tempdir:
2378 target_repo_config = Butler.makeRepo(tempdir)
2379 refs = [repo.ref1, repo.ref2]
2380 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2381 self.assertEqual(len(datasets), 2)
2382 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2383 for dataset in datasets:
2384 path = ResourcePath(dataset.path, forceAbsolute=False)
2385 # Paths should be relative paths to the target datastore.
2386 self.assertFalse(path.isabs())
2387 # Files should have been copied into the target datastore
2388 self.assertTrue(ResourcePath(tempdir).join(path).exists())
2390 # Make sure the target Butler can ingest the datasets.
2391 target_butler = Butler(target_repo_config, writeable=True)
2392 self.enterContext(target_butler)
2393 target_butler.transfer_dimension_records_from(source_butler, refs)
2394 target_butler.ingest(*datasets, transfer=None)
2395 self.assertIsNotNone(target_butler.get(repo.ref1))
2396 self.assertIsNotNone(target_butler.get(repo.ref2))
2398 # Giving an empty list of files is a no-op.
2399 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), [])
2400 self.assertEqual(len(no_datasets), 0)
2402 # Test writing outputs to a ChainedDatastore.
2403 with tempfile.TemporaryDirectory() as tempdir:
2404 # Set up a second dataset type, so we can split the files across
2405 # multiple datastore roots.
2406 dt1 = repo.datasetType
2407 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass)
2408 source_butler.registry.registerDatasetType(dt2)
2409 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2)
2410 config = Config.fromString(
2411 f"""
2412 datastore:
2413 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore
2414 datastore_constraints:
2415 - constraints:
2416 accept:
2417 - {dt1.name}
2418 - constraints:
2419 accept:
2420 - {dt2.name}
2421 datastores:
2422 - datastore:
2423 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2424 root: <butlerRoot>/FileDatastore_0
2425 - datastore:
2426 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2427 root: <butlerRoot>/FileDatastore_1
2428 """
2429 )
2430 target_repo_config = Butler.makeRepo(tempdir, config)
2431 refs = [repo.ref1, repo.ref2, other_ref]
2432 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2433 self.assertEqual(len(datasets), 3)
2434 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2435 for dataset in datasets:
2436 path = ResourcePath(dataset.path, forceAbsolute=False)
2437 # Paths should be relative paths to the target datastore.
2438 self.assertFalse(path.isabs())
2439 # Files should have been split up between the two datastores
2440 # in the chain.
2441 datastore_root = ResourcePath(tempdir)
2442 if dataset.refs[0].datasetType.name == dt1.name:
2443 datastore_root = datastore_root.join("FileDatastore_0")
2444 else:
2445 datastore_root = datastore_root.join("FileDatastore_1")
2446 self.assertTrue(datastore_root.join(path).exists())
2448 # Make sure the target Butler can ingest the datasets.
2449 target_butler = Butler(target_repo_config, writeable=True)
2450 self.enterContext(target_butler)
2451 target_butler.transfer_dimension_records_from(source_butler, refs)
2452 target_butler.ingest(*datasets, transfer=None)
2453 self.assertIsNotNone(target_butler.get(repo.ref1))
2454 self.assertIsNotNone(target_butler.get(repo.ref2))
2455 self.assertIsNotNone(target_butler.get(other_ref))
2457 def test_temporary_for_ingest(self) -> None:
2458 """Test the `lsst.daf.butler._rubin.ingest_from_temporary` module."""
2459 with self.create_empty_butler("example_run") as butler:
2460 dataset_type = DatasetType("example", butler.dimensions.empty, "StructuredDataDict")
2461 butler.registry.registerDatasetType(dataset_type)
2462 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(butler.dimensions), "example_run")
2463 with TemporaryForIngest(butler, ref) as temporary:
2464 temporary.path.write(b"three: 3")
2465 found = TemporaryForIngest.find_orphaned_temporaries_by_ref(ref, butler)
2466 self.assertEqual(found, [temporary.path])
2467 self.assertIn(".tmp", temporary.ospath)
2468 temporary.ingest()
2469 loaded = butler.get(ref)
2470 self.assertEqual(loaded, {"three": 3})
2473class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2474 """PosixDatastore specialization of a butler using Postgres"""
2476 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2477 fullConfigKey = ".datastore.formatters"
2478 validationCanFail = True
2479 datastoreStr = ["/tmp"]
2480 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2481 registryStr = "PostgreSQL@test"
2483 @classmethod
2484 def setUpClass(cls) -> None:
2485 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
2486 super().setUpClass()
2488 def setUp(self) -> None:
2489 # Need to add a registry section to the config.
2490 self._temp_config = False
2491 config = Config(self.configFile)
2492 self.postgresql.patch_butler_config(config)
2493 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
2494 config.dump(fh)
2495 self.configFile = fh.name
2496 self._temp_config = True
2497 super().setUp()
2499 def tearDown(self) -> None:
2500 if self._temp_config and os.path.exists(self.configFile):
2501 os.remove(self.configFile)
2502 super().tearDown()
2504 def testMakeRepo(self) -> None:
2505 # The base class test assumes that it's using sqlite and assumes
2506 # the config file is acceptable to sqlite.
2507 raise unittest.SkipTest("Postgres config is not compatible with this test.")
2510class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase):
2511 """Test that Butler with a Postgres registry still works after cloning."""
2513 def create_butler(
2514 self,
2515 run: str,
2516 storageClass: StorageClass | str,
2517 datasetTypeName: str,
2518 metrics: ButlerMetrics | None = None,
2519 ) -> tuple[DirectButler, DatasetType]:
2520 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2521 return butler.clone(run=run, metrics=metrics), datasetType
2524class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
2525 """InMemoryDatastore specialization of a butler"""
2527 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")
2528 fullConfigKey = None
2529 useTempRoot = False
2530 validationCanFail = False
2531 datastoreStr = ["datastore='InMemory"]
2532 datastoreName = ["InMemoryDatastore@"]
2533 registryStr = "/gen3.sqlite3"
2535 def testIngest(self) -> None:
2536 pass
2538 def test_ingest_zip(self) -> None:
2539 pass
2542class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase):
2543 """Test that a Butler with a Sqlite registry still works after cloning."""
2545 def create_butler(
2546 self,
2547 run: str,
2548 storageClass: StorageClass | str,
2549 datasetTypeName: str,
2550 metrics: ButlerMetrics | None = None,
2551 ) -> tuple[DirectButler, DatasetType]:
2552 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2553 return butler.clone(run=run), datasetType
2556class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2557 """PosixDatastore specialization"""
2559 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
2560 fullConfigKey = ".datastore.datastores.1.formatters"
2561 validationCanFail = True
2562 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"]
2563 datastoreName = [
2564 "InMemoryDatastore@",
2565 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1",
2566 "SecondDatastore",
2567 ]
2568 registryStr = "/gen3.sqlite3"
2570 def testPruneDatasets(self) -> None:
2571 # This test relies on manipulating files out-of-band, which is
2572 # impossible for this configuration because of the InMemoryDatastore in
2573 # the ChainedDatastore.
2574 pass
2577class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase):
2578 """Test that a yaml file in one location can refer to a root in another."""
2580 datastoreStr = ["dir1"]
2581 # Disable the makeRepo test since we are deliberately not using
2582 # butler.yaml as the config name.
2583 fullConfigKey = None
2585 def setUp(self) -> None:
2586 self.root = makeTestTempDir(TESTDIR)
2588 # Make a new repository in one place
2589 self.dir1 = os.path.join(self.root, "dir1")
2590 Butler.makeRepo(self.dir1, config=Config(self.configFile))
2592 # Move the yaml file to a different place and add a "root"
2593 self.dir2 = os.path.join(self.root, "dir2")
2594 os.makedirs(self.dir2, exist_ok=True)
2595 configFile1 = os.path.join(self.dir1, "butler.yaml")
2596 config = Config(configFile1)
2597 config["root"] = self.dir1
2598 configFile2 = os.path.join(self.dir2, "butler2.yaml")
2599 config.dumpToUri(configFile2)
2600 os.remove(configFile1)
2601 self.tmpConfigFile = configFile2
2603 def testFileLocations(self) -> None:
2604 self.assertNotEqual(self.dir1, self.dir2)
2605 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml")))
2606 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml")))
2607 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3")))
2610class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase):
2611 """Test that a config file created by makeRepo outside of repo works."""
2613 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2615 def setUp(self) -> None:
2616 self.root = makeTestTempDir(TESTDIR)
2617 self.root2 = makeTestTempDir(TESTDIR)
2619 self.tmpConfigFile = os.path.join(self.root2, "different.yaml")
2620 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2622 def tearDown(self) -> None:
2623 if os.path.exists(self.root2):
2624 shutil.rmtree(self.root2, ignore_errors=True)
2625 super().tearDown()
2627 def testConfigExistence(self) -> None:
2628 c = Config(self.tmpConfigFile)
2629 uri_config = ResourcePath(c["root"])
2630 uri_expected = ResourcePath(self.root, forceDirectory=True)
2631 self.assertEqual(uri_config.geturl(), uri_expected.geturl())
2632 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path")
2634 def testPutGet(self) -> None:
2635 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2636 self.runPutGetTest(storageClass, "test_metric")
2639class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase):
2640 """Test that a config file created by makeRepo outside of repo works."""
2642 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2644 def setUp(self) -> None:
2645 self.root = makeTestTempDir(TESTDIR)
2646 self.root2 = makeTestTempDir(TESTDIR)
2648 self.tmpConfigFile = self.root2
2649 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2651 def testConfigExistence(self) -> None:
2652 # Append the yaml file else Config constructor does not know the file
2653 # type.
2654 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml")
2655 super().testConfigExistence()
2658class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase):
2659 """Test that a config file created by makeRepo outside of repo works."""
2661 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2663 def setUp(self) -> None:
2664 self.root = makeTestTempDir(TESTDIR)
2665 self.root2 = makeTestTempDir(TESTDIR)
2667 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl()
2668 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2671@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!")
2672class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2673 """S3Datastore specialization of a butler; an S3 storage Datastore +
2674 a local in-memory SqlRegistry.
2675 """
2677 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")
2678 fullConfigKey = None
2679 validationCanFail = True
2681 bucketName = "anybucketname"
2682 """Name of the Bucket that will be used in the tests. The name is read from
2683 the config file used with the tests during set-up.
2684 """
2686 root = "butlerRoot/"
2687 """Root repository directory expected to be used in case useTempRoot=False.
2688 Otherwise the root is set to a 20 characters long randomly generated string
2689 during set-up.
2690 """
2692 datastoreStr = [f"datastore={root}"]
2693 """Contains all expected root locations in a format expected to be
2694 returned by Butler stringification.
2695 """
2697 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]
2698 """The expected format of the S3 Datastore string."""
2700 registryStr = "/gen3.sqlite3"
2701 """Expected format of the Registry string."""
2703 mock_aws = mock_aws()
2704 """The mocked s3 interface from moto."""
2706 def genRoot(self) -> str:
2707 """Return a random string of len 20 to serve as a root
2708 name for the temporary bucket repo.
2710 This is equivalent to tempfile.mkdtemp as this is what self.root
2711 becomes when useTempRoot is True.
2712 """
2713 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20))
2714 return rndstr + "/"
2716 def setUp(self) -> None:
2717 config = Config(self.configFile)
2718 uri = ResourcePath(config[".datastore.datastore.root"])
2719 self.bucketName = uri.netloc
2721 # Enable S3 mocking of tests.
2722 self.enterContext(clean_test_environment_for_s3())
2723 self.mock_aws.start()
2725 if self.useTempRoot:
2726 self.root = self.genRoot()
2727 rooturi = f"s3://{self.bucketName}/{self.root}"
2728 config.update({"datastore": {"datastore": {"root": rooturi}}})
2730 # need local folder to store registry database
2731 self.reg_dir = makeTestTempDir(TESTDIR)
2732 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"
2734 # MOTO needs to know that we expect Bucket bucketname to exist
2735 # (this used to be the class attribute bucketName)
2736 s3 = boto3.resource("s3")
2737 s3.create_bucket(Bucket=self.bucketName)
2739 self.datastoreStr = [f"datastore='{rooturi}'"]
2740 self.datastoreName = [f"FileDatastore@{rooturi}"]
2741 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)
2742 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")
2744 def tearDown(self) -> None:
2745 s3 = boto3.resource("s3")
2746 bucket = s3.Bucket(self.bucketName)
2747 try:
2748 bucket.objects.all().delete()
2749 except botocore.exceptions.ClientError as e:
2750 if e.response["Error"]["Code"] == "404":
2751 # the key was not reachable - pass
2752 pass
2753 else:
2754 raise
2756 bucket = s3.Bucket(self.bucketName)
2757 bucket.delete()
2759 # Stop the S3 mock.
2760 self.mock_aws.stop()
2762 if self.reg_dir is not None and os.path.exists(self.reg_dir):
2763 shutil.rmtree(self.reg_dir, ignore_errors=True)
2765 if self.useTempRoot and os.path.exists(self.root):
2766 shutil.rmtree(self.root, ignore_errors=True)
2768 super().tearDown()
2771class DatastoreTransfers(TestCaseMixin):
2772 """Base test setup for data transfers between butlers. The concrete tests
2773 for specific configurations are in other classes, below.
2774 """
2776 storageClassFactory: StorageClassFactory
2778 @classmethod
2779 def setUpClass(cls) -> None:
2780 cls.storageClassFactory = StorageClassFactory()
2782 def setUp(self) -> None:
2783 self.root = makeTestTempDir(TESTDIR)
2784 self.config = Config(self.configFile)
2786 # Some tests cause convertors to be replaced so ensure
2787 # the storage class factory is reset each time.
2788 self.storageClassFactory.reset()
2789 self.storageClassFactory.addFromConfig(self.configFile)
2791 def tearDown(self) -> None:
2792 removeTestTempDir(self.root)
2794 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler:
2795 if manager is None:
2796 manager = (
2797 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
2798 )
2799 config = Config(config_file if config_file is not None else self.configFile)
2800 config["registry", "managers", "datasets"] = manager
2801 butler = Butler.from_config(
2802 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True
2803 )
2804 self.enterContext(butler)
2805 return butler
2807 def assertButlerTransfers(
2808 self,
2809 purge: bool = False,
2810 storageClassName: str = "StructuredData",
2811 storageClassNameTarget: str | None = None,
2812 ) -> None:
2813 """Test that a run can be transferred to another butler."""
2814 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
2815 if storageClassNameTarget is not None:
2816 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget)
2817 else:
2818 storageClassTarget = storageClass
2820 datasetTypeName = "random_data"
2822 # Test will create 3 collections and we will want to transfer
2823 # two of those three.
2824 runs = ["run1", "run2", "other"]
2826 # Also want to use two different dataset types to ensure that
2827 # grouping works.
2828 datasetTypeNames = ["random_data", "random_data_2"]
2830 # Create the run collections in the source butler.
2831 for run in runs:
2832 self.source_butler.collections.register(run)
2834 # Create dimensions in source butler.
2835 n_exposures = 30
2836 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
2837 self.source_butler.registry.insertDimensionData(
2838 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
2839 )
2840 self.source_butler.registry.insertDimensionData(
2841 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
2842 )
2843 self.source_butler.registry.insertDimensionData(
2844 "day_obs",
2845 {
2846 "instrument": "DummyCamComp",
2847 "id": 20250101,
2848 },
2849 )
2851 for i in range(n_exposures):
2852 self.source_butler.registry.insertDimensionData(
2853 "group", {"instrument": "DummyCamComp", "name": f"group{i}"}
2854 )
2855 self.source_butler.registry.insertDimensionData(
2856 "exposure",
2857 {
2858 "instrument": "DummyCamComp",
2859 "id": i,
2860 "obs_id": f"exp{i}",
2861 "physical_filter": "d-r",
2862 "group": f"group{i}",
2863 "day_obs": 20250101,
2864 },
2865 )
2867 # Create dataset types in the source butler.
2868 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"])
2869 for datasetTypeName in datasetTypeNames:
2870 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
2871 self.source_butler.registry.registerDatasetType(datasetType)
2873 # Write a dataset to an unrelated run -- this will ensure that
2874 # we are rewriting integer dataset ids in the target if necessary.
2875 # Will not be relevant for UUID.
2876 run = "distraction"
2877 butler = Butler.from_config(butler=self.source_butler, run=run)
2878 self.enterContext(butler)
2879 butler.put(
2880 makeExampleMetrics(),
2881 datasetTypeName,
2882 exposure=1,
2883 instrument="DummyCamComp",
2884 physical_filter="d-r",
2885 )
2887 # Write some example metrics to the source
2888 butler = Butler.from_config(butler=self.source_butler)
2889 self.enterContext(butler)
2891 # Set of DatasetRefs that should be in the list of refs to transfer
2892 # but which will not be transferred.
2893 deleted: set[DatasetRef] = set()
2895 n_expected = 20 # Number of datasets expected to be transferred
2896 source_refs = []
2897 for i in range(n_exposures):
2898 # Put a third of datasets into each collection, only retain
2899 # two thirds.
2900 index = i % 3
2901 run = runs[index]
2902 datasetTypeName = datasetTypeNames[i % 2]
2904 metric = MetricsExample(
2905 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)]
2906 )
2907 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"}
2908 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run)
2910 # Remove the datastore record using low-level API, but only
2911 # for a specific index.
2912 if purge and index == 1:
2913 # For one of these delete the file as well.
2914 # This allows the "missing" code to filter the
2915 # file out.
2916 # Access the individual datastores.
2917 datastores = []
2918 if hasattr(butler._datastore, "datastores"):
2919 datastores.extend(butler._datastore.datastores)
2920 else:
2921 datastores.append(butler._datastore)
2923 if not deleted:
2924 # For a chained datastore we need to remove
2925 # files in each chain.
2926 for datastore in datastores:
2927 # The file might not be known to the datastore
2928 # if constraints are used.
2929 try:
2930 primary, uris = datastore.getURIs(ref)
2931 except FileNotFoundError:
2932 continue
2933 if primary and primary.scheme != "mem":
2934 primary.remove()
2935 for uri in uris.values():
2936 if uri.scheme != "mem":
2937 uri.remove()
2938 n_expected -= 1
2939 deleted.add(ref)
2941 # Remove the datastore record.
2942 for datastore in datastores:
2943 if hasattr(datastore, "removeStoredItemInfo"):
2944 datastore.removeStoredItemInfo(ref)
2946 if index < 2:
2947 source_refs.append(ref)
2948 if ref not in deleted:
2949 new_metric = butler.get(ref)
2950 self.assertEqual(new_metric, metric)
2952 # Create some bad dataset types to ensure we check for inconsistent
2953 # definitions.
2954 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList")
2955 for datasetTypeName in datasetTypeNames:
2956 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass)
2957 self.target_butler.registry.registerDatasetType(datasetType)
2958 with self.assertRaises(ConflictingDefinitionError) as cm:
2959 self.target_butler.transfer_from(self.source_butler, source_refs)
2960 self.assertIn("dataset type differs", str(cm.exception))
2962 # And remove the bad definitions.
2963 for datasetTypeName in datasetTypeNames:
2964 self.target_butler.registry.removeDatasetType(datasetTypeName)
2966 # Transfer without creating dataset types should fail.
2967 with self.assertRaises(KeyError):
2968 self.target_butler.transfer_from(self.source_butler, source_refs)
2970 # Transfer without creating dimensions should fail.
2971 with self.assertRaises(ConflictingDefinitionError) as cm:
2972 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True)
2973 self.assertIn("dimension", str(cm.exception))
2975 # The dry run test requires dataset types to exist. If we have
2976 # been given distinct storage classes for the target we have
2977 # to redefine at least one of the dataset types in the target butler.
2978 if storageClass != storageClassTarget:
2979 self.target_butler.registry.removeDatasetType(datasetTypeNames[0])
2980 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget)
2981 self.target_butler.registry.registerDatasetType(datasetType)
2983 # The failed transfer above leaves registry in an inconsistent
2984 # state because the run is created but then rolled back without
2985 # the collection cache being cleared. For now force a refresh.
2986 # Can remove with DM-35498.
2987 self.target_butler.registry.refresh()
2989 # Do a dry run -- this should not have any effect on the target butler.
2990 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True)
2992 # Transfer the records for one ref to test the alternative API.
2993 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2994 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]])
2995 self.assertIn("number of records transferred: 1", ";".join(log_cm.output))
2997 # Now transfer them to the second butler, including dimensions.
2998 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2999 transferred = self.target_butler.transfer_from(
3000 self.source_butler,
3001 source_refs,
3002 register_dataset_types=True,
3003 transfer_dimensions=True,
3004 )
3005 self.assertEqual(len(transferred), n_expected)
3006 log_output = ";".join(log_cm.output)
3008 # A ChainedDatastore will use the in-memory datastore for mexists
3009 # so we can not rely on the mexists log message.
3010 self.assertIn("Number of datastore records found in source", log_output)
3011 self.assertIn("Creating output run", log_output)
3013 # Do the transfer twice to ensure that it will do nothing extra.
3014 # Only do this if purge=True because it does not work for int
3015 # dataset_id.
3016 if purge:
3017 # This should not need to register dataset types.
3018 transferred = self.target_butler.transfer_from(self.source_butler, source_refs)
3019 self.assertEqual(len(transferred), n_expected)
3021 with self.assertRaises((TypeError, AttributeError)):
3022 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore
3024 with self.assertRaises(ValueError):
3025 self.target_butler._datastore.transfer_from(
3026 self.source_butler._datastore, source_refs, transfer="split"
3027 )
3029 # Now try to get the same refs from the new butler.
3030 for ref in source_refs:
3031 if ref not in deleted:
3032 new_metric = self.target_butler.get(ref)
3033 old_metric = self.source_butler.get(ref)
3034 self.assertEqual(new_metric, old_metric)
3036 # Try again without implicit storage class conversion
3037 # triggered by using the source ref. This will do conversion
3038 # since the formatter will be returning the source python type.
3039 target_ref = self.target_butler.get_dataset(ref.id)
3040 if target_ref.datasetType.storageClass != ref.datasetType.storageClass:
3041 new_metric = self.target_butler.get(target_ref)
3042 self.assertNotEqual(type(new_metric), type(old_metric))
3044 # Remove the dataset from the target and put it again
3045 # as if it was the right type all along for this butler.
3046 self.target_butler.pruneDatasets(
3047 [target_ref], unstore=True, purge=True, disassociate=True
3048 )
3049 self.target_butler.put(new_metric, target_ref)
3050 new_new_metric = self.target_butler.get(target_ref)
3051 new_old_metric = self.target_butler.get(
3052 target_ref, storageClass=ref.datasetType.storageClass
3053 )
3054 self.assertEqual(new_new_metric, new_metric)
3055 self.assertEqual(new_old_metric, old_metric)
3057 # Now prune run2 collection and create instead a CHAINED collection.
3058 # This should block the transfer.
3059 self.target_butler.removeRuns(["run2"])
3060 self.target_butler.collections.register("run2", CollectionType.CHAINED)
3061 with self.assertRaises(CollectionTypeError):
3062 # Re-importing the run1 datasets can be problematic if they
3063 # use integer IDs so filter those out.
3064 to_transfer = [ref for ref in source_refs if ref.run == "run2"]
3065 self.target_butler.transfer_from(self.source_butler, to_transfer)
3068class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3069 """Test data transfers between butlers.
3071 Test for different managers. UUID to UUID and integer to integer are
3072 tested. UUID to integer is not supported since we do not currently
3073 want to allow that. Integer to UUID is supported with the caveat
3074 that UUID4 will be generated and this will be incorrect for raw
3075 dataset types. The test ignores that.
3076 """
3078 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3080 def create_butlers(
3081 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None
3082 ) -> None:
3083 self.source_butler = self.create_butler(manager1, "1", config_file=source_config)
3084 self.target_butler = self.create_butler(manager2, "2")
3086 def testTransferUuidToUuid(self) -> None:
3087 self.create_butlers()
3088 self.assertButlerTransfers()
3090 def testTransferFromChainedUuidToUuid(self) -> None:
3091 """Force the source butler to be a ChainedDatastore."""
3092 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml"))
3093 self.assertButlerTransfers()
3095 def testTransferFromIncompatibleUuidToUuid(self) -> None:
3096 """Force the source butler to be a incompatible datastore."""
3097 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml"))
3098 with self.assertRaises(NotImplementedError):
3099 self.assertButlerTransfers()
3101 def testTransferFromIncompatibleChainUuidToUuid(self) -> None:
3102 """Force the source butler to be a incompatible datastore."""
3103 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml"))
3104 with self.assertRaises(TypeError):
3105 self.assertButlerTransfers()
3107 def testTransferFromFileUuidToUuid(self) -> None:
3108 """Force the source butler to be a FileDatastore."""
3109 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml"))
3110 self.assertButlerTransfers()
3112 def testTransferMissing(self) -> None:
3113 """Test transfers where datastore records are missing.
3115 This is how execution butler works.
3116 """
3117 self.create_butlers()
3119 # Configure the source butler to allow trust.
3120 self.source_butler._datastore._set_trust_mode(True)
3122 self.assertButlerTransfers(purge=True)
3124 def testTransferMissingDisassembly(self) -> None:
3125 """Test transfers where datastore records are missing.
3127 This is how execution butler works.
3128 """
3129 self.create_butlers()
3131 # Configure the source butler to allow trust.
3132 self.source_butler._datastore._set_trust_mode(True)
3134 # Test disassembly.
3135 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite")
3137 def testTransferDifferingStorageClasses(self) -> None:
3138 """Test transfers when the source butler dataset type has a different
3139 but compatible storage class.
3140 """
3141 self.create_butlers()
3143 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion")
3145 def testTransferDifferingStorageClassesDisassembly(self) -> None:
3146 """Test transfers when the source butler dataset type has a different
3147 but compatible storage class and where the source butler has
3148 disassembled.
3149 """
3150 self.create_butlers()
3152 self.assertButlerTransfers(
3153 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion"
3154 )
3156 def testAbsoluteURITransferDirect(self) -> None:
3157 """Test transfer using an absolute URI."""
3158 self._absolute_transfer("auto")
3160 def testAbsoluteURITransferCopy(self) -> None:
3161 """Test transfer using an absolute URI."""
3162 self._absolute_transfer("copy")
3164 def _absolute_transfer(self, transfer: str) -> None:
3165 self.create_butlers()
3167 storageClassName = "StructuredData"
3168 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
3169 datasetTypeName = "random_data"
3170 run = "run1"
3171 self.source_butler.collections.register(run)
3173 dimensions = self.source_butler.dimensions.conform(())
3174 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3175 self.source_butler.registry.registerDatasetType(datasetType)
3177 metrics = makeExampleMetrics()
3178 with ResourcePath.temporary_uri(suffix=".json") as temp:
3179 dataId = DataCoordinate.make_empty(self.source_butler.dimensions)
3180 source_refs = [DatasetRef(datasetType, dataId, run=run)]
3181 temp.write(json.dumps(metrics.exportAsDict()).encode())
3182 dataset = FileDataset(path=temp, refs=source_refs)
3183 self.source_butler.ingest(dataset, transfer="direct")
3185 self.target_butler.transfer_from(
3186 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer
3187 )
3189 uri = self.target_butler.getURI(dataset.refs[0])
3190 if transfer == "auto":
3191 self.assertEqual(uri, temp)
3192 else:
3193 self.assertNotEqual(uri, temp)
3195 def test_shared_dimension_group(self):
3196 """Test internal logic that divides dataset types by dimension group
3197 when doing registry updates.
3198 """
3199 self.create_butlers()
3200 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
3201 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
3203 source_butler = self.source_butler
3204 target_butler = self.target_butler
3206 # Create a dataset type with the same dimensions as the 'bias' dataset
3207 # type from base.yaml
3208 dataset_type = DatasetType(
3209 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
3210 )
3211 source_butler.registry.registerDatasetType(dataset_type)
3212 # This has the same data ID as one of the bias datasets in
3213 # datasets.yaml.
3214 test_ref = source_butler.registry.insertDatasets(
3215 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
3216 )[0]
3218 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
3219 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
3220 refs = [test_ref, *biases, *flats]
3222 # Test setup will be even more convoluted if we want the datastore to
3223 # actually transfer files. For testing the dimension group behavior,
3224 # we really only care about the registry.
3225 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
3226 mock.return_value = (set(refs), set())
3227 target_butler.transfer_from(
3228 source_butler,
3229 refs,
3230 transfer=None,
3231 register_dataset_types=True,
3232 skip_missing=False,
3233 transfer_dimensions=True,
3234 )
3236 transferred_test_ref = target_butler.find_dataset(
3237 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3238 )
3239 self.assertEqual(transferred_test_ref.id, test_ref.id)
3241 transferred_bias = target_butler.find_dataset(
3242 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3243 )
3244 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
3246 transferred_flat = target_butler.find_dataset(
3247 "flat",
3248 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
3249 collections="imported_r",
3250 )
3251 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
3254class ChainedDatastoreTransfers(PosixDatastoreTransfers):
3255 """Test transfers using a chained datastore."""
3257 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
3260@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3261class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3262 """Test ``transfer_from`` involving Butler server."""
3264 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3266 def test_transfers_from_remote_to_direct(self) -> None:
3267 from lsst.daf.butler.remote_butler._remote_file_transfer_source import (
3268 mock_file_transfer_uris_for_unit_test,
3269 )
3271 self.target_butler = self.create_butler(None, "2")
3272 with create_test_server(TESTDIR) as server:
3273 self.source_butler = server.hybrid_butler
3275 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath:
3276 # The Butler server returns HTTP URIs with a domain name that
3277 # is not resolvable because there is no actual HTTP server
3278 # involved in these tests. Strip this first layer of
3279 # indirection, and return the target of the redirect instead.
3280 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers)
3281 return ResourcePath(str(response.next_request.url))
3283 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url):
3284 self.assertButlerTransfers()
3287class NullDatastoreTestCase(unittest.TestCase):
3288 """Test that we can fall back to a null datastore."""
3290 # Need a good config to create the repo.
3291 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3292 storageClassFactory: StorageClassFactory
3294 @classmethod
3295 def setUpClass(cls) -> None:
3296 cls.storageClassFactory = StorageClassFactory()
3297 cls.storageClassFactory.addFromConfig(cls.configFile)
3299 def setUp(self) -> None:
3300 """Create a new butler root for each test."""
3301 self.root = makeTestTempDir(TESTDIR)
3302 Butler.makeRepo(self.root, config=Config(self.configFile))
3304 def tearDown(self) -> None:
3305 removeTestTempDir(self.root)
3307 def test_fallback(self) -> None:
3308 # Read the butler config and mess with the datastore section.
3309 config_path = os.path.join(self.root, "butler.yaml")
3310 bad_config = Config(config_path)
3311 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore"
3312 bad_config.dumpToUri(config_path)
3314 with self.assertRaises(RuntimeError):
3315 Butler(self.root, without_datastore=False)
3317 with self.assertRaises(RuntimeError):
3318 Butler.from_config(self.root, without_datastore=False)
3320 butler = Butler.from_config(self.root, writeable=True, without_datastore=True)
3321 self.enterContext(butler)
3322 self.assertIsInstance(butler._datastore, NullDatastore)
3324 # Check that registry is working.
3325 butler.collections.register("MYRUN")
3326 collections = butler.collections.query("*")
3327 self.assertIn("MYRUN", set(collections))
3329 # Create a ref.
3330 dimensions = butler.dimensions.conform([])
3331 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
3332 datasetTypeName = "metric"
3333 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3334 butler.registry.registerDatasetType(datasetType)
3335 ref = DatasetRef(datasetType, {}, run="MYRUN")
3337 # Check that datastore will complain.
3338 with self.assertRaises(FileNotFoundError):
3339 butler.get(ref)
3340 with self.assertRaises(FileNotFoundError):
3341 butler.getURI(ref)
3344@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3345class ButlerServerTests(FileDatastoreButlerTests):
3346 """Test RemoteButler and Butler server."""
3348 configFile = None
3349 predictionSupported = False
3350 trustModeSupported = False
3352 postgres: TemporaryPostgresInstance | None
3354 def setUp(self):
3355 self.server_instance = self.enterContext(create_test_server(TESTDIR))
3357 def tearDown(self):
3358 pass
3360 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
3361 # S3 pre-signed URLs may end up with differing expiration times in the
3362 # query parameters, so ignore query parameters when comparing.
3363 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path
3365 def create_empty_butler(
3366 self,
3367 run: str | None = None,
3368 writeable: bool | None = None,
3369 metrics: ButlerMetrics | None = None,
3370 cleanup: bool = True,
3371 ) -> Butler:
3372 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics)
3374 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
3375 # Can't delete a file via S3 signed URLs, so we need to reach in
3376 # through DirectButler to delete the dataset.
3377 uri = self.server_instance.direct_butler.getURI(ref)
3378 uri.remove()
3380 def testConstructor(self):
3381 # RemoteButler constructor is tested in test_server.py and
3382 # test_remote_butler.py.
3383 pass
3385 def testDafButlerRepositories(self):
3386 # Loading of RemoteButler via repository index is tested in
3387 # test_server.py.
3388 pass
3390 def testGetDatasetTypes(self) -> None:
3391 # This is mostly a test of validateConfiguration, which is for
3392 # validating Datastore configuration and thus isn't relevant to
3393 # RemoteButler.
3394 pass
3396 def testMakeRepo(self) -> None:
3397 # Only applies to DirectButler.
3398 pass
3400 # Pickling not yet implemented for RemoteButler/HybridButler.
3401 @unittest.expectedFailure
3402 def testPickle(self) -> None:
3403 return super().testPickle()
3405 def testStringification(self) -> None:
3406 self.assertEqual(
3407 str(self.server_instance.remote_butler),
3408 "RemoteButler(https://test.example/api/butler/repo/testrepo/)",
3409 )
3411 def testTransaction(self) -> None:
3412 # Transactions will never be supported for RemoteButler.
3413 pass
3415 def testPutTemplates(self) -> None:
3416 # The Butler server instance is configured with different file naming
3417 # templates than this test is expecting.
3418 pass
3421@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3422class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase):
3423 """Tests for RemoteButler's registry shim, with a SQLite DB backing the
3424 server.
3425 """
3427 postgres = None
3430@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3431class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase):
3432 """Tests for RemoteButler's registry shim, with a Postgres DB backing the
3433 server.
3434 """
3436 @classmethod
3437 def setUpClass(cls):
3438 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
3439 super().setUpClass()
3442def setup_module(module: types.ModuleType) -> None:
3443 """Set up the module for pytest."""
3444 clean_environment()
3447def _get_test_data_path(filename: str) -> ResourcePath:
3448 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}")
3451if __name__ == "__main__":
3452 clean_environment()
3453 unittest.main()