Coverage for tests / test_butler.py: 13%
1779 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:11 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
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 butler.put(data, datasetType, visit=425, instrument="DummyCamComp")
2175 butler.get(data_ref)
2176 self.assertEqual(new.n_get, 1)
2177 self.assertEqual(new.n_put, 1)
2180class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2181 """PosixDatastore specialization of a butler"""
2183 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2184 fullConfigKey: str | None = ".datastore.formatters"
2185 validationCanFail = True
2186 datastoreStr = ["/tmp"]
2187 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2188 registryStr = "/gen3.sqlite3"
2190 def testPathConstructor(self) -> None:
2191 """Independent test of constructor using PathLike."""
2192 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
2193 self.enterContext(butler)
2194 self.assertIsInstance(butler, Butler)
2196 # And again with a Path object with the butler yaml
2197 path = pathlib.Path(self.tmpConfigFile)
2198 butler = Butler.from_config(path, writeable=False)
2199 self.enterContext(butler)
2200 self.assertIsInstance(butler, Butler)
2202 # And again with a Path object without the butler yaml
2203 # (making sure we skip it if the tmp config doesn't end
2204 # in butler.yaml -- which is the case for a subclass)
2205 if self.tmpConfigFile.endswith("butler.yaml"):
2206 path = pathlib.Path(os.path.dirname(self.tmpConfigFile))
2207 butler = Butler.from_config(path, writeable=False)
2208 self.enterContext(butler)
2209 self.assertIsInstance(butler, Butler)
2211 def testExportTransferCopy(self) -> None:
2212 """Test local export using all transfer modes"""
2213 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2214 exportButler = self.runPutGetTest(storageClass, "test_metric")
2215 # Test that the repo actually has at least one dataset.
2216 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
2217 self.assertGreater(len(datasets), 0)
2218 uris = [exportButler.getURI(d) for d in datasets]
2219 assert isinstance(exportButler._datastore, FileDatastore)
2220 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]]
2222 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris]
2224 for path in pathsInStore:
2225 # Assume local file system
2226 assert path is not None
2227 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}")
2229 for transfer in ("copy", "link", "symlink", "relsymlink"):
2230 with safeTestTempDir(TESTDIR) as exportDir:
2231 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export:
2232 export.saveDatasets(datasets)
2233 for path in pathsInStore:
2234 assert path is not None
2235 self.assertTrue(
2236 self.checkFileExists(exportDir, path),
2237 f"Check that mode {transfer} exported files",
2238 )
2240 def testPytypeCoercion(self) -> None:
2241 """Test python type coercion on Butler.get and put."""
2242 # Store some data with the normal example storage class.
2243 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2244 datasetTypeName = "test_metric"
2245 butler = self.runPutGetTest(storageClass, datasetTypeName)
2247 dataId = {"instrument": "DummyCamComp", "visit": 423}
2248 metric = butler.get(datasetTypeName, dataId=dataId)
2249 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
2251 datasetType_ori = butler.get_dataset_type(datasetTypeName)
2252 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
2254 # Now need to hack the registry dataset type definition.
2255 # There is no API for this.
2256 assert isinstance(butler._registry, SqlRegistry)
2257 manager = butler._registry._managers.datasets
2258 assert hasattr(manager, "_db") and hasattr(manager, "_static")
2259 manager._db.update(
2260 manager._static.dataset_type,
2261 {"name": datasetTypeName},
2262 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
2263 )
2265 # Force reset of dataset type cache
2266 butler.registry.refresh()
2268 datasetType_new = butler.get_dataset_type(datasetTypeName)
2269 self.assertEqual(datasetType_new.name, datasetType_ori.name)
2270 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
2272 metric_model = butler.get(datasetTypeName, dataId=dataId)
2273 self.assertNotEqual(type(metric_model), type(metric))
2274 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
2276 # Put the model and read it back to show that everything now
2277 # works as normal.
2278 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
2279 metric_model_new = butler.get(metric_ref)
2280 self.assertEqual(metric_model_new, metric_model)
2282 # Hack the storage class again to something that will fail on the
2283 # get with no conversion class.
2284 manager._db.update(
2285 manager._static.dataset_type,
2286 {"name": datasetTypeName},
2287 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
2288 )
2289 butler.registry.refresh()
2291 with self.assertRaises(ValueError):
2292 butler.get(datasetTypeName, dataId=dataId)
2294 def test_provenance(self):
2295 """Test that provenance is attached on put."""
2296 run = "test_run"
2297 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
2298 metric = MetricsExampleModel(
2299 summary={"AM1": 5.2, "AM2": 30.6},
2300 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2301 data=[563, 234, 456.7, 752, 8, 9, 27],
2302 )
2303 # Provenance can be attached to the object being put. Whether
2304 # it is or not is dependent on the formatter. For this test we
2305 # copy on adding provenance to ensure they differ.
2306 self.assertIsNone(metric.dataset_id)
2307 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
2308 self.assertIsNone(metric.dataset_id)
2309 metric_2 = butler.get(metric_ref)
2310 self.assertEqual(metric_2.data, metric.data)
2311 self.assertEqual(metric_2.dataset_id, metric_ref.id)
2312 self.assertIsNone(metric_2.provenance)
2314 # Put with provenance.
2315 prov = DatasetProvenance(quantum_id=uuid.uuid4())
2316 prov.add_input(metric_ref)
2317 prov.add_extra_provenance(metric_ref.id, {"answer": 42})
2318 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
2319 metric_3 = butler.get(metric_ref2)
2320 self.assertEqual(metric_3.provenance, prov)
2322 # Check that we can extract provenance from dict form.
2323 prov_dict = prov.to_flat_dict(metric_ref2)
2324 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2325 self.assertEqual(ref_from_prov, metric_ref2)
2326 # Direct __eq__ of the provenance does not work because one side
2327 # includes dimension records.
2328 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2329 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id)
2330 self.assertEqual(prov_from_prov.extras, prov.extras)
2332 # Force a bad ID into the dict.
2333 prov_dict["id"] = uuid.uuid4()
2334 with self.assertRaises(ValueError):
2335 DatasetProvenance.from_flat_dict(prov_dict, butler)
2336 del prov_dict["id"]
2337 prov_dict["input 0 id"] = uuid.uuid4()
2338 with self.assertRaises(ValueError):
2339 DatasetProvenance.from_flat_dict(prov_dict, butler)
2341 # Check that simple types can be reconstructed with non-standard
2342 # separators.
2343 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True)
2344 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2345 self.assertEqual(ref_from_prov, metric_ref2)
2346 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2348 with self.assertRaises(ValueError):
2349 DatasetProvenance.from_flat_dict({"unknown": 42}, butler)
2351 def test_specialized_file_datasets_functions(self):
2352 """Test a workflow used in Prompt Processing where we export datasets
2353 from one repository and write them in-place to the datastore of
2354 another, without immediately inserting registry entries for the
2355 datasets.
2356 """
2357 repo = MetricTestRepo.create_from_butler(
2358 self.create_empty_butler(writeable=True),
2359 self.tmpConfigFile,
2360 "StructuredCompositeReadCompNoDisassembly",
2361 )
2362 source_butler = repo.butler
2364 # Test writing outputs to a FileDatastore.
2365 with tempfile.TemporaryDirectory() as tempdir:
2366 target_repo_config = Butler.makeRepo(tempdir)
2367 refs = [repo.ref1, repo.ref2]
2368 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2369 self.assertEqual(len(datasets), 2)
2370 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2371 for dataset in datasets:
2372 path = ResourcePath(dataset.path, forceAbsolute=False)
2373 # Paths should be relative paths to the target datastore.
2374 self.assertFalse(path.isabs())
2375 # Files should have been copied into the target datastore
2376 self.assertTrue(ResourcePath(tempdir).join(path).exists())
2378 # Make sure the target Butler can ingest the datasets.
2379 target_butler = Butler(target_repo_config, writeable=True)
2380 self.enterContext(target_butler)
2381 target_butler.transfer_dimension_records_from(source_butler, refs)
2382 target_butler.ingest(*datasets, transfer=None)
2383 self.assertIsNotNone(target_butler.get(repo.ref1))
2384 self.assertIsNotNone(target_butler.get(repo.ref2))
2386 # Giving an empty list of files is a no-op.
2387 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), [])
2388 self.assertEqual(len(no_datasets), 0)
2390 # Test writing outputs to a ChainedDatastore.
2391 with tempfile.TemporaryDirectory() as tempdir:
2392 # Set up a second dataset type, so we can split the files across
2393 # multiple datastore roots.
2394 dt1 = repo.datasetType
2395 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass)
2396 source_butler.registry.registerDatasetType(dt2)
2397 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2)
2398 config = Config.fromString(
2399 f"""
2400 datastore:
2401 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore
2402 datastore_constraints:
2403 - constraints:
2404 accept:
2405 - {dt1.name}
2406 - constraints:
2407 accept:
2408 - {dt2.name}
2409 datastores:
2410 - datastore:
2411 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2412 root: <butlerRoot>/FileDatastore_0
2413 - datastore:
2414 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2415 root: <butlerRoot>/FileDatastore_1
2416 """
2417 )
2418 target_repo_config = Butler.makeRepo(tempdir, config)
2419 refs = [repo.ref1, repo.ref2, other_ref]
2420 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2421 self.assertEqual(len(datasets), 3)
2422 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2423 for dataset in datasets:
2424 path = ResourcePath(dataset.path, forceAbsolute=False)
2425 # Paths should be relative paths to the target datastore.
2426 self.assertFalse(path.isabs())
2427 # Files should have been split up between the two datastores
2428 # in the chain.
2429 datastore_root = ResourcePath(tempdir)
2430 if dataset.refs[0].datasetType.name == dt1.name:
2431 datastore_root = datastore_root.join("FileDatastore_0")
2432 else:
2433 datastore_root = datastore_root.join("FileDatastore_1")
2434 self.assertTrue(datastore_root.join(path).exists())
2436 # Make sure the target Butler can ingest the datasets.
2437 target_butler = Butler(target_repo_config, writeable=True)
2438 self.enterContext(target_butler)
2439 target_butler.transfer_dimension_records_from(source_butler, refs)
2440 target_butler.ingest(*datasets, transfer=None)
2441 self.assertIsNotNone(target_butler.get(repo.ref1))
2442 self.assertIsNotNone(target_butler.get(repo.ref2))
2443 self.assertIsNotNone(target_butler.get(other_ref))
2445 def test_temporary_for_ingest(self) -> None:
2446 """Test the `lsst.daf.butler._rubin.ingest_from_temporary` module."""
2447 with self.create_empty_butler("example_run") as butler:
2448 dataset_type = DatasetType("example", butler.dimensions.empty, "StructuredDataDict")
2449 butler.registry.registerDatasetType(dataset_type)
2450 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(butler.dimensions), "example_run")
2451 with TemporaryForIngest(butler, ref) as temporary:
2452 temporary.path.write(b"three: 3")
2453 found = TemporaryForIngest.find_orphaned_temporaries_by_ref(ref, butler)
2454 self.assertEqual(found, [temporary.path])
2455 self.assertIn(".tmp", temporary.ospath)
2456 temporary.ingest()
2457 loaded = butler.get(ref)
2458 self.assertEqual(loaded, {"three": 3})
2461class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2462 """PosixDatastore specialization of a butler using Postgres"""
2464 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2465 fullConfigKey = ".datastore.formatters"
2466 validationCanFail = True
2467 datastoreStr = ["/tmp"]
2468 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2469 registryStr = "PostgreSQL@test"
2471 @classmethod
2472 def setUpClass(cls) -> None:
2473 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
2474 super().setUpClass()
2476 def setUp(self) -> None:
2477 # Need to add a registry section to the config.
2478 self._temp_config = False
2479 config = Config(self.configFile)
2480 self.postgresql.patch_butler_config(config)
2481 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
2482 config.dump(fh)
2483 self.configFile = fh.name
2484 self._temp_config = True
2485 super().setUp()
2487 def tearDown(self) -> None:
2488 if self._temp_config and os.path.exists(self.configFile):
2489 os.remove(self.configFile)
2490 super().tearDown()
2492 def testMakeRepo(self) -> None:
2493 # The base class test assumes that it's using sqlite and assumes
2494 # the config file is acceptable to sqlite.
2495 raise unittest.SkipTest("Postgres config is not compatible with this test.")
2498class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase):
2499 """Test that Butler with a Postgres registry still works after cloning."""
2501 def create_butler(
2502 self,
2503 run: str,
2504 storageClass: StorageClass | str,
2505 datasetTypeName: str,
2506 metrics: ButlerMetrics | None = None,
2507 ) -> tuple[DirectButler, DatasetType]:
2508 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2509 return butler.clone(run=run, metrics=metrics), datasetType
2512class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
2513 """InMemoryDatastore specialization of a butler"""
2515 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")
2516 fullConfigKey = None
2517 useTempRoot = False
2518 validationCanFail = False
2519 datastoreStr = ["datastore='InMemory"]
2520 datastoreName = ["InMemoryDatastore@"]
2521 registryStr = "/gen3.sqlite3"
2523 def testIngest(self) -> None:
2524 pass
2526 def test_ingest_zip(self) -> None:
2527 pass
2530class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase):
2531 """Test that a Butler with a Sqlite registry still works after cloning."""
2533 def create_butler(
2534 self,
2535 run: str,
2536 storageClass: StorageClass | str,
2537 datasetTypeName: str,
2538 metrics: ButlerMetrics | None = None,
2539 ) -> tuple[DirectButler, DatasetType]:
2540 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2541 return butler.clone(run=run), datasetType
2544class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2545 """PosixDatastore specialization"""
2547 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
2548 fullConfigKey = ".datastore.datastores.1.formatters"
2549 validationCanFail = True
2550 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"]
2551 datastoreName = [
2552 "InMemoryDatastore@",
2553 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1",
2554 "SecondDatastore",
2555 ]
2556 registryStr = "/gen3.sqlite3"
2558 def testPruneDatasets(self) -> None:
2559 # This test relies on manipulating files out-of-band, which is
2560 # impossible for this configuration because of the InMemoryDatastore in
2561 # the ChainedDatastore.
2562 pass
2565class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase):
2566 """Test that a yaml file in one location can refer to a root in another."""
2568 datastoreStr = ["dir1"]
2569 # Disable the makeRepo test since we are deliberately not using
2570 # butler.yaml as the config name.
2571 fullConfigKey = None
2573 def setUp(self) -> None:
2574 self.root = makeTestTempDir(TESTDIR)
2576 # Make a new repository in one place
2577 self.dir1 = os.path.join(self.root, "dir1")
2578 Butler.makeRepo(self.dir1, config=Config(self.configFile))
2580 # Move the yaml file to a different place and add a "root"
2581 self.dir2 = os.path.join(self.root, "dir2")
2582 os.makedirs(self.dir2, exist_ok=True)
2583 configFile1 = os.path.join(self.dir1, "butler.yaml")
2584 config = Config(configFile1)
2585 config["root"] = self.dir1
2586 configFile2 = os.path.join(self.dir2, "butler2.yaml")
2587 config.dumpToUri(configFile2)
2588 os.remove(configFile1)
2589 self.tmpConfigFile = configFile2
2591 def testFileLocations(self) -> None:
2592 self.assertNotEqual(self.dir1, self.dir2)
2593 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml")))
2594 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml")))
2595 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3")))
2598class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase):
2599 """Test that a config file created by makeRepo outside of repo works."""
2601 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2603 def setUp(self) -> None:
2604 self.root = makeTestTempDir(TESTDIR)
2605 self.root2 = makeTestTempDir(TESTDIR)
2607 self.tmpConfigFile = os.path.join(self.root2, "different.yaml")
2608 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2610 def tearDown(self) -> None:
2611 if os.path.exists(self.root2):
2612 shutil.rmtree(self.root2, ignore_errors=True)
2613 super().tearDown()
2615 def testConfigExistence(self) -> None:
2616 c = Config(self.tmpConfigFile)
2617 uri_config = ResourcePath(c["root"])
2618 uri_expected = ResourcePath(self.root, forceDirectory=True)
2619 self.assertEqual(uri_config.geturl(), uri_expected.geturl())
2620 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path")
2622 def testPutGet(self) -> None:
2623 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2624 self.runPutGetTest(storageClass, "test_metric")
2627class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase):
2628 """Test that a config file created by makeRepo outside of repo works."""
2630 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2632 def setUp(self) -> None:
2633 self.root = makeTestTempDir(TESTDIR)
2634 self.root2 = makeTestTempDir(TESTDIR)
2636 self.tmpConfigFile = self.root2
2637 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2639 def testConfigExistence(self) -> None:
2640 # Append the yaml file else Config constructor does not know the file
2641 # type.
2642 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml")
2643 super().testConfigExistence()
2646class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase):
2647 """Test that a config file created by makeRepo outside of repo works."""
2649 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2651 def setUp(self) -> None:
2652 self.root = makeTestTempDir(TESTDIR)
2653 self.root2 = makeTestTempDir(TESTDIR)
2655 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl()
2656 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2659@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!")
2660class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2661 """S3Datastore specialization of a butler; an S3 storage Datastore +
2662 a local in-memory SqlRegistry.
2663 """
2665 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")
2666 fullConfigKey = None
2667 validationCanFail = True
2669 bucketName = "anybucketname"
2670 """Name of the Bucket that will be used in the tests. The name is read from
2671 the config file used with the tests during set-up.
2672 """
2674 root = "butlerRoot/"
2675 """Root repository directory expected to be used in case useTempRoot=False.
2676 Otherwise the root is set to a 20 characters long randomly generated string
2677 during set-up.
2678 """
2680 datastoreStr = [f"datastore={root}"]
2681 """Contains all expected root locations in a format expected to be
2682 returned by Butler stringification.
2683 """
2685 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]
2686 """The expected format of the S3 Datastore string."""
2688 registryStr = "/gen3.sqlite3"
2689 """Expected format of the Registry string."""
2691 mock_aws = mock_aws()
2692 """The mocked s3 interface from moto."""
2694 def genRoot(self) -> str:
2695 """Return a random string of len 20 to serve as a root
2696 name for the temporary bucket repo.
2698 This is equivalent to tempfile.mkdtemp as this is what self.root
2699 becomes when useTempRoot is True.
2700 """
2701 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20))
2702 return rndstr + "/"
2704 def setUp(self) -> None:
2705 config = Config(self.configFile)
2706 uri = ResourcePath(config[".datastore.datastore.root"])
2707 self.bucketName = uri.netloc
2709 # Enable S3 mocking of tests.
2710 self.enterContext(clean_test_environment_for_s3())
2711 self.mock_aws.start()
2713 if self.useTempRoot:
2714 self.root = self.genRoot()
2715 rooturi = f"s3://{self.bucketName}/{self.root}"
2716 config.update({"datastore": {"datastore": {"root": rooturi}}})
2718 # need local folder to store registry database
2719 self.reg_dir = makeTestTempDir(TESTDIR)
2720 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"
2722 # MOTO needs to know that we expect Bucket bucketname to exist
2723 # (this used to be the class attribute bucketName)
2724 s3 = boto3.resource("s3")
2725 s3.create_bucket(Bucket=self.bucketName)
2727 self.datastoreStr = [f"datastore='{rooturi}'"]
2728 self.datastoreName = [f"FileDatastore@{rooturi}"]
2729 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)
2730 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")
2732 def tearDown(self) -> None:
2733 s3 = boto3.resource("s3")
2734 bucket = s3.Bucket(self.bucketName)
2735 try:
2736 bucket.objects.all().delete()
2737 except botocore.exceptions.ClientError as e:
2738 if e.response["Error"]["Code"] == "404":
2739 # the key was not reachable - pass
2740 pass
2741 else:
2742 raise
2744 bucket = s3.Bucket(self.bucketName)
2745 bucket.delete()
2747 # Stop the S3 mock.
2748 self.mock_aws.stop()
2750 if self.reg_dir is not None and os.path.exists(self.reg_dir):
2751 shutil.rmtree(self.reg_dir, ignore_errors=True)
2753 if self.useTempRoot and os.path.exists(self.root):
2754 shutil.rmtree(self.root, ignore_errors=True)
2756 super().tearDown()
2759class DatastoreTransfers(TestCaseMixin):
2760 """Base test setup for data transfers between butlers. The concrete tests
2761 for specific configurations are in other classes, below.
2762 """
2764 storageClassFactory: StorageClassFactory
2766 @classmethod
2767 def setUpClass(cls) -> None:
2768 cls.storageClassFactory = StorageClassFactory()
2770 def setUp(self) -> None:
2771 self.root = makeTestTempDir(TESTDIR)
2772 self.config = Config(self.configFile)
2774 # Some tests cause convertors to be replaced so ensure
2775 # the storage class factory is reset each time.
2776 self.storageClassFactory.reset()
2777 self.storageClassFactory.addFromConfig(self.configFile)
2779 def tearDown(self) -> None:
2780 removeTestTempDir(self.root)
2782 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler:
2783 if manager is None:
2784 manager = (
2785 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
2786 )
2787 config = Config(config_file if config_file is not None else self.configFile)
2788 config["registry", "managers", "datasets"] = manager
2789 butler = Butler.from_config(
2790 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True
2791 )
2792 self.enterContext(butler)
2793 return butler
2795 def assertButlerTransfers(
2796 self,
2797 purge: bool = False,
2798 storageClassName: str = "StructuredData",
2799 storageClassNameTarget: str | None = None,
2800 ) -> None:
2801 """Test that a run can be transferred to another butler."""
2802 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
2803 if storageClassNameTarget is not None:
2804 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget)
2805 else:
2806 storageClassTarget = storageClass
2808 datasetTypeName = "random_data"
2810 # Test will create 3 collections and we will want to transfer
2811 # two of those three.
2812 runs = ["run1", "run2", "other"]
2814 # Also want to use two different dataset types to ensure that
2815 # grouping works.
2816 datasetTypeNames = ["random_data", "random_data_2"]
2818 # Create the run collections in the source butler.
2819 for run in runs:
2820 self.source_butler.collections.register(run)
2822 # Create dimensions in source butler.
2823 n_exposures = 30
2824 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
2825 self.source_butler.registry.insertDimensionData(
2826 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
2827 )
2828 self.source_butler.registry.insertDimensionData(
2829 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
2830 )
2831 self.source_butler.registry.insertDimensionData(
2832 "day_obs",
2833 {
2834 "instrument": "DummyCamComp",
2835 "id": 20250101,
2836 },
2837 )
2839 for i in range(n_exposures):
2840 self.source_butler.registry.insertDimensionData(
2841 "group", {"instrument": "DummyCamComp", "name": f"group{i}"}
2842 )
2843 self.source_butler.registry.insertDimensionData(
2844 "exposure",
2845 {
2846 "instrument": "DummyCamComp",
2847 "id": i,
2848 "obs_id": f"exp{i}",
2849 "physical_filter": "d-r",
2850 "group": f"group{i}",
2851 "day_obs": 20250101,
2852 },
2853 )
2855 # Create dataset types in the source butler.
2856 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"])
2857 for datasetTypeName in datasetTypeNames:
2858 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
2859 self.source_butler.registry.registerDatasetType(datasetType)
2861 # Write a dataset to an unrelated run -- this will ensure that
2862 # we are rewriting integer dataset ids in the target if necessary.
2863 # Will not be relevant for UUID.
2864 run = "distraction"
2865 butler = Butler.from_config(butler=self.source_butler, run=run)
2866 self.enterContext(butler)
2867 butler.put(
2868 makeExampleMetrics(),
2869 datasetTypeName,
2870 exposure=1,
2871 instrument="DummyCamComp",
2872 physical_filter="d-r",
2873 )
2875 # Write some example metrics to the source
2876 butler = Butler.from_config(butler=self.source_butler)
2877 self.enterContext(butler)
2879 # Set of DatasetRefs that should be in the list of refs to transfer
2880 # but which will not be transferred.
2881 deleted: set[DatasetRef] = set()
2883 n_expected = 20 # Number of datasets expected to be transferred
2884 source_refs = []
2885 for i in range(n_exposures):
2886 # Put a third of datasets into each collection, only retain
2887 # two thirds.
2888 index = i % 3
2889 run = runs[index]
2890 datasetTypeName = datasetTypeNames[i % 2]
2892 metric = MetricsExample(
2893 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)]
2894 )
2895 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"}
2896 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run)
2898 # Remove the datastore record using low-level API, but only
2899 # for a specific index.
2900 if purge and index == 1:
2901 # For one of these delete the file as well.
2902 # This allows the "missing" code to filter the
2903 # file out.
2904 # Access the individual datastores.
2905 datastores = []
2906 if hasattr(butler._datastore, "datastores"):
2907 datastores.extend(butler._datastore.datastores)
2908 else:
2909 datastores.append(butler._datastore)
2911 if not deleted:
2912 # For a chained datastore we need to remove
2913 # files in each chain.
2914 for datastore in datastores:
2915 # The file might not be known to the datastore
2916 # if constraints are used.
2917 try:
2918 primary, uris = datastore.getURIs(ref)
2919 except FileNotFoundError:
2920 continue
2921 if primary and primary.scheme != "mem":
2922 primary.remove()
2923 for uri in uris.values():
2924 if uri.scheme != "mem":
2925 uri.remove()
2926 n_expected -= 1
2927 deleted.add(ref)
2929 # Remove the datastore record.
2930 for datastore in datastores:
2931 if hasattr(datastore, "removeStoredItemInfo"):
2932 datastore.removeStoredItemInfo(ref)
2934 if index < 2:
2935 source_refs.append(ref)
2936 if ref not in deleted:
2937 new_metric = butler.get(ref)
2938 self.assertEqual(new_metric, metric)
2940 # Create some bad dataset types to ensure we check for inconsistent
2941 # definitions.
2942 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList")
2943 for datasetTypeName in datasetTypeNames:
2944 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass)
2945 self.target_butler.registry.registerDatasetType(datasetType)
2946 with self.assertRaises(ConflictingDefinitionError) as cm:
2947 self.target_butler.transfer_from(self.source_butler, source_refs)
2948 self.assertIn("dataset type differs", str(cm.exception))
2950 # And remove the bad definitions.
2951 for datasetTypeName in datasetTypeNames:
2952 self.target_butler.registry.removeDatasetType(datasetTypeName)
2954 # Transfer without creating dataset types should fail.
2955 with self.assertRaises(KeyError):
2956 self.target_butler.transfer_from(self.source_butler, source_refs)
2958 # Transfer without creating dimensions should fail.
2959 with self.assertRaises(ConflictingDefinitionError) as cm:
2960 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True)
2961 self.assertIn("dimension", str(cm.exception))
2963 # The dry run test requires dataset types to exist. If we have
2964 # been given distinct storage classes for the target we have
2965 # to redefine at least one of the dataset types in the target butler.
2966 if storageClass != storageClassTarget:
2967 self.target_butler.registry.removeDatasetType(datasetTypeNames[0])
2968 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget)
2969 self.target_butler.registry.registerDatasetType(datasetType)
2971 # The failed transfer above leaves registry in an inconsistent
2972 # state because the run is created but then rolled back without
2973 # the collection cache being cleared. For now force a refresh.
2974 # Can remove with DM-35498.
2975 self.target_butler.registry.refresh()
2977 # Do a dry run -- this should not have any effect on the target butler.
2978 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True)
2980 # Transfer the records for one ref to test the alternative API.
2981 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2982 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]])
2983 self.assertIn("number of records transferred: 1", ";".join(log_cm.output))
2985 # Now transfer them to the second butler, including dimensions.
2986 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2987 transferred = self.target_butler.transfer_from(
2988 self.source_butler,
2989 source_refs,
2990 register_dataset_types=True,
2991 transfer_dimensions=True,
2992 )
2993 self.assertEqual(len(transferred), n_expected)
2994 log_output = ";".join(log_cm.output)
2996 # A ChainedDatastore will use the in-memory datastore for mexists
2997 # so we can not rely on the mexists log message.
2998 self.assertIn("Number of datastore records found in source", log_output)
2999 self.assertIn("Creating output run", log_output)
3001 # Do the transfer twice to ensure that it will do nothing extra.
3002 # Only do this if purge=True because it does not work for int
3003 # dataset_id.
3004 if purge:
3005 # This should not need to register dataset types.
3006 transferred = self.target_butler.transfer_from(self.source_butler, source_refs)
3007 self.assertEqual(len(transferred), n_expected)
3009 with self.assertRaises((TypeError, AttributeError)):
3010 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore
3012 with self.assertRaises(ValueError):
3013 self.target_butler._datastore.transfer_from(
3014 self.source_butler._datastore, source_refs, transfer="split"
3015 )
3017 # Now try to get the same refs from the new butler.
3018 for ref in source_refs:
3019 if ref not in deleted:
3020 new_metric = self.target_butler.get(ref)
3021 old_metric = self.source_butler.get(ref)
3022 self.assertEqual(new_metric, old_metric)
3024 # Try again without implicit storage class conversion
3025 # triggered by using the source ref. This will do conversion
3026 # since the formatter will be returning the source python type.
3027 target_ref = self.target_butler.get_dataset(ref.id)
3028 if target_ref.datasetType.storageClass != ref.datasetType.storageClass:
3029 new_metric = self.target_butler.get(target_ref)
3030 self.assertNotEqual(type(new_metric), type(old_metric))
3032 # Remove the dataset from the target and put it again
3033 # as if it was the right type all along for this butler.
3034 self.target_butler.pruneDatasets(
3035 [target_ref], unstore=True, purge=True, disassociate=True
3036 )
3037 self.target_butler.put(new_metric, target_ref)
3038 new_new_metric = self.target_butler.get(target_ref)
3039 new_old_metric = self.target_butler.get(
3040 target_ref, storageClass=ref.datasetType.storageClass
3041 )
3042 self.assertEqual(new_new_metric, new_metric)
3043 self.assertEqual(new_old_metric, old_metric)
3045 # Now prune run2 collection and create instead a CHAINED collection.
3046 # This should block the transfer.
3047 self.target_butler.removeRuns(["run2"])
3048 self.target_butler.collections.register("run2", CollectionType.CHAINED)
3049 with self.assertRaises(CollectionTypeError):
3050 # Re-importing the run1 datasets can be problematic if they
3051 # use integer IDs so filter those out.
3052 to_transfer = [ref for ref in source_refs if ref.run == "run2"]
3053 self.target_butler.transfer_from(self.source_butler, to_transfer)
3056class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3057 """Test data transfers between butlers.
3059 Test for different managers. UUID to UUID and integer to integer are
3060 tested. UUID to integer is not supported since we do not currently
3061 want to allow that. Integer to UUID is supported with the caveat
3062 that UUID4 will be generated and this will be incorrect for raw
3063 dataset types. The test ignores that.
3064 """
3066 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3068 def create_butlers(
3069 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None
3070 ) -> None:
3071 self.source_butler = self.create_butler(manager1, "1", config_file=source_config)
3072 self.target_butler = self.create_butler(manager2, "2")
3074 def testTransferUuidToUuid(self) -> None:
3075 self.create_butlers()
3076 self.assertButlerTransfers()
3078 def testTransferFromChainedUuidToUuid(self) -> None:
3079 """Force the source butler to be a ChainedDatastore."""
3080 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml"))
3081 self.assertButlerTransfers()
3083 def testTransferFromIncompatibleUuidToUuid(self) -> None:
3084 """Force the source butler to be a incompatible datastore."""
3085 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml"))
3086 with self.assertRaises(NotImplementedError):
3087 self.assertButlerTransfers()
3089 def testTransferFromIncompatibleChainUuidToUuid(self) -> None:
3090 """Force the source butler to be a incompatible datastore."""
3091 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml"))
3092 with self.assertRaises(TypeError):
3093 self.assertButlerTransfers()
3095 def testTransferFromFileUuidToUuid(self) -> None:
3096 """Force the source butler to be a FileDatastore."""
3097 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml"))
3098 self.assertButlerTransfers()
3100 def testTransferMissing(self) -> None:
3101 """Test transfers where datastore records are missing.
3103 This is how execution butler works.
3104 """
3105 self.create_butlers()
3107 # Configure the source butler to allow trust.
3108 self.source_butler._datastore._set_trust_mode(True)
3110 self.assertButlerTransfers(purge=True)
3112 def testTransferMissingDisassembly(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 # Test disassembly.
3123 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite")
3125 def testTransferDifferingStorageClasses(self) -> None:
3126 """Test transfers when the source butler dataset type has a different
3127 but compatible storage class.
3128 """
3129 self.create_butlers()
3131 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion")
3133 def testTransferDifferingStorageClassesDisassembly(self) -> None:
3134 """Test transfers when the source butler dataset type has a different
3135 but compatible storage class and where the source butler has
3136 disassembled.
3137 """
3138 self.create_butlers()
3140 self.assertButlerTransfers(
3141 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion"
3142 )
3144 def testAbsoluteURITransferDirect(self) -> None:
3145 """Test transfer using an absolute URI."""
3146 self._absolute_transfer("auto")
3148 def testAbsoluteURITransferCopy(self) -> None:
3149 """Test transfer using an absolute URI."""
3150 self._absolute_transfer("copy")
3152 def _absolute_transfer(self, transfer: str) -> None:
3153 self.create_butlers()
3155 storageClassName = "StructuredData"
3156 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
3157 datasetTypeName = "random_data"
3158 run = "run1"
3159 self.source_butler.collections.register(run)
3161 dimensions = self.source_butler.dimensions.conform(())
3162 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3163 self.source_butler.registry.registerDatasetType(datasetType)
3165 metrics = makeExampleMetrics()
3166 with ResourcePath.temporary_uri(suffix=".json") as temp:
3167 dataId = DataCoordinate.make_empty(self.source_butler.dimensions)
3168 source_refs = [DatasetRef(datasetType, dataId, run=run)]
3169 temp.write(json.dumps(metrics.exportAsDict()).encode())
3170 dataset = FileDataset(path=temp, refs=source_refs)
3171 self.source_butler.ingest(dataset, transfer="direct")
3173 self.target_butler.transfer_from(
3174 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer
3175 )
3177 uri = self.target_butler.getURI(dataset.refs[0])
3178 if transfer == "auto":
3179 self.assertEqual(uri, temp)
3180 else:
3181 self.assertNotEqual(uri, temp)
3183 def test_shared_dimension_group(self):
3184 """Test internal logic that divides dataset types by dimension group
3185 when doing registry updates.
3186 """
3187 self.create_butlers()
3188 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
3189 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
3191 source_butler = self.source_butler
3192 target_butler = self.target_butler
3194 # Create a dataset type with the same dimensions as the 'bias' dataset
3195 # type from base.yaml
3196 dataset_type = DatasetType(
3197 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
3198 )
3199 source_butler.registry.registerDatasetType(dataset_type)
3200 # This has the same data ID as one of the bias datasets in
3201 # datasets.yaml.
3202 test_ref = source_butler.registry.insertDatasets(
3203 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
3204 )[0]
3206 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
3207 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
3208 refs = [test_ref, *biases, *flats]
3210 # Test setup will be even more convoluted if we want the datastore to
3211 # actually transfer files. For testing the dimension group behavior,
3212 # we really only care about the registry.
3213 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
3214 mock.return_value = (set(refs), set())
3215 target_butler.transfer_from(
3216 source_butler,
3217 refs,
3218 transfer=None,
3219 register_dataset_types=True,
3220 skip_missing=False,
3221 transfer_dimensions=True,
3222 )
3224 transferred_test_ref = target_butler.find_dataset(
3225 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3226 )
3227 self.assertEqual(transferred_test_ref.id, test_ref.id)
3229 transferred_bias = target_butler.find_dataset(
3230 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3231 )
3232 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
3234 transferred_flat = target_butler.find_dataset(
3235 "flat",
3236 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
3237 collections="imported_r",
3238 )
3239 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
3242class ChainedDatastoreTransfers(PosixDatastoreTransfers):
3243 """Test transfers using a chained datastore."""
3245 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
3248@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3249class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3250 """Test ``transfer_from`` involving Butler server."""
3252 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3254 def test_transfers_from_remote_to_direct(self) -> None:
3255 from lsst.daf.butler.remote_butler._remote_file_transfer_source import (
3256 mock_file_transfer_uris_for_unit_test,
3257 )
3259 self.target_butler = self.create_butler(None, "2")
3260 with create_test_server(TESTDIR) as server:
3261 self.source_butler = server.hybrid_butler
3263 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath:
3264 # The Butler server returns HTTP URIs with a domain name that
3265 # is not resolvable because there is no actual HTTP server
3266 # involved in these tests. Strip this first layer of
3267 # indirection, and return the target of the redirect instead.
3268 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers)
3269 return ResourcePath(str(response.next_request.url))
3271 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url):
3272 self.assertButlerTransfers()
3275class NullDatastoreTestCase(unittest.TestCase):
3276 """Test that we can fall back to a null datastore."""
3278 # Need a good config to create the repo.
3279 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3280 storageClassFactory: StorageClassFactory
3282 @classmethod
3283 def setUpClass(cls) -> None:
3284 cls.storageClassFactory = StorageClassFactory()
3285 cls.storageClassFactory.addFromConfig(cls.configFile)
3287 def setUp(self) -> None:
3288 """Create a new butler root for each test."""
3289 self.root = makeTestTempDir(TESTDIR)
3290 Butler.makeRepo(self.root, config=Config(self.configFile))
3292 def tearDown(self) -> None:
3293 removeTestTempDir(self.root)
3295 def test_fallback(self) -> None:
3296 # Read the butler config and mess with the datastore section.
3297 config_path = os.path.join(self.root, "butler.yaml")
3298 bad_config = Config(config_path)
3299 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore"
3300 bad_config.dumpToUri(config_path)
3302 with self.assertRaises(RuntimeError):
3303 Butler(self.root, without_datastore=False)
3305 with self.assertRaises(RuntimeError):
3306 Butler.from_config(self.root, without_datastore=False)
3308 butler = Butler.from_config(self.root, writeable=True, without_datastore=True)
3309 self.enterContext(butler)
3310 self.assertIsInstance(butler._datastore, NullDatastore)
3312 # Check that registry is working.
3313 butler.collections.register("MYRUN")
3314 collections = butler.collections.query("*")
3315 self.assertIn("MYRUN", set(collections))
3317 # Create a ref.
3318 dimensions = butler.dimensions.conform([])
3319 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
3320 datasetTypeName = "metric"
3321 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3322 butler.registry.registerDatasetType(datasetType)
3323 ref = DatasetRef(datasetType, {}, run="MYRUN")
3325 # Check that datastore will complain.
3326 with self.assertRaises(FileNotFoundError):
3327 butler.get(ref)
3328 with self.assertRaises(FileNotFoundError):
3329 butler.getURI(ref)
3332@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3333class ButlerServerTests(FileDatastoreButlerTests):
3334 """Test RemoteButler and Butler server."""
3336 configFile = None
3337 predictionSupported = False
3338 trustModeSupported = False
3340 postgres: TemporaryPostgresInstance | None
3342 def setUp(self):
3343 self.server_instance = self.enterContext(create_test_server(TESTDIR))
3345 def tearDown(self):
3346 pass
3348 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
3349 # S3 pre-signed URLs may end up with differing expiration times in the
3350 # query parameters, so ignore query parameters when comparing.
3351 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path
3353 def create_empty_butler(
3354 self,
3355 run: str | None = None,
3356 writeable: bool | None = None,
3357 metrics: ButlerMetrics | None = None,
3358 cleanup: bool = True,
3359 ) -> Butler:
3360 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics)
3362 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
3363 # Can't delete a file via S3 signed URLs, so we need to reach in
3364 # through DirectButler to delete the dataset.
3365 uri = self.server_instance.direct_butler.getURI(ref)
3366 uri.remove()
3368 def testConstructor(self):
3369 # RemoteButler constructor is tested in test_server.py and
3370 # test_remote_butler.py.
3371 pass
3373 def testDafButlerRepositories(self):
3374 # Loading of RemoteButler via repository index is tested in
3375 # test_server.py.
3376 pass
3378 def testGetDatasetTypes(self) -> None:
3379 # This is mostly a test of validateConfiguration, which is for
3380 # validating Datastore configuration and thus isn't relevant to
3381 # RemoteButler.
3382 pass
3384 def testMakeRepo(self) -> None:
3385 # Only applies to DirectButler.
3386 pass
3388 # Pickling not yet implemented for RemoteButler/HybridButler.
3389 @unittest.expectedFailure
3390 def testPickle(self) -> None:
3391 return super().testPickle()
3393 def testStringification(self) -> None:
3394 self.assertEqual(
3395 str(self.server_instance.remote_butler),
3396 "RemoteButler(https://test.example/api/butler/repo/testrepo/)",
3397 )
3399 def testTransaction(self) -> None:
3400 # Transactions will never be supported for RemoteButler.
3401 pass
3403 def testPutTemplates(self) -> None:
3404 # The Butler server instance is configured with different file naming
3405 # templates than this test is expecting.
3406 pass
3409@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3410class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase):
3411 """Tests for RemoteButler's registry shim, with a SQLite DB backing the
3412 server.
3413 """
3415 postgres = None
3418@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3419class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase):
3420 """Tests for RemoteButler's registry shim, with a Postgres DB backing the
3421 server.
3422 """
3424 @classmethod
3425 def setUpClass(cls):
3426 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
3427 super().setUpClass()
3430def setup_module(module: types.ModuleType) -> None:
3431 """Set up the module for pytest."""
3432 clean_environment()
3435def _get_test_data_path(filename: str) -> ResourcePath:
3436 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}")
3439if __name__ == "__main__":
3440 clean_environment()
3441 unittest.main()