Coverage for tests / test_butler.py: 13%

1838 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-15 06:22 -0300

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 

872 def testCompositePutGetVirtual(self) -> None: 

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

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

875 

876 # Should be disassembled 

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

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

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

880 

881 if butler._datastore.isEphemeral: 

882 # Never disassemble in-memory datastore 

883 self.assertIsInstance(uri, ResourcePath) 

884 self.assertFalse(components) 

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

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

887 else: 

888 self.assertIsNone(uri) 

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

890 for compuri in components.values(): 

891 self.assertIsInstance(compuri, ResourcePath) 

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

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

894 

895 if self.predictionSupported: 

896 # Predicted dataset 

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

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

899 

900 if butler._datastore.isEphemeral: 

901 # Never disassembled 

902 self.assertIsInstance(uri, ResourcePath) 

903 self.assertFalse(components) 

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

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

906 else: 

907 self.assertIsNone(uri) 

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

909 for compuri in components.values(): 

910 self.assertIsInstance(compuri, ResourcePath) 

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

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

913 

914 def testStorageClassOverrideGet(self) -> None: 

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

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

917 datasetTypeName = "anything" 

918 run = self.default_run 

919 

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

921 

922 # Create and store a dataset. 

923 metric = makeExampleMetrics() 

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

925 

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

927 

928 # Return native type. 

929 retrieved = butler.get(ref) 

930 self.assertEqual(retrieved, metric) 

931 

932 # Specify an override. 

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

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

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

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

937 self.assertEqual(retrieved, model) 

938 

939 # Defer but override later. 

940 deferred = butler.getDeferred(ref) 

941 model = deferred.get(storageClass=new_sc) 

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

943 self.assertEqual(retrieved, model) 

944 

945 # Defer but override up front. 

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

947 model = deferred.get() 

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

949 self.assertEqual(retrieved, model) 

950 

951 # Retrieve a component. Should be a tuple. 

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

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

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

955 

956 # Parameter on the write storage class should work regardless 

957 # of read storage class. 

958 data = butler.get( 

959 "anything.data", 

960 dataId, 

961 storageClass="StructuredDataDataTestTuple", 

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

963 ) 

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

965 

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

967 # the write storage class. 

968 with self.assertRaises(KeyError): 

969 butler.get( 

970 "anything.data", 

971 dataId, 

972 storageClass="StructuredDataDataTestTuple", 

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

974 ) 

975 

976 def testPytypePutCoercion(self) -> None: 

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

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

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

980 datasetTypeName = "test_metric" 

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

982 

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

984 

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

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

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

988 test_metric = butler.get(metric_ref) 

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

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

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

992 

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

994 # a definition matching this python type. 

995 registry_type = butler.get_dataset_type(datasetTypeName) 

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

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

998 self.assertEqual(metric2_ref.datasetType, registry_type) 

999 

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

1001 test_metric2 = butler.get(metric2_ref) 

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

1003 

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

1005 # This should now return a dict. 

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

1007 test_dict2 = butler.get(new_ref) 

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

1009 

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

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

1012 # behavior and return the type of the DatasetType. 

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

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

1015 

1016 def test_ingest_zip(self) -> None: 

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

1018 butler, dataset_type = self.create_butler( 

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

1020 ) 

1021 

1022 metric = makeExampleMetrics() 

1023 refs = [] 

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

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

1026 refs.append(ref) 

1027 

1028 # Retrieve a Zip file. 

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

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

1031 

1032 # Ingest will fail. 

1033 with self.assertRaises(ConflictingDefinitionError): 

1034 butler.ingest_zip(zip) 

1035 

1036 # Clear out the collection. 

1037 butler.removeRuns([self.default_run]) 

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

1039 

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

1041 

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

1043 new_butler_cfg = Butler.makeRepo(tmpdir) 

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

1045 self.enterContext(new_butler) 

1046 

1047 # This will fail since dimensions records are missing. 

1048 with self.assertRaises(ConflictingDefinitionError): 

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

1050 

1051 # Dry run should work. 

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

1053 

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

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

1056 

1057 # Check that the refs can be read again. 

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

1059 

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

1061 self.assertTrue(uri.exists()) 

1062 

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

1064 # remaining refs to be read. 

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

1066 self.assertTrue(uri.exists()) 

1067 

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

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

1070 

1071 butler.removeRuns([self.default_run]) 

1072 self.assertFalse(uri.exists()) 

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

1074 

1075 with self.assertRaises(ValueError): 

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

1077 

1078 def testIngest(self) -> None: 

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

1080 

1081 # Create and register a DatasetType 

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

1083 

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

1085 datasetTypeName = "metric" 

1086 

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

1088 

1089 # Add needed Dimensions 

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

1091 butler.registry.insertDimensionData( 

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

1093 ) 

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

1095 for detector in (1, 2): 

1096 butler.registry.insertDimensionData( 

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

1098 ) 

1099 

1100 butler.registry.insertDimensionData( 

1101 "visit", 

1102 { 

1103 "instrument": "DummyCamComp", 

1104 "id": 423, 

1105 "name": "fourtwentythree", 

1106 "physical_filter": "d-r", 

1107 "day_obs": 20250101, 

1108 }, 

1109 { 

1110 "instrument": "DummyCamComp", 

1111 "id": 424, 

1112 "name": "fourtwentyfour", 

1113 "physical_filter": "d-r", 

1114 "day_obs": 20250101, 

1115 }, 

1116 ) 

1117 

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

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

1120 datasets = [] 

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

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

1123 # required. 

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

1125 for detector in (1, 2): 

1126 detector_name = f"detector_{detector}" 

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

1128 dataId = butler.registry.expandDataId( 

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

1130 ) 

1131 # Create a DatasetRef for ingest 

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

1133 

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

1135 

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

1137 

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

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

1140 

1141 metrics1 = butler.get(datasetTypeName, dataId1) 

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

1143 self.assertNotEqual(metrics1, metrics2) 

1144 

1145 # Compare URIs 

1146 uri1 = butler.getURI(datasetTypeName, dataId1) 

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

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

1149 

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

1151 # skip_existing=False. 

1152 with self.assertRaises(ConflictingDefinitionError): 

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

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

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

1156 

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

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

1159 refs = [] 

1160 for detector in (1, 2): 

1161 detector_name = f"detector_{detector}" 

1162 dataId = butler.registry.expandDataId( 

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

1164 ) 

1165 # Create a DatasetRef for ingest 

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

1167 

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

1169 # have disappeared following ingest. 

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

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

1172 

1173 datasets = [] 

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

1175 

1176 # For first ingest use copy. 

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

1178 

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

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

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

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

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

1184 datasets[0].refs = [ 

1185 cast( 

1186 DatasetRef, 

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

1188 ) 

1189 for ref in datasets[0].refs 

1190 ] 

1191 all_refs = [] 

1192 for dataset in datasets: 

1193 refs = [] 

1194 for ref in dataset.refs: 

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

1196 new_data_id = dict(ref.dataId.required) 

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

1198 assert new_ref is not None 

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

1200 refs.append(new_ref) 

1201 dataset.refs = refs 

1202 all_refs.extend(dataset.refs) 

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

1204 

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

1206 # disable recording of file size. 

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

1208 

1209 # Check that every ref now has records. 

1210 for dataset in datasets: 

1211 for ref in dataset.refs: 

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

1213 

1214 # Ensure that the file has disappeared. 

1215 self.assertFalse(tempFile.exists()) 

1216 

1217 # Check that the datastore recorded no file size. 

1218 # Not all datastores can support this. 

1219 try: 

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

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

1222 except AttributeError: 

1223 pass 

1224 

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

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

1227 

1228 multi1 = butler.get(datasetTypeName, dataId1) 

1229 multi2 = butler.get(datasetTypeName, dataId2) 

1230 

1231 self.assertEqual(multi1, metrics1) 

1232 self.assertEqual(multi2, metrics2) 

1233 

1234 # Compare URIs 

1235 uri1 = butler.getURI(datasetTypeName, dataId1) 

1236 uri2 = butler.getURI(datasetTypeName, dataId2) 

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

1238 

1239 # Test that removing one does not break the second 

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

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

1242 # files. 

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

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

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

1246 multi2b = butler.get(datasetTypeName, dataId2) 

1247 self.assertEqual(multi2, multi2b) 

1248 

1249 # Ensure we can ingest 0 datasets 

1250 datasets = [] 

1251 butler.ingest(*datasets) 

1252 

1253 def testPickle(self) -> None: 

1254 """Test pickle support.""" 

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

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

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

1258 self.enterContext(butlerOut) 

1259 self.assertIsInstance(butlerOut, Butler) 

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

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

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

1263 

1264 def testGetDatasetTypes(self) -> None: 

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

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

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

1268 ( 

1269 "instrument", 

1270 [ 

1271 {"instrument": "DummyCam"}, 

1272 {"instrument": "DummyHSC"}, 

1273 {"instrument": "DummyCamComp"}, 

1274 ], 

1275 ), 

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

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

1278 ( 

1279 "visit", 

1280 [ 

1281 { 

1282 "instrument": "DummyCam", 

1283 "id": 42, 

1284 "name": "fortytwo", 

1285 "physical_filter": "d-r", 

1286 "day_obs": 20250101, 

1287 } 

1288 ], 

1289 ), 

1290 ] 

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

1292 # Add needed Dimensions 

1293 for element, data in dimensionEntries: 

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

1295 

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

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

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

1299 components = set() 

1300 for datasetTypeName in datasetTypeNames: 

1301 # Create and register a DatasetType 

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

1303 

1304 for componentName in storageClass.components: 

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

1306 

1307 fromRegistry: set[DatasetType] = set() 

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

1309 fromRegistry.add(parent_dataset_type) 

1310 fromRegistry.update(parent_dataset_type.makeAllComponentDatasetTypes()) 

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

1312 

1313 # Query with wildcard. 

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

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

1316 # but not regex. 

1317 with self.assertRaises(DatasetTypeExpressionError): 

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

1319 

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

1321 butler.validateConfiguration( 

1322 ignore=[ 

1323 "test_metric_comp", 

1324 "metric3", 

1325 "metric5", 

1326 "calexp", 

1327 "DummySC", 

1328 "datasetType.component", 

1329 "random_data", 

1330 "random_data_2", 

1331 ] 

1332 ) 

1333 

1334 # Add a new datasetType that will fail template validation 

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

1336 if self.validationCanFail: 

1337 with self.assertRaises(ValidationError): 

1338 butler.validateConfiguration() 

1339 

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

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

1342 

1343 # Rerun validation but ignore the bad datasetType 

1344 butler.validateConfiguration( 

1345 ignore=[ 

1346 "test_metric_comp", 

1347 "metric3", 

1348 "metric5", 

1349 "calexp", 

1350 "DummySC", 

1351 "datasetType.component", 

1352 "random_data", 

1353 "random_data_2", 

1354 ] 

1355 ) 

1356 

1357 def testTransaction(self) -> None: 

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

1359 datasetTypeName = "test_metric" 

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

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

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

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

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

1365 ( 

1366 "visit", 

1367 { 

1368 "instrument": "DummyCam", 

1369 "id": 42, 

1370 "name": "fortytwo", 

1371 "physical_filter": "d-r", 

1372 "day_obs": 20250101, 

1373 }, 

1374 ), 

1375 ) 

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

1377 metric = makeExampleMetrics() 

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

1379 # Create and register a DatasetType 

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

1381 with self.assertRaises(TransactionTestError): 

1382 with butler.transaction(): 

1383 # Add needed Dimensions 

1384 for args in dimensionEntries: 

1385 butler.registry.insertDimensionData(*args) 

1386 # Store a dataset 

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

1388 self.assertIsInstance(ref, DatasetRef) 

1389 # Test get of a ref. 

1390 metricOut = butler.get(ref) 

1391 self.assertEqual(metric, metricOut) 

1392 # Test get 

1393 metricOut = butler.get(datasetTypeName, dataId) 

1394 self.assertEqual(metric, metricOut) 

1395 # Check we can get components 

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

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

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

1399 butler.registry.expandDataId(dataId) 

1400 # Should raise LookupError for missing data ID value 

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

1402 butler.get(datasetTypeName, dataId) 

1403 # Also check explicitly if Dataset entry is missing 

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

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

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

1407 butler.get(ref) 

1408 

1409 def testMakeRepo(self) -> None: 

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

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

1412 repo root. 

1413 """ 

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

1415 # not support a file system root 

1416 if self.fullConfigKey is None: 

1417 return 

1418 

1419 # create two separate directories 

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

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

1422 

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

1424 limited = Config(self.configFile) 

1425 butler1 = Butler.from_config(butlerConfig) 

1426 self.enterContext(butler1) 

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

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

1429 full = Config(self.tmpConfigFile) 

1430 butler2 = Butler.from_config(butlerConfig) 

1431 self.enterContext(butler2) 

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

1433 # Butlers should have the same configuration regardless of whether 

1434 # defaults were expanded. 

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

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

1437 self.assertNotEqual(limited, full) 

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

1439 # inheriting from defaults. 

1440 self.assertIn(self.fullConfigKey, full) 

1441 self.assertNotIn(self.fullConfigKey, limited) 

1442 

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

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

1445 self.assertEqual(collections1, set()) 

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

1447 

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

1449 # work properly with relocatable Butler repo 

1450 butlerConfig.configFile = None 

1451 with self.assertRaises(ValueError): 

1452 Butler.from_config(butlerConfig) 

1453 

1454 with self.assertRaises(FileExistsError): 

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

1456 

1457 def testStringification(self) -> None: 

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

1459 self.enterContext(butler) 

1460 butlerStr = str(butler) 

1461 

1462 if self.datastoreStr is not None: 

1463 for testStr in self.datastoreStr: 

1464 self.assertIn(testStr, butlerStr) 

1465 if self.registryStr is not None: 

1466 self.assertIn(self.registryStr, butlerStr) 

1467 

1468 datastoreName = butler._datastore.name 

1469 if self.datastoreName is not None: 

1470 for testStr in self.datastoreName: 

1471 self.assertIn(testStr, datastoreName) 

1472 

1473 def testButlerRewriteDataId(self) -> None: 

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

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

1476 

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

1478 datasetTypeName = "random_data" 

1479 

1480 # Create dimension records. 

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

1482 butler.registry.insertDimensionData( 

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

1484 ) 

1485 butler.registry.insertDimensionData( 

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

1487 ) 

1488 

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

1490 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

1491 butler.registry.registerDatasetType(datasetType) 

1492 

1493 n_exposures = 5 

1494 dayobs = 20210530 

1495 

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

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

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

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

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

1501 

1502 for i in range(n_exposures): 

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

1504 butler.registry.insertDimensionData( 

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

1506 ) 

1507 butler.registry.insertDimensionData( 

1508 "exposure", 

1509 { 

1510 "instrument": "DummyCamComp", 

1511 "id": day_obs + i, 

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

1513 "seq_num": i, 

1514 "day_obs": day_obs, 

1515 "physical_filter": "d-r", 

1516 "group": group_name, 

1517 }, 

1518 ) 

1519 

1520 # Write some data. 

1521 for i in range(n_exposures): 

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

1523 

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

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

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

1527 

1528 # Check that the exposure is correct in the dataId 

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

1530 

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

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

1533 self.assertEqual(new_metric, metric) 

1534 

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

1536 # exposure.day_obs. 

1537 datasets_1 = list( 

1538 butler.registry.queryDatasets( 

1539 datasetType, 

1540 collections=self.default_run, 

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

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

1543 ) 

1544 ) 

1545 datasets_2 = list( 

1546 butler.registry.queryDatasets( 

1547 datasetType, 

1548 collections=self.default_run, 

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

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

1551 ) 

1552 ) 

1553 self.assertEqual(datasets_1, datasets_2) 

1554 

1555 def testGetDatasetCollectionCaching(self): 

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

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

1558 # after the collection cache was last updated. 

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

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

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

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

1563 get_ref = reader_butler.get_dataset(put_ref.id) 

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

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

1566 # instance. 

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

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

1569 

1570 def testCollectionChainRedefine(self): 

1571 butler = self._setup_to_test_collection_chain() 

1572 

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

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

1575 

1576 # Duplicates are removed from the list of children 

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

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

1579 

1580 # Empty list clears the chain 

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

1582 self._check_chain(butler, []) 

1583 

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

1585 

1586 def testCollectionChainPrepend(self): 

1587 butler = self._setup_to_test_collection_chain() 

1588 

1589 # Duplicates are removed from the list of children 

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

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

1592 

1593 # Prepend goes on the front of existing chain 

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

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

1596 

1597 # Empty prepend does nothing 

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

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

1600 

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

1602 # their current position. 

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

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

1605 

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

1607 

1608 def testCollectionChainExtend(self): 

1609 butler = self._setup_to_test_collection_chain() 

1610 

1611 # Duplicates are removed from the list of children 

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

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

1614 

1615 # Extend goes on the end of existing chain 

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

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

1618 

1619 # Empty extend does nothing 

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

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

1622 

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

1624 # their current position. 

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

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

1627 

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

1629 

1630 def testCollectionChainRemove(self) -> None: 

1631 butler = self._setup_to_test_collection_chain() 

1632 

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

1634 

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

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

1637 

1638 # Duplicates are allowed in the list of children 

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

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

1641 

1642 # Empty remove does nothing 

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

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

1645 

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

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

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

1649 

1650 self._test_common_chain_functionality( 

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

1652 ) 

1653 

1654 def _setup_to_test_collection_chain(self) -> Butler: 

1655 butler = self.create_empty_butler(writeable=True) 

1656 

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

1658 

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

1660 for run in runs: 

1661 butler.collections.register(run) 

1662 

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

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

1665 

1666 return butler 

1667 

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

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

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

1671 

1672 def _test_common_chain_functionality( 

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

1674 ) -> None: 

1675 # Missing parent collection 

1676 with self.assertRaises(MissingCollectionError): 

1677 func("doesnotexist", []) 

1678 # Missing child collection 

1679 with self.assertRaises(MissingCollectionError): 

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

1681 # Forbid operations on non-chained collections 

1682 with self.assertRaises(CollectionTypeError): 

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

1684 

1685 # Prevent collection cycles 

1686 if not skip_cycle_check: 

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

1688 func("chain2", "chain") 

1689 with self.assertRaises(CollectionCycleError): 

1690 func("chain", "chain2") 

1691 

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

1693 # chains. 

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

1695 

1696 with butler._caching_context(): 

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

1698 func("chain", "a") 

1699 

1700 def test_transfer_dimension_records_from(self) -> None: 

1701 source_butler = self.create_empty_butler(writeable=True) 

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

1703 

1704 visit_id = 2025120200439 

1705 exposure_id = visit_id 

1706 target_butler = self.enterContext(create_populated_sqlite_registry()) 

1707 target_butler.transfer_dimension_records_from( 

1708 source_butler, 

1709 [ 

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

1711 # "populated_by" records (visit_detector_region, 

1712 # visit_definition, etc.) 

1713 DataCoordinate.standardize( 

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

1715 universe=source_butler.dimensions, 

1716 ), 

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

1718 DataCoordinate.make_empty(source_butler.dimensions), 

1719 ], 

1720 ) 

1721 

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

1723 records = target_butler.query_dimension_records(dimension) 

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

1725 return records[0] 

1726 

1727 visit = _fetch_record("visit") 

1728 self.assertEqual(visit.id, visit_id) 

1729 self.assertEqual(visit.day_obs, 20251202) 

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

1731 self.assertEqual(visit.seq_num, 439) 

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

1733 0 

1734 ] 

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

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

1737 

1738 visit_detector_region = _fetch_record("visit_detector_region") 

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

1740 self.assertEqual(visit_detector_region.detector, 10) 

1741 self.assertEqual(visit_detector_region.visit, visit_id) 

1742 original_visit_detector_region = source_butler.query_dimension_records( 

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

1744 )[0] 

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

1746 

1747 visit_definition = _fetch_record("visit_definition") 

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

1749 self.assertEqual(visit_definition.exposure, 2025120200439) 

1750 self.assertEqual(visit_definition.visit, visit_id) 

1751 

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

1753 # visit -> visit_definition. 

1754 exposure = _fetch_record("exposure") 

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

1756 self.assertEqual(exposure.id, 2025120200439) 

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

1758 original_exposure = source_butler.query_dimension_records( 

1759 "exposure", instrument="LSSTCam", exposure=exposure_id 

1760 )[0] 

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

1762 

1763 group = _fetch_record("group") 

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

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

1766 

1767 visit_system_memberships = target_butler.query_dimension_records("visit_system_membership") 

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

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

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

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

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

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

1774 

1775 visit_systems = target_butler.query_dimension_records("visit_system") 

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

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

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

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

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

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

1782 

1783 

1784class FileDatastoreButlerTests(ButlerTests): 

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

1786 by datastores that inherit from FileDatastore. 

1787 """ 

1788 

1789 trustModeSupported = True 

1790 

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

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

1793 

1794 Test testPutTemplates verifies actual physical existance of the files 

1795 in the requested location. 

1796 """ 

1797 uri = ResourcePath(root, forceDirectory=True) 

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

1799 

1800 def testPutTemplates(self) -> None: 

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

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

1803 

1804 # Add needed Dimensions 

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

1806 butler.registry.insertDimensionData( 

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

1808 ) 

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

1810 butler.registry.insertDimensionData( 

1811 "visit", 

1812 { 

1813 "instrument": "DummyCamComp", 

1814 "id": 423, 

1815 "name": "v423", 

1816 "physical_filter": "d-r", 

1817 "day_obs": 20250101, 

1818 }, 

1819 ) 

1820 butler.registry.insertDimensionData( 

1821 "visit", 

1822 { 

1823 "instrument": "DummyCamComp", 

1824 "id": 425, 

1825 "name": "v425", 

1826 "physical_filter": "d-r", 

1827 "day_obs": 20250101, 

1828 }, 

1829 ) 

1830 

1831 # Create and store a dataset 

1832 metric = makeExampleMetrics() 

1833 

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

1835 # template) 

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

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

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

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

1840 

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

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

1843 

1844 # Put with exactly the data ID keys needed 

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

1846 uri = butler.getURI(ref) 

1847 self.assertTrue(uri.exists()) 

1848 self.assertTrue( 

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

1850 ) 

1851 

1852 # Check the template based on dimensions 

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

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

1855 

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

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

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

1859 # must be consistent). 

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

1861 uri = butler.getURI(ref) 

1862 self.assertTrue(uri.exists()) 

1863 self.assertTrue( 

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

1865 ) 

1866 

1867 # Check the template based on dimensions 

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

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

1870 

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

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

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

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

1875 path = template.format(ref) 

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

1877 

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

1879 with self.assertRaises(KeyError): 

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

1881 template.format(ref) 

1882 

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

1884 with self.assertRaises(FileTemplateValidationError): 

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

1886 

1887 def testImportExport(self) -> None: 

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

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

1890 self.runImportExportTest(storageClass) 

1891 

1892 @unittest.expectedFailure 

1893 def testImportExportVirtualComposite(self) -> None: 

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

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

1896 self.runImportExportTest(storageClass) 

1897 

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

1899 """Test exporting and importing. 

1900 

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

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

1903 """ 

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

1905 

1906 # Test that we must have a file extension. 

1907 with self.assertRaises(ValueError): 

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

1909 pass 

1910 

1911 # Test that unknown format is not allowed. 

1912 with self.assertRaises(ValueError): 

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

1914 pass 

1915 

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

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

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

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

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

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

1922 # Export and then import datasets. 

1923 with safeTestTempDir(TESTDIR) as exportDir: 

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

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

1926 export.saveDatasets(datasets) 

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

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

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

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

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

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

1933 # because of internal deduplication. 

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

1935 # Save some dimension records directly. 

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

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

1938 with safeTestTempDir(TESTDIR) as importDir: 

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

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

1941 # Calling script.butlerImport tests the implementation of the 

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

1943 # in the script folder are generally considered protected and 

1944 # should not be used as public api. 

1945 with open(exportFile) as f: 

1946 script.butlerImport( 

1947 importDir, 

1948 export_file=f, 

1949 directory=exportDir, 

1950 transfer="auto", 

1951 skip_dimensions=None, 

1952 ) 

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

1954 self.enterContext(importButler) 

1955 for ref in datasets: 

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

1957 # Test for existence by passing in the DatasetType and 

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

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

1960 self.assertEqual( 

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

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

1963 ) 

1964 

1965 def testRemoveRuns(self) -> None: 

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

1967 butler = self.create_empty_butler(writeable=True) 

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

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

1970 # Add some RUN-type collection. 

1971 run1 = "run1" 

1972 butler.collections.register(run1) 

1973 run2 = "run2" 

1974 butler.collections.register(run2) 

1975 # put a dataset in each 

1976 metric = makeExampleMetrics() 

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

1978 datasetType = self.addDatasetType( 

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

1980 ) 

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

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

1983 uri1 = butler.getURI(ref1) 

1984 uri2 = butler.getURI(ref2) 

1985 

1986 # Put one of the runs in a chain. 

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

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

1989 

1990 with self.assertRaises(OrphanedRecordError): 

1991 butler.registry.removeDatasetType(datasetType.name) 

1992 

1993 # Remove a non-run. 

1994 with self.assertRaises(TypeError): 

1995 butler.removeRuns(["Chain"]) 

1996 

1997 # Remove without unlinking from chain should fail. 

1998 with self.assertRaises(IntegrityError): 

1999 butler.removeRuns([run1]) 

2000 

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

2002 # always purges. 

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

2004 

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

2006 # not think either exists. 

2007 with self.assertRaises(MissingCollectionError): 

2008 butler.collections.get_info(run1) 

2009 with self.assertRaises(MissingCollectionError): 

2010 butler.collections.get_info(run1) 

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

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

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

2014 self.assertFalse(uri1.exists()) 

2015 self.assertFalse(uri2.exists()) 

2016 

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

2018 # dataset type 

2019 butler.registry.removeDatasetType(datasetType.name) 

2020 

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

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

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

2024 

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

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

2027 knowledge. 

2028 

2029 Subclasses may override to handle more complicated datastore 

2030 configurations. 

2031 """ 

2032 uri = butler.getURI(ref) 

2033 uri.remove() 

2034 datastore = cast(FileDatastore, butler._datastore) 

2035 datastore.cacheManager.remove_from_cache(ref) 

2036 

2037 def testPruneDatasets(self) -> None: 

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

2039 butler = self.create_empty_butler(writeable=True) 

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

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

2042 # Add some RUN-type collections. 

2043 run1 = "run1" 

2044 butler.collections.register(run1) 

2045 run2 = "run2" 

2046 butler.collections.register(run2) 

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

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

2049 metric = makeExampleMetrics() 

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

2051 datasetType = self.addDatasetType( 

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

2053 ) 

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

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

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

2057 

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

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

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

2061 

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

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

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

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

2066 

2067 # Simple prune. 

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

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

2070 

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

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

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

2074 

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

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

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

2078 

2079 # Put data back. 

2080 ref1_new = butler.put(metric, ref1) 

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

2082 ref2 = butler.put(metric, ref2) 

2083 

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

2085 self.assertTrue(many_stored[ref1]) 

2086 self.assertTrue(many_stored[ref2]) 

2087 self.assertFalse(many_stored[ref3]) 

2088 

2089 ref3 = butler.put(metric, ref3) 

2090 

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

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

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

2094 

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

2096 refs = [ref1, ref2, ref3] 

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

2098 for ref in refs: 

2099 butler.put(metric, ref) 

2100 

2101 # Confirm we can retrieve deferred. 

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

2103 metric1 = dref1.get() 

2104 self.assertEqual(metric1, metric) 

2105 

2106 # Test different forms of file availability. 

2107 # Need to be in a state where: 

2108 # - one ref just has registry record. 

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

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

2111 # - one ref does not exist anywhere. 

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

2113 # above. 

2114 ref0 = DatasetRef( 

2115 datasetType, 

2116 DataCoordinate.standardize( 

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

2118 ), 

2119 run=run1, 

2120 ) 

2121 

2122 # Delete from datastore and retain in Registry. 

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

2124 

2125 # File has been removed. 

2126 self.remove_dataset_out_of_band(butler, ref2) 

2127 

2128 # Datastore has lost track. 

2129 butler._datastore.forget([ref3]) 

2130 

2131 # First test with a standard butler. 

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

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

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

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

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

2137 

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

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

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

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

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

2143 self.assertTrue(exists_many[ref2]) 

2144 

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

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

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

2148 

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

2150 # retrieved. 

2151 with self.assertRaises(LookupError): 

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

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

2154 with self.assertRaises(FileNotFoundError): 

2155 dref2.get() 

2156 

2157 # Test again with a trusting butler. 

2158 if self.trustModeSupported: 

2159 butler._datastore.trustGetRequest = True 

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

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

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

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

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

2165 

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

2167 # known but does exist. 

2168 dref3 = butler.getDeferred(ref3) 

2169 metric3 = dref3.get() 

2170 self.assertEqual(metric3, metric) 

2171 

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

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

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

2175 

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

2177 # but is not the same. 

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

2179 with self.assertRaises(ValueError): 

2180 butler.exists(ref_bad) 

2181 

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

2183 ref_compat = ref2.overrideStorageClass("StructuredDataDict") 

2184 exists = butler.exists(ref_compat) 

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

2186 

2187 # Remove everything and start from scratch. 

2188 butler._datastore.trustGetRequest = False 

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

2190 for ref in refs: 

2191 butler.put(metric, ref) 

2192 

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

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

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

2196 # trash not touching the file. 

2197 uri1 = butler.getURI(ref1) 

2198 butler._datastore.bridge.moveToTrash( 

2199 [ref1], transaction=None 

2200 ) # Update the dataset_location table 

2201 butler._datastore.forget([ref1]) 

2202 butler._datastore.trash(ref1) 

2203 butler._datastore.emptyTrash() 

2204 self.assertTrue(uri1.exists()) 

2205 uri1.remove() # Clean it up. 

2206 

2207 # Simulate execution butler setup by deleting the datastore 

2208 # record but keeping the file around and trusting. 

2209 butler._datastore.trustGetRequest = True 

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

2211 uri2 = uris[ref2].primaryURI 

2212 uri3 = uris[ref3].primaryURI 

2213 self.assertTrue(uri2.exists()) 

2214 self.assertTrue(uri3.exists()) 

2215 

2216 # Remove the datastore record. 

2217 butler._datastore.bridge.moveToTrash( 

2218 [ref2], transaction=None 

2219 ) # Update the dataset_location table 

2220 butler._datastore.forget([ref2]) 

2221 self.assertTrue(uri2.exists()) 

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

2223 # Immediate removal for ref2 file 

2224 self.assertFalse(uri2.exists()) 

2225 # But ref3 has to wait for the empty. 

2226 self.assertTrue(uri3.exists()) 

2227 butler._datastore.emptyTrash() 

2228 self.assertFalse(uri3.exists()) 

2229 

2230 # Clear out the datasets from registry. 

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

2232 

2233 def test_butler_metrics(self): 

2234 """Test that metrics are collected.""" 

2235 run = "test_run" 

2236 metrics = ButlerMetrics() 

2237 butler, datasetType = self.create_butler( 

2238 run, "MetricsExampleModelProvenance", "prov_metric", metrics=metrics 

2239 ) 

2240 data = MetricsExampleModel( 

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

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

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

2244 ) 

2245 

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

2247 butler.get(data_ref) 

2248 butler.get(data_ref) 

2249 self.assertEqual(metrics.n_get, 2) 

2250 self.assertGreater(metrics.time_in_get, 0.0) 

2251 self.assertEqual(metrics.n_put, 1) 

2252 self.assertGreater(metrics.time_in_put, 0.0) 

2253 

2254 deferred = butler.getDeferred(data_ref) 

2255 deferred.get() 

2256 self.assertEqual(metrics.n_get, 3) 

2257 

2258 with butler.record_metrics() as new: 

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

2260 butler.get(data_ref) 

2261 

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

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

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

2265 refs = [ 

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

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

2268 ] 

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

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

2271 

2272 self.assertEqual(new.n_get, 1) 

2273 self.assertEqual(new.n_put, 1) 

2274 self.assertEqual(new.n_ingest, 2) 

2275 

2276 

2277class PosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2278 """PosixDatastore specialization of a butler""" 

2279 

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

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

2282 validationCanFail = True 

2283 datastoreStr = ["/tmp"] 

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

2285 registryStr = "/gen3.sqlite3" 

2286 

2287 def testPathConstructor(self) -> None: 

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

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

2290 self.enterContext(butler) 

2291 self.assertIsInstance(butler, Butler) 

2292 

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

2294 path = pathlib.Path(self.tmpConfigFile) 

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

2296 self.enterContext(butler) 

2297 self.assertIsInstance(butler, Butler) 

2298 

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

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

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

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

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

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

2305 self.enterContext(butler) 

2306 self.assertIsInstance(butler, Butler) 

2307 

2308 def testExportTransferCopy(self) -> None: 

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

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

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

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

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

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

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

2316 assert isinstance(exportButler._datastore, FileDatastore) 

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

2318 

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

2320 

2321 for path in pathsInStore: 

2322 # Assume local file system 

2323 assert path is not None 

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

2325 

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

2327 with safeTestTempDir(TESTDIR) as exportDir: 

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

2329 export.saveDatasets(datasets) 

2330 for path in pathsInStore: 

2331 assert path is not None 

2332 self.assertTrue( 

2333 self.checkFileExists(exportDir, path), 

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

2335 ) 

2336 

2337 def testPytypeCoercion(self) -> None: 

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

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

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

2341 datasetTypeName = "test_metric" 

2342 butler = self.runPutGetTest(storageClass, datasetTypeName) 

2343 

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

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

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

2347 

2348 datasetType_ori = butler.get_dataset_type(datasetTypeName) 

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

2350 

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

2352 # There is no API for this. 

2353 assert isinstance(butler._registry, SqlRegistry) 

2354 manager = butler._registry._managers.datasets 

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

2356 manager._db.update( 

2357 manager._static.dataset_type, 

2358 {"name": datasetTypeName}, 

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

2360 ) 

2361 

2362 # Force reset of dataset type cache 

2363 butler.registry.refresh() 

2364 

2365 datasetType_new = butler.get_dataset_type(datasetTypeName) 

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

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

2368 

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

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

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

2372 

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

2374 # works as normal. 

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

2376 metric_model_new = butler.get(metric_ref) 

2377 self.assertEqual(metric_model_new, metric_model) 

2378 

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

2380 # get with no conversion class. 

2381 manager._db.update( 

2382 manager._static.dataset_type, 

2383 {"name": datasetTypeName}, 

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

2385 ) 

2386 butler.registry.refresh() 

2387 

2388 with self.assertRaises(ValueError): 

2389 butler.get(datasetTypeName, dataId=dataId) 

2390 

2391 def test_provenance(self): 

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

2393 run = "test_run" 

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

2395 metric = MetricsExampleModel( 

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

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

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

2399 ) 

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

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

2402 # copy on adding provenance to ensure they differ. 

2403 self.assertIsNone(metric.dataset_id) 

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

2405 self.assertIsNone(metric.dataset_id) 

2406 metric_2 = butler.get(metric_ref) 

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

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

2409 self.assertIsNone(metric_2.provenance) 

2410 

2411 # Put with provenance. 

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

2413 prov.add_input(metric_ref) 

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

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

2416 metric_3 = butler.get(metric_ref2) 

2417 self.assertEqual(metric_3.provenance, prov) 

2418 

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

2420 prov_dict = prov.to_flat_dict(metric_ref2) 

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

2422 self.assertEqual(ref_from_prov, metric_ref2) 

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

2424 # includes dimension records. 

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

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

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

2428 

2429 # Force a bad ID into the dict. 

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

2431 with self.assertRaises(ValueError): 

2432 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2433 del prov_dict["id"] 

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

2435 with self.assertRaises(ValueError): 

2436 DatasetProvenance.from_flat_dict(prov_dict, butler) 

2437 

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

2439 # separators. 

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

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

2442 self.assertEqual(ref_from_prov, metric_ref2) 

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

2444 

2445 with self.assertRaises(ValueError): 

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

2447 

2448 def test_specialized_file_datasets_functions(self): 

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

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

2451 another, without immediately inserting registry entries for the 

2452 datasets. 

2453 """ 

2454 repo = MetricTestRepo.create_from_butler( 

2455 self.create_empty_butler(writeable=True), 

2456 self.tmpConfigFile, 

2457 "StructuredCompositeReadCompNoDisassembly", 

2458 ) 

2459 source_butler = repo.butler 

2460 

2461 # Test writing outputs to a FileDatastore. 

2462 with tempfile.TemporaryDirectory() as tempdir: 

2463 target_repo_config = Butler.makeRepo(tempdir) 

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

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

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

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

2468 for dataset in datasets: 

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

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

2471 self.assertFalse(path.isabs()) 

2472 # Files should have been copied into the target datastore 

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

2474 

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

2476 target_butler = Butler(target_repo_config, writeable=True) 

2477 self.enterContext(target_butler) 

2478 target_butler.transfer_dimension_records_from(source_butler, refs) 

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

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

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

2482 

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

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

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

2486 

2487 # Test writing outputs to a ChainedDatastore. 

2488 with tempfile.TemporaryDirectory() as tempdir: 

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

2490 # multiple datastore roots. 

2491 dt1 = repo.datasetType 

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

2493 source_butler.registry.registerDatasetType(dt2) 

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

2495 config = Config.fromString( 

2496 f""" 

2497 datastore: 

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

2499 datastore_constraints: 

2500 - constraints: 

2501 accept: 

2502 - {dt1.name} 

2503 - constraints: 

2504 accept: 

2505 - {dt2.name} 

2506 datastores: 

2507 - datastore: 

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

2509 root: <butlerRoot>/FileDatastore_0 

2510 - datastore: 

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

2512 root: <butlerRoot>/FileDatastore_1 

2513 """ 

2514 ) 

2515 target_repo_config = Butler.makeRepo(tempdir, config) 

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

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

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

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

2520 for dataset in datasets: 

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

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

2523 self.assertFalse(path.isabs()) 

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

2525 # in the chain. 

2526 datastore_root = ResourcePath(tempdir) 

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

2528 datastore_root = datastore_root.join("FileDatastore_0") 

2529 else: 

2530 datastore_root = datastore_root.join("FileDatastore_1") 

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

2532 

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

2534 target_butler = Butler(target_repo_config, writeable=True) 

2535 self.enterContext(target_butler) 

2536 target_butler.transfer_dimension_records_from(source_butler, refs) 

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

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

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

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

2541 

2542 def test_temporary_for_ingest(self) -> None: 

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

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

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

2546 butler.registry.registerDatasetType(dataset_type) 

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

2548 with TemporaryForIngest(butler, ref) as temporary: 

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

2550 found = TemporaryForIngest.find_orphaned_temporaries_by_ref(ref, butler) 

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

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

2553 temporary.ingest() 

2554 loaded = butler.get(ref) 

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

2556 

2557 

2558class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

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

2560 

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

2562 fullConfigKey = ".datastore.formatters" 

2563 validationCanFail = True 

2564 datastoreStr = ["/tmp"] 

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

2566 registryStr = "PostgreSQL@test" 

2567 

2568 @classmethod 

2569 def setUpClass(cls) -> None: 

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

2571 super().setUpClass() 

2572 

2573 def setUp(self) -> None: 

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

2575 self._temp_config = False 

2576 config = Config(self.configFile) 

2577 self.postgresql.patch_butler_config(config) 

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

2579 config.dump(fh) 

2580 self.configFile = fh.name 

2581 self._temp_config = True 

2582 super().setUp() 

2583 

2584 def tearDown(self) -> None: 

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

2586 os.remove(self.configFile) 

2587 super().tearDown() 

2588 

2589 def testMakeRepo(self) -> None: 

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

2591 # the config file is acceptable to sqlite. 

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

2593 

2594 

2595class ClonedPostgresPosixDatastoreButlerTestCase(PostgresPosixDatastoreButlerTestCase, unittest.TestCase): 

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

2597 

2598 def create_butler( 

2599 self, 

2600 run: str, 

2601 storageClass: StorageClass | str, 

2602 datasetTypeName: str, 

2603 metrics: ButlerMetrics | None = None, 

2604 ) -> tuple[DirectButler, DatasetType]: 

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

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

2607 

2608 

2609class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase): 

2610 """InMemoryDatastore specialization of a butler""" 

2611 

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

2613 fullConfigKey = None 

2614 useTempRoot = False 

2615 validationCanFail = False 

2616 datastoreStr = ["datastore='InMemory"] 

2617 datastoreName = ["InMemoryDatastore@"] 

2618 registryStr = "/gen3.sqlite3" 

2619 

2620 def testIngest(self) -> None: 

2621 pass 

2622 

2623 def test_ingest_zip(self) -> None: 

2624 pass 

2625 

2626 

2627class ClonedSqliteButlerTestCase(InMemoryDatastoreButlerTestCase, unittest.TestCase): 

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

2629 

2630 def create_butler( 

2631 self, 

2632 run: str, 

2633 storageClass: StorageClass | str, 

2634 datasetTypeName: str, 

2635 metrics: ButlerMetrics | None = None, 

2636 ) -> tuple[DirectButler, DatasetType]: 

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

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

2639 

2640 

2641class ChainedDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

2642 """PosixDatastore specialization""" 

2643 

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

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

2646 validationCanFail = True 

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

2648 datastoreName = [ 

2649 "InMemoryDatastore@", 

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

2651 "SecondDatastore", 

2652 ] 

2653 registryStr = "/gen3.sqlite3" 

2654 

2655 def testPruneDatasets(self) -> None: 

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

2657 # impossible for this configuration because of the InMemoryDatastore in 

2658 # the ChainedDatastore. 

2659 pass 

2660 

2661 

2662class ButlerExplicitRootTestCase(PosixDatastoreButlerTestCase): 

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

2664 

2665 datastoreStr = ["dir1"] 

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

2667 # butler.yaml as the config name. 

2668 fullConfigKey = None 

2669 

2670 def setUp(self) -> None: 

2671 self.root = makeTestTempDir(TESTDIR) 

2672 

2673 # Make a new repository in one place 

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

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

2676 

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

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

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

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

2681 config = Config(configFile1) 

2682 config["root"] = self.dir1 

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

2684 config.dumpToUri(configFile2) 

2685 os.remove(configFile1) 

2686 self.tmpConfigFile = configFile2 

2687 

2688 def testFileLocations(self) -> None: 

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

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

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

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

2693 

2694 

2695class ButlerMakeRepoOutfileTestCase(ButlerPutGetTests, unittest.TestCase): 

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

2697 

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

2699 

2700 def setUp(self) -> None: 

2701 self.root = makeTestTempDir(TESTDIR) 

2702 self.root2 = makeTestTempDir(TESTDIR) 

2703 

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

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

2706 

2707 def tearDown(self) -> None: 

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

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

2710 super().tearDown() 

2711 

2712 def testConfigExistence(self) -> None: 

2713 c = Config(self.tmpConfigFile) 

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

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

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

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

2718 

2719 def testPutGet(self) -> None: 

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

2721 self.runPutGetTest(storageClass, "test_metric") 

2722 

2723 

2724class ButlerMakeRepoOutfileDirTestCase(ButlerMakeRepoOutfileTestCase): 

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

2726 

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

2728 

2729 def setUp(self) -> None: 

2730 self.root = makeTestTempDir(TESTDIR) 

2731 self.root2 = makeTestTempDir(TESTDIR) 

2732 

2733 self.tmpConfigFile = self.root2 

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

2735 

2736 def testConfigExistence(self) -> None: 

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

2738 # type. 

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

2740 super().testConfigExistence() 

2741 

2742 

2743class ButlerMakeRepoOutfileUriTestCase(ButlerMakeRepoOutfileTestCase): 

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

2745 

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

2747 

2748 def setUp(self) -> None: 

2749 self.root = makeTestTempDir(TESTDIR) 

2750 self.root2 = makeTestTempDir(TESTDIR) 

2751 

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

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

2754 

2755 

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

2757class S3DatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase): 

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

2759 a local in-memory SqlRegistry. 

2760 """ 

2761 

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

2763 fullConfigKey = None 

2764 validationCanFail = True 

2765 

2766 bucketName = "anybucketname" 

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

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

2769 """ 

2770 

2771 root = "butlerRoot/" 

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

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

2774 during set-up. 

2775 """ 

2776 

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

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

2779 returned by Butler stringification. 

2780 """ 

2781 

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

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

2784 

2785 registryStr = "/gen3.sqlite3" 

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

2787 

2788 mock_aws = mock_aws() 

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

2790 

2791 def genRoot(self) -> str: 

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

2793 name for the temporary bucket repo. 

2794 

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

2796 becomes when useTempRoot is True. 

2797 """ 

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

2799 return rndstr + "/" 

2800 

2801 def setUp(self) -> None: 

2802 config = Config(self.configFile) 

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

2804 self.bucketName = uri.netloc 

2805 

2806 # Enable S3 mocking of tests. 

2807 self.enterContext(clean_test_environment_for_s3()) 

2808 self.mock_aws.start() 

2809 

2810 if self.useTempRoot: 

2811 self.root = self.genRoot() 

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

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

2814 

2815 # need local folder to store registry database 

2816 self.reg_dir = makeTestTempDir(TESTDIR) 

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

2818 

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

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

2821 s3 = boto3.resource("s3") 

2822 s3.create_bucket(Bucket=self.bucketName) 

2823 

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

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

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

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

2828 

2829 def tearDown(self) -> None: 

2830 s3 = boto3.resource("s3") 

2831 bucket = s3.Bucket(self.bucketName) 

2832 try: 

2833 bucket.objects.all().delete() 

2834 except botocore.exceptions.ClientError as e: 

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

2836 # the key was not reachable - pass 

2837 pass 

2838 else: 

2839 raise 

2840 

2841 bucket = s3.Bucket(self.bucketName) 

2842 bucket.delete() 

2843 

2844 # Stop the S3 mock. 

2845 self.mock_aws.stop() 

2846 

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

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

2849 

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

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

2852 

2853 super().tearDown() 

2854 

2855 

2856class DatastoreTransfers(TestCaseMixin): 

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

2858 for specific configurations are in other classes, below. 

2859 """ 

2860 

2861 storageClassFactory: StorageClassFactory 

2862 

2863 @classmethod 

2864 def setUpClass(cls) -> None: 

2865 cls.storageClassFactory = StorageClassFactory() 

2866 

2867 def setUp(self) -> None: 

2868 self.root = makeTestTempDir(TESTDIR) 

2869 self.config = Config(self.configFile) 

2870 

2871 # Some tests cause convertors to be replaced so ensure 

2872 # the storage class factory is reset each time. 

2873 self.storageClassFactory.reset() 

2874 self.storageClassFactory.addFromConfig(self.configFile) 

2875 

2876 def tearDown(self) -> None: 

2877 removeTestTempDir(self.root) 

2878 

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

2880 if manager is None: 

2881 manager = ( 

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

2883 ) 

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

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

2886 butler = Butler.from_config( 

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

2888 ) 

2889 self.enterContext(butler) 

2890 return butler 

2891 

2892 def assertButlerTransfers( 

2893 self, 

2894 purge: bool = False, 

2895 storageClassName: str = "StructuredData", 

2896 storageClassNameTarget: str | None = None, 

2897 ) -> None: 

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

2899 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

2900 if storageClassNameTarget is not None: 

2901 storageClassTarget = self.storageClassFactory.getStorageClass(storageClassNameTarget) 

2902 else: 

2903 storageClassTarget = storageClass 

2904 

2905 datasetTypeName = "random_data" 

2906 

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

2908 # two of those three. 

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

2910 

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

2912 # grouping works. 

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

2914 

2915 # Create the run collections in the source butler. 

2916 for run in runs: 

2917 self.source_butler.collections.register(run) 

2918 

2919 # Create dimensions in source butler. 

2920 n_exposures = 30 

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

2922 self.source_butler.registry.insertDimensionData( 

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

2924 ) 

2925 self.source_butler.registry.insertDimensionData( 

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

2927 ) 

2928 self.source_butler.registry.insertDimensionData( 

2929 "day_obs", 

2930 { 

2931 "instrument": "DummyCamComp", 

2932 "id": 20250101, 

2933 }, 

2934 ) 

2935 

2936 for i in range(n_exposures): 

2937 self.source_butler.registry.insertDimensionData( 

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

2939 ) 

2940 self.source_butler.registry.insertDimensionData( 

2941 "exposure", 

2942 { 

2943 "instrument": "DummyCamComp", 

2944 "id": i, 

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

2946 "physical_filter": "d-r", 

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

2948 "day_obs": 20250101, 

2949 }, 

2950 ) 

2951 

2952 # Create dataset types in the source butler. 

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

2954 for datasetTypeName in datasetTypeNames: 

2955 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

2956 self.source_butler.registry.registerDatasetType(datasetType) 

2957 

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

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

2960 # Will not be relevant for UUID. 

2961 run = "distraction" 

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

2963 self.enterContext(butler) 

2964 butler.put( 

2965 makeExampleMetrics(), 

2966 datasetTypeName, 

2967 exposure=1, 

2968 instrument="DummyCamComp", 

2969 physical_filter="d-r", 

2970 ) 

2971 

2972 # Write some example metrics to the source 

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

2974 self.enterContext(butler) 

2975 

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

2977 # but which will not be transferred. 

2978 deleted: set[DatasetRef] = set() 

2979 

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

2981 source_refs = [] 

2982 for i in range(n_exposures): 

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

2984 # two thirds. 

2985 index = i % 3 

2986 run = runs[index] 

2987 datasetTypeName = datasetTypeNames[i % 2] 

2988 

2989 metric = MetricsExample( 

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

2991 ) 

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

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

2994 

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

2996 # for a specific index. 

2997 if purge and index == 1: 

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

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

3000 # file out. 

3001 # Access the individual datastores. 

3002 datastores = [] 

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

3004 datastores.extend(butler._datastore.datastores) 

3005 else: 

3006 datastores.append(butler._datastore) 

3007 

3008 if not deleted: 

3009 # For a chained datastore we need to remove 

3010 # files in each chain. 

3011 for datastore in datastores: 

3012 # The file might not be known to the datastore 

3013 # if constraints are used. 

3014 try: 

3015 primary, uris = datastore.getURIs(ref) 

3016 except FileNotFoundError: 

3017 continue 

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

3019 primary.remove() 

3020 for uri in uris.values(): 

3021 if uri.scheme != "mem": 

3022 uri.remove() 

3023 n_expected -= 1 

3024 deleted.add(ref) 

3025 

3026 # Remove the datastore record. 

3027 for datastore in datastores: 

3028 if hasattr(datastore, "removeStoredItemInfo"): 

3029 datastore.removeStoredItemInfo(ref) 

3030 

3031 if index < 2: 

3032 source_refs.append(ref) 

3033 if ref not in deleted: 

3034 new_metric = butler.get(ref) 

3035 self.assertEqual(new_metric, metric) 

3036 

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

3038 # definitions. 

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

3040 for datasetTypeName in datasetTypeNames: 

3041 datasetType = DatasetType(datasetTypeName, dimensions, badStorageClass) 

3042 self.target_butler.registry.registerDatasetType(datasetType) 

3043 with self.assertRaises(ConflictingDefinitionError) as cm: 

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

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

3046 

3047 # And remove the bad definitions. 

3048 for datasetTypeName in datasetTypeNames: 

3049 self.target_butler.registry.removeDatasetType(datasetTypeName) 

3050 

3051 # Transfer without creating dataset types should fail. 

3052 with self.assertRaises(KeyError): 

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

3054 

3055 # Transfer without creating dimensions should fail. 

3056 with self.assertRaises(ConflictingDefinitionError) as cm: 

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

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

3059 

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

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

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

3063 if storageClass != storageClassTarget: 

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

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

3066 self.target_butler.registry.registerDatasetType(datasetType) 

3067 

3068 # The failed transfer above leaves registry in an inconsistent 

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

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

3071 # Can remove with DM-35498. 

3072 self.target_butler.registry.refresh() 

3073 

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

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

3076 

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

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

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

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

3081 

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

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

3084 transferred = self.target_butler.transfer_from( 

3085 self.source_butler, 

3086 source_refs, 

3087 register_dataset_types=True, 

3088 transfer_dimensions=True, 

3089 ) 

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

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

3092 

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

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

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

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

3097 

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

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

3100 # dataset_id. 

3101 if purge: 

3102 # This should not need to register dataset types. 

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

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

3105 

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

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

3108 

3109 with self.assertRaises(ValueError): 

3110 self.target_butler._datastore.transfer_from( 

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

3112 ) 

3113 

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

3115 for ref in source_refs: 

3116 if ref not in deleted: 

3117 new_metric = self.target_butler.get(ref) 

3118 old_metric = self.source_butler.get(ref) 

3119 self.assertEqual(new_metric, old_metric) 

3120 

3121 # Try again without implicit storage class conversion 

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

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

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

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

3126 new_metric = self.target_butler.get(target_ref) 

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

3128 

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

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

3131 self.target_butler.pruneDatasets( 

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

3133 ) 

3134 self.target_butler.put(new_metric, target_ref) 

3135 new_new_metric = self.target_butler.get(target_ref) 

3136 new_old_metric = self.target_butler.get( 

3137 target_ref, storageClass=ref.datasetType.storageClass 

3138 ) 

3139 self.assertEqual(new_new_metric, new_metric) 

3140 self.assertEqual(new_old_metric, old_metric) 

3141 

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

3143 # This should block the transfer. 

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

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

3146 with self.assertRaises(CollectionTypeError): 

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

3148 # use integer IDs so filter those out. 

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

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

3151 

3152 

3153class PosixDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3154 """Test data transfers between butlers. 

3155 

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

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

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

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

3160 dataset types. The test ignores that. 

3161 """ 

3162 

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

3164 

3165 def create_butlers( 

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

3167 ) -> None: 

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

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

3170 

3171 def testTransferUuidToUuid(self) -> None: 

3172 self.create_butlers() 

3173 self.assertButlerTransfers() 

3174 

3175 def testTransferFromChainedUuidToUuid(self) -> None: 

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

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

3178 self.assertButlerTransfers() 

3179 

3180 def testTransferFromIncompatibleUuidToUuid(self) -> None: 

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

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

3183 with self.assertRaises(NotImplementedError): 

3184 self.assertButlerTransfers() 

3185 

3186 def testTransferFromIncompatibleChainUuidToUuid(self) -> None: 

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

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

3189 with self.assertRaises(TypeError): 

3190 self.assertButlerTransfers() 

3191 

3192 def testTransferFromFileUuidToUuid(self) -> None: 

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

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

3195 self.assertButlerTransfers() 

3196 

3197 def testTransferMissing(self) -> None: 

3198 """Test transfers where datastore records are missing. 

3199 

3200 This is how execution butler works. 

3201 """ 

3202 self.create_butlers() 

3203 

3204 # Configure the source butler to allow trust. 

3205 self.source_butler._datastore._set_trust_mode(True) 

3206 

3207 self.assertButlerTransfers(purge=True) 

3208 

3209 def testTransferMissingDisassembly(self) -> None: 

3210 """Test transfers where datastore records are missing. 

3211 

3212 This is how execution butler works. 

3213 """ 

3214 self.create_butlers() 

3215 

3216 # Configure the source butler to allow trust. 

3217 self.source_butler._datastore._set_trust_mode(True) 

3218 

3219 # Test disassembly. 

3220 self.assertButlerTransfers(purge=True, storageClassName="StructuredComposite") 

3221 

3222 def testTransferDifferingStorageClasses(self) -> None: 

3223 """Test transfers when the source butler dataset type has a different 

3224 but compatible storage class. 

3225 """ 

3226 self.create_butlers() 

3227 

3228 self.assertButlerTransfers(storageClassNameTarget="MetricsConversion") 

3229 

3230 def testTransferDifferingStorageClassesDisassembly(self) -> None: 

3231 """Test transfers when the source butler dataset type has a different 

3232 but compatible storage class and where the source butler has 

3233 disassembled. 

3234 """ 

3235 self.create_butlers() 

3236 

3237 self.assertButlerTransfers( 

3238 storageClassName="StructuredComposite", storageClassNameTarget="MetricsConversion" 

3239 ) 

3240 

3241 def testAbsoluteURITransferDirect(self) -> None: 

3242 """Test transfer using an absolute URI.""" 

3243 self._absolute_transfer("auto") 

3244 

3245 def testAbsoluteURITransferCopy(self) -> None: 

3246 """Test transfer using an absolute URI.""" 

3247 self._absolute_transfer("copy") 

3248 

3249 def _absolute_transfer(self, transfer: str) -> None: 

3250 self.create_butlers() 

3251 

3252 storageClassName = "StructuredData" 

3253 storageClass = self.storageClassFactory.getStorageClass(storageClassName) 

3254 datasetTypeName = "random_data" 

3255 run = "run1" 

3256 self.source_butler.collections.register(run) 

3257 

3258 dimensions = self.source_butler.dimensions.conform(()) 

3259 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3260 self.source_butler.registry.registerDatasetType(datasetType) 

3261 

3262 metrics = makeExampleMetrics() 

3263 with ResourcePath.temporary_uri(suffix=".json") as temp: 

3264 dataId = DataCoordinate.make_empty(self.source_butler.dimensions) 

3265 source_refs = [DatasetRef(datasetType, dataId, run=run)] 

3266 temp.write(json.dumps(metrics.exportAsDict()).encode()) 

3267 dataset = FileDataset(path=temp, refs=source_refs) 

3268 self.source_butler.ingest(dataset, transfer="direct") 

3269 

3270 self.target_butler.transfer_from( 

3271 self.source_butler, dataset.refs, register_dataset_types=True, transfer=transfer 

3272 ) 

3273 

3274 uri = self.target_butler.getURI(dataset.refs[0]) 

3275 if transfer == "auto": 

3276 self.assertEqual(uri, temp) 

3277 else: 

3278 self.assertNotEqual(uri, temp) 

3279 

3280 def test_shared_dimension_group(self): 

3281 """Test internal logic that divides dataset types by dimension group 

3282 when doing registry updates. 

3283 """ 

3284 self.create_butlers() 

3285 self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True) 

3286 self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True) 

3287 

3288 source_butler = self.source_butler 

3289 target_butler = self.target_butler 

3290 

3291 # Create a dataset type with the same dimensions as the 'bias' dataset 

3292 # type from base.yaml 

3293 dataset_type = DatasetType( 

3294 "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions 

3295 ) 

3296 source_butler.registry.registerDatasetType(dataset_type) 

3297 # This has the same data ID as one of the bias datasets in 

3298 # datasets.yaml. 

3299 test_ref = source_butler.registry.insertDatasets( 

3300 "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g" 

3301 )[0] 

3302 

3303 biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"]) 

3304 flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"]) 

3305 refs = [test_ref, *biases, *flats] 

3306 

3307 # Test setup will be even more convoluted if we want the datastore to 

3308 # actually transfer files. For testing the dimension group behavior, 

3309 # we really only care about the registry. 

3310 with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock: 

3311 mock.return_value = (set(refs), set()) 

3312 target_butler.transfer_from( 

3313 source_butler, 

3314 refs, 

3315 transfer=None, 

3316 register_dataset_types=True, 

3317 skip_missing=False, 

3318 transfer_dimensions=True, 

3319 ) 

3320 

3321 transferred_test_ref = target_butler.find_dataset( 

3322 "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3323 ) 

3324 self.assertEqual(transferred_test_ref.id, test_ref.id) 

3325 

3326 transferred_bias = target_butler.find_dataset( 

3327 "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g" 

3328 ) 

3329 self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd")) 

3330 

3331 transferred_flat = target_butler.find_dataset( 

3332 "flat", 

3333 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"}, 

3334 collections="imported_r", 

3335 ) 

3336 self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840")) 

3337 

3338 

3339class ChainedDatastoreTransfers(PosixDatastoreTransfers): 

3340 """Test transfers using a chained datastore.""" 

3341 

3342 configFile = os.path.join(TESTDIR, "config/basic/butler-chained.yaml") 

3343 

3344 

3345@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3346class ButlerServerDatastoreTransfers(DatastoreTransfers, unittest.TestCase): 

3347 """Test ``transfer_from`` involving Butler server.""" 

3348 

3349 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

3350 

3351 def test_transfers_from_remote_to_direct(self) -> None: 

3352 from lsst.daf.butler.remote_butler._remote_file_transfer_source import ( 

3353 mock_file_transfer_uris_for_unit_test, 

3354 ) 

3355 

3356 self.target_butler = self.create_butler(None, "2") 

3357 with create_test_server(TESTDIR) as server: 

3358 self.source_butler = server.hybrid_butler 

3359 

3360 def _remap_transfer_url(path: HttpResourcePath) -> HttpResourcePath: 

3361 # The Butler server returns HTTP URIs with a domain name that 

3362 # is not resolvable because there is no actual HTTP server 

3363 # involved in these tests. Strip this first layer of 

3364 # indirection, and return the target of the redirect instead. 

3365 response = server.client.get(str(path), follow_redirects=False, headers=path._extra_headers) 

3366 return ResourcePath(str(response.next_request.url)) 

3367 

3368 with mock_file_transfer_uris_for_unit_test(_remap_transfer_url): 

3369 self.assertButlerTransfers() 

3370 

3371 

3372class NullDatastoreTestCase(unittest.TestCase): 

3373 """Test that we can fall back to a null datastore.""" 

3374 

3375 # Need a good config to create the repo. 

3376 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

3377 storageClassFactory: StorageClassFactory 

3378 

3379 @classmethod 

3380 def setUpClass(cls) -> None: 

3381 cls.storageClassFactory = StorageClassFactory() 

3382 cls.storageClassFactory.addFromConfig(cls.configFile) 

3383 

3384 def setUp(self) -> None: 

3385 """Create a new butler root for each test.""" 

3386 self.root = makeTestTempDir(TESTDIR) 

3387 Butler.makeRepo(self.root, config=Config(self.configFile)) 

3388 

3389 def tearDown(self) -> None: 

3390 removeTestTempDir(self.root) 

3391 

3392 def test_fallback(self) -> None: 

3393 # Read the butler config and mess with the datastore section. 

3394 config_path = os.path.join(self.root, "butler.yaml") 

3395 bad_config = Config(config_path) 

3396 bad_config["datastore", "cls"] = "lsst.not.a.datastore.Datastore" 

3397 bad_config.dumpToUri(config_path) 

3398 

3399 with self.assertRaises(RuntimeError): 

3400 Butler(self.root, without_datastore=False) 

3401 

3402 with self.assertRaises(RuntimeError): 

3403 Butler.from_config(self.root, without_datastore=False) 

3404 

3405 butler = Butler.from_config(self.root, writeable=True, without_datastore=True) 

3406 self.enterContext(butler) 

3407 self.assertIsInstance(butler._datastore, NullDatastore) 

3408 

3409 # Check that registry is working. 

3410 butler.collections.register("MYRUN") 

3411 collections = butler.collections.query("*") 

3412 self.assertIn("MYRUN", set(collections)) 

3413 

3414 # Create a ref. 

3415 dimensions = butler.dimensions.conform([]) 

3416 storageClass = self.storageClassFactory.getStorageClass("StructuredDataDict") 

3417 datasetTypeName = "metric" 

3418 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

3419 butler.registry.registerDatasetType(datasetType) 

3420 ref = DatasetRef(datasetType, {}, run="MYRUN") 

3421 

3422 # Check that datastore will complain. 

3423 with self.assertRaises(FileNotFoundError): 

3424 butler.get(ref) 

3425 with self.assertRaises(FileNotFoundError): 

3426 butler.getURI(ref) 

3427 

3428 

3429@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3430class ButlerServerTests(FileDatastoreButlerTests): 

3431 """Test RemoteButler and Butler server.""" 

3432 

3433 configFile = None 

3434 predictionSupported = False 

3435 trustModeSupported = False 

3436 

3437 postgres: TemporaryPostgresInstance | None 

3438 

3439 def setUp(self): 

3440 self.server_instance = self.enterContext(create_test_server(TESTDIR)) 

3441 

3442 def tearDown(self): 

3443 pass 

3444 

3445 def are_uris_equivalent(self, uri1: ResourcePath, uri2: ResourcePath) -> bool: 

3446 # S3 pre-signed URLs may end up with differing expiration times in the 

3447 # query parameters, so ignore query parameters when comparing. 

3448 return uri1.scheme == uri2.scheme and uri1.netloc == uri2.netloc and uri1.path == uri2.path 

3449 

3450 def create_empty_butler( 

3451 self, 

3452 run: str | None = None, 

3453 writeable: bool | None = None, 

3454 metrics: ButlerMetrics | None = None, 

3455 cleanup: bool = True, 

3456 ) -> Butler: 

3457 return self.server_instance.hybrid_butler.clone(run=run, metrics=metrics) 

3458 

3459 def remove_dataset_out_of_band(self, butler: Butler, ref: DatasetRef) -> None: 

3460 # Can't delete a file via S3 signed URLs, so we need to reach in 

3461 # through DirectButler to delete the dataset. 

3462 uri = self.server_instance.direct_butler.getURI(ref) 

3463 uri.remove() 

3464 

3465 def testConstructor(self): 

3466 # RemoteButler constructor is tested in test_server.py and 

3467 # test_remote_butler.py. 

3468 pass 

3469 

3470 def testDafButlerRepositories(self): 

3471 # Loading of RemoteButler via repository index is tested in 

3472 # test_server.py. 

3473 pass 

3474 

3475 def testGetDatasetTypes(self) -> None: 

3476 # This is mostly a test of validateConfiguration, which is for 

3477 # validating Datastore configuration and thus isn't relevant to 

3478 # RemoteButler. 

3479 pass 

3480 

3481 def testMakeRepo(self) -> None: 

3482 # Only applies to DirectButler. 

3483 pass 

3484 

3485 # Pickling not yet implemented for RemoteButler/HybridButler. 

3486 @unittest.expectedFailure 

3487 def testPickle(self) -> None: 

3488 return super().testPickle() 

3489 

3490 def testStringification(self) -> None: 

3491 self.assertEqual( 

3492 str(self.server_instance.remote_butler), 

3493 "RemoteButler(https://test.example/api/butler/repo/testrepo/)", 

3494 ) 

3495 

3496 def testTransaction(self) -> None: 

3497 # Transactions will never be supported for RemoteButler. 

3498 pass 

3499 

3500 def testPutTemplates(self) -> None: 

3501 # The Butler server instance is configured with different file naming 

3502 # templates than this test is expecting. 

3503 pass 

3504 

3505 

3506@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3507class ButlerServerSqliteTests(ButlerServerTests, unittest.TestCase): 

3508 """Tests for RemoteButler's registry shim, with a SQLite DB backing the 

3509 server. 

3510 """ 

3511 

3512 postgres = None 

3513 

3514 

3515@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

3516class ButlerServerPostgresTests(ButlerServerTests, unittest.TestCase): 

3517 """Tests for RemoteButler's registry shim, with a Postgres DB backing the 

3518 server. 

3519 """ 

3520 

3521 @classmethod 

3522 def setUpClass(cls): 

3523 cls.postgres = cls.enterClassContext(setup_postgres_test_db()) 

3524 super().setUpClass() 

3525 

3526 

3527def setup_module(module: types.ModuleType) -> None: 

3528 """Set up the module for pytest.""" 

3529 clean_environment() 

3530 

3531 

3532def _get_test_data_path(filename: str) -> ResourcePath: 

3533 return ResourcePath(f"resource://lsst.daf.butler/tests/registry_data/{filename}") 

3534 

3535 

3536if __name__ == "__main__": 

3537 clean_environment() 

3538 unittest.main()