Coverage for tests / test_butler.py: 13%

1858 statements  

« prev     ^ index     » next       coverage.py v7.13.2, created at 2026-02-27 01:21 -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/>. 

27 

28"""Tests for Butler.""" 

29 

30from __future__ import annotations 

31 

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 

50 

51try: 

52 import boto3 

53 import botocore 

54 

55 from lsst.resources.s3utils import clean_test_environment_for_s3 

56 

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 

63 

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 

67 

68 

69import astropy.time 

70from sqlalchemy.exc import IntegrityError 

71 

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 DimensionRecord, 

87 FileDataset, 

88 NoDefaultCollectionError, 

89 StorageClassFactory, 

90 ValidationError, 

91 script, 

92) 

93from lsst.daf.butler._rubin.file_datasets import transfer_datasets_to_datastore 

94from lsst.daf.butler._rubin.temporary_for_ingest import TemporaryForIngest 

95from lsst.daf.butler.datastore import NullDatastore 

96from lsst.daf.butler.datastore.file_templates import FileTemplate, FileTemplateValidationError 

97from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex 

98from lsst.daf.butler.datastores.fileDatastore import FileDatastore 

99from lsst.daf.butler.direct_butler import DirectButler 

100from lsst.daf.butler.registry import ( 

101 CollectionError, 

102 CollectionTypeError, 

103 ConflictingDefinitionError, 

104 DataIdValueError, 

105 DatasetTypeExpressionError, 

106 MissingCollectionError, 

107 OrphanedRecordError, 

108) 

109from lsst.daf.butler.registry.sql_registry import SqlRegistry 

110from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG 

111from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter 

112from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db 

113from lsst.daf.butler.tests.server_available import butler_server_import_error, butler_server_is_available 

114from lsst.daf.butler.tests.utils import ( 

115 MetricTestRepo, 

116 TestCaseMixin, 

117 create_populated_sqlite_registry, 

118 makeTestTempDir, 

119 removeTestTempDir, 

120 safeTestTempDir, 

121) 

122from lsst.resources import ResourcePath 

123from lsst.resources.http import HttpResourcePath 

124from lsst.utils import doImportType 

125from lsst.utils.introspection import get_full_type_name 

126 

127if butler_server_is_available: 

128 from lsst.daf.butler.tests.server import create_test_server 

129 

130 

131if TYPE_CHECKING: 

132 import types 

133 

134 from lsst.daf.butler import DimensionGroup, Registry, StorageClass 

135 

136TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

137 

138 

139def clean_environment() -> None: 

140 """Remove external environment variables that affect the tests.""" 

141 for k in ("DAF_BUTLER_REPOSITORY_INDEX",): 

142 os.environ.pop(k, None) 

143 

144 

145def makeExampleMetrics() -> MetricsExample: 

146 """Return example dataset suitable for tests.""" 

147 return MetricsExample( 

148 {"AM1": 5.2, "AM2": 30.6}, 

149 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

150 [563, 234, 456.7, 752, 8, 9, 27], 

151 ) 

152 

153 

154class TransactionTestError(Exception): 

155 """Specific error for testing transactions, to prevent misdiagnosing 

156 that might otherwise occur when a standard exception is used. 

157 """ 

158 

159 pass 

160 

161 

162class ButlerConfigTests(unittest.TestCase): 

163 """Simple tests for ButlerConfig that are not tested in any other test 

164 cases. 

165 """ 

166 

167 def testSearchPath(self) -> None: 

168 configFile = os.path.join(TESTDIR, "config", "basic", "butler.yaml") 

169 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm: 

170 config1 = ButlerConfig(configFile) 

171 self.assertNotIn("testConfigs", "\n".join(cm.output)) 

172 

173 overrideDirectory = os.path.join(TESTDIR, "config", "testConfigs") 

174 with self.assertLogs("lsst.daf.butler", level="DEBUG") as cm: 

175 config2 = ButlerConfig(configFile, searchPaths=[overrideDirectory]) 

176 self.assertIn("testConfigs", "\n".join(cm.output)) 

177 

178 key = ("datastore", "records", "table") 

179 self.assertNotEqual(config1[key], config2[key]) 

180 self.assertEqual(config2[key], "override_record") 

181 

182 

183class ButlerPutGetTests(TestCaseMixin): 

184 """Helper method for running a suite of put/get tests from different 

185 butler configurations. 

186 """ 

187 

188 root: str 

189 default_run = "ingésτ😺" 

190 storageClassFactory: StorageClassFactory 

191 configFile: str | None 

192 tmpConfigFile: str 

193 

194 @staticmethod 

195 def addDatasetType( 

196 datasetTypeName: str, dimensions: DimensionGroup, storageClass: StorageClass | str, registry: Registry 

197 ) -> DatasetType: 

198 """Create a DatasetType and register it""" 

199 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

200 registry.registerDatasetType(datasetType) 

201 return datasetType 

202 

203 @classmethod 

204 def setUpClass(cls) -> None: 

205 cls.storageClassFactory = StorageClassFactory() 

206 if cls.configFile is not None: 

207 cls.storageClassFactory.addFromConfig(cls.configFile) 

208 

209 def assertGetComponents( 

210 self, 

211 butler: Butler, 

212 datasetRef: DatasetRef, 

213 components: tuple[str, ...], 

214 reference: Any, 

215 collections: Any = None, 

216 ) -> None: 

217 datasetType = datasetRef.datasetType 

218 dataId = datasetRef.dataId 

219 deferred = butler.getDeferred(datasetRef) 

220 

221 for component in components: 

222 compTypeName = datasetType.componentTypeName(component) 

223 result = butler.get(compTypeName, dataId, collections=collections) 

224 self.assertEqual(result, getattr(reference, component)) 

225 result_deferred = deferred.get(component=component) 

226 self.assertEqual(result_deferred, result) 

227 

228 def tearDown(self) -> None: 

229 if self.root is not None: 

230 removeTestTempDir(self.root) 

231 

232 def create_empty_butler( 

233 self, 

234 run: str | None = None, 

235 writeable: bool | None = None, 

236 metrics: ButlerMetrics | None = None, 

237 cleanup: bool = True, 

238 ): 

239 """Create a Butler for the test repository, without inserting test 

240 data. 

241 """ 

242 butler = Butler.from_config(self.tmpConfigFile, run=run, writeable=writeable, metrics=metrics) 

243 if cleanup: 

244 self.enterContext(butler) 

245 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration" 

246 return butler 

247 

248 def create_butler( 

249 self, 

250 run: str, 

251 storageClass: StorageClass | str, 

252 datasetTypeName: str, 

253 metrics: ButlerMetrics | None = None, 

254 ) -> tuple[Butler, DatasetType]: 

255 """Create a Butler for the test repository and insert some test data 

256 into it. 

257 """ 

258 butler = self.create_empty_butler(run=run, metrics=metrics) 

259 

260 collections = set(butler.collections.query("*")) 

261 self.assertEqual(collections, {run}) 

262 # Create and register a DatasetType 

263 dimensions = butler.dimensions.conform(["instrument", "visit"]) 

264 

265 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry) 

266 

267 # Add needed Dimensions 

268 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

269 butler.registry.insertDimensionData( 

270 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

271 ) 

272 butler.registry.insertDimensionData( 

273 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"} 

274 ) 

275 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101}) 

276 visit_start = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai") 

277 visit_end = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai") 

278 butler.registry.insertDimensionData( 

279 "visit", 

280 { 

281 "instrument": "DummyCamComp", 

282 "id": 423, 

283 "name": "fourtwentythree", 

284 "physical_filter": "d-r", 

285 "datetime_begin": visit_start, 

286 "datetime_end": visit_end, 

287 "day_obs": 20200101, 

288 }, 

289 ) 

290 

291 # Add more visits for some later tests 

292 for visit_id in (424, 425): 

293 butler.registry.insertDimensionData( 

294 "visit", 

295 { 

296 "instrument": "DummyCamComp", 

297 "id": visit_id, 

298 "name": f"fourtwentyfour_{visit_id}", 

299 "physical_filter": "d-r", 

300 "day_obs": 20200101, 

301 }, 

302 ) 

303 return butler, datasetType 

304 

305 def runPutGetTest(self, storageClass: StorageClass, datasetTypeName: str) -> Butler: 

306 # New datasets will be added to run and tag, but we will only look in 

307 # tag when looking up datasets. 

308 run = self.default_run 

309 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName) 

310 assert butler.run is not None 

311 

312 # Create and store a dataset 

313 metric = makeExampleMetrics() 

314 dataId = butler.registry.expandDataId({"instrument": "DummyCamComp", "visit": 423}) 

315 

316 # Dataset should not exist if we haven't added it 

317 with self.assertRaises(DatasetNotFoundError): 

318 butler.get(datasetTypeName, dataId) 

319 

320 # Put and remove the dataset once as a DatasetRef, once as a dataId, 

321 # and once with a DatasetType 

322 

323 # Keep track of any collections we add and do not clean up 

324 expected_collections = {run} 

325 

326 counter = 0 

327 ref = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run="put_run_1") 

328 args = tuple[DatasetRef] | tuple[str | DatasetType, DataCoordinate] 

329 for args in ((ref,), (datasetTypeName, dataId), (datasetType, dataId)): 

330 # Since we are using subTest we can get cascading failures 

331 # here with the first attempt failing and the others failing 

332 # immediately because the dataset already exists. Work around 

333 # this by using a distinct run collection each time 

334 counter += 1 

335 this_run = f"put_run_{counter}" 

336 butler.collections.register(this_run) 

337 expected_collections.update({this_run}) 

338 

339 with self.subTest(args=repr(args)): 

340 kwargs: dict[str, Any] = {} 

341 if not isinstance(args[0], DatasetRef): # type: ignore 

342 kwargs["run"] = this_run 

343 ref = butler.put(metric, *args, **kwargs) 

344 self.assertIsInstance(ref, DatasetRef) 

345 

346 # Test get of a ref. 

347 metricOut = butler.get(ref) 

348 self.assertEqual(metric, metricOut) 

349 # Test get 

350 metricOut = butler.get(ref.datasetType.name, dataId, collections=this_run) 

351 self.assertEqual(metric, metricOut) 

352 # Test get with a datasetRef 

353 metricOut = butler.get(ref) 

354 self.assertEqual(metric, metricOut) 

355 # Test getDeferred with dataId 

356 metricOut = butler.getDeferred(ref.datasetType.name, dataId, collections=this_run).get() 

357 self.assertEqual(metric, metricOut) 

358 # Test getDeferred with a ref 

359 metricOut = butler.getDeferred(ref).get() 

360 self.assertEqual(metric, metricOut) 

361 

362 # Check we can get components 

363 if storageClass.isComposite(): 

364 self.assertGetComponents( 

365 butler, ref, ("summary", "data", "output"), metric, collections=this_run 

366 ) 

367 

368 primary_uri, secondary_uris = butler.getURIs(ref) 

369 n_uris = len(secondary_uris) 

370 if primary_uri: 

371 n_uris += 1 

372 

373 # Can the artifacts themselves be retrieved? 

374 if not butler._datastore.isEphemeral: 

375 # Create a temporary directory to hold the retrieved 

376 # artifacts. 

377 with tempfile.TemporaryDirectory( 

378 prefix="butler-artifacts-", ignore_cleanup_errors=True 

379 ) as artifact_root: 

380 root_uri = ResourcePath(artifact_root, forceDirectory=True) 

381 

382 for preserve_path in (True, False): 

383 destination = root_uri.join(f"{preserve_path}_{counter}/") 

384 log = logging.getLogger("lsst.x") 

385 log.debug("Using destination %s for args %s", destination, args) 

386 # Use copy so that we can test that overwrite 

387 # protection works (using "auto" for File URIs 

388 # would use hard links and subsequent transfer 

389 # would work because it knows they are the same 

390 # file). 

391 transferred = butler.retrieveArtifacts( 

392 [ref], destination, preserve_path=preserve_path, transfer="copy" 

393 ) 

394 self.assertGreater(len(transferred), 0) 

395 artifacts = list(ResourcePath.findFileResources([destination])) 

396 # Filter out the index file. 

397 artifacts = [a for a in artifacts if a.basename() != ZipIndex.index_name] 

398 self.assertEqual(set(transferred), set(artifacts)) 

399 

400 for artifact in transferred: 

401 path_in_destination = artifact.relative_to(destination) 

402 self.assertIsNotNone(path_in_destination) 

403 assert path_in_destination is not None 

404 

405 # When path is not preserved there should not 

406 # be any path separators. 

407 num_seps = path_in_destination.count("/") 

408 if preserve_path: 

409 self.assertGreater(num_seps, 0) 

410 else: 

411 self.assertEqual(num_seps, 0) 

412 

413 self.assertEqual( 

414 len(artifacts), 

415 n_uris, 

416 "Comparing expected artifacts vs actual:" 

417 f" {artifacts} vs {primary_uri} and {secondary_uris}", 

418 ) 

419 

420 if preserve_path: 

421 # No need to run these twice 

422 with self.assertRaises(ValueError): 

423 butler.retrieveArtifacts([ref], destination, transfer="move") 

424 

425 with self.assertRaisesRegex( 

426 ValueError, "^Destination location must refer to a directory" 

427 ): 

428 butler.retrieveArtifacts( 

429 [ref], ResourcePath("/some/file.txt", forceDirectory=False) 

430 ) 

431 

432 with self.assertRaises(FileExistsError): 

433 butler.retrieveArtifacts([ref], destination) 

434 

435 transferred_again = butler.retrieveArtifacts( 

436 [ref], destination, preserve_path=preserve_path, overwrite=True 

437 ) 

438 self.assertEqual(set(transferred_again), set(transferred)) 

439 

440 # Now remove the dataset completely. 

441 butler.pruneDatasets([ref], purge=True, unstore=True) 

442 # Lookup with original args should still fail. 

443 kwargs = {"collections": this_run} 

444 if isinstance(args[0], DatasetRef): 

445 kwargs = {} # Prevent warning from being issued. 

446 self.assertFalse(butler.exists(*args, **kwargs)) 

447 # get() should still fail. 

448 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)): 

449 butler.get(ref) 

450 # Registry shouldn't be able to find it by dataset_id anymore. 

451 self.assertIsNone(butler.get_dataset(ref.id)) 

452 

453 # Do explicit registry removal since we know they are 

454 # empty 

455 butler.collections.x_remove(this_run) 

456 expected_collections.remove(this_run) 

457 

458 # Create DatasetRef for put using default run. 

459 refIn = DatasetRef(datasetType, dataId, id=uuid.UUID(int=1), run=butler.run) 

460 

461 # Check that getDeferred fails with standalone ref. 

462 with self.assertRaises(LookupError): 

463 butler.getDeferred(refIn) 

464 

465 # Put the dataset again, since the last thing we did was remove it 

466 # and we want to use the default collection. 

467 ref = butler.put(metric, refIn) 

468 

469 # Get with parameters 

470 stop = 4 

471 sliced = butler.get(ref, parameters={"slice": slice(stop)}) 

472 self.assertNotEqual(metric, sliced) 

473 self.assertEqual(metric.summary, sliced.summary) 

474 self.assertEqual(metric.output, sliced.output) 

475 assert metric.data is not None # for mypy 

476 self.assertEqual(metric.data[:stop], sliced.data) 

477 # getDeferred with parameters 

478 sliced = butler.getDeferred(ref, parameters={"slice": slice(stop)}).get() 

479 self.assertNotEqual(metric, sliced) 

480 self.assertEqual(metric.summary, sliced.summary) 

481 self.assertEqual(metric.output, sliced.output) 

482 self.assertEqual(metric.data[:stop], sliced.data) 

483 # getDeferred with deferred parameters 

484 sliced = butler.getDeferred(ref).get(parameters={"slice": slice(stop)}) 

485 self.assertNotEqual(metric, sliced) 

486 self.assertEqual(metric.summary, sliced.summary) 

487 self.assertEqual(metric.output, sliced.output) 

488 self.assertEqual(metric.data[:stop], sliced.data) 

489 

490 if storageClass.isComposite(): 

491 # Check that components can be retrieved 

492 metricOut = butler.get(ref.datasetType.name, dataId) 

493 compNameS = ref.datasetType.componentTypeName("summary") 

494 compNameD = ref.datasetType.componentTypeName("data") 

495 summary = butler.get(compNameS, dataId) 

496 self.assertEqual(summary, metric.summary) 

497 data = butler.get(compNameD, dataId) 

498 self.assertEqual(data, metric.data) 

499 

500 if "counter" in storageClass.derivedComponents: 

501 count = butler.get(ref.datasetType.componentTypeName("counter"), dataId) 

502 self.assertEqual(count, len(data)) 

503 

504 count = butler.get( 

505 ref.datasetType.componentTypeName("counter"), dataId, parameters={"slice": slice(stop)} 

506 ) 

507 self.assertEqual(count, stop) 

508 

509 compRef = butler.find_dataset(compNameS, dataId, collections=butler.collections.defaults) 

510 assert compRef is not None 

511 summary = butler.get(compRef) 

512 self.assertEqual(summary, metric.summary) 

513 

514 # Create a Dataset type that has the same name but is inconsistent. 

515 inconsistentDatasetType = DatasetType( 

516 datasetTypeName, datasetType.dimensions, self.storageClassFactory.getStorageClass("Config") 

517 ) 

518 

519 # Getting with a dataset type that does not match registry fails 

520 with self.assertRaisesRegex( 

521 ValueError, 

522 "(Supplied dataset type .* inconsistent with registry)" 

523 "|(The new storage class .* is not compatible with the existing storage class)", 

524 ): 

525 butler.get(inconsistentDatasetType, dataId) 

526 

527 # Combining a DatasetRef with a dataId should fail 

528 with self.assertRaisesRegex(ValueError, "DatasetRef given, cannot use dataId as well"): 

529 butler.get(ref, dataId) 

530 # Getting with an explicit ref should fail if the id doesn't match. 

531 with self.assertRaises((FileNotFoundError, DatasetNotFoundError)): 

532 butler.get(DatasetRef(ref.datasetType, ref.dataId, id=uuid.UUID(int=101), run=butler.run)) 

533 

534 # Getting a dataset with unknown parameters should fail 

535 with self.assertRaisesRegex(KeyError, "Parameter 'unsupported' not understood"): 

536 butler.get(ref, parameters={"unsupported": True}) 

537 

538 # Check we have a collection 

539 collections = set(butler.collections.query("*")) 

540 self.assertEqual(collections, expected_collections) 

541 

542 # Clean up to check that we can remove something that may have 

543 # already had a component removed 

544 butler.pruneDatasets([ref], unstore=True, purge=True) 

545 

546 # Add the same ref again, so we can check that duplicate put fails. 

547 ref = butler.put(metric, datasetType, dataId) 

548 

549 # Repeat put will fail. 

550 with self.assertRaisesRegex( 

551 ConflictingDefinitionError, "A database constraint failure was triggered" 

552 ): 

553 butler.put(metric, datasetType, dataId) 

554 

555 # Remove the datastore entry. 

556 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False) 

557 

558 # Put will still fail 

559 with self.assertRaisesRegex( 

560 ConflictingDefinitionError, "A database constraint failure was triggered" 

561 ): 

562 butler.put(metric, datasetType, dataId) 

563 

564 # Repeat the same sequence with resolved ref. 

565 butler.pruneDatasets([ref], unstore=True, purge=True) 

566 ref = butler.put(metric, refIn) 

567 

568 # Repeat put will fail. 

569 with self.assertRaisesRegex(ConflictingDefinitionError, "Datastore already contains dataset"): 

570 butler.put(metric, refIn) 

571 

572 # Remove the datastore entry. 

573 butler.pruneDatasets([ref], unstore=True, purge=False, disassociate=False) 

574 

575 # In case of resolved ref this write will succeed. 

576 ref = butler.put(metric, refIn) 

577 

578 # Leave the dataset in place since some downstream tests require 

579 # something to be present 

580 

581 return butler 

582 

583 def testDeferredCollectionPassing(self) -> None: 

584 # Construct a butler with no run or collection, but make it writeable. 

585 butler = self.create_empty_butler(writeable=True) 

586 # Create and register a DatasetType 

587 dimensions = butler.dimensions.conform(["instrument", "visit"]) 

588 datasetType = self.addDatasetType( 

589 "example", dimensions, self.storageClassFactory.getStorageClass("StructuredData"), butler.registry 

590 ) 

591 # Add needed Dimensions 

592 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

593 butler.registry.insertDimensionData( 

594 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

595 ) 

596 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101}) 

597 butler.registry.insertDimensionData( 

598 "visit", 

599 { 

600 "instrument": "DummyCamComp", 

601 "id": 423, 

602 "name": "fourtwentythree", 

603 "physical_filter": "d-r", 

604 "day_obs": 20250101, 

605 }, 

606 ) 

607 dataId = {"instrument": "DummyCamComp", "visit": 423} 

608 # Create dataset. 

609 metric = makeExampleMetrics() 

610 # Register a new run and put dataset. 

611 run = "deferred" 

612 self.assertTrue(butler.collections.register(run)) 

613 # Second time it will be allowed but indicate no-op 

614 self.assertFalse(butler.collections.register(run)) 

615 ref = butler.put(metric, datasetType, dataId, run=run) 

616 # Putting with no run should fail with TypeError. 

617 with self.assertRaises(CollectionError): 

618 butler.put(metric, datasetType, dataId) 

619 # Dataset should exist. 

620 self.assertTrue(butler.exists(datasetType, dataId, collections=[run])) 

621 # We should be able to get the dataset back, but with and without 

622 # a deferred dataset handle. 

623 self.assertEqual(metric, butler.get(datasetType, dataId, collections=[run])) 

624 self.assertEqual(metric, butler.getDeferred(datasetType, dataId, collections=[run]).get()) 

625 # Trying to find the dataset without any collection is an error. 

626 with self.assertRaises(NoDefaultCollectionError): 

627 butler.exists(datasetType, dataId) 

628 with self.assertRaises(CollectionError): 

629 butler.get(datasetType, dataId) 

630 # Associate the dataset with a different collection. 

631 butler.collections.register("tagged", type=CollectionType.TAGGED) 

632 butler.registry.associate("tagged", [ref]) 

633 # Deleting the dataset from the new collection should make it findable 

634 # in the original collection. 

635 butler.pruneDatasets([ref], tags=["tagged"]) 

636 self.assertTrue(butler.exists(datasetType, dataId, collections=[run])) 

637 

638 

639class ButlerTests(ButlerPutGetTests): 

640 """Tests for Butler.""" 

641 

642 useTempRoot = True 

643 validationCanFail: bool 

644 fullConfigKey: str | None 

645 registryStr: str | None 

646 datastoreName: list[str] | None 

647 datastoreStr: list[str] 

648 predictionSupported = True 

649 """Does getURIs support 'prediction mode'?""" 

650 

651 def setUp(self) -> None: 

652 """Create a new butler root for each test.""" 

653 self.root = makeTestTempDir(TESTDIR) 

654 Butler.makeRepo(self.root, config=Config(self.configFile)) 

655 self.tmpConfigFile = os.path.join(self.root, "butler.yaml") 

656 

657 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool: 

658 """Return True if two URIs refer to the same resource. 

659 

660 Subclasses may override to handle unique requirements. 

661 """ 

662 return uri1 == uri2 

663 

664 def testConstructor(self) -> None: 

665 """Independent test of constructor.""" 

666 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run) 

667 self.enterContext(butler) 

668 self.assertIsInstance(butler, Butler) 

669 

670 # Check that butler.yaml is added automatically. 

671 if self.tmpConfigFile.endswith(end := "/butler.yaml"): 

672 config_dir = self.tmpConfigFile[: -len(end)] 

673 butler = Butler.from_config(config_dir, run=self.default_run) 

674 self.enterContext(butler) 

675 self.assertIsInstance(butler, Butler) 

676 

677 # Even with a ResourcePath. 

678 butler = Butler.from_config(ResourcePath(config_dir, forceDirectory=True), run=self.default_run) 

679 self.enterContext(butler) 

680 self.assertIsInstance(butler, Butler) 

681 

682 collections = set(butler.collections.query("*")) 

683 self.assertEqual(collections, {self.default_run}) 

684 

685 # Check that some special characters can be included in run name. 

686 special_run = "u@b.c-A" 

687 butler_special = Butler.from_config(butler=butler, run=special_run) 

688 self.enterContext(butler_special) 

689 collections = set(butler_special.registry.queryCollections("*@*")) 

690 self.assertEqual(collections, {special_run}) 

691 

692 butler2 = Butler.from_config(butler=butler, collections=["other"]) 

693 self.enterContext(butler2) 

694 self.assertEqual(butler2.collections.defaults, ("other",)) 

695 self.assertIsNone(butler2.run) 

696 self.assertEqual(type(butler._datastore), type(butler2._datastore)) 

697 self.assertEqual(butler._datastore.config, butler2._datastore.config) 

698 

699 # Test that we can use an environment variable to find this 

700 # repository. 

701 butler_index = Config() 

702 butler_index["label"] = self.tmpConfigFile 

703 for suffix in (".yaml", ".json"): 

704 # Ensure that the content differs so that we know that 

705 # we aren't reusing the cache. 

706 bad_label = f"file://bucket/not_real{suffix}" 

707 butler_index["bad_label"] = bad_label 

708 with ResourcePath.temporary_uri(suffix=suffix) as temp_file: 

709 butler_index.dumpToUri(temp_file) 

710 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}): 

711 self.assertEqual(Butler.get_known_repos(), {"label", "bad_label"}) 

712 uri = Butler.get_repo_uri("bad_label") 

713 self.assertEqual(uri, ResourcePath(bad_label)) 

714 uri = Butler.get_repo_uri("label") 

715 butler = Butler.from_config(uri, writeable=False) 

716 self.assertIsInstance(butler, Butler) 

717 butler.close() 

718 butler = Butler.from_config("label", writeable=False) 

719 self.assertIsInstance(butler, Butler) 

720 butler.close() 

721 with self.assertRaisesRegex(FileNotFoundError, "aliases:.*bad_label"): 

722 Butler.from_config("not_there", writeable=False) 

723 with self.assertRaisesRegex(FileNotFoundError, "resolved from alias 'bad_label'"): 

724 Butler.from_config("bad_label") 

725 with self.assertRaises(FileNotFoundError): 

726 # Should ignore aliases. 

727 Butler.from_config(ResourcePath("label", forceAbsolute=False)) 

728 with self.assertRaises(KeyError) as cm: 

729 Butler.get_repo_uri("missing") 

730 self.assertEqual( 

731 Butler.get_repo_uri("missing", True), ResourcePath("missing", forceAbsolute=False) 

732 ) 

733 self.assertIn("not known to", str(cm.exception)) 

734 # Should report no failure. 

735 self.assertEqual(ButlerRepoIndex.get_failure_reason(), "") 

736 with ResourcePath.temporary_uri(suffix=suffix) as temp_file: 

737 # Now with empty configuration. 

738 butler_index = Config() 

739 butler_index.dumpToUri(temp_file) 

740 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}): 

741 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases)"): 

742 Butler.from_config("label") 

743 with ResourcePath.temporary_uri(suffix=suffix) as temp_file: 

744 # Now with bad contents. 

745 with open(temp_file.ospath, "w") as fh: 

746 print("'", file=fh) 

747 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(temp_file)}): 

748 with self.assertRaisesRegex(FileNotFoundError, "(no known aliases:.*could not be read)"): 

749 Butler.from_config("label") 

750 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": "file://not_found/x.yaml"}): 

751 with self.assertRaises(FileNotFoundError): 

752 Butler.get_repo_uri("label") 

753 self.assertEqual(Butler.get_known_repos(), set()) 

754 

755 with self.assertRaisesRegex(FileNotFoundError, "index file not found"): 

756 Butler.from_config("label") 

757 

758 # Check that we can create Butler when the alias file is not found. 

759 butler = Butler.from_config(self.tmpConfigFile, writeable=False) 

760 self.enterContext(butler) 

761 self.assertIsInstance(butler, Butler) 

762 with self.assertRaises(RuntimeError) as cm: 

763 # No environment variable set. 

764 Butler.get_repo_uri("label") 

765 self.assertEqual(Butler.get_repo_uri("label", True), ResourcePath("label", forceAbsolute=False)) 

766 self.assertIn("No repository index defined", str(cm.exception)) 

767 with self.assertRaisesRegex(FileNotFoundError, "no known aliases.*No repository index"): 

768 # No aliases registered. 

769 Butler.from_config("not_there") 

770 self.assertEqual(Butler.get_known_repos(), set()) 

771 

772 def testClose(self): 

773 butler = self.create_empty_butler(cleanup=False) 

774 is_direct_butler = isinstance(butler, DirectButler) 

775 if is_direct_butler: 

776 self.assertFalse(butler._closed) 

777 

778 with butler as butler_from_context_manager: 

779 self.assertIs(butler, butler_from_context_manager) 

780 if is_direct_butler: 

781 self.assertTrue(butler._closed) 

782 with self.assertRaisesRegex(RuntimeError, "has been closed"): 

783 butler.get_dataset_type("raw") 

784 

785 # Close may be called multiple times. 

786 butler.close() 

787 if is_direct_butler: 

788 self.assertTrue(butler._closed) 

789 

790 def testGarbageCollection(self): 

791 """Test that Butler does not have any circular references that prevent 

792 it from being garbage collected immediately when it goes out of scope. 

793 """ 

794 butler = self.create_empty_butler(cleanup=False) 

795 is_direct_butler = isinstance(butler, DirectButler) 

796 butler_ref = weakref.ref(butler) 

797 if is_direct_butler: 

798 registry_ref = weakref.ref(butler._registry) 

799 managers_ref = weakref.ref(butler._registry._managers) 

800 datastore_ref = weakref.ref(butler._datastore) 

801 db_ref = weakref.ref(butler._registry._db) 

802 engine_ref = weakref.ref(butler._registry._db._engine) 

803 

804 with warnings.catch_warnings(): 

805 # Hide warnings from unclosed database handles. 

806 warnings.simplefilter("ignore", ResourceWarning) 

807 del butler 

808 self.assertIsNone(butler_ref(), "Butler should have been garbage collected") 

809 if is_direct_butler: 

810 self.assertIsNone(registry_ref(), "SqlRegistry should have been garbage collected") 

811 self.assertIsNone(managers_ref(), "Registry managers should have been garbage collected") 

812 self.assertIsNone(datastore_ref(), "Datastore should have been garbage collected") 

813 self.assertIsNone(db_ref(), "Database should have been garbage collected") 

814 # SQLAlchemy has internal reference cycles, so the Engine instance 

815 # is not cleaned up promptly even if we release our reference to 

816 # it. Explicitly clean it up here to avoid file handles leaking. 

817 if is_direct_butler: 

818 engine = engine_ref() 

819 if engine is not None: 

820 engine.dispose() 

821 

822 def testDafButlerRepositories(self): 

823 with unittest.mock.patch.dict( 

824 os.environ, 

825 {"DAF_BUTLER_REPOSITORIES": "label: 'https://someuri.com'\notherLabel: 'https://otheruri.com'\n"}, 

826 ): 

827 self.assertEqual(str(Butler.get_repo_uri("label")), "https://someuri.com") 

828 

829 with unittest.mock.patch.dict( 

830 os.environ, 

831 { 

832 "DAF_BUTLER_REPOSITORIES": "label: https://someuri.com", 

833 "DAF_BUTLER_REPOSITORY_INDEX": "https://someuri.com", 

834 }, 

835 ): 

836 with self.assertRaisesRegex(RuntimeError, "Only one of the environment variables"): 

837 Butler.get_repo_uri("label") 

838 

839 with unittest.mock.patch.dict( 

840 os.environ, 

841 {"DAF_BUTLER_REPOSITORIES": "invalid"}, 

842 ): 

843 with self.assertRaisesRegex(ValueError, "Repository index not in expected format"): 

844 Butler.get_repo_uri("label") 

845 

846 def testBasicPutGet(self) -> None: 

847 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

848 self.runPutGetTest(storageClass, "test_metric") 

849 

850 def testCompositePutGetConcrete(self) -> None: 

851 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadCompNoDisassembly") 

852 butler = self.runPutGetTest(storageClass, "test_metric") 

853 

854 # Should *not* be disassembled 

855 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run)) 

856 self.assertEqual(len(datasets), 1) 

857 uri, components = butler.getURIs(datasets[0]) 

858 self.assertIsInstance(uri, ResourcePath) 

859 self.assertFalse(components) 

860 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}") 

861 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}") 

862 

863 # Predicted dataset 

864 if self.predictionSupported: 

865 dataId = {"instrument": "DummyCamComp", "visit": 424} 

866 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True) 

867 self.assertFalse(components) 

868 self.assertIsInstance(uri, ResourcePath) 

869 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}") 

870 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}") 

871 # Repeat with a DatasetRef to test that code path. 

872 ref = DatasetRef( 

873 datasets[0].datasetType, 

874 dataId=DataCoordinate.standardize(dataId, universe=butler.dimensions), 

875 run=self.default_run, 

876 ) 

877 uri2, components2 = butler.getURIs(ref, predict=True) 

878 self.assertFalse(components2) 

879 self.assertEqual(uri, uri2) 

880 

881 def testCompositePutGetVirtual(self) -> None: 

882 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp") 

883 butler = self.runPutGetTest(storageClass, "test_metric_comp") 

884 

885 # Should be disassembled 

886 datasets = list(butler.registry.queryDatasets(..., collections=self.default_run)) 

887 self.assertEqual(len(datasets), 1) 

888 uri, components = butler.getURIs(datasets[0]) 

889 

890 if butler._datastore.isEphemeral: 

891 # Never disassemble in-memory datastore 

892 self.assertIsInstance(uri, ResourcePath) 

893 self.assertFalse(components) 

894 self.assertEqual(uri.fragment, "", f"Checking absence of fragment in {uri}") 

895 self.assertIn("423", str(uri), f"Checking visit is in URI {uri}") 

896 else: 

897 self.assertIsNone(uri) 

898 self.assertEqual(set(components), set(storageClass.components)) 

899 for compuri in components.values(): 

900 self.assertIsInstance(compuri, ResourcePath) 

901 self.assertIn("423", str(compuri), f"Checking visit is in URI {compuri}") 

902 self.assertEqual(compuri.fragment, "", f"Checking absence of fragment in {compuri}") 

903 

904 if self.predictionSupported: 

905 # Predicted dataset 

906 dataId = {"instrument": "DummyCamComp", "visit": 424} 

907 uri, components = butler.getURIs(datasets[0].datasetType, dataId=dataId, predict=True) 

908 

909 if butler._datastore.isEphemeral: 

910 # Never disassembled 

911 self.assertIsInstance(uri, ResourcePath) 

912 self.assertFalse(components) 

913 self.assertIn("424", str(uri), f"Checking visit is in URI {uri}") 

914 self.assertEqual(uri.fragment, "predicted", f"Checking for fragment in {uri}") 

915 else: 

916 self.assertIsNone(uri) 

917 self.assertEqual(set(components), set(storageClass.components)) 

918 for compuri in components.values(): 

919 self.assertIsInstance(compuri, ResourcePath) 

920 self.assertIn("424", str(compuri), f"Checking visit is in URI {compuri}") 

921 self.assertEqual(compuri.fragment, "predicted", f"Checking for fragment in {compuri}") 

922 

923 def testStorageClassOverrideGet(self) -> None: 

924 """Test storage class conversion on get with override.""" 

925 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

926 datasetTypeName = "anything" 

927 run = self.default_run 

928 

929 butler, datasetType = self.create_butler(run, storageClass, datasetTypeName) 

930 

931 # Create and store a dataset. 

932 metric = makeExampleMetrics() 

933 dataId = {"instrument": "DummyCamComp", "visit": 423} 

934 

935 ref = butler.put(metric, datasetType, dataId) 

936 

937 # Return native type. 

938 retrieved = butler.get(ref) 

939 self.assertEqual(retrieved, metric) 

940 

941 # Specify an override. 

942 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion") 

943 model = butler.get(ref, storageClass=new_sc) 

944 self.assertNotEqual(type(model), type(retrieved)) 

945 self.assertIs(type(model), new_sc.pytype) 

946 self.assertEqual(retrieved, model) 

947 

948 # Defer but override later. 

949 deferred = butler.getDeferred(ref) 

950 model = deferred.get(storageClass=new_sc) 

951 self.assertIs(type(model), new_sc.pytype) 

952 self.assertEqual(retrieved, model) 

953 

954 # Defer but override up front. 

955 deferred = butler.getDeferred(ref, storageClass=new_sc) 

956 model = deferred.get() 

957 self.assertIs(type(model), new_sc.pytype) 

958 self.assertEqual(retrieved, model) 

959 

960 # Retrieve a component. Should be a tuple. 

961 data = butler.get("anything.data", dataId, storageClass="StructuredDataDataTestTuple") 

962 self.assertIs(type(data), tuple) 

963 self.assertEqual(data, tuple(retrieved.data)) 

964 

965 # Parameter on the write storage class should work regardless 

966 # of read storage class. 

967 data = butler.get( 

968 "anything.data", 

969 dataId, 

970 storageClass="StructuredDataDataTestTuple", 

971 parameters={"slice": slice(2, 4)}, 

972 ) 

973 self.assertEqual(len(data), 2) 

974 

975 # Try a parameter that is known to the read storage class but not 

976 # the write storage class. 

977 with self.assertRaises(KeyError): 

978 butler.get( 

979 "anything.data", 

980 dataId, 

981 storageClass="StructuredDataDataTestTuple", 

982 parameters={"xslice": slice(2, 4)}, 

983 ) 

984 

985 def testPytypePutCoercion(self) -> None: 

986 """Test python type coercion on Butler.get and put.""" 

987 # Store some data with the normal example storage class. 

988 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

989 datasetTypeName = "test_metric" 

990 butler, _ = self.create_butler(self.default_run, storageClass, datasetTypeName) 

991 

992 dataId = {"instrument": "DummyCamComp", "visit": 423} 

993 

994 # Put a dict and this should coerce to a MetricsExample 

995 test_dict = {"summary": {"a": 1}, "output": {"b": 2}} 

996 metric_ref = butler.put(test_dict, datasetTypeName, dataId=dataId, visit=424) 

997 test_metric = butler.get(metric_ref) 

998 self.assertEqual(get_full_type_name(test_metric), "lsst.daf.butler.tests.MetricsExample") 

999 self.assertEqual(test_metric.summary, test_dict["summary"]) 

1000 self.assertEqual(test_metric.output, test_dict["output"]) 

1001 

1002 # Check that the put still works if a DatasetType is given with 

1003 # a definition matching this python type. 

1004 registry_type = butler.get_dataset_type(datasetTypeName) 

1005 this_type = DatasetType(datasetTypeName, registry_type.dimensions, "StructuredDataDictJson") 

1006 metric2_ref = butler.put(test_dict, this_type, dataId=dataId, visit=425) 

1007 self.assertEqual(metric2_ref.datasetType, registry_type) 

1008 

1009 # The get will return the type expected by registry. 

1010 test_metric2 = butler.get(metric2_ref) 

1011 self.assertEqual(get_full_type_name(test_metric2), "lsst.daf.butler.tests.MetricsExample") 

1012 

1013 # Make a new DatasetRef with the compatible but different DatasetType. 

1014 # This should now return a dict. 

1015 new_ref = DatasetRef(this_type, metric2_ref.dataId, id=metric2_ref.id, run=metric2_ref.run) 

1016 test_dict2 = butler.get(new_ref) 

1017 self.assertEqual(get_full_type_name(test_dict2), "dict") 

1018 

1019 # Get it again with the wrong dataset type definition using get() 

1020 # rather than get(). This should be consistent with get() 

1021 # behavior and return the type of the DatasetType. 

1022 test_dict3 = butler.get(this_type, dataId=dataId, visit=425) 

1023 self.assertEqual(get_full_type_name(test_dict3), "dict") 

1024 

1025 def test_ingest_zip(self) -> None: 

1026 """Create butler, export data, delete data, import from Zip.""" 

1027 butler, dataset_type = self.create_butler( 

1028 run=self.default_run, storageClass="StructuredData", datasetTypeName="metrics" 

1029 ) 

1030 

1031 metric = makeExampleMetrics() 

1032 refs = [] 

1033 for visit in (423, 424, 425): 

1034 ref = butler.put(metric, dataset_type, instrument="DummyCamComp", visit=visit) 

1035 refs.append(ref) 

1036 

1037 # Retrieve a Zip file. 

1038 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

1039 zip = butler.retrieve_artifacts_zip(refs, destination=tmpdir) 

1040 

1041 # Ingest will fail. 

1042 with self.assertRaises(ConflictingDefinitionError): 

1043 butler.ingest_zip(zip) 

1044 

1045 # Clear out the collection. 

1046 butler.removeRuns([self.default_run]) 

1047 self.assertFalse(butler.exists(refs[0])) 

1048 

1049 butler.ingest_zip(zip, transfer="copy") 

1050 self.assertGreater(butler._metrics.time_in_ingest, 0.0) 

1051 self.assertEqual(butler._metrics.n_ingest, len(refs)) 

1052 

1053 # Check that it fails if we try it again. 

1054 with self.assertRaises(ConflictingDefinitionError): 

1055 butler.ingest_zip(zip, transfer="copy") 

1056 

1057 # This will be a no-op. 

1058 butler.ingest_zip(zip, transfer="copy", skip_existing=True) 

1059 

1060 # Create an entirely new local file butler in this temp directory. 

1061 new_butler_cfg = Butler.makeRepo(tmpdir) 

1062 new_butler = Butler.from_config(new_butler_cfg, writeable=True) 

1063 self.enterContext(new_butler) 

1064 

1065 # This will fail since dimensions records are missing. 

1066 with self.assertRaises(ConflictingDefinitionError): 

1067 new_butler.ingest_zip(zip, transfer="copy") 

1068 

1069 # Dry run should work. 

1070 new_butler.ingest_zip(zip, transfer="copy", dry_run=True) 

1071 

1072 new_butler.ingest_zip(zip, transfer="copy", transfer_dimensions=True) 

1073 self.assertTrue(butler.exists(refs[0])) 

1074 

1075 # Check that the refs can be read again. 

1076 _ = [butler.get(ref) for ref in refs] 

1077 

1078 uri = butler.getURI(refs[2]) 

1079 self.assertTrue(uri.exists()) 

1080 

1081 # Delete one dataset. The Zip file should still exist and allow 

1082 # remaining refs to be read. 

1083 butler.pruneDatasets([refs[0]], purge=True, unstore=True) 

1084 self.assertTrue(uri.exists()) 

1085 

1086 metric2 = butler.get(refs[1]) 

1087 self.assertEqual(metric2, metric, msg=f"{metric2} != {metric}") 

1088 

1089 butler.removeRuns([self.default_run]) 

1090 self.assertFalse(uri.exists()) 

1091 self.assertFalse(butler.exists(refs[-1])) 

1092 

1093 with self.assertRaises(ValueError): 

1094 butler.retrieve_artifacts_zip([], destination=".") 

1095 

1096 def testIngest(self) -> None: 

1097 butler = self.create_empty_butler(run=self.default_run) 

1098 

1099 # Create and register a DatasetType 

1100 dimensions = butler.dimensions.conform(["instrument", "visit", "detector"]) 

1101 

1102 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDictYaml") 

1103 datasetTypeName = "metric" 

1104 

1105 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry) 

1106 

1107 # Add needed Dimensions 

1108 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

1109 butler.registry.insertDimensionData( 

1110 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

1111 ) 

1112 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101}) 

1113 for detector in (1, 2): 

1114 butler.registry.insertDimensionData( 

1115 "detector", {"instrument": "DummyCamComp", "id": detector, "full_name": f"detector{detector}"} 

1116 ) 

1117 

1118 butler.registry.insertDimensionData( 

1119 "visit", 

1120 { 

1121 "instrument": "DummyCamComp", 

1122 "id": 423, 

1123 "name": "fourtwentythree", 

1124 "physical_filter": "d-r", 

1125 "day_obs": 20250101, 

1126 }, 

1127 { 

1128 "instrument": "DummyCamComp", 

1129 "id": 424, 

1130 "name": "fourtwentyfour", 

1131 "physical_filter": "d-r", 

1132 "day_obs": 20250101, 

1133 }, 

1134 ) 

1135 

1136 formatter = doImportType("lsst.daf.butler.formatters.yaml.YamlFormatter") 

1137 dataRoot = os.path.join(TESTDIR, "data", "basic") 

1138 datasets = [] 

1139 # Test one DatasetRef with a run that exists, and the other with a run 

1140 # that doesn't exist, to verify that run collections are created when 

1141 # required. 

1142 runs = {1: self.default_run, 2: "a/new/run"} 

1143 for detector in (1, 2): 

1144 detector_name = f"detector_{detector}" 

1145 metricFile = os.path.join(dataRoot, f"{detector_name}.yaml") 

1146 dataId = butler.registry.expandDataId( 

1147 {"instrument": "DummyCamComp", "visit": 423, "detector": detector} 

1148 ) 

1149 # Create a DatasetRef for ingest 

1150 refIn = DatasetRef(datasetType, dataId, run=runs[detector]) 

1151 

1152 datasets.append(FileDataset(path=metricFile, refs=[refIn], formatter=formatter)) 

1153 

1154 butler.ingest(*datasets, transfer="copy") 

1155 

1156 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 423} 

1157 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 423} 

1158 

1159 metrics1 = butler.get(datasetTypeName, dataId1) 

1160 metrics2 = butler.get(datasetTypeName, dataId2, collections="a/new/run") 

1161 self.assertNotEqual(metrics1, metrics2) 

1162 

1163 # Compare URIs 

1164 uri1 = butler.getURI(datasetTypeName, dataId1) 

1165 uri2 = butler.getURI(datasetTypeName, dataId2, collections="a/new/run") 

1166 self.assertFalse(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}") 

1167 

1168 # Re-ingesting the same datasets raises an error with 

1169 # skip_existing=False. 

1170 with self.assertRaises(ConflictingDefinitionError): 

1171 butler.ingest(*datasets, transfer="copy") 

1172 # skip_existing=True makes it a no-op to re-ingest the same datasets. 

1173 butler.ingest(*datasets, transfer="copy", skip_existing=True) 

1174 

1175 # Now do a multi-dataset but single file ingest 

1176 metricFile = os.path.join(dataRoot, "detectors.yaml") 

1177 refs = [] 

1178 for detector in (1, 2): 

1179 detector_name = f"detector_{detector}" 

1180 dataId = butler.registry.expandDataId( 

1181 {"instrument": "DummyCamComp", "visit": 424, "detector": detector} 

1182 ) 

1183 # Create a DatasetRef for ingest 

1184 refs.append(DatasetRef(datasetType, dataId, run=self.default_run)) 

1185 

1186 # Test "move" transfer to ensure that the files themselves 

1187 # have disappeared following ingest. 

1188 with ResourcePath.temporary_uri(suffix=".yaml") as tempFile: 

1189 tempFile.transfer_from(ResourcePath(metricFile), transfer="copy") 

1190 

1191 datasets = [] 

1192 datasets.append(FileDataset(path=tempFile, refs=refs, formatter=MultiDetectorFormatter)) 

1193 

1194 # For first ingest use copy. 

1195 butler.ingest(*datasets, transfer="copy", record_validation_info=False) 

1196 

1197 # Now try to ingest again in "execution butler" mode where 

1198 # the registry entries exist but the datastore does not have 

1199 # the files. We also need to strip the dimension records to ensure 

1200 # that they will be re-added by the ingest. 

1201 ref = datasets[0].refs[0] 

1202 datasets[0].refs = [ 

1203 cast( 

1204 DatasetRef, 

1205 butler.find_dataset(ref.datasetType, data_id=ref.dataId, collections=ref.run), 

1206 ) 

1207 for ref in datasets[0].refs 

1208 ] 

1209 all_refs = [] 

1210 for dataset in datasets: 

1211 refs = [] 

1212 for ref in dataset.refs: 

1213 # Create a dict from the dataId to drop the records. 

1214 new_data_id = dict(ref.dataId.required) 

1215 new_ref = butler.find_dataset(ref.datasetType, new_data_id, collections=ref.run) 

1216 assert new_ref is not None 

1217 self.assertFalse(new_ref.dataId.hasRecords()) 

1218 refs.append(new_ref) 

1219 dataset.refs = refs 

1220 all_refs.extend(dataset.refs) 

1221 butler.pruneDatasets(all_refs, disassociate=False, unstore=True, purge=False) 

1222 

1223 # Use move mode to test that the file is deleted. Also 

1224 # disable recording of file size. 

1225 butler.ingest(*datasets, transfer="move", record_validation_info=False) 

1226 

1227 # Check that every ref now has records. 

1228 for dataset in datasets: 

1229 for ref in dataset.refs: 

1230 self.assertTrue(ref.dataId.hasRecords()) 

1231 

1232 # Ensure that the file has disappeared. 

1233 self.assertFalse(tempFile.exists()) 

1234 

1235 # Check that the datastore recorded no file size. 

1236 # Not all datastores can support this. 

1237 try: 

1238 infos = butler._datastore.getStoredItemsInfo(datasets[0].refs[0]) # type: ignore[attr-defined] 

1239 self.assertEqual(infos[0].file_size, -1) 

1240 except AttributeError: 

1241 pass 

1242 

1243 dataId1 = {"instrument": "DummyCamComp", "detector": 1, "visit": 424} 

1244 dataId2 = {"instrument": "DummyCamComp", "detector": 2, "visit": 424} 

1245 

1246 multi1 = butler.get(datasetTypeName, dataId1) 

1247 multi2 = butler.get(datasetTypeName, dataId2) 

1248 

1249 self.assertEqual(multi1, metrics1) 

1250 self.assertEqual(multi2, metrics2) 

1251 

1252 # Compare URIs 

1253 uri1 = butler.getURI(datasetTypeName, dataId1) 

1254 uri2 = butler.getURI(datasetTypeName, dataId2) 

1255 self.assertTrue(self.are_uris_equivalent(uri1, uri2), f"Cf. {uri1} with {uri2}") 

1256 

1257 # Test that removing one does not break the second 

1258 # This line will issue a warning log message for a ChainedDatastore 

1259 # that uses an InMemoryDatastore since in-memory can not ingest 

1260 # files. 

1261 butler.pruneDatasets([datasets[0].refs[0]], unstore=True, disassociate=False) 

1262 self.assertFalse(butler.exists(datasetTypeName, dataId1)) 

1263 self.assertTrue(butler.exists(datasetTypeName, dataId2)) 

1264 multi2b = butler.get(datasetTypeName, dataId2) 

1265 self.assertEqual(multi2, multi2b) 

1266 

1267 # Ensure we can ingest 0 datasets 

1268 datasets = [] 

1269 butler.ingest(*datasets) 

1270 

1271 def testPickle(self) -> None: 

1272 """Test pickle support.""" 

1273 butler = self.create_empty_butler(run=self.default_run) 

1274 assert isinstance(butler, DirectButler), "Expect DirectButler in configuration" 

1275 butlerOut = pickle.loads(pickle.dumps(butler)) 

1276 self.enterContext(butlerOut) 

1277 self.assertIsInstance(butlerOut, Butler) 

1278 self.assertEqual(butlerOut._config, butler._config) 

1279 self.assertEqual(list(butlerOut.collections.defaults), list(butler.collections.defaults)) 

1280 self.assertEqual(butlerOut.run, butler.run) 

1281 

1282 def testGetDatasetTypes(self) -> None: 

1283 butler = self.create_empty_butler(run=self.default_run) 

1284 dimensions = butler.dimensions.conform(["instrument", "visit", "physical_filter"]) 

1285 dimensionEntries: list[tuple[str, list[Mapping[str, Any]]]] = [ 

1286 ( 

1287 "instrument", 

1288 [ 

1289 {"instrument": "DummyCam"}, 

1290 {"instrument": "DummyHSC"}, 

1291 {"instrument": "DummyCamComp"}, 

1292 ], 

1293 ), 

1294 ("physical_filter", [{"instrument": "DummyCam", "name": "d-r", "band": "R"}]), 

1295 ("day_obs", [{"instrument": "DummyCam", "id": 20250101}]), 

1296 ( 

1297 "visit", 

1298 [ 

1299 { 

1300 "instrument": "DummyCam", 

1301 "id": 42, 

1302 "name": "fortytwo", 

1303 "physical_filter": "d-r", 

1304 "day_obs": 20250101, 

1305 } 

1306 ], 

1307 ), 

1308 ] 

1309 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1310 # Add needed Dimensions 

1311 for element, data in dimensionEntries: 

1312 butler.registry.insertDimensionData(element, *data) 

1313 

1314 # When a DatasetType is added to the registry entries are not created 

1315 # for components but querying them can return the components. 

1316 datasetTypeNames = {"metric", "metric2", "metric4", "metric33", "pvi", "paramtest"} 

1317 components = set() 

1318 for datasetTypeName in datasetTypeNames: 

1319 # Create and register a DatasetType 

1320 self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry) 

1321 

1322 for componentName in storageClass.components: 

1323 components.add(DatasetType.nameWithComponent(datasetTypeName, componentName)) 

1324 

1325 fromRegistry: set[DatasetType] = set() 

1326 for parent_dataset_type in butler.registry.queryDatasetTypes(): 

1327 fromRegistry.add(parent_dataset_type) 

1328 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes()) 

1329 self.assertEqual({d.name for d in fromRegistry}, datasetTypeNames | components) 

1330 

1331 # Query with wildcard. 

1332 dataset_types = butler.registry.queryDatasetTypes("metric*") 

1333 self.assertEqual(len(dataset_types), 4, f"Got: {dataset_types}") 

1334 # but not regex. 

1335 with self.assertRaises(DatasetTypeExpressionError): 

1336 butler.registry.queryDatasetTypes(["pvi", re.compile("metric.*")]) 

1337 

1338 # Now that we have some dataset types registered, validate them 

1339 butler.validateConfiguration( 

1340 ignore=[ 

1341 "test_metric_comp", 

1342 "metric3", 

1343 "metric5", 

1344 "calexp", 

1345 "DummySC", 

1346 "datasetType.component", 

1347 "random_data", 

1348 "random_data_2", 

1349 ] 

1350 ) 

1351 

1352 # Add a new datasetType that will fail template validation 

1353 self.addDatasetType("test_metric_comp", dimensions, storageClass, butler.registry) 

1354 if self.validationCanFail: 

1355 with self.assertRaises(ValidationError): 

1356 butler.validateConfiguration() 

1357 

1358 # Rerun validation but with a subset of dataset type names 

1359 butler.validateConfiguration(datasetTypeNames=["metric4"]) 

1360 

1361 # Rerun validation but ignore the bad datasetType 

1362 butler.validateConfiguration( 

1363 ignore=[ 

1364 "test_metric_comp", 

1365 "metric3", 

1366 "metric5", 

1367 "calexp", 

1368 "DummySC", 

1369 "datasetType.component", 

1370 "random_data", 

1371 "random_data_2", 

1372 ] 

1373 ) 

1374 

1375 def testTransaction(self) -> None: 

1376 butler = self.create_empty_butler(run=self.default_run) 

1377 datasetTypeName = "test_metric" 

1378 dimensions = butler.dimensions.conform(["instrument", "visit"]) 

1379 dimensionEntries: tuple[tuple[str, Mapping[str, Any]], ...] = ( 

1380 ("instrument", {"instrument": "DummyCam"}), 

1381 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}), 

1382 ("day_obs", {"instrument": "DummyCam", "id": 20250101}), 

1383 ( 

1384 "visit", 

1385 { 

1386 "instrument": "DummyCam", 

1387 "id": 42, 

1388 "name": "fortytwo", 

1389 "physical_filter": "d-r", 

1390 "day_obs": 20250101, 

1391 }, 

1392 ), 

1393 ) 

1394 storageClass = self.storageClassFactory.getStorageClass("StructuredData") 

1395 metric = makeExampleMetrics() 

1396 dataId = {"instrument": "DummyCam", "visit": 42} 

1397 # Create and register a DatasetType 

1398 datasetType = self.addDatasetType(datasetTypeName, dimensions, storageClass, butler.registry) 

1399 with self.assertRaises(TransactionTestError): 

1400 with butler.transaction(): 

1401 # Add needed Dimensions 

1402 for args in dimensionEntries: 

1403 butler.registry.insertDimensionData(*args) 

1404 # Store a dataset 

1405 ref = butler.put(metric, datasetTypeName, dataId) 

1406 self.assertIsInstance(ref, DatasetRef) 

1407 # Test get of a ref. 

1408 metricOut = butler.get(ref) 

1409 self.assertEqual(metric, metricOut) 

1410 # Test get 

1411 metricOut = butler.get(datasetTypeName, dataId) 

1412 self.assertEqual(metric, metricOut) 

1413 # Check we can get components 

1414 self.assertGetComponents(butler, ref, ("summary", "data", "output"), metric) 

1415 raise TransactionTestError("This should roll back the entire transaction") 

1416 with self.assertRaises(DataIdValueError, msg=f"Check can't expand DataId {dataId}"): 

1417 butler.registry.expandDataId(dataId) 

1418 # Should raise LookupError for missing data ID value 

1419 with self.assertRaises(LookupError, msg=f"Check can't get by {datasetTypeName} and {dataId}"): 

1420 butler.get(datasetTypeName, dataId) 

1421 # Also check explicitly if Dataset entry is missing 

1422 self.assertIsNone(butler.find_dataset(datasetType, dataId, collections=butler.collections.defaults)) 

1423 # Direct retrieval should not find the file in the Datastore 

1424 with self.assertRaises(FileNotFoundError, msg=f"Check {ref} can't be retrieved directly"): 

1425 butler.get(ref) 

1426 

1427 def testMakeRepo(self) -> None: 

1428 """Test that we can write butler configuration to a new repository via 

1429 the Butler.makeRepo interface and then instantiate a butler from the 

1430 repo root. 

1431 """ 

1432 # Do not run the test if we know this datastore configuration does 

1433 # not support a file system root 

1434 if self.fullConfigKey is None: 

1435 return 

1436 

1437 # create two separate directories 

1438 root1 = tempfile.mkdtemp(dir=self.root) 

1439 root2 = tempfile.mkdtemp(dir=self.root) 

1440 

1441 butlerConfig = Butler.makeRepo(root1, config=Config(self.configFile)) 

1442 limited = Config(self.configFile) 

1443 butler1 = Butler.from_config(butlerConfig) 

1444 self.enterContext(butler1) 

1445 assert isinstance(butler1, DirectButler), "Expect DirectButler in configuration" 

1446 butlerConfig = Butler.makeRepo(root2, standalone=True, config=Config(self.configFile)) 

1447 full = Config(self.tmpConfigFile) 

1448 butler2 = Butler.from_config(butlerConfig) 

1449 self.enterContext(butler2) 

1450 assert isinstance(butler2, DirectButler), "Expect DirectButler in configuration" 

1451 # Butlers should have the same configuration regardless of whether 

1452 # defaults were expanded. 

1453 self.assertEqual(butler1._config, butler2._config) 

1454 # Config files loaded directly should not be the same. 

1455 self.assertNotEqual(limited, full) 

1456 # Make sure "limited" doesn't have a few keys we know it should be 

1457 # inheriting from defaults. 

1458 self.assertIn(self.fullConfigKey, full) 

1459 self.assertNotIn(self.fullConfigKey, limited) 

1460 

1461 # Collections don't appear until something is put in them 

1462 collections1 = set(butler1.registry.queryCollections()) 

1463 self.assertEqual(collections1, set()) 

1464 self.assertEqual(set(butler2.registry.queryCollections()), collections1) 

1465 

1466 # Check that a config with no associated file name will not 

1467 # work properly with relocatable Butler repo 

1468 butlerConfig.configFile = None 

1469 with self.assertRaises(ValueError): 

1470 Butler.from_config(butlerConfig) 

1471 

1472 with self.assertRaises(FileExistsError): 

1473 Butler.makeRepo(self.root, standalone=True, config=Config(self.configFile), overwrite=False) 

1474 

1475 def testStringification(self) -> None: 

1476 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run) 

1477 self.enterContext(butler) 

1478 butlerStr = str(butler) 

1479 

1480 if self.datastoreStr is not None: 

1481 for testStr in self.datastoreStr: 

1482 self.assertIn(testStr, butlerStr) 

1483 if self.registryStr is not None: 

1484 self.assertIn(self.registryStr, butlerStr) 

1485 

1486 datastoreName = butler._datastore.name 

1487 if self.datastoreName is not None: 

1488 for testStr in self.datastoreName: 

1489 self.assertIn(testStr, datastoreName) 

1490 

1491 def testButlerRewriteDataId(self) -> None: 

1492 """Test that dataIds can be rewritten based on dimension records.""" 

1493 butler = self.create_empty_butler(run=self.default_run) 

1494 

1495 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

1496 datasetTypeName = "random_data" 

1497 

1498 # Create dimension records. 

1499 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

1500 butler.registry.insertDimensionData( 

1501 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

1502 ) 

1503 butler.registry.insertDimensionData( 

1504 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"} 

1505 ) 

1506 

1507 dimensions = butler.dimensions.conform(["instrument", "exposure"]) 

1508 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

1509 butler.registry.registerDatasetType(datasetType) 

1510 

1511 n_exposures = 5 

1512 dayobs = 20210530 

1513 

1514 # Create records for multiple day_obs but same seq_num to test that 

1515 # we are constraining gets properly when day_obs/seq_num is used 

1516 # for an exposure. Second day is year in future but is not used. 

1517 for day_obs in (dayobs, dayobs + 1_00_00): 

1518 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": day_obs}) 

1519 

1520 for i in range(n_exposures): 

1521 group_name = f"group_{day_obs}_{i}" 

1522 butler.registry.insertDimensionData( 

1523 "group", {"instrument": "DummyCamComp", "name": group_name} 

1524 ) 

1525 butler.registry.insertDimensionData( 

1526 "exposure", 

1527 { 

1528 "instrument": "DummyCamComp", 

1529 "id": day_obs + i, 

1530 "obs_id": f"exp_{day_obs}_{i}", 

1531 "seq_num": i, 

1532 "day_obs": day_obs, 

1533 "physical_filter": "d-r", 

1534 "group": group_name, 

1535 }, 

1536 ) 

1537 

1538 # Write some data. 

1539 for i in range(n_exposures): 

1540 metric = {"something": i, "other": "metric", "list": [2 * x for x in range(i)]} 

1541 

1542 # Use the seq_num for the put to test rewriting. 

1543 dataId = {"seq_num": i, "day_obs": dayobs, "instrument": "DummyCamComp", "physical_filter": "d-r"} 

1544 ref = butler.put(metric, datasetTypeName, dataId=dataId) 

1545 

1546 # Check that the exposure is correct in the dataId 

1547 self.assertEqual(ref.dataId["exposure"], dayobs + i) 

1548 

1549 # and check that we can get the dataset back with the same dataId 

1550 new_metric = butler.get(datasetTypeName, dataId=dataId) 

1551 self.assertEqual(new_metric, metric) 

1552 

1553 # Check that we can find the datasets using the day_obs or the 

1554 # exposure.day_obs. 

1555 datasets_1 = list( 

1556 butler.registry.queryDatasets( 

1557 datasetType, 

1558 collections=self.default_run, 

1559 where="day_obs = :dayObs AND instrument = :instr", 

1560 bind={"dayObs": dayobs, "instr": "DummyCamComp"}, 

1561 ) 

1562 ) 

1563 datasets_2 = list( 

1564 butler.registry.queryDatasets( 

1565 datasetType, 

1566 collections=self.default_run, 

1567 where="exposure.day_obs = :dayObs AND instrument = :instr", 

1568 bind={"dayObs": dayobs, "instr": "DummyCamComp"}, 

1569 ) 

1570 ) 

1571 self.assertEqual(datasets_1, datasets_2) 

1572 

1573 def testGetDatasetCollectionCaching(self): 

1574 # Prior to DM-41117, there was a bug where get_dataset would throw 

1575 # MissingCollectionError if you tried to fetch a dataset that was added 

1576 # after the collection cache was last updated. 

1577 reader_butler, datasetType = self.create_butler(self.default_run, "int", "datasettypename") 

1578 writer_butler = self.create_empty_butler(writeable=True, run="new_run") 

1579 dataId = {"instrument": "DummyCamComp", "visit": 423} 

1580 put_ref = writer_butler.put(123, datasetType, dataId) 

1581 get_ref = reader_butler.get_dataset(put_ref.id) 

1582 self.assertEqual(get_ref.id, put_ref.id) 

1583 # Also works when looking up via a hexadecimal string instead of a UUID 

1584 # instance. 

1585 hex_ref = reader_butler.get_dataset(put_ref.id.hex) 

1586 self.assertEqual(hex_ref.id, put_ref.id) 

1587 

1588 def testCollectionChainRedefine(self): 

1589 butler = self._setup_to_test_collection_chain() 

1590 

1591 butler.collections.redefine_chain("chain", "a") 

1592 self._check_chain(butler, ["a"]) 

1593 

1594 # Duplicates are removed from the list of children 

1595 butler.collections.redefine_chain("chain", ["c", "b", "c"]) 

1596 self._check_chain(butler, ["c", "b"]) 

1597 

1598 # Empty list clears the chain 

1599 butler.collections.redefine_chain("chain", []) 

1600 self._check_chain(butler, []) 

1601 

1602 self._test_common_chain_functionality(butler, butler.collections.redefine_chain) 

1603 

1604 def testCollectionChainPrepend(self): 

1605 butler = self._setup_to_test_collection_chain() 

1606 

1607 # Duplicates are removed from the list of children 

1608 butler.collections.prepend_chain("chain", ["c", "b", "c"]) 

1609 self._check_chain(butler, ["c", "b"]) 

1610 

1611 # Prepend goes on the front of existing chain 

1612 butler.collections.prepend_chain("chain", ["a"]) 

1613 self._check_chain(butler, ["a", "c", "b"]) 

1614 

1615 # Empty prepend does nothing 

1616 butler.collections.prepend_chain("chain", []) 

1617 self._check_chain(butler, ["a", "c", "b"]) 

1618 

1619 # Prepending children that already exist in the chain removes them from 

1620 # their current position. 

1621 butler.collections.prepend_chain("chain", ["d", "b", "c"]) 

1622 self._check_chain(butler, ["d", "b", "c", "a"]) 

1623 

1624 self._test_common_chain_functionality(butler, butler.collections.prepend_chain) 

1625 

1626 def testCollectionChainExtend(self): 

1627 butler = self._setup_to_test_collection_chain() 

1628 

1629 # Duplicates are removed from the list of children 

1630 butler.collections.extend_chain("chain", ["c", "b", "c"]) 

1631 self._check_chain(butler, ["c", "b"]) 

1632 

1633 # Extend goes on the end of existing chain 

1634 butler.collections.extend_chain("chain", ["a"]) 

1635 self._check_chain(butler, ["c", "b", "a"]) 

1636 

1637 # Empty extend does nothing 

1638 butler.collections.extend_chain("chain", []) 

1639 self._check_chain(butler, ["c", "b", "a"]) 

1640 

1641 # Extending children that already exist in the chain removes them from 

1642 # their current position. 

1643 butler.collections.extend_chain("chain", ["d", "b", "c"]) 

1644 self._check_chain(butler, ["a", "d", "b", "c"]) 

1645 

1646 self._test_common_chain_functionality(butler, butler.collections.extend_chain) 

1647 

1648 def testCollectionChainRemove(self) -> None: 

1649 butler = self._setup_to_test_collection_chain() 

1650 

1651 butler.collections.redefine_chain("chain", ["a", "b", "c", "d"]) 

1652 

1653 butler.collections.remove_from_chain("chain", "c") 

1654 self._check_chain(butler, ["a", "b", "d"]) 

1655 

1656 # Duplicates are allowed in the list of children 

1657 butler.collections.remove_from_chain("chain", ["b", "b", "a"]) 

1658 self._check_chain(butler, ["d"]) 

1659 

1660 # Empty remove does nothing 

1661 butler.collections.remove_from_chain("chain", []) 

1662 self._check_chain(butler, ["d"]) 

1663 

1664 # Removing children that aren't in the chain does nothing 

1665 butler.collections.remove_from_chain("chain", ["a", "chain"]) 

1666 self._check_chain(butler, ["d"]) 

1667 

1668 self._test_common_chain_functionality( 

1669 butler, butler.collections.remove_from_chain, skip_cycle_check=True 

1670 ) 

1671 

1672 def _setup_to_test_collection_chain(self) -> Butler: 

1673 butler = self.create_empty_butler(writeable=True) 

1674 

1675 butler.collections.register("chain", CollectionType.CHAINED) 

1676 

1677 runs = ["a", "b", "c", "d"] 

1678 for run in runs: 

1679 butler.collections.register(run) 

1680 

1681 butler.collections.register("staticchain", CollectionType.CHAINED) 

1682 butler.collections.redefine_chain("staticchain", ["a", "b"]) 

1683 

1684 return butler 

1685 

1686 def _check_chain(self, butler: Butler, expected: list[str]) -> None: 

1687 children = butler.collections.get_info("chain").children 

1688 self.assertEqual(expected, list(children)) 

1689 

1690 def _test_common_chain_functionality( 

1691 self, butler, func: Callable[[str, str | list[str]], Any], *, skip_cycle_check=False 

1692 ) -> None: 

1693 # Missing parent collection 

1694 with self.assertRaises(MissingCollectionError): 

1695 func("doesnotexist", []) 

1696 # Missing child collection 

1697 with self.assertRaises(MissingCollectionError): 

1698 func("chain", ["doesnotexist"]) 

1699 # Forbid operations on non-chained collections 

1700 with self.assertRaises(CollectionTypeError): 

1701 func("d", ["a"]) 

1702 

1703 # Prevent collection cycles 

1704 if not skip_cycle_check: 

1705 butler.collections.register("chain2", CollectionType.CHAINED) 

1706 func("chain2", "chain") 

1707 with self.assertRaises(CollectionCycleError): 

1708 func("chain", "chain2") 

1709 

1710 # Make sure none of the earlier operations interfered with unrelated 

1711 # chains. 

1712 self.assertEqual(["a", "b"], list(butler.collections.get_info("staticchain").children)) 

1713 

1714 with butler._caching_context(): 

1715 with self.assertRaisesRegex(RuntimeError, "Chained collection modification not permitted"): 

1716 func("chain", "a") 

1717 

1718 def test_transfer_dimension_records_from(self) -> None: 

1719 source_butler = self.create_empty_butler(writeable=True) 

1720 source_butler.import_(filename=_get_test_data_path("lsstcam-subset.yaml")) 

1721 

1722 visit_id = 2025120200439 

1723 exposure_id = visit_id 

1724 target_butler = self.enterContext(create_populated_sqlite_registry()) 

1725 target_butler.transfer_dimension_records_from( 

1726 source_butler, 

1727 [ 

1728 # Should trigger the lookup of visit and all its associated 

1729 # "populated_by" records (visit_detector_region, 

1730 # visit_definition, etc.) 

1731 DataCoordinate.standardize( 

1732 {"instrument": "LSSTCam", "visit": visit_id, "detector": 10}, 

1733 universe=source_butler.dimensions, 

1734 ), 

1735 # Shouldn't add any records to the lookup. 

1736 DataCoordinate.make_empty(source_butler.dimensions), 

1737 ], 

1738 ) 

1739 

1740 def _fetch_record(dimension: str) -> DimensionRecord: 

1741 records = target_butler.query_dimension_records(dimension) 

1742 self.assertEqual(len(records), 1) 

1743 return records[0] 

1744 

1745 visit = _fetch_record("visit") 

1746 self.assertEqual(visit.id, visit_id) 

1747 self.assertEqual(visit.day_obs, 20251202) 

1748 self.assertEqual(visit.target_name, "lowdust") 

1749 self.assertEqual(visit.seq_num, 439) 

1750 original_visit = source_butler.query_dimension_records("visit", instrument="LSSTCam", visit=visit_id)[ 

1751 0 

1752 ] 

1753 self.assertEqual(visit.region, original_visit.region) 

1754 self.assertEqual(visit.timespan, original_visit.timespan) 

1755 

1756 visit_detector_region = _fetch_record("visit_detector_region") 

1757 self.assertEqual(visit_detector_region.instrument, "LSSTCam") 

1758 self.assertEqual(visit_detector_region.detector, 10) 

1759 self.assertEqual(visit_detector_region.visit, visit_id) 

1760 original_visit_detector_region = source_butler.query_dimension_records( 

1761 "visit_detector_region", instrument="LSSTCam", visit=visit_id, detector=10 

1762 )[0] 

1763 self.assertEqual(visit_detector_region.region, original_visit_detector_region.region) 

1764 

1765 visit_definition = _fetch_record("visit_definition") 

1766 self.assertEqual(visit_definition.instrument, "LSSTCam") 

1767 self.assertEqual(visit_definition.exposure, 2025120200439) 

1768 self.assertEqual(visit_definition.visit, visit_id) 

1769 

1770 # The matching exposure record should have been pulled in via 

1771 # visit -> visit_definition. 

1772 exposure = _fetch_record("exposure") 

1773 self.assertEqual(exposure.instrument, "LSSTCam") 

1774 self.assertEqual(exposure.id, 2025120200439) 

1775 self.assertEqual(exposure.obs_id, "MC_O_20251202_000439") 

1776 original_exposure = source_butler.query_dimension_records( 

1777 "exposure", instrument="LSSTCam", exposure=exposure_id 

1778 )[0] 

1779 self.assertEqual(exposure.timespan, original_exposure.timespan) 

1780 

1781 group = _fetch_record("group") 

1782 self.assertEqual(group.instrument, "LSSTCam") 

1783 self.assertEqual(group.name, "2025-12-03T07:58:10.858") 

1784 

1785 visit_system_memberships = target_butler.query_dimension_records("visit_system_membership") 

1786 visit_system_memberships.sort(key=lambda record: record.visit_system) 

1787 self.assertEqual(len(visit_system_memberships), 2) 

1788 self.assertEqual(visit_system_memberships[0].visit_system, 0) 

1789 self.assertEqual(visit_system_memberships[1].visit_system, 2) 

1790 self.assertEqual(visit_system_memberships[0].visit, visit_id) 

1791 self.assertEqual(visit_system_memberships[1].visit, visit_id) 

1792 

1793 visit_systems = target_butler.query_dimension_records("visit_system") 

1794 visit_systems.sort(key=lambda record: record.id) 

1795 visit_system_memberships.sort(key=lambda record: record.visit_system) 

1796 self.assertEqual(visit_systems[0].id, 0) 

1797 self.assertEqual(visit_systems[1].id, 2) 

1798 self.assertEqual(visit_systems[0].name, "one-to-one") 

1799 self.assertEqual(visit_systems[1].name, "by-seq-start-end") 

1800 

1801 

1802class FileDatastoreButlerTests(ButlerTests): 

1803 """Common tests and specialization of ButlerTests for butlers backed 

1804 by datastores that inherit from FileDatastore. 

1805 """ 

1806 

1807 trustModeSupported = True 

1808 

1809 def checkFileExists(self, root: str | ResourcePath, relpath: str | ResourcePath) -> bool: 

1810 """Check if file exists at a given path (relative to root). 

1811 

1812 Test testPutTemplates verifies actual physical existance of the files 

1813 in the requested location. 

1814 """ 

1815 uri = ResourcePath(root, forceDirectory=True) 

1816 return uri.join(relpath).exists() 

1817 

1818 def testPutTemplates(self) -> None: 

1819 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

1820 butler = self.create_empty_butler(run=self.default_run) 

1821 

1822 # Add needed Dimensions 

1823 butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

1824 butler.registry.insertDimensionData( 

1825 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

1826 ) 

1827 butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20250101}) 

1828 butler.registry.insertDimensionData( 

1829 "visit", 

1830 { 

1831 "instrument": "DummyCamComp", 

1832 "id": 423, 

1833 "name": "v423", 

1834 "physical_filter": "d-r", 

1835 "day_obs": 20250101, 

1836 }, 

1837 ) 

1838 butler.registry.insertDimensionData( 

1839 "visit", 

1840 { 

1841 "instrument": "DummyCamComp", 

1842 "id": 425, 

1843 "name": "v425", 

1844 "physical_filter": "d-r", 

1845 "day_obs": 20250101, 

1846 }, 

1847 ) 

1848 

1849 # Create and store a dataset 

1850 metric = makeExampleMetrics() 

1851 

1852 # Create two almost-identical DatasetTypes (both will use default 

1853 # template) 

1854 dimensions = butler.dimensions.conform(["instrument", "visit"]) 

1855 butler.registry.registerDatasetType(DatasetType("metric1", dimensions, storageClass)) 

1856 butler.registry.registerDatasetType(DatasetType("metric2", dimensions, storageClass)) 

1857 butler.registry.registerDatasetType(DatasetType("metric3", dimensions, storageClass)) 

1858 

1859 dataId1 = {"instrument": "DummyCamComp", "visit": 423} 

1860 dataId2 = {"instrument": "DummyCamComp", "visit": 423, "physical_filter": "d-r"} 

1861 

1862 # Put with exactly the data ID keys needed 

1863 ref = butler.put(metric, "metric1", dataId1) 

1864 uri = butler.getURI(ref) 

1865 self.assertTrue(uri.exists()) 

1866 self.assertTrue( 

1867 uri.unquoted_path.endswith(f"{self.default_run}/metric1/??#?/d-r/DummyCamComp_423.pickle") 

1868 ) 

1869 

1870 # Check the template based on dimensions 

1871 if hasattr(butler._datastore, "templates"): 

1872 butler._datastore.templates.validateTemplates([ref]) 

1873 

1874 # Put with extra data ID keys (physical_filter is an optional 

1875 # dependency); should not change template (at least the way we're 

1876 # defining them to behave now; the important thing is that they 

1877 # must be consistent). 

1878 ref = butler.put(metric, "metric2", dataId2) 

1879 uri = butler.getURI(ref) 

1880 self.assertTrue(uri.exists()) 

1881 self.assertTrue( 

1882 uri.unquoted_path.endswith(f"{self.default_run}/metric2/d-r/DummyCamComp_v423.pickle") 

1883 ) 

1884 

1885 # Check the template based on dimensions 

1886 if hasattr(butler._datastore, "templates"): 

1887 butler._datastore.templates.validateTemplates([ref]) 

1888 

1889 # Use a template that has a typo in dimension record metadata. 

1890 # Easier to test with a butler that has a ref with records attached. 

1891 template = FileTemplate("a/{visit.name}/{id}_{visit.namex:?}.fits") 

1892 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"): 

1893 path = template.format(ref) 

1894 self.assertEqual(path, f"a/v423/{ref.id}_fits") 

1895 

1896 template = FileTemplate("a/{visit.name}/{id}_{visit.namex}.fits") 

1897 with self.assertRaises(KeyError): 

1898 with self.assertLogs("lsst.daf.butler.datastore.file_templates", "INFO"): 

1899 template.format(ref) 

1900 

1901 # Now use a file template that will not result in unique filenames 

1902 with self.assertRaises(FileTemplateValidationError): 

1903 butler.put(metric, "metric3", dataId1) 

1904 

1905 def testImportExport(self) -> None: 

1906 # Run put/get tests just to create and populate a repo. 

1907 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

1908 self.runImportExportTest(storageClass) 

1909 

1910 @unittest.expectedFailure 

1911 def testImportExportVirtualComposite(self) -> None: 

1912 # Run put/get tests just to create and populate a repo. 

1913 storageClass = self.storageClassFactory.getStorageClass("StructuredComposite") 

1914 self.runImportExportTest(storageClass) 

1915 

1916 def runImportExportTest(self, storageClass: StorageClass) -> None: 

1917 """Test exporting and importing. 

1918 

1919 This test does an export to a temp directory and an import back 

1920 into a new temp directory repo. It does not assume a posix datastore. 

1921 """ 

1922 exportButler = self.runPutGetTest(storageClass, "test_metric") 

1923 

1924 # Test that we must have a file extension. 

1925 with self.assertRaises(ValueError): 

1926 with exportButler.export(filename="dump", directory=".") as export: 

1927 pass 

1928 

1929 # Test that unknown format is not allowed. 

1930 with self.assertRaises(ValueError): 

1931 with exportButler.export(filename="dump.fits", directory=".") as export: 

1932 pass 

1933 

1934 # Test that the repo actually has at least one dataset. 

1935 datasets = list(exportButler.registry.queryDatasets(..., collections=...)) 

1936 self.assertGreater(len(datasets), 0) 

1937 # Add a DimensionRecord that's unused by those datasets. 

1938 skymapRecord = {"name": "example_skymap", "hash": (50).to_bytes(8, byteorder="little")} 

1939 exportButler.registry.insertDimensionData("skymap", skymapRecord) 

1940 # Export and then import datasets. 

1941 with safeTestTempDir(TESTDIR) as exportDir: 

1942 exportFile = os.path.join(exportDir, "exports.yaml") 

1943 with exportButler.export(filename=exportFile, directory=exportDir, transfer="auto") as export: 

1944 export.saveDatasets(datasets) 

1945 # Export the same datasets again. This should quietly do 

1946 # nothing because of internal deduplication, and it shouldn't 

1947 # complain about being asked to export the "htm7" elements even 

1948 # though there aren't any in these datasets or in the database. 

1949 export.saveDatasets(datasets, elements=["htm7"]) 

1950 # Save one of the data IDs again; this should be harmless 

1951 # because of internal deduplication. 

1952 export.saveDataIds([datasets[0].dataId]) 

1953 # Save some dimension records directly. 

1954 export.saveDimensionData("skymap", [skymapRecord]) 

1955 self.assertTrue(os.path.exists(exportFile)) 

1956 with safeTestTempDir(TESTDIR) as importDir: 

1957 # We always want this to be a local posix butler 

1958 Butler.makeRepo(importDir, config=Config(os.path.join(TESTDIR, "config/basic/butler.yaml"))) 

1959 # Calling script.butlerImport tests the implementation of the 

1960 # butler command line interface "import" subcommand. Functions 

1961 # in the script folder are generally considered protected and 

1962 # should not be used as public api. 

1963 with open(exportFile) as f: 

1964 script.butlerImport( 

1965 importDir, 

1966 export_file=f, 

1967 directory=exportDir, 

1968 transfer="auto", 

1969 skip_dimensions=None, 

1970 ) 

1971 importButler = Butler.from_config(importDir, run=self.default_run) 

1972 self.enterContext(importButler) 

1973 for ref in datasets: 

1974 with self.subTest(ref=repr(ref)): 

1975 # Test for existence by passing in the DatasetType and 

1976 # data ID separately, to avoid lookup by dataset_id. 

1977 self.assertTrue(importButler.exists(ref.datasetType, ref.dataId)) 

1978 self.assertEqual( 

1979 list(importButler.registry.queryDimensionRecords("skymap")), 

1980 [importButler.dimensions["skymap"].RecordClass(**skymapRecord)], 

1981 ) 

1982 

1983 def testRemoveRuns(self) -> None: 

1984 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

1985 butler = self.create_empty_butler(writeable=True) 

1986 # Load registry data with dimensions to hang datasets off of. 

1987 butler.import_(filename=ResourcePath("resource://lsst.daf.butler/tests/registry_data/base.yaml")) 

1988 # Add some RUN-type collection. 

1989 run1 = "run1" 

1990 butler.collections.register(run1) 

1991 run2 = "run2" 

1992 butler.collections.register(run2) 

1993 # put a dataset in each 

1994 metric = makeExampleMetrics() 

1995 dimensions = butler.dimensions.conform(["instrument", "physical_filter"]) 

1996 datasetType = self.addDatasetType( 

1997 "prune_collections_test_dataset", dimensions, storageClass, butler.registry 

1998 ) 

1999 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1) 

2000 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2) 

2001 uri1 = butler.getURI(ref1) 

2002 uri2 = butler.getURI(ref2) 

2003 

2004 # Put one of the runs in a chain. 

2005 butler.collections.register("Chain", CollectionType.CHAINED) 

2006 butler.collections.extend_chain("Chain", run1) 

2007 

2008 with self.assertRaises(OrphanedRecordError): 

2009 butler.registry.removeDatasetType(datasetType.name) 

2010 

2011 # Remove a non-run. 

2012 with self.assertRaises(TypeError): 

2013 butler.removeRuns(["Chain"]) 

2014 

2015 # Remove without unlinking from chain should fail. 

2016 with self.assertRaises(IntegrityError): 

2017 butler.removeRuns([run1]) 

2018 

2019 # Remove from both runs. No longer use unstore parameter since it 

2020 # always purges. 

2021 butler.removeRuns([run1, run2], unlink_from_chains=True) 

2022 

2023 # Should be nothing in registry for either one, and datastore should 

2024 # not think either exists. 

2025 with self.assertRaises(MissingCollectionError): 

2026 butler.collections.get_info(run1) 

2027 with self.assertRaises(MissingCollectionError): 

2028 butler.collections.get_info(run1) 

2029 self.assertFalse(butler.stored(ref1)) 

2030 self.assertFalse(butler.stored(ref2)) 

2031 # We always unstore so both URIs should be gone. 

2032 self.assertFalse(uri1.exists()) 

2033 self.assertFalse(uri2.exists()) 

2034 

2035 # Now that the collections have been pruned we can remove the 

2036 # dataset type 

2037 butler.registry.removeDatasetType(datasetType.name) 

2038 

2039 with self.assertLogs("lsst.daf.butler.registry", "INFO") as cm: 

2040 butler.registry.removeDatasetType(("test*", "test*")) 

2041 self.assertIn("not defined", "\n".join(cm.output)) 

2042 

2043 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None: 

2044 """Simulate an external actor removing a file outside of Butler's 

2045 knowledge. 

2046 

2047 Subclasses may override to handle more complicated datastore 

2048 configurations. 

2049 """ 

2050 uri = butler.getURI(ref) 

2051 uri.remove() 

2052 datastore = cast(FileDatastore, butler._datastore) 

2053 datastore.cacheManager.remove_from_cache(ref) 

2054 

2055 def testPruneDatasets(self) -> None: 

2056 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

2057 butler = self.create_empty_butler(writeable=True) 

2058 # Load registry data with dimensions to hang datasets off of. 

2059 butler.import_(filename=_get_test_data_path("base.yaml")) 

2060 # Add some RUN-type collections. 

2061 run1 = "run1" 

2062 butler.collections.register(run1) 

2063 run2 = "run2" 

2064 butler.collections.register(run2) 

2065 # put some datasets. ref1 and ref2 have the same data ID, and are in 

2066 # different runs. ref3 has a different data ID. 

2067 metric = makeExampleMetrics() 

2068 dimensions = butler.dimensions.conform(["instrument", "physical_filter"]) 

2069 datasetType = self.addDatasetType( 

2070 "prune_collections_test_dataset", dimensions, storageClass, butler.registry 

2071 ) 

2072 ref1 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run1) 

2073 ref2 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-G"}, run=run2) 

2074 ref3 = butler.put(metric, datasetType, {"instrument": "Cam1", "physical_filter": "Cam1-R1"}, run=run1) 

2075 

2076 many_stored = butler.stored_many([ref1, ref2, ref3]) 

2077 for ref, stored in many_stored.items(): 

2078 self.assertTrue(stored, f"Ref {ref} should be stored") 

2079 

2080 many_exists = butler._exists_many([ref1, ref2, ref3]) 

2081 for ref, exists in many_exists.items(): 

2082 self.assertTrue(exists, f"Checking ref {ref} exists.") 

2083 self.assertEqual(exists, DatasetExistence.VERIFIED, f"Ref {ref} should be stored") 

2084 

2085 # Simple prune. 

2086 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True) 

2087 self.assertFalse(butler.exists(ref1.datasetType, ref1.dataId, collections=run1)) 

2088 

2089 many_stored = butler.stored_many([ref1, ref2, ref3]) 

2090 for ref, stored in many_stored.items(): 

2091 self.assertFalse(stored, f"Ref {ref} should not be stored") 

2092 

2093 many_exists = butler._exists_many([ref1, ref2, ref3]) 

2094 for ref, exists in many_exists.items(): 

2095 self.assertEqual(exists, DatasetExistence.UNRECOGNIZED, f"Ref {ref} should not be stored") 

2096 

2097 # Put data back. 

2098 ref1_new = butler.put(metric, ref1) 

2099 self.assertEqual(ref1_new, ref1) # Reuses original ID. 

2100 ref2 = butler.put(metric, ref2) 

2101 

2102 many_stored = butler.stored_many([ref1, ref2, ref3]) 

2103 self.assertTrue(many_stored[ref1]) 

2104 self.assertTrue(many_stored[ref2]) 

2105 self.assertFalse(many_stored[ref3]) 

2106 

2107 ref3 = butler.put(metric, ref3) 

2108 

2109 many_exists = butler._exists_many([ref1, ref2, ref3]) 

2110 for ref, exists in many_exists.items(): 

2111 self.assertTrue(exists, f"Ref {ref} should not be stored") 

2112 

2113 # Clear out the datasets from registry and start again. 

2114 refs = [ref1, ref2, ref3] 

2115 butler.pruneDatasets(refs, purge=True, unstore=True) 

2116 for ref in refs: 

2117 butler.put(metric, ref) 

2118 

2119 # Confirm we can retrieve deferred. 

2120 dref1 = butler.getDeferred(ref1) # known and exists 

2121 metric1 = dref1.get() 

2122 self.assertEqual(metric1, metric) 

2123 

2124 # Test different forms of file availability. 

2125 # Need to be in a state where: 

2126 # - one ref just has registry record. 

2127 # - one ref has a missing file but a datastore record. 

2128 # - one ref has a missing datastore record but file is there. 

2129 # - one ref does not exist anywhere. 

2130 # Do not need to test a ref that has everything since that is tested 

2131 # above. 

2132 ref0 = DatasetRef( 

2133 datasetType, 

2134 DataCoordinate.standardize( 

2135 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, universe=butler.dimensions 

2136 ), 

2137 run=run1, 

2138 ) 

2139 

2140 # Delete from datastore and retain in Registry. 

2141 butler.pruneDatasets([ref1], purge=False, unstore=True, disassociate=False) 

2142 

2143 # File has been removed. 

2144 self.remove_dataset_out_of_band(butler, ref2) 

2145 

2146 # Datastore has lost track. 

2147 butler._datastore.forget([ref3]) 

2148 

2149 # First test with a standard butler. 

2150 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True) 

2151 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED) 

2152 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED) 

2153 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE) 

2154 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED) 

2155 

2156 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=False) 

2157 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED) 

2158 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED | DatasetExistence._ASSUMED) 

2159 self.assertEqual(exists_many[ref2], DatasetExistence.KNOWN) 

2160 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ASSUMED) 

2161 self.assertTrue(exists_many[ref2]) 

2162 

2163 # Check that per-ref query gives the same answer as many query. 

2164 for ref, exists in exists_many.items(): 

2165 self.assertEqual(butler.exists(ref, full_check=False), exists) 

2166 

2167 # Get deferred checks for existence before it allows it to be 

2168 # retrieved. 

2169 with self.assertRaises(LookupError): 

2170 butler.getDeferred(ref3) # not known, file exists 

2171 dref2 = butler.getDeferred(ref2) # known but file missing 

2172 with self.assertRaises(FileNotFoundError): 

2173 dref2.get() 

2174 

2175 # Test again with a trusting butler. 

2176 if self.trustModeSupported: 

2177 butler._datastore.trustGetRequest = True 

2178 exists_many = butler._exists_many([ref0, ref1, ref2, ref3], full_check=True) 

2179 self.assertEqual(exists_many[ref0], DatasetExistence.UNRECOGNIZED) 

2180 self.assertEqual(exists_many[ref1], DatasetExistence.RECORDED) 

2181 self.assertEqual(exists_many[ref2], DatasetExistence.RECORDED | DatasetExistence.DATASTORE) 

2182 self.assertEqual(exists_many[ref3], DatasetExistence.RECORDED | DatasetExistence._ARTIFACT) 

2183 

2184 # When trusting we can get a deferred dataset handle that is not 

2185 # known but does exist. 

2186 dref3 = butler.getDeferred(ref3) 

2187 metric3 = dref3.get() 

2188 self.assertEqual(metric3, metric) 

2189 

2190 # Check that per-ref query gives the same answer as many query. 

2191 for ref, exists in exists_many.items(): 

2192 self.assertEqual(butler.exists(ref, full_check=True), exists) 

2193 

2194 # Create a ref that surprisingly has the UUID of an existing ref 

2195 # but is not the same. 

2196 ref_bad = DatasetRef(datasetType, dataId=ref3.dataId, run=ref3.run, id=ref2.id) 

2197 with self.assertRaises(ValueError): 

2198 butler.exists(ref_bad) 

2199 

2200 # Create a ref that has a compatible storage class. 

2201 ref_compat = ref2.overrideStorageClass("StructuredDataDict") 

2202 exists = butler.exists(ref_compat) 

2203 self.assertEqual(exists, exists_many[ref2]) 

2204 

2205 # Remove everything and start from scratch. 

2206 butler._datastore.trustGetRequest = False 

2207 butler.pruneDatasets(refs, purge=True, unstore=True) 

2208 for ref in refs: 

2209 butler.put(metric, ref) 

2210 

2211 # These tests mess directly with the trash table and can leave the 

2212 # datastore in an odd state. Do them at the end. 

2213 # Check that in normal mode, deleting the record will lead to 

2214 # trash not touching the file. 

2215 uri1 = butler.getURI(ref1) 

2216 butler._datastore.bridge.moveToTrash( 

2217 [ref1], transaction=None 

2218 ) # Update the dataset_location table 

2219 butler._datastore.forget([ref1]) 

2220 butler._datastore.trash(ref1) 

2221 butler._datastore.emptyTrash() 

2222 self.assertTrue(uri1.exists()) 

2223 uri1.remove() # Clean it up. 

2224 

2225 # Simulate execution butler setup by deleting the datastore 

2226 # record but keeping the file around and trusting. 

2227 butler._datastore.trustGetRequest = True 

2228 uris = butler.get_many_uris([ref2, ref3]) 

2229 uri2 = uris[ref2].primaryURI 

2230 uri3 = uris[ref3].primaryURI 

2231 self.assertTrue(uri2.exists()) 

2232 self.assertTrue(uri3.exists()) 

2233 

2234 # Remove the datastore record. 

2235 butler._datastore.bridge.moveToTrash( 

2236 [ref2], transaction=None 

2237 ) # Update the dataset_location table 

2238 butler._datastore.forget([ref2]) 

2239 self.assertTrue(uri2.exists()) 

2240 butler._datastore.trash([ref2, ref3]) 

2241 # Immediate removal for ref2 file 

2242 self.assertFalse(uri2.exists()) 

2243 # But ref3 has to wait for the empty. 

2244 self.assertTrue(uri3.exists()) 

2245 butler._datastore.emptyTrash() 

2246 self.assertFalse(uri3.exists()) 

2247 

2248 # Clear out the datasets from registry. 

2249 butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True) 

2250 

2251 def test_butler_metrics(self): 

2252 """Test that metrics are collected.""" 

2253 run = "test_run" 

2254 metrics = ButlerMetrics() 

2255 butler, datasetType = self.create_butler( 

2256 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics 

2257 ) 

2258 data = MetricsExampleModel( 

2259 summary={"AM1": 5.2, "AM2": 30.6}, 

2260 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

2261 data=[563, 234, 456.7, 752, 8, 9, 27], 

2262 ) 

2263 

2264 data_ref = butler.put(data, datasetType, visit=424, instrument="DummyCamComp") 

2265 butler.get(data_ref) 

2266 butler.get(data_ref) 

2267 self.assertEqual(metrics.n_get, 2) 

2268 self.assertGreater(metrics.time_in_get, 0.0) 

2269 self.assertEqual(metrics.n_put, 1) 

2270 self.assertGreater(metrics.time_in_put, 0.0) 

2271 

2272 deferred = butler.getDeferred(data_ref) 

2273 deferred.get() 

2274 self.assertEqual(metrics.n_get, 3) 

2275 

2276 with butler.record_metrics() as new: 

2277 data_ref_2 = butler.put(data, datasetType, visit=425, instrument="DummyCamComp") 

2278 butler.get(data_ref) 

2279 

2280 butler.pruneDatasets([data_ref, data_ref_2], purge=True, unstore=True) 

2281 with ResourcePath.temporary_uri(suffix=".json") as tmpFile: 

2282 tmpFile.write(data.model_dump_json().encode()) 

2283 refs = [ 

2284 DatasetRef(datasetType, data_ref_2.dataId, run), 

2285 DatasetRef(datasetType, data_ref.dataId, run), 

2286 ] 

2287 datasets = [FileDataset(path=tmpFile, refs=refs)] 

2288 butler.ingest(*datasets, transfer="copy") 

2289 

2290 self.assertEqual(new.n_get, 1) 

2291 self.assertEqual(new.n_put, 1) 

2292 self.assertEqual(new.n_ingest, 2) 

2293 

2294 

2295class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2296 """PosixDatastore specialization of a butler""" 

2297 

2298 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

2299 fullConfigKey: str | None = ".datastore.formatters" 

2300 validationCanFail = True 

2301 datastoreStr = ["/tmp"] 

2302 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"] 

2303 registryStr = "/gen3.sqlite3" 

2304 

2305 def testPathConstructor(self) -> None: 

2306 """Independent test of constructor using PathLike.""" 

2307 butler = Butler.from_config(self.tmpConfigFile, run=self.default_run) 

2308 self.enterContext(butler) 

2309 self.assertIsInstance(butler, Butler) 

2310 

2311 # And again with a Path object with the butler yaml 

2312 path = pathlib.Path(self.tmpConfigFile) 

2313 butler = Butler.from_config(path, writeable=False) 

2314 self.enterContext(butler) 

2315 self.assertIsInstance(butler, Butler) 

2316 

2317 # And again with a Path object without the butler yaml 

2318 # (making sure we skip it if the tmp config doesn't end 

2319 # in butler.yaml -- which is the case for a subclass) 

2320 if self.tmpConfigFile.endswith("butler.yaml"): 

2321 path = pathlib.Path(os.path.dirname(self.tmpConfigFile)) 

2322 butler = Butler.from_config(path, writeable=False) 

2323 self.enterContext(butler) 

2324 self.assertIsInstance(butler, Butler) 

2325 

2326 def testExportTransferCopy(self) -> None: 

2327 """Test local export using all transfer modes""" 

2328 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

2329 exportButler = self.runPutGetTest(storageClass, "test_metric") 

2330 # Test that the repo actually has at least one dataset. 

2331 datasets = list(exportButler.registry.queryDatasets(..., collections=...)) 

2332 self.assertGreater(len(datasets), 0) 

2333 uris = [exportButler.getURI(d) for d in datasets] 

2334 assert isinstance(exportButler._datastore, FileDatastore) 

2335 datastoreRoot = exportButler.get_datastore_roots()[exportButler.get_datastore_names()[0]] 

2336 

2337 pathsInStore = [uri.relative_to(datastoreRoot) for uri in uris] 

2338 

2339 for path in pathsInStore: 

2340 # Assume local file system 

2341 assert path is not None 

2342 self.assertTrue(self.checkFileExists(datastoreRoot, path), f"Checking path {path}") 

2343 

2344 for transfer in ("copy", "link", "symlink", "relsymlink"): 

2345 with safeTestTempDir(TESTDIR) as exportDir: 

2346 with exportButler.export(directory=exportDir, format="yaml", transfer=transfer) as export: 

2347 export.saveDatasets(datasets) 

2348 for path in pathsInStore: 

2349 assert path is not None 

2350 self.assertTrue( 

2351 self.checkFileExists(exportDir, path), 

2352 f"Check that mode {transfer} exported files", 

2353 ) 

2354 

2355 def testPytypeCoercion(self) -> None: 

2356 """Test python type coercion on Butler.get and put.""" 

2357 # Store some data with the normal example storage class. 

2358 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

2359 datasetTypeName = "test_metric" 

2360 butler = self.runPutGetTest(storageClass, datasetTypeName) 

2361 

2362 dataId = {"instrument": "DummyCamComp", "visit": 423} 

2363 metric = butler.get(datasetTypeName, dataId=dataId) 

2364 self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample") 

2365 

2366 datasetType_ori = butler.get_dataset_type(datasetTypeName) 

2367 self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents") 

2368 

2369 # Now need to hack the registry dataset type definition. 

2370 # There is no API for this. 

2371 assert isinstance(butler._registry, SqlRegistry) 

2372 manager = butler._registry._managers.datasets 

2373 assert hasattr(manager, "_db") and hasattr(manager, "_static") 

2374 manager._db.update( 

2375 manager._static.dataset_type, 

2376 {"name": datasetTypeName}, 

2377 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"}, 

2378 ) 

2379 

2380 # Force reset of dataset type cache 

2381 butler.registry.refresh() 

2382 

2383 datasetType_new = butler.get_dataset_type(datasetTypeName) 

2384 self.assertEqual(datasetType_new.name, datasetType_ori.name) 

2385 self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel") 

2386 

2387 metric_model = butler.get(datasetTypeName, dataId=dataId) 

2388 self.assertNotEqual(type(metric_model), type(metric)) 

2389 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel") 

2390 

2391 # Put the model and read it back to show that everything now 

2392 # works as normal. 

2393 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424) 

2394 metric_model_new = butler.get(metric_ref) 

2395 self.assertEqual(metric_model_new, metric_model) 

2396 

2397 # Hack the storage class again to something that will fail on the 

2398 # get with no conversion class. 

2399 manager._db.update( 

2400 manager._static.dataset_type, 

2401 {"name": datasetTypeName}, 

2402 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"}, 

2403 ) 

2404 butler.registry.refresh() 

2405 

2406 with self.assertRaises(ValueError): 

2407 butler.get(datasetTypeName, dataId=dataId) 

2408 

2409 def test_provenance(self): 

2410 """Test that provenance is attached on put.""" 

2411 run = "test_run" 

2412 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric") 

2413 metric = MetricsExampleModel( 

2414 summary={"AM1": 5.2, "AM2": 30.6}, 

2415 output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

2416 data=[563, 234, 456.7, 752, 8, 9, 27], 

2417 ) 

2418 # Provenance can be attached to the object being put. Whether 

2419 # it is or not is dependent on the formatter. For this test we 

2420 # copy on adding provenance to ensure they differ. 

2421 self.assertIsNone(metric.dataset_id) 

2422 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp") 

2423 self.assertIsNone(metric.dataset_id) 

2424 metric_2 = butler.get(metric_ref) 

2425 self.assertEqual(metric_2.data, metric.data) 

2426 self.assertEqual(metric_2.dataset_id, metric_ref.id) 

2427 self.assertIsNone(metric_2.provenance) 

2428 

2429 # Put with provenance. 

2430 prov = DatasetProvenance(quantum_id=uuid.uuid4()) 

2431 prov.add_input(metric_ref) 

2432 prov.add_extra_provenance(metric_ref.id, {"answer": 42}) 

2433 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov) 

2434 metric_3 = butler.get(metric_ref2) 

2435 self.assertEqual(metric_3.provenance, prov) 

2436 

2437 # Check that we can extract provenance from dict form. 

2438 prov_dict = prov.to_flat_dict(metric_ref2) 

2439 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler) 

2440 self.assertEqual(ref_from_prov, metric_ref2) 

2441 # Direct __eq__ of the provenance does not work because one side 

2442 # includes dimension records. 

2443 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs}) 

2444 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id) 

2445 self.assertEqual(prov_from_prov.extras, prov.extras) 

2446 

2447 # Force a bad ID into the dict. 

2448 prov_dict["id"] = uuid.uuid4() 

2449 with self.assertRaises(ValueError): 

2450 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2451 del prov_dict["id"] 

2452 prov_dict["input 0 id"] = uuid.uuid4() 

2453 with self.assertRaises(ValueError): 

2454 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2455 

2456 # Check that simple types can be reconstructed with non-standard 

2457 # separators. 

2458 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True) 

2459 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler) 

2460 self.assertEqual(ref_from_prov, metric_ref2) 

2461 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs}) 

2462 

2463 with self.assertRaises(ValueError): 

2464 DatasetProvenance.from_flat_dict({"unknown": 42}, butler) 

2465 

2466 def test_specialized_file_datasets_functions(self): 

2467 """Test a workflow used in Prompt Processing where we export datasets 

2468 from one repository and write them in-place to the datastore of 

2469 another, without immediately inserting registry entries for the 

2470 datasets. 

2471 """ 

2472 repo = MetricTestRepo.create_from_butler( 

2473 self.create_empty_butler(writeable=True), 

2474 self.tmpConfigFile, 

2475 "StructuredCompositeReadCompNoDisassembly", 

2476 ) 

2477 source_butler = repo.butler 

2478 

2479 # Test writing outputs to a FileDatastore. 

2480 with tempfile.TemporaryDirectory() as tempdir: 

2481 target_repo_config = Butler.makeRepo(tempdir) 

2482 refs = [repo.ref1, repo.ref2] 

2483 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs) 

2484 self.assertEqual(len(datasets), 2) 

2485 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets}) 

2486 for dataset in datasets: 

2487 path = ResourcePath(dataset.path, forceAbsolute=False) 

2488 # Paths should be relative paths to the target datastore. 

2489 self.assertFalse(path.isabs()) 

2490 # Files should have been copied into the target datastore 

2491 self.assertTrue(ResourcePath(tempdir).join(path).exists()) 

2492 

2493 # Make sure the target Butler can ingest the datasets. 

2494 target_butler = Butler(target_repo_config, writeable=True) 

2495 self.enterContext(target_butler) 

2496 target_butler.transfer_dimension_records_from(source_butler, refs) 

2497 target_butler.ingest(*datasets, transfer=None) 

2498 self.assertIsNotNone(target_butler.get(repo.ref1)) 

2499 self.assertIsNotNone(target_butler.get(repo.ref2)) 

2500 

2501 # Giving an empty list of files is a no-op. 

2502 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), []) 

2503 self.assertEqual(len(no_datasets), 0) 

2504 

2505 # Test writing outputs to a ChainedDatastore. 

2506 with tempfile.TemporaryDirectory() as tempdir: 

2507 # Set up a second dataset type, so we can split the files across 

2508 # multiple datastore roots. 

2509 dt1 = repo.datasetType 

2510 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass) 

2511 source_butler.registry.registerDatasetType(dt2) 

2512 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2) 

2513 config = Config.fromString( 

2514 f""" 

2515 datastore: 

2516 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore 

2517 datastore_constraints: 

2518 - constraints: 

2519 accept: 

2520 - {dt1.name} 

2521 - constraints: 

2522 accept: 

2523 - {dt2.name} 

2524 datastores: 

2525 - datastore: 

2526 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore 

2527 root: <butlerRoot>/FileDatastore_0 

2528 - datastore: 

2529 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore 

2530 root: <butlerRoot>/FileDatastore_1 

2531 """ 

2532 ) 

2533 target_repo_config = Butler.makeRepo(tempdir, config) 

2534 refs = [repo.ref1, repo.ref2, other_ref] 

2535 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs) 

2536 self.assertEqual(len(datasets), 3) 

2537 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets}) 

2538 for dataset in datasets: 

2539 path = ResourcePath(dataset.path, forceAbsolute=False) 

2540 # Paths should be relative paths to the target datastore. 

2541 self.assertFalse(path.isabs()) 

2542 # Files should have been split up between the two datastores 

2543 # in the chain. 

2544 datastore_root = ResourcePath(tempdir) 

2545 if dataset.refs[0].datasetType.name == dt1.name: 

2546 datastore_root = datastore_root.join("FileDatastore_0") 

2547 else: 

2548 datastore_root = datastore_root.join("FileDatastore_1") 

2549 self.assertTrue(datastore_root.join(path).exists()) 

2550 

2551 # Make sure the target Butler can ingest the datasets. 

2552 target_butler = Butler(target_repo_config, writeable=True) 

2553 self.enterContext(target_butler) 

2554 target_butler.transfer_dimension_records_from(source_butler, refs) 

2555 target_butler.ingest(*datasets, transfer=None) 

2556 self.assertIsNotNone(target_butler.get(repo.ref1)) 

2557 self.assertIsNotNone(target_butler.get(repo.ref2)) 

2558 self.assertIsNotNone(target_butler.get(other_ref)) 

2559 

2560 def test_temporary_for_ingest(self) -> None: 

2561 """Test the `lsst.daf.butler._rubin.ingest_from_temporary` module.""" 

2562 with self.create_empty_butler("example_run") as butler: 

2563 dataset_type = DatasetType("example", butler.dimensions.empty, "StructuredDataDict") 

2564 butler.registry.registerDatasetType(dataset_type) 

2565 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(butler.dimensions), "example_run") 

2566 with TemporaryForIngest(butler, ref) as temporary: 

2567 temporary.path.write(b"three: 3") 

2568 found = TemporaryForIngest.find_orphaned_temporaries_by_ref(ref, butler) 

2569 self.assertEqual(found, [temporary.path]) 

2570 self.assertIn(".tmp", temporary.ospath) 

2571 temporary.ingest() 

2572 loaded = butler.get(ref) 

2573 self.assertEqual(loaded, {"three": 3}) 

2574 

2575 

2576class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2577 """PosixDatastore specialization of a butler using Postgres""" 

2578 

2579 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

2580 fullConfigKey = ".datastore.formatters" 

2581 validationCanFail = True 

2582 datastoreStr = ["/tmp"] 

2583 datastoreName = [f"FileDatastore@{BUTLER_ROOT_TAG}"] 

2584 registryStr = "PostgreSQL@test" 

2585 

2586 @classmethod 

2587 def setUpClass(cls) -> None: 

2588 cls.postgresql = cls.enterClassContext(setup_postgres_test_db()) 

2589 super().setUpClass() 

2590 

2591 def setUp(self) -> None: 

2592 # Need to add a registry section to the config. 

2593 self._temp_config = False 

2594 config = Config(self.configFile) 

2595 self.postgresql.patch_butler_config(config) 

2596 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh: 

2597 config.dump(fh) 

2598 self.configFile = fh.name 

2599 self._temp_config = True 

2600 super().setUp() 

2601 

2602 def tearDown(self) -> None: 

2603 if self._temp_config and os.path.exists(self.configFile): 

2604 os.remove(self.configFile) 

2605 super().tearDown() 

2606 

2607 def testMakeRepo(self) -> None: 

2608 # The base class test assumes that it's using sqlite and assumes 

2609 # the config file is acceptable to sqlite. 

2610 raise unittest.SkipTest("Postgres config is not compatible with this test.") 

2611 

2612 

2613class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase): 

2614 """Test that Butler with a Postgres registry still works after cloning.""" 

2615 

2616 def create_butler( 

2617 self, 

2618 run: str, 

2619 storageClass: StorageClass | str, 

2620 datasetTypeName: str, 

2621 metrics: ButlerMetrics | None = None, 

2622 ) -> tuple[DirectButler, DatasetType]: 

2623 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics) 

2624 return butler.clone(run=run, metrics=metrics), datasetType 

2625 

2626 

2627class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase): 

2628 """InMemoryDatastore specialization of a butler""" 

2629 

2630 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml") 

2631 fullConfigKey = None 

2632 useTempRoot = False 

2633 validationCanFail = False 

2634 datastoreStr = ["datastore='InMemory"] 

2635 datastoreName = ["InMemoryDatastore@"] 

2636 registryStr = "/gen3.sqlite3" 

2637 

2638 def testIngest(self) -> None: 

2639 pass 

2640 

2641 def test_ingest_zip(self) -> None: 

2642 pass 

2643 

2644 

2645class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase): 

2646 """Test that a Butler with a Sqlite registry still works after cloning.""" 

2647 

2648 def create_butler( 

2649 self, 

2650 run: str, 

2651 storageClass: StorageClass | str, 

2652 datasetTypeName: str, 

2653 metrics: ButlerMetrics | None = None, 

2654 ) -> tuple[DirectButler, DatasetType]: 

2655 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics) 

2656 return butler.clone(run=run), datasetType 

2657 

2658 

2659class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2660 """PosixDatastore specialization""" 

2661 

2662 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml") 

2663 fullConfigKey = ".datastore.datastores.1.formatters" 

2664 validationCanFail = True 

2665 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"] 

2666 datastoreName = [ 

2667 "InMemoryDatastore@", 

2668 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1", 

2669 "SecondDatastore", 

2670 ] 

2671 registryStr = "/gen3.sqlite3" 

2672 

2673 def testPruneDatasets(self) -> None: 

2674 # This test relies on manipulating files out-of-band, which is 

2675 # impossible for this configuration because of the InMemoryDatastore in 

2676 # the ChainedDatastore. 

2677 pass 

2678 

2679 

2680class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase): 

2681 """Test that a yaml file in one location can refer to a root in another.""" 

2682 

2683 datastoreStr = ["dir1"] 

2684 # Disable the makeRepo test since we are deliberately not using 

2685 # butler.yaml as the config name. 

2686 fullConfigKey = None 

2687 

2688 def setUp(self) -> None: 

2689 self.root = makeTestTempDir(TESTDIR) 

2690 

2691 # Make a new repository in one place 

2692 self.dir1 = os.path.join(self.root, "dir1") 

2693 Butler.makeRepo(self.dir1, config=Config(self.configFile)) 

2694 

2695 # Move the yaml file to a different place and add a "root" 

2696 self.dir2 = os.path.join(self.root, "dir2") 

2697 os.makedirs(self.dir2, exist_ok=True) 

2698 configFile1 = os.path.join(self.dir1, "butler.yaml") 

2699 config = Config(configFile1) 

2700 config["root"] = self.dir1 

2701 configFile2 = os.path.join(self.dir2, "butler2.yaml") 

2702 config.dumpToUri(configFile2) 

2703 os.remove(configFile1) 

2704 self.tmpConfigFile = configFile2 

2705 

2706 def testFileLocations(self) -> None: 

2707 self.assertNotEqual(self.dir1, self.dir2) 

2708 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml"))) 

2709 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml"))) 

2710 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3"))) 

2711 

2712 

2713class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase): 

2714 """Test that a config file created by makeRepo outside of repo works.""" 

2715 

2716 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

2717 

2718 def setUp(self) -> None: 

2719 self.root = makeTestTempDir(TESTDIR) 

2720 self.root2 = makeTestTempDir(TESTDIR) 

2721 

2722 self.tmpConfigFile = os.path.join(self.root2, "different.yaml") 

2723 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2724 

2725 def tearDown(self) -> None: 

2726 if os.path.exists(self.root2): 

2727 shutil.rmtree(self.root2, ignore_errors=True) 

2728 super().tearDown() 

2729 

2730 def testConfigExistence(self) -> None: 

2731 c = Config(self.tmpConfigFile) 

2732 uri_config = ResourcePath(c["root"]) 

2733 uri_expected = ResourcePath(self.root, forceDirectory=True) 

2734 self.assertEqual(uri_config.geturl(), uri_expected.geturl()) 

2735 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path") 

2736 

2737 def testPutGet(self) -> None: 

2738 storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents") 

2739 self.runPutGetTest(storageClass, "test_metric") 

2740 

2741 

2742class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase): 

2743 """Test that a config file created by makeRepo outside of repo works.""" 

2744 

2745 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

2746 

2747 def setUp(self) -> None: 

2748 self.root = makeTestTempDir(TESTDIR) 

2749 self.root2 = makeTestTempDir(TESTDIR) 

2750 

2751 self.tmpConfigFile = self.root2 

2752 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2753 

2754 def testConfigExistence(self) -> None: 

2755 # Append the yaml file else Config constructor does not know the file 

2756 # type. 

2757 self.tmpConfigFile = os.path.join(self.tmpConfigFile, "butler.yaml") 

2758 super().testConfigExistence() 

2759 

2760 

2761class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase): 

2762 """Test that a config file created by makeRepo outside of repo works.""" 

2763 

2764 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

2765 

2766 def setUp(self) -> None: 

2767 self.root = makeTestTempDir(TESTDIR) 

2768 self.root2 = makeTestTempDir(TESTDIR) 

2769 

2770 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl() 

2771 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2772 

2773 

2774@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!") 

2775class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2776 """S3Datastore specialization of a butler; an S3 storage Datastore + 

2777 a local in-memory SqlRegistry. 

2778 """ 

2779 

2780 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml") 

2781 fullConfigKey = None 

2782 validationCanFail = True 

2783 

2784 bucketName = "anybucketname" 

2785 """Name of the Bucket that will be used in the tests. The name is read from 

2786 the config file used with the tests during set-up. 

2787 """ 

2788 

2789 root = "butlerRoot/" 

2790 """Root repository directory expected to be used in case useTempRoot=False. 

2791 Otherwise the root is set to a 20 characters long randomly generated string 

2792 during set-up. 

2793 """ 

2794 

2795 datastoreStr = [f"datastore={root}"] 

2796 """Contains all expected root locations in a format expected to be 

2797 returned by Butler stringification. 

2798 """ 

2799 

2800 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"] 

2801 """The expected format of the S3 Datastore string.""" 

2802 

2803 registryStr = "/gen3.sqlite3" 

2804 """Expected format of the Registry string.""" 

2805 

2806 mock_aws = mock_aws() 

2807 """The mocked s3 interface from moto.""" 

2808 

2809 def genRoot(self) -> str: 

2810 """Return a random string of len 20 to serve as a root 

2811 name for the temporary bucket repo. 

2812 

2813 This is equivalent to tempfile.mkdtemp as this is what self.root 

2814 becomes when useTempRoot is True. 

2815 """ 

2816 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20)) 

2817 return rndstr + "/" 

2818 

2819 def setUp(self) -> None: 

2820 config = Config(self.configFile) 

2821 uri = ResourcePath(config[".datastore.datastore.root"]) 

2822 self.bucketName = uri.netloc 

2823 

2824 # Enable S3 mocking of tests. 

2825 self.enterContext(clean_test_environment_for_s3()) 

2826 self.mock_aws.start() 

2827 

2828 if self.useTempRoot: 

2829 self.root = self.genRoot() 

2830 rooturi = f"s3://{self.bucketName}/{self.root}" 

2831 config.update({"datastore": {"datastore": {"root": rooturi}}}) 

2832 

2833 # need local folder to store registry database 

2834 self.reg_dir = makeTestTempDir(TESTDIR) 

2835 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3" 

2836 

2837 # MOTO needs to know that we expect Bucket bucketname to exist 

2838 # (this used to be the class attribute bucketName) 

2839 s3 = boto3.resource("s3") 

2840 s3.create_bucket(Bucket=self.bucketName) 

2841 

2842 self.datastoreStr = [f"datastore='{rooturi}'"] 

2843 self.datastoreName = [f"FileDatastore@{rooturi}"] 

2844 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False) 

2845 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml") 

2846 

2847 def tearDown(self) -> None: 

2848 s3 = boto3.resource("s3") 

2849 bucket = s3.Bucket(self.bucketName) 

2850 try: 

2851 bucket.objects.all().delete() 

2852 except botocore.exceptions.ClientError as e: 

2853 if e.response["Error"]["Code"] == "404": 

2854 # the key was not reachable - pass 

2855 pass 

2856 else: 

2857 raise 

2858 

2859 bucket = s3.Bucket(self.bucketName) 

2860 bucket.delete() 

2861 

2862 # Stop the S3 mock. 

2863 self.mock_aws.stop() 

2864 

2865 if self.reg_dir is not None and os.path.exists(self.reg_dir): 

2866 shutil.rmtree(self.reg_dir, ignore_errors=True) 

2867 

2868 if self.useTempRoot and os.path.exists(self.root): 

2869 shutil.rmtree(self.root, ignore_errors=True) 

2870 

2871 super().tearDown() 

2872 

2873 

2874class DatastoreTransfers(TestCaseMixin): 

2875 """Base test setup for data transfers between butlers. The concrete tests 

2876 for specific configurations are in other classes, below. 

2877 """ 

2878 

2879 storageClassFactory: StorageClassFactory 

2880 

2881 @classmethod 

2882 def setUpClass(cls) -> None: 

2883 cls.storageClassFactory = StorageClassFactory() 

2884 

2885 def setUp(self) -> None: 

2886 self.root = makeTestTempDir(TESTDIR) 

2887 self.config = Config(self.configFile) 

2888 

2889 # Some tests cause convertors to be replaced so ensure 

2890 # the storage class factory is reset each time. 

2891 self.storageClassFactory.reset() 

2892 self.storageClassFactory.addFromConfig(self.configFile) 

2893 

2894 def tearDown(self) -> None: 

2895 removeTestTempDir(self.root) 

2896 

2897 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler: 

2898 if manager is None: 

2899 manager = ( 

2900 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID" 

2901 ) 

2902 config = Config(config_file if config_file is not None else self.configFile) 

2903 config["registry", "managers", "datasets"] = manager 

2904 butler = Butler.from_config( 

2905 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True 

2906 ) 

2907 self.enterContext(butler) 

2908 return butler 

2909 

2910 def assertButlerTransfers( 

2911 self, 

2912 purge: bool = False, 

2913 storageClassName: str = "StructuredData", 

2914 storageClassNameTarget: str | None = None, 

2915 ) -> None: 

2916 """Test that a run can be transferred to another butler.""" 

2917 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

2918 if storageClassNameTarget is not None: 

2919 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget) 

2920 else: 

2921 storageClassTarget = storageClass 

2922 

2923 datasetTypeName = "random_data" 

2924 

2925 # Test will create 3 collections and we will want to transfer 

2926 # two of those three. 

2927 runs = ["run1", "run2", "other"] 

2928 

2929 # Also want to use two different dataset types to ensure that 

2930 # grouping works. 

2931 datasetTypeNames = ["random_data", "random_data_2"] 

2932 

2933 # Create the run collections in the source butler. 

2934 for run in runs: 

2935 self.source_butler.collections.register(run) 

2936 

2937 # Create dimensions in source butler. 

2938 n_exposures = 30 

2939 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

2940 self.source_butler.registry.insertDimensionData( 

2941 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

2942 ) 

2943 self.source_butler.registry.insertDimensionData( 

2944 "detector", {"instrument": "DummyCamComp", "id": 1, "full_name": "det1"} 

2945 ) 

2946 self.source_butler.registry.insertDimensionData( 

2947 "day_obs", 

2948 { 

2949 "instrument": "DummyCamComp", 

2950 "id": 20250101, 

2951 }, 

2952 ) 

2953 

2954 for i in range(n_exposures): 

2955 self.source_butler.registry.insertDimensionData( 

2956 "group", {"instrument": "DummyCamComp", "name": f"group{i}"} 

2957 ) 

2958 self.source_butler.registry.insertDimensionData( 

2959 "exposure", 

2960 { 

2961 "instrument": "DummyCamComp", 

2962 "id": i, 

2963 "obs_id": f"exp{i}", 

2964 "physical_filter": "d-r", 

2965 "group": f"group{i}", 

2966 "day_obs": 20250101, 

2967 }, 

2968 ) 

2969 

2970 # Create dataset types in the source butler. 

2971 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"]) 

2972 for datasetTypeName in datasetTypeNames: 

2973 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

2974 self.source_butler.registry.registerDatasetType(datasetType) 

2975 

2976 # Write a dataset to an unrelated run -- this will ensure that 

2977 # we are rewriting integer dataset ids in the target if necessary. 

2978 # Will not be relevant for UUID. 

2979 run = "distraction" 

2980 butler = Butler.from_config(butler=self.source_butler, run=run) 

2981 self.enterContext(butler) 

2982 butler.put( 

2983 makeExampleMetrics(), 

2984 datasetTypeName, 

2985 exposure=1, 

2986 instrument="DummyCamComp", 

2987 physical_filter="d-r", 

2988 ) 

2989 

2990 # Write some example metrics to the source 

2991 butler = Butler.from_config(butler=self.source_butler) 

2992 self.enterContext(butler) 

2993 

2994 # Set of DatasetRefs that should be in the list of refs to transfer 

2995 # but which will not be transferred. 

2996 deleted: set[DatasetRef] = set() 

2997 

2998 n_expected = 20 # Number of datasets expected to be transferred 

2999 source_refs = [] 

3000 for i in range(n_exposures): 

3001 # Put a third of datasets into each collection, only retain 

3002 # two thirds. 

3003 index = i % 3 

3004 run = runs[index] 

3005 datasetTypeName = datasetTypeNames[i % 2] 

3006 

3007 metric = MetricsExample( 

3008 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)] 

3009 ) 

3010 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"} 

3011 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run) 

3012 

3013 # Remove the datastore record using low-level API, but only 

3014 # for a specific index. 

3015 if purge and index == 1: 

3016 # For one of these delete the file as well. 

3017 # This allows the "missing" code to filter the 

3018 # file out. 

3019 # Access the individual datastores. 

3020 datastores = [] 

3021 if hasattr(butler._datastore, "datastores"): 

3022 datastores.extend(butler._datastore.datastores) 

3023 else: 

3024 datastores.append(butler._datastore) 

3025 

3026 if not deleted: 

3027 # For a chained datastore we need to remove 

3028 # files in each chain. 

3029 for datastore in datastores: 

3030 # The file might not be known to the datastore 

3031 # if constraints are used. 

3032 try: 

3033 primary, uris = datastore.getURIs(ref) 

3034 except FileNotFoundError: 

3035 continue 

3036 if primary and primary.scheme != "mem": 

3037 primary.remove() 

3038 for uri in uris.values(): 

3039 if uri.scheme != "mem": 

3040 uri.remove() 

3041 n_expected -= 1 

3042 deleted.add(ref) 

3043 

3044 # Remove the datastore record. 

3045 for datastore in datastores: 

3046 if hasattr(datastore, "removeStoredItemInfo"): 

3047 datastore.removeStoredItemInfo(ref) 

3048 

3049 if index < 2: 

3050 source_refs.append(ref) 

3051 if ref not in deleted: 

3052 new_metric = butler.get(ref) 

3053 self.assertEqual(new_metric, metric) 

3054 

3055 # Create some bad dataset types to ensure we check for inconsistent 

3056 # definitions. 

3057 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList") 

3058 for datasetTypeName in datasetTypeNames: 

3059 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass) 

3060 self.target_butler.registry.registerDatasetType(datasetType) 

3061 with self.assertRaises(ConflictingDefinitionError) as cm: 

3062 self.target_butler.transfer_from(self.source_butler, source_refs) 

3063 self.assertIn("dataset type differs", str(cm.exception)) 

3064 

3065 # And remove the bad definitions. 

3066 for datasetTypeName in datasetTypeNames: 

3067 self.target_butler.registry.removeDatasetType(datasetTypeName) 

3068 

3069 # Transfer without creating dataset types should fail. 

3070 with self.assertRaises(KeyError): 

3071 self.target_butler.transfer_from(self.source_butler, source_refs) 

3072 

3073 # Transfer without creating dimensions should fail. 

3074 with self.assertRaises(ConflictingDefinitionError) as cm: 

3075 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True) 

3076 self.assertIn("dimension", str(cm.exception)) 

3077 

3078 # The dry run test requires dataset types to exist. If we have 

3079 # been given distinct storage classes for the target we have 

3080 # to redefine at least one of the dataset types in the target butler. 

3081 if storageClass != storageClassTarget: 

3082 self.target_butler.registry.removeDatasetType(datasetTypeNames[0]) 

3083 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget) 

3084 self.target_butler.registry.registerDatasetType(datasetType) 

3085 

3086 # The failed transfer above leaves registry in an inconsistent 

3087 # state because the run is created but then rolled back without 

3088 # the collection cache being cleared. For now force a refresh. 

3089 # Can remove with DM-35498. 

3090 self.target_butler.registry.refresh() 

3091 

3092 # Do a dry run -- this should not have any effect on the target butler. 

3093 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True) 

3094 

3095 # Transfer the records for one ref to test the alternative API. 

3096 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm: 

3097 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]]) 

3098 self.assertIn("number of records transferred: 1", ";".join(log_cm.output)) 

3099 

3100 # Now transfer them to the second butler, including dimensions. 

3101 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm: 

3102 transferred = self.target_butler.transfer_from( 

3103 self.source_butler, 

3104 source_refs, 

3105 register_dataset_types=True, 

3106 transfer_dimensions=True, 

3107 ) 

3108 self.assertEqual(len(transferred), n_expected) 

3109 log_output = ";".join(log_cm.output) 

3110 

3111 # A ChainedDatastore will use the in-memory datastore for mexists 

3112 # so we can not rely on the mexists log message. 

3113 self.assertIn("Number of datastore records found in source", log_output) 

3114 self.assertIn("Creating output run", log_output) 

3115 

3116 # Do the transfer twice to ensure that it will do nothing extra. 

3117 # Only do this if purge=True because it does not work for int 

3118 # dataset_id. 

3119 if purge: 

3120 # This should not need to register dataset types. 

3121 transferred = self.target_butler.transfer_from(self.source_butler, source_refs) 

3122 self.assertEqual(len(transferred), n_expected) 

3123 

3124 with self.assertRaises((TypeError, AttributeError)): 

3125 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore 

3126 

3127 with self.assertRaises(ValueError): 

3128 self.target_butler._datastore.transfer_from( 

3129 self.source_butler._datastore, source_refs, transfer="split" 

3130 ) 

3131 

3132 # Now try to get the same refs from the new butler. 

3133 for ref in source_refs: 

3134 if ref not in deleted: 

3135 new_metric = self.target_butler.get(ref) 

3136 old_metric = self.source_butler.get(ref) 

3137 self.assertEqual(new_metric, old_metric) 

3138 

3139 # Try again without implicit storage class conversion 

3140 # triggered by using the source ref. This will do conversion 

3141 # since the formatter will be returning the source python type. 

3142 target_ref = self.target_butler.get_dataset(ref.id) 

3143 if target_ref.datasetType.storageClass != ref.datasetType.storageClass: 

3144 new_metric = self.target_butler.get(target_ref) 

3145 self.assertNotEqual(type(new_metric), type(old_metric)) 

3146 

3147 # Remove the dataset from the target and put it again 

3148 # as if it was the right type all along for this butler. 

3149 self.target_butler.pruneDatasets( 

3150 [target_ref], unstore=True, purge=True, disassociate=True 

3151 ) 

3152 self.target_butler.put(new_metric, target_ref) 

3153 new_new_metric = self.target_butler.get(target_ref) 

3154 new_old_metric = self.target_butler.get( 

3155 target_ref, storageClass=ref.datasetType.storageClass 

3156 ) 

3157 self.assertEqual(new_new_metric, new_metric) 

3158 self.assertEqual(new_old_metric, old_metric) 

3159 

3160 # Now prune run2 collection and create instead a CHAINED collection. 

3161 # This should block the transfer. 

3162 self.target_butler.removeRuns(["run2"]) 

3163 self.target_butler.collections.register("run2", CollectionType.CHAINED) 

3164 with self.assertRaises(CollectionTypeError): 

3165 # Re-importing the run1 datasets can be problematic if they 

3166 # use integer IDs so filter those out. 

3167 to_transfer = [ref for ref in source_refs if ref.run == "run2"] 

3168 self.target_butler.transfer_from(self.source_butler, to_transfer) 

3169 

3170 

3171class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3172 """Test data transfers between butlers. 

3173 

3174 Test for different managers. UUID to UUID and integer to integer are 

3175 tested. UUID to integer is not supported since we do not currently 

3176 want to allow that. Integer to UUID is supported with the caveat 

3177 that UUID4 will be generated and this will be incorrect for raw 

3178 dataset types. The test ignores that. 

3179 """ 

3180 

3181 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

3182 

3183 def create_butlers( 

3184 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None 

3185 ) -> None: 

3186 self.source_butler = self.create_butler(manager1, "1", config_file=source_config) 

3187 self.target_butler = self.create_butler(manager2, "2") 

3188 

3189 def testTransferUuidToUuid(self) -> None: 

3190 self.create_butlers() 

3191 self.assertButlerTransfers() 

3192 

3193 def testTransferFromChainedUuidToUuid(self) -> None: 

3194 """Force the source butler to be a ChainedDatastore.""" 

3195 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml")) 

3196 self.assertButlerTransfers() 

3197 

3198 def testTransferFromIncompatibleUuidToUuid(self) -> None: 

3199 """Force the source butler to be a incompatible datastore.""" 

3200 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")) 

3201 with self.assertRaises(NotImplementedError): 

3202 self.assertButlerTransfers() 

3203 

3204 def testTransferFromIncompatibleChainUuidToUuid(self) -> None: 

3205 """Force the source butler to be a incompatible datastore.""" 

3206 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml")) 

3207 with self.assertRaises(TypeError): 

3208 self.assertButlerTransfers() 

3209 

3210 def testTransferFromFileUuidToUuid(self) -> None: 

3211 """Force the source butler to be a FileDatastore.""" 

3212 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml")) 

3213 self.assertButlerTransfers() 

3214 

3215 def testTransferMissing(self) -> None: 

3216 """Test transfers where datastore records are missing. 

3217 

3218 This is how execution butler works. 

3219 """ 

3220 self.create_butlers() 

3221 

3222 # Configure the source butler to allow trust. 

3223 self.source_butler._datastore._set_trust_mode(True) 

3224 

3225 self.assertButlerTransfers(purge=True) 

3226 

3227 def testTransferMissingDisassembly(self) -> None: 

3228 """Test transfers where datastore records are missing. 

3229 

3230 This is how execution butler works. 

3231 """ 

3232 self.create_butlers() 

3233 

3234 # Configure the source butler to allow trust. 

3235 self.source_butler._datastore._set_trust_mode(True) 

3236 

3237 # Test disassembly. 

3238 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite") 

3239 

3240 def testTransferDifferingStorageClasses(self) -> None: 

3241 """Test transfers when the source butler dataset type has a different 

3242 but compatible storage class. 

3243 """ 

3244 self.create_butlers() 

3245 

3246 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion") 

3247 

3248 def testTransferDifferingStorageClassesDisassembly(self) -> None: 

3249 """Test transfers when the source butler dataset type has a different 

3250 but compatible storage class and where the source butler has 

3251 disassembled. 

3252 """ 

3253 self.create_butlers() 

3254 

3255 self.assertButlerTransfers( 

3256 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion" 

3257 ) 

3258 

3259 def testUnsafeDirectTransfer(self) -> None: 

3260 """Test that transfer='unsafe_direct' records the absolute URI of 

3261 source files in the target datastore. 

3262 """ 

3263 self.create_butlers() 

3264 dataset_type = DatasetType("dt", [], "int", universe=self.source_butler.dimensions) 

3265 self.source_butler.registry.registerDatasetType(dataset_type) 

3266 self.source_butler.collections.register("run") 

3267 ref = self.source_butler.put(123, "dt", [], run="run") 

3268 self.target_butler.transfer_from( 

3269 self.source_butler, [ref], transfer="unsafe_direct", register_dataset_types=True 

3270 ) 

3271 self.assertEqual(self.target_butler.get(ref), 123) 

3272 self.assertEqual(self.source_butler.getURI(ref), self.target_butler.getURI(ref)) 

3273 

3274 def testAbsoluteURITransferDirect(self) -> None: 

3275 """Test transfer using an absolute URI.""" 

3276 self._absolute_transfer("auto") 

3277 

3278 def testAbsoluteURITransferUnsafeDirect(self) -> None: 

3279 """Test transfer using an absolute URI.""" 

3280 self._absolute_transfer("unsafe_direct") 

3281 

3282 def testAbsoluteURITransferCopy(self) -> None: 

3283 """Test transfer using an absolute URI.""" 

3284 self._absolute_transfer("copy") 

3285 

3286 def _absolute_transfer(self, transfer: str) -> None: 

3287 self.create_butlers() 

3288 

3289 storageClassName = "StructuredData" 

3290 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

3291 datasetTypeName = "random_data" 

3292 run = "run1" 

3293 self.source_butler.collections.register(run) 

3294 

3295 dimensions = self.source_butler.dimensions.conform(()) 

3296 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3297 self.source_butler.registry.registerDatasetType(datasetType) 

3298 

3299 metrics = makeExampleMetrics() 

3300 with ResourcePath.temporary_uri(suffix=".json") as temp: 

3301 dataId = DataCoordinate.make_empty(self.source_butler.dimensions) 

3302 source_refs = [DatasetRef(datasetType, dataId, run=run)] 

3303 temp.write(json.dumps(metrics.exportAsDict()).encode()) 

3304 dataset = FileDataset(path=temp, refs=source_refs) 

3305 self.source_butler.ingest(dataset, transfer="direct") 

3306 

3307 self.target_butler.transfer_from( 

3308 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer 

3309 ) 

3310 

3311 uri = self.target_butler.getURI(dataset.refs[0]) 

3312 if transfer == "auto" or transfer == "unsafe_direct": 

3313 self.assertEqual(uri, temp) 

3314 else: 

3315 self.assertNotEqual(uri, temp) 

3316 

3317 def test_shared_dimension_group(self): 

3318 """Test internal logic that divides dataset types by dimension group 

3319 when doing registry updates. 

3320 """ 

3321 self.create_butlers() 

3322 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True) 

3323 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True) 

3324 

3325 source_butler = self.source_butler 

3326 target_butler = self.target_butler 

3327 

3328 # Create a dataset type with the same dimensions as the 'bias' dataset 

3329 # type from base.yaml 

3330 dataset_type = DatasetType( 

3331 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions 

3332 ) 

3333 source_butler.registry.registerDatasetType(dataset_type) 

3334 # This has the same data ID as one of the bias datasets in 

3335 # datasets.yaml. 

3336 test_ref = source_butler.registry.insertDatasets( 

3337 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g" 

3338 )[0] 

3339 

3340 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"]) 

3341 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"]) 

3342 refs = [test_ref, *biases, *flats] 

3343 

3344 # Test setup will be even more convoluted if we want the datastore to 

3345 # actually transfer files. For testing the dimension group behavior, 

3346 # we really only care about the registry. 

3347 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock: 

3348 mock.return_value = (set(refs), set()) 

3349 target_butler.transfer_from( 

3350 source_butler, 

3351 refs, 

3352 transfer=None, 

3353 register_dataset_types=True, 

3354 skip_missing=False, 

3355 transfer_dimensions=True, 

3356 ) 

3357 

3358 transferred_test_ref = target_butler.find_dataset( 

3359 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3360 ) 

3361 self.assertEqual(transferred_test_ref.id, test_ref.id) 

3362 

3363 transferred_bias = target_butler.find_dataset( 

3364 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3365 ) 

3366 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd")) 

3367 

3368 transferred_flat = target_butler.find_dataset( 

3369 "flat", 

3370 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"}, 

3371 collections="imported_r", 

3372 ) 

3373 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840")) 

3374 

3375 

3376class ChainedDatastoreTransfers(PosixDatastoreTransfers): 

3377 """Test transfers using a chained datastore.""" 

3378 

3379 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml") 

3380 

3381 

3382@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3383class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3384 """Test ``transfer_from`` involving Butler server.""" 

3385 

3386 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

3387 

3388 def test_transfers_from_remote_to_direct(self) -> None: 

3389 from lsst.daf.butler.remote_butler._remote_file_transfer_source import ( 

3390 mock_file_transfer_uris_for_unit_test, 

3391 ) 

3392 

3393 self.target_butler = self.create_butler(None, "2") 

3394 with create_test_server(TESTDIR) as server: 

3395 self.source_butler = server.hybrid_butler 

3396 

3397 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath: 

3398 # The Butler server returns HTTP URIs with a domain name that 

3399 # is not resolvable because there is no actual HTTP server 

3400 # involved in these tests. Strip this first layer of 

3401 # indirection, and return the target of the redirect instead. 

3402 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers) 

3403 return ResourcePath(str(response.next_request.url)) 

3404 

3405 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url): 

3406 self.assertButlerTransfers() 

3407 

3408 

3409class NullDatastoreTestCase(unittest.TestCase): 

3410 """Test that we can fall back to a null datastore.""" 

3411 

3412 # Need a good config to create the repo. 

3413 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

3414 storageClassFactory: StorageClassFactory 

3415 

3416 @classmethod 

3417 def setUpClass(cls) -> None: 

3418 cls.storageClassFactory = StorageClassFactory() 

3419 cls.storageClassFactory.addFromConfig(cls.configFile) 

3420 

3421 def setUp(self) -> None: 

3422 """Create a new butler root for each test.""" 

3423 self.root = makeTestTempDir(TESTDIR) 

3424 Butler.makeRepo(self.root, config=Config(self.configFile)) 

3425 

3426 def tearDown(self) -> None: 

3427 removeTestTempDir(self.root) 

3428 

3429 def test_fallback(self) -> None: 

3430 # Read the butler config and mess with the datastore section. 

3431 config_path = os.path.join(self.root, "butler.yaml") 

3432 bad_config = Config(config_path) 

3433 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore" 

3434 bad_config.dumpToUri(config_path) 

3435 

3436 with self.assertRaises(RuntimeError): 

3437 Butler(self.root, without_datastore=False) 

3438 

3439 with self.assertRaises(RuntimeError): 

3440 Butler.from_config(self.root, without_datastore=False) 

3441 

3442 butler = Butler.from_config(self.root, writeable=True, without_datastore=True) 

3443 self.enterContext(butler) 

3444 self.assertIsInstance(butler._datastore, NullDatastore) 

3445 

3446 # Check that registry is working. 

3447 butler.collections.register("MYRUN") 

3448 collections = butler.collections.query("*") 

3449 self.assertIn("MYRUN", set(collections)) 

3450 

3451 # Create a ref. 

3452 dimensions = butler.dimensions.conform([]) 

3453 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

3454 datasetTypeName = "metric" 

3455 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3456 butler.registry.registerDatasetType(datasetType) 

3457 ref = DatasetRef(datasetType, {}, run="MYRUN") 

3458 

3459 # Check that datastore will complain. 

3460 with self.assertRaises(FileNotFoundError): 

3461 butler.get(ref) 

3462 with self.assertRaises(FileNotFoundError): 

3463 butler.getURI(ref) 

3464 

3465 

3466@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3467class ButlerServerTests(FileDatastoreButlerTests): 

3468 """Test RemoteButler and Butler server.""" 

3469 

3470 configFile = None 

3471 predictionSupported = False 

3472 trustModeSupported = False 

3473 

3474 postgres: TemporaryPostgresInstance | None 

3475 

3476 def setUp(self): 

3477 self.server_instance = self.enterContext(create_test_server(TESTDIR)) 

3478 

3479 def tearDown(self): 

3480 pass 

3481 

3482 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool: 

3483 # S3 pre-signed URLs may end up with differing expiration times in the 

3484 # query parameters, so ignore query parameters when comparing. 

3485 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path 

3486 

3487 def create_empty_butler( 

3488 self, 

3489 run: str | None = None, 

3490 writeable: bool | None = None, 

3491 metrics: ButlerMetrics | None = None, 

3492 cleanup: bool = True, 

3493 ) -> Butler: 

3494 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics) 

3495 

3496 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None: 

3497 # Can't delete a file via S3 signed URLs, so we need to reach in 

3498 # through DirectButler to delete the dataset. 

3499 uri = self.server_instance.direct_butler.getURI(ref) 

3500 uri.remove() 

3501 

3502 def testConstructor(self): 

3503 # RemoteButler constructor is tested in test_server.py and 

3504 # test_remote_butler.py. 

3505 pass 

3506 

3507 def testDafButlerRepositories(self): 

3508 # Loading of RemoteButler via repository index is tested in 

3509 # test_server.py. 

3510 pass 

3511 

3512 def testGetDatasetTypes(self) -> None: 

3513 # This is mostly a test of validateConfiguration, which is for 

3514 # validating Datastore configuration and thus isn't relevant to 

3515 # RemoteButler. 

3516 pass 

3517 

3518 def testMakeRepo(self) -> None: 

3519 # Only applies to DirectButler. 

3520 pass 

3521 

3522 # Pickling not yet implemented for RemoteButler/HybridButler. 

3523 @unittest.expectedFailure 

3524 def testPickle(self) -> None: 

3525 return super().testPickle() 

3526 

3527 def testStringification(self) -> None: 

3528 self.assertEqual( 

3529 str(self.server_instance.remote_butler), 

3530 "RemoteButler(https://test.example/api/butler/repo/testrepo/)", 

3531 ) 

3532 

3533 def testTransaction(self) -> None: 

3534 # Transactions will never be supported for RemoteButler. 

3535 pass 

3536 

3537 def testPutTemplates(self) -> None: 

3538 # The Butler server instance is configured with different file naming 

3539 # templates than this test is expecting. 

3540 pass 

3541 

3542 

3543@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3544class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase): 

3545 """Tests for RemoteButler's registry shim, with a SQLite DB backing the 

3546 server. 

3547 """ 

3548 

3549 postgres = None 

3550 

3551 

3552@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3553class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase): 

3554 """Tests for RemoteButler's registry shim, with a Postgres DB backing the 

3555 server. 

3556 """ 

3557 

3558 @classmethod 

3559 def setUpClass(cls): 

3560 cls.postgres = cls.enterClassContext(setup_postgres_test_db()) 

3561 super().setUpClass() 

3562 

3563 

3564def setup_module(module: types.ModuleType) -> None: 

3565 """Set up the module for pytest.""" 

3566 clean_environment() 

3567 

3568 

3569def _get_test_data_path(filename: str) -> ResourcePath: 

3570 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}") 

3571 

3572 

3573if __name__ == "__main__": 

3574 clean_environment() 

3575 unittest.main()