Coverage for tests / test_butler.py: 13%

1842 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-22 01:49 -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 

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

1052 new_butler_cfg = Butler.makeRepo(tmpdir) 

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

1054 self.enterContext(new_butler) 

1055 

1056 # This will fail since dimensions records are missing. 

1057 with self.assertRaises(ConflictingDefinitionError): 

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

1059 

1060 # Dry run should work. 

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

1062 

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

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

1065 

1066 # Check that the refs can be read again. 

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

1068 

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

1070 self.assertTrue(uri.exists()) 

1071 

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

1073 # remaining refs to be read. 

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

1075 self.assertTrue(uri.exists()) 

1076 

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

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

1079 

1080 butler.removeRuns([self.default_run]) 

1081 self.assertFalse(uri.exists()) 

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

1083 

1084 with self.assertRaises(ValueError): 

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

1086 

1087 def testIngest(self) -> None: 

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

1089 

1090 # Create and register a DatasetType 

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

1092 

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

1094 datasetTypeName = "metric" 

1095 

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

1097 

1098 # Add needed Dimensions 

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

1100 butler.registry.insertDimensionData( 

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

1102 ) 

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

1104 for detector in (1, 2): 

1105 butler.registry.insertDimensionData( 

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

1107 ) 

1108 

1109 butler.registry.insertDimensionData( 

1110 "visit", 

1111 { 

1112 "instrument": "DummyCamComp", 

1113 "id": 423, 

1114 "name": "fourtwentythree", 

1115 "physical_filter": "d-r", 

1116 "day_obs": 20250101, 

1117 }, 

1118 { 

1119 "instrument": "DummyCamComp", 

1120 "id": 424, 

1121 "name": "fourtwentyfour", 

1122 "physical_filter": "d-r", 

1123 "day_obs": 20250101, 

1124 }, 

1125 ) 

1126 

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

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

1129 datasets = [] 

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

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

1132 # required. 

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

1134 for detector in (1, 2): 

1135 detector_name = f"detector_{detector}" 

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

1137 dataId = butler.registry.expandDataId( 

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

1139 ) 

1140 # Create a DatasetRef for ingest 

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

1142 

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

1144 

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

1146 

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

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

1149 

1150 metrics1 = butler.get(datasetTypeName, dataId1) 

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

1152 self.assertNotEqual(metrics1, metrics2) 

1153 

1154 # Compare URIs 

1155 uri1 = butler.getURI(datasetTypeName, dataId1) 

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

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

1158 

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

1160 # skip_existing=False. 

1161 with self.assertRaises(ConflictingDefinitionError): 

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

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

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

1165 

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

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

1168 refs = [] 

1169 for detector in (1, 2): 

1170 detector_name = f"detector_{detector}" 

1171 dataId = butler.registry.expandDataId( 

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

1173 ) 

1174 # Create a DatasetRef for ingest 

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

1176 

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

1178 # have disappeared following ingest. 

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

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

1181 

1182 datasets = [] 

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

1184 

1185 # For first ingest use copy. 

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

1187 

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

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

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

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

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

1193 datasets[0].refs = [ 

1194 cast( 

1195 DatasetRef, 

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

1197 ) 

1198 for ref in datasets[0].refs 

1199 ] 

1200 all_refs = [] 

1201 for dataset in datasets: 

1202 refs = [] 

1203 for ref in dataset.refs: 

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

1205 new_data_id = dict(ref.dataId.required) 

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

1207 assert new_ref is not None 

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

1209 refs.append(new_ref) 

1210 dataset.refs = refs 

1211 all_refs.extend(dataset.refs) 

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

1213 

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

1215 # disable recording of file size. 

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

1217 

1218 # Check that every ref now has records. 

1219 for dataset in datasets: 

1220 for ref in dataset.refs: 

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

1222 

1223 # Ensure that the file has disappeared. 

1224 self.assertFalse(tempFile.exists()) 

1225 

1226 # Check that the datastore recorded no file size. 

1227 # Not all datastores can support this. 

1228 try: 

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

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

1231 except AttributeError: 

1232 pass 

1233 

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

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

1236 

1237 multi1 = butler.get(datasetTypeName, dataId1) 

1238 multi2 = butler.get(datasetTypeName, dataId2) 

1239 

1240 self.assertEqual(multi1, metrics1) 

1241 self.assertEqual(multi2, metrics2) 

1242 

1243 # Compare URIs 

1244 uri1 = butler.getURI(datasetTypeName, dataId1) 

1245 uri2 = butler.getURI(datasetTypeName, dataId2) 

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

1247 

1248 # Test that removing one does not break the second 

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

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

1251 # files. 

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

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

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

1255 multi2b = butler.get(datasetTypeName, dataId2) 

1256 self.assertEqual(multi2, multi2b) 

1257 

1258 # Ensure we can ingest 0 datasets 

1259 datasets = [] 

1260 butler.ingest(*datasets) 

1261 

1262 def testPickle(self) -> None: 

1263 """Test pickle support.""" 

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

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

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

1267 self.enterContext(butlerOut) 

1268 self.assertIsInstance(butlerOut, Butler) 

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

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

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

1272 

1273 def testGetDatasetTypes(self) -> None: 

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

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

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

1277 ( 

1278 "instrument", 

1279 [ 

1280 {"instrument": "DummyCam"}, 

1281 {"instrument": "DummyHSC"}, 

1282 {"instrument": "DummyCamComp"}, 

1283 ], 

1284 ), 

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

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

1287 ( 

1288 "visit", 

1289 [ 

1290 { 

1291 "instrument": "DummyCam", 

1292 "id": 42, 

1293 "name": "fortytwo", 

1294 "physical_filter": "d-r", 

1295 "day_obs": 20250101, 

1296 } 

1297 ], 

1298 ), 

1299 ] 

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

1301 # Add needed Dimensions 

1302 for element, data in dimensionEntries: 

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

1304 

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

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

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

1308 components = set() 

1309 for datasetTypeName in datasetTypeNames: 

1310 # Create and register a DatasetType 

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

1312 

1313 for componentName in storageClass.components: 

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

1315 

1316 fromRegistry: set[DatasetType] = set() 

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

1318 fromRegistry.add(parent_dataset_type) 

1319 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes()) 

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

1321 

1322 # Query with wildcard. 

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

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

1325 # but not regex. 

1326 with self.assertRaises(DatasetTypeExpressionError): 

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

1328 

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

1330 butler.validateConfiguration( 

1331 ignore=[ 

1332 "test_metric_comp", 

1333 "metric3", 

1334 "metric5", 

1335 "calexp", 

1336 "DummySC", 

1337 "datasetType.component", 

1338 "random_data", 

1339 "random_data_2", 

1340 ] 

1341 ) 

1342 

1343 # Add a new datasetType that will fail template validation 

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

1345 if self.validationCanFail: 

1346 with self.assertRaises(ValidationError): 

1347 butler.validateConfiguration() 

1348 

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

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

1351 

1352 # Rerun validation but ignore the bad datasetType 

1353 butler.validateConfiguration( 

1354 ignore=[ 

1355 "test_metric_comp", 

1356 "metric3", 

1357 "metric5", 

1358 "calexp", 

1359 "DummySC", 

1360 "datasetType.component", 

1361 "random_data", 

1362 "random_data_2", 

1363 ] 

1364 ) 

1365 

1366 def testTransaction(self) -> None: 

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

1368 datasetTypeName = "test_metric" 

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

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

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

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

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

1374 ( 

1375 "visit", 

1376 { 

1377 "instrument": "DummyCam", 

1378 "id": 42, 

1379 "name": "fortytwo", 

1380 "physical_filter": "d-r", 

1381 "day_obs": 20250101, 

1382 }, 

1383 ), 

1384 ) 

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

1386 metric = makeExampleMetrics() 

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

1388 # Create and register a DatasetType 

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

1390 with self.assertRaises(TransactionTestError): 

1391 with butler.transaction(): 

1392 # Add needed Dimensions 

1393 for args in dimensionEntries: 

1394 butler.registry.insertDimensionData(*args) 

1395 # Store a dataset 

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

1397 self.assertIsInstance(ref, DatasetRef) 

1398 # Test get of a ref. 

1399 metricOut = butler.get(ref) 

1400 self.assertEqual(metric, metricOut) 

1401 # Test get 

1402 metricOut = butler.get(datasetTypeName, dataId) 

1403 self.assertEqual(metric, metricOut) 

1404 # Check we can get components 

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

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

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

1408 butler.registry.expandDataId(dataId) 

1409 # Should raise LookupError for missing data ID value 

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

1411 butler.get(datasetTypeName, dataId) 

1412 # Also check explicitly if Dataset entry is missing 

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

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

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

1416 butler.get(ref) 

1417 

1418 def testMakeRepo(self) -> None: 

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

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

1421 repo root. 

1422 """ 

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

1424 # not support a file system root 

1425 if self.fullConfigKey is None: 

1426 return 

1427 

1428 # create two separate directories 

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

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

1431 

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

1433 limited = Config(self.configFile) 

1434 butler1 = Butler.from_config(butlerConfig) 

1435 self.enterContext(butler1) 

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

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

1438 full = Config(self.tmpConfigFile) 

1439 butler2 = Butler.from_config(butlerConfig) 

1440 self.enterContext(butler2) 

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

1442 # Butlers should have the same configuration regardless of whether 

1443 # defaults were expanded. 

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

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

1446 self.assertNotEqual(limited, full) 

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

1448 # inheriting from defaults. 

1449 self.assertIn(self.fullConfigKey, full) 

1450 self.assertNotIn(self.fullConfigKey, limited) 

1451 

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

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

1454 self.assertEqual(collections1, set()) 

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

1456 

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

1458 # work properly with relocatable Butler repo 

1459 butlerConfig.configFile = None 

1460 with self.assertRaises(ValueError): 

1461 Butler.from_config(butlerConfig) 

1462 

1463 with self.assertRaises(FileExistsError): 

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

1465 

1466 def testStringification(self) -> None: 

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

1468 self.enterContext(butler) 

1469 butlerStr = str(butler) 

1470 

1471 if self.datastoreStr is not None: 

1472 for testStr in self.datastoreStr: 

1473 self.assertIn(testStr, butlerStr) 

1474 if self.registryStr is not None: 

1475 self.assertIn(self.registryStr, butlerStr) 

1476 

1477 datastoreName = butler._datastore.name 

1478 if self.datastoreName is not None: 

1479 for testStr in self.datastoreName: 

1480 self.assertIn(testStr, datastoreName) 

1481 

1482 def testButlerRewriteDataId(self) -> None: 

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

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

1485 

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

1487 datasetTypeName = "random_data" 

1488 

1489 # Create dimension records. 

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

1491 butler.registry.insertDimensionData( 

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

1493 ) 

1494 butler.registry.insertDimensionData( 

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

1496 ) 

1497 

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

1499 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

1500 butler.registry.registerDatasetType(datasetType) 

1501 

1502 n_exposures = 5 

1503 dayobs = 20210530 

1504 

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

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

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

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

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

1510 

1511 for i in range(n_exposures): 

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

1513 butler.registry.insertDimensionData( 

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

1515 ) 

1516 butler.registry.insertDimensionData( 

1517 "exposure", 

1518 { 

1519 "instrument": "DummyCamComp", 

1520 "id": day_obs + i, 

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

1522 "seq_num": i, 

1523 "day_obs": day_obs, 

1524 "physical_filter": "d-r", 

1525 "group": group_name, 

1526 }, 

1527 ) 

1528 

1529 # Write some data. 

1530 for i in range(n_exposures): 

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

1532 

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

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

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

1536 

1537 # Check that the exposure is correct in the dataId 

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

1539 

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

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

1542 self.assertEqual(new_metric, metric) 

1543 

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

1545 # exposure.day_obs. 

1546 datasets_1 = list( 

1547 butler.registry.queryDatasets( 

1548 datasetType, 

1549 collections=self.default_run, 

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

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

1552 ) 

1553 ) 

1554 datasets_2 = list( 

1555 butler.registry.queryDatasets( 

1556 datasetType, 

1557 collections=self.default_run, 

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

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

1560 ) 

1561 ) 

1562 self.assertEqual(datasets_1, datasets_2) 

1563 

1564 def testGetDatasetCollectionCaching(self): 

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

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

1567 # after the collection cache was last updated. 

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

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

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

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

1572 get_ref = reader_butler.get_dataset(put_ref.id) 

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

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

1575 # instance. 

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

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

1578 

1579 def testCollectionChainRedefine(self): 

1580 butler = self._setup_to_test_collection_chain() 

1581 

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

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

1584 

1585 # Duplicates are removed from the list of children 

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

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

1588 

1589 # Empty list clears the chain 

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

1591 self._check_chain(butler, []) 

1592 

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

1594 

1595 def testCollectionChainPrepend(self): 

1596 butler = self._setup_to_test_collection_chain() 

1597 

1598 # Duplicates are removed from the list of children 

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

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

1601 

1602 # Prepend goes on the front of existing chain 

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

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

1605 

1606 # Empty prepend does nothing 

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

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

1609 

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

1611 # their current position. 

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

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

1614 

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

1616 

1617 def testCollectionChainExtend(self): 

1618 butler = self._setup_to_test_collection_chain() 

1619 

1620 # Duplicates are removed from the list of children 

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

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

1623 

1624 # Extend goes on the end of existing chain 

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

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

1627 

1628 # Empty extend does nothing 

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

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

1631 

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

1633 # their current position. 

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

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

1636 

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

1638 

1639 def testCollectionChainRemove(self) -> None: 

1640 butler = self._setup_to_test_collection_chain() 

1641 

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

1643 

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

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

1646 

1647 # Duplicates are allowed in the list of children 

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

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

1650 

1651 # Empty remove does nothing 

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

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

1654 

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

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

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

1658 

1659 self._test_common_chain_functionality( 

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

1661 ) 

1662 

1663 def _setup_to_test_collection_chain(self) -> Butler: 

1664 butler = self.create_empty_butler(writeable=True) 

1665 

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

1667 

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

1669 for run in runs: 

1670 butler.collections.register(run) 

1671 

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

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

1674 

1675 return butler 

1676 

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

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

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

1680 

1681 def _test_common_chain_functionality( 

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

1683 ) -> None: 

1684 # Missing parent collection 

1685 with self.assertRaises(MissingCollectionError): 

1686 func("doesnotexist", []) 

1687 # Missing child collection 

1688 with self.assertRaises(MissingCollectionError): 

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

1690 # Forbid operations on non-chained collections 

1691 with self.assertRaises(CollectionTypeError): 

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

1693 

1694 # Prevent collection cycles 

1695 if not skip_cycle_check: 

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

1697 func("chain2", "chain") 

1698 with self.assertRaises(CollectionCycleError): 

1699 func("chain", "chain2") 

1700 

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

1702 # chains. 

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

1704 

1705 with butler._caching_context(): 

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

1707 func("chain", "a") 

1708 

1709 def test_transfer_dimension_records_from(self) -> None: 

1710 source_butler = self.create_empty_butler(writeable=True) 

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

1712 

1713 visit_id = 2025120200439 

1714 exposure_id = visit_id 

1715 target_butler = self.enterContext(create_populated_sqlite_registry()) 

1716 target_butler.transfer_dimension_records_from( 

1717 source_butler, 

1718 [ 

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

1720 # "populated_by" records (visit_detector_region, 

1721 # visit_definition, etc.) 

1722 DataCoordinate.standardize( 

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

1724 universe=source_butler.dimensions, 

1725 ), 

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

1727 DataCoordinate.make_empty(source_butler.dimensions), 

1728 ], 

1729 ) 

1730 

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

1732 records = target_butler.query_dimension_records(dimension) 

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

1734 return records[0] 

1735 

1736 visit = _fetch_record("visit") 

1737 self.assertEqual(visit.id, visit_id) 

1738 self.assertEqual(visit.day_obs, 20251202) 

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

1740 self.assertEqual(visit.seq_num, 439) 

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

1742 0 

1743 ] 

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

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

1746 

1747 visit_detector_region = _fetch_record("visit_detector_region") 

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

1749 self.assertEqual(visit_detector_region.detector, 10) 

1750 self.assertEqual(visit_detector_region.visit, visit_id) 

1751 original_visit_detector_region = source_butler.query_dimension_records( 

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

1753 )[0] 

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

1755 

1756 visit_definition = _fetch_record("visit_definition") 

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

1758 self.assertEqual(visit_definition.exposure, 2025120200439) 

1759 self.assertEqual(visit_definition.visit, visit_id) 

1760 

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

1762 # visit -> visit_definition. 

1763 exposure = _fetch_record("exposure") 

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

1765 self.assertEqual(exposure.id, 2025120200439) 

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

1767 original_exposure = source_butler.query_dimension_records( 

1768 "exposure", instrument="LSSTCam", exposure=exposure_id 

1769 )[0] 

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

1771 

1772 group = _fetch_record("group") 

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

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

1775 

1776 visit_system_memberships = target_butler.query_dimension_records("visit_system_membership") 

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

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

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

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

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

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

1783 

1784 visit_systems = target_butler.query_dimension_records("visit_system") 

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

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

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

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

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

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

1791 

1792 

1793class FileDatastoreButlerTests(ButlerTests): 

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

1795 by datastores that inherit from FileDatastore. 

1796 """ 

1797 

1798 trustModeSupported = True 

1799 

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

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

1802 

1803 Test testPutTemplates verifies actual physical existance of the files 

1804 in the requested location. 

1805 """ 

1806 uri = ResourcePath(root, forceDirectory=True) 

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

1808 

1809 def testPutTemplates(self) -> None: 

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

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

1812 

1813 # Add needed Dimensions 

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

1815 butler.registry.insertDimensionData( 

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

1817 ) 

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

1819 butler.registry.insertDimensionData( 

1820 "visit", 

1821 { 

1822 "instrument": "DummyCamComp", 

1823 "id": 423, 

1824 "name": "v423", 

1825 "physical_filter": "d-r", 

1826 "day_obs": 20250101, 

1827 }, 

1828 ) 

1829 butler.registry.insertDimensionData( 

1830 "visit", 

1831 { 

1832 "instrument": "DummyCamComp", 

1833 "id": 425, 

1834 "name": "v425", 

1835 "physical_filter": "d-r", 

1836 "day_obs": 20250101, 

1837 }, 

1838 ) 

1839 

1840 # Create and store a dataset 

1841 metric = makeExampleMetrics() 

1842 

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

1844 # template) 

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

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

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

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

1849 

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

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

1852 

1853 # Put with exactly the data ID keys needed 

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

1855 uri = butler.getURI(ref) 

1856 self.assertTrue(uri.exists()) 

1857 self.assertTrue( 

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

1859 ) 

1860 

1861 # Check the template based on dimensions 

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

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

1864 

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

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

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

1868 # must be consistent). 

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

1870 uri = butler.getURI(ref) 

1871 self.assertTrue(uri.exists()) 

1872 self.assertTrue( 

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

1874 ) 

1875 

1876 # Check the template based on dimensions 

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

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

1879 

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

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

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

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

1884 path = template.format(ref) 

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

1886 

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

1888 with self.assertRaises(KeyError): 

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

1890 template.format(ref) 

1891 

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

1893 with self.assertRaises(FileTemplateValidationError): 

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

1895 

1896 def testImportExport(self) -> None: 

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

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

1899 self.runImportExportTest(storageClass) 

1900 

1901 @unittest.expectedFailure 

1902 def testImportExportVirtualComposite(self) -> None: 

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

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

1905 self.runImportExportTest(storageClass) 

1906 

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

1908 """Test exporting and importing. 

1909 

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

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

1912 """ 

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

1914 

1915 # Test that we must have a file extension. 

1916 with self.assertRaises(ValueError): 

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

1918 pass 

1919 

1920 # Test that unknown format is not allowed. 

1921 with self.assertRaises(ValueError): 

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

1923 pass 

1924 

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

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

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

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

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

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

1931 # Export and then import datasets. 

1932 with safeTestTempDir(TESTDIR) as exportDir: 

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

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

1935 export.saveDatasets(datasets) 

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

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

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

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

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

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

1942 # because of internal deduplication. 

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

1944 # Save some dimension records directly. 

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

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

1947 with safeTestTempDir(TESTDIR) as importDir: 

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

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

1950 # Calling script.butlerImport tests the implementation of the 

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

1952 # in the script folder are generally considered protected and 

1953 # should not be used as public api. 

1954 with open(exportFile) as f: 

1955 script.butlerImport( 

1956 importDir, 

1957 export_file=f, 

1958 directory=exportDir, 

1959 transfer="auto", 

1960 skip_dimensions=None, 

1961 ) 

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

1963 self.enterContext(importButler) 

1964 for ref in datasets: 

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

1966 # Test for existence by passing in the DatasetType and 

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

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

1969 self.assertEqual( 

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

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

1972 ) 

1973 

1974 def testRemoveRuns(self) -> None: 

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

1976 butler = self.create_empty_butler(writeable=True) 

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

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

1979 # Add some RUN-type collection. 

1980 run1 = "run1" 

1981 butler.collections.register(run1) 

1982 run2 = "run2" 

1983 butler.collections.register(run2) 

1984 # put a dataset in each 

1985 metric = makeExampleMetrics() 

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

1987 datasetType = self.addDatasetType( 

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

1989 ) 

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

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

1992 uri1 = butler.getURI(ref1) 

1993 uri2 = butler.getURI(ref2) 

1994 

1995 # Put one of the runs in a chain. 

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

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

1998 

1999 with self.assertRaises(OrphanedRecordError): 

2000 butler.registry.removeDatasetType(datasetType.name) 

2001 

2002 # Remove a non-run. 

2003 with self.assertRaises(TypeError): 

2004 butler.removeRuns(["Chain"]) 

2005 

2006 # Remove without unlinking from chain should fail. 

2007 with self.assertRaises(IntegrityError): 

2008 butler.removeRuns([run1]) 

2009 

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

2011 # always purges. 

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

2013 

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

2015 # not think either exists. 

2016 with self.assertRaises(MissingCollectionError): 

2017 butler.collections.get_info(run1) 

2018 with self.assertRaises(MissingCollectionError): 

2019 butler.collections.get_info(run1) 

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

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

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

2023 self.assertFalse(uri1.exists()) 

2024 self.assertFalse(uri2.exists()) 

2025 

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

2027 # dataset type 

2028 butler.registry.removeDatasetType(datasetType.name) 

2029 

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

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

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

2033 

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

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

2036 knowledge. 

2037 

2038 Subclasses may override to handle more complicated datastore 

2039 configurations. 

2040 """ 

2041 uri = butler.getURI(ref) 

2042 uri.remove() 

2043 datastore = cast(FileDatastore, butler._datastore) 

2044 datastore.cacheManager.remove_from_cache(ref) 

2045 

2046 def testPruneDatasets(self) -> None: 

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

2048 butler = self.create_empty_butler(writeable=True) 

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

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

2051 # Add some RUN-type collections. 

2052 run1 = "run1" 

2053 butler.collections.register(run1) 

2054 run2 = "run2" 

2055 butler.collections.register(run2) 

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

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

2058 metric = makeExampleMetrics() 

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

2060 datasetType = self.addDatasetType( 

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

2062 ) 

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

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

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

2066 

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

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

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

2070 

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

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

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

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

2075 

2076 # Simple prune. 

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

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

2079 

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

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

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

2083 

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

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

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

2087 

2088 # Put data back. 

2089 ref1_new = butler.put(metric, ref1) 

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

2091 ref2 = butler.put(metric, ref2) 

2092 

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

2094 self.assertTrue(many_stored[ref1]) 

2095 self.assertTrue(many_stored[ref2]) 

2096 self.assertFalse(many_stored[ref3]) 

2097 

2098 ref3 = butler.put(metric, ref3) 

2099 

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

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

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

2103 

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

2105 refs = [ref1, ref2, ref3] 

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

2107 for ref in refs: 

2108 butler.put(metric, ref) 

2109 

2110 # Confirm we can retrieve deferred. 

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

2112 metric1 = dref1.get() 

2113 self.assertEqual(metric1, metric) 

2114 

2115 # Test different forms of file availability. 

2116 # Need to be in a state where: 

2117 # - one ref just has registry record. 

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

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

2120 # - one ref does not exist anywhere. 

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

2122 # above. 

2123 ref0 = DatasetRef( 

2124 datasetType, 

2125 DataCoordinate.standardize( 

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

2127 ), 

2128 run=run1, 

2129 ) 

2130 

2131 # Delete from datastore and retain in Registry. 

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

2133 

2134 # File has been removed. 

2135 self.remove_dataset_out_of_band(butler, ref2) 

2136 

2137 # Datastore has lost track. 

2138 butler._datastore.forget([ref3]) 

2139 

2140 # First test with a standard butler. 

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

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

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

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

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

2146 

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

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

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

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

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

2152 self.assertTrue(exists_many[ref2]) 

2153 

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

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

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

2157 

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

2159 # retrieved. 

2160 with self.assertRaises(LookupError): 

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

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

2163 with self.assertRaises(FileNotFoundError): 

2164 dref2.get() 

2165 

2166 # Test again with a trusting butler. 

2167 if self.trustModeSupported: 

2168 butler._datastore.trustGetRequest = True 

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

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

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

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

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

2174 

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

2176 # known but does exist. 

2177 dref3 = butler.getDeferred(ref3) 

2178 metric3 = dref3.get() 

2179 self.assertEqual(metric3, metric) 

2180 

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

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

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

2184 

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

2186 # but is not the same. 

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

2188 with self.assertRaises(ValueError): 

2189 butler.exists(ref_bad) 

2190 

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

2192 ref_compat = ref2.overrideStorageClass("StructuredDataDict") 

2193 exists = butler.exists(ref_compat) 

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

2195 

2196 # Remove everything and start from scratch. 

2197 butler._datastore.trustGetRequest = False 

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

2199 for ref in refs: 

2200 butler.put(metric, ref) 

2201 

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

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

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

2205 # trash not touching the file. 

2206 uri1 = butler.getURI(ref1) 

2207 butler._datastore.bridge.moveToTrash( 

2208 [ref1], transaction=None 

2209 ) # Update the dataset_location table 

2210 butler._datastore.forget([ref1]) 

2211 butler._datastore.trash(ref1) 

2212 butler._datastore.emptyTrash() 

2213 self.assertTrue(uri1.exists()) 

2214 uri1.remove() # Clean it up. 

2215 

2216 # Simulate execution butler setup by deleting the datastore 

2217 # record but keeping the file around and trusting. 

2218 butler._datastore.trustGetRequest = True 

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

2220 uri2 = uris[ref2].primaryURI 

2221 uri3 = uris[ref3].primaryURI 

2222 self.assertTrue(uri2.exists()) 

2223 self.assertTrue(uri3.exists()) 

2224 

2225 # Remove the datastore record. 

2226 butler._datastore.bridge.moveToTrash( 

2227 [ref2], transaction=None 

2228 ) # Update the dataset_location table 

2229 butler._datastore.forget([ref2]) 

2230 self.assertTrue(uri2.exists()) 

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

2232 # Immediate removal for ref2 file 

2233 self.assertFalse(uri2.exists()) 

2234 # But ref3 has to wait for the empty. 

2235 self.assertTrue(uri3.exists()) 

2236 butler._datastore.emptyTrash() 

2237 self.assertFalse(uri3.exists()) 

2238 

2239 # Clear out the datasets from registry. 

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

2241 

2242 def test_butler_metrics(self): 

2243 """Test that metrics are collected.""" 

2244 run = "test_run" 

2245 metrics = ButlerMetrics() 

2246 butler, datasetType = self.create_butler( 

2247 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics 

2248 ) 

2249 data = MetricsExampleModel( 

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

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

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

2253 ) 

2254 

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

2256 butler.get(data_ref) 

2257 butler.get(data_ref) 

2258 self.assertEqual(metrics.n_get, 2) 

2259 self.assertGreater(metrics.time_in_get, 0.0) 

2260 self.assertEqual(metrics.n_put, 1) 

2261 self.assertGreater(metrics.time_in_put, 0.0) 

2262 

2263 deferred = butler.getDeferred(data_ref) 

2264 deferred.get() 

2265 self.assertEqual(metrics.n_get, 3) 

2266 

2267 with butler.record_metrics() as new: 

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

2269 butler.get(data_ref) 

2270 

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

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

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

2274 refs = [ 

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

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

2277 ] 

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

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

2280 

2281 self.assertEqual(new.n_get, 1) 

2282 self.assertEqual(new.n_put, 1) 

2283 self.assertEqual(new.n_ingest, 2) 

2284 

2285 

2286class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2287 """PosixDatastore specialization of a butler""" 

2288 

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

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

2291 validationCanFail = True 

2292 datastoreStr = ["/tmp"] 

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

2294 registryStr = "/gen3.sqlite3" 

2295 

2296 def testPathConstructor(self) -> None: 

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

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

2299 self.enterContext(butler) 

2300 self.assertIsInstance(butler, Butler) 

2301 

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

2303 path = pathlib.Path(self.tmpConfigFile) 

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

2305 self.enterContext(butler) 

2306 self.assertIsInstance(butler, Butler) 

2307 

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

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

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

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

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

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

2314 self.enterContext(butler) 

2315 self.assertIsInstance(butler, Butler) 

2316 

2317 def testExportTransferCopy(self) -> None: 

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

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

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

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

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

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

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

2325 assert isinstance(exportButler._datastore, FileDatastore) 

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

2327 

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

2329 

2330 for path in pathsInStore: 

2331 # Assume local file system 

2332 assert path is not None 

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

2334 

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

2336 with safeTestTempDir(TESTDIR) as exportDir: 

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

2338 export.saveDatasets(datasets) 

2339 for path in pathsInStore: 

2340 assert path is not None 

2341 self.assertTrue( 

2342 self.checkFileExists(exportDir, path), 

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

2344 ) 

2345 

2346 def testPytypeCoercion(self) -> None: 

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

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

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

2350 datasetTypeName = "test_metric" 

2351 butler = self.runPutGetTest(storageClass, datasetTypeName) 

2352 

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

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

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

2356 

2357 datasetType_ori = butler.get_dataset_type(datasetTypeName) 

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

2359 

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

2361 # There is no API for this. 

2362 assert isinstance(butler._registry, SqlRegistry) 

2363 manager = butler._registry._managers.datasets 

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

2365 manager._db.update( 

2366 manager._static.dataset_type, 

2367 {"name": datasetTypeName}, 

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

2369 ) 

2370 

2371 # Force reset of dataset type cache 

2372 butler.registry.refresh() 

2373 

2374 datasetType_new = butler.get_dataset_type(datasetTypeName) 

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

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

2377 

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

2379 self.assertNotEqual(type(metric_model), type(metric)) 

2380 self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel") 

2381 

2382 # Put the model and read it back to show that everything now 

2383 # works as normal. 

2384 metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424) 

2385 metric_model_new = butler.get(metric_ref) 

2386 self.assertEqual(metric_model_new, metric_model) 

2387 

2388 # Hack the storage class again to something that will fail on the 

2389 # get with no conversion class. 

2390 manager._db.update( 

2391 manager._static.dataset_type, 

2392 {"name": datasetTypeName}, 

2393 {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"}, 

2394 ) 

2395 butler.registry.refresh() 

2396 

2397 with self.assertRaises(ValueError): 

2398 butler.get(datasetTypeName, dataId=dataId) 

2399 

2400 def test_provenance(self): 

2401 """Test that provenance is attached on put.""" 

2402 run = "test_run" 

2403 butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric") 

2404 metric = MetricsExampleModel( 

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

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

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

2408 ) 

2409 # Provenance can be attached to the object being put. Whether 

2410 # it is or not is dependent on the formatter. For this test we 

2411 # copy on adding provenance to ensure they differ. 

2412 self.assertIsNone(metric.dataset_id) 

2413 metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp") 

2414 self.assertIsNone(metric.dataset_id) 

2415 metric_2 = butler.get(metric_ref) 

2416 self.assertEqual(metric_2.data, metric.data) 

2417 self.assertEqual(metric_2.dataset_id, metric_ref.id) 

2418 self.assertIsNone(metric_2.provenance) 

2419 

2420 # Put with provenance. 

2421 prov = DatasetProvenance(quantum_id=uuid.uuid4()) 

2422 prov.add_input(metric_ref) 

2423 prov.add_extra_provenance(metric_ref.id, {"answer": 42}) 

2424 metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov) 

2425 metric_3 = butler.get(metric_ref2) 

2426 self.assertEqual(metric_3.provenance, prov) 

2427 

2428 # Check that we can extract provenance from dict form. 

2429 prov_dict = prov.to_flat_dict(metric_ref2) 

2430 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler) 

2431 self.assertEqual(ref_from_prov, metric_ref2) 

2432 # Direct __eq__ of the provenance does not work because one side 

2433 # includes dimension records. 

2434 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs}) 

2435 self.assertEqual(prov_from_prov.quantum_id, prov.quantum_id) 

2436 self.assertEqual(prov_from_prov.extras, prov.extras) 

2437 

2438 # Force a bad ID into the dict. 

2439 prov_dict["id"] = uuid.uuid4() 

2440 with self.assertRaises(ValueError): 

2441 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2442 del prov_dict["id"] 

2443 prov_dict["input 0 id"] = uuid.uuid4() 

2444 with self.assertRaises(ValueError): 

2445 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2446 

2447 # Check that simple types can be reconstructed with non-standard 

2448 # separators. 

2449 prov_dict = prov.to_flat_dict(metric_ref2, prefix="XYZ", sep="😎", simple_types=True) 

2450 prov_from_prov, ref_from_prov = DatasetProvenance.from_flat_dict(prov_dict, butler) 

2451 self.assertEqual(ref_from_prov, metric_ref2) 

2452 self.assertEqual({ref.id for ref in prov_from_prov.inputs}, {ref.id for ref in prov.inputs}) 

2453 

2454 with self.assertRaises(ValueError): 

2455 DatasetProvenance.from_flat_dict({"unknown": 42}, butler) 

2456 

2457 def test_specialized_file_datasets_functions(self): 

2458 """Test a workflow used in Prompt Processing where we export datasets 

2459 from one repository and write them in-place to the datastore of 

2460 another, without immediately inserting registry entries for the 

2461 datasets. 

2462 """ 

2463 repo = MetricTestRepo.create_from_butler( 

2464 self.create_empty_butler(writeable=True), 

2465 self.tmpConfigFile, 

2466 "StructuredCompositeReadCompNoDisassembly", 

2467 ) 

2468 source_butler = repo.butler 

2469 

2470 # Test writing outputs to a FileDatastore. 

2471 with tempfile.TemporaryDirectory() as tempdir: 

2472 target_repo_config = Butler.makeRepo(tempdir) 

2473 refs = [repo.ref1, repo.ref2] 

2474 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs) 

2475 self.assertEqual(len(datasets), 2) 

2476 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets}) 

2477 for dataset in datasets: 

2478 path = ResourcePath(dataset.path, forceAbsolute=False) 

2479 # Paths should be relative paths to the target datastore. 

2480 self.assertFalse(path.isabs()) 

2481 # Files should have been copied into the target datastore 

2482 self.assertTrue(ResourcePath(tempdir).join(path).exists()) 

2483 

2484 # Make sure the target Butler can ingest the datasets. 

2485 target_butler = Butler(target_repo_config, writeable=True) 

2486 self.enterContext(target_butler) 

2487 target_butler.transfer_dimension_records_from(source_butler, refs) 

2488 target_butler.ingest(*datasets, transfer=None) 

2489 self.assertIsNotNone(target_butler.get(repo.ref1)) 

2490 self.assertIsNotNone(target_butler.get(repo.ref2)) 

2491 

2492 # Giving an empty list of files is a no-op. 

2493 no_datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), []) 

2494 self.assertEqual(len(no_datasets), 0) 

2495 

2496 # Test writing outputs to a ChainedDatastore. 

2497 with tempfile.TemporaryDirectory() as tempdir: 

2498 # Set up a second dataset type, so we can split the files across 

2499 # multiple datastore roots. 

2500 dt1 = repo.datasetType 

2501 dt2 = DatasetType("other", dt1.dimensions, dt1.storageClass) 

2502 source_butler.registry.registerDatasetType(dt2) 

2503 other_ref = repo.addDataset(repo.ref1.dataId, datasetType=dt2) 

2504 config = Config.fromString( 

2505 f""" 

2506 datastore: 

2507 cls: lsst.daf.butler.datastores.chainedDatastore.ChainedDatastore 

2508 datastore_constraints: 

2509 - constraints: 

2510 accept: 

2511 - {dt1.name} 

2512 - constraints: 

2513 accept: 

2514 - {dt2.name} 

2515 datastores: 

2516 - datastore: 

2517 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore 

2518 root: <butlerRoot>/FileDatastore_0 

2519 - datastore: 

2520 cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore 

2521 root: <butlerRoot>/FileDatastore_1 

2522 """ 

2523 ) 

2524 target_repo_config = Butler.makeRepo(tempdir, config) 

2525 refs = [repo.ref1, repo.ref2, other_ref] 

2526 datasets = transfer_datasets_to_datastore(source_butler, ButlerConfig(target_repo_config), refs) 

2527 self.assertEqual(len(datasets), 3) 

2528 self.assertEqual({ref.id for ref in refs}, {dataset.refs[0].id for dataset in datasets}) 

2529 for dataset in datasets: 

2530 path = ResourcePath(dataset.path, forceAbsolute=False) 

2531 # Paths should be relative paths to the target datastore. 

2532 self.assertFalse(path.isabs()) 

2533 # Files should have been split up between the two datastores 

2534 # in the chain. 

2535 datastore_root = ResourcePath(tempdir) 

2536 if dataset.refs[0].datasetType.name == dt1.name: 

2537 datastore_root = datastore_root.join("FileDatastore_0") 

2538 else: 

2539 datastore_root = datastore_root.join("FileDatastore_1") 

2540 self.assertTrue(datastore_root.join(path).exists()) 

2541 

2542 # Make sure the target Butler can ingest the datasets. 

2543 target_butler = Butler(target_repo_config, writeable=True) 

2544 self.enterContext(target_butler) 

2545 target_butler.transfer_dimension_records_from(source_butler, refs) 

2546 target_butler.ingest(*datasets, transfer=None) 

2547 self.assertIsNotNone(target_butler.get(repo.ref1)) 

2548 self.assertIsNotNone(target_butler.get(repo.ref2)) 

2549 self.assertIsNotNone(target_butler.get(other_ref)) 

2550 

2551 def test_temporary_for_ingest(self) -> None: 

2552 """Test the `lsst.daf.butler._rubin.ingest_from_temporary` module.""" 

2553 with self.create_empty_butler("example_run") as butler: 

2554 dataset_type = DatasetType("example", butler.dimensions.empty, "StructuredDataDict") 

2555 butler.registry.registerDatasetType(dataset_type) 

2556 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(butler.dimensions), "example_run") 

2557 with TemporaryForIngest(butler, ref) as temporary: 

2558 temporary.path.write(b"three: 3") 

2559 found = TemporaryForIngest.find_orphaned_temporaries_by_ref(ref, butler) 

2560 self.assertEqual(found, [temporary.path]) 

2561 self.assertIn(".tmp", temporary.ospath) 

2562 temporary.ingest() 

2563 loaded = butler.get(ref) 

2564 self.assertEqual(loaded, {"three": 3}) 

2565 

2566 

2567class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2568 """PosixDatastore specialization of a butler using Postgres""" 

2569 

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

2571 fullConfigKey = ".datastore.formatters" 

2572 validationCanFail = True 

2573 datastoreStr = ["/tmp"] 

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

2575 registryStr = "PostgreSQL@test" 

2576 

2577 @classmethod 

2578 def setUpClass(cls) -> None: 

2579 cls.postgresql = cls.enterClassContext(setup_postgres_test_db()) 

2580 super().setUpClass() 

2581 

2582 def setUp(self) -> None: 

2583 # Need to add a registry section to the config. 

2584 self._temp_config = False 

2585 config = Config(self.configFile) 

2586 self.postgresql.patch_butler_config(config) 

2587 with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh: 

2588 config.dump(fh) 

2589 self.configFile = fh.name 

2590 self._temp_config = True 

2591 super().setUp() 

2592 

2593 def tearDown(self) -> None: 

2594 if self._temp_config and os.path.exists(self.configFile): 

2595 os.remove(self.configFile) 

2596 super().tearDown() 

2597 

2598 def testMakeRepo(self) -> None: 

2599 # The base class test assumes that it's using sqlite and assumes 

2600 # the config file is acceptable to sqlite. 

2601 raise unittest.SkipTest("Postgres config is not compatible with this test.") 

2602 

2603 

2604class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase): 

2605 """Test that Butler with a Postgres registry still works after cloning.""" 

2606 

2607 def create_butler( 

2608 self, 

2609 run: str, 

2610 storageClass: StorageClass | str, 

2611 datasetTypeName: str, 

2612 metrics: ButlerMetrics | None = None, 

2613 ) -> tuple[DirectButler, DatasetType]: 

2614 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics) 

2615 return butler.clone(run=run, metrics=metrics), datasetType 

2616 

2617 

2618class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase): 

2619 """InMemoryDatastore specialization of a butler""" 

2620 

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

2622 fullConfigKey = None 

2623 useTempRoot = False 

2624 validationCanFail = False 

2625 datastoreStr = ["datastore='InMemory"] 

2626 datastoreName = ["InMemoryDatastore@"] 

2627 registryStr = "/gen3.sqlite3" 

2628 

2629 def testIngest(self) -> None: 

2630 pass 

2631 

2632 def test_ingest_zip(self) -> None: 

2633 pass 

2634 

2635 

2636class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase): 

2637 """Test that a Butler with a Sqlite registry still works after cloning.""" 

2638 

2639 def create_butler( 

2640 self, 

2641 run: str, 

2642 storageClass: StorageClass | str, 

2643 datasetTypeName: str, 

2644 metrics: ButlerMetrics | None = None, 

2645 ) -> tuple[DirectButler, DatasetType]: 

2646 butler, datasetType = super().create_butler(run, storageClass, datasetTypeName, metrics=metrics) 

2647 return butler.clone(run=run), datasetType 

2648 

2649 

2650class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2651 """PosixDatastore specialization""" 

2652 

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

2654 fullConfigKey = ".datastore.datastores.1.formatters" 

2655 validationCanFail = True 

2656 datastoreStr = ["datastore='InMemory", "/FileDatastore_1/,", "/FileDatastore_2/'"] 

2657 datastoreName = [ 

2658 "InMemoryDatastore@", 

2659 f"FileDatastore@{BUTLER_ROOT_TAG}/FileDatastore_1", 

2660 "SecondDatastore", 

2661 ] 

2662 registryStr = "/gen3.sqlite3" 

2663 

2664 def testPruneDatasets(self) -> None: 

2665 # This test relies on manipulating files out-of-band, which is 

2666 # impossible for this configuration because of the InMemoryDatastore in 

2667 # the ChainedDatastore. 

2668 pass 

2669 

2670 

2671class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase): 

2672 """Test that a yaml file in one location can refer to a root in another.""" 

2673 

2674 datastoreStr = ["dir1"] 

2675 # Disable the makeRepo test since we are deliberately not using 

2676 # butler.yaml as the config name. 

2677 fullConfigKey = None 

2678 

2679 def setUp(self) -> None: 

2680 self.root = makeTestTempDir(TESTDIR) 

2681 

2682 # Make a new repository in one place 

2683 self.dir1 = os.path.join(self.root, "dir1") 

2684 Butler.makeRepo(self.dir1, config=Config(self.configFile)) 

2685 

2686 # Move the yaml file to a different place and add a "root" 

2687 self.dir2 = os.path.join(self.root, "dir2") 

2688 os.makedirs(self.dir2, exist_ok=True) 

2689 configFile1 = os.path.join(self.dir1, "butler.yaml") 

2690 config = Config(configFile1) 

2691 config["root"] = self.dir1 

2692 configFile2 = os.path.join(self.dir2, "butler2.yaml") 

2693 config.dumpToUri(configFile2) 

2694 os.remove(configFile1) 

2695 self.tmpConfigFile = configFile2 

2696 

2697 def testFileLocations(self) -> None: 

2698 self.assertNotEqual(self.dir1, self.dir2) 

2699 self.assertTrue(os.path.exists(os.path.join(self.dir2, "butler2.yaml"))) 

2700 self.assertFalse(os.path.exists(os.path.join(self.dir1, "butler.yaml"))) 

2701 self.assertTrue(os.path.exists(os.path.join(self.dir1, "gen3.sqlite3"))) 

2702 

2703 

2704class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase): 

2705 """Test that a config file created by makeRepo outside of repo works.""" 

2706 

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

2708 

2709 def setUp(self) -> None: 

2710 self.root = makeTestTempDir(TESTDIR) 

2711 self.root2 = makeTestTempDir(TESTDIR) 

2712 

2713 self.tmpConfigFile = os.path.join(self.root2, "different.yaml") 

2714 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2715 

2716 def tearDown(self) -> None: 

2717 if os.path.exists(self.root2): 

2718 shutil.rmtree(self.root2, ignore_errors=True) 

2719 super().tearDown() 

2720 

2721 def testConfigExistence(self) -> None: 

2722 c = Config(self.tmpConfigFile) 

2723 uri_config = ResourcePath(c["root"]) 

2724 uri_expected = ResourcePath(self.root, forceDirectory=True) 

2725 self.assertEqual(uri_config.geturl(), uri_expected.geturl()) 

2726 self.assertNotIn(":", uri_config.path, "Check for URI concatenated with normal path") 

2727 

2728 def testPutGet(self) -> None: 

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

2730 self.runPutGetTest(storageClass, "test_metric") 

2731 

2732 

2733class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase): 

2734 """Test that a config file created by makeRepo outside of repo works.""" 

2735 

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

2737 

2738 def setUp(self) -> None: 

2739 self.root = makeTestTempDir(TESTDIR) 

2740 self.root2 = makeTestTempDir(TESTDIR) 

2741 

2742 self.tmpConfigFile = self.root2 

2743 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2744 

2745 def testConfigExistence(self) -> None: 

2746 # Append the yaml file else Config constructor does not know the file 

2747 # type. 

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

2749 super().testConfigExistence() 

2750 

2751 

2752class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase): 

2753 """Test that a config file created by makeRepo outside of repo works.""" 

2754 

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

2756 

2757 def setUp(self) -> None: 

2758 self.root = makeTestTempDir(TESTDIR) 

2759 self.root2 = makeTestTempDir(TESTDIR) 

2760 

2761 self.tmpConfigFile = ResourcePath(os.path.join(self.root2, "something.yaml")).geturl() 

2762 Butler.makeRepo(self.root, config=Config(self.configFile), outfile=self.tmpConfigFile) 

2763 

2764 

2765@unittest.skipIf(not boto3, "Warning: boto3 AWS SDK not found!") 

2766class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2767 """S3Datastore specialization of a butler; an S3 storage Datastore + 

2768 a local in-memory SqlRegistry. 

2769 """ 

2770 

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

2772 fullConfigKey = None 

2773 validationCanFail = True 

2774 

2775 bucketName = "anybucketname" 

2776 """Name of the Bucket that will be used in the tests. The name is read from 

2777 the config file used with the tests during set-up. 

2778 """ 

2779 

2780 root = "butlerRoot/" 

2781 """Root repository directory expected to be used in case useTempRoot=False. 

2782 Otherwise the root is set to a 20 characters long randomly generated string 

2783 during set-up. 

2784 """ 

2785 

2786 datastoreStr = [f"datastore={root}"] 

2787 """Contains all expected root locations in a format expected to be 

2788 returned by Butler stringification. 

2789 """ 

2790 

2791 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"] 

2792 """The expected format of the S3 Datastore string.""" 

2793 

2794 registryStr = "/gen3.sqlite3" 

2795 """Expected format of the Registry string.""" 

2796 

2797 mock_aws = mock_aws() 

2798 """The mocked s3 interface from moto.""" 

2799 

2800 def genRoot(self) -> str: 

2801 """Return a random string of len 20 to serve as a root 

2802 name for the temporary bucket repo. 

2803 

2804 This is equivalent to tempfile.mkdtemp as this is what self.root 

2805 becomes when useTempRoot is True. 

2806 """ 

2807 rndstr = "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20)) 

2808 return rndstr + "/" 

2809 

2810 def setUp(self) -> None: 

2811 config = Config(self.configFile) 

2812 uri = ResourcePath(config[".datastore.datastore.root"]) 

2813 self.bucketName = uri.netloc 

2814 

2815 # Enable S3 mocking of tests. 

2816 self.enterContext(clean_test_environment_for_s3()) 

2817 self.mock_aws.start() 

2818 

2819 if self.useTempRoot: 

2820 self.root = self.genRoot() 

2821 rooturi = f"s3://{self.bucketName}/{self.root}" 

2822 config.update({"datastore": {"datastore": {"root": rooturi}}}) 

2823 

2824 # need local folder to store registry database 

2825 self.reg_dir = makeTestTempDir(TESTDIR) 

2826 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3" 

2827 

2828 # MOTO needs to know that we expect Bucket bucketname to exist 

2829 # (this used to be the class attribute bucketName) 

2830 s3 = boto3.resource("s3") 

2831 s3.create_bucket(Bucket=self.bucketName) 

2832 

2833 self.datastoreStr = [f"datastore='{rooturi}'"] 

2834 self.datastoreName = [f"FileDatastore@{rooturi}"] 

2835 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False) 

2836 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml") 

2837 

2838 def tearDown(self) -> None: 

2839 s3 = boto3.resource("s3") 

2840 bucket = s3.Bucket(self.bucketName) 

2841 try: 

2842 bucket.objects.all().delete() 

2843 except botocore.exceptions.ClientError as e: 

2844 if e.response["Error"]["Code"] == "404": 

2845 # the key was not reachable - pass 

2846 pass 

2847 else: 

2848 raise 

2849 

2850 bucket = s3.Bucket(self.bucketName) 

2851 bucket.delete() 

2852 

2853 # Stop the S3 mock. 

2854 self.mock_aws.stop() 

2855 

2856 if self.reg_dir is not None and os.path.exists(self.reg_dir): 

2857 shutil.rmtree(self.reg_dir, ignore_errors=True) 

2858 

2859 if self.useTempRoot and os.path.exists(self.root): 

2860 shutil.rmtree(self.root, ignore_errors=True) 

2861 

2862 super().tearDown() 

2863 

2864 

2865class DatastoreTransfers(TestCaseMixin): 

2866 """Base test setup for data transfers between butlers. The concrete tests 

2867 for specific configurations are in other classes, below. 

2868 """ 

2869 

2870 storageClassFactory: StorageClassFactory 

2871 

2872 @classmethod 

2873 def setUpClass(cls) -> None: 

2874 cls.storageClassFactory = StorageClassFactory() 

2875 

2876 def setUp(self) -> None: 

2877 self.root = makeTestTempDir(TESTDIR) 

2878 self.config = Config(self.configFile) 

2879 

2880 # Some tests cause convertors to be replaced so ensure 

2881 # the storage class factory is reset each time. 

2882 self.storageClassFactory.reset() 

2883 self.storageClassFactory.addFromConfig(self.configFile) 

2884 

2885 def tearDown(self) -> None: 

2886 removeTestTempDir(self.root) 

2887 

2888 def create_butler(self, manager: str | None, label: str, config_file: str | None = None) -> Butler: 

2889 if manager is None: 

2890 manager = ( 

2891 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID" 

2892 ) 

2893 config = Config(config_file if config_file is not None else self.configFile) 

2894 config["registry", "managers", "datasets"] = manager 

2895 butler = Butler.from_config( 

2896 Butler.makeRepo(f"{self.root}/butler{label}", config=config), writeable=True 

2897 ) 

2898 self.enterContext(butler) 

2899 return butler 

2900 

2901 def assertButlerTransfers( 

2902 self, 

2903 purge: bool = False, 

2904 storageClassName: str = "StructuredData", 

2905 storageClassNameTarget: str | None = None, 

2906 ) -> None: 

2907 """Test that a run can be transferred to another butler.""" 

2908 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

2909 if storageClassNameTarget is not None: 

2910 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget) 

2911 else: 

2912 storageClassTarget = storageClass 

2913 

2914 datasetTypeName = "random_data" 

2915 

2916 # Test will create 3 collections and we will want to transfer 

2917 # two of those three. 

2918 runs = ["run1", "run2", "other"] 

2919 

2920 # Also want to use two different dataset types to ensure that 

2921 # grouping works. 

2922 datasetTypeNames = ["random_data", "random_data_2"] 

2923 

2924 # Create the run collections in the source butler. 

2925 for run in runs: 

2926 self.source_butler.collections.register(run) 

2927 

2928 # Create dimensions in source butler. 

2929 n_exposures = 30 

2930 self.source_butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

2931 self.source_butler.registry.insertDimensionData( 

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

2933 ) 

2934 self.source_butler.registry.insertDimensionData( 

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

2936 ) 

2937 self.source_butler.registry.insertDimensionData( 

2938 "day_obs", 

2939 { 

2940 "instrument": "DummyCamComp", 

2941 "id": 20250101, 

2942 }, 

2943 ) 

2944 

2945 for i in range(n_exposures): 

2946 self.source_butler.registry.insertDimensionData( 

2947 "group", {"instrument": "DummyCamComp", "name": f"group{i}"} 

2948 ) 

2949 self.source_butler.registry.insertDimensionData( 

2950 "exposure", 

2951 { 

2952 "instrument": "DummyCamComp", 

2953 "id": i, 

2954 "obs_id": f"exp{i}", 

2955 "physical_filter": "d-r", 

2956 "group": f"group{i}", 

2957 "day_obs": 20250101, 

2958 }, 

2959 ) 

2960 

2961 # Create dataset types in the source butler. 

2962 dimensions = self.source_butler.dimensions.conform(["instrument", "exposure"]) 

2963 for datasetTypeName in datasetTypeNames: 

2964 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

2965 self.source_butler.registry.registerDatasetType(datasetType) 

2966 

2967 # Write a dataset to an unrelated run -- this will ensure that 

2968 # we are rewriting integer dataset ids in the target if necessary. 

2969 # Will not be relevant for UUID. 

2970 run = "distraction" 

2971 butler = Butler.from_config(butler=self.source_butler, run=run) 

2972 self.enterContext(butler) 

2973 butler.put( 

2974 makeExampleMetrics(), 

2975 datasetTypeName, 

2976 exposure=1, 

2977 instrument="DummyCamComp", 

2978 physical_filter="d-r", 

2979 ) 

2980 

2981 # Write some example metrics to the source 

2982 butler = Butler.from_config(butler=self.source_butler) 

2983 self.enterContext(butler) 

2984 

2985 # Set of DatasetRefs that should be in the list of refs to transfer 

2986 # but which will not be transferred. 

2987 deleted: set[DatasetRef] = set() 

2988 

2989 n_expected = 20 # Number of datasets expected to be transferred 

2990 source_refs = [] 

2991 for i in range(n_exposures): 

2992 # Put a third of datasets into each collection, only retain 

2993 # two thirds. 

2994 index = i % 3 

2995 run = runs[index] 

2996 datasetTypeName = datasetTypeNames[i % 2] 

2997 

2998 metric = MetricsExample( 

2999 summary={"counter": i}, output={"text": "metric"}, data=[2 * x for x in range(i)] 

3000 ) 

3001 dataId = {"exposure": i, "instrument": "DummyCamComp", "physical_filter": "d-r"} 

3002 ref = butler.put(metric, datasetTypeName, dataId=dataId, run=run) 

3003 

3004 # Remove the datastore record using low-level API, but only 

3005 # for a specific index. 

3006 if purge and index == 1: 

3007 # For one of these delete the file as well. 

3008 # This allows the "missing" code to filter the 

3009 # file out. 

3010 # Access the individual datastores. 

3011 datastores = [] 

3012 if hasattr(butler._datastore, "datastores"): 

3013 datastores.extend(butler._datastore.datastores) 

3014 else: 

3015 datastores.append(butler._datastore) 

3016 

3017 if not deleted: 

3018 # For a chained datastore we need to remove 

3019 # files in each chain. 

3020 for datastore in datastores: 

3021 # The file might not be known to the datastore 

3022 # if constraints are used. 

3023 try: 

3024 primary, uris = datastore.getURIs(ref) 

3025 except FileNotFoundError: 

3026 continue 

3027 if primary and primary.scheme != "mem": 

3028 primary.remove() 

3029 for uri in uris.values(): 

3030 if uri.scheme != "mem": 

3031 uri.remove() 

3032 n_expected -= 1 

3033 deleted.add(ref) 

3034 

3035 # Remove the datastore record. 

3036 for datastore in datastores: 

3037 if hasattr(datastore, "removeStoredItemInfo"): 

3038 datastore.removeStoredItemInfo(ref) 

3039 

3040 if index < 2: 

3041 source_refs.append(ref) 

3042 if ref not in deleted: 

3043 new_metric = butler.get(ref) 

3044 self.assertEqual(new_metric, metric) 

3045 

3046 # Create some bad dataset types to ensure we check for inconsistent 

3047 # definitions. 

3048 badStorageClass = self.storageClassFactory.getStorageClass("StructuredDataList") 

3049 for datasetTypeName in datasetTypeNames: 

3050 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass) 

3051 self.target_butler.registry.registerDatasetType(datasetType) 

3052 with self.assertRaises(ConflictingDefinitionError) as cm: 

3053 self.target_butler.transfer_from(self.source_butler, source_refs) 

3054 self.assertIn("dataset type differs", str(cm.exception)) 

3055 

3056 # And remove the bad definitions. 

3057 for datasetTypeName in datasetTypeNames: 

3058 self.target_butler.registry.removeDatasetType(datasetTypeName) 

3059 

3060 # Transfer without creating dataset types should fail. 

3061 with self.assertRaises(KeyError): 

3062 self.target_butler.transfer_from(self.source_butler, source_refs) 

3063 

3064 # Transfer without creating dimensions should fail. 

3065 with self.assertRaises(ConflictingDefinitionError) as cm: 

3066 self.target_butler.transfer_from(self.source_butler, source_refs, register_dataset_types=True) 

3067 self.assertIn("dimension", str(cm.exception)) 

3068 

3069 # The dry run test requires dataset types to exist. If we have 

3070 # been given distinct storage classes for the target we have 

3071 # to redefine at least one of the dataset types in the target butler. 

3072 if storageClass != storageClassTarget: 

3073 self.target_butler.registry.removeDatasetType(datasetTypeNames[0]) 

3074 datasetType = DatasetType(datasetTypeNames[0], dimensions, storageClassTarget) 

3075 self.target_butler.registry.registerDatasetType(datasetType) 

3076 

3077 # The failed transfer above leaves registry in an inconsistent 

3078 # state because the run is created but then rolled back without 

3079 # the collection cache being cleared. For now force a refresh. 

3080 # Can remove with DM-35498. 

3081 self.target_butler.registry.refresh() 

3082 

3083 # Do a dry run -- this should not have any effect on the target butler. 

3084 self.target_butler.transfer_from(self.source_butler, source_refs, dry_run=True) 

3085 

3086 # Transfer the records for one ref to test the alternative API. 

3087 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm: 

3088 self.target_butler.transfer_dimension_records_from(self.source_butler, [source_refs[0]]) 

3089 self.assertIn("number of records transferred: 1", ";".join(log_cm.output)) 

3090 

3091 # Now transfer them to the second butler, including dimensions. 

3092 with self.assertLogs(logger="lsst", level=logging.DEBUG) as log_cm: 

3093 transferred = self.target_butler.transfer_from( 

3094 self.source_butler, 

3095 source_refs, 

3096 register_dataset_types=True, 

3097 transfer_dimensions=True, 

3098 ) 

3099 self.assertEqual(len(transferred), n_expected) 

3100 log_output = ";".join(log_cm.output) 

3101 

3102 # A ChainedDatastore will use the in-memory datastore for mexists 

3103 # so we can not rely on the mexists log message. 

3104 self.assertIn("Number of datastore records found in source", log_output) 

3105 self.assertIn("Creating output run", log_output) 

3106 

3107 # Do the transfer twice to ensure that it will do nothing extra. 

3108 # Only do this if purge=True because it does not work for int 

3109 # dataset_id. 

3110 if purge: 

3111 # This should not need to register dataset types. 

3112 transferred = self.target_butler.transfer_from(self.source_butler, source_refs) 

3113 self.assertEqual(len(transferred), n_expected) 

3114 

3115 with self.assertRaises((TypeError, AttributeError)): 

3116 self.target_butler._datastore.transfer_from(self.source_butler, source_refs) # type: ignore 

3117 

3118 with self.assertRaises(ValueError): 

3119 self.target_butler._datastore.transfer_from( 

3120 self.source_butler._datastore, source_refs, transfer="split" 

3121 ) 

3122 

3123 # Now try to get the same refs from the new butler. 

3124 for ref in source_refs: 

3125 if ref not in deleted: 

3126 new_metric = self.target_butler.get(ref) 

3127 old_metric = self.source_butler.get(ref) 

3128 self.assertEqual(new_metric, old_metric) 

3129 

3130 # Try again without implicit storage class conversion 

3131 # triggered by using the source ref. This will do conversion 

3132 # since the formatter will be returning the source python type. 

3133 target_ref = self.target_butler.get_dataset(ref.id) 

3134 if target_ref.datasetType.storageClass != ref.datasetType.storageClass: 

3135 new_metric = self.target_butler.get(target_ref) 

3136 self.assertNotEqual(type(new_metric), type(old_metric)) 

3137 

3138 # Remove the dataset from the target and put it again 

3139 # as if it was the right type all along for this butler. 

3140 self.target_butler.pruneDatasets( 

3141 [target_ref], unstore=True, purge=True, disassociate=True 

3142 ) 

3143 self.target_butler.put(new_metric, target_ref) 

3144 new_new_metric = self.target_butler.get(target_ref) 

3145 new_old_metric = self.target_butler.get( 

3146 target_ref, storageClass=ref.datasetType.storageClass 

3147 ) 

3148 self.assertEqual(new_new_metric, new_metric) 

3149 self.assertEqual(new_old_metric, old_metric) 

3150 

3151 # Now prune run2 collection and create instead a CHAINED collection. 

3152 # This should block the transfer. 

3153 self.target_butler.removeRuns(["run2"]) 

3154 self.target_butler.collections.register("run2", CollectionType.CHAINED) 

3155 with self.assertRaises(CollectionTypeError): 

3156 # Re-importing the run1 datasets can be problematic if they 

3157 # use integer IDs so filter those out. 

3158 to_transfer = [ref for ref in source_refs if ref.run == "run2"] 

3159 self.target_butler.transfer_from(self.source_butler, to_transfer) 

3160 

3161 

3162class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3163 """Test data transfers between butlers. 

3164 

3165 Test for different managers. UUID to UUID and integer to integer are 

3166 tested. UUID to integer is not supported since we do not currently 

3167 want to allow that. Integer to UUID is supported with the caveat 

3168 that UUID4 will be generated and this will be incorrect for raw 

3169 dataset types. The test ignores that. 

3170 """ 

3171 

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

3173 

3174 def create_butlers( 

3175 self, manager1: str | None = None, manager2: str | None = None, source_config: str | None = None 

3176 ) -> None: 

3177 self.source_butler = self.create_butler(manager1, "1", config_file=source_config) 

3178 self.target_butler = self.create_butler(manager2, "2") 

3179 

3180 def testTransferUuidToUuid(self) -> None: 

3181 self.create_butlers() 

3182 self.assertButlerTransfers() 

3183 

3184 def testTransferFromChainedUuidToUuid(self) -> None: 

3185 """Force the source butler to be a ChainedDatastore.""" 

3186 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-chained.yaml")) 

3187 self.assertButlerTransfers() 

3188 

3189 def testTransferFromIncompatibleUuidToUuid(self) -> None: 

3190 """Force the source butler to be a incompatible datastore.""" 

3191 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")) 

3192 with self.assertRaises(NotImplementedError): 

3193 self.assertButlerTransfers() 

3194 

3195 def testTransferFromIncompatibleChainUuidToUuid(self) -> None: 

3196 """Force the source butler to be a incompatible datastore.""" 

3197 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler-inmemory-chain.yaml")) 

3198 with self.assertRaises(TypeError): 

3199 self.assertButlerTransfers() 

3200 

3201 def testTransferFromFileUuidToUuid(self) -> None: 

3202 """Force the source butler to be a FileDatastore.""" 

3203 self.create_butlers(source_config=os.path.join(TESTDIR, "config/basic/butler.yaml")) 

3204 self.assertButlerTransfers() 

3205 

3206 def testTransferMissing(self) -> None: 

3207 """Test transfers where datastore records are missing. 

3208 

3209 This is how execution butler works. 

3210 """ 

3211 self.create_butlers() 

3212 

3213 # Configure the source butler to allow trust. 

3214 self.source_butler._datastore._set_trust_mode(True) 

3215 

3216 self.assertButlerTransfers(purge=True) 

3217 

3218 def testTransferMissingDisassembly(self) -> None: 

3219 """Test transfers where datastore records are missing. 

3220 

3221 This is how execution butler works. 

3222 """ 

3223 self.create_butlers() 

3224 

3225 # Configure the source butler to allow trust. 

3226 self.source_butler._datastore._set_trust_mode(True) 

3227 

3228 # Test disassembly. 

3229 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite") 

3230 

3231 def testTransferDifferingStorageClasses(self) -> None: 

3232 """Test transfers when the source butler dataset type has a different 

3233 but compatible storage class. 

3234 """ 

3235 self.create_butlers() 

3236 

3237 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion") 

3238 

3239 def testTransferDifferingStorageClassesDisassembly(self) -> None: 

3240 """Test transfers when the source butler dataset type has a different 

3241 but compatible storage class and where the source butler has 

3242 disassembled. 

3243 """ 

3244 self.create_butlers() 

3245 

3246 self.assertButlerTransfers( 

3247 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion" 

3248 ) 

3249 

3250 def testAbsoluteURITransferDirect(self) -> None: 

3251 """Test transfer using an absolute URI.""" 

3252 self._absolute_transfer("auto") 

3253 

3254 def testAbsoluteURITransferCopy(self) -> None: 

3255 """Test transfer using an absolute URI.""" 

3256 self._absolute_transfer("copy") 

3257 

3258 def _absolute_transfer(self, transfer: str) -> None: 

3259 self.create_butlers() 

3260 

3261 storageClassName = "StructuredData" 

3262 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

3263 datasetTypeName = "random_data" 

3264 run = "run1" 

3265 self.source_butler.collections.register(run) 

3266 

3267 dimensions = self.source_butler.dimensions.conform(()) 

3268 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3269 self.source_butler.registry.registerDatasetType(datasetType) 

3270 

3271 metrics = makeExampleMetrics() 

3272 with ResourcePath.temporary_uri(suffix=".json") as temp: 

3273 dataId = DataCoordinate.make_empty(self.source_butler.dimensions) 

3274 source_refs = [DatasetRef(datasetType, dataId, run=run)] 

3275 temp.write(json.dumps(metrics.exportAsDict()).encode()) 

3276 dataset = FileDataset(path=temp, refs=source_refs) 

3277 self.source_butler.ingest(dataset, transfer="direct") 

3278 

3279 self.target_butler.transfer_from( 

3280 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer 

3281 ) 

3282 

3283 uri = self.target_butler.getURI(dataset.refs[0]) 

3284 if transfer == "auto": 

3285 self.assertEqual(uri, temp) 

3286 else: 

3287 self.assertNotEqual(uri, temp) 

3288 

3289 def test_shared_dimension_group(self): 

3290 """Test internal logic that divides dataset types by dimension group 

3291 when doing registry updates. 

3292 """ 

3293 self.create_butlers() 

3294 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True) 

3295 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True) 

3296 

3297 source_butler = self.source_butler 

3298 target_butler = self.target_butler 

3299 

3300 # Create a dataset type with the same dimensions as the 'bias' dataset 

3301 # type from base.yaml 

3302 dataset_type = DatasetType( 

3303 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions 

3304 ) 

3305 source_butler.registry.registerDatasetType(dataset_type) 

3306 # This has the same data ID as one of the bias datasets in 

3307 # datasets.yaml. 

3308 test_ref = source_butler.registry.insertDatasets( 

3309 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g" 

3310 )[0] 

3311 

3312 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"]) 

3313 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"]) 

3314 refs = [test_ref, *biases, *flats] 

3315 

3316 # Test setup will be even more convoluted if we want the datastore to 

3317 # actually transfer files. For testing the dimension group behavior, 

3318 # we really only care about the registry. 

3319 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock: 

3320 mock.return_value = (set(refs), set()) 

3321 target_butler.transfer_from( 

3322 source_butler, 

3323 refs, 

3324 transfer=None, 

3325 register_dataset_types=True, 

3326 skip_missing=False, 

3327 transfer_dimensions=True, 

3328 ) 

3329 

3330 transferred_test_ref = target_butler.find_dataset( 

3331 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3332 ) 

3333 self.assertEqual(transferred_test_ref.id, test_ref.id) 

3334 

3335 transferred_bias = target_butler.find_dataset( 

3336 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3337 ) 

3338 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd")) 

3339 

3340 transferred_flat = target_butler.find_dataset( 

3341 "flat", 

3342 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"}, 

3343 collections="imported_r", 

3344 ) 

3345 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840")) 

3346 

3347 

3348class ChainedDatastoreTransfers(PosixDatastoreTransfers): 

3349 """Test transfers using a chained datastore.""" 

3350 

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

3352 

3353 

3354@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3355class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3356 """Test ``transfer_from`` involving Butler server.""" 

3357 

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

3359 

3360 def test_transfers_from_remote_to_direct(self) -> None: 

3361 from lsst.daf.butler.remote_butler._remote_file_transfer_source import ( 

3362 mock_file_transfer_uris_for_unit_test, 

3363 ) 

3364 

3365 self.target_butler = self.create_butler(None, "2") 

3366 with create_test_server(TESTDIR) as server: 

3367 self.source_butler = server.hybrid_butler 

3368 

3369 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath: 

3370 # The Butler server returns HTTP URIs with a domain name that 

3371 # is not resolvable because there is no actual HTTP server 

3372 # involved in these tests. Strip this first layer of 

3373 # indirection, and return the target of the redirect instead. 

3374 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers) 

3375 return ResourcePath(str(response.next_request.url)) 

3376 

3377 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url): 

3378 self.assertButlerTransfers() 

3379 

3380 

3381class NullDatastoreTestCase(unittest.TestCase): 

3382 """Test that we can fall back to a null datastore.""" 

3383 

3384 # Need a good config to create the repo. 

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

3386 storageClassFactory: StorageClassFactory 

3387 

3388 @classmethod 

3389 def setUpClass(cls) -> None: 

3390 cls.storageClassFactory = StorageClassFactory() 

3391 cls.storageClassFactory.addFromConfig(cls.configFile) 

3392 

3393 def setUp(self) -> None: 

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

3395 self.root = makeTestTempDir(TESTDIR) 

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

3397 

3398 def tearDown(self) -> None: 

3399 removeTestTempDir(self.root) 

3400 

3401 def test_fallback(self) -> None: 

3402 # Read the butler config and mess with the datastore section. 

3403 config_path = os.path.join(self.root, "butler.yaml") 

3404 bad_config = Config(config_path) 

3405 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore" 

3406 bad_config.dumpToUri(config_path) 

3407 

3408 with self.assertRaises(RuntimeError): 

3409 Butler(self.root, without_datastore=False) 

3410 

3411 with self.assertRaises(RuntimeError): 

3412 Butler.from_config(self.root, without_datastore=False) 

3413 

3414 butler = Butler.from_config(self.root, writeable=True, without_datastore=True) 

3415 self.enterContext(butler) 

3416 self.assertIsInstance(butler._datastore, NullDatastore) 

3417 

3418 # Check that registry is working. 

3419 butler.collections.register("MYRUN") 

3420 collections = butler.collections.query("*") 

3421 self.assertIn("MYRUN", set(collections)) 

3422 

3423 # Create a ref. 

3424 dimensions = butler.dimensions.conform([]) 

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

3426 datasetTypeName = "metric" 

3427 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3428 butler.registry.registerDatasetType(datasetType) 

3429 ref = DatasetRef(datasetType, {}, run="MYRUN") 

3430 

3431 # Check that datastore will complain. 

3432 with self.assertRaises(FileNotFoundError): 

3433 butler.get(ref) 

3434 with self.assertRaises(FileNotFoundError): 

3435 butler.getURI(ref) 

3436 

3437 

3438@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3439class ButlerServerTests(FileDatastoreButlerTests): 

3440 """Test RemoteButler and Butler server.""" 

3441 

3442 configFile = None 

3443 predictionSupported = False 

3444 trustModeSupported = False 

3445 

3446 postgres: TemporaryPostgresInstance | None 

3447 

3448 def setUp(self): 

3449 self.server_instance = self.enterContext(create_test_server(TESTDIR)) 

3450 

3451 def tearDown(self): 

3452 pass 

3453 

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

3455 # S3 pre-signed URLs may end up with differing expiration times in the 

3456 # query parameters, so ignore query parameters when comparing. 

3457 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path 

3458 

3459 def create_empty_butler( 

3460 self, 

3461 run: str | None = None, 

3462 writeable: bool | None = None, 

3463 metrics: ButlerMetrics | None = None, 

3464 cleanup: bool = True, 

3465 ) -> Butler: 

3466 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics) 

3467 

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

3469 # Can't delete a file via S3 signed URLs, so we need to reach in 

3470 # through DirectButler to delete the dataset. 

3471 uri = self.server_instance.direct_butler.getURI(ref) 

3472 uri.remove() 

3473 

3474 def testConstructor(self): 

3475 # RemoteButler constructor is tested in test_server.py and 

3476 # test_remote_butler.py. 

3477 pass 

3478 

3479 def testDafButlerRepositories(self): 

3480 # Loading of RemoteButler via repository index is tested in 

3481 # test_server.py. 

3482 pass 

3483 

3484 def testGetDatasetTypes(self) -> None: 

3485 # This is mostly a test of validateConfiguration, which is for 

3486 # validating Datastore configuration and thus isn't relevant to 

3487 # RemoteButler. 

3488 pass 

3489 

3490 def testMakeRepo(self) -> None: 

3491 # Only applies to DirectButler. 

3492 pass 

3493 

3494 # Pickling not yet implemented for RemoteButler/HybridButler. 

3495 @unittest.expectedFailure 

3496 def testPickle(self) -> None: 

3497 return super().testPickle() 

3498 

3499 def testStringification(self) -> None: 

3500 self.assertEqual( 

3501 str(self.server_instance.remote_butler), 

3502 "RemoteButler(https://test.example/api/butler/repo/testrepo/)", 

3503 ) 

3504 

3505 def testTransaction(self) -> None: 

3506 # Transactions will never be supported for RemoteButler. 

3507 pass 

3508 

3509 def testPutTemplates(self) -> None: 

3510 # The Butler server instance is configured with different file naming 

3511 # templates than this test is expecting. 

3512 pass 

3513 

3514 

3515@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3516class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase): 

3517 """Tests for RemoteButler's registry shim, with a SQLite DB backing the 

3518 server. 

3519 """ 

3520 

3521 postgres = None 

3522 

3523 

3524@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3525class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase): 

3526 """Tests for RemoteButler's registry shim, with a Postgres DB backing the 

3527 server. 

3528 """ 

3529 

3530 @classmethod 

3531 def setUpClass(cls): 

3532 cls.postgres = cls.enterClassContext(setup_postgres_test_db()) 

3533 super().setUpClass() 

3534 

3535 

3536def setup_module(module: types.ModuleType) -> None: 

3537 """Set up the module for pytest.""" 

3538 clean_environment() 

3539 

3540 

3541def _get_test_data_path(filename: str) -> ResourcePath: 

3542 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}") 

3543 

3544 

3545if __name__ == "__main__": 

3546 clean_environment() 

3547 unittest.main()