Coverage for tests / test_butler.py: 13%
1765 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:12 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
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.datastore import NullDatastore
94from lsst.daf.butler.datastore.file_templates import FileTemplate, FileTemplateValidationError
95from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex
96from lsst.daf.butler.datastores.fileDatastore import FileDatastore
97from lsst.daf.butler.direct_butler import DirectButler
98from lsst.daf.butler.registry import (
99 CollectionError,
100 CollectionTypeError,
101 ConflictingDefinitionError,
102 DataIdValueError,
103 DatasetTypeExpressionError,
104 MissingCollectionError,
105 OrphanedRecordError,
106)
107from lsst.daf.butler.registry.sql_registry import SqlRegistry
108from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG
109from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter
110from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db
111from lsst.daf.butler.tests.server_available import butler_server_import_error, butler_server_is_available
112from lsst.daf.butler.tests.utils import (
113 MetricTestRepo,
114 TestCaseMixin,
115 makeTestTempDir,
116 removeTestTempDir,
117 safeTestTempDir,
118)
119from lsst.resources import ResourcePath
120from lsst.resources.http import HttpResourcePath
121from lsst.utils import doImportType
122from lsst.utils.introspection import get_full_type_name
124if butler_server_is_available:
125 from lsst.daf.butler.tests.server import create_test_server
128if TYPE_CHECKING:
129 import types
131 from lsst.daf.butler import DimensionGroup, Registry, StorageClass
133TESTDIR = os.path.abspath(os.path.dirname(__file__))
136def clean_environment() -> None:
137 """Remove external environment variables that affect the tests."""
138 for k in ("DAF_BUTLER_REPOSITORY_INDEX",):
139 os.environ.pop(k, None)
142def makeExampleMetrics() -> MetricsExample:
143 """Return example dataset suitable for tests."""
144 return MetricsExample(
145 {"AM1": 5.2, "AM2": 30.6},
146 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
147 [563, 234, 456.7, 752, 8, 9, 27],
148 )
151class TransactionTestError(Exception):
152 """Specific error for testing transactions, to prevent misdiagnosing
153 that might otherwise occur when a standard exception is used.
154 """
156 pass
159class ButlerConfigTests(unittest.TestCase):
160 """Simple tests for ButlerConfig that are not tested in any other test
161 cases.
162 """
164 def testSearchPath(self) -> None:
165 configFile = os.path.join(TESTDIR, "config", "basic", "butler.yaml")
166 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
167 config1 = ButlerConfig(configFile)
168 self.assertNotIn("testConfigs", "\n".join(cm.output))
170 overrideDirectory = os.path.join(TESTDIR, "config", "testConfigs")
171 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm:
172 config2 = ButlerConfig(configFile, searchPaths=[overrideDirectory])
173 self.assertIn("testConfigs", "\n".join(cm.output))
175 key = ("datastore", "records", "table")
176 self.assertNotEqual(config1[key], config2[key])
177 self.assertEqual(config2[key], "override_record")
180class ButlerPutGetTests(TestCaseMixin):
181 """Helper method for running a suite of put/get tests from different
182 butler configurations.
183 """
185 root: str
186 default_run = "ingésτ😺"
187 storageClassFactory: StorageClassFactory
188 configFile: str | None
189 tmpConfigFile: str
191 @staticmethod
192 def addDatasetType(
193 datasetTypeName: str, dimensions: DimensionGroup, storageClass: StorageClass | str, registry: Registry
194 ) -> DatasetType:
195 """Create a DatasetType and register it"""
196 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
197 registry.registerDatasetType(datasetType)
198 return datasetType
200 @classmethod
201 def setUpClass(cls) -> None:
202 cls.storageClassFactory = StorageClassFactory()
203 if cls.configFile is not None:
204 cls.storageClassFactory.addFromConfig(cls.configFile)
206 def assertGetComponents(
207 self,
208 butler: Butler,
209 datasetRef: DatasetRef,
210 components: tuple[str, ...],
211 reference: Any,
212 collections: Any = None,
213 ) -> None:
214 datasetType = datasetRef.datasetType
215 dataId = datasetRef.dataId
216 deferred = butler.getDeferred(datasetRef)
218 for component in components:
219 compTypeName = datasetType.componentTypeName(component)
220 result = butler.get(compTypeName, dataId, collections=collections)
221 self.assertEqual(result, getattr(reference, component))
222 result_deferred = deferred.get(component=component)
223 self.assertEqual(result_deferred, result)
225 def tearDown(self) -> None:
226 if self.root is not None:
227 removeTestTempDir(self.root)
229 def create_empty_butler(
230 self,
231 run: str | None = None,
232 writeable: bool | None = None,
233 metrics: ButlerMetrics | None = None,
234 cleanup: bool = True,
235 ):
236 """Create a Butler for the test repository, without inserting test
237 data.
238 """
239 butler = Butler.from_config(self.tmpConfigFile, run=run, writeable=writeable, metrics=metrics)
240 if cleanup:
241 self.enterContext(butler)
242 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
243 return butler
245 def create_butler(
246 self,
247 run: str,
248 storageClass: StorageClass | str,
249 datasetTypeName: str,
250 metrics: ButlerMetrics | None = None,
251 ) -> tuple[Butler, DatasetType]:
252 """Create a Butler for the test repository and insert some test data
253 into it.
254 """
255 butler = self.create_empty_butler(run=run, metrics=metrics)
257 collections = set(butler.collections.query("*"))
258 self.assertEqual(collections, {run})
259 # Create and register a DatasetType
260 dimensions = butler.dimensions.conform(["instrument", "visit"])
262 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
264 # Add needed Dimensions
265 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
266 butler.registry.insertDimensionData(
267 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
268 )
269 butler.registry.insertDimensionData(
270 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
271 )
272 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101})
273 visit_start = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
274 visit_end = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
275 butler.registry.insertDimensionData(
276 "visit",
277 {
278 "instrument": "DummyCamComp",
279 "id": 423,
280 "name": "fourtwentythree",
281 "physical_filter": "d-r",
282 "datetime_begin": visit_start,
283 "datetime_end": visit_end,
284 "day_obs": 20200101,
285 },
286 )
288 # Add more visits for some later tests
289 for visit_id in (424, 425):
290 butler.registry.insertDimensionData(
291 "visit",
292 {
293 "instrument": "DummyCamComp",
294 "id": visit_id,
295 "name": f"fourtwentyfour_{visit_id}",
296 "physical_filter": "d-r",
297 "day_obs": 20200101,
298 },
299 )
300 return butler, datasetType
302 def runPutGetTest(self, storageClass: StorageClass, datasetTypeName: str) -> Butler:
303 # New datasets will be added to run and tag, but we will only look in
304 # tag when looking up datasets.
305 run = self.default_run
306 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
307 assert butler.run is not None
309 # Create and store a dataset
310 metric = makeExampleMetrics()
311 dataId = butler.registry.expandDataId({"instrument": "DummyCamComp", "visit": 423})
313 # Dataset should not exist if we haven't added it
314 with self.assertRaises(DatasetNotFoundError):
315 butler.get(datasetTypeName, dataId)
317 # Put and remove the dataset once as a DatasetRef, once as a dataId,
318 # and once with a DatasetType
320 # Keep track of any collections we add and do not clean up
321 expected_collections = {run}
323 counter = 0
324 ref = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run="put_run_1")
325 args = tuple[DatasetRef] | tuple[str | DatasetType, DataCoordinate]
326 for args in ((ref,), (datasetTypeName, dataId), (datasetType, dataId)):
327 # Since we are using subTest we can get cascading failures
328 # here with the first attempt failing and the others failing
329 # immediately because the dataset already exists. Work around
330 # this by using a distinct run collection each time
331 counter += 1
332 this_run = f"put_run_{counter}"
333 butler.collections.register(this_run)
334 expected_collections.update({this_run})
336 with self.subTest(args=repr(args)):
337 kwargs: dict[str, Any] = {}
338 if not isinstance(args[0], DatasetRef): # type: ignore
339 kwargs["run"] = this_run
340 ref = butler.put(metric, *args, **kwargs)
341 self.assertIsInstance(ref, DatasetRef)
343 # Test get of a ref.
344 metricOut = butler.get(ref)
345 self.assertEqual(metric, metricOut)
346 # Test get
347 metricOut = butler.get(ref.datasetType.name, dataId, collections=this_run)
348 self.assertEqual(metric, metricOut)
349 # Test get with a datasetRef
350 metricOut = butler.get(ref)
351 self.assertEqual(metric, metricOut)
352 # Test getDeferred with dataId
353 metricOut = butler.getDeferred(ref.datasetType.name, dataId, collections=this_run).get()
354 self.assertEqual(metric, metricOut)
355 # Test getDeferred with a ref
356 metricOut = butler.getDeferred(ref).get()
357 self.assertEqual(metric, metricOut)
359 # Check we can get components
360 if storageClass.isComposite():
361 self.assertGetComponents(
362 butler, ref, ("summary", "data", "output"), metric, collections=this_run
363 )
365 primary_uri, secondary_uris = butler.getURIs(ref)
366 n_uris = len(secondary_uris)
367 if primary_uri:
368 n_uris += 1
370 # Can the artifacts themselves be retrieved?
371 if not butler._datastore.isEphemeral:
372 # Create a temporary directory to hold the retrieved
373 # artifacts.
374 with tempfile.TemporaryDirectory(
375 prefix="butler-artifacts-", ignore_cleanup_errors=True
376 ) as artifact_root:
377 root_uri = ResourcePath(artifact_root, forceDirectory=True)
379 for preserve_path in (True, False):
380 destination = root_uri.join(f"{preserve_path}_{counter}/")
381 log = logging.getLogger("lsst.x")
382 log.debug("Using destination %s for args %s", destination, args)
383 # Use copy so that we can test that overwrite
384 # protection works (using "auto" for File URIs
385 # would use hard links and subsequent transfer
386 # would work because it knows they are the same
387 # file).
388 transferred = butler.retrieveArtifacts(
389 [ref], destination, preserve_path=preserve_path, transfer="copy"
390 )
391 self.assertGreater(len(transferred), 0)
392 artifacts = list(ResourcePath.findFileResources([destination]))
393 # Filter out the index file.
394 artifacts = [a for a in artifacts if a.basename() != ZipIndex.index_name]
395 self.assertEqual(set(transferred), set(artifacts))
397 for artifact in transferred:
398 path_in_destination = artifact.relative_to(destination)
399 self.assertIsNotNone(path_in_destination)
400 assert path_in_destination is not None
402 # When path is not preserved there should not
403 # be any path separators.
404 num_seps = path_in_destination.count("/")
405 if preserve_path:
406 self.assertGreater(num_seps, 0)
407 else:
408 self.assertEqual(num_seps, 0)
410 self.assertEqual(
411 len(artifacts),
412 n_uris,
413 "Comparing expected artifacts vs actual:"
414 f" {artifacts} vs {primary_uri} and {secondary_uris}",
415 )
417 if preserve_path:
418 # No need to run these twice
419 with self.assertRaises(ValueError):
420 butler.retrieveArtifacts([ref], destination, transfer="move")
422 with self.assertRaisesRegex(
423 ValueError, "^Destination location must refer to a directory"
424 ):
425 butler.retrieveArtifacts(
426 [ref], ResourcePath("/some/file.txt", forceDirectory=False)
427 )
429 with self.assertRaises(FileExistsError):
430 butler.retrieveArtifacts([ref], destination)
432 transferred_again = butler.retrieveArtifacts(
433 [ref], destination, preserve_path=preserve_path, overwrite=True
434 )
435 self.assertEqual(set(transferred_again), set(transferred))
437 # Now remove the dataset completely.
438 butler.pruneDatasets([ref], purge=True, unstore=True)
439 # Lookup with original args should still fail.
440 kwargs = {"collections": this_run}
441 if isinstance(args[0], DatasetRef):
442 kwargs = {} # Prevent warning from being issued.
443 self.assertFalse(butler.exists(*args, **kwargs))
444 # get() should still fail.
445 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
446 butler.get(ref)
447 # Registry shouldn't be able to find it by dataset_id anymore.
448 self.assertIsNone(butler.get_dataset(ref.id))
450 # Do explicit registry removal since we know they are
451 # empty
452 butler.collections.x_remove(this_run)
453 expected_collections.remove(this_run)
455 # Create DatasetRef for put using default run.
456 refIn = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run=butler.run)
458 # Check that getDeferred fails with standalone ref.
459 with self.assertRaises(LookupError):
460 butler.getDeferred(refIn)
462 # Put the dataset again, since the last thing we did was remove it
463 # and we want to use the default collection.
464 ref = butler.put(metric, refIn)
466 # Get with parameters
467 stop = 4
468 sliced = butler.get(ref, parameters={"slice": slice(stop)})
469 self.assertNotEqual(metric, sliced)
470 self.assertEqual(metric.summary, sliced.summary)
471 self.assertEqual(metric.output, sliced.output)
472 assert metric.data is not None # for mypy
473 self.assertEqual(metric.data[:stop], sliced.data)
474 # getDeferred with parameters
475 sliced = butler.getDeferred(ref, parameters={"slice": slice(stop)}).get()
476 self.assertNotEqual(metric, sliced)
477 self.assertEqual(metric.summary, sliced.summary)
478 self.assertEqual(metric.output, sliced.output)
479 self.assertEqual(metric.data[:stop], sliced.data)
480 # getDeferred with deferred parameters
481 sliced = butler.getDeferred(ref).get(parameters={"slice": slice(stop)})
482 self.assertNotEqual(metric, sliced)
483 self.assertEqual(metric.summary, sliced.summary)
484 self.assertEqual(metric.output, sliced.output)
485 self.assertEqual(metric.data[:stop], sliced.data)
487 if storageClass.isComposite():
488 # Check that components can be retrieved
489 metricOut = butler.get(ref.datasetType.name, dataId)
490 compNameS = ref.datasetType.componentTypeName("summary")
491 compNameD = ref.datasetType.componentTypeName("data")
492 summary = butler.get(compNameS, dataId)
493 self.assertEqual(summary, metric.summary)
494 data = butler.get(compNameD, dataId)
495 self.assertEqual(data, metric.data)
497 if "counter" in storageClass.derivedComponents:
498 count = butler.get(ref.datasetType.componentTypeName("counter"), dataId)
499 self.assertEqual(count, len(data))
501 count = butler.get(
502 ref.datasetType.componentTypeName("counter"), dataId, parameters={"slice": slice(stop)}
503 )
504 self.assertEqual(count, stop)
506 compRef = butler.find_dataset(compNameS, dataId, collections=butler.collections.defaults)
507 assert compRef is not None
508 summary = butler.get(compRef)
509 self.assertEqual(summary, metric.summary)
511 # Create a Dataset type that has the same name but is inconsistent.
512 inconsistentDatasetType = DatasetType(
513 datasetTypeName, datasetType.dimensions, self.storageClassFactory.getStorageClass("Config")
514 )
516 # Getting with a dataset type that does not match registry fails
517 with self.assertRaisesRegex(
518 ValueError,
519 "(Supplied dataset type .* inconsistent with registry)"
520 "|(The new storage class .* is not compatible with the existing storage class)",
521 ):
522 butler.get(inconsistentDatasetType, dataId)
524 # Combining a DatasetRef with a dataId should fail
525 with self.assertRaisesRegex(ValueError, "DatasetRef given, cannot use dataId as well"):
526 butler.get(ref, dataId)
527 # Getting with an explicit ref should fail if the id doesn't match.
528 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)):
529 butler.get(DatasetRef(ref.datasetType, ref.dataId, id=uuid.UUID(int=101), run=butler.run))
531 # Getting a dataset with unknown parameters should fail
532 with self.assertRaisesRegex(KeyError, "Parameter 'unsupported' not understood"):
533 butler.get(ref, parameters={"unsupported": True})
535 # Check we have a collection
536 collections = set(butler.collections.query("*"))
537 self.assertEqual(collections, expected_collections)
539 # Clean up to check that we can remove something that may have
540 # already had a component removed
541 butler.pruneDatasets([ref], unstore=True, purge=True)
543 # Add the same ref again, so we can check that duplicate put fails.
544 ref = butler.put(metric, datasetType, dataId)
546 # Repeat put will fail.
547 with self.assertRaisesRegex(
548 ConflictingDefinitionError, "A database constraint failure was triggered"
549 ):
550 butler.put(metric, datasetType, dataId)
552 # Remove the datastore entry.
553 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
555 # Put will still fail
556 with self.assertRaisesRegex(
557 ConflictingDefinitionError, "A database constraint failure was triggered"
558 ):
559 butler.put(metric, datasetType, dataId)
561 # Repeat the same sequence with resolved ref.
562 butler.pruneDatasets([ref], unstore=True, purge=True)
563 ref = butler.put(metric, refIn)
565 # Repeat put will fail.
566 with self.assertRaisesRegex(ConflictingDefinitionError, "Datastore already contains dataset"):
567 butler.put(metric, refIn)
569 # Remove the datastore entry.
570 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False)
572 # In case of resolved ref this write will succeed.
573 ref = butler.put(metric, refIn)
575 # Leave the dataset in place since some downstream tests require
576 # something to be present
578 return butler
580 def testDeferredCollectionPassing(self) -> None:
581 # Construct a butler with no run or collection, but make it writeable.
582 butler = self.create_empty_butler(writeable=True)
583 # Create and register a DatasetType
584 dimensions = butler.dimensions.conform(["instrument", "visit"])
585 datasetType = self.addDatasetType(
586 "example", dimensions, self.storageClassFactory.getStorageClass("StructuredData"), butler.registry
587 )
588 # Add needed Dimensions
589 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
590 butler.registry.insertDimensionData(
591 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
592 )
593 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
594 butler.registry.insertDimensionData(
595 "visit",
596 {
597 "instrument": "DummyCamComp",
598 "id": 423,
599 "name": "fourtwentythree",
600 "physical_filter": "d-r",
601 "day_obs": 20250101,
602 },
603 )
604 dataId = {"instrument": "DummyCamComp", "visit": 423}
605 # Create dataset.
606 metric = makeExampleMetrics()
607 # Register a new run and put dataset.
608 run = "deferred"
609 self.assertTrue(butler.collections.register(run))
610 # Second time it will be allowed but indicate no-op
611 self.assertFalse(butler.collections.register(run))
612 ref = butler.put(metric, datasetType, dataId, run=run)
613 # Putting with no run should fail with TypeError.
614 with self.assertRaises(CollectionError):
615 butler.put(metric, datasetType, dataId)
616 # Dataset should exist.
617 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
618 # We should be able to get the dataset back, but with and without
619 # a deferred dataset handle.
620 self.assertEqual(metric, butler.get(datasetType, dataId, collections=[run]))
621 self.assertEqual(metric, butler.getDeferred(datasetType, dataId, collections=[run]).get())
622 # Trying to find the dataset without any collection is an error.
623 with self.assertRaises(NoDefaultCollectionError):
624 butler.exists(datasetType, dataId)
625 with self.assertRaises(CollectionError):
626 butler.get(datasetType, dataId)
627 # Associate the dataset with a different collection.
628 butler.collections.register("tagged", type=CollectionType.TAGGED)
629 butler.registry.associate("tagged", [ref])
630 # Deleting the dataset from the new collection should make it findable
631 # in the original collection.
632 butler.pruneDatasets([ref], tags=["tagged"])
633 self.assertTrue(butler.exists(datasetType, dataId, collections=[run]))
636class ButlerTests(ButlerPutGetTests):
637 """Tests for Butler."""
639 useTempRoot = True
640 validationCanFail: bool
641 fullConfigKey: str | None
642 registryStr: str | None
643 datastoreName: list[str] | None
644 datastoreStr: list[str]
645 predictionSupported = True
646 """Does getURIs support 'prediction mode'?"""
648 def setUp(self) -> None:
649 """Create a new butler root for each test."""
650 self.root = makeTestTempDir(TESTDIR)
651 Butler.makeRepo(self.root, config=Config(self.configFile))
652 self.tmpConfigFile = os.path.join(self.root, "butler.yaml")
654 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
655 """Return True if two URIs refer to the same resource.
657 Subclasses may override to handle unique requirements.
658 """
659 return uri1 == uri2
661 def testConstructor(self) -> None:
662 """Independent test of constructor."""
663 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
664 self.enterContext(butler)
665 self.assertIsInstance(butler, Butler)
667 # Check that butler.yaml is added automatically.
668 if self.tmpConfigFile.endswith(end := "/butler.yaml"):
669 config_dir = self.tmpConfigFile[: -len(end)]
670 butler = Butler.from_config(config_dir, run=self.default_run)
671 self.enterContext(butler)
672 self.assertIsInstance(butler, Butler)
674 # Even with a ResourcePath.
675 butler = Butler.from_config(ResourcePath(config_dir, forceDirectory=True), run=self.default_run)
676 self.enterContext(butler)
677 self.assertIsInstance(butler, Butler)
679 collections = set(butler.collections.query("*"))
680 self.assertEqual(collections, {self.default_run})
682 # Check that some special characters can be included in run name.
683 special_run = "u@b.c-A"
684 butler_special = Butler.from_config(butler=butler, run=special_run)
685 self.enterContext(butler_special)
686 collections = set(butler_special.registry.queryCollections("*@*"))
687 self.assertEqual(collections, {special_run})
689 butler2 = Butler.from_config(butler=butler, collections=["other"])
690 self.enterContext(butler2)
691 self.assertEqual(butler2.collections.defaults, ("other",))
692 self.assertIsNone(butler2.run)
693 self.assertEqual(type(butler._datastore), type(butler2._datastore))
694 self.assertEqual(butler._datastore.config, butler2._datastore.config)
696 # Test that we can use an environment variable to find this
697 # repository.
698 butler_index = Config()
699 butler_index["label"] = self.tmpConfigFile
700 for suffix in (".yaml", ".json"):
701 # Ensure that the content differs so that we know that
702 # we aren't reusing the cache.
703 bad_label = f"file://bucket/not_real{suffix}"
704 butler_index["bad_label"] = bad_label
705 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
706 butler_index.dumpToUri(temp_file)
707 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
708 self.assertEqual(Butler.get_known_repos(), {"label", "bad_label"})
709 uri = Butler.get_repo_uri("bad_label")
710 self.assertEqual(uri, ResourcePath(bad_label))
711 uri = Butler.get_repo_uri("label")
712 butler = Butler.from_config(uri, writeable=False)
713 self.assertIsInstance(butler, Butler)
714 butler.close()
715 butler = Butler.from_config("label", writeable=False)
716 self.assertIsInstance(butler, Butler)
717 butler.close()
718 with self.assertRaisesRegex(FileNotFoundError, "aliases:.*bad_label"):
719 Butler.from_config("not_there", writeable=False)
720 with self.assertRaisesRegex(FileNotFoundError, "resolved from alias 'bad_label'"):
721 Butler.from_config("bad_label")
722 with self.assertRaises(FileNotFoundError):
723 # Should ignore aliases.
724 Butler.from_config(ResourcePath("label", forceAbsolute=False))
725 with self.assertRaises(KeyError) as cm:
726 Butler.get_repo_uri("missing")
727 self.assertEqual(
728 Butler.get_repo_uri("missing", True), ResourcePath("missing", forceAbsolute=False)
729 )
730 self.assertIn("not known to", str(cm.exception))
731 # Should report no failure.
732 self.assertEqual(ButlerRepoIndex.get_failure_reason(), "")
733 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
734 # Now with empty configuration.
735 butler_index = Config()
736 butler_index.dumpToUri(temp_file)
737 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
738 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases)"):
739 Butler.from_config("label")
740 with ResourcePath.temporary_uri(suffix=suffix) as temp_file:
741 # Now with bad contents.
742 with open(temp_file.ospath, "w") as fh:
743 print("'", file=fh)
744 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}):
745 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases:.*could not be read)"):
746 Butler.from_config("label")
747 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": "file://not_found/x.yaml"}):
748 with self.assertRaises(FileNotFoundError):
749 Butler.get_repo_uri("label")
750 self.assertEqual(Butler.get_known_repos(), set())
752 with self.assertRaisesRegex(FileNotFoundError, "index file not found"):
753 Butler.from_config("label")
755 # Check that we can create Butler when the alias file is not found.
756 butler = Butler.from_config(self.tmpConfigFile, writeable=False)
757 self.enterContext(butler)
758 self.assertIsInstance(butler, Butler)
759 with self.assertRaises(RuntimeError) as cm:
760 # No environment variable set.
761 Butler.get_repo_uri("label")
762 self.assertEqual(Butler.get_repo_uri("label", True), ResourcePath("label", forceAbsolute=False))
763 self.assertIn("No repository index defined", str(cm.exception))
764 with self.assertRaisesRegex(FileNotFoundError, "no known aliases.*No repository index"):
765 # No aliases registered.
766 Butler.from_config("not_there")
767 self.assertEqual(Butler.get_known_repos(), set())
769 def testClose(self):
770 butler = self.create_empty_butler(cleanup=False)
771 is_direct_butler = isinstance(butler, DirectButler)
772 if is_direct_butler:
773 self.assertFalse(butler._closed)
775 with butler as butler_from_context_manager:
776 self.assertIs(butler, butler_from_context_manager)
777 if is_direct_butler:
778 self.assertTrue(butler._closed)
779 with self.assertRaisesRegex(RuntimeError, "has been closed"):
780 butler.get_dataset_type("raw")
782 # Close may be called multiple times.
783 butler.close()
784 if is_direct_butler:
785 self.assertTrue(butler._closed)
787 def testGarbageCollection(self):
788 """Test that Butler does not have any circular references that prevent
789 it from being garbage collected immediately when it goes out of scope.
790 """
791 butler = self.create_empty_butler(cleanup=False)
792 is_direct_butler = isinstance(butler, DirectButler)
793 butler_ref = weakref.ref(butler)
794 if is_direct_butler:
795 registry_ref = weakref.ref(butler._registry)
796 managers_ref = weakref.ref(butler._registry._managers)
797 datastore_ref = weakref.ref(butler._datastore)
798 db_ref = weakref.ref(butler._registry._db)
799 engine_ref = weakref.ref(butler._registry._db._engine)
801 with warnings.catch_warnings():
802 # Hide warnings from unclosed database handles.
803 warnings.simplefilter("ignore", ResourceWarning)
804 del butler
805 self.assertIsNone(butler_ref(), "Butler should have been garbage collected")
806 if is_direct_butler:
807 self.assertIsNone(registry_ref(), "SqlRegistry should have been garbage collected")
808 self.assertIsNone(managers_ref(), "Registry managers should have been garbage collected")
809 self.assertIsNone(datastore_ref(), "Datastore should have been garbage collected")
810 self.assertIsNone(db_ref(), "Database should have been garbage collected")
811 # SQLAlchemy has internal reference cycles, so the Engine instance
812 # is not cleaned up promptly even if we release our reference to
813 # it. Explicitly clean it up here to avoid file handles leaking.
814 if is_direct_butler:
815 engine = engine_ref()
816 if engine is not None:
817 engine.dispose()
819 def testDafButlerRepositories(self):
820 with unittest.mock.patch.dict(
821 os.environ,
822 {"DAF_BUTLER_REPOSITORIES": "label: 'https://someuri.com'\notherLabel: 'https://otheruri.com'\n"},
823 ):
824 self.assertEqual(str(Butler.get_repo_uri("label")), "https://someuri.com")
826 with unittest.mock.patch.dict(
827 os.environ,
828 {
829 "DAF_BUTLER_REPOSITORIES": "label: https://someuri.com",
830 "DAF_BUTLER_REPOSITORY_INDEX": "https://someuri.com",
831 },
832 ):
833 with self.assertRaisesRegex(RuntimeError, "Only one of the environment variables"):
834 Butler.get_repo_uri("label")
836 with unittest.mock.patch.dict(
837 os.environ,
838 {"DAF_BUTLER_REPOSITORIES": "invalid"},
839 ):
840 with self.assertRaisesRegex(ValueError, "Repository index not in expected format"):
841 Butler.get_repo_uri("label")
843 def testBasicPutGet(self) -> None:
844 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
845 self.runPutGetTest(storageClass, "test_metric")
847 def testCompositePutGetConcrete(self) -> None:
848 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadCompNoDisassembly")
849 butler = self.runPutGetTest(storageClass, "test_metric")
851 # Should *not* be disassembled
852 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
853 self.assertEqual(len(datasets), 1)
854 uri, components = butler.getURIs(datasets[0])
855 self.assertIsInstance(uri, ResourcePath)
856 self.assertFalse(components)
857 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
858 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
860 # Predicted dataset
861 if self.predictionSupported:
862 dataId = {"instrument": "DummyCamComp", "visit": 424}
863 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
864 self.assertFalse(components)
865 self.assertIsInstance(uri, ResourcePath)
866 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
867 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
869 def testCompositePutGetVirtual(self) -> None:
870 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp")
871 butler = self.runPutGetTest(storageClass, "test_metric_comp")
873 # Should be disassembled
874 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run))
875 self.assertEqual(len(datasets), 1)
876 uri, components = butler.getURIs(datasets[0])
878 if butler._datastore.isEphemeral:
879 # Never disassemble in-memory datastore
880 self.assertIsInstance(uri, ResourcePath)
881 self.assertFalse(components)
882 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}")
883 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}")
884 else:
885 self.assertIsNone(uri)
886 self.assertEqual(set(components), set(storageClass.components))
887 for compuri in components.values():
888 self.assertIsInstance(compuri, ResourcePath)
889 self.assertIn("423", str(compuri), f"Checking visit is in URI {compuri}")
890 self.assertEqual(compuri.fragment, "", f"Checking absence of fragment in {compuri}")
892 if self.predictionSupported:
893 # Predicted dataset
894 dataId = {"instrument": "DummyCamComp", "visit": 424}
895 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True)
897 if butler._datastore.isEphemeral:
898 # Never disassembled
899 self.assertIsInstance(uri, ResourcePath)
900 self.assertFalse(components)
901 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}")
902 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}")
903 else:
904 self.assertIsNone(uri)
905 self.assertEqual(set(components), set(storageClass.components))
906 for compuri in components.values():
907 self.assertIsInstance(compuri, ResourcePath)
908 self.assertIn("424", str(compuri), f"Checking visit is in URI {compuri}")
909 self.assertEqual(compuri.fragment, "predicted", f"Checking for fragment in {compuri}")
911 def testStorageClassOverrideGet(self) -> None:
912 """Test storage class conversion on get with override."""
913 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
914 datasetTypeName = "anything"
915 run = self.default_run
917 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName)
919 # Create and store a dataset.
920 metric = makeExampleMetrics()
921 dataId = {"instrument": "DummyCamComp", "visit": 423}
923 ref = butler.put(metric, datasetType, dataId)
925 # Return native type.
926 retrieved = butler.get(ref)
927 self.assertEqual(retrieved, metric)
929 # Specify an override.
930 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
931 model = butler.get(ref, storageClass=new_sc)
932 self.assertNotEqual(type(model), type(retrieved))
933 self.assertIs(type(model), new_sc.pytype)
934 self.assertEqual(retrieved, model)
936 # Defer but override later.
937 deferred = butler.getDeferred(ref)
938 model = deferred.get(storageClass=new_sc)
939 self.assertIs(type(model), new_sc.pytype)
940 self.assertEqual(retrieved, model)
942 # Defer but override up front.
943 deferred = butler.getDeferred(ref, storageClass=new_sc)
944 model = deferred.get()
945 self.assertIs(type(model), new_sc.pytype)
946 self.assertEqual(retrieved, model)
948 # Retrieve a component. Should be a tuple.
949 data = butler.get("anything.data", dataId, storageClass="StructuredDataDataTestTuple")
950 self.assertIs(type(data), tuple)
951 self.assertEqual(data, tuple(retrieved.data))
953 # Parameter on the write storage class should work regardless
954 # of read storage class.
955 data = butler.get(
956 "anything.data",
957 dataId,
958 storageClass="StructuredDataDataTestTuple",
959 parameters={"slice": slice(2, 4)},
960 )
961 self.assertEqual(len(data), 2)
963 # Try a parameter that is known to the read storage class but not
964 # the write storage class.
965 with self.assertRaises(KeyError):
966 butler.get(
967 "anything.data",
968 dataId,
969 storageClass="StructuredDataDataTestTuple",
970 parameters={"xslice": slice(2, 4)},
971 )
973 def testPytypePutCoercion(self) -> None:
974 """Test python type coercion on Butler.get and put."""
975 # Store some data with the normal example storage class.
976 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
977 datasetTypeName = "test_metric"
978 butler, _ = self.create_butler(self.default_run, storageClass, datasetTypeName)
980 dataId = {"instrument": "DummyCamComp", "visit": 423}
982 # Put a dict and this should coerce to a MetricsExample
983 test_dict = {"summary": {"a": 1}, "output": {"b": 2}}
984 metric_ref = butler.put(test_dict, datasetTypeName, dataId=dataId, visit=424)
985 test_metric = butler.get(metric_ref)
986 self.assertEqual(get_full_type_name(test_metric), "lsst.daf.butler.tests.MetricsExample")
987 self.assertEqual(test_metric.summary, test_dict["summary"])
988 self.assertEqual(test_metric.output, test_dict["output"])
990 # Check that the put still works if a DatasetType is given with
991 # a definition matching this python type.
992 registry_type = butler.get_dataset_type(datasetTypeName)
993 this_type = DatasetType(datasetTypeName, registry_type.dimensions, "StructuredDataDictJson")
994 metric2_ref = butler.put(test_dict, this_type, dataId=dataId, visit=425)
995 self.assertEqual(metric2_ref.datasetType, registry_type)
997 # The get will return the type expected by registry.
998 test_metric2 = butler.get(metric2_ref)
999 self.assertEqual(get_full_type_name(test_metric2), "lsst.daf.butler.tests.MetricsExample")
1001 # Make a new DatasetRef with the compatible but different DatasetType.
1002 # This should now return a dict.
1003 new_ref = DatasetRef(this_type, metric2_ref.dataId, id=metric2_ref.id, run=metric2_ref.run)
1004 test_dict2 = butler.get(new_ref)
1005 self.assertEqual(get_full_type_name(test_dict2), "dict")
1007 # Get it again with the wrong dataset type definition using get()
1008 # rather than get(). This should be consistent with get()
1009 # behavior and return the type of the DatasetType.
1010 test_dict3 = butler.get(this_type, dataId=dataId, visit=425)
1011 self.assertEqual(get_full_type_name(test_dict3), "dict")
1013 def test_ingest_zip(self) -> None:
1014 """Create butler, export data, delete data, import from Zip."""
1015 butler, dataset_type = self.create_butler(
1016 run=self.default_run, storageClass="StructuredData", datasetTypeName="metrics"
1017 )
1019 metric = makeExampleMetrics()
1020 refs = []
1021 for visit in (423, 424, 425):
1022 ref = butler.put(metric, dataset_type, instrument="DummyCamComp", visit=visit)
1023 refs.append(ref)
1025 # Retrieve a Zip file.
1026 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
1027 zip = butler.retrieve_artifacts_zip(refs, destination=tmpdir)
1029 # Ingest will fail.
1030 with self.assertRaises(ConflictingDefinitionError):
1031 butler.ingest_zip(zip)
1033 # Clear out the collection.
1034 butler.removeRuns([self.default_run])
1035 self.assertFalse(butler.exists(refs[0]))
1037 butler.ingest_zip(zip, transfer="copy")
1039 # Create an entirely new local file butler in this temp directory.
1040 new_butler_cfg = Butler.makeRepo(tmpdir)
1041 new_butler = Butler.from_config(new_butler_cfg, writeable=True)
1042 self.enterContext(new_butler)
1044 # This will fail since dimensions records are missing.
1045 with self.assertRaises(ConflictingDefinitionError):
1046 new_butler.ingest_zip(zip, transfer="copy")
1048 # Dry run should work.
1049 new_butler.ingest_zip(zip, transfer="copy", dry_run=True)
1051 new_butler.ingest_zip(zip, transfer="copy", transfer_dimensions=True)
1052 self.assertTrue(butler.exists(refs[0]))
1054 # Check that the refs can be read again.
1055 _ = [butler.get(ref) for ref in refs]
1057 uri = butler.getURI(refs[2])
1058 self.assertTrue(uri.exists())
1060 # Delete one dataset. The Zip file should still exist and allow
1061 # remaining refs to be read.
1062 butler.pruneDatasets([refs[0]], purge=True, unstore=True)
1063 self.assertTrue(uri.exists())
1065 metric2 = butler.get(refs[1])
1066 self.assertEqual(metric2, metric, msg=f"{metric2} != {metric}")
1068 butler.removeRuns([self.default_run])
1069 self.assertFalse(uri.exists())
1070 self.assertFalse(butler.exists(refs[-1]))
1072 with self.assertRaises(ValueError):
1073 butler.retrieve_artifacts_zip([], destination=".")
1075 def testIngest(self) -> None:
1076 butler = self.create_empty_butler(run=self.default_run)
1078 # Create and register a DatasetType
1079 dimensions = butler.dimensions.conform(["instrument", "visit", "detector"])
1081 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDictYaml")
1082 datasetTypeName = "metric"
1084 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1086 # Add needed Dimensions
1087 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1088 butler.registry.insertDimensionData(
1089 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1090 )
1091 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1092 for detector in (1, 2):
1093 butler.registry.insertDimensionData(
1094 "detector", {"instrument": "DummyCamComp", "id": detector, "full_name": f"detector{detector}"}
1095 )
1097 butler.registry.insertDimensionData(
1098 "visit",
1099 {
1100 "instrument": "DummyCamComp",
1101 "id": 423,
1102 "name": "fourtwentythree",
1103 "physical_filter": "d-r",
1104 "day_obs": 20250101,
1105 },
1106 {
1107 "instrument": "DummyCamComp",
1108 "id": 424,
1109 "name": "fourtwentyfour",
1110 "physical_filter": "d-r",
1111 "day_obs": 20250101,
1112 },
1113 )
1115 formatter = doImportType("lsst.daf.butler.formatters.yaml.YamlFormatter")
1116 dataRoot = os.path.join(TESTDIR, "data", "basic")
1117 datasets = []
1118 # Test one DatasetRef with a run that exists, and the other with a run
1119 # that doesn't exist, to verify that run collections are created when
1120 # required.
1121 runs = {1: self.default_run, 2: "a/new/run"}
1122 for detector in (1, 2):
1123 detector_name = f"detector_{detector}"
1124 metricFile = os.path.join(dataRoot, f"{detector_name}.yaml")
1125 dataId = butler.registry.expandDataId(
1126 {"instrument": "DummyCamComp", "visit": 423, "detector": detector}
1127 )
1128 # Create a DatasetRef for ingest
1129 refIn = DatasetRef(datasetType, dataId, run=runs[detector])
1131 datasets.append(FileDataset(path=metricFile, refs=[refIn], formatter=formatter))
1133 butler.ingest(*datasets, transfer="copy")
1135 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 423}
1136 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 423}
1138 metrics1 = butler.get(datasetTypeName, dataId1)
1139 metrics2 = butler.get(datasetTypeName, dataId2, collections="a/new/run")
1140 self.assertNotEqual(metrics1, metrics2)
1142 # Compare URIs
1143 uri1 = butler.getURI(datasetTypeName, dataId1)
1144 uri2 = butler.getURI(datasetTypeName, dataId2, collections="a/new/run")
1145 self.assertFalse(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1147 # Re-ingesting the same datasets raises an error with
1148 # skip_existing=False.
1149 with self.assertRaises(ConflictingDefinitionError):
1150 butler.ingest(*datasets, transfer="copy")
1151 # skip_existing=True makes it a no-op to re-ingest the same datasets.
1152 butler.ingest(*datasets, transfer="copy", skip_existing=True)
1154 # Now do a multi-dataset but single file ingest
1155 metricFile = os.path.join(dataRoot, "detectors.yaml")
1156 refs = []
1157 for detector in (1, 2):
1158 detector_name = f"detector_{detector}"
1159 dataId = butler.registry.expandDataId(
1160 {"instrument": "DummyCamComp", "visit": 424, "detector": detector}
1161 )
1162 # Create a DatasetRef for ingest
1163 refs.append(DatasetRef(datasetType, dataId, run=self.default_run))
1165 # Test "move" transfer to ensure that the files themselves
1166 # have disappeared following ingest.
1167 with ResourcePath.temporary_uri(suffix=".yaml") as tempFile:
1168 tempFile.transfer_from(ResourcePath(metricFile), transfer="copy")
1170 datasets = []
1171 datasets.append(FileDataset(path=tempFile, refs=refs, formatter=MultiDetectorFormatter))
1173 # For first ingest use copy.
1174 butler.ingest(*datasets, transfer="copy", record_validation_info=False)
1176 # Now try to ingest again in "execution butler" mode where
1177 # the registry entries exist but the datastore does not have
1178 # the files. We also need to strip the dimension records to ensure
1179 # that they will be re-added by the ingest.
1180 ref = datasets[0].refs[0]
1181 datasets[0].refs = [
1182 cast(
1183 DatasetRef,
1184 butler.find_dataset(ref.datasetType, data_id=ref.dataId, collections=ref.run),
1185 )
1186 for ref in datasets[0].refs
1187 ]
1188 all_refs = []
1189 for dataset in datasets:
1190 refs = []
1191 for ref in dataset.refs:
1192 # Create a dict from the dataId to drop the records.
1193 new_data_id = dict(ref.dataId.required)
1194 new_ref = butler.find_dataset(ref.datasetType, new_data_id, collections=ref.run)
1195 assert new_ref is not None
1196 self.assertFalse(new_ref.dataId.hasRecords())
1197 refs.append(new_ref)
1198 dataset.refs = refs
1199 all_refs.extend(dataset.refs)
1200 butler.pruneDatasets(all_refs, disassociate=False, unstore=True, purge=False)
1202 # Use move mode to test that the file is deleted. Also
1203 # disable recording of file size.
1204 butler.ingest(*datasets, transfer="move", record_validation_info=False)
1206 # Check that every ref now has records.
1207 for dataset in datasets:
1208 for ref in dataset.refs:
1209 self.assertTrue(ref.dataId.hasRecords())
1211 # Ensure that the file has disappeared.
1212 self.assertFalse(tempFile.exists())
1214 # Check that the datastore recorded no file size.
1215 # Not all datastores can support this.
1216 try:
1217 infos = butler._datastore.getStoredItemsInfo(datasets[0].refs[0]) # type: ignore[attr-defined]
1218 self.assertEqual(infos[0].file_size, -1)
1219 except AttributeError:
1220 pass
1222 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 424}
1223 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 424}
1225 multi1 = butler.get(datasetTypeName, dataId1)
1226 multi2 = butler.get(datasetTypeName, dataId2)
1228 self.assertEqual(multi1, metrics1)
1229 self.assertEqual(multi2, metrics2)
1231 # Compare URIs
1232 uri1 = butler.getURI(datasetTypeName, dataId1)
1233 uri2 = butler.getURI(datasetTypeName, dataId2)
1234 self.assertTrue(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}")
1236 # Test that removing one does not break the second
1237 # This line will issue a warning log message for a ChainedDatastore
1238 # that uses an InMemoryDatastore since in-memory can not ingest
1239 # files.
1240 butler.pruneDatasets([datasets[0].refs[0]], unstore=True, disassociate=False)
1241 self.assertFalse(butler.exists(datasetTypeName, dataId1))
1242 self.assertTrue(butler.exists(datasetTypeName, dataId2))
1243 multi2b = butler.get(datasetTypeName, dataId2)
1244 self.assertEqual(multi2, multi2b)
1246 # Ensure we can ingest 0 datasets
1247 datasets = []
1248 butler.ingest(*datasets)
1250 def testPickle(self) -> None:
1251 """Test pickle support."""
1252 butler = self.create_empty_butler(run=self.default_run)
1253 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration"
1254 butlerOut = pickle.loads(pickle.dumps(butler))
1255 self.enterContext(butlerOut)
1256 self.assertIsInstance(butlerOut, Butler)
1257 self.assertEqual(butlerOut._config, butler._config)
1258 self.assertEqual(list(butlerOut.collections.defaults), list(butler.collections.defaults))
1259 self.assertEqual(butlerOut.run, butler.run)
1261 def testGetDatasetTypes(self) -> None:
1262 butler = self.create_empty_butler(run=self.default_run)
1263 dimensions = butler.dimensions.conform(["instrument", "visit", "physical_filter"])
1264 dimensionEntries: list[tuple[str, list[Mapping[str, Any]]]] = [
1265 (
1266 "instrument",
1267 [
1268 {"instrument": "DummyCam"},
1269 {"instrument": "DummyHSC"},
1270 {"instrument": "DummyCamComp"},
1271 ],
1272 ),
1273 ("physical_filter", [{"instrument": "DummyCam", "name": "d-r", "band": "R"}]),
1274 ("day_obs", [{"instrument": "DummyCam", "id": 20250101}]),
1275 (
1276 "visit",
1277 [
1278 {
1279 "instrument": "DummyCam",
1280 "id": 42,
1281 "name": "fortytwo",
1282 "physical_filter": "d-r",
1283 "day_obs": 20250101,
1284 }
1285 ],
1286 ),
1287 ]
1288 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1289 # Add needed Dimensions
1290 for element, data in dimensionEntries:
1291 butler.registry.insertDimensionData(element, *data)
1293 # When a DatasetType is added to the registry entries are not created
1294 # for components but querying them can return the components.
1295 datasetTypeNames = {"metric", "metric2", "metric4", "metric33", "pvi", "paramtest"}
1296 components = set()
1297 for datasetTypeName in datasetTypeNames:
1298 # Create and register a DatasetType
1299 self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1301 for componentName in storageClass.components:
1302 components.add(DatasetType.nameWithComponent(datasetTypeName, componentName))
1304 fromRegistry: set[DatasetType] = set()
1305 for parent_dataset_type in butler.registry.queryDatasetTypes():
1306 fromRegistry.add(parent_dataset_type)
1307 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes())
1308 self.assertEqual({d.name for d in fromRegistry}, datasetTypeNames | components)
1310 # Query with wildcard.
1311 dataset_types = butler.registry.queryDatasetTypes("metric*")
1312 self.assertEqual(len(dataset_types), 4, f"Got: {dataset_types}")
1313 # but not regex.
1314 with self.assertRaises(DatasetTypeExpressionError):
1315 butler.registry.queryDatasetTypes(["pvi", re.compile("metric.*")])
1317 # Now that we have some dataset types registered, validate them
1318 butler.validateConfiguration(
1319 ignore=[
1320 "test_metric_comp",
1321 "metric3",
1322 "metric5",
1323 "calexp",
1324 "DummySC",
1325 "datasetType.component",
1326 "random_data",
1327 "random_data_2",
1328 ]
1329 )
1331 # Add a new datasetType that will fail template validation
1332 self.addDatasetType("test_metric_comp", dimensions, storageClass, butler.registry)
1333 if self.validationCanFail:
1334 with self.assertRaises(ValidationError):
1335 butler.validateConfiguration()
1337 # Rerun validation but with a subset of dataset type names
1338 butler.validateConfiguration(datasetTypeNames=["metric4"])
1340 # Rerun validation but ignore the bad datasetType
1341 butler.validateConfiguration(
1342 ignore=[
1343 "test_metric_comp",
1344 "metric3",
1345 "metric5",
1346 "calexp",
1347 "DummySC",
1348 "datasetType.component",
1349 "random_data",
1350 "random_data_2",
1351 ]
1352 )
1354 def testTransaction(self) -> None:
1355 butler = self.create_empty_butler(run=self.default_run)
1356 datasetTypeName = "test_metric"
1357 dimensions = butler.dimensions.conform(["instrument", "visit"])
1358 dimensionEntries: tuple[tuple[str, Mapping[str, Any]], ...] = (
1359 ("instrument", {"instrument": "DummyCam"}),
1360 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}),
1361 ("day_obs", {"instrument": "DummyCam", "id": 20250101}),
1362 (
1363 "visit",
1364 {
1365 "instrument": "DummyCam",
1366 "id": 42,
1367 "name": "fortytwo",
1368 "physical_filter": "d-r",
1369 "day_obs": 20250101,
1370 },
1371 ),
1372 )
1373 storageClass = self.storageClassFactory.getStorageClass("StructuredData")
1374 metric = makeExampleMetrics()
1375 dataId = {"instrument": "DummyCam", "visit": 42}
1376 # Create and register a DatasetType
1377 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry)
1378 with self.assertRaises(TransactionTestError):
1379 with butler.transaction():
1380 # Add needed Dimensions
1381 for args in dimensionEntries:
1382 butler.registry.insertDimensionData(*args)
1383 # Store a dataset
1384 ref = butler.put(metric, datasetTypeName, dataId)
1385 self.assertIsInstance(ref, DatasetRef)
1386 # Test get of a ref.
1387 metricOut = butler.get(ref)
1388 self.assertEqual(metric, metricOut)
1389 # Test get
1390 metricOut = butler.get(datasetTypeName, dataId)
1391 self.assertEqual(metric, metricOut)
1392 # Check we can get components
1393 self.assertGetComponents(butler, ref, ("summary", "data", "output"), metric)
1394 raise TransactionTestError("This should roll back the entire transaction")
1395 with self.assertRaises(DataIdValueError, msg=f"Check can't expand DataId {dataId}"):
1396 butler.registry.expandDataId(dataId)
1397 # Should raise LookupError for missing data ID value
1398 with self.assertRaises(LookupError, msg=f"Check can't get by {datasetTypeName} and {dataId}"):
1399 butler.get(datasetTypeName, dataId)
1400 # Also check explicitly if Dataset entry is missing
1401 self.assertIsNone(butler.find_dataset(datasetType, dataId, collections=butler.collections.defaults))
1402 # Direct retrieval should not find the file in the Datastore
1403 with self.assertRaises(FileNotFoundError, msg=f"Check {ref} can't be retrieved directly"):
1404 butler.get(ref)
1406 def testMakeRepo(self) -> None:
1407 """Test that we can write butler configuration to a new repository via
1408 the Butler.makeRepo interface and then instantiate a butler from the
1409 repo root.
1410 """
1411 # Do not run the test if we know this datastore configuration does
1412 # not support a file system root
1413 if self.fullConfigKey is None:
1414 return
1416 # create two separate directories
1417 root1 = tempfile.mkdtemp(dir=self.root)
1418 root2 = tempfile.mkdtemp(dir=self.root)
1420 butlerConfig = Butler.makeRepo(root1, config=Config(self.configFile))
1421 limited = Config(self.configFile)
1422 butler1 = Butler.from_config(butlerConfig)
1423 self.enterContext(butler1)
1424 assert isinstance(butler1, DirectButler), "Expect DirectButler in configuration"
1425 butlerConfig = Butler.makeRepo(root2, standalone=True, config=Config(self.configFile))
1426 full = Config(self.tmpConfigFile)
1427 butler2 = Butler.from_config(butlerConfig)
1428 self.enterContext(butler2)
1429 assert isinstance(butler2, DirectButler), "Expect DirectButler in configuration"
1430 # Butlers should have the same configuration regardless of whether
1431 # defaults were expanded.
1432 self.assertEqual(butler1._config, butler2._config)
1433 # Config files loaded directly should not be the same.
1434 self.assertNotEqual(limited, full)
1435 # Make sure "limited" doesn't have a few keys we know it should be
1436 # inheriting from defaults.
1437 self.assertIn(self.fullConfigKey, full)
1438 self.assertNotIn(self.fullConfigKey, limited)
1440 # Collections don't appear until something is put in them
1441 collections1 = set(butler1.registry.queryCollections())
1442 self.assertEqual(collections1, set())
1443 self.assertEqual(set(butler2.registry.queryCollections()), collections1)
1445 # Check that a config with no associated file name will not
1446 # work properly with relocatable Butler repo
1447 butlerConfig.configFile = None
1448 with self.assertRaises(ValueError):
1449 Butler.from_config(butlerConfig)
1451 with self.assertRaises(FileExistsError):
1452 Butler.makeRepo(self.root, standalone=True, config=Config(self.configFile), overwrite=False)
1454 def testStringification(self) -> None:
1455 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
1456 self.enterContext(butler)
1457 butlerStr = str(butler)
1459 if self.datastoreStr is not None:
1460 for testStr in self.datastoreStr:
1461 self.assertIn(testStr, butlerStr)
1462 if self.registryStr is not None:
1463 self.assertIn(self.registryStr, butlerStr)
1465 datastoreName = butler._datastore.name
1466 if self.datastoreName is not None:
1467 for testStr in self.datastoreName:
1468 self.assertIn(testStr, datastoreName)
1470 def testButlerRewriteDataId(self) -> None:
1471 """Test that dataIds can be rewritten based on dimension records."""
1472 butler = self.create_empty_butler(run=self.default_run)
1474 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
1475 datasetTypeName = "random_data"
1477 # Create dimension records.
1478 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1479 butler.registry.insertDimensionData(
1480 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1481 )
1482 butler.registry.insertDimensionData(
1483 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
1484 )
1486 dimensions = butler.dimensions.conform(["instrument", "exposure"])
1487 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
1488 butler.registry.registerDatasetType(datasetType)
1490 n_exposures = 5
1491 dayobs = 20210530
1493 # Create records for multiple day_obs but same seq_num to test that
1494 # we are constraining gets properly when day_obs/seq_num is used
1495 # for an exposure. Second day is year in future but is not used.
1496 for day_obs in (dayobs, dayobs + 1_00_00):
1497 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": day_obs})
1499 for i in range(n_exposures):
1500 group_name = f"group_{day_obs}_{i}"
1501 butler.registry.insertDimensionData(
1502 "group", {"instrument": "DummyCamComp", "name": group_name}
1503 )
1504 butler.registry.insertDimensionData(
1505 "exposure",
1506 {
1507 "instrument": "DummyCamComp",
1508 "id": day_obs + i,
1509 "obs_id": f"exp_{day_obs}_{i}",
1510 "seq_num": i,
1511 "day_obs": day_obs,
1512 "physical_filter": "d-r",
1513 "group": group_name,
1514 },
1515 )
1517 # Write some data.
1518 for i in range(n_exposures):
1519 metric = {"something": i, "other": "metric", "list": [2 * x for x in range(i)]}
1521 # Use the seq_num for the put to test rewriting.
1522 dataId = {"seq_num": i, "day_obs": dayobs, "instrument": "DummyCamComp", "physical_filter": "d-r"}
1523 ref = butler.put(metric, datasetTypeName, dataId=dataId)
1525 # Check that the exposure is correct in the dataId
1526 self.assertEqual(ref.dataId["exposure"], dayobs + i)
1528 # and check that we can get the dataset back with the same dataId
1529 new_metric = butler.get(datasetTypeName, dataId=dataId)
1530 self.assertEqual(new_metric, metric)
1532 # Check that we can find the datasets using the day_obs or the
1533 # exposure.day_obs.
1534 datasets_1 = list(
1535 butler.registry.queryDatasets(
1536 datasetType,
1537 collections=self.default_run,
1538 where="day_obs = :dayObs AND instrument = :instr",
1539 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1540 )
1541 )
1542 datasets_2 = list(
1543 butler.registry.queryDatasets(
1544 datasetType,
1545 collections=self.default_run,
1546 where="exposure.day_obs = :dayObs AND instrument = :instr",
1547 bind={"dayObs": dayobs, "instr": "DummyCamComp"},
1548 )
1549 )
1550 self.assertEqual(datasets_1, datasets_2)
1552 def testGetDatasetCollectionCaching(self):
1553 # Prior to DM-41117, there was a bug where get_dataset would throw
1554 # MissingCollectionError if you tried to fetch a dataset that was added
1555 # after the collection cache was last updated.
1556 reader_butler, datasetType = self.create_butler(self.default_run, "int", "datasettypename")
1557 writer_butler = self.create_empty_butler(writeable=True, run="new_run")
1558 dataId = {"instrument": "DummyCamComp", "visit": 423}
1559 put_ref = writer_butler.put(123, datasetType, dataId)
1560 get_ref = reader_butler.get_dataset(put_ref.id)
1561 self.assertEqual(get_ref.id, put_ref.id)
1562 # Also works when looking up via a hexadecimal string instead of a UUID
1563 # instance.
1564 hex_ref = reader_butler.get_dataset(put_ref.id.hex)
1565 self.assertEqual(hex_ref.id, put_ref.id)
1567 def testCollectionChainRedefine(self):
1568 butler = self._setup_to_test_collection_chain()
1570 butler.collections.redefine_chain("chain", "a")
1571 self._check_chain(butler, ["a"])
1573 # Duplicates are removed from the list of children
1574 butler.collections.redefine_chain("chain", ["c", "b", "c"])
1575 self._check_chain(butler, ["c", "b"])
1577 # Empty list clears the chain
1578 butler.collections.redefine_chain("chain", [])
1579 self._check_chain(butler, [])
1581 self._test_common_chain_functionality(butler, butler.collections.redefine_chain)
1583 def testCollectionChainPrepend(self):
1584 butler = self._setup_to_test_collection_chain()
1586 # Duplicates are removed from the list of children
1587 butler.collections.prepend_chain("chain", ["c", "b", "c"])
1588 self._check_chain(butler, ["c", "b"])
1590 # Prepend goes on the front of existing chain
1591 butler.collections.prepend_chain("chain", ["a"])
1592 self._check_chain(butler, ["a", "c", "b"])
1594 # Empty prepend does nothing
1595 butler.collections.prepend_chain("chain", [])
1596 self._check_chain(butler, ["a", "c", "b"])
1598 # Prepending children that already exist in the chain removes them from
1599 # their current position.
1600 butler.collections.prepend_chain("chain", ["d", "b", "c"])
1601 self._check_chain(butler, ["d", "b", "c", "a"])
1603 self._test_common_chain_functionality(butler, butler.collections.prepend_chain)
1605 def testCollectionChainExtend(self):
1606 butler = self._setup_to_test_collection_chain()
1608 # Duplicates are removed from the list of children
1609 butler.collections.extend_chain("chain", ["c", "b", "c"])
1610 self._check_chain(butler, ["c", "b"])
1612 # Extend goes on the end of existing chain
1613 butler.collections.extend_chain("chain", ["a"])
1614 self._check_chain(butler, ["c", "b", "a"])
1616 # Empty extend does nothing
1617 butler.collections.extend_chain("chain", [])
1618 self._check_chain(butler, ["c", "b", "a"])
1620 # Extending children that already exist in the chain removes them from
1621 # their current position.
1622 butler.collections.extend_chain("chain", ["d", "b", "c"])
1623 self._check_chain(butler, ["a", "d", "b", "c"])
1625 self._test_common_chain_functionality(butler, butler.collections.extend_chain)
1627 def testCollectionChainRemove(self) -> None:
1628 butler = self._setup_to_test_collection_chain()
1630 butler.collections.redefine_chain("chain", ["a", "b", "c", "d"])
1632 butler.collections.remove_from_chain("chain", "c")
1633 self._check_chain(butler, ["a", "b", "d"])
1635 # Duplicates are allowed in the list of children
1636 butler.collections.remove_from_chain("chain", ["b", "b", "a"])
1637 self._check_chain(butler, ["d"])
1639 # Empty remove does nothing
1640 butler.collections.remove_from_chain("chain", [])
1641 self._check_chain(butler, ["d"])
1643 # Removing children that aren't in the chain does nothing
1644 butler.collections.remove_from_chain("chain", ["a", "chain"])
1645 self._check_chain(butler, ["d"])
1647 self._test_common_chain_functionality(
1648 butler, butler.collections.remove_from_chain, skip_cycle_check=True
1649 )
1651 def _setup_to_test_collection_chain(self) -> Butler:
1652 butler = self.create_empty_butler(writeable=True)
1654 butler.collections.register("chain", CollectionType.CHAINED)
1656 runs = ["a", "b", "c", "d"]
1657 for run in runs:
1658 butler.collections.register(run)
1660 butler.collections.register("staticchain", CollectionType.CHAINED)
1661 butler.collections.redefine_chain("staticchain", ["a", "b"])
1663 return butler
1665 def _check_chain(self, butler: Butler, expected: list[str]) -> None:
1666 children = butler.collections.get_info("chain").children
1667 self.assertEqual(expected, list(children))
1669 def _test_common_chain_functionality(
1670 self, butler, func: Callable[[str, str | list[str]], Any], *, skip_cycle_check=False
1671 ) -> None:
1672 # Missing parent collection
1673 with self.assertRaises(MissingCollectionError):
1674 func("doesnotexist", [])
1675 # Missing child collection
1676 with self.assertRaises(MissingCollectionError):
1677 func("chain", ["doesnotexist"])
1678 # Forbid operations on non-chained collections
1679 with self.assertRaises(CollectionTypeError):
1680 func("d", ["a"])
1682 # Prevent collection cycles
1683 if not skip_cycle_check:
1684 butler.collections.register("chain2", CollectionType.CHAINED)
1685 func("chain2", "chain")
1686 with self.assertRaises(CollectionCycleError):
1687 func("chain", "chain2")
1689 # Make sure none of the earlier operations interfered with unrelated
1690 # chains.
1691 self.assertEqual(["a", "b"], list(butler.collections.get_info("staticchain").children))
1693 with butler._caching_context():
1694 with self.assertRaisesRegex(RuntimeError, "Chained collection modification not permitted"):
1695 func("chain", "a")
1698class FileDatastoreButlerTests(ButlerTests):
1699 """Common tests and specialization of ButlerTests for butlers backed
1700 by datastores that inherit from FileDatastore.
1701 """
1703 trustModeSupported = True
1705 def checkFileExists(self, root: str | ResourcePath, relpath: str | ResourcePath) -> bool:
1706 """Check if file exists at a given path (relative to root).
1708 Test testPutTemplates verifies actual physical existance of the files
1709 in the requested location.
1710 """
1711 uri = ResourcePath(root, forceDirectory=True)
1712 return uri.join(relpath).exists()
1714 def testPutTemplates(self) -> None:
1715 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1716 butler = self.create_empty_butler(run=self.default_run)
1718 # Add needed Dimensions
1719 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
1720 butler.registry.insertDimensionData(
1721 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
1722 )
1723 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101})
1724 butler.registry.insertDimensionData(
1725 "visit",
1726 {
1727 "instrument": "DummyCamComp",
1728 "id": 423,
1729 "name": "v423",
1730 "physical_filter": "d-r",
1731 "day_obs": 20250101,
1732 },
1733 )
1734 butler.registry.insertDimensionData(
1735 "visit",
1736 {
1737 "instrument": "DummyCamComp",
1738 "id": 425,
1739 "name": "v425",
1740 "physical_filter": "d-r",
1741 "day_obs": 20250101,
1742 },
1743 )
1745 # Create and store a dataset
1746 metric = makeExampleMetrics()
1748 # Create two almost-identical DatasetTypes (both will use default
1749 # template)
1750 dimensions = butler.dimensions.conform(["instrument", "visit"])
1751 butler.registry.registerDatasetType(DatasetType("metric1", dimensions, storageClass))
1752 butler.registry.registerDatasetType(DatasetType("metric2", dimensions, storageClass))
1753 butler.registry.registerDatasetType(DatasetType("metric3", dimensions, storageClass))
1755 dataId1 = {"instrument": "DummyCamComp", "visit": 423}
1756 dataId2 = {"instrument": "DummyCamComp", "visit": 423, "physical_filter": "d-r"}
1758 # Put with exactly the data ID keys needed
1759 ref = butler.put(metric, "metric1", dataId1)
1760 uri = butler.getURI(ref)
1761 self.assertTrue(uri.exists())
1762 self.assertTrue(
1763 uri.unquoted_path.endswith(f"{self.default_run}/metric1/??#?/d-r/DummyCamComp_423.pickle")
1764 )
1766 # Check the template based on dimensions
1767 if hasattr(butler._datastore, "templates"):
1768 butler._datastore.templates.validateTemplates([ref])
1770 # Put with extra data ID keys (physical_filter is an optional
1771 # dependency); should not change template (at least the way we're
1772 # defining them to behave now; the important thing is that they
1773 # must be consistent).
1774 ref = butler.put(metric, "metric2", dataId2)
1775 uri = butler.getURI(ref)
1776 self.assertTrue(uri.exists())
1777 self.assertTrue(
1778 uri.unquoted_path.endswith(f"{self.default_run}/metric2/d-r/DummyCamComp_v423.pickle")
1779 )
1781 # Check the template based on dimensions
1782 if hasattr(butler._datastore, "templates"):
1783 butler._datastore.templates.validateTemplates([ref])
1785 # Use a template that has a typo in dimension record metadata.
1786 # Easier to test with a butler that has a ref with records attached.
1787 template = FileTemplate("a/{visit.name}/{id}_{visit.namex:?}.fits")
1788 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1789 path = template.format(ref)
1790 self.assertEqual(path, f"a/v423/{ref.id}_fits")
1792 template = FileTemplate("a/{visit.name}/{id}_{visit.namex}.fits")
1793 with self.assertRaises(KeyError):
1794 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"):
1795 template.format(ref)
1797 # Now use a file template that will not result in unique filenames
1798 with self.assertRaises(FileTemplateValidationError):
1799 butler.put(metric, "metric3", dataId1)
1801 def testImportExport(self) -> None:
1802 # Run put/get tests just to create and populate a repo.
1803 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1804 self.runImportExportTest(storageClass)
1806 @unittest.expectedFailure
1807 def testImportExportVirtualComposite(self) -> None:
1808 # Run put/get tests just to create and populate a repo.
1809 storageClass = self.storageClassFactory.getStorageClass("StructuredComposite")
1810 self.runImportExportTest(storageClass)
1812 def runImportExportTest(self, storageClass: StorageClass) -> None:
1813 """Test exporting and importing.
1815 This test does an export to a temp directory and an import back
1816 into a new temp directory repo. It does not assume a posix datastore.
1817 """
1818 exportButler = self.runPutGetTest(storageClass, "test_metric")
1820 # Test that we must have a file extension.
1821 with self.assertRaises(ValueError):
1822 with exportButler.export(filename="dump", directory=".") as export:
1823 pass
1825 # Test that unknown format is not allowed.
1826 with self.assertRaises(ValueError):
1827 with exportButler.export(filename="dump.fits", directory=".") as export:
1828 pass
1830 # Test that the repo actually has at least one dataset.
1831 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
1832 self.assertGreater(len(datasets), 0)
1833 # Add a DimensionRecord that's unused by those datasets.
1834 skymapRecord = {"name": "example_skymap", "hash": (50).to_bytes(8, byteorder="little")}
1835 exportButler.registry.insertDimensionData("skymap", skymapRecord)
1836 # Export and then import datasets.
1837 with safeTestTempDir(TESTDIR) as exportDir:
1838 exportFile = os.path.join(exportDir, "exports.yaml")
1839 with exportButler.export(filename=exportFile, directory=exportDir, transfer="auto") as export:
1840 export.saveDatasets(datasets)
1841 # Export the same datasets again. This should quietly do
1842 # nothing because of internal deduplication, and it shouldn't
1843 # complain about being asked to export the "htm7" elements even
1844 # though there aren't any in these datasets or in the database.
1845 export.saveDatasets(datasets, elements=["htm7"])
1846 # Save one of the data IDs again; this should be harmless
1847 # because of internal deduplication.
1848 export.saveDataIds([datasets[0].dataId])
1849 # Save some dimension records directly.
1850 export.saveDimensionData("skymap", [skymapRecord])
1851 self.assertTrue(os.path.exists(exportFile))
1852 with safeTestTempDir(TESTDIR) as importDir:
1853 # We always want this to be a local posix butler
1854 Butler.makeRepo(importDir, config=Config(os.path.join(TESTDIR, "config/basic/butler.yaml")))
1855 # Calling script.butlerImport tests the implementation of the
1856 # butler command line interface "import" subcommand. Functions
1857 # in the script folder are generally considered protected and
1858 # should not be used as public api.
1859 with open(exportFile) as f:
1860 script.butlerImport(
1861 importDir,
1862 export_file=f,
1863 directory=exportDir,
1864 transfer="auto",
1865 skip_dimensions=None,
1866 )
1867 importButler = Butler.from_config(importDir, run=self.default_run)
1868 self.enterContext(importButler)
1869 for ref in datasets:
1870 with self.subTest(ref=repr(ref)):
1871 # Test for existence by passing in the DatasetType and
1872 # data ID separately, to avoid lookup by dataset_id.
1873 self.assertTrue(importButler.exists(ref.datasetType, ref.dataId))
1874 self.assertEqual(
1875 list(importButler.registry.queryDimensionRecords("skymap")),
1876 [importButler.dimensions["skymap"].RecordClass(**skymapRecord)],
1877 )
1879 def testRemoveRuns(self) -> None:
1880 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1881 butler = self.create_empty_butler(writeable=True)
1882 # Load registry data with dimensions to hang datasets off of.
1883 butler.import_(filename=ResourcePath("resource://lsst.daf.butler/tests/registry_data/base.yaml"))
1884 # Add some RUN-type collection.
1885 run1 = "run1"
1886 butler.collections.register(run1)
1887 run2 = "run2"
1888 butler.collections.register(run2)
1889 # put a dataset in each
1890 metric = makeExampleMetrics()
1891 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1892 datasetType = self.addDatasetType(
1893 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1894 )
1895 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1896 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1897 uri1 = butler.getURI(ref1)
1898 uri2 = butler.getURI(ref2)
1900 # Put one of the runs in a chain.
1901 butler.collections.register("Chain", CollectionType.CHAINED)
1902 butler.collections.extend_chain("Chain", run1)
1904 with self.assertRaises(OrphanedRecordError):
1905 butler.registry.removeDatasetType(datasetType.name)
1907 # Remove a non-run.
1908 with self.assertRaises(TypeError):
1909 butler.removeRuns(["Chain"])
1911 # Remove without unlinking from chain should fail.
1912 with self.assertRaises(IntegrityError):
1913 butler.removeRuns([run1])
1915 # Remove from both runs. No longer use unstore parameter since it
1916 # always purges.
1917 butler.removeRuns([run1, run2], unlink_from_chains=True)
1919 # Should be nothing in registry for either one, and datastore should
1920 # not think either exists.
1921 with self.assertRaises(MissingCollectionError):
1922 butler.collections.get_info(run1)
1923 with self.assertRaises(MissingCollectionError):
1924 butler.collections.get_info(run1)
1925 self.assertFalse(butler.stored(ref1))
1926 self.assertFalse(butler.stored(ref2))
1927 # We always unstore so both URIs should be gone.
1928 self.assertFalse(uri1.exists())
1929 self.assertFalse(uri2.exists())
1931 # Now that the collections have been pruned we can remove the
1932 # dataset type
1933 butler.registry.removeDatasetType(datasetType.name)
1935 with self.assertLogs("lsst.daf.butler.registry", "INFO") as cm:
1936 butler.registry.removeDatasetType(("test*", "test*"))
1937 self.assertIn("not defined", "\n".join(cm.output))
1939 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
1940 """Simulate an external actor removing a file outside of Butler's
1941 knowledge.
1943 Subclasses may override to handle more complicated datastore
1944 configurations.
1945 """
1946 uri = butler.getURI(ref)
1947 uri.remove()
1948 datastore = cast(FileDatastore, butler._datastore)
1949 datastore.cacheManager.remove_from_cache(ref)
1951 def testPruneDatasets(self) -> None:
1952 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
1953 butler = self.create_empty_butler(writeable=True)
1954 # Load registry data with dimensions to hang datasets off of.
1955 butler.import_(filename=_get_test_data_path("base.yaml"))
1956 # Add some RUN-type collections.
1957 run1 = "run1"
1958 butler.collections.register(run1)
1959 run2 = "run2"
1960 butler.collections.register(run2)
1961 # put some datasets. ref1 and ref2 have the same data ID, and are in
1962 # different runs. ref3 has a different data ID.
1963 metric = makeExampleMetrics()
1964 dimensions = butler.dimensions.conform(["instrument", "physical_filter"])
1965 datasetType = self.addDatasetType(
1966 "prune_collections_test_dataset", dimensions, storageClass, butler.registry
1967 )
1968 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1)
1969 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2)
1970 ref3 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, run=run1)
1972 many_stored = butler.stored_many([ref1, ref2, ref3])
1973 for ref, stored in many_stored.items():
1974 self.assertTrue(stored, f"Ref {ref} should be stored")
1976 many_exists = butler._exists_many([ref1, ref2, ref3])
1977 for ref, exists in many_exists.items():
1978 self.assertTrue(exists, f"Checking ref {ref} exists.")
1979 self.assertEqual(exists, DatasetExistence.VERIFIED, f"Ref {ref} should be stored")
1981 # Simple prune.
1982 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
1983 self.assertFalse(butler.exists(ref1.datasetType, ref1.dataId, collections=run1))
1985 many_stored = butler.stored_many([ref1, ref2, ref3])
1986 for ref, stored in many_stored.items():
1987 self.assertFalse(stored, f"Ref {ref} should not be stored")
1989 many_exists = butler._exists_many([ref1, ref2, ref3])
1990 for ref, exists in many_exists.items():
1991 self.assertEqual(exists, DatasetExistence.UNRECOGNIZED, f"Ref {ref} should not be stored")
1993 # Put data back.
1994 ref1_new = butler.put(metric, ref1)
1995 self.assertEqual(ref1_new, ref1) # Reuses original ID.
1996 ref2 = butler.put(metric, ref2)
1998 many_stored = butler.stored_many([ref1, ref2, ref3])
1999 self.assertTrue(many_stored[ref1])
2000 self.assertTrue(many_stored[ref2])
2001 self.assertFalse(many_stored[ref3])
2003 ref3 = butler.put(metric, ref3)
2005 many_exists = butler._exists_many([ref1, ref2, ref3])
2006 for ref, exists in many_exists.items():
2007 self.assertTrue(exists, f"Ref {ref} should not be stored")
2009 # Clear out the datasets from registry and start again.
2010 refs = [ref1, ref2, ref3]
2011 butler.pruneDatasets(refs, purge=True, unstore=True)
2012 for ref in refs:
2013 butler.put(metric, ref)
2015 # Confirm we can retrieve deferred.
2016 dref1 = butler.getDeferred(ref1) # known and exists
2017 metric1 = dref1.get()
2018 self.assertEqual(metric1, metric)
2020 # Test different forms of file availability.
2021 # Need to be in a state where:
2022 # - one ref just has registry record.
2023 # - one ref has a missing file but a datastore record.
2024 # - one ref has a missing datastore record but file is there.
2025 # - one ref does not exist anywhere.
2026 # Do not need to test a ref that has everything since that is tested
2027 # above.
2028 ref0 = DatasetRef(
2029 datasetType,
2030 DataCoordinate.standardize(
2031 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, universe=butler.dimensions
2032 ),
2033 run=run1,
2034 )
2036 # Delete from datastore and retain in Registry.
2037 butler.pruneDatasets([ref1], purge=False, unstore=True, disassociate=False)
2039 # File has been removed.
2040 self.remove_dataset_out_of_band(butler, ref2)
2042 # Datastore has lost track.
2043 butler._datastore.forget([ref3])
2045 # First test with a standard butler.
2046 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
2047 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2048 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2049 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2050 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED)
2052 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=False)
2053 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2054 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
2055 self.assertEqual(exists_many[ref2], DatasetExistence.KNOWN)
2056 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ASSUMED)
2057 self.assertTrue(exists_many[ref2])
2059 # Check that per-ref query gives the same answer as many query.
2060 for ref, exists in exists_many.items():
2061 self.assertEqual(butler.exists(ref, full_check=False), exists)
2063 # Get deferred checks for existence before it allows it to be
2064 # retrieved.
2065 with self.assertRaises(LookupError):
2066 butler.getDeferred(ref3) # not known, file exists
2067 dref2 = butler.getDeferred(ref2) # known but file missing
2068 with self.assertRaises(FileNotFoundError):
2069 dref2.get()
2071 # Test again with a trusting butler.
2072 if self.trustModeSupported:
2073 butler._datastore.trustGetRequest = True
2074 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True)
2075 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED)
2076 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED)
2077 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE)
2078 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ARTIFACT)
2080 # When trusting we can get a deferred dataset handle that is not
2081 # known but does exist.
2082 dref3 = butler.getDeferred(ref3)
2083 metric3 = dref3.get()
2084 self.assertEqual(metric3, metric)
2086 # Check that per-ref query gives the same answer as many query.
2087 for ref, exists in exists_many.items():
2088 self.assertEqual(butler.exists(ref, full_check=True), exists)
2090 # Create a ref that surprisingly has the UUID of an existing ref
2091 # but is not the same.
2092 ref_bad = DatasetRef(datasetType, dataId=ref3.dataId, run=ref3.run, id=ref2.id)
2093 with self.assertRaises(ValueError):
2094 butler.exists(ref_bad)
2096 # Create a ref that has a compatible storage class.
2097 ref_compat = ref2.overrideStorageClass("StructuredDataDict")
2098 exists = butler.exists(ref_compat)
2099 self.assertEqual(exists, exists_many[ref2])
2101 # Remove everything and start from scratch.
2102 butler._datastore.trustGetRequest = False
2103 butler.pruneDatasets(refs, purge=True, unstore=True)
2104 for ref in refs:
2105 butler.put(metric, ref)
2107 # These tests mess directly with the trash table and can leave the
2108 # datastore in an odd state. Do them at the end.
2109 # Check that in normal mode, deleting the record will lead to
2110 # trash not touching the file.
2111 uri1 = butler.getURI(ref1)
2112 butler._datastore.bridge.moveToTrash(
2113 [ref1], transaction=None
2114 ) # Update the dataset_location table
2115 butler._datastore.forget([ref1])
2116 butler._datastore.trash(ref1)
2117 butler._datastore.emptyTrash()
2118 self.assertTrue(uri1.exists())
2119 uri1.remove() # Clean it up.
2121 # Simulate execution butler setup by deleting the datastore
2122 # record but keeping the file around and trusting.
2123 butler._datastore.trustGetRequest = True
2124 uris = butler.get_many_uris([ref2, ref3])
2125 uri2 = uris[ref2].primaryURI
2126 uri3 = uris[ref3].primaryURI
2127 self.assertTrue(uri2.exists())
2128 self.assertTrue(uri3.exists())
2130 # Remove the datastore record.
2131 butler._datastore.bridge.moveToTrash(
2132 [ref2], transaction=None
2133 ) # Update the dataset_location table
2134 butler._datastore.forget([ref2])
2135 self.assertTrue(uri2.exists())
2136 butler._datastore.trash([ref2, ref3])
2137 # Immediate removal for ref2 file
2138 self.assertFalse(uri2.exists())
2139 # But ref3 has to wait for the empty.
2140 self.assertTrue(uri3.exists())
2141 butler._datastore.emptyTrash()
2142 self.assertFalse(uri3.exists())
2144 # Clear out the datasets from registry.
2145 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
2147 def test_butler_metrics(self):
2148 """Test that metrics are collected."""
2149 run = "test_run"
2150 metrics = ButlerMetrics()
2151 butler, datasetType = self.create_butler(
2152 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics
2153 )
2154 data = MetricsExampleModel(
2155 summary={"AM1": 5.2, "AM2": 30.6},
2156 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2157 data=[563, 234, 456.7, 752, 8, 9, 27],
2158 )
2160 data_ref = butler.put(data, datasetType, visit=424, instrument="DummyCamComp")
2161 butler.get(data_ref)
2162 butler.get(data_ref)
2163 self.assertEqual(metrics.n_get, 2)
2164 self.assertGreater(metrics.time_in_get, 0.0)
2165 self.assertEqual(metrics.n_put, 1)
2166 self.assertGreater(metrics.time_in_put, 0.0)
2168 deferred = butler.getDeferred(data_ref)
2169 deferred.get()
2170 self.assertEqual(metrics.n_get, 3)
2172 with butler.record_metrics() as new:
2173 butler.put(data, datasetType, visit=425, instrument="DummyCamComp")
2174 butler.get(data_ref)
2175 self.assertEqual(new.n_get, 1)
2176 self.assertEqual(new.n_put, 1)
2179class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2180 """PosixDatastore specialization of a butler"""
2182 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2183 fullConfigKey: str | None = ".datastore.formatters"
2184 validationCanFail = True
2185 datastoreStr = ["/tmp"]
2186 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2187 registryStr = "/gen3.sqlite3"
2189 def testPathConstructor(self) -> None:
2190 """Independent test of constructor using PathLike."""
2191 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run)
2192 self.enterContext(butler)
2193 self.assertIsInstance(butler, Butler)
2195 # And again with a Path object with the butler yaml
2196 path = pathlib.Path(self.tmpConfigFile)
2197 butler = Butler.from_config(path, writeable=False)
2198 self.enterContext(butler)
2199 self.assertIsInstance(butler, Butler)
2201 # And again with a Path object without the butler yaml
2202 # (making sure we skip it if the tmp config doesn't end
2203 # in butler.yaml -- which is the case for a subclass)
2204 if self.tmpConfigFile.endswith("butler.yaml"):
2205 path = pathlib.Path(os.path.dirname(self.tmpConfigFile))
2206 butler = Butler.from_config(path, writeable=False)
2207 self.enterContext(butler)
2208 self.assertIsInstance(butler, Butler)
2210 def testExportTransferCopy(self) -> None:
2211 """Test local export using all transfer modes"""
2212 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2213 exportButler = self.runPutGetTest(storageClass, "test_metric")
2214 # Test that the repo actually has at least one dataset.
2215 datasets = list(exportButler.registry.queryDatasets(..., collections=...))
2216 self.assertGreater(len(datasets), 0)
2217 uris = [exportButler.getURI(d) for d in datasets]
2218 assert isinstance(exportButler._datastore, FileDatastore)
2219 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]]
2221 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris]
2223 for path in pathsInStore:
2224 # Assume local file system
2225 assert path is not None
2226 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}")
2228 for transfer in ("copy", "link", "symlink", "relsymlink"):
2229 with safeTestTempDir(TESTDIR) as exportDir:
2230 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export:
2231 export.saveDatasets(datasets)
2232 for path in pathsInStore:
2233 assert path is not None
2234 self.assertTrue(
2235 self.checkFileExists(exportDir, path),
2236 f"Check that mode {transfer} exported files",
2237 )
2239 def testPytypeCoercion(self) -> None:
2240 """Test python type coercion on Butler.get and put."""
2241 # Store some data with the normal example storage class.
2242 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2243 datasetTypeName = "test_metric"
2244 butler = self.runPutGetTest(storageClass, datasetTypeName)
2246 dataId = {"instrument": "DummyCamComp", "visit": 423}
2247 metric = butler.get(datasetTypeName, dataId=dataId)
2248 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
2250 datasetType_ori = butler.get_dataset_type(datasetTypeName)
2251 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
2253 # Now need to hack the registry dataset type definition.
2254 # There is no API for this.
2255 assert isinstance(butler._registry, SqlRegistry)
2256 manager = butler._registry._managers.datasets
2257 assert hasattr(manager, "_db") and hasattr(manager, "_static")
2258 manager._db.update(
2259 manager._static.dataset_type,
2260 {"name": datasetTypeName},
2261 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
2262 )
2264 # Force reset of dataset type cache
2265 butler.registry.refresh()
2267 datasetType_new = butler.get_dataset_type(datasetTypeName)
2268 self.assertEqual(datasetType_new.name, datasetType_ori.name)
2269 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
2271 metric_model = butler.get(datasetTypeName, dataId=dataId)
2272 self.assertNotEqual(type(metric_model), type(metric))
2273 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
2275 # Put the model and read it back to show that everything now
2276 # works as normal.
2277 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
2278 metric_model_new = butler.get(metric_ref)
2279 self.assertEqual(metric_model_new, metric_model)
2281 # Hack the storage class again to something that will fail on the
2282 # get with no conversion class.
2283 manager._db.update(
2284 manager._static.dataset_type,
2285 {"name": datasetTypeName},
2286 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
2287 )
2288 butler.registry.refresh()
2290 with self.assertRaises(ValueError):
2291 butler.get(datasetTypeName, dataId=dataId)
2293 def test_provenance(self):
2294 """Test that provenance is attached on put."""
2295 run = "test_run"
2296 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
2297 metric = MetricsExampleModel(
2298 summary={"AM1": 5.2, "AM2": 30.6},
2299 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
2300 data=[563, 234, 456.7, 752, 8, 9, 27],
2301 )
2302 # Provenance can be attached to the object being put. Whether
2303 # it is or not is dependent on the formatter. For this test we
2304 # copy on adding provenance to ensure they differ.
2305 self.assertIsNone(metric.dataset_id)
2306 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
2307 self.assertIsNone(metric.dataset_id)
2308 metric_2 = butler.get(metric_ref)
2309 self.assertEqual(metric_2.data, metric.data)
2310 self.assertEqual(metric_2.dataset_id, metric_ref.id)
2311 self.assertIsNone(metric_2.provenance)
2313 # Put with provenance.
2314 prov = DatasetProvenance(quantum_id=uuid.uuid4())
2315 prov.add_input(metric_ref)
2316 prov.add_extra_provenance(metric_ref.id, {"answer": 42})
2317 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
2318 metric_3 = butler.get(metric_ref2)
2319 self.assertEqual(metric_3.provenance, prov)
2321 # Check that we can extract provenance from dict form.
2322 prov_dict = prov.to_flat_dict(metric_ref2)
2323 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2324 self.assertEqual(ref_from_prov, metric_ref2)
2325 # Direct __eq__ of the provenance does not work because one side
2326 # includes dimension records.
2327 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2328 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id)
2329 self.assertEqual(prov_from_prov.extras, prov.extras)
2331 # Force a bad ID into the dict.
2332 prov_dict["id"] = uuid.uuid4()
2333 with self.assertRaises(ValueError):
2334 DatasetProvenance.from_flat_dict(prov_dict, butler)
2335 del prov_dict["id"]
2336 prov_dict["input 0 id"] = uuid.uuid4()
2337 with self.assertRaises(ValueError):
2338 DatasetProvenance.from_flat_dict(prov_dict, butler)
2340 # Check that simple types can be reconstructed with non-standard
2341 # separators.
2342 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True)
2343 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler)
2344 self.assertEqual(ref_from_prov, metric_ref2)
2345 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs})
2347 with self.assertRaises(ValueError):
2348 DatasetProvenance.from_flat_dict({"unknown": 42}, butler)
2350 def test_specialized_file_datasets_functions(self):
2351 """Test a workflow used in Prompt Processing where we export datasets
2352 from one repository and write them in-place to the datastore of
2353 another, without immediately inserting registry entries for the
2354 datasets.
2355 """
2356 repo = MetricTestRepo.create_from_butler(
2357 self.create_empty_butler(writeable=True),
2358 self.tmpConfigFile,
2359 "StructuredCompositeReadCompNoDisassembly",
2360 )
2361 source_butler = repo.butler
2363 # Test writing outputs to a FileDatastore.
2364 with tempfile.TemporaryDirectory() as tempdir:
2365 target_repo_config = Butler.makeRepo(tempdir)
2366 refs = [repo.ref1, repo.ref2]
2367 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2368 self.assertEqual(len(datasets), 2)
2369 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2370 for dataset in datasets:
2371 path = ResourcePath(dataset.path, forceAbsolute=False)
2372 # Paths should be relative paths to the target datastore.
2373 self.assertFalse(path.isabs())
2374 # Files should have been copied into the target datastore
2375 self.assertTrue(ResourcePath(tempdir).join(path).exists())
2377 # Make sure the target Butler can ingest the datasets.
2378 target_butler = Butler(target_repo_config, writeable=True)
2379 self.enterContext(target_butler)
2380 target_butler.transfer_dimension_records_from(source_butler, refs)
2381 target_butler.ingest(*datasets, transfer=None)
2382 self.assertIsNotNone(target_butler.get(repo.ref1))
2383 self.assertIsNotNone(target_butler.get(repo.ref2))
2385 # Giving an empty list of files is a no-op.
2386 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), [])
2387 self.assertEqual(len(no_datasets), 0)
2389 # Test writing outputs to a ChainedDatastore.
2390 with tempfile.TemporaryDirectory() as tempdir:
2391 # Set up a second dataset type, so we can split the files across
2392 # multiple datastore roots.
2393 dt1 = repo.datasetType
2394 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass)
2395 source_butler.registry.registerDatasetType(dt2)
2396 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2)
2397 config = Config.fromString(
2398 f"""
2399 datastore:
2400 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore
2401 datastore_constraints:
2402 - constraints:
2403 accept:
2404 - {dt1.name}
2405 - constraints:
2406 accept:
2407 - {dt2.name}
2408 datastores:
2409 - datastore:
2410 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2411 root: <butlerRoot>/FileDatastore_0
2412 - datastore:
2413 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
2414 root: <butlerRoot>/FileDatastore_1
2415 """
2416 )
2417 target_repo_config = Butler.makeRepo(tempdir, config)
2418 refs = [repo.ref1, repo.ref2, other_ref]
2419 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs)
2420 self.assertEqual(len(datasets), 3)
2421 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets})
2422 for dataset in datasets:
2423 path = ResourcePath(dataset.path, forceAbsolute=False)
2424 # Paths should be relative paths to the target datastore.
2425 self.assertFalse(path.isabs())
2426 # Files should have been split up between the two datastores
2427 # in the chain.
2428 datastore_root = ResourcePath(tempdir)
2429 if dataset.refs[0].datasetType.name == dt1.name:
2430 datastore_root = datastore_root.join("FileDatastore_0")
2431 else:
2432 datastore_root = datastore_root.join("FileDatastore_1")
2433 self.assertTrue(datastore_root.join(path).exists())
2435 # Make sure the target Butler can ingest the datasets.
2436 target_butler = Butler(target_repo_config, writeable=True)
2437 self.enterContext(target_butler)
2438 target_butler.transfer_dimension_records_from(source_butler, refs)
2439 target_butler.ingest(*datasets, transfer=None)
2440 self.assertIsNotNone(target_butler.get(repo.ref1))
2441 self.assertIsNotNone(target_butler.get(repo.ref2))
2442 self.assertIsNotNone(target_butler.get(other_ref))
2445class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2446 """PosixDatastore specialization of a butler using Postgres"""
2448 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2449 fullConfigKey = ".datastore.formatters"
2450 validationCanFail = True
2451 datastoreStr = ["/tmp"]
2452 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"]
2453 registryStr = "PostgreSQL@test"
2455 @classmethod
2456 def setUpClass(cls) -> None:
2457 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
2458 super().setUpClass()
2460 def setUp(self) -> None:
2461 # Need to add a registry section to the config.
2462 self._temp_config = False
2463 config = Config(self.configFile)
2464 self.postgresql.patch_butler_config(config)
2465 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
2466 config.dump(fh)
2467 self.configFile = fh.name
2468 self._temp_config = True
2469 super().setUp()
2471 def tearDown(self) -> None:
2472 if self._temp_config and os.path.exists(self.configFile):
2473 os.remove(self.configFile)
2474 super().tearDown()
2476 def testMakeRepo(self) -> None:
2477 # The base class test assumes that it's using sqlite and assumes
2478 # the config file is acceptable to sqlite.
2479 raise unittest.SkipTest("Postgres config is not compatible with this test.")
2482class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase):
2483 """Test that Butler with a Postgres registry still works after cloning."""
2485 def create_butler(
2486 self,
2487 run: str,
2488 storageClass: StorageClass | str,
2489 datasetTypeName: str,
2490 metrics: ButlerMetrics | None = None,
2491 ) -> tuple[DirectButler, DatasetType]:
2492 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2493 return butler.clone(run=run, metrics=metrics), datasetType
2496class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
2497 """InMemoryDatastore specialization of a butler"""
2499 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")
2500 fullConfigKey = None
2501 useTempRoot = False
2502 validationCanFail = False
2503 datastoreStr = ["datastore='InMemory"]
2504 datastoreName = ["InMemoryDatastore@"]
2505 registryStr = "/gen3.sqlite3"
2507 def testIngest(self) -> None:
2508 pass
2510 def test_ingest_zip(self) -> None:
2511 pass
2514class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase):
2515 """Test that a Butler with a Sqlite registry still works after cloning."""
2517 def create_butler(
2518 self,
2519 run: str,
2520 storageClass: StorageClass | str,
2521 datasetTypeName: str,
2522 metrics: ButlerMetrics | None = None,
2523 ) -> tuple[DirectButler, DatasetType]:
2524 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics)
2525 return butler.clone(run=run), datasetType
2528class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2529 """PosixDatastore specialization"""
2531 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
2532 fullConfigKey = ".datastore.datastores.1.formatters"
2533 validationCanFail = True
2534 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"]
2535 datastoreName = [
2536 "InMemoryDatastore@",
2537 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1",
2538 "SecondDatastore",
2539 ]
2540 registryStr = "/gen3.sqlite3"
2542 def testPruneDatasets(self) -> None:
2543 # This test relies on manipulating files out-of-band, which is
2544 # impossible for this configuration because of the InMemoryDatastore in
2545 # the ChainedDatastore.
2546 pass
2549class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase):
2550 """Test that a yaml file in one location can refer to a root in another."""
2552 datastoreStr = ["dir1"]
2553 # Disable the makeRepo test since we are deliberately not using
2554 # butler.yaml as the config name.
2555 fullConfigKey = None
2557 def setUp(self) -> None:
2558 self.root = makeTestTempDir(TESTDIR)
2560 # Make a new repository in one place
2561 self.dir1 = os.path.join(self.root, "dir1")
2562 Butler.makeRepo(self.dir1, config=Config(self.configFile))
2564 # Move the yaml file to a different place and add a "root"
2565 self.dir2 = os.path.join(self.root, "dir2")
2566 os.makedirs(self.dir2, exist_ok=True)
2567 configFile1 = os.path.join(self.dir1, "butler.yaml")
2568 config = Config(configFile1)
2569 config["root"] = self.dir1
2570 configFile2 = os.path.join(self.dir2, "butler2.yaml")
2571 config.dumpToUri(configFile2)
2572 os.remove(configFile1)
2573 self.tmpConfigFile = configFile2
2575 def testFileLocations(self) -> None:
2576 self.assertNotEqual(self.dir1, self.dir2)
2577 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml")))
2578 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml")))
2579 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3")))
2582class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase):
2583 """Test that a config file created by makeRepo outside of repo works."""
2585 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2587 def setUp(self) -> None:
2588 self.root = makeTestTempDir(TESTDIR)
2589 self.root2 = makeTestTempDir(TESTDIR)
2591 self.tmpConfigFile = os.path.join(self.root2, "different.yaml")
2592 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2594 def tearDown(self) -> None:
2595 if os.path.exists(self.root2):
2596 shutil.rmtree(self.root2, ignore_errors=True)
2597 super().tearDown()
2599 def testConfigExistence(self) -> None:
2600 c = Config(self.tmpConfigFile)
2601 uri_config = ResourcePath(c["root"])
2602 uri_expected = ResourcePath(self.root, forceDirectory=True)
2603 self.assertEqual(uri_config.geturl(), uri_expected.geturl())
2604 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path")
2606 def testPutGet(self) -> None:
2607 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
2608 self.runPutGetTest(storageClass, "test_metric")
2611class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase):
2612 """Test that a config file created by makeRepo outside of repo works."""
2614 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2616 def setUp(self) -> None:
2617 self.root = makeTestTempDir(TESTDIR)
2618 self.root2 = makeTestTempDir(TESTDIR)
2620 self.tmpConfigFile = self.root2
2621 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2623 def testConfigExistence(self) -> None:
2624 # Append the yaml file else Config constructor does not know the file
2625 # type.
2626 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml")
2627 super().testConfigExistence()
2630class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase):
2631 """Test that a config file created by makeRepo outside of repo works."""
2633 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
2635 def setUp(self) -> None:
2636 self.root = makeTestTempDir(TESTDIR)
2637 self.root2 = makeTestTempDir(TESTDIR)
2639 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl()
2640 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile)
2643@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!")
2644class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
2645 """S3Datastore specialization of a butler; an S3 storage Datastore +
2646 a local in-memory SqlRegistry.
2647 """
2649 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")
2650 fullConfigKey = None
2651 validationCanFail = True
2653 bucketName = "anybucketname"
2654 """Name of the Bucket that will be used in the tests. The name is read from
2655 the config file used with the tests during set-up.
2656 """
2658 root = "butlerRoot/"
2659 """Root repository directory expected to be used in case useTempRoot=False.
2660 Otherwise the root is set to a 20 characters long randomly generated string
2661 during set-up.
2662 """
2664 datastoreStr = [f"datastore={root}"]
2665 """Contains all expected root locations in a format expected to be
2666 returned by Butler stringification.
2667 """
2669 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]
2670 """The expected format of the S3 Datastore string."""
2672 registryStr = "/gen3.sqlite3"
2673 """Expected format of the Registry string."""
2675 mock_aws = mock_aws()
2676 """The mocked s3 interface from moto."""
2678 def genRoot(self) -> str:
2679 """Return a random string of len 20 to serve as a root
2680 name for the temporary bucket repo.
2682 This is equivalent to tempfile.mkdtemp as this is what self.root
2683 becomes when useTempRoot is True.
2684 """
2685 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20))
2686 return rndstr + "/"
2688 def setUp(self) -> None:
2689 config = Config(self.configFile)
2690 uri = ResourcePath(config[".datastore.datastore.root"])
2691 self.bucketName = uri.netloc
2693 # Enable S3 mocking of tests.
2694 self.enterContext(clean_test_environment_for_s3())
2695 self.mock_aws.start()
2697 if self.useTempRoot:
2698 self.root = self.genRoot()
2699 rooturi = f"s3://{self.bucketName}/{self.root}"
2700 config.update({"datastore": {"datastore": {"root": rooturi}}})
2702 # need local folder to store registry database
2703 self.reg_dir = makeTestTempDir(TESTDIR)
2704 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"
2706 # MOTO needs to know that we expect Bucket bucketname to exist
2707 # (this used to be the class attribute bucketName)
2708 s3 = boto3.resource("s3")
2709 s3.create_bucket(Bucket=self.bucketName)
2711 self.datastoreStr = [f"datastore='{rooturi}'"]
2712 self.datastoreName = [f"FileDatastore@{rooturi}"]
2713 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)
2714 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")
2716 def tearDown(self) -> None:
2717 s3 = boto3.resource("s3")
2718 bucket = s3.Bucket(self.bucketName)
2719 try:
2720 bucket.objects.all().delete()
2721 except botocore.exceptions.ClientError as e:
2722 if e.response["Error"]["Code"] == "404":
2723 # the key was not reachable - pass
2724 pass
2725 else:
2726 raise
2728 bucket = s3.Bucket(self.bucketName)
2729 bucket.delete()
2731 # Stop the S3 mock.
2732 self.mock_aws.stop()
2734 if self.reg_dir is not None and os.path.exists(self.reg_dir):
2735 shutil.rmtree(self.reg_dir, ignore_errors=True)
2737 if self.useTempRoot and os.path.exists(self.root):
2738 shutil.rmtree(self.root, ignore_errors=True)
2740 super().tearDown()
2743class DatastoreTransfers(TestCaseMixin):
2744 """Base test setup for data transfers between butlers. The concrete tests
2745 for specific configurations are in other classes, below.
2746 """
2748 storageClassFactory: StorageClassFactory
2750 @classmethod
2751 def setUpClass(cls) -> None:
2752 cls.storageClassFactory = StorageClassFactory()
2754 def setUp(self) -> None:
2755 self.root = makeTestTempDir(TESTDIR)
2756 self.config = Config(self.configFile)
2758 # Some tests cause convertors to be replaced so ensure
2759 # the storage class factory is reset each time.
2760 self.storageClassFactory.reset()
2761 self.storageClassFactory.addFromConfig(self.configFile)
2763 def tearDown(self) -> None:
2764 removeTestTempDir(self.root)
2766 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler:
2767 if manager is None:
2768 manager = (
2769 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
2770 )
2771 config = Config(config_file if config_file is not None else self.configFile)
2772 config["registry", "managers", "datasets"] = manager
2773 butler = Butler.from_config(
2774 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True
2775 )
2776 self.enterContext(butler)
2777 return butler
2779 def assertButlerTransfers(
2780 self,
2781 purge: bool = False,
2782 storageClassName: str = "StructuredData",
2783 storageClassNameTarget: str | None = None,
2784 ) -> None:
2785 """Test that a run can be transferred to another butler."""
2786 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
2787 if storageClassNameTarget is not None:
2788 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget)
2789 else:
2790 storageClassTarget = storageClass
2792 datasetTypeName = "random_data"
2794 # Test will create 3 collections and we will want to transfer
2795 # two of those three.
2796 runs = ["run1", "run2", "other"]
2798 # Also want to use two different dataset types to ensure that
2799 # grouping works.
2800 datasetTypeNames = ["random_data", "random_data_2"]
2802 # Create the run collections in the source butler.
2803 for run in runs:
2804 self.source_butler.collections.register(run)
2806 # Create dimensions in source butler.
2807 n_exposures = 30
2808 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
2809 self.source_butler.registry.insertDimensionData(
2810 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
2811 )
2812 self.source_butler.registry.insertDimensionData(
2813 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"}
2814 )
2815 self.source_butler.registry.insertDimensionData(
2816 "day_obs",
2817 {
2818 "instrument": "DummyCamComp",
2819 "id": 20250101,
2820 },
2821 )
2823 for i in range(n_exposures):
2824 self.source_butler.registry.insertDimensionData(
2825 "group", {"instrument": "DummyCamComp", "name": f"group{i}"}
2826 )
2827 self.source_butler.registry.insertDimensionData(
2828 "exposure",
2829 {
2830 "instrument": "DummyCamComp",
2831 "id": i,
2832 "obs_id": f"exp{i}",
2833 "physical_filter": "d-r",
2834 "group": f"group{i}",
2835 "day_obs": 20250101,
2836 },
2837 )
2839 # Create dataset types in the source butler.
2840 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"])
2841 for datasetTypeName in datasetTypeNames:
2842 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
2843 self.source_butler.registry.registerDatasetType(datasetType)
2845 # Write a dataset to an unrelated run -- this will ensure that
2846 # we are rewriting integer dataset ids in the target if necessary.
2847 # Will not be relevant for UUID.
2848 run = "distraction"
2849 butler = Butler.from_config(butler=self.source_butler, run=run)
2850 self.enterContext(butler)
2851 butler.put(
2852 makeExampleMetrics(),
2853 datasetTypeName,
2854 exposure=1,
2855 instrument="DummyCamComp",
2856 physical_filter="d-r",
2857 )
2859 # Write some example metrics to the source
2860 butler = Butler.from_config(butler=self.source_butler)
2861 self.enterContext(butler)
2863 # Set of DatasetRefs that should be in the list of refs to transfer
2864 # but which will not be transferred.
2865 deleted: set[DatasetRef] = set()
2867 n_expected = 20 # Number of datasets expected to be transferred
2868 source_refs = []
2869 for i in range(n_exposures):
2870 # Put a third of datasets into each collection, only retain
2871 # two thirds.
2872 index = i % 3
2873 run = runs[index]
2874 datasetTypeName = datasetTypeNames[i % 2]
2876 metric = MetricsExample(
2877 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)]
2878 )
2879 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"}
2880 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run)
2882 # Remove the datastore record using low-level API, but only
2883 # for a specific index.
2884 if purge and index == 1:
2885 # For one of these delete the file as well.
2886 # This allows the "missing" code to filter the
2887 # file out.
2888 # Access the individual datastores.
2889 datastores = []
2890 if hasattr(butler._datastore, "datastores"):
2891 datastores.extend(butler._datastore.datastores)
2892 else:
2893 datastores.append(butler._datastore)
2895 if not deleted:
2896 # For a chained datastore we need to remove
2897 # files in each chain.
2898 for datastore in datastores:
2899 # The file might not be known to the datastore
2900 # if constraints are used.
2901 try:
2902 primary, uris = datastore.getURIs(ref)
2903 except FileNotFoundError:
2904 continue
2905 if primary and primary.scheme != "mem":
2906 primary.remove()
2907 for uri in uris.values():
2908 if uri.scheme != "mem":
2909 uri.remove()
2910 n_expected -= 1
2911 deleted.add(ref)
2913 # Remove the datastore record.
2914 for datastore in datastores:
2915 if hasattr(datastore, "removeStoredItemInfo"):
2916 datastore.removeStoredItemInfo(ref)
2918 if index < 2:
2919 source_refs.append(ref)
2920 if ref not in deleted:
2921 new_metric = butler.get(ref)
2922 self.assertEqual(new_metric, metric)
2924 # Create some bad dataset types to ensure we check for inconsistent
2925 # definitions.
2926 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList")
2927 for datasetTypeName in datasetTypeNames:
2928 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass)
2929 self.target_butler.registry.registerDatasetType(datasetType)
2930 with self.assertRaises(ConflictingDefinitionError) as cm:
2931 self.target_butler.transfer_from(self.source_butler, source_refs)
2932 self.assertIn("dataset type differs", str(cm.exception))
2934 # And remove the bad definitions.
2935 for datasetTypeName in datasetTypeNames:
2936 self.target_butler.registry.removeDatasetType(datasetTypeName)
2938 # Transfer without creating dataset types should fail.
2939 with self.assertRaises(KeyError):
2940 self.target_butler.transfer_from(self.source_butler, source_refs)
2942 # Transfer without creating dimensions should fail.
2943 with self.assertRaises(ConflictingDefinitionError) as cm:
2944 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True)
2945 self.assertIn("dimension", str(cm.exception))
2947 # The dry run test requires dataset types to exist. If we have
2948 # been given distinct storage classes for the target we have
2949 # to redefine at least one of the dataset types in the target butler.
2950 if storageClass != storageClassTarget:
2951 self.target_butler.registry.removeDatasetType(datasetTypeNames[0])
2952 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget)
2953 self.target_butler.registry.registerDatasetType(datasetType)
2955 # The failed transfer above leaves registry in an inconsistent
2956 # state because the run is created but then rolled back without
2957 # the collection cache being cleared. For now force a refresh.
2958 # Can remove with DM-35498.
2959 self.target_butler.registry.refresh()
2961 # Do a dry run -- this should not have any effect on the target butler.
2962 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True)
2964 # Transfer the records for one ref to test the alternative API.
2965 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2966 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]])
2967 self.assertIn("number of records transferred: 1", ";".join(log_cm.output))
2969 # Now transfer them to the second butler, including dimensions.
2970 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm:
2971 transferred = self.target_butler.transfer_from(
2972 self.source_butler,
2973 source_refs,
2974 register_dataset_types=True,
2975 transfer_dimensions=True,
2976 )
2977 self.assertEqual(len(transferred), n_expected)
2978 log_output = ";".join(log_cm.output)
2980 # A ChainedDatastore will use the in-memory datastore for mexists
2981 # so we can not rely on the mexists log message.
2982 self.assertIn("Number of datastore records found in source", log_output)
2983 self.assertIn("Creating output run", log_output)
2985 # Do the transfer twice to ensure that it will do nothing extra.
2986 # Only do this if purge=True because it does not work for int
2987 # dataset_id.
2988 if purge:
2989 # This should not need to register dataset types.
2990 transferred = self.target_butler.transfer_from(self.source_butler, source_refs)
2991 self.assertEqual(len(transferred), n_expected)
2993 with self.assertRaises((TypeError, AttributeError)):
2994 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore
2996 with self.assertRaises(ValueError):
2997 self.target_butler._datastore.transfer_from(
2998 self.source_butler._datastore, source_refs, transfer="split"
2999 )
3001 # Now try to get the same refs from the new butler.
3002 for ref in source_refs:
3003 if ref not in deleted:
3004 new_metric = self.target_butler.get(ref)
3005 old_metric = self.source_butler.get(ref)
3006 self.assertEqual(new_metric, old_metric)
3008 # Try again without implicit storage class conversion
3009 # triggered by using the source ref. This will do conversion
3010 # since the formatter will be returning the source python type.
3011 target_ref = self.target_butler.get_dataset(ref.id)
3012 if target_ref.datasetType.storageClass != ref.datasetType.storageClass:
3013 new_metric = self.target_butler.get(target_ref)
3014 self.assertNotEqual(type(new_metric), type(old_metric))
3016 # Remove the dataset from the target and put it again
3017 # as if it was the right type all along for this butler.
3018 self.target_butler.pruneDatasets(
3019 [target_ref], unstore=True, purge=True, disassociate=True
3020 )
3021 self.target_butler.put(new_metric, target_ref)
3022 new_new_metric = self.target_butler.get(target_ref)
3023 new_old_metric = self.target_butler.get(
3024 target_ref, storageClass=ref.datasetType.storageClass
3025 )
3026 self.assertEqual(new_new_metric, new_metric)
3027 self.assertEqual(new_old_metric, old_metric)
3029 # Now prune run2 collection and create instead a CHAINED collection.
3030 # This should block the transfer.
3031 self.target_butler.removeRuns(["run2"])
3032 self.target_butler.collections.register("run2", CollectionType.CHAINED)
3033 with self.assertRaises(CollectionTypeError):
3034 # Re-importing the run1 datasets can be problematic if they
3035 # use integer IDs so filter those out.
3036 to_transfer = [ref for ref in source_refs if ref.run == "run2"]
3037 self.target_butler.transfer_from(self.source_butler, to_transfer)
3040class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3041 """Test data transfers between butlers.
3043 Test for different managers. UUID to UUID and integer to integer are
3044 tested. UUID to integer is not supported since we do not currently
3045 want to allow that. Integer to UUID is supported with the caveat
3046 that UUID4 will be generated and this will be incorrect for raw
3047 dataset types. The test ignores that.
3048 """
3050 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3052 def create_butlers(
3053 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None
3054 ) -> None:
3055 self.source_butler = self.create_butler(manager1, "1", config_file=source_config)
3056 self.target_butler = self.create_butler(manager2, "2")
3058 def testTransferUuidToUuid(self) -> None:
3059 self.create_butlers()
3060 self.assertButlerTransfers()
3062 def testTransferFromChainedUuidToUuid(self) -> None:
3063 """Force the source butler to be a ChainedDatastore."""
3064 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml"))
3065 self.assertButlerTransfers()
3067 def testTransferFromIncompatibleUuidToUuid(self) -> None:
3068 """Force the source butler to be a incompatible datastore."""
3069 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml"))
3070 with self.assertRaises(NotImplementedError):
3071 self.assertButlerTransfers()
3073 def testTransferFromIncompatibleChainUuidToUuid(self) -> None:
3074 """Force the source butler to be a incompatible datastore."""
3075 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml"))
3076 with self.assertRaises(TypeError):
3077 self.assertButlerTransfers()
3079 def testTransferFromFileUuidToUuid(self) -> None:
3080 """Force the source butler to be a FileDatastore."""
3081 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml"))
3082 self.assertButlerTransfers()
3084 def testTransferMissing(self) -> None:
3085 """Test transfers where datastore records are missing.
3087 This is how execution butler works.
3088 """
3089 self.create_butlers()
3091 # Configure the source butler to allow trust.
3092 self.source_butler._datastore._set_trust_mode(True)
3094 self.assertButlerTransfers(purge=True)
3096 def testTransferMissingDisassembly(self) -> None:
3097 """Test transfers where datastore records are missing.
3099 This is how execution butler works.
3100 """
3101 self.create_butlers()
3103 # Configure the source butler to allow trust.
3104 self.source_butler._datastore._set_trust_mode(True)
3106 # Test disassembly.
3107 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite")
3109 def testTransferDifferingStorageClasses(self) -> None:
3110 """Test transfers when the source butler dataset type has a different
3111 but compatible storage class.
3112 """
3113 self.create_butlers()
3115 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion")
3117 def testTransferDifferingStorageClassesDisassembly(self) -> None:
3118 """Test transfers when the source butler dataset type has a different
3119 but compatible storage class and where the source butler has
3120 disassembled.
3121 """
3122 self.create_butlers()
3124 self.assertButlerTransfers(
3125 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion"
3126 )
3128 def testAbsoluteURITransferDirect(self) -> None:
3129 """Test transfer using an absolute URI."""
3130 self._absolute_transfer("auto")
3132 def testAbsoluteURITransferCopy(self) -> None:
3133 """Test transfer using an absolute URI."""
3134 self._absolute_transfer("copy")
3136 def _absolute_transfer(self, transfer: str) -> None:
3137 self.create_butlers()
3139 storageClassName = "StructuredData"
3140 storageClass = self.storageClassFactory.getStorageClass(storageClassName)
3141 datasetTypeName = "random_data"
3142 run = "run1"
3143 self.source_butler.collections.register(run)
3145 dimensions = self.source_butler.dimensions.conform(())
3146 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3147 self.source_butler.registry.registerDatasetType(datasetType)
3149 metrics = makeExampleMetrics()
3150 with ResourcePath.temporary_uri(suffix=".json") as temp:
3151 dataId = DataCoordinate.make_empty(self.source_butler.dimensions)
3152 source_refs = [DatasetRef(datasetType, dataId, run=run)]
3153 temp.write(json.dumps(metrics.exportAsDict()).encode())
3154 dataset = FileDataset(path=temp, refs=source_refs)
3155 self.source_butler.ingest(dataset, transfer="direct")
3157 self.target_butler.transfer_from(
3158 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer
3159 )
3161 uri = self.target_butler.getURI(dataset.refs[0])
3162 if transfer == "auto":
3163 self.assertEqual(uri, temp)
3164 else:
3165 self.assertNotEqual(uri, temp)
3167 def test_shared_dimension_group(self):
3168 """Test internal logic that divides dataset types by dimension group
3169 when doing registry updates.
3170 """
3171 self.create_butlers()
3172 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
3173 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
3175 source_butler = self.source_butler
3176 target_butler = self.target_butler
3178 # Create a dataset type with the same dimensions as the 'bias' dataset
3179 # type from base.yaml
3180 dataset_type = DatasetType(
3181 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
3182 )
3183 source_butler.registry.registerDatasetType(dataset_type)
3184 # This has the same data ID as one of the bias datasets in
3185 # datasets.yaml.
3186 test_ref = source_butler.registry.insertDatasets(
3187 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
3188 )[0]
3190 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
3191 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
3192 refs = [test_ref, *biases, *flats]
3194 # Test setup will be even more convoluted if we want the datastore to
3195 # actually transfer files. For testing the dimension group behavior,
3196 # we really only care about the registry.
3197 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
3198 mock.return_value = (set(refs), set())
3199 target_butler.transfer_from(
3200 source_butler,
3201 refs,
3202 transfer=None,
3203 register_dataset_types=True,
3204 skip_missing=False,
3205 transfer_dimensions=True,
3206 )
3208 transferred_test_ref = target_butler.find_dataset(
3209 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3210 )
3211 self.assertEqual(transferred_test_ref.id, test_ref.id)
3213 transferred_bias = target_butler.find_dataset(
3214 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
3215 )
3216 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
3218 transferred_flat = target_butler.find_dataset(
3219 "flat",
3220 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
3221 collections="imported_r",
3222 )
3223 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
3226class ChainedDatastoreTransfers(PosixDatastoreTransfers):
3227 """Test transfers using a chained datastore."""
3229 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml")
3232@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3233class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase):
3234 """Test ``transfer_from`` involving Butler server."""
3236 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3238 def test_transfers_from_remote_to_direct(self) -> None:
3239 from lsst.daf.butler.remote_butler._remote_file_transfer_source import (
3240 mock_file_transfer_uris_for_unit_test,
3241 )
3243 self.target_butler = self.create_butler(None, "2")
3244 with create_test_server(TESTDIR) as server:
3245 self.source_butler = server.hybrid_butler
3247 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath:
3248 # The Butler server returns HTTP URIs with a domain name that
3249 # is not resolvable because there is no actual HTTP server
3250 # involved in these tests. Strip this first layer of
3251 # indirection, and return the target of the redirect instead.
3252 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers)
3253 return ResourcePath(str(response.next_request.url))
3255 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url):
3256 self.assertButlerTransfers()
3259class NullDatastoreTestCase(unittest.TestCase):
3260 """Test that we can fall back to a null datastore."""
3262 # Need a good config to create the repo.
3263 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
3264 storageClassFactory: StorageClassFactory
3266 @classmethod
3267 def setUpClass(cls) -> None:
3268 cls.storageClassFactory = StorageClassFactory()
3269 cls.storageClassFactory.addFromConfig(cls.configFile)
3271 def setUp(self) -> None:
3272 """Create a new butler root for each test."""
3273 self.root = makeTestTempDir(TESTDIR)
3274 Butler.makeRepo(self.root, config=Config(self.configFile))
3276 def tearDown(self) -> None:
3277 removeTestTempDir(self.root)
3279 def test_fallback(self) -> None:
3280 # Read the butler config and mess with the datastore section.
3281 config_path = os.path.join(self.root, "butler.yaml")
3282 bad_config = Config(config_path)
3283 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore"
3284 bad_config.dumpToUri(config_path)
3286 with self.assertRaises(RuntimeError):
3287 Butler(self.root, without_datastore=False)
3289 with self.assertRaises(RuntimeError):
3290 Butler.from_config(self.root, without_datastore=False)
3292 butler = Butler.from_config(self.root, writeable=True, without_datastore=True)
3293 self.enterContext(butler)
3294 self.assertIsInstance(butler._datastore, NullDatastore)
3296 # Check that registry is working.
3297 butler.collections.register("MYRUN")
3298 collections = butler.collections.query("*")
3299 self.assertIn("MYRUN", set(collections))
3301 # Create a ref.
3302 dimensions = butler.dimensions.conform([])
3303 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict")
3304 datasetTypeName = "metric"
3305 datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
3306 butler.registry.registerDatasetType(datasetType)
3307 ref = DatasetRef(datasetType, {}, run="MYRUN")
3309 # Check that datastore will complain.
3310 with self.assertRaises(FileNotFoundError):
3311 butler.get(ref)
3312 with self.assertRaises(FileNotFoundError):
3313 butler.getURI(ref)
3316@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3317class ButlerServerTests(FileDatastoreButlerTests):
3318 """Test RemoteButler and Butler server."""
3320 configFile = None
3321 predictionSupported = False
3322 trustModeSupported = False
3324 postgres: TemporaryPostgresInstance | None
3326 def setUp(self):
3327 self.server_instance = self.enterContext(create_test_server(TESTDIR))
3329 def tearDown(self):
3330 pass
3332 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool:
3333 # S3 pre-signed URLs may end up with differing expiration times in the
3334 # query parameters, so ignore query parameters when comparing.
3335 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path
3337 def create_empty_butler(
3338 self,
3339 run: str | None = None,
3340 writeable: bool | None = None,
3341 metrics: ButlerMetrics | None = None,
3342 cleanup: bool = True,
3343 ) -> Butler:
3344 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics)
3346 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None:
3347 # Can't delete a file via S3 signed URLs, so we need to reach in
3348 # through DirectButler to delete the dataset.
3349 uri = self.server_instance.direct_butler.getURI(ref)
3350 uri.remove()
3352 def testConstructor(self):
3353 # RemoteButler constructor is tested in test_server.py and
3354 # test_remote_butler.py.
3355 pass
3357 def testDafButlerRepositories(self):
3358 # Loading of RemoteButler via repository index is tested in
3359 # test_server.py.
3360 pass
3362 def testGetDatasetTypes(self) -> None:
3363 # This is mostly a test of validateConfiguration, which is for
3364 # validating Datastore configuration and thus isn't relevant to
3365 # RemoteButler.
3366 pass
3368 def testMakeRepo(self) -> None:
3369 # Only applies to DirectButler.
3370 pass
3372 # Pickling not yet implemented for RemoteButler/HybridButler.
3373 @unittest.expectedFailure
3374 def testPickle(self) -> None:
3375 return super().testPickle()
3377 def testStringification(self) -> None:
3378 self.assertEqual(
3379 str(self.server_instance.remote_butler),
3380 "RemoteButler(https://test.example/api/butler/repo/testrepo/)",
3381 )
3383 def testTransaction(self) -> None:
3384 # Transactions will never be supported for RemoteButler.
3385 pass
3387 def testPutTemplates(self) -> None:
3388 # The Butler server instance is configured with different file naming
3389 # templates than this test is expecting.
3390 pass
3393@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3394class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase):
3395 """Tests for RemoteButler's registry shim, with a SQLite DB backing the
3396 server.
3397 """
3399 postgres = None
3402@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
3403class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase):
3404 """Tests for RemoteButler's registry shim, with a Postgres DB backing the
3405 server.
3406 """
3408 @classmethod
3409 def setUpClass(cls):
3410 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
3411 super().setUpClass()
3414def setup_module(module: types.ModuleType) -> None:
3415 """Set up the module for pytest."""
3416 clean_environment()
3419def _get_test_data_path(filename: str) -> ResourcePath:
3420 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}")
3423if __name__ == "__main__":
3424 clean_environment()
3425 unittest.main()