Coverage for python/lsst/daf/butler/registry/tests/

1# This file is part of daf_butler.

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.

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/>.

27from __future__ import annotations

29from ... import ddl

31__all__ = ["RegistryTests"]

33import contextlib

34import datetime

35import itertools

36import re

37import time

38import unittest

39import uuid

40from abc import ABC, abstractmethod

41from collections import defaultdict, namedtuple

42from collections.abc import Callable, Iterator

43from concurrent.futures import ThreadPoolExecutor

44from contextlib import ExitStack

45from datetime import timedelta

46from threading import Barrier

47from typing import TypeVar

49import astropy.time

50import sqlalchemy

52try:

53 import numpy as np

54except ImportError:

55 np = None

57import lsst.sphgeom

59from ... import Butler

60from ..._collection_type import CollectionType

61from ..._dataset_association import DatasetAssociation

62from ..._dataset_ref import DatasetIdFactory, DatasetIdGenEnum, DatasetRef

63from ..._dataset_type import DatasetType

64from ..._exceptions import (

65 CalibrationLookupError,

66 CollectionTypeError,

67 DataIdValueError,

68 InconsistentDataIdError,

69 InvalidQueryError,

70 MissingCollectionError,

71 MissingDatasetTypeError,

72)

73from ..._exceptions_legacy import DatasetTypeError

74from ..._storage_class import StorageClass

75from ..._timespan import Timespan

76from ...dimensions import DataCoordinate, DataCoordinateSet, DimensionUniverse, SkyPixDimension

77from ...direct_butler import DirectButler

78from .._collection_summary import CollectionSummary

79from .._config import RegistryConfig

80from .._exceptions import (

81 ArgumentError,

82 CollectionError,

83 ConflictingDefinitionError,

84 DatasetTypeExpressionError,

85 NoDefaultCollectionError,

86 OrphanedRecordError,

87)

88from ..interfaces import ButlerAttributeExistsError, ReadOnlyDatabaseError

89from ..sql_registry import SqlRegistry

91_T = TypeVar("_T")

94class RegistryTests(ABC):

95 """Generic tests for the `SqlRegistry` class that can be subclassed to

96 generate tests for different configurations.

97 """

99 collectionsManager: str | None = None

100 """Name of the collections manager class, if subclass provides value for

101 this member then it overrides name specified in default configuration

102 (`str`).

103 """

104

105 datasetsManager: str | dict[str, str] | None = None

106 """Name or configuration dictionary of the datasets manager class, if

107 subclass provides value for this member then it overrides name specified

108 in default configuration (`str` or `dict`).

109 """

110

111 supportsCollectionRegex: bool = False

112 """True if the registry class being tested supports regex searches for

113 collections."""

114

115 def makeRegistryConfig(self) -> RegistryConfig:

116 """Create RegistryConfig used to create a registry.

117

118 This method should be called by a subclass from `makeRegistry`.

119 Returned instance will be pre-configured based on the values of class

120 members, and default-configured for all other parameters. Subclasses

121 that need default configuration should just instantiate

122 `RegistryConfig` directly.

123 """

124 config = RegistryConfig()

125 if self.collectionsManager:

126 config["managers", "collections"] = self.collectionsManager

127 if self.datasetsManager:

128 config["managers", "datasets"] = self.datasetsManager

129 return config

130

131 @abstractmethod

132 def make_butler(self, registry_config: RegistryConfig | None = None) -> Butler:

133 """Return the butler to be tested.

134

135 Parameters

136 ----------

137 registry_config : `RegistryConfig`, optional

138 Registry configuration used when instantiating the Butler.

139

140 Returns

141 -------

142 butler : `~lsst.daf.butler.Butler`

143 The butler with a registry to be tested.

144 """

145 raise NotImplementedError()

146

147 def load_data(self, butler: Butler, *filenames: str) -> None:

148 """Load registry test data from

149 ``resource://lsst.daf.butler/tests/registry_data/<filename>``,

150 which should be a YAML import/export file.

151

152 Parameters

153 ----------

154 butler : `Butler`

155 The butler to load into.

156 *filenames : `str`

157 The names of the files to load.

158 """

159 for filename in filenames:

160 butler.import_(

161 filename=f"resource://lsst.daf.butler/tests/registry_data/{filename}", without_datastore=True

162 )

163

164 def checkQueryResults(self, results, expected):

165 """Check that a query results object contains expected values.

166

167 Parameters

168 ----------

169 results : `DataCoordinateQueryResults` or `DatasetQueryResults`

170 A lazy-evaluation query results object.

171 expected : `list`

172 A list of `DataCoordinate` o `DatasetRef` objects that should be

173 equal to results of the query, aside from ordering.

174 """

175 self.assertCountEqual(list(results), expected)

176 self.assertEqual(results.count(), len(expected))

177 if expected:

178 self.assertTrue(results.any())

179 else:

180 self.assertFalse(results.any())

181

182 def testOpaque(self):

183 """Tests for `SqlRegistry.registerOpaqueTable`,

184 `SqlRegistry.insertOpaqueData`, `SqlRegistry.fetchOpaqueData`, and

185 `SqlRegistry.deleteOpaqueData`.

186 """

187 butler = self.make_butler()

188 registry = butler._registry

189 table = "opaque_table_for_testing"

190 registry.registerOpaqueTable(

191 table,

192 spec=ddl.TableSpec(

193 fields=[

194 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, primaryKey=True),

195 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=16, nullable=False),

196 ddl.FieldSpec("count", dtype=sqlalchemy.SmallInteger, nullable=True),

197 ],

198 ),

199 )

200 rows = [

201 {"id": 1, "name": "one", "count": None},

202 {"id": 2, "name": "two", "count": 5},

203 {"id": 3, "name": "three", "count": 6},

204 ]

205 registry.insertOpaqueData(table, *rows)

206 self.assertCountEqual(rows, list(registry.fetchOpaqueData(table)))

207 self.assertEqual(rows[0:1], list(registry.fetchOpaqueData(table, id=1)))

208 self.assertEqual(rows[1:2], list(registry.fetchOpaqueData(table, name="two")))

209 self.assertEqual(rows[0:1], list(registry.fetchOpaqueData(table, id=(1, 3), name=("one", "two"))))

210 self.assertEqual(rows, list(registry.fetchOpaqueData(table, id=(1, 2, 3))))

211 # Test very long IN clause which exceeds sqlite limit on number of

212 # parameters. SQLite says the limit is 32k but it looks like it is

213 # much higher.

214 self.assertEqual(rows, list(registry.fetchOpaqueData(table, id=list(range(300_000)))))

215 # Two IN clauses, each longer than 1k batch size, first with

216 # duplicates, second has matching elements in different batches (after

217 # sorting).

218 self.assertEqual(

219 rows[0:2],

220 list(

221 registry.fetchOpaqueData(

222 table,

223 id=list(range(1000)) + list(range(100, 0, -1)),

224 name=["one"] + [f"q{i}" for i in range(2200)] + ["two"],

225 )

226 ),

227 )

228 self.assertEqual([], list(registry.fetchOpaqueData(table, id=1, name="two")))

229 registry.deleteOpaqueData(table, id=3)

230 self.assertCountEqual(rows[:2], list(registry.fetchOpaqueData(table)))

231 registry.deleteOpaqueData(table)

232 self.assertEqual([], list(registry.fetchOpaqueData(table)))

233

234 def testDatasetType(self):

235 """Tests for `SqlRegistry.registerDatasetType` and

236 `SqlRegistry.getDatasetType`.

237 """

238 butler = self.make_butler()

239 registry = butler.registry

240 # Check valid insert

241 datasetTypeName = "test"

242 storageClass = StorageClass("testDatasetType")

243 registry.storageClasses.registerStorageClass(storageClass)

244 dimensions = registry.dimensions.conform(("instrument", "visit"))

245 differentDimensions = registry.dimensions.conform(("instrument", "patch"))

246 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass)

247 # Inserting for the first time should return True

248 self.assertTrue(registry.registerDatasetType(inDatasetType))

249 outDatasetType1 = registry.getDatasetType(datasetTypeName)

250 self.assertEqual(outDatasetType1, inDatasetType)

251

252 # Re-inserting should work

253 self.assertFalse(registry.registerDatasetType(inDatasetType))

254 # Except when they are not identical

255 with self.assertRaises(ConflictingDefinitionError):

256 nonIdenticalDatasetType = DatasetType(datasetTypeName, differentDimensions, storageClass)

257 registry.registerDatasetType(nonIdenticalDatasetType)

258

259 # Template can be None

260 datasetTypeName = "testNoneTemplate"

261 storageClass = StorageClass("testDatasetType2")

262 registry.storageClasses.registerStorageClass(storageClass)

263 dimensions = registry.dimensions.conform(("instrument", "visit"))

264 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass)

265 registry.registerDatasetType(inDatasetType)

266 outDatasetType2 = registry.getDatasetType(datasetTypeName)

267 self.assertEqual(outDatasetType2, inDatasetType)

268

269 allTypes = set(registry.queryDatasetTypes())

270 self.assertEqual(allTypes, {outDatasetType1, outDatasetType2})

271

272 # Test some basic queryDatasetTypes functionality

273 missing: list[str] = []

274 types = registry.queryDatasetTypes(["te*", "notarealdatasettype"], missing=missing)

275 self.assertCountEqual([dt.name for dt in types], ["test", "testNoneTemplate"])

276 self.assertEqual(missing, ["notarealdatasettype"])

277

278 # Trying to register a dataset type with different universe version or

279 # namespace will raise.

280 wrong_universes = (DimensionUniverse(version=-1), DimensionUniverse(namespace="🔭"))

281 for universe in wrong_universes:

282 storageClass = StorageClass("testDatasetType")

283 dataset_type = DatasetType(

284 "wrong_universe", ("instrument", "visit"), storageClass, universe=universe

285 )

286 with self.assertRaisesRegex(ValueError, "Incompatible dimension universe versions"):

287 registry.registerDatasetType(dataset_type)

288

289 def testDatasetTypeCache(self):

290 """Test for dataset type cache update logic after a cache miss."""

291 butler1 = self.make_butler()

292 butler2 = butler1.clone()

293 self.load_data(butler1, "base.yaml")

294

295 # Trigger full cache load.

296 butler2.get_dataset_type("flat")

297 # Have an external process register a dataset type.

298 butler1.registry.registerDatasetType(

299 DatasetType("test_type", ["instrument"], "int", universe=butler1.dimensions)

300 )

301 # Try to read the new dataset type -- this is a cache miss that

302 # triggers fetch of a single dataset type.

303 dt = butler2.get_dataset_type("test_type")

304 self.assertEqual(dt.name, "test_type")

305 self.assertEqual(list(dt.dimensions.names), ["instrument"])

306 # Read it again -- this time it should pull from the cache.

307 dt = butler2.get_dataset_type("test_type")

308 self.assertEqual(dt.name, "test_type")

309 self.assertEqual(list(dt.dimensions.names), ["instrument"])

310 # Do a query that uses the dataset type's tags table.

311 self.assertEqual(

312 butler2.query_datasets("test_type", collections="*", find_first=False, explain=False), []

313 )

314

315 def testDimensions(self):

316 """Tests for `SqlRegistry.insertDimensionData`,

317 `SqlRegistry.syncDimensionData`, and `SqlRegistry.expandDataId`.

318 """

319 butler = self.make_butler()

320 registry = butler.registry

321 dimensionName = "instrument"

322 dimension = registry.dimensions[dimensionName]

323 dimensionValue = {

324 "name": "DummyCam",

325 "visit_max": 10,

326 "visit_system": 0,

327 "exposure_max": 10,

328 "detector_max": 2,

329 "class_name": "lsst.pipe.base.Instrument",

330 }

331 registry.insertDimensionData(dimensionName, dimensionValue)

332 # Inserting the same value twice should fail

333 with self.assertRaises(sqlalchemy.exc.IntegrityError):

334 registry.insertDimensionData(dimensionName, dimensionValue)

335 # expandDataId should retrieve the record we just inserted

336 self.assertEqual(

337 registry.expandDataId(instrument="DummyCam", dimensions=dimension.minimal_group)

338 .records[dimensionName]

339 .toDict(),

340 dimensionValue,

341 )

342 # expandDataId should raise if there is no record with the given ID.

343 with self.assertRaises(DataIdValueError):

344 registry.expandDataId({"instrument": "Unknown"}, dimensions=dimension.minimal_group)

345 # band doesn't have a table; insert should fail.

346 with self.assertRaises(TypeError):

347 registry.insertDimensionData("band", {"band": "i"})

348 dimensionName2 = "physical_filter"

349 dimension2 = registry.dimensions[dimensionName2]

350 dimensionValue2 = {"name": "DummyCam_i", "band": "i"}

351 # Missing required dependency ("instrument") should fail

352 with self.assertRaises(KeyError):

353 registry.insertDimensionData(dimensionName2, dimensionValue2)

354 # Adding required dependency should fix the failure

355 dimensionValue2["instrument"] = "DummyCam"

356 registry.insertDimensionData(dimensionName2, dimensionValue2)

357 # expandDataId should retrieve the record we just inserted.

358 self.assertEqual(

359 registry.expandDataId(

360 instrument="DummyCam", physical_filter="DummyCam_i", dimensions=dimension2.minimal_group

361 )

362 .records[dimensionName2]

363 .toDict(),

364 dimensionValue2,

365 )

366 # Use syncDimensionData to insert a new record successfully.

367 dimensionName3 = "detector"

368 dimensionValue3 = {

369 "instrument": "DummyCam",

370 "id": 1,

371 "full_name": "one",

372 "name_in_raft": "zero",

373 "purpose": "SCIENCE",

374 }

375 self.assertTrue(registry.syncDimensionData(dimensionName3, dimensionValue3))

376 # Sync that again. Note that one field ("raft") is NULL, and that

377 # should be okay.

378 self.assertFalse(registry.syncDimensionData(dimensionName3, dimensionValue3))

379 # Now try that sync with the same primary key but a different value.

380 # This should fail.

381 with self.assertRaises(ConflictingDefinitionError):

382 registry.syncDimensionData(

383 dimensionName3,

384 {

385 "instrument": "DummyCam",

386 "id": 1,

387 "full_name": "one",

388 "name_in_raft": "four",

389 "purpose": "SCIENCE",

390 },

391 )

392

393 @unittest.skipIf(np is None, "numpy not available.")

394 def testNumpyDataId(self):

395 """Test that we can use a numpy int in a dataId."""

396 butler = self.make_butler()

397 registry = butler.registry

398 dimensionEntries = [

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

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

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

402 # Using an np.int64 here fails unless Records.fromDict is also

403 # patched to look for numbers.Integral

404 (

405 "visit",

406 {

407 "instrument": "DummyCam",

408 "id": 42,

409 "name": "fortytwo",

410 "physical_filter": "d-r",

411 "day_obs": 20250101,

412 },

413 ),

414 ]

415 for args in dimensionEntries:

416 registry.insertDimensionData(*args)

417

418 # Try a normal integer and something that looks like an int but

419 # is not.

420 for visit_id in (42, np.int64(42)):

421 with self.subTest(visit_id=visit_id, id_type=type(visit_id).__name__):

422 expanded = registry.expandDataId({"instrument": "DummyCam", "visit": visit_id})

423 self.assertEqual(expanded["visit"], int(visit_id))

424 self.assertIsInstance(expanded["visit"], int)

425

426 def testDataIdRelationships(self):

427 """Test that `SqlRegistry.expandDataId` raises an exception when the

428 given keys are inconsistent.

429 """

430 butler = self.make_butler()

431 self.load_data(butler, "base.yaml")

432 registry = butler.registry

433 # Insert a few more dimension records for the next test.

434 registry.insertDimensionData(

435 "day_obs",

436 {"instrument": "Cam1", "id": 20250101},

437 )

438 registry.insertDimensionData(

439 "group",

440 {"instrument": "Cam1", "name": "group1"},

441 )

442 registry.insertDimensionData(

443 "exposure",

444 {

445 "instrument": "Cam1",

446 "id": 1,

447 "obs_id": "one",

448 "physical_filter": "Cam1-G",

449 "group": "group1",

450 "day_obs": 20250101,

451 },

452 )

453 registry.insertDimensionData(

454 "group",

455 {"instrument": "Cam1", "name": "group2"},

456 )

457 registry.insertDimensionData(

458 "exposure",

459 {

460 "instrument": "Cam1",

461 "id": 2,

462 "obs_id": "two",

463 "physical_filter": "Cam1-G",

464 "group": "group2",

465 "day_obs": 20250101,

466 },

467 )

468 registry.insertDimensionData(

469 "visit_system",

470 {"instrument": "Cam1", "id": 0, "name": "one-to-one"},

471 )

472 registry.insertDimensionData(

473 "visit",

474 {"instrument": "Cam1", "id": 1, "name": "one", "physical_filter": "Cam1-G", "day_obs": 20250101},

475 )

476 registry.insertDimensionData(

477 "visit_definition",

478 {"instrument": "Cam1", "visit": 1, "exposure": 1},

479 )

480 with self.assertRaises(InconsistentDataIdError):

481 registry.expandDataId(

482 {"instrument": "Cam1", "visit": 1, "exposure": 2},

483 )

484

485 def testDataset(self):

486 """Basic tests for `SqlRegistry.insertDatasets`,

487 `SqlRegistry.getDataset`, and `SqlRegistry.removeDatasets`.

488 """

489 butler = self.make_butler()

490 registry = butler.registry

491 self.load_data(butler, "base.yaml")

492 run = "tésτ"

493 registry.registerRun(run)

494 datasetType = registry.getDatasetType("bias")

495 dataId = {"instrument": "Cam1", "detector": 2}

496 (ref,) = registry.insertDatasets(datasetType, dataIds=[dataId], run=run)

497 outRef = registry.getDataset(ref.id)

498 self.assertIsNotNone(ref.id)

499 self.assertEqual(ref, outRef)

500 with self.assertRaises(ConflictingDefinitionError):

501 registry.insertDatasets(datasetType, dataIds=[dataId], run=run)

502 registry.removeDatasets([ref])

503 self.assertIsNone(registry.findDataset(datasetType, dataId, collections=[run]))

504

505 def testFindDataset(self):

506 """Tests for `SqlRegistry.findDataset`."""

507 butler = self.make_butler()

508 registry = butler.registry

509 self.load_data(butler, "base.yaml")

510 run = "tésτ"

511 datasetType = registry.getDatasetType("bias")

512 dataId = {"instrument": "Cam1", "detector": 4}

513 registry.registerRun(run)

514 (inputRef,) = registry.insertDatasets(datasetType, dataIds=[dataId], run=run)

515 outputRef = registry.findDataset(datasetType, dataId, collections=[run])

516 self.assertEqual(outputRef, inputRef)

517 # Check that retrieval with invalid dataId raises

518 with self.assertRaises(LookupError):

519 dataId = {"instrument": "Cam1"} # no detector

520 registry.findDataset(datasetType, dataId, collections=run)

521 # Check that different dataIds match to different datasets

522 dataId1 = {"instrument": "Cam1", "detector": 1}

523 (inputRef1,) = registry.insertDatasets(datasetType, dataIds=[dataId1], run=run)

524 dataId2 = {"instrument": "Cam1", "detector": 2}

525 (inputRef2,) = registry.insertDatasets(datasetType, dataIds=[dataId2], run=run)

526 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef1)

527 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef2)

528 self.assertNotEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef2)

529 self.assertNotEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef1)

530 # Check that requesting a non-existing dataId returns None

531 nonExistingDataId = {"instrument": "Cam1", "detector": 3}

532 self.assertIsNone(registry.findDataset(datasetType, nonExistingDataId, collections=run))

533 # Search more than one collection, in which two have the right

534 # dataset type and another does not.

535 registry.registerRun("empty")

536 self.load_data(butler, "datasets.yaml")

537 bias1 = registry.findDataset("bias", instrument="Cam1", detector=2, collections=["imported_g"])

538 self.assertIsNotNone(bias1)

539 bias2 = registry.findDataset("bias", instrument="Cam1", detector=2, collections=["imported_r"])

540 self.assertIsNotNone(bias2)

541 self.assertEqual(

542 bias1,

543 registry.findDataset(

544 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_g", "imported_r"]

545 ),

546 )

547 self.assertEqual(

548 bias2,

549 registry.findDataset(

550 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_r", "imported_g"]

551 ),

552 )

553 # If the input data ID was an expanded DataCoordinate with records,

554 # then the output ref has records, too.

555 expanded_id = registry.expandDataId({"instrument": "Cam1", "detector": 2})

556 expanded_ref = registry.findDataset("bias", expanded_id, collections=["imported_r"])

557 self.assertTrue(expanded_ref.dataId.hasRecords())

558 # Search more than one collection, with one of them a CALIBRATION

559 # collection.

560 registry.registerCollection("Cam1/calib", CollectionType.CALIBRATION)

561 timespan = Timespan(

562 begin=astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai"),

563 end=astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai"),

564 )

565 registry.certify("Cam1/calib", [bias2], timespan=timespan)

566 self.assertEqual(

567 bias1,

568 registry.findDataset(

569 "bias",

570 instrument="Cam1",

571 detector=2,

572 collections=["empty", "imported_g", "Cam1/calib"],

573 timespan=timespan,

574 ),

575 )

576 self.assertEqual(

577 bias1,

578 registry.findDataset(

579 "bias",

580 instrument="Cam1",

581 detector=2,

582 # Calibration dataset type, with no calibration collection, but

583 # a timespan was provided.

584 collections=["imported_g"],

585 timespan=timespan,

586 ),

587 )

588 self.assertEqual(

589 bias2,

590 registry.findDataset(

591 "bias",

592 instrument="Cam1",

593 detector=2,

594 collections=["empty", "Cam1/calib", "imported_g"],

595 timespan=timespan,

596 ),

597 )

598 # If we try to search those same collections without a timespan, it

599 # should still work, since the CALIBRATION collection is ignored.

600 self.assertEqual(

601 bias1,

602 registry.findDataset(

603 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_g", "Cam1/calib"]

604 ),

605 )

606 self.assertEqual(

607 bias1,

608 registry.findDataset(

609 "bias", instrument="Cam1", detector=2, collections=["empty", "Cam1/calib", "imported_g"]

610 ),

611 )

612 self.assertIsNone(

613 registry.findDataset("bias", instrument="Cam1", detector=2, collections=["Cam1/calib"])

614 )

615 # Test non-calibration dataset type.

616 registry.registerDatasetType(

617 DatasetType("noncalibration", ["instrument", "detector"], "int", universe=butler.dimensions)

618 )

619 (non_calibration_ref,) = registry.insertDatasets("noncalibration", dataIds=[dataId2], run=run)

620 self.assertIsNone(

621 registry.findDataset("noncalibration", instrument="Cam1", detector=2, collections=["imported_g"])

622 )

623 self.assertEqual(

624 non_calibration_ref,

625 registry.findDataset("noncalibration", instrument="Cam1", detector=2, collections=[run]),

626 )

627 # Timespan parameter is ignored for non-calibration dataset types.

628 self.assertIsNone(

629 registry.findDataset(

630 "noncalibration", instrument="Cam1", detector=2, collections=["imported_g"], timespan=timespan

631 )

632 )

633 self.assertEqual(

634 non_calibration_ref,

635 registry.findDataset(

636 "noncalibration", instrument="Cam1", detector=2, collections=[run], timespan=timespan

637 ),

638 )

639 self.assertEqual(

640 non_calibration_ref,

641 registry.findDataset(

642 "noncalibration",

643 instrument="Cam1",

644 detector=2,

645 collections=["Cam1/calib", run],

646 timespan=timespan,

647 ),

648 )

649 # Add a dataset type whose dimension group involves an "implied"

650 # dimension. ("physical_filter" implies "band".)

651 registry.registerDatasetType(

652 DatasetType(

653 "dt_with_implied",

654 [

655 "instrument",

656 "physical_filter",

657 ],

658 "int",

659 universe=butler.dimensions,

660 )

661 )

662 data_id = {"instrument": "Cam1", "physical_filter": "Cam1-G"}

663 (implied_ref,) = registry.insertDatasets("dt_with_implied", dataIds=[data_id], run=run)

664 found_ref = registry.findDataset("dt_with_implied", data_id, collections=[run])

665 self.assertEqual(implied_ref, found_ref)

666 # The "full" data ID with implied values is looked up, even though we

667 # provided only the "required" values.

668 self.assertTrue(found_ref.dataId.hasFull())

669 # The search ignores excess data ID values beyond the 'required' set.

670 # This is not the correct band value for this physical_filter, but

671 # the mismatch is ignored.

672 self.assertEqual(

673 implied_ref,

674 registry.findDataset(

675 "dt_with_implied",

676 {"instrument": "Cam1", "physical_filter": "Cam1-G", "band": "r"},

677 collections=[run],

678 ),

679 )

680 # Correct band value, wrong physical_filter.

681 self.assertIsNone(

682 registry.findDataset(

683 "dt_with_implied",

684 {"instrument": "Cam1", "physical_filter": "Cam1-R1", "band": "g"},

685 collections=[run],

686 ),

687 )

688

689 def testRemoveDatasetTypeSuccess(self):

690 """Test that SqlRegistry.removeDatasetType works when there are no

691 datasets of that type present.

692 """

693 butler = self.make_butler()

694 registry = butler.registry

695 self.load_data(butler, "base.yaml")

696 registry.removeDatasetType("flat")

697 with self.assertRaises(MissingDatasetTypeError):

698 registry.getDatasetType("flat")

699

700 def testRemoveDatasetTypeFailure(self):

701 """Test that SqlRegistry.removeDatasetType raises when there are

702 datasets of that type present or if the dataset type is for a

703 component.

704 """

705 butler = self.make_butler()

706 registry = butler.registry

707 self.load_data(butler, "base.yaml", "datasets.yaml")

708 with self.assertRaises(OrphanedRecordError):

709 registry.removeDatasetType("flat")

710 with self.assertRaises(DatasetTypeError):

711 registry.removeDatasetType(DatasetType.nameWithComponent("flat", "image"))

712

713 def testImportDatasetsUUID(self):

714 """Test for `SqlRegistry._importDatasets` with UUID dataset ID."""

715 if isinstance(self.datasetsManager, str):

716 if not self.datasetsManager.endswith(".ByDimensionsDatasetRecordStorageManagerUUID"):

717 self.skipTest(f"Unexpected dataset manager {self.datasetsManager}")

718 elif isinstance(self.datasetsManager, dict) and not self.datasetsManager["cls"].endswith(

719 ".ByDimensionsDatasetRecordStorageManagerUUID"

720 ):

721 self.skipTest(f"Unexpected dataset manager {self.datasetsManager['cls']}")

722

723 butler = self.make_butler()

724 registry = butler.registry

725 self.load_data(butler, "base.yaml")

726 for run in range(6):

727 registry.registerRun(f"run{run}")

728 datasetTypeBias = registry.getDatasetType("bias")

729 datasetTypeFlat = registry.getDatasetType("flat")

730 dataIdBias1 = {"instrument": "Cam1", "detector": 1}

731 dataIdBias2 = {"instrument": "Cam1", "detector": 2}

732 dataIdFlat1 = {"instrument": "Cam1", "detector": 1, "physical_filter": "Cam1-G", "band": "g"}

733

734 ref = DatasetRef(datasetTypeBias, dataIdBias1, run="run0")

735 (ref1,) = registry._importDatasets([ref])

736 # UUID is used without change

737 self.assertEqual(ref.id, ref1.id)

738

739 # All different failure modes

740 refs = (

741 # Importing same DatasetRef with different dataset ID is an error

742 DatasetRef(datasetTypeBias, dataIdBias1, run="run0"),

743 # Same DatasetId but different DataId

744 DatasetRef(datasetTypeBias, dataIdBias2, id=ref1.id, run="run0"),

745 DatasetRef(datasetTypeFlat, dataIdFlat1, id=ref1.id, run="run0"),

746 # Same DatasetRef and DatasetId but different run

747 DatasetRef(datasetTypeBias, dataIdBias1, id=ref1.id, run="run1"),

748 )

749 for ref in refs:

750 with self.assertRaises(ConflictingDefinitionError):

751 registry._importDatasets([ref])

752

753 # Test for non-unique IDs, they can be re-imported multiple times.

754 for run, idGenMode in ((2, DatasetIdGenEnum.DATAID_TYPE), (4, DatasetIdGenEnum.DATAID_TYPE_RUN)):

755 with self.subTest(idGenMode=idGenMode):

756 # Make dataset ref with reproducible dataset ID.

757 ref = DatasetRef(datasetTypeBias, dataIdBias1, run=f"run{run}", id_generation_mode=idGenMode)

758 (ref1,) = registry._importDatasets([ref])

759 self.assertIsInstance(ref1.id, uuid.UUID)

760 self.assertEqual(ref1.id.version, 5)

761 self.assertEqual(ref1.id, ref.id)

762

763 # Importing it again is OK

764 (ref2,) = registry._importDatasets([ref1])

765 self.assertEqual(ref2.id, ref1.id)

766

767 # Cannot import to different run with the same ID

768 ref = DatasetRef(datasetTypeBias, dataIdBias1, id=ref1.id, run=f"run{run + 1}")

769 with self.assertRaises(ConflictingDefinitionError):

770 registry._importDatasets([ref])

771

772 ref = DatasetRef(

773 datasetTypeBias, dataIdBias1, run=f"run{run + 1}", id_generation_mode=idGenMode

774 )

775 if idGenMode is DatasetIdGenEnum.DATAID_TYPE:

776 # Cannot import same DATAID_TYPE ref into a new run

777 with self.assertRaises(ConflictingDefinitionError):

778 (ref2,) = registry._importDatasets([ref])

779 else:

780 # DATAID_TYPE_RUN ref can be imported into a new run

781 (ref2,) = registry._importDatasets([ref])

782

783 def testComponentLookups(self):

784 """Test searching for component datasets via their parents.

785

786 Components can no longer be found by registry. This test checks

787 that this now fails.

788 """

789 butler = self.make_butler()

790 registry = butler.registry

791 self.load_data(butler, "base.yaml", "datasets.yaml")

792 # Test getting the child dataset type (which does still exist in the

793 # Registry), and check for consistency with

794 # DatasetRef.makeComponentRef.

795 collection = "imported_g"

796 parentType = registry.getDatasetType("bias")

797 childType = registry.getDatasetType("bias.wcs")

798 parentRefResolved = registry.findDataset(

799 parentType, collections=collection, instrument="Cam1", detector=1

800 )

801 self.assertIsInstance(parentRefResolved, DatasetRef)

802 self.assertEqual(childType, parentRefResolved.makeComponentRef("wcs").datasetType)

803 # Search for a single dataset with findDataset.

804 with self.assertRaises(DatasetTypeError):

805 registry.findDataset("bias.wcs", collections=collection, dataId=parentRefResolved.dataId)

806

807 def testCollections(self):

808 """Tests for registry methods that manage collections."""

809 butler = self.make_butler()

810 registry = butler.registry

811 other_registry = butler.clone().registry

812 self.load_data(butler, "base.yaml", "datasets.yaml")

813 run1 = "imported_g"

814 run2 = "imported_r"

815 # Test setting a collection docstring after it has been created.

816 registry.setCollectionDocumentation(run1, "doc for run1")

817 self.assertEqual(registry.getCollectionDocumentation(run1), "doc for run1")

818 registry.setCollectionDocumentation(run1, None)

819 self.assertIsNone(registry.getCollectionDocumentation(run1))

820 datasetType = "bias"

821 # Find some datasets via their run's collection.

822 dataId1 = {"instrument": "Cam1", "detector": 1}

823 ref1 = registry.findDataset(datasetType, dataId1, collections=run1)

824 self.assertIsNotNone(ref1)

825 dataId2 = {"instrument": "Cam1", "detector": 2}

826 ref2 = registry.findDataset(datasetType, dataId2, collections=run1)

827 self.assertIsNotNone(ref2)

828 # Associate those into a new collection, then look for them there.

829 tag1 = "tag1"

830 registry.registerCollection(tag1, type=CollectionType.TAGGED, doc="doc for tag1")

831 # Check that we can query for old and new collections by type.

832 self.assertEqual(set(registry.queryCollections(collectionTypes=CollectionType.RUN)), {run1, run2})

833 self.assertEqual(

834 set(registry.queryCollections(collectionTypes={CollectionType.TAGGED, CollectionType.RUN})),

835 {tag1, run1, run2},

836 )

837 self.assertEqual(registry.getCollectionDocumentation(tag1), "doc for tag1")

838 registry.associate(tag1, [ref1, ref2])

839 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)

840 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)

841 # Disassociate one and verify that we can't it there anymore...

842 registry.disassociate(tag1, [ref1])

843 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=tag1))

844 # ...but we can still find ref2 in tag1, and ref1 in the run.

845 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run1), ref1)

846 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)

847 collections = set(registry.queryCollections())

848 self.assertEqual(collections, {run1, run2, tag1})

849 # Associate both refs into tag1 again; ref2 is already there, but that

850 # should be a harmless no-op.

851 registry.associate(tag1, [ref1, ref2])

852 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)

853 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)

854 # Get a different dataset (from a different run) that has the same

855 # dataset type and data ID as ref2.

856 ref2b = registry.findDataset(datasetType, dataId2, collections=run2)

857 self.assertNotEqual(ref2, ref2b)

858 # Attempting to associate that into tag1 should be an error.

859 with self.assertRaises(ConflictingDefinitionError):

860 registry.associate(tag1, [ref2b])

861 # That error shouldn't have messed up what we had before.

862 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)

863 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)

864 # Attempt to associate the conflicting dataset again, this time with

865 # a dataset that isn't in the collection and won't cause a conflict.

866 # Should also fail without modifying anything.

867 dataId3 = {"instrument": "Cam1", "detector": 3}

868 ref3 = registry.findDataset(datasetType, dataId3, collections=run1)

869 with self.assertRaises(ConflictingDefinitionError):

870 registry.associate(tag1, [ref3, ref2b])

871 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)

872 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)

873 self.assertIsNone(registry.findDataset(datasetType, dataId3, collections=tag1))

874 # Register a chained collection that searches [tag1, run2]

875 chain1 = "chain1"

876 registry.registerCollection(chain1, type=CollectionType.CHAINED)

877 self.assertIs(registry.getCollectionType(chain1), CollectionType.CHAINED)

878 # Chained collection exists, but has no collections in it.

879 self.assertFalse(registry.getCollectionChain(chain1))

880 # If we query for all collections, we should get the chained collection

881 # if we don't ask to flatten it (i.e. yield only its children) or if we

882 # explicitly ask to include it too.

883 self.assertEqual(set(registry.queryCollections(flattenChains=False)), {tag1, run1, run2, chain1})

884 self.assertEqual(set(registry.queryCollections(flattenChains=True)), {tag1, run1, run2})

885 self.assertEqual(

886 set(registry.queryCollections(flattenChains=True, includeChains=True)), {tag1, run1, run2, chain1}

887 )

888 # Attempt to set its child collections to something circular; that

889 # should fail.

890 with self.assertRaises(ValueError):

891 registry.setCollectionChain(chain1, [tag1, chain1])

892 # Add the child collections.

893 registry.setCollectionChain(chain1, [tag1, run2])

894 self.assertEqual(list(registry.getCollectionChain(chain1)), [tag1, run2])

895 self.assertEqual(registry.getCollectionParentChains(tag1), {chain1})

896 self.assertEqual(registry.getCollectionParentChains(run2), {chain1})

897 # Refresh the other registry that points to the same repo, and make

898 # sure it can see the things we've done (note that this does require

899 # an explicit refresh(); that's the documented behavior, because

900 # caching is ~impossible otherwise).

901 if other_registry is not None:

902 other_registry.refresh()

903 self.assertEqual(list(other_registry.getCollectionChain(chain1)), [tag1, run2])

904 self.assertEqual(other_registry.getCollectionParentChains(tag1), {chain1})

905 self.assertEqual(other_registry.getCollectionParentChains(run2), {chain1})

906 # Searching for dataId1 or dataId2 in the chain should return ref1 and

907 # ref2, because both are in tag1.

908 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain1), ref1)

909 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain1), ref2)

910 # Now disassociate ref2 from tag1. The search (for bias) with

911 # dataId2 in chain1 should then:

912 # 1. not find it in tag1

913 # 2. find a different dataset in run2

914 registry.disassociate(tag1, [ref2])

915 ref2b = registry.findDataset(datasetType, dataId2, collections=chain1)

916 self.assertNotEqual(ref2b, ref2)

917 self.assertEqual(ref2b, registry.findDataset(datasetType, dataId2, collections=run2))

918 # Define a new chain so we can test recursive chains.

919 chain2 = "chain2"

920 registry.registerCollection(chain2, type=CollectionType.CHAINED)

921 registry.setCollectionChain(chain2, [run2, chain1])

922 self.assertEqual(registry.getCollectionParentChains(chain1), {chain2})

923 self.assertEqual(registry.getCollectionParentChains(run2), {chain1, chain2})

924

925 if self.supportsCollectionRegex:

926 # Query for collections matching a regex.

927 with self.assertWarns(FutureWarning):

928 self.assertCountEqual(

929 list(registry.queryCollections(re.compile("imported_."), flattenChains=False)),

930 ["imported_r", "imported_g"],

931 )

932 # Query for collections matching a regex or an explicit str.

933 with self.assertWarns(FutureWarning):

934 self.assertCountEqual(

935 list(

936 registry.queryCollections([re.compile("imported_."), "chain1"], flattenChains=False)

937 ),

938 ["imported_r", "imported_g", "chain1"],

939 )

940 # Same queries as the regex ones above, but using globs instead of

941 # regex.

942 self.assertCountEqual(

943 list(registry.queryCollections("imported_*", flattenChains=False)),

944 ["imported_r", "imported_g"],

945 )

946 # Query for collections matching a regex or an explicit str.

947 self.assertCountEqual(

948 list(registry.queryCollections(["imported_*", "chain1"], flattenChains=False)),

949 ["imported_r", "imported_g", "chain1"],

950 )

951 # Query for collection matching chain names, by flattening it should

952 # only return non-chain names.

953 self.assertCountEqual(list(registry.queryCollections("chain?", flattenChains=True)), [tag1, run2])

954 # Query for collection matching chain name, by flattening and

955 # asking to include chains it should return everything.

956 self.assertCountEqual(

957 list(registry.queryCollections("chain*", flattenChains=True, includeChains=True)),

958 [tag1, run2, chain1, chain2],

959 )

960 # Order of children in chained collections is preserved.

961 self.assertEqual(list(registry.queryCollections("chain1", flattenChains=True)), [tag1, run2])

962 self.assertEqual(list(registry.queryCollections("cha*2", flattenChains=True)), [run2, tag1])

963 self.assertEqual(

964 list(registry.queryCollections("chain1", flattenChains=True, includeChains=True)),

965 [chain1, tag1, run2],

966 )

967 self.assertEqual(

968 list(registry.queryCollections("chain2", flattenChains=True, includeChains=True)),

969 [chain2, run2, chain1, tag1],

970 )

971

972 # Search for bias with dataId1 should find it via tag1 in chain2,

973 # recursing, because is not in run1.

974 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=run2))

975 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain2), ref1)

976 # Search for bias with dataId2 should find it in run2 (ref2b).

977 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain2), ref2b)

978 # Search for a flat that is in run2. That should not be found

979 # at the front of chain2, because of the restriction to bias

980 # on run2 there, but it should be found in at the end of chain1.

981 dataId4 = {"instrument": "Cam1", "detector": 3, "physical_filter": "Cam1-R2"}

982 ref4 = registry.findDataset("flat", dataId4, collections=run2)

983 self.assertIsNotNone(ref4)

984 self.assertEqual(ref4, registry.findDataset("flat", dataId4, collections=chain2))

985 # Deleting a collection that's part of a CHAINED collection is not

986 # allowed, and is exception-safe.

987 with self.assertRaises(sqlalchemy.exc.IntegrityError):

988 registry.removeCollection(run2)

989 self.assertEqual(registry.getCollectionType(run2), CollectionType.RUN)

990 with self.assertRaises(sqlalchemy.exc.IntegrityError):

991 registry.removeCollection(chain1)

992 self.assertEqual(registry.getCollectionType(chain1), CollectionType.CHAINED)

993 # Actually remove chain2, test that it's gone by asking for its type.

994 registry.removeCollection(chain2)

995 with self.assertRaises(MissingCollectionError):

996 registry.getCollectionType(chain2)

997 # Actually remove run2 and chain1, which should work now.

998 registry.removeCollection(chain1)

999 registry.removeCollection(run2)

1000 with self.assertRaises(MissingCollectionError):

1001 registry.getCollectionType(run2)

1002 with self.assertRaises(MissingCollectionError):

1003 registry.getCollectionType(chain1)

1004 # Remove tag1 as well, just to test that we can remove TAGGED

1005 # collections.

1006 registry.removeCollection(tag1)

1007 with self.assertRaises(MissingCollectionError):

1008 registry.getCollectionType(tag1)

1009

1010 def test_collection_clearing(self) -> None:

1011 """Test that we can delete TAGGED and CALIBRATION collections without

1012 manually removing all associated datasets first.

1013 """

1014 butler = self.make_butler()

1015 self.load_data(butler, "base.yaml", "datasets.yaml")

1016

1017 # This brings in datasets of two different types, with the same

1018 # dimension group.

1019 original_datasets = tuple(butler.query_all_datasets("imported_r", instrument="Cam1", detector=2))

1020 self.assertEqual(len(original_datasets), 2)

1021

1022 # Test tagged collections.

1023 butler.collections.register("tag1", CollectionType.TAGGED)

1024 butler.collections.register("tag2", CollectionType.TAGGED)

1025 butler.registry.associate("tag1", original_datasets)

1026 butler.registry.associate("tag2", original_datasets)

1027 butler.collections.x_remove("tag1")

1028 with self.assertRaises(MissingCollectionError):

1029 butler.collections.get_info("tag1")

1030 # Make sure there was no collateral damage -- tag2 should still be

1031 # intact.

1032 self.assertEqual(set(butler.query_all_datasets("tag2")), set(original_datasets))

1033

1034 # Test calibration collections.

1035 butler.collections.register("calib1", CollectionType.CALIBRATION)

1036 butler.collections.register("calib2", CollectionType.CALIBRATION)

1037 butler.registry.certify("calib1", original_datasets, Timespan(None, None))

1038 butler.registry.certify("calib2", original_datasets, Timespan(None, None))

1039 butler.collections.x_remove("calib1")

1040 with self.assertRaises(MissingCollectionError):

1041 butler.collections.get_info("calib1")

1042 # Make sure there was no collateral damage -- calib2 should still be

1043 # intact.

1044 self.assertEqual(set(butler.query_all_datasets("calib2")), set(original_datasets))

1045

1046 def testCollectionChainCaching(self):

1047 butler = self.make_butler()

1048 registry = butler.registry

1049 with registry.caching_context():

1050 registry.registerCollection("a")

1051 registry.registerCollection("chain", CollectionType.CHAINED)

1052 # There used to be a caching bug (DM-43750) that would throw an

1053 # exception if you modified a collection chain for a collection

1054 # that was already in the cache.

1055 registry.setCollectionChain("chain", ["a"])

1056 self.assertEqual(list(registry.getCollectionChain("chain")), ["a"])

1057

1058 def testCollectionChainFlatten(self):

1059 """Test that `SqlRegistry.setCollectionChain` obeys its 'flatten'

1060 option.

1061 """

1062 butler = self.make_butler()

1063 registry = butler.registry

1064 registry.registerCollection("inner", CollectionType.CHAINED)

1065 registry.registerCollection("innermost", CollectionType.RUN)

1066 registry.setCollectionChain("inner", ["innermost"])

1067 registry.registerCollection("outer", CollectionType.CHAINED)

1068 registry.setCollectionChain("outer", ["inner"], flatten=False)

1069 self.assertEqual(list(registry.getCollectionChain("outer")), ["inner"])

1070 registry.setCollectionChain("outer", ["inner"], flatten=True)

1071 self.assertEqual(list(registry.getCollectionChain("outer")), ["innermost"])

1072

1073 def testCollectionChainPrependConcurrency(self):

1074 """Verify that locking via database row locks is working as

1075 expected.

1076 """

1077

1078 def blocked_thread_func(butler: Butler):

1079 # This call will become blocked after it has decided on positions

1080 # for the new children in the collection chain, but before

1081 # inserting them.

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

1083

1084 def unblocked_thread_func(butler: Butler):

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

1086

1087 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)

1088

1089 # blocked_thread_func should have finished first, inserting "a".

1090 # unblocked_thread_func should have finished second, prepending "b".

1091 self.assertEqual(("b", "a"), registry.getCollectionChain("chain"))

1092

1093 def testCollectionChainReplaceConcurrency(self):

1094 """Verify that locking via database row locks is working as

1095 expected.

1096 """

1097

1098 def blocked_thread_func(butler: Butler):

1099 # This call will become blocked after deleting children, but before

1100 # inserting new ones.

1101 butler.collections.redefine_chain("chain", ["a"])

1102

1103 def unblocked_thread_func(butler: Butler):

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

1105

1106 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)

1107

1108 # blocked_thread_func should have finished first.

1109 # unblocked_thread_func should have finished second, overwriting the

1110 # chain with "b".

1111 self.assertEqual(("b",), registry.getCollectionChain("chain"))

1112

1113 def testCollectionChainRemoveConcurrency(self):

1114 def blocked_thread_func(butler: Butler):

1115 # This call will become blocked after taking the lock, but before

1116 # deleting the children.

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

1118

1119 def unblocked_thread_func(butler: Butler):

1120 butler.collections.redefine_chain("chain", ["b", "a"])

1121

1122 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)

1123

1124 # blocked_thread_func should have finished first, removing "b".

1125 # unblocked_thread_func should have finished second, putting "b" back.

1126 self.assertEqual(("b", "a"), registry.getCollectionChain("chain"))

1127

1128 def _do_collection_concurrency_test(

1129 self, blocked_thread_func: Callable[[Butler]], unblocked_thread_func: Callable[[Butler]]

1130 ) -> SqlRegistry:

1131 # This function:

1132 # 1. Sets up two registries pointing at the same database.

1133 # 2. Start running 'blocked_thread_func' in a background thread,

1134 # arranging for it to become blocked during a critical section in

1135 # the collections manager.

1136 # 3. Wait for 'blocked_thread_func' to reach the critical section

1137 # 4. Start running 'unblocked_thread_func'.

1138 # 5. Allow both functions to run to completion.

1139

1140 # Set up two registries pointing to the same DB

1141 butler1 = self.make_butler()

1142 butler2 = butler1.clone()

1143 registry1 = butler1._registry

1144 assert isinstance(registry1, SqlRegistry)

1145 registry2 = butler2._registry

1146

1147 with contextlib.suppress(AttributeError):

1148 if ":memory:" in str(registry2._db):

1149 raise unittest.SkipTest("Testing concurrency requires two connections to the same DB.")

1150

1151 registry1.registerCollection("chain", CollectionType.CHAINED)

1152 for collection in ["a", "b"]:

1153 registry1.registerCollection(collection)

1154

1155 # Arrange for registry1 to block during its critical section, allowing

1156 # us to detect this and control when it becomes unblocked.

1157 enter_barrier = Barrier(2, timeout=60)

1158 exit_barrier = Barrier(2, timeout=60)

1159

1160 def wait_for_barrier():

1161 enter_barrier.wait()

1162 exit_barrier.wait()

1163

1164 registry1._managers.collections._block_for_concurrency_test = wait_for_barrier

1165

1166 with ThreadPoolExecutor(max_workers=1) as exec1:

1167 with ThreadPoolExecutor(max_workers=1) as exec2:

1168 future1 = exec1.submit(blocked_thread_func, butler1)

1169 enter_barrier.wait()

1170

1171 # At this point registry 1 has entered the critical section and

1172 # is waiting for us to release it. Start the other thread.

1173 future2 = exec2.submit(unblocked_thread_func, butler2)

1174 # thread2 should block inside a database call, but we have no

1175 # way to detect when it is in this state.

1176 time.sleep(0.200)

1177

1178 # Let the threads run to completion.

1179 exit_barrier.wait()

1180 future1.result()

1181 future2.result()

1182

1183 return registry1

1184

1185 def testBasicTransaction(self):

1186 """Test that all operations within a single transaction block are

1187 rolled back if an exception propagates out of the block.

1188 """

1189 butler = self.make_butler()

1190 registry = butler.registry

1191 storageClass = StorageClass("testDatasetType")

1192 registry.storageClasses.registerStorageClass(storageClass)

1193 with registry.transaction():

1194 registry.insertDimensionData("instrument", {"name": "Cam1", "class_name": "A"})

1195 with self.assertRaises(ValueError):

1196 with registry.transaction():

1197 registry.insertDimensionData("instrument", {"name": "Cam2"})

1198 raise ValueError("Oops, something went wrong")

1199 # Cam1 should exist

1200 self.assertEqual(registry.expandDataId(instrument="Cam1").records["instrument"].class_name, "A")

1201 # But Cam2 and Cam3 should both not exist

1202 with self.assertRaises(DataIdValueError):

1203 registry.expandDataId(instrument="Cam2")

1204 with self.assertRaises(DataIdValueError):

1205 registry.expandDataId(instrument="Cam3")

1206

1207 def testNestedTransaction(self):

1208 """Test that operations within a transaction block are not rolled back

1209 if an exception propagates out of an inner transaction block and is

1210 then caught.

1211 """

1212 butler = self.make_butler()

1213 registry = butler.registry

1214 dimension = registry.dimensions["instrument"]

1215 dataId1 = {"instrument": "DummyCam"}

1216 dataId2 = {"instrument": "DummyCam2"}

1217 checkpointReached = False

1218 with registry.transaction():

1219 # This should be added and (ultimately) committed.

1220 registry.insertDimensionData(dimension, dataId1)

1221 with self.assertRaises(sqlalchemy.exc.IntegrityError):

1222 with registry.transaction(savepoint=True):

1223 # This does not conflict, and should succeed (but not

1224 # be committed).

1225 registry.insertDimensionData(dimension, dataId2)

1226 checkpointReached = True

1227 # This should conflict and raise, triggering a rollback

1228 # of the previous insertion within the same transaction

1229 # context, but not the original insertion in the outer

1230 # block.

1231 registry.insertDimensionData(dimension, dataId1)

1232 self.assertTrue(checkpointReached)

1233 self.assertIsNotNone(registry.expandDataId(dataId1, dimensions=dimension.minimal_group))

1234 with self.assertRaises(DataIdValueError):

1235 registry.expandDataId(dataId2, dimensions=dimension.minimal_group)

1236

1237 def testInstrumentDimensions(self):

1238 """Test queries involving only instrument dimensions, with no joins to

1239 skymap.

1240 """

1241 butler = self.make_butler()

1242 registry = butler.registry

1243

1244 # need a bunch of dimensions and datasets for test

1245 registry.insertDimensionData(

1246 "instrument", dict(name="DummyCam", visit_max=25, exposure_max=300, detector_max=6)

1247 )

1248 registry.insertDimensionData("day_obs", dict(instrument="DummyCam", id=20250101))

1249 registry.insertDimensionData(

1250 "physical_filter",

1251 dict(instrument="DummyCam", name="dummy_r", band="r"),

1252 dict(instrument="DummyCam", name="dummy_i", band="i"),

1253 )

1254 registry.insertDimensionData(

1255 "detector", *[dict(instrument="DummyCam", id=i, full_name=str(i)) for i in range(1, 6)]

1256 )

1257 registry.insertDimensionData(

1258 "visit",

1259 dict(instrument="DummyCam", id=10, name="ten", physical_filter="dummy_i", day_obs=20250101),

1260 dict(instrument="DummyCam", id=11, name="eleven", physical_filter="dummy_r", day_obs=20250101),

1261 dict(instrument="DummyCam", id=20, name="twelve", physical_filter="dummy_r", day_obs=20250101),

1262 )

1263 registry.insertDimensionData(

1264 "group",

1265 dict(instrument="DummyCam", name="ten"),

1266 dict(instrument="DummyCam", name="eleven"),

1267 dict(instrument="DummyCam", name="twelve"),

1268 )

1269 for i in range(1, 6):

1270 registry.insertDimensionData(

1271 "visit_detector_region",

1272 dict(instrument="DummyCam", visit=10, detector=i),

1273 dict(instrument="DummyCam", visit=11, detector=i),

1274 dict(instrument="DummyCam", visit=20, detector=i),

1275 )

1276 registry.insertDimensionData(

1277 "exposure",

1278 dict(

1279 instrument="DummyCam",

1280 id=100,

1281 obs_id="100",

1282 physical_filter="dummy_i",

1283 group="ten",

1284 day_obs=20250101,

1285 ),

1286 dict(

1287 instrument="DummyCam",

1288 id=101,

1289 obs_id="101",

1290 physical_filter="dummy_i",

1291 group="ten",

1292 day_obs=20250101,

1293 ),

1294 dict(

1295 instrument="DummyCam",

1296 id=110,

1297 obs_id="110",

1298 physical_filter="dummy_r",

1299 group="eleven",

1300 day_obs=20250101,

1301 ),

1302 dict(

1303 instrument="DummyCam",

1304 id=111,

1305 obs_id="111",

1306 physical_filter="dummy_r",

1307 group="eleven",

1308 day_obs=20250101,

1309 ),

1310 dict(

1311 instrument="DummyCam",

1312 id=200,

1313 obs_id="200",

1314 physical_filter="dummy_r",

1315 group="twelve",

1316 day_obs=20250101,

1317 ),

1318 dict(

1319 instrument="DummyCam",

1320 id=201,

1321 obs_id="201",

1322 physical_filter="dummy_r",

1323 group="twelve",

1324 day_obs=20250101,

1325 ),

1326 )

1327 registry.insertDimensionData(

1328 "visit_definition",

1329 dict(instrument="DummyCam", exposure=100, visit=10),

1330 dict(instrument="DummyCam", exposure=101, visit=10),

1331 dict(instrument="DummyCam", exposure=110, visit=11),

1332 dict(instrument="DummyCam", exposure=111, visit=11),

1333 dict(instrument="DummyCam", exposure=200, visit=20),

1334 dict(instrument="DummyCam", exposure=201, visit=20),

1335 )

1336 # dataset types

1337 run1 = "test1_r"

1338 run2 = "test2_r"

1339 tagged2 = "test2_t"

1340 registry.registerRun(run1)

1341 registry.registerRun(run2)

1342 registry.registerCollection(tagged2)

1343 storageClass = StorageClass("testDataset")

1344 registry.storageClasses.registerStorageClass(storageClass)

1345 rawType = DatasetType(

1346 name="RAW",

1347 dimensions=registry.dimensions.conform(("instrument", "exposure", "detector")),

1348 storageClass=storageClass,

1349 )

1350 registry.registerDatasetType(rawType)

1351 calexpType = DatasetType(

1352 name="CALEXP",

1353 dimensions=registry.dimensions.conform(("instrument", "visit", "detector")),

1354 storageClass=storageClass,

1355 )

1356 registry.registerDatasetType(calexpType)

1357

1358 # add pre-existing datasets

1359 for exposure in (100, 101, 110, 111):

1360 for detector in (1, 2, 3):

1361 # note that only 3 of 5 detectors have datasets

1362 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector)

1363 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run1)

1364 # exposures 100 and 101 appear in both run1 and tagged2.

1365 # 100 has different datasets in the different collections

1366 # 101 has the same dataset in both collections.

1367 if exposure == 100:

1368 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run2)

1369 if exposure in (100, 101):

1370 registry.associate(tagged2, [ref])

1371 # Add pre-existing datasets to tagged2.

1372 for exposure in (200, 201):

1373 for detector in (3, 4, 5):

1374 # note that only 3 of 5 detectors have datasets

1375 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector)

1376 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run2)

1377 registry.associate(tagged2, [ref])

1378

1379 dimensions = registry.dimensions.conform(rawType.dimensions.required | calexpType.dimensions.required)

1380 # Test that single dim string works as well as list of str

1381 rows = registry.queryDataIds("visit", datasets=rawType, collections=run1).expanded().toSet()

1382 rowsI = registry.queryDataIds(["visit"], datasets=rawType, collections=run1).expanded().toSet()

1383 self.assertEqual(rows, rowsI)

1384 # with empty expression

1385 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1).expanded().toSet()

1386 self.assertEqual(len(rows), 4 * 3) # 4 exposures times 3 detectors

1387 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 110, 111))

1388 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 11))

1389 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))

1390

1391 # second collection

1392 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=tagged2).toSet()

1393 self.assertEqual(len(rows), 4 * 3) # 4 exposures times 3 detectors

1394 for dataId in rows:

1395 self.assertCountEqual(dataId.dimensions.required, ("instrument", "detector", "exposure", "visit"))

1396 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 200, 201))

1397 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 20))

1398 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3, 4, 5))

1399

1400 # with two input datasets

1401 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=[run1, tagged2]).toSet()

1402 self.assertEqual(len(set(rows)), 6 * 3) # 6 exposures times 3 detectors; set needed to de-dupe

1403 for dataId in rows:

1404 self.assertCountEqual(dataId.dimensions.required, ("instrument", "detector", "exposure", "visit"))

1405 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 110, 111, 200, 201))

1406 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 11, 20))

1407 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3, 4, 5))

1408

1409 # limit to single visit

1410 rows = registry.queryDataIds(

1411 dimensions, datasets=rawType, collections=run1, where="visit = 10", instrument="DummyCam"

1412 ).toSet()

1413 self.assertEqual(len(rows), 2 * 3) # 2 exposures times 3 detectors

1414 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101))

1415 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10,))

1416 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))

1417

1418 # more limiting expression, using link names instead of Table.column

1419 rows = registry.queryDataIds(

1420 dimensions,

1421 datasets=rawType,

1422 collections=run1,

1423 where="visit = 10 and detector > 1 and 'DummyCam'=instrument",

1424 ).toSet()

1425 self.assertEqual(len(rows), 2 * 2) # 2 exposures times 2 detectors

1426 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101))

1427 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10,))

1428 self.assertCountEqual({dataId["detector"] for dataId in rows}, (2, 3))

1429

1430 # queryDataIds with only one of `datasets` and `collections` is an

1431 # error.

1432 with self.assertRaises(CollectionError):

1433 registry.queryDataIds(dimensions, datasets=rawType)

1434 with self.assertRaises(ArgumentError):

1435 registry.queryDataIds(dimensions, collections=run1)

1436

1437 # expression excludes everything

1438 rows = registry.queryDataIds(

1439 dimensions, datasets=rawType, collections=run1, where="visit > 1000", instrument="DummyCam"

1440 ).toSet()

1441 self.assertEqual(len(rows), 0)

1442

1443 # Selecting by physical_filter, this is not in the dimensions, but it

1444 # is a part of the full expression so it should work too.

1445 rows = registry.queryDataIds(

1446 dimensions,

1447 datasets=rawType,

1448 collections=run1,

1449 where="physical_filter = 'dummy_r'",

1450 instrument="DummyCam",

1451 ).toSet()

1452 self.assertEqual(len(rows), 2 * 3) # 2 exposures times 3 detectors

1453 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (110, 111))

1454 self.assertCountEqual({dataId["visit"] for dataId in rows}, (11,))

1455 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))

1456

1457 def testSkyMapDimensions(self):

1458 """Tests involving only skymap dimensions, no joins to instrument."""

1459 butler = self.make_butler()

1460 registry = butler.registry

1461

1462 # need a bunch of dimensions and datasets for test, we want

1463 # "band" in the test so also have to add physical_filter

1464 # dimensions

1465 registry.insertDimensionData("instrument", dict(instrument="DummyCam"))

1466 registry.insertDimensionData(

1467 "physical_filter",

1468 dict(instrument="DummyCam", name="dummy_r", band="r"),

1469 dict(instrument="DummyCam", name="dummy_i", band="i"),

1470 )

1471 registry.insertDimensionData("skymap", dict(name="DummyMap", hash=b"sha!"))

1472 for tract in range(10):

1473 registry.insertDimensionData("tract", dict(skymap="DummyMap", id=tract))

1474 registry.insertDimensionData(

1475 "patch",

1476 *[dict(skymap="DummyMap", tract=tract, id=patch, cell_x=0, cell_y=0) for patch in range(10)],

1477 )

1478

1479 # dataset types

1480 run = "tésτ"

1481 registry.registerRun(run)

1482 storageClass = StorageClass("testDataset")

1483 registry.storageClasses.registerStorageClass(storageClass)

1484 calexpType = DatasetType(

1485 name="deepCoadd_calexp",

1486 dimensions=registry.dimensions.conform(("skymap", "tract", "patch", "band")),

1487 storageClass=storageClass,

1488 )

1489 registry.registerDatasetType(calexpType)

1490 mergeType = DatasetType(

1491 name="deepCoadd_mergeDet",

1492 dimensions=registry.dimensions.conform(("skymap", "tract", "patch")),

1493 storageClass=storageClass,

1494 )

1495 registry.registerDatasetType(mergeType)

1496 measType = DatasetType(

1497 name="deepCoadd_meas",

1498 dimensions=registry.dimensions.conform(("skymap", "tract", "patch", "band")),

1499 storageClass=storageClass,

1500 )

1501 registry.registerDatasetType(measType)

1502

1503 dimensions = registry.dimensions.conform(

1504 calexpType.dimensions.required | mergeType.dimensions.required | measType.dimensions.required

1505 )

1506

1507 # add pre-existing datasets

1508 for tract in (1, 3, 5):

1509 for patch in (2, 4, 6, 7):

1510 dataId = dict(skymap="DummyMap", tract=tract, patch=patch)

1511 registry.insertDatasets(mergeType, dataIds=[dataId], run=run)

1512 for aFilter in ("i", "r"):

1513 dataId = dict(skymap="DummyMap", tract=tract, patch=patch, band=aFilter)

1514 registry.insertDatasets(calexpType, dataIds=[dataId], run=run)

1515

1516 # with empty expression

1517 rows = registry.queryDataIds(dimensions, datasets=[calexpType, mergeType], collections=run).toSet()

1518 self.assertEqual(len(rows), 3 * 4 * 2) # 4 tracts x 4 patches x 2 filters

1519 for dataId in rows:

1520 self.assertCountEqual(dataId.dimensions.required, ("skymap", "tract", "patch", "band"))

1521 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 3, 5))

1522 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 4, 6, 7))

1523 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i", "r"))

1524

1525 # limit to 2 tracts and 2 patches

1526 rows = registry.queryDataIds(

1527 dimensions,

1528 datasets=[calexpType, mergeType],

1529 collections=run,

1530 where="tract IN (1, 5) AND patch IN (2, 7)",

1531 skymap="DummyMap",

1532 ).toSet()

1533 self.assertEqual(len(rows), 2 * 2 * 2) # 2 tracts x 2 patches x 2 filters

1534 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 5))

1535 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 7))

1536 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i", "r"))

1537

1538 # limit to single filter

1539 rows = registry.queryDataIds(

1540 dimensions, datasets=[calexpType, mergeType], collections=run, where="band = 'i'"

1541 ).toSet()

1542 self.assertEqual(len(rows), 3 * 4 * 1) # 4 tracts x 4 patches x 2 filters

1543 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 3, 5))

1544 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 4, 6, 7))

1545 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i",))

1546

1547 def do_query():

1548 return registry.queryDataIds(

1549 dimensions, datasets=[calexpType, mergeType], collections=run, where="skymap = 'Mars'"

1550 ).toSet()

1551

1552 self.assertEqual(len(do_query()), 0)

1553

1554 def testSpatialJoin(self):

1555 """Test queries that involve spatial overlap joins."""

1556 butler = self.make_butler()

1557 registry = butler.registry

1558 self.load_data(butler, "base.yaml", "spatial.yaml")

1559

1560 # Dictionary of spatial DatabaseDimensionElements, keyed by the name of

1561 # the TopologicalFamily they belong to. We'll relate all elements in

1562 # each family to all of the elements in each other family.

1563 families = defaultdict(set)

1564 # Dictionary of {element.name: {dataId: region}}.

1565 regions = {}

1566 for element in registry.dimensions.database_elements:

1567 if element.spatial is not None:

1568 families[element.spatial.name].add(element)

1569 regions[element.name] = {

1570 record.dataId: record.region for record in registry.queryDimensionRecords(element)

1571 }

1572

1573 # If this check fails, it's not necessarily a problem - it may just be

1574 # a reasonable change to the default dimension definitions - but the

1575 # test below depends on there being more than one family to do anything

1576 # useful.

1577 self.assertEqual(len(families), 2)

1578

1579 # Overlap DatabaseDimensionElements with each other.

1580 for family1, family2 in itertools.combinations(families, 2):

1581 for element1, element2 in itertools.product(families[family1], families[family2]):

1582 dimensions = element1.minimal_group | element2.minimal_group

1583 # Construct expected set of overlapping data IDs via a

1584 # brute-force comparison of the regions we've already fetched.

1585 expected = {

1586 DataCoordinate.standardize(

1587 {**dataId1.required, **dataId2.required}, dimensions=dimensions

1588 )

1589 for (dataId1, region1), (dataId2, region2) in itertools.product(

1590 regions[element1.name].items(), regions[element2.name].items()

1591 )

1592 if not region1.isDisjointFrom(region2)

1593 }

1594 self.assertGreater(len(expected), 2, msg="Test that we aren't just comparing empty sets.")

1595 queried = set(registry.queryDataIds(dimensions))

1596 self.assertEqual(expected, queried)

1597

1598 # Overlap each DatabaseDimensionElement with the commonSkyPix system.

1599 commonSkyPix = registry.dimensions.commonSkyPix

1600 for elementName, these_regions in regions.items():

1601 dimensions = registry.dimensions[elementName].minimal_group | commonSkyPix.minimal_group

1602 expected = set()

1603 for dataId, region in these_regions.items():

1604 for begin, end in commonSkyPix.pixelization.envelope(region):

1605 expected.update(

1606 DataCoordinate.standardize(

1607 {commonSkyPix.name: index, **dataId.required}, dimensions=dimensions

1608 )

1609 for index in range(begin, end)

1610 )

1611 self.assertGreater(len(expected), 2, msg="Test that we aren't just comparing empty sets.")

1612 queried = set(registry.queryDataIds(dimensions))

1613 self.assertEqual(expected, queried)

1614

1615 def testAbstractQuery(self):

1616 """Test that we can run a query that just lists the known

1617 bands. This is tricky because band is

1618 backed by a query against physical_filter.

1619 """

1620 butler = self.make_butler()

1621 registry = butler.registry

1622 registry.insertDimensionData("instrument", dict(name="DummyCam"))

1623 registry.insertDimensionData(

1624 "physical_filter",

1625 dict(instrument="DummyCam", name="dummy_i", band="i"),

1626 dict(instrument="DummyCam", name="dummy_i2", band="i"),

1627 dict(instrument="DummyCam", name="dummy_r", band="r"),

1628 )

1629 rows = registry.queryDataIds(["band"]).toSet()

1630 self.assertCountEqual(

1631 rows,

1632 [

1633 DataCoordinate.standardize(band="i", universe=registry.dimensions),

1634 DataCoordinate.standardize(band="r", universe=registry.dimensions),

1635 ],

1636 )

1637

1638 def testAttributeManager(self):

1639 """Test basic functionality of attribute manager."""

1640 # number of attributes with schema versions in a fresh database,

1641 # 6 managers with 2 records per manager, plus config for dimensions

1642 VERSION_COUNT = 6 * 2 + 1

1643

1644 butler = self.make_butler()

1645 registry = butler._registry

1646 attributes = registry._managers.attributes

1647

1648 # check what get() returns for non-existing key

1649 self.assertIsNone(attributes.get("attr"))

1650 self.assertEqual(attributes.get("attr", ""), "")

1651 self.assertEqual(attributes.get("attr", "Value"), "Value")

1652 self.assertEqual(len(list(attributes.items())), VERSION_COUNT)

1653

1654 # cannot store empty key or value

1655 with self.assertRaises(ValueError):

1656 attributes.set("", "value")

1657 with self.assertRaises(ValueError):

1658 attributes.set("attr", "")

1659

1660 # set value of non-existing key

1661 attributes.set("attr", "value")

1662 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)

1663 self.assertEqual(attributes.get("attr"), "value")

1664

1665 # update value of existing key

1666 with self.assertRaises(ButlerAttributeExistsError):

1667 attributes.set("attr", "value2")

1668

1669 attributes.set("attr", "value2", force=True)

1670 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)

1671 self.assertEqual(attributes.get("attr"), "value2")

1672

1673 # delete existing key

1674 self.assertTrue(attributes.delete("attr"))

1675 self.assertEqual(len(list(attributes.items())), VERSION_COUNT)

1676

1677 # delete non-existing key

1678 self.assertFalse(attributes.delete("non-attr"))

1679

1680 # store bunch of keys and get the list back

1681 data = [

1682 ("version.core", "1.2.3"),

1683 ("version.dimensions", "3.2.1"),

1684 ("config.managers.opaque", "ByNameOpaqueTableStorageManager"),

1685 ]

1686 for key, value in data:

1687 attributes.set(key, value)

1688 items = dict(attributes.items())

1689 for key, value in data:

1690 self.assertEqual(items[key], value)

1691

1692 def testQueryDatasetsDeduplication(self):

1693 """Test that the findFirst option to queryDatasets selects datasets

1694 from collections in the order given".

1695 """

1696 butler = self.make_butler()

1697 registry = butler.registry

1698 self.load_data(butler, "base.yaml", "datasets.yaml")

1699 self.assertCountEqual(

1700 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"])),

1701 [

1702 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),

1703 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g"),

1704 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g"),

1705 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r"),

1706 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r"),

1707 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),

1708 ],

1709 )

1710 self.assertCountEqual(

1711 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"], findFirst=True)),

1712 [

1713 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),

1714 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g"),

1715 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g"),

1716 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),

1717 ],

1718 )

1719 self.assertCountEqual(

1720 list(registry.queryDatasets("bias", collections=["imported_r", "imported_g"], findFirst=True)),

1721 [

1722 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),

1723 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r"),

1724 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r"),

1725 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),

1726 ],

1727 )

1728

1729 with self.assertRaises(TypeError):

1730 # Collection wildcards not allowed in find-first searches because

1731 # they do not guarantee the ordering of collections.

1732 registry.queryDatasets("bias", collections="imported_*", findFirst=True)

1733

1734 def testQueryDatasetsExtraDimensions(self):

1735 butler = self.make_butler()

1736 registry = butler.registry

1737 self.load_data(butler, "base.yaml", "datasets.yaml")

1738 # Bias dataset type does not include physical filter. By adding

1739 # "physical_filter" to dimensions, we are effectively searching here

1740 # for bias datasets with an instrument that has a specific filter

1741 # available, even though that filter has nothing to do with the bias

1742 # datasets we are finding.

1743 self.assertEqual(

1744 0,

1745 registry.queryDatasets(

1746 "bias",

1747 collections=...,

1748 dimensions=["physical_filter"],

1749 dataId={

1750 "instrument": "Cam1",

1751 "band": "not_a_real_band",

1752 },

1753 ).count(),

1754 )

1755 self.assertEqual(

1756 6,

1757 len(

1758 set(

1759 registry.queryDatasets(

1760 "bias",

1761 collections=...,

1762 dimensions=["physical_filter"],

1763 dataId={

1764 "instrument": "Cam1",

1765 "band": "r",

1766 },

1767 )

1768 )

1769 ),

1770 )

1771

1772 def testQueryResults(self):

1773 """Test querying for data IDs and then manipulating the QueryResults

1774 object returned to perform other queries.

1775 """

1776 butler = self.make_butler()

1777 registry = butler.registry

1778 self.load_data(butler, "base.yaml", "datasets.yaml")

1779 bias = registry.getDatasetType("bias")

1780 flat = registry.getDatasetType("flat")

1781 # Obtain expected results from methods other than those we're testing

1782 # here. That includes:

1783 # - the dimensions of the data IDs we want to query:

1784 expected_dimensions = registry.dimensions.conform(["detector", "physical_filter"])

1785 # - the dimensions of some other data IDs we'll extract from that:

1786 expected_subset_dimensions = registry.dimensions.conform(["detector"])

1787 # - the data IDs we expect to obtain from the first queries:

1788 expectedDataIds = DataCoordinateSet(

1789 {

1790 DataCoordinate.standardize(

1791 instrument="Cam1", detector=d, physical_filter=p, universe=registry.dimensions

1792 )

1793 for d, p in itertools.product({1, 2, 3}, {"Cam1-G", "Cam1-R1", "Cam1-R2"})

1794 },

1795 dimensions=expected_dimensions,

1796 hasFull=False,

1797 hasRecords=False,

1798 )

1799 # - the flat datasets we expect to find from those data IDs, in just

1800 # one collection (so deduplication is irrelevant):

1801 expectedFlats = [

1802 registry.findDataset(

1803 flat, instrument="Cam1", detector=1, physical_filter="Cam1-R1", collections="imported_r"

1804 ),

1805 registry.findDataset(

1806 flat, instrument="Cam1", detector=2, physical_filter="Cam1-R1", collections="imported_r"

1807 ),

1808 registry.findDataset(

1809 flat, instrument="Cam1", detector=3, physical_filter="Cam1-R2", collections="imported_r"

1810 ),

1811 ]

1812 # - the data IDs we expect to extract from that:

1813 expectedSubsetDataIds = expectedDataIds.subset(expected_subset_dimensions)

1814 # - the bias datasets we expect to find from those data IDs, after we

1815 # subset-out the physical_filter dimension, both with duplicates:

1816 expectedAllBiases = [

1817 registry.findDataset(bias, instrument="Cam1", detector=1, collections="imported_g"),

1818 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_g"),

1819 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_g"),

1820 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_r"),

1821 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_r"),

1822 ]

1823 # - ...and without duplicates:

1824 expectedDeduplicatedBiases = [

1825 registry.findDataset(bias, instrument="Cam1", detector=1, collections="imported_g"),

1826 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_r"),

1827 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_r"),

1828 ]

1829 # Test against those expected results, using a "lazy" query for the

1830 # data IDs (which re-executes that query each time we use it to do

1831 # something new).

1832 dataIds = registry.queryDataIds(

1833 ["detector", "physical_filter"],

1834 where="detector.purpose = 'SCIENCE'", # this rejects detector=4

1835 instrument="Cam1",

1836 )

1837 self.assertEqual(dataIds.dimensions, expected_dimensions)

1838 self.assertEqual(dataIds.toSet(), expectedDataIds)

1839 self.assertCountEqual(

1840 list(

1841 dataIds.findDatasets(

1842 flat,

1843 collections=["imported_r"],

1844 )

1845 ),

1846 expectedFlats,

1847 )

1848 subsetDataIds = dataIds.subset(expected_subset_dimensions, unique=True)

1849 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)

1850 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)

1851 self.assertCountEqual(

1852 list(subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=False)),

1853 expectedAllBiases,

1854 )

1855 self.assertCountEqual(

1856 list(subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)),

1857 expectedDeduplicatedBiases,

1858 )

1859

1860 # Searching for a dataset with dimensions we had projected away

1861 # restores those dimensions.

1862 self.assertCountEqual(

1863 list(subsetDataIds.findDatasets("flat", collections=["imported_r"], findFirst=True)),

1864 expectedFlats,

1865 )

1866

1867 # Use a named dataset type that does not exist and a dataset type

1868 # object that does not exist.

1869 unknown_type = DatasetType("not_known", dimensions=bias.dimensions, storageClass="Exposure")

1870

1871 # Test both string name and dataset type object.

1872 test_type: str | DatasetType

1873 for test_type, test_type_name in (

1874 (unknown_type, unknown_type.name),

1875 (unknown_type.name, unknown_type.name),

1876 ):

1877 with self.assertRaisesRegex(DatasetTypeError, expected_regex=test_type_name):

1878 list(

1879 subsetDataIds.findDatasets(

1880 test_type, collections=["imported_r", "imported_g"], findFirst=True

1881 )

1882 )

1883

1884 # Materialize the data ID subset query, but not the dataset queries.

1885 with subsetDataIds.materialize() as subsetDataIds:

1886 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)

1887 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)

1888 self.assertCountEqual(

1889 list(

1890 subsetDataIds.findDatasets(

1891 bias, collections=["imported_r", "imported_g"], findFirst=False

1892 )

1893 ),

1894 expectedAllBiases,

1895 )

1896 self.assertCountEqual(

1897 list(

1898 subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)

1899 ),

1900 expectedDeduplicatedBiases,

1901 )

1902 # Materialize the original query, but none of the follow-up queries.

1903 with dataIds.materialize() as dataIds:

1904 self.assertEqual(dataIds.dimensions, expected_dimensions)

1905 self.assertEqual(dataIds.toSet(), expectedDataIds)

1906 self.assertCountEqual(

1907 list(

1908 dataIds.findDatasets(

1909 flat,

1910 collections=["imported_r"],

1911 )

1912 ),

1913 expectedFlats,

1914 )

1915 subsetDataIds = dataIds.subset(expected_subset_dimensions, unique=True)

1916 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)

1917 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)

1918 self.assertCountEqual(

1919 list(

1920 subsetDataIds.findDatasets(

1921 bias, collections=["imported_r", "imported_g"], findFirst=False

1922 )

1923 ),

1924 expectedAllBiases,

1925 )

1926 self.assertCountEqual(

1927 list(

1928 subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)

1929 ),

1930 expectedDeduplicatedBiases,

1931 )

1932 # Materialize the subset data ID query, but not the dataset

1933 # queries.

1934 with subsetDataIds.materialize() as subsetDataIds:

1935 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)

1936 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)

1937 self.assertCountEqual(

1938 list(

1939 subsetDataIds.findDatasets(

1940 bias, collections=["imported_r", "imported_g"], findFirst=False

1941 )

1942 ),

1943 expectedAllBiases,

1944 )

1945 self.assertCountEqual(

1946 list(

1947 subsetDataIds.findDatasets(

1948 bias, collections=["imported_r", "imported_g"], findFirst=True

1949 )

1950 ),

1951 expectedDeduplicatedBiases,

1952 )

1953

1954 def testStorageClassPropagation(self):

1955 """Test that queries for datasets respect the storage class passed in

1956 as part of a full dataset type.

1957 """

1958 butler = self.make_butler()

1959 registry = butler.registry

1960 self.load_data(butler, "base.yaml")

1961 dataset_type_in_registry = DatasetType(

1962 "tbl", dimensions=["instrument"], storageClass="Packages", universe=registry.dimensions

1963 )

1964 registry.registerDatasetType(dataset_type_in_registry)

1965 run = "run1"

1966 registry.registerRun(run)

1967 (inserted_ref,) = registry.insertDatasets(

1968 dataset_type_in_registry, [registry.expandDataId(instrument="Cam1")], run=run

1969 )

1970 self.assertEqual(inserted_ref.datasetType, dataset_type_in_registry)

1971 query_dataset_type = DatasetType(

1972 "tbl", dimensions=["instrument"], storageClass="StructuredDataDict", universe=registry.dimensions

1973 )

1974 self.assertNotEqual(dataset_type_in_registry, query_dataset_type)

1975 query_datasets_result = registry.queryDatasets(query_dataset_type, collections=[run])

1976 self.assertEqual(query_datasets_result.parentDatasetType, query_dataset_type) # type: ignore

1977 (query_datasets_ref,) = query_datasets_result

1978 self.assertEqual(query_datasets_ref.datasetType, query_dataset_type)

1979 query_data_ids_find_datasets_result = registry.queryDataIds(["instrument"]).findDatasets(

1980 query_dataset_type, collections=[run]

1981 )

1982 self.assertEqual(query_data_ids_find_datasets_result.parentDatasetType, query_dataset_type)

1983 (query_data_ids_find_datasets_ref,) = query_data_ids_find_datasets_result

1984 self.assertEqual(query_data_ids_find_datasets_ref.datasetType, query_dataset_type)

1985 query_dataset_types_result = registry.queryDatasetTypes(query_dataset_type)

1986 self.assertEqual(list(query_dataset_types_result), [query_dataset_type])

1987 find_dataset_ref = registry.findDataset(query_dataset_type, instrument="Cam1", collections=[run])

1988 self.assertEqual(find_dataset_ref.datasetType, query_dataset_type)

1989

1990 def testEmptyDimensionsQueries(self):

1991 """Test Query and QueryResults objects in the case where there are no

1992 dimensions.

1993 """

1994 # Set up test data: one dataset type, two runs, one dataset in each.

1995 butler = self.make_butler()

1996 registry = butler.registry

1997 self.load_data(butler, "base.yaml")

1998 schema = DatasetType("schema", dimensions=registry.dimensions.empty, storageClass="Catalog")

1999 registry.registerDatasetType(schema)

2000 dataId = DataCoordinate.make_empty(registry.dimensions)

2001 run1 = "run1"

2002 run2 = "run2"

2003 registry.registerRun(run1)

2004 registry.registerRun(run2)

2005 (dataset1,) = registry.insertDatasets(schema, dataIds=[dataId], run=run1)

2006 (dataset2,) = registry.insertDatasets(schema, dataIds=[dataId], run=run2)

2007 # Query directly for both of the datasets, and each one, one at a time.

2008 self.checkQueryResults(

2009 registry.queryDatasets(schema, collections=[run1, run2], findFirst=False), [dataset1, dataset2]

2010 )

2011 self.checkQueryResults(

2012 registry.queryDatasets(schema, collections=[run1, run2], findFirst=True),

2013 [dataset1],

2014 )

2015 self.checkQueryResults(

2016 registry.queryDatasets(schema, collections=[run2, run1], findFirst=True),

2017 [dataset2],

2018 )

2019 # Query for data IDs with no dimensions.

2020 dataIds = registry.queryDataIds([])

2021 self.checkQueryResults(dataIds, [dataId])

2022 # Use queried data IDs to find the datasets.

2023 self.checkQueryResults(

2024 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),

2025 [dataset1, dataset2],

2026 )

2027 self.checkQueryResults(

2028 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2029 [dataset1],

2030 )

2031 self.checkQueryResults(

2032 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2033 [dataset2],

2034 )

2035 # Now materialize the data ID query results and repeat those tests.

2036 with dataIds.materialize() as dataIds:

2037 self.checkQueryResults(dataIds, [dataId])

2038 self.checkQueryResults(

2039 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2040 [dataset1],

2041 )

2042 self.checkQueryResults(

2043 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2044 [dataset2],

2045 )

2046 # Query for non-empty data IDs, then subset that to get the empty one.

2047 # Repeat the above tests starting from that.

2048 dataIds = registry.queryDataIds(["instrument"]).subset(registry.dimensions.empty, unique=True)

2049 self.checkQueryResults(dataIds, [dataId])

2050 self.checkQueryResults(

2051 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),

2052 [dataset1, dataset2],

2053 )

2054 self.checkQueryResults(

2055 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2056 [dataset1],

2057 )

2058 self.checkQueryResults(

2059 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2060 [dataset2],

2061 )

2062 with dataIds.materialize() as dataIds:

2063 self.checkQueryResults(dataIds, [dataId])

2064 self.checkQueryResults(

2065 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),

2066 [dataset1, dataset2],

2067 )

2068 self.checkQueryResults(

2069 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2070 [dataset1],

2071 )

2072 self.checkQueryResults(

2073 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2074 [dataset2],

2075 )

2076 # Query for non-empty data IDs, then materialize, then subset to get

2077 # the empty one. Repeat again.

2078 with registry.queryDataIds(["instrument"]).materialize() as nonEmptyDataIds:

2079 dataIds = nonEmptyDataIds.subset(registry.dimensions.empty, unique=True)

2080 self.checkQueryResults(dataIds, [dataId])

2081 self.checkQueryResults(

2082 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),

2083 [dataset1, dataset2],

2084 )

2085 self.checkQueryResults(

2086 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2087 [dataset1],

2088 )

2089 self.checkQueryResults(

2090 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2091 [dataset2],

2092 )

2093 with dataIds.materialize() as dataIds:

2094 self.checkQueryResults(dataIds, [dataId])

2095 self.checkQueryResults(

2096 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),

2097 [dataset1, dataset2],

2098 )

2099 self.checkQueryResults(

2100 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),

2101 [dataset1],

2102 )

2103 self.checkQueryResults(

2104 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),

2105 [dataset2],

2106 )

2107 # Query for non-empty data IDs with a constraint on an empty-data-ID

2108 # dataset that exists.

2109 dataIds = registry.queryDataIds(["instrument"], datasets="schema", collections=...)

2110 self.checkQueryResults(

2111 dataIds.subset(unique=True),

2112 [DataCoordinate.standardize(instrument="Cam1", universe=registry.dimensions)],

2113 )

2114 # Again query for non-empty data IDs with a constraint on empty-data-ID

2115 # datasets, but when the datasets don't exist. We delete the existing

2116 # dataset and query just that collection rather than creating a new

2117 # empty collection because this is a bit less likely for our build-time

2118 # logic to shortcut-out (via the collection summaries), and such a

2119 # shortcut would make this test a bit more trivial than we'd like.

2120 registry.removeDatasets([dataset2])

2121 dataIds = registry.queryDataIds(["instrument"], datasets="schema", collections=run2)

2122 self.checkQueryResults(dataIds, [])

2123

2124 def testDimensionDataModifications(self):

2125 """Test that modifying dimension records via:

2126 syncDimensionData(..., update=True) and

2127 insertDimensionData(..., replace=True) works as expected, even in the

2128 presence of datasets using those dimensions and spatial overlap

2129 relationships.

2130 """

2131

2132 def _unpack_range_set(ranges: lsst.sphgeom.RangeSet) -> Iterator[int]:

2133 """Unpack a sphgeom.RangeSet into the integers it contains."""

2134 for begin, end in ranges:

2135 yield from range(begin, end)

2136

2137 def _range_set_hull(

2138 ranges: lsst.sphgeom.RangeSet,

2139 pixelization: lsst.sphgeom.HtmPixelization,

2140 ) -> lsst.sphgeom.ConvexPolygon:

2141 """Create a ConvexPolygon hull of the region defined by a set of

2142 HTM pixelization index ranges.

2143 """

2144 points = []

2145 for index in _unpack_range_set(ranges):

2146 points.extend(pixelization.triangle(index).getVertices())

2147 return lsst.sphgeom.ConvexPolygon(points)

2148

2149 # Use HTM to set up an initial parent region (one arbitrary trixel)

2150 # and four child regions (the trixels within the parent at the next

2151 # level. We'll use the parent as a tract/visit region and the children

2152 # as its patch/visit_detector regions.

2153 butler = self.make_butler()

2154 registry = butler.registry

2155 htm6 = registry.dimensions.skypix["htm"][6].pixelization

2156 commonSkyPix = registry.dimensions.commonSkyPix.pixelization

2157 index = 12288

2158 child_ranges_small = lsst.sphgeom.RangeSet(index).scaled(4)

2159 assert htm6.universe().contains(child_ranges_small)

2160 child_regions_small = [htm6.triangle(i) for i in _unpack_range_set(child_ranges_small)]

2161 parent_region_small = lsst.sphgeom.ConvexPolygon(

2162 list(itertools.chain.from_iterable(c.getVertices() for c in child_regions_small))

2163 )

2164 assert all(parent_region_small.contains(c) for c in child_regions_small)

2165 # Make a larger version of each child region, defined to be the set of

2166 # htm6 trixels that overlap the original's bounding circle. Make a new

2167 # parent that's the convex hull of the new children.

2168 child_regions_large = [

2169 _range_set_hull(htm6.envelope(c.getBoundingCircle()), htm6) for c in child_regions_small

2170 ]

2171 assert all(

2172 large.contains(small)

2173 for large, small in zip(child_regions_large, child_regions_small, strict=True)

2174 )

2175 parent_region_large = lsst.sphgeom.ConvexPolygon(

2176 list(itertools.chain.from_iterable(c.getVertices() for c in child_regions_large))

2177 )

2178 assert all(parent_region_large.contains(c) for c in child_regions_large)

2179 assert parent_region_large.contains(parent_region_small)

2180 assert not parent_region_small.contains(parent_region_large)

2181 assert not all(parent_region_small.contains(c) for c in child_regions_large)

2182 # Find some commonSkyPix indices that overlap the large regions but not

2183 # overlap the small regions. We use commonSkyPix here to make sure the

2184 # real tests later involve what's in the database, not just post-query

2185 # filtering of regions.

2186 child_difference_indices = []

2187 for large, small in zip(child_regions_large, child_regions_small, strict=True):

2188 difference = list(_unpack_range_set(commonSkyPix.envelope(large) - commonSkyPix.envelope(small)))

2189 assert difference, "if this is empty, we can't test anything useful with these regions"

2190 assert all(

2191 not commonSkyPix.triangle(d).isDisjointFrom(large)

2192 and commonSkyPix.triangle(d).isDisjointFrom(small)

2193 for d in difference

2194 )

2195 child_difference_indices.append(difference)

2196 parent_difference_indices = list(

2197 _unpack_range_set(

2198 commonSkyPix.envelope(parent_region_large) - commonSkyPix.envelope(parent_region_small)

2199 )

2200 )

2201 assert parent_difference_indices, "if this is empty, we can't test anything useful with these regions"

2202 assert all(

2203 (

2204 not commonSkyPix.triangle(d).isDisjointFrom(parent_region_large)

2205 and commonSkyPix.triangle(d).isDisjointFrom(parent_region_small)

2206 )

2207 for d in parent_difference_indices

2208 )

2209 # Now that we've finally got those regions, we'll insert the large ones

2210 # as tract/patch dimension records.

2211 skymap_name = "testing_v1"

2212 registry.insertDimensionData(

2213 "skymap",

2214 {

2215 "name": skymap_name,

2216 "hash": bytes([42]),

2217 "tract_max": 1,

2218 "patch_nx_max": 2,

2219 "patch_ny_max": 2,

2220 },

2221 )

2222 registry.insertDimensionData("tract", {"skymap": skymap_name, "id": 0, "region": parent_region_large})

2223 registry.insertDimensionData(

2224 "patch",

2225 *[

2226 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}

2227 for n, c in enumerate(child_regions_large)

2228 ],

2229 )

2230 # Add at dataset that uses these dimensions to make sure that modifying

2231 # them doesn't disrupt foreign keys (need to make sure DB doesn't

2232 # implement insert with replace=True as delete-then-insert).

2233 dataset_type = DatasetType(

2234 "coadd",

2235 dimensions=["tract", "patch"],

2236 universe=registry.dimensions,

2237 storageClass="Exposure",

2238 )

2239 registry.registerDatasetType(dataset_type)

2240 registry.registerCollection("the_run", CollectionType.RUN)

2241 registry.insertDatasets(

2242 dataset_type,

2243 [{"skymap": skymap_name, "tract": 0, "patch": 2}],

2244 run="the_run",

2245 )

2246 # Query for tracts and patches that overlap some "difference" htm9

2247 # pixels; there should be overlaps, because the database has

2248 # the "large" suite of regions.

2249 self.assertEqual(

2250 {0},

2251 {

2252 data_id["tract"]

2253 for data_id in registry.queryDataIds(

2254 ["tract"],

2255 skymap=skymap_name,

2256 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},

2257 )

2258 },

2259 )

2260 for patch_id, patch_difference_indices in enumerate(child_difference_indices):

2261 self.assertIn(

2262 patch_id,

2263 {

2264 data_id["patch"]

2265 for data_id in registry.queryDataIds(

2266 ["patch"],

2267 skymap=skymap_name,

2268 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},

2269 )

2270 },

2271 )

2272 # Use sync to update the tract region and insert to update the regions

2273 # of the patches, to the "small" suite.

2274 updated = registry.syncDimensionData(

2275 "tract",

2276 {"skymap": skymap_name, "id": 0, "region": parent_region_small},

2277 update=True,

2278 )

2279 self.assertEqual(updated, {"region": parent_region_large})

2280 registry.insertDimensionData(

2281 "patch",

2282 *[

2283 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}

2284 for n, c in enumerate(child_regions_small)

2285 ],

2286 replace=True,

2287 )

2288 # Query again; there now should be no such overlaps, because the

2289 # database has the "small" suite of regions.

2290 self.assertFalse(

2291 set(

2292 registry.queryDataIds(

2293 ["tract"],

2294 skymap=skymap_name,

2295 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},

2296 )

2297 )

2298 )

2299 for patch_id, patch_difference_indices in enumerate(child_difference_indices):

2300 self.assertNotIn(

2301 patch_id,

2302 {

2303 data_id["patch"]

2304 for data_id in registry.queryDataIds(

2305 ["patch"],

2306 skymap=skymap_name,

2307 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},

2308 )

2309 },

2310 )

2311 # Update back to the large regions and query one more time.

2312 updated = registry.syncDimensionData(

2313 "tract",

2314 {"skymap": skymap_name, "id": 0, "region": parent_region_large},

2315 update=True,

2316 )

2317 self.assertEqual(updated, {"region": parent_region_small})

2318 registry.insertDimensionData(

2319 "patch",

2320 *[

2321 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}

2322 for n, c in enumerate(child_regions_large)

2323 ],

2324 replace=True,

2325 )

2326 self.assertEqual(

2327 {0},

2328 {

2329 data_id["tract"]

2330 for data_id in registry.queryDataIds(

2331 ["tract"],

2332 skymap=skymap_name,

2333 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},

2334 )

2335 },

2336 )

2337 for patch_id, patch_difference_indices in enumerate(child_difference_indices):

2338 self.assertIn(

2339 patch_id,

2340 {

2341 data_id["patch"]

2342 for data_id in registry.queryDataIds(

2343 ["patch"],

2344 skymap=skymap_name,

2345 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},

2346 )

2347 },

2348 )

2349

2350 def testCalibrationCollections(self):

2351 """Test operations on `~CollectionType.CALIBRATION` collections,

2352 including `SqlRegistry.certify`, `SqlRegistry.decertify`,

2353 `SqlRegistry.findDataset`, and

2354 `DataCoordinateQueryResults.findRelatedDatasets`.

2355 """

2356 # Setup - make a Registry, fill it with some datasets in

2357 # non-calibration collections.

2358 butler = self.make_butler()

2359 registry = butler.registry

2360 self.load_data(butler, "base.yaml", "datasets.yaml")

2361 # Set up some timestamps.

2362 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")

2363 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai")

2364 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai")

2365 t4 = astropy.time.Time("2020-01-01T04:00:00", format="isot", scale="tai")

2366 t5 = astropy.time.Time("2020-01-01T05:00:00", format="isot", scale="tai")

2367 allTimespans = [

2368 Timespan(a, b) for a, b in itertools.combinations([None, t1, t2, t3, t4, t5, None], r=2)

2369 ]

2370 # Insert some exposure records with timespans between each sequential

2371 # pair of those.

2372 registry.insertDimensionData(

2373 "day_obs", {"instrument": "Cam1", "id": 20200101, "timespan": Timespan(t1, t5)}

2374 )

2375 registry.insertDimensionData(

2376 "group",

2377 {"instrument": "Cam1", "name": "group0"},

2378 {"instrument": "Cam1", "name": "group1"},

2379 {"instrument": "Cam1", "name": "group2"},

2380 {"instrument": "Cam1", "name": "group3"},

2381 )

2382 registry.insertDimensionData(

2383 "exposure",

2384 {

2385 "instrument": "Cam1",

2386 "id": 0,

2387 "group": "group0",

2388 "obs_id": "zero",

2389 "physical_filter": "Cam1-G",

2390 "day_obs": 20200101,

2391 "timespan": Timespan(t1, t2),

2392 },

2393 {

2394 "instrument": "Cam1",

2395 "id": 1,

2396 "group": "group1",

2397 "obs_id": "one",

2398 "physical_filter": "Cam1-G",

2399 "day_obs": 20200101,

2400 "timespan": Timespan(t2, t3),

2401 },

2402 {

2403 "instrument": "Cam1",

2404 "id": 2,

2405 "group": "group2",

2406 "obs_id": "two",

2407 "physical_filter": "Cam1-G",

2408 "day_obs": 20200101,

2409 "timespan": Timespan(t3, t4),

2410 },

2411 {

2412 "instrument": "Cam1",

2413 "id": 3,

2414 "group": "group3",

2415 "obs_id": "three",

2416 "physical_filter": "Cam1-G",

2417 "day_obs": 20200101,

2418 "timespan": Timespan(t4, t5),

2419 },

2420 )

2421 # Get references to some datasets.

2422 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g")

2423 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g")

2424 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r")

2425 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r")

2426 # Register the main calibration collection we'll be working with.

2427 collection = "Cam1/calibs/default"

2428 registry.registerCollection(collection, type=CollectionType.CALIBRATION)

2429 # Cannot associate into a calibration collection (no timespan).

2430 with self.assertRaises(CollectionTypeError):

2431 registry.associate(collection, [bias2a])

2432 # Certify 2a dataset with [t2, t4) validity.

2433 registry.certify(collection, [bias2a], Timespan(begin=t2, end=t4))

2434 # Test that we can query for this dataset via the new collection, both

2435 # on its own and with a RUN collection.

2436 self.assertEqual(

2437 set(registry.queryDatasets("bias", findFirst=False, collections=collection)),

2438 {bias2a},

2439 )

2440 self.assertEqual(

2441 set(registry.queryDatasets("bias", findFirst=False, collections=[collection, "imported_r"])),

2442 {

2443 bias2a,

2444 bias2b,

2445 bias3b,

2446 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),

2447 },

2448 )

2449 self.assertEqual(

2450 set(registry.queryDataIds("detector", datasets="bias", collections=collection)),

2451 {registry.expandDataId(instrument="Cam1", detector=2)},

2452 )

2453 self.assertEqual(

2454 set(registry.queryDataIds("detector", datasets="bias", collections=[collection, "imported_r"])),

2455 {

2456 registry.expandDataId(instrument="Cam1", detector=2),

2457 registry.expandDataId(instrument="Cam1", detector=3),

2458 registry.expandDataId(instrument="Cam1", detector=4),

2459 },

2460 )

2461 self.assertEqual(

2462 set(

2463 registry.queryDataIds(["exposure", "detector"]).findRelatedDatasets(

2464 "bias", findFirst=True, collections=[collection]

2465 )

2466 ),

2467 {

2468 (registry.expandDataId(instrument="Cam1", detector=2, exposure=1), bias2a),

2469 (registry.expandDataId(instrument="Cam1", detector=2, exposure=2), bias2a),

2470 },

2471 )

2472

2473 # We should not be able to certify 2b with anything overlapping that

2474 # window.

2475 with self.assertRaises(ConflictingDefinitionError):

2476 registry.certify(collection, [bias2b], Timespan(begin=None, end=t3))

2477 with self.assertRaises(ConflictingDefinitionError):

2478 registry.certify(collection, [bias2b], Timespan(begin=None, end=t5))

2479 with self.assertRaises(ConflictingDefinitionError):

2480 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t3))

2481 with self.assertRaises(ConflictingDefinitionError):

2482 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t5))

2483 with self.assertRaises(ConflictingDefinitionError):

2484 registry.certify(collection, [bias2b], Timespan(begin=t1, end=None))

2485 with self.assertRaises(ConflictingDefinitionError):

2486 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t3))

2487 with self.assertRaises(ConflictingDefinitionError):

2488 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t5))

2489 with self.assertRaises(ConflictingDefinitionError):

2490 registry.certify(collection, [bias2b], Timespan(begin=t2, end=None))

2491 # We should be able to certify 3a with a range overlapping that window,

2492 # because it's for a different detector.

2493 # We'll certify 3a over [t1, t3).

2494 registry.certify(collection, [bias3a], Timespan(begin=t1, end=t3))

2495 # Now we'll certify 2b and 3b together over [t4, ∞).

2496 registry.certify(collection, [bias2b, bias3b], Timespan(begin=t4, end=None))

2497

2498 # Fetch all associations and check that they are what we expect.

2499 self.assertCountEqual(

2500 list(

2501 registry.queryDatasetAssociations(

2502 "bias",

2503 collections=[collection, "imported_g", "imported_r"],

2504 )

2505 ),

2506 [

2507 DatasetAssociation(

2508 ref=registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),

2509 collection="imported_g",

2510 timespan=None,

2511 ),

2512 DatasetAssociation(

2513 ref=registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),

2514 collection="imported_r",

2515 timespan=None,

2516 ),

2517 DatasetAssociation(ref=bias2a, collection="imported_g", timespan=None),

2518 DatasetAssociation(ref=bias3a, collection="imported_g", timespan=None),

2519 DatasetAssociation(ref=bias2b, collection="imported_r", timespan=None),

2520 DatasetAssociation(ref=bias3b, collection="imported_r", timespan=None),

2521 DatasetAssociation(ref=bias2a, collection=collection, timespan=Timespan(begin=t2, end=t4)),

2522 DatasetAssociation(ref=bias3a, collection=collection, timespan=Timespan(begin=t1, end=t3)),

2523 DatasetAssociation(ref=bias2b, collection=collection, timespan=Timespan(begin=t4, end=None)),

2524 DatasetAssociation(ref=bias3b, collection=collection, timespan=Timespan(begin=t4, end=None)),

2525 ],

2526 )

2527

2528 class Ambiguous:

2529 """Tag class to denote lookups that should be ambiguous."""

2530

2531 pass

2532

2533 def _assertLookup(

2534 detector: int, timespan: Timespan, expected: DatasetRef | type[Ambiguous] | None

2535 ) -> None:

2536 """Local function that asserts that a bias lookup returns the given

2537 expected result.

2538 """

2539 if expected is Ambiguous:

2540 with self.assertRaises((DatasetTypeError, LookupError)):

2541 registry.findDataset(

2542 "bias",

2543 collections=collection,

2544 instrument="Cam1",

2545 detector=detector,

2546 timespan=timespan,

2547 )

2548 else:

2549 self.assertEqual(

2550 expected,

2551 registry.findDataset(

2552 "bias",

2553 collections=collection,

2554 instrument="Cam1",

2555 detector=detector,

2556 timespan=timespan,

2557 ),

2558 )

2559

2560 # Systematically test lookups against expected results.

2561 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)

2562 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)

2563 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)

2564 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)

2565 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=Ambiguous)

2566 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)

2567 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)

2568 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)

2569 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)

2570 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=Ambiguous)

2571 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)

2572 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)

2573 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)

2574 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=Ambiguous)

2575 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)

2576 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=bias2a)

2577 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=Ambiguous)

2578 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=Ambiguous)

2579 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=bias2b)

2580 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)

2581 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)

2582 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)

2583 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)

2584 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)

2585 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)

2586 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=Ambiguous)

2587 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)

2588 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)

2589 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)

2590 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)

2591 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=Ambiguous)

2592 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)

2593 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)

2594 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)

2595 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=Ambiguous)

2596 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)

2597 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)

2598 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=bias3b)

2599 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)

2600 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=bias3b)

2601 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)

2602 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)

2603

2604 # Test lookups via temporal joins to exposures.

2605 self.assertEqual(

2606 set(

2607 registry.queryDataIds(

2608 ["exposure", "detector"], instrument="Cam1", detector=2

2609 ).findRelatedDatasets("bias", collections=[collection])

2610 ),

2611 {

2612 (registry.expandDataId(instrument="Cam1", exposure=1, detector=2), bias2a),

2613 (registry.expandDataId(instrument="Cam1", exposure=2, detector=2), bias2a),

2614 (registry.expandDataId(instrument="Cam1", exposure=3, detector=2), bias2b),

2615 },

2616 )

2617 self.assertEqual(

2618 set(

2619 registry.queryDataIds(

2620 ["exposure", "detector"], instrument="Cam1", detector=3

2621 ).findRelatedDatasets("bias", collections=[collection])

2622 ),

2623 {

2624 (registry.expandDataId(instrument="Cam1", exposure=0, detector=3), bias3a),

2625 (registry.expandDataId(instrument="Cam1", exposure=1, detector=3), bias3a),

2626 (registry.expandDataId(instrument="Cam1", exposure=3, detector=3), bias3b),

2627 },

2628 )

2629

2630 # Decertify [t3, t5) for all data IDs, and do test lookups again.

2631 # This should truncate bias2a to [t2, t3), leave bias3a unchanged at

2632 # [t1, t3), and truncate bias2b and bias3b to [t5, ∞).

2633 registry.decertify(collection=collection, datasetType="bias", timespan=Timespan(t3, t5))

2634 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)

2635 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)

2636 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)

2637 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)

2638 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=bias2a)

2639 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)

2640 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)

2641 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)

2642 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)

2643 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=bias2a)

2644 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)

2645 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)

2646 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)

2647 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=bias2a)

2648 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)

2649 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=None)

2650 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=None)

2651 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=bias2b)

2652 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=None)

2653 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)

2654 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)

2655 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)

2656 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)

2657 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)

2658 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)

2659 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=bias3a)

2660 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)

2661 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)

2662 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)

2663 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)

2664 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=bias3a)

2665 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)

2666 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)

2667 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)

2668 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=bias3a)

2669 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)

2670 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)

2671 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=None)

2672 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)

2673 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=None)

2674 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)

2675 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)

2676

2677 # Decertify everything, this time with explicit data IDs, then check

2678 # that no lookups succeed.

2679 registry.decertify(

2680 collection,

2681 "bias",

2682 Timespan(None, None),

2683 dataIds=[

2684 dict(instrument="Cam1", detector=2),

2685 dict(instrument="Cam1", detector=3),

2686 ],

2687 )

2688 for detector in (2, 3):

2689 for timespan in allTimespans:

2690 _assertLookup(detector=detector, timespan=timespan, expected=None)

2691 # Certify bias2a and bias3a over (-∞, ∞), check that all lookups return

2692 # those.

2693 registry.certify(

2694 collection,

2695 [bias2a, bias3a],

2696 Timespan(None, None),

2697 )

2698 for timespan in allTimespans:

2699 _assertLookup(detector=2, timespan=timespan, expected=bias2a)

2700 _assertLookup(detector=3, timespan=timespan, expected=bias3a)

2701 # Decertify just bias2 over [t2, t4).

2702 # This should split a single certification row into two (and leave the

2703 # other existing row, for bias3a, alone).

2704 registry.decertify(

2705 collection, "bias", Timespan(t2, t4), dataIds=[dict(instrument="Cam1", detector=2)]

2706 )

2707 for timespan in allTimespans:

2708 _assertLookup(detector=3, timespan=timespan, expected=bias3a)

2709 overlapsBefore = timespan.overlaps(Timespan(None, t2))

2710 overlapsAfter = timespan.overlaps(Timespan(t4, None))

2711 if overlapsBefore and overlapsAfter:

2712 expected = Ambiguous

2713 elif overlapsBefore or overlapsAfter:

2714 expected = bias2a

2715 else:

2716 expected = None

2717 _assertLookup(detector=2, timespan=timespan, expected=expected)

2718

2719 def testSkipCalibs(self):

2720 """Test how queries handle skipping of calibration collections."""

2721 butler = self.make_butler()

2722 registry = butler.registry

2723 self.load_data(butler, "base.yaml", "datasets.yaml")

2724

2725 coll_calib = "Cam1/calibs/default"

2726 registry.registerCollection(coll_calib, type=CollectionType.CALIBRATION)

2727

2728 # Add all biases to the calibration collection.

2729 # Without this, the logic that prunes dataset subqueries based on

2730 # datasetType-collection summary information will fire before the logic

2731 # we want to test below. This is a good thing (it avoids the dreaded

2732 # NotImplementedError a bit more often) everywhere but here.

2733 registry.certify(coll_calib, registry.queryDatasets("bias", collections=...), Timespan(None, None))

2734

2735 coll_list = [coll_calib, "imported_g", "imported_r"]

2736 chain = "Cam1/chain"

2737 registry.registerCollection(chain, type=CollectionType.CHAINED)

2738 registry.setCollectionChain(chain, coll_list)

2739

2740 # Lookup is ambiguous due to multiple datasets with the same data ID

2741 # in the calibration collection.

2742 with self.assertRaises(CalibrationLookupError):

2743 list(registry.queryDatasets("bias", collections=coll_list, findFirst=True))

2744

2745 # chain will skip

2746 datasets = list(registry.queryDatasets("bias", collections=chain))

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

2748

2749 dataIds = list(registry.queryDataIds(["instrument", "detector"], datasets="bias", collections=chain))

2750 self.assertGreater(len(dataIds), 0)

2751

2752 # glob will skip too

2753 datasets = list(registry.queryDatasets("bias", collections="*d*"))

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

2755

2756 # regular expression will skip too

2757 if self.supportsCollectionRegex:

2758 pattern = re.compile(".*")

2759 with self.assertWarns(FutureWarning):

2760 datasets = list(registry.queryDatasets("bias", collections=pattern))

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

2762

2763 # ellipsis should work as usual

2764 datasets = list(registry.queryDatasets("bias", collections=...))

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

2766

2767 # New query system correctly determines that this search is

2768 # ambiguous, because there are multiple datasets with the same

2769 # {instrument=Cam1, detector=2} data ID in the calibration

2770 # collection at the beginning of the chain.

2771 with self.assertRaises(CalibrationLookupError):

2772 datasets = list(registry.queryDatasets("bias", collections=chain, findFirst=True))

2773

2774 def testIngestTimeQuery(self):

2775 butler = self.make_butler()

2776 registry = butler.registry

2777 dt0 = datetime.datetime.now(datetime.UTC)

2778 self.load_data(butler, "base.yaml", "datasets.yaml")

2779 dt1 = datetime.datetime.now(datetime.UTC)

2780

2781 datasets = list(registry.queryDatasets(..., collections=...))

2782 len0 = len(datasets)

2783 self.assertGreater(len0, 0)

2784

2785 for where in ("ingest_date > T'2000-01-01'", "T'2000-01-01' < ingest_date"):

2786 datasets = list(registry.queryDatasets(..., collections=..., where=where))

2787 len1 = len(datasets)

2788 self.assertEqual(len0, len1)

2789

2790 # no one will ever use this piece of software in 30 years

2791 for where in ("ingest_date > T'2050-01-01'", "T'2050-01-01' < ingest_date"):

2792 datasets = list(registry.queryDatasets(..., collections=..., where=where))

2793 len2 = len(datasets)

2794 self.assertEqual(len2, 0)

2795

2796 # Check more exact timing to make sure there is no 37 seconds offset

2797 # (after fixing DM-30124). SQLite time precision is 1 second, make

2798 # sure that we don't test with higher precision.

2799 tests = [

2800 # format: (timestamp, operator, expected_len)

2801 (dt0 - timedelta(seconds=1), ">", len0),

2802 (dt0 - timedelta(seconds=1), "<", 0),

2803 (dt1 + timedelta(seconds=1), "<", len0),

2804 (dt1 + timedelta(seconds=1), ">", 0),

2805 ]

2806 for dt, op, expect_len in tests:

2807 dt_str = dt.isoformat(sep=" ")

2808

2809 where = f"ingest_date {op} T'{dt_str}'"

2810 datasets = list(registry.queryDatasets(..., collections=..., where=where))

2811 self.assertEqual(len(datasets), expect_len)

2812

2813 # same with bind using datetime or astropy Time

2814 where = f"ingest_date {op} :ingest_time"

2815 datasets = list(

2816 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt})

2817 )

2818 self.assertEqual(len(datasets), expect_len)

2819

2820 dt_astropy = astropy.time.Time(dt, format="datetime")

2821 datasets = list(

2822 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt_astropy})

2823 )

2824 self.assertEqual(len(datasets), expect_len)

2825

2826 def testTimespanQueries(self):

2827 """Test query expressions involving timespans."""

2828 butler = self.make_butler()

2829 registry = butler.registry

2830 self.load_data(butler, "ci_hsc-subset.yaml")

2831 # All exposures in the database; mapping from ID to timespan.

2832 visits = {record.id: record.timespan for record in registry.queryDimensionRecords("visit")}

2833 # Just those IDs, sorted (which is also temporal sorting, because HSC

2834 # exposure IDs are monotonically increasing).

2835 ids = sorted(visits.keys())

2836 self.assertEqual(len(ids), 11)

2837 # Pick some quasi-random indexes into `ids` to play with.

2838 i1 = int(len(ids) * 0.1)

2839 i2 = int(len(ids) * 0.3)

2840 i3 = int(len(ids) * 0.6)

2841 i4 = int(len(ids) * 0.8)

2842 # Extract some times from those: just before the beginning of i1 (which

2843 # should be after the end of the exposure before), exactly the

2844 # beginning of i2, just after the beginning of i3 (and before its end),

2845 # and the exact end of i4.

2846 t1 = visits[ids[i1]].begin - astropy.time.TimeDelta(1.0, format="sec")

2847 self.assertGreater(t1, visits[ids[i1 - 1]].end)

2848 t2 = visits[ids[i2]].begin

2849 t3 = visits[ids[i3]].begin + astropy.time.TimeDelta(1.0, format="sec")

2850 self.assertLess(t3, visits[ids[i3]].end)

2851 t4 = visits[ids[i4]].end

2852 # Make sure those are actually in order.

2853 self.assertEqual([t1, t2, t3, t4], sorted([t4, t3, t2, t1]))

2854

2855 bind = {

2856 "t1": t1,

2857 "t2": t2,

2858 "t3": t3,

2859 "t4": t4,

2860 "ts23": Timespan(t2, t3),

2861 }

2862

2863 def query(where):

2864 """Return results as a sorted, deduplicated list of visit IDs.

2865

2866 Parameters

2867 ----------

2868 where : `str`

2869 The WHERE clause for the query.

2870 """

2871 return sorted(

2872 {

2873 dataId["visit"]

2874 for dataId in registry.queryDataIds("visit", instrument="HSC", bind=bind, where=where)

2875 }

2876 )

2877

2878 # Try a bunch of timespan queries, mixing up the bounds themselves,

2879 # where they appear in the expression, and how we get the timespan into

2880 # the expression.

2881

2882 # t1 is before the start of i1, so this should not include i1.

2883 self.assertEqual(ids[:i1], query("visit.timespan OVERLAPS (null, :t1)"))

2884 # t2 is exactly at the start of i2, but ends are exclusive, so these

2885 # should not include i2.

2886 self.assertEqual(ids[i1:i2], query("(:t1, :t2) OVERLAPS visit.timespan"))

2887 # t3 is in the middle of i3, so this should include i3.

2888 self.assertEqual(ids[i2 : i3 + 1], query("visit.timespan OVERLAPS :ts23"))

2889 # This one should not include t3 by the same reasoning.

2890 # t4 is exactly at the end of i4, so this should include i4.

2891 self.assertEqual(ids[i3 : i4 + 1], query(f"visit.timespan OVERLAPS (T'{t3.tai.isot}/tai', :t4)"))

2892 # i4's upper bound of t4 is exclusive so this should not include t4.

2893 self.assertEqual(ids[i4 + 1 :], query("visit.timespan OVERLAPS (:t4, NULL)"))

2894

2895 # Now some timespan vs. time scalar queries.

2896 self.assertEqual(ids[i3 : i3 + 1], query("visit.timespan OVERLAPS :t3"))

2897 self.assertEqual(ids[i3 : i3 + 1], query(f"T'{t3.tai.isot}/tai' OVERLAPS visit.timespan"))

2898

2899 # Empty timespans should not overlap anything.

2900 self.assertEqual([], query("visit.timespan OVERLAPS (:t3, :t2)"))

2901

2902 # Make sure that expanded data IDs include the timespans.

2903 results = list(

2904 registry.queryDataIds(["visit"], dataId={"instrument": "HSC", "visit": ids[1]}).expanded()

2905 )

2906 self.assertEqual(len(results), 1)

2907 visit_timespan = visits[ids[1]]

2908 self.assertEqual(results[0].timespan, visit_timespan)

2909 visit_record = results[0].records["visit"]

2910 assert visit_record is not None

2911 self.assertEqual(visit_record.timespan, visit_timespan)

2912 day_obs_record = results[0].records["day_obs"]

2913 assert day_obs_record is not None

2914 self.assertEqual(day_obs_record.id, 20130617)

2915 self.assertEqual(

2916 day_obs_record.timespan,

2917 Timespan(

2918 astropy.time.Time("2013-06-17T00:00:00", scale="tai"),

2919 astropy.time.Time("2013-06-18T00:00:00", scale="tai"),

2920 ),

2921 )

2922

2923 def testCollectionSummaries(self):

2924 """Test recording and retrieval of collection summaries."""

2925 self.maxDiff = None

2926 butler = self.make_butler()

2927 registry = butler.registry

2928 # Importing datasets from yaml should go through the code path where

2929 # we update collection summaries as we insert datasets.

2930 self.load_data(butler, "base.yaml", "datasets.yaml")

2931 flat = registry.getDatasetType("flat")

2932 expected1 = CollectionSummary()

2933 expected1.dataset_types.add(registry.getDatasetType("bias"))

2934 expected1.add_data_ids(

2935 flat, [DataCoordinate.standardize(instrument="Cam1", universe=registry.dimensions)]

2936 )

2937 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)

2938 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)

2939 # Create a chained collection with both of the imported runs; the

2940 # summary should be the same, because it's a union with itself.

2941 chain = "chain"

2942 registry.registerCollection(chain, CollectionType.CHAINED)

2943 registry.setCollectionChain(chain, ["imported_r", "imported_g"])

2944 self.assertEqual(registry.getCollectionSummary(chain), expected1)

2945 # Associate flats only into a tagged collection and a calibration

2946 # collection to check summaries of those.

2947 tag = "tag"

2948 registry.registerCollection(tag, CollectionType.TAGGED)

2949 registry.associate(tag, registry.queryDatasets(flat, collections="imported_g"))

2950 calibs = "calibs"

2951 registry.registerCollection(calibs, CollectionType.CALIBRATION)

2952 registry.certify(

2953 calibs, registry.queryDatasets(flat, collections="imported_g"), timespan=Timespan(None, None)

2954 )

2955 expected2 = expected1.copy()

2956 expected2.dataset_types.discard("bias")

2957 self.assertEqual(registry.getCollectionSummary(tag), expected2)

2958 self.assertEqual(registry.getCollectionSummary(calibs), expected2)

2959 # Explicitly calling SqlRegistry.refresh() should load those same

2960 # summaries, via a totally different code path.

2961 registry.refresh()

2962 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)

2963 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)

2964 self.assertEqual(registry.getCollectionSummary(tag), expected2)

2965 self.assertEqual(registry.getCollectionSummary(calibs), expected2)

2966

2967 def testBindInQueryDatasets(self):

2968 """Test that the bind parameter is correctly forwarded in

2969 queryDatasets recursion.

2970 """

2971 butler = self.make_butler()

2972 registry = butler.registry

2973 # Importing datasets from yaml should go through the code path where

2974 # we update collection summaries as we insert datasets.

2975 self.load_data(butler, "base.yaml", "datasets.yaml")

2976 self.assertEqual(

2977 set(registry.queryDatasets("flat", band="r", collections=...)),

2978 set(

2979 registry.queryDatasets("flat", where="band=:my_band", bind={"my_band": "r"}, collections=...)

2980 ),

2981 )

2982

2983 def testQueryIntRangeExpressions(self):

2984 """Test integer range expressions in ``where`` arguments.

2985

2986 Note that our expressions use inclusive stop values, unlike Python's.

2987 """

2988 butler = self.make_butler()

2989 registry = butler.registry

2990 self.load_data(butler, "base.yaml")

2991 self.assertEqual(

2992 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..2)")),

2993 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 2]},

2994 )

2995 self.assertEqual(

2996 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..4:2)")),

2997 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 3]},

2998 )

2999 self.assertEqual(

3000 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (2..4:2)")),

3001 {registry.expandDataId(instrument="Cam1", detector=n) for n in [2, 4]},

3002 )

3003

3004 def testQueryResultSummaries(self):

3005 """Test summary methods like `count`, `any`, and `explain_no_results`

3006 on `DataCoordinateQueryResults` and `DatasetQueryResults`.

3007 """

3008 butler = self.make_butler()

3009 registry = butler.registry

3010 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3011 # Default test dataset has two collections, each with both flats and

3012 # biases. Add a new collection with only biases.

3013 registry.registerCollection("biases", CollectionType.TAGGED)

3014 registry.associate("biases", registry.queryDatasets("bias", collections=["imported_g"]))

3015 # First query yields two results, and involves no postprocessing.

3016 query1 = registry.queryDataIds(["physical_filter"], band="r")

3017 self.assertTrue(query1.any(execute=False, exact=False))

3018 self.assertTrue(query1.any(execute=True, exact=False))

3019 self.assertTrue(query1.any(execute=True, exact=True))

3020 self.assertEqual(query1.count(exact=False), 2)

3021 self.assertEqual(query1.count(exact=True), 2)

3022 self.assertFalse(list(query1.explain_no_results()))

3023 # Second query should yield no results, which we should see when

3024 # we attempt to expand the data ID.

3025 query2 = registry.queryDataIds(["physical_filter"], band="h")

3026 # There's no execute=False, exact=False test here because the behavior

3027 # not something we want to guarantee in this case (and exact=False

3028 # says either answer is legal).

3029 self.assertFalse(query2.any(execute=True, exact=False))

3030 self.assertFalse(query2.any(execute=True, exact=True))

3031 self.assertEqual(query2.count(exact=False), 0)

3032 self.assertEqual(query2.count(exact=True), 0)

3033 # These queries yield no results due to various problems that can be

3034 # spotted prior to execution, yielding helpful diagnostics.

3035 base_query = registry.queryDataIds(["detector", "physical_filter"])

3036 queries_and_snippets = [

3037 (

3038 # Dataset type name doesn't match any existing dataset types.

3039 registry.queryDatasets("nonexistent", collections=...),

3040 ["nonexistent"],

3041 ),

3042 (

3043 # Dataset type object isn't registered.

3044 registry.queryDatasets(

3045 DatasetType(

3046 "nonexistent",

3047 dimensions=["instrument"],

3048 universe=registry.dimensions,

3049 storageClass="Image",

3050 ),

3051 collections=...,

3052 ),

3053 ["nonexistent"],

3054 ),

3055 (

3056 # No datasets of this type in this collection.

3057 registry.queryDatasets("flat", collections=["biases"]),

3058 ["flat", "biases"],

3059 ),

3060 (

3061 # No datasets of this type in this collection.

3062 base_query.findDatasets("flat", collections=["biases"]),

3063 ["flat", "biases"],

3064 ),

3065 (

3066 # No collections matching at all.

3067 registry.queryDatasets("flat", collections="potato*"),

3068 ["potato"],

3069 ),

3070 ]

3071 with self.assertRaises(MissingDatasetTypeError):

3072 # Dataset type name doesn't match any existing dataset types.

3073 list(registry.queryDataIds(["detector"], datasets=["nonexistent"], collections=...))

3074 with self.assertRaises(MissingDatasetTypeError):

3075 # Dataset type name doesn't match any existing dataset types.

3076 registry.queryDimensionRecords("detector", datasets=["nonexistent"], collections=...).any()

3077 with self.assertRaises(DatasetTypeExpressionError):

3078 # queryDimensionRecords does not allow dataset type wildcards.

3079 registry.queryDimensionRecords("detector", datasets=["f*"], collections=...).any()

3080 for query, snippets in queries_and_snippets:

3081 self.assertFalse(query.any(execute=False, exact=False))

3082 self.assertFalse(query.any(execute=True, exact=False))

3083 self.assertFalse(query.any(execute=True, exact=True))

3084 self.assertEqual(query.count(exact=False), 0)

3085 self.assertEqual(query.count(exact=True), 0)

3086 messages = list(query.explain_no_results())

3087 self.assertTrue(messages)

3088 # Want all expected snippets to appear in at least one message.

3089 self.assertTrue(

3090 any(

3091 all(snippet in message for snippet in snippets) for message in query.explain_no_results()

3092 ),

3093 messages,

3094 )

3095

3096 # Wildcards on dataset types are not permitted in queryDataIds.

3097 with self.assertRaises(DatasetTypeExpressionError):

3098 registry.queryDataIds(["detector"], datasets=re.compile("^nonexistent$"), collections=...)

3099

3100 # This query yields four overlaps in the database, but one is filtered

3101 # out in postprocessing. The count queries aren't accurate because

3102 # they don't account for duplication that happens due to an internal

3103 # join against commonSkyPix.

3104 query3 = registry.queryDataIds(["visit", "tract"], instrument="Cam1", skymap="SkyMap1")

3105 self.assertEqual(

3106 {

3107 DataCoordinate.standardize(

3108 instrument="Cam1",

3109 skymap="SkyMap1",

3110 visit=v,

3111 tract=t,

3112 universe=registry.dimensions,

3113 )

3114 for v, t in [(1, 0), (2, 0), (2, 1)]

3115 },

3116 set(query3),

3117 )

3118 self.assertTrue(query3.any(execute=False, exact=False))

3119 self.assertTrue(query3.any(execute=True, exact=False))

3120 self.assertTrue(query3.any(execute=True, exact=True))

3121 self.assertGreaterEqual(query3.count(exact=False), 3)

3122 self.assertGreaterEqual(query3.count(exact=True, discard=True), 3)

3123 self.assertFalse(list(query3.explain_no_results()))

3124 # This query yields overlaps in the database, but all are filtered

3125 # out in postprocessing. The count queries again aren't very useful.

3126 # We have to use `where=` here to avoid an optimization that

3127 # (currently) skips the spatial postprocess-filtering because it

3128 # recognizes that no spatial join is necessary. That's not ideal, but

3129 # fixing it is out of scope for this ticket.

3130 query4 = registry.queryDataIds(

3131 ["visit", "tract"],

3132 instrument="Cam1",

3133 skymap="SkyMap1",

3134 where="visit=1 AND detector=1 AND tract=0 AND patch=4",

3135 )

3136 self.assertFalse(set(query4))

3137 self.assertTrue(query4.any(execute=False, exact=False))

3138 self.assertTrue(query4.any(execute=True, exact=False))

3139 self.assertFalse(query4.any(execute=True, exact=True))

3140 self.assertGreaterEqual(query4.count(exact=False), 1)

3141 self.assertEqual(query4.count(exact=True, discard=True), 0)

3142 # This query should yield results from one dataset type but not the

3143 # other, which is not registered.

3144 query5 = registry.queryDatasets(["bias", "nonexistent"], collections=["biases"])

3145 self.assertTrue(set(query5))

3146 self.assertTrue(query5.any(execute=False, exact=False))

3147 self.assertTrue(query5.any(execute=True, exact=False))

3148 self.assertTrue(query5.any(execute=True, exact=True))

3149 self.assertGreaterEqual(query5.count(exact=False), 1)

3150 self.assertGreaterEqual(query5.count(exact=True), 1)

3151 # This query applies a selection that yields no results, fully in the

3152 # database. Explaining why it fails involves traversing the relation

3153 # tree and running a LIMIT 1 query at each level that has the potential

3154 # to remove rows.

3155 query6 = registry.queryDimensionRecords(

3156 "detector", where="detector.purpose = 'no-purpose'", instrument="Cam1"

3157 )

3158 self.assertEqual(query6.count(exact=True), 0)

3159 self.assertFalse(query6.any())

3160

3161 def testQueryDataIdsExpressionError(self):

3162 """Test error checking of 'where' expressions in queryDataIds."""

3163 butler = self.make_butler()

3164 registry = butler.registry

3165 self.load_data(butler, "base.yaml")

3166 bind = {"time": astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")}

3167 # The diagnostics raised are slightly different between the old query

3168 # system (ValueError, first error string) and the new query system

3169 # (InvalidQueryError, second error string).

3170 with self.assertRaisesRegex(

3171 (LookupError, InvalidQueryError),

3172 r"(No dimension element with name 'foo' in 'foo\.bar'\.)|(Unrecognized identifier 'foo.bar')",

3173 ):

3174 list(registry.queryDataIds(["detector"], where="foo.bar = 12"))

3175 with self.assertRaisesRegex(

3176 (LookupError, InvalidQueryError),

3177 "(Dimension element name cannot be inferred in this context.)"

3178 "|(Unrecognized identifier 'timespan')",

3179 ):

3180 list(registry.queryDataIds(["detector"], where="timespan.end < :time", bind=bind))

3181

3182 def testQueryDataIdsOrderBy(self):

3183 """Test order_by and limit on result returned by queryDataIds()."""

3184 butler = self.make_butler()

3185 registry = butler.registry

3186 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3187

3188 def do_query(dimensions=("visit", "tract"), datasets=None, collections=None):

3189 return registry.queryDataIds(

3190 dimensions, datasets=datasets, collections=collections, instrument="Cam1", skymap="SkyMap1"

3191 )

3192

3193 Test = namedtuple(

3194 "testQueryDataIdsOrderByTest",

3195 ("order_by", "keys", "result", "limit", "datasets", "collections"),

3196 defaults=(None, None, None),

3197 )

3198

3199 test_data = [

3200 Test("tract,visit", "tract,visit", ((0, 1), (0, 2), (1, 2))),

3201 Test("-tract,visit", "tract,visit", ((1, 2), (0, 1), (0, 2))),

3202 Test("tract,-visit", "tract,visit", ((0, 2), (0, 1), (1, 2))),

3203 Test("-tract,-visit", "tract,visit", ((1, 2), (0, 2), (0, 1))),

3204 Test("tract.id,visit.id", "tract,visit", ((0, 1),), limit=(1,)),

3205 Test("-tract,-visit", "tract,visit", ((1, 2),), limit=(1,)),

3206 Test("tract,visit.exposure_time", "tract,visit", ((0, 2), (0, 1), (1, 2))),

3207 Test("-tract,-visit.exposure_time", "tract,visit", ((1, 2), (0, 1), (0, 2))),

3208 Test("tract,-exposure_time", "tract,visit", ((0, 1), (0, 2), (1, 2))),

3209 Test("tract,visit.name", "tract,visit", ((0, 1), (0, 2), (1, 2))),

3210 Test(

3211 "tract,-visit.timespan.begin,visit.timespan.end",

3212 "tract,visit",

3213 ((0, 2), (0, 1), (1, 2)),

3214 ),

3215 Test("visit.day_obs,exposure.day_obs", "visit,exposure", ()),

3216 Test("visit.timespan.begin,-exposure.timespan.begin", "visit,exposure", ()),

3217 Test(

3218 "tract,detector",

3219 "tract,detector",

3220 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),

3221 datasets="flat",

3222 collections="imported_r",

3223 ),

3224 Test(

3225 "tract,detector.full_name",

3226 "tract,detector",

3227 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),

3228 datasets="flat",

3229 collections="imported_r",

3230 ),

3231 Test(

3232 "tract,detector.raft,detector.name_in_raft",

3233 "tract,detector",

3234 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),

3235 datasets="flat",

3236 collections="imported_r",

3237 ),

3238 ]

3239

3240 for test in test_data:

3241 with self.subTest(test=test):

3242 order_by = test.order_by.split(",")

3243 keys = test.keys.split(",")

3244 query = do_query(keys, test.datasets, test.collections).order_by(*order_by)

3245 if test.limit is not None:

3246 query = query.limit(*test.limit)

3247 dataIds = tuple(tuple(dataId[k] for k in keys) for dataId in query)

3248 self.assertEqual(dataIds, test.result)

3249

3250 # and materialize

3251 query = do_query(keys).order_by(*order_by)

3252 if test.limit is not None:

3253 query = query.limit(*test.limit)

3254

3255 # Test exceptions for errors in a name.

3256 # Many of these raise slightly different diagnostics in the old query

3257 # system (ValueError, first error string) than the new query system

3258 # (InvalidQueryError, second error string).

3259 for order_by in ("", "-"):

3260 with self.assertRaisesRegex((ValueError, InvalidQueryError), "Empty dimension name in ORDER BY"):

3261 list(do_query().order_by(order_by))

3262

3263 for order_by in ("undimension.name", "-undimension.name"):

3264 with self.assertRaisesRegex(

3265 (ValueError, InvalidQueryError),

3266 "(Unknown dimension element 'undimension')|(Unrecognized identifier 'undimension.name')",

3267 ):

3268 list(do_query().order_by(order_by))

3269

3270 for order_by in ("attract", "-attract"):

3271 with self.assertRaisesRegex(

3272 (ValueError, InvalidQueryError),

3273 "(Metadata 'attract' cannot be found in any dimension)|(Unrecognized identifier 'attract')",

3274 ):

3275 list(do_query().order_by(order_by))

3276

3277 with self.assertRaisesRegex(

3278 (ValueError, InvalidQueryError),

3279 "(Metadata 'exposure_time' exists in more than one dimension)"

3280 "|(Ambiguous identifier 'exposure_time' matches multiple fields)",

3281 ):

3282 list(do_query(("exposure", "visit")).order_by("exposure_time"))

3283

3284 with self.assertRaisesRegex(

3285 (ValueError, InvalidQueryError),

3286 r"(Timespan exists in more than one dimension element $day_obs, exposure, visit$; "

3287 r"qualify timespan with specific dimension name\.)|"

3288 r"(Ambiguous identifier 'timespan' matches multiple fields)",

3289 ):

3290 list(do_query(("exposure", "visit")).order_by("timespan.begin"))

3291

3292 with self.assertRaisesRegex(

3293 (ValueError, InvalidQueryError),

3294 "(Cannot find any temporal dimension element for 'timespan.begin')"

3295 "|(Unrecognized identifier 'timespan')",

3296 ):

3297 list(do_query("tract").order_by("timespan.begin"))

3298

3299 with self.assertRaisesRegex(

3300 (ValueError, InvalidQueryError),

3301 "(Cannot use 'timespan.begin' with non-temporal element)"

3302 "|(Unrecognized field 'timespan' for tract)",

3303 ):

3304 list(do_query("tract").order_by("tract.timespan.begin"))

3305

3306 with self.assertRaisesRegex(

3307 (ValueError, InvalidQueryError),

3308 "(Field 'name' does not exist in 'tract')|(Unrecognized field 'name' for tract.)",

3309 ):

3310 list(do_query("tract").order_by("tract.name"))

3311

3312 with self.assertRaisesRegex(

3313 (ValueError, InvalidQueryError),

3314 r"(Unknown dimension element 'timestamp'; perhaps you meant 'timespan.begin'\?)"

3315 r"|(Unrecognized identifier 'timestamp.begin')",

3316 ):

3317 list(do_query("visit").order_by("timestamp.begin"))

3318

3319 def testQueryDataIdsGovernorExceptions(self):

3320 """Test exceptions raised by queryDataIds() for incorrect governors."""

3321 butler = self.make_butler()

3322 registry = butler.registry

3323 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3324

3325 def do_query(dimensions, dataId=None, where="", bind=None, **kwargs):

3326 return registry.queryDataIds(dimensions, dataId=dataId, where=where, bind=bind, **kwargs)

3327

3328 Test = namedtuple(

3329 "testQueryDataIdExceptionsTest",

3330 ("dimensions", "dataId", "where", "bind", "kwargs", "exception", "count"),

3331 defaults=(None, None, None, {}, None, 0),

3332 )

3333

3334 test_data = (

3335 Test("tract,visit", count=3),

3336 Test("tract,visit", kwargs={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),

3337 Test("tract,visit", kwargs={"instrument": "Cam2", "skymap": "SkyMap1"}, count=0),

3338 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),

3339 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap2"}, count=0),

3340 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap1'", count=3),

3341 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap5'", count=0),

3342 Test(

3343 "tract,visit",

3344 where="instrument=:cam AND skymap=:map",

3345 bind={"cam": "Cam1", "map": "SkyMap1"},

3346 count=3,

3347 ),

3348 Test(

3349 "tract,visit",

3350 where="instrument=:cam AND skymap=:map",

3351 bind={"cam": "Cam", "map": "SkyMap"},

3352 count=0,

3353 ),

3354 )

3355

3356 for test in test_data:

3357 print(test)

3358 dimensions = test.dimensions.split(",")

3359 if test.exception:

3360 with self.assertRaises(test.exception):

3361 with ExitStack() as stack:

3362 if test.exception == DataIdValueError:

3363 stack.enter_context(self.assertWarns(FutureWarning))

3364 do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs).count()

3365 else:

3366 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)

3367 print(list(query))

3368 self.assertEqual(query.count(discard=True), test.count)

3369

3370 # and materialize

3371 if test.exception:

3372 with self.assertRaises(test.exception):

3373 with ExitStack() as stack:

3374 if test.exception == DataIdValueError:

3375 stack.enter_context(self.assertWarns(FutureWarning))

3376 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)

3377 with query.materialize() as materialized:

3378 materialized.count(discard=True)

3379 else:

3380 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)

3381 with query.materialize() as materialized:

3382 self.assertEqual(materialized.count(discard=True), test.count)

3383

3384 def testQueryDimensionRecordsOrderBy(self):

3385 """Test order_by and limit on result returned by

3386 queryDimensionRecords().

3387 """

3388 butler = self.make_butler()

3389 registry = butler.registry

3390 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3391

3392 def do_query(element, datasets=None, collections=None):

3393 return registry.queryDimensionRecords(

3394 element, instrument="Cam1", datasets=datasets, collections=collections

3395 )

3396

3397 query = do_query("detector")

3398 self.assertEqual(len(list(query)), 4)

3399

3400 Test = namedtuple(

3401 "testQueryDataIdsOrderByTest",

3402 ("element", "order_by", "result", "limit", "datasets", "collections"),

3403 defaults=(None, None, None),

3404 )

3405

3406 test_data = [

3407 Test("detector", "detector", (1, 2, 3, 4)),

3408 Test("detector", "-detector", (4, 3, 2, 1)),

3409 Test("detector", "raft,-name_in_raft", (2, 1, 4, 3)),

3410 Test("detector", "-detector.purpose", (4,), limit=(1,)),

3411 Test("visit", "visit", (1, 2)),

3412 Test("visit", "-visit.id", (2, 1)),

3413 Test("visit", "zenith_angle", (1, 2)),

3414 Test("visit", "-visit.name", (2, 1)),

3415 Test("visit", "day_obs,-visit.timespan.begin", (2, 1)),

3416 ]

3417

3418 def do_test(test: Test):

3419 order_by = test.order_by.split(",")

3420 query = do_query(test.element).order_by(*order_by)

3421 if test.limit is not None:

3422 query = query.limit(*test.limit)

3423 dataIds = tuple(rec.id for rec in query)

3424 self.assertEqual(dataIds, test.result)

3425

3426 for test in test_data:

3427 do_test(test)

3428

3429 # errors in a name

3430 for order_by in ("", "-"):

3431 with self.assertRaisesRegex(

3432 (ValueError, InvalidQueryError),

3433 "(Empty dimension name in ORDER BY)|(Unrecognized identifier)",

3434 ):

3435 list(do_query("detector").order_by(order_by))

3436

3437 for order_by in ("undimension.name", "-undimension.name"):

3438 with self.assertRaisesRegex(

3439 (ValueError, InvalidQueryError),

3440 "(Element name mismatch: 'undimension')|(Unrecognized identifier)",

3441 ):

3442 list(do_query("detector").order_by(order_by))

3443

3444 for order_by in ("attract", "-attract"):

3445 with self.assertRaisesRegex(

3446 (ValueError, InvalidQueryError),

3447 "(Field 'attract' does not exist in 'detector'.)|(Unrecognized identifier)",

3448 ):

3449 list(do_query("detector").order_by(order_by))

3450

3451 for order_by in ("timestamp.begin", "-timestamp.begin"):

3452 with self.assertRaisesRegex(

3453 (ValueError, InvalidQueryError),

3454 r"(Element name mismatch: 'timestamp' instead of 'visit'; "

3455 r"perhaps you meant 'timespan.begin'\?)"

3456 r"|(Unrecognized identifier)",

3457 ):

3458 list(do_query("visit").order_by(order_by))

3459

3460 def testQueryDimensionRecordsExceptions(self):

3461 """Test exceptions raised by queryDimensionRecords()."""

3462 butler = self.make_butler()

3463 registry = butler.registry

3464 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3465

3466 result = registry.queryDimensionRecords("detector")

3467 self.assertEqual(result.count(), 4)

3468 result = registry.queryDimensionRecords("detector", instrument="Cam1")

3469 self.assertEqual(result.count(), 4)

3470 result = registry.queryDimensionRecords("detector", dataId={"instrument": "Cam1"})

3471 self.assertEqual(result.count(), 4)

3472

3473 # Test that values specified in kwargs override those specified in

3474 # dataId.

3475 result = registry.queryDimensionRecords(

3476 "detector", dataId={"instrument": "NotCam1"}, instrument="Cam1"

3477 )

3478 self.assertEqual(result.count(), 4)

3479

3480 result = registry.queryDimensionRecords("detector", where="instrument='Cam1'")

3481 self.assertEqual(result.count(), 4)

3482 result = registry.queryDimensionRecords("detector", where="instrument=:instr", bind={"instr": "Cam1"})

3483 self.assertTrue(result.any())

3484 self.assertEqual(result.count(), 4)

3485

3486 def testDatasetConstrainedDimensionRecordQueries(self):

3487 """Test that queryDimensionRecords works even when given a dataset

3488 constraint whose dimensions extend beyond the requested dimension

3489 element's.

3490 """

3491 butler = self.make_butler()

3492 registry = butler.registry

3493 self.load_data(butler, "base.yaml", "datasets.yaml")

3494 # Query for physical_filter dimension records, using a dataset that

3495 # has both physical_filter and dataset dimensions.

3496 records = registry.queryDimensionRecords(

3497 "physical_filter",

3498 datasets=["flat"],

3499 collections="imported_r",

3500 )

3501 self.assertEqual({record.name for record in records}, {"Cam1-R1", "Cam1-R2"})

3502 # Trying to constrain by all dataset types is an error.

3503 with self.assertRaises(TypeError):

3504 list(registry.queryDimensionRecords("physical_filter", datasets=..., collections="imported_r"))

3505

3506 def testSkyPixDatasetQueries(self):

3507 """Test that we can build queries involving skypix dimensions as long

3508 as a dataset type that uses those dimensions is included.

3509 """

3510 butler = self.make_butler()

3511 registry = butler.registry

3512 self.load_data(butler, "base.yaml")

3513 dataset_type = DatasetType(

3514 "a", dimensions=["htm7", "instrument"], universe=registry.dimensions, storageClass="int"

3515 )

3516 registry.registerDatasetType(dataset_type)

3517 run = "r"

3518 registry.registerRun(run)

3519 # First try queries where there are no datasets; the concern is whether

3520 # we can even build and execute these queries without raising, even

3521 # when "doomed" query shortcuts are in play.

3522 self.assertFalse(

3523 list(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run))

3524 )

3525 self.assertFalse(list(registry.queryDatasets(dataset_type, collections=run)))

3526 # Now add a dataset and see that we can get it back.

3527 htm7 = registry.dimensions.skypix["htm"][7].pixelization

3528 data_id = registry.expandDataId(instrument="Cam1", htm7=htm7.universe()[0][0])

3529 (ref,) = registry.insertDatasets(dataset_type, [data_id], run=run)

3530 self.assertEqual(

3531 set(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run)),

3532 {data_id},

3533 )

3534 self.assertEqual(set(registry.queryDatasets(dataset_type, collections=run)), {ref})

3535

3536 def testDatasetIdFactory(self):

3537 """Simple test for DatasetIdFactory, mostly to catch potential changes

3538 in its API.

3539 """

3540 butler = self.make_butler()

3541 registry = butler.registry

3542 factory = DatasetIdFactory()

3543 dataset_type = DatasetType(

3544 "datasetType",

3545 dimensions=["detector", "instrument"],

3546 universe=registry.dimensions,

3547 storageClass="int",

3548 )

3549 run = "run"

3550 data_id = DataCoordinate.standardize(

3551 instrument="Cam1", detector=1, dimensions=dataset_type.dimensions

3552 )

3553

3554 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.UNIQUE)

3555 self.assertIsInstance(datasetId, uuid.UUID)

3556 self.assertEqual(datasetId.version, 7)

3557

3558 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE)

3559 self.assertIsInstance(datasetId, uuid.UUID)

3560 self.assertEqual(datasetId.version, 5)

3561

3562 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE_RUN)

3563 self.assertIsInstance(datasetId, uuid.UUID)

3564 self.assertEqual(datasetId.version, 5)

3565

3566 def testExposureQueries(self):

3567 """Test query methods using arguments sourced from the exposure log

3568 service.

3569

3570 The most complete test dataset currently available to daf_butler tests

3571 is ci_hsc-subset.yaml export , but that does not have 'exposure'

3572 dimension records. So in this test we need to translate queries that

3573 originally used the exposure dimension to use the (very similar) visit

3574 dimension instead.

3575 """

3576 butler = self.make_butler()

3577 registry = butler.registry

3578 self.load_data(butler, "ci_hsc-subset.yaml")

3579 self.assertEqual(

3580 [

3581 record.id

3582 for record in registry.queryDimensionRecords("visit", instrument="HSC")

3583 .order_by("visit")

3584 .limit(5)

3585 ],

3586 [903334, 903336, 903338, 903342, 903344],

3587 )

3588 self.assertEqual(

3589 [

3590 data_id["visit"]

3591 for data_id in registry.queryDataIds(["visit"], instrument="HSC").order_by("visit").limit(5)

3592 ],

3593 [903334, 903336, 903338, 903342, 903344],

3594 )

3595 self.assertEqual(

3596 [

3597 record.id

3598 for record in registry.queryDimensionRecords("detector", instrument="HSC")

3599 .order_by("full_name")

3600 .limit(5)

3601 ],

3602 [25, 24, 23, 22, 18],

3603 )

3604 self.assertEqual(

3605 [

3606 data_id["detector"]

3607 for data_id in registry.queryDataIds(["detector"], instrument="HSC")

3608 .order_by("full_name")

3609 .limit(5)

3610 ],

3611 [25, 24, 23, 22, 18],

3612 )

3613

3614 def test_long_query_names(self) -> None:

3615 """Test that queries involving very long names are handled correctly.

3616

3617 This is especially important for PostgreSQL, which truncates symbols

3618 longer than 64 chars, but it's worth testing for all DBs.

3619 """

3620 butler = self.make_butler()

3621 registry = butler.registry

3622 name = "abcd" * 17

3623 registry.registerDatasetType(

3624 DatasetType(

3625 name,

3626 dimensions=(),

3627 storageClass="Exposure",

3628 universe=registry.dimensions,

3629 )

3630 )

3631 # Need to search more than one collection actually containing a

3632 # matching dataset to avoid optimizations that sidestep bugs due to

3633 # truncation by making findFirst=True a no-op.

3634 run1 = "run1"

3635 registry.registerRun(run1)

3636 run2 = "run2"

3637 registry.registerRun(run2)

3638 (ref1,) = registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run1)

3639 registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run2)

3640 self.assertEqual(

3641 set(registry.queryDatasets(name, collections=[run1, run2], findFirst=True)),

3642 {ref1},

3643 )

3644

3645 def test_skypix_constraint_queries(self) -> None:

3646 """Test queries spatially constrained by a skypix data ID."""

3647 butler = self.make_butler()

3648 registry = butler.registry

3649 self.load_data(butler, "base.yaml", "spatial.yaml")

3650 patch_regions = {

3651 (data_id["tract"], data_id["patch"]): data_id.region

3652 for data_id in registry.queryDataIds(["patch"]).expanded()

3653 }

3654 skypix_dimension: SkyPixDimension = registry.dimensions["htm11"]

3655 # This check ensures the test doesn't become trivial due to a config

3656 # change; if it does, just pick a different HTML level.

3657 self.assertNotEqual(skypix_dimension, registry.dimensions.commonSkyPix)

3658 # Gather all skypix IDs that definitely overlap at least one of these

3659 # patches.

3660 relevant_skypix_ids = lsst.sphgeom.RangeSet()

3661 for patch_region in patch_regions.values():

3662 relevant_skypix_ids |= skypix_dimension.pixelization.interior(patch_region)

3663 # Look for a "nontrivial" skypix_id that overlaps at least one patch

3664 # and does not overlap at least one other patch.

3665 for skypix_id in itertools.chain.from_iterable(

3666 range(begin, end) for begin, end in relevant_skypix_ids

3667 ):

3668 skypix_region = skypix_dimension.pixelization.pixel(skypix_id)

3669 overlapping_patches = {

3670 patch_key

3671 for patch_key, patch_region in patch_regions.items()

3672 if not patch_region.isDisjointFrom(skypix_region)

3673 }

3674 if overlapping_patches and overlapping_patches != patch_regions.keys():

3675 break

3676 else:

3677 raise RuntimeError("Could not find usable skypix ID for this dimension configuration.")

3678 # Test that a three-way join that includes the common skypix system in

3679 # the dimensions doesn't generate redundant join terms in the query.

3680 with self.assertRaises(InvalidQueryError):

3681 set(

3682 registry.queryDataIds(

3683 ["tract", "visit", "htm7"], skymap="SkyMap1", instrument="Cam1"

3684 ).expanded()

3685 )

3686

3687 def test_spatial_constraint_queries(self) -> None:

3688 """Test queries in which one spatial dimension in the constraint (data

3689 ID or ``where`` string) constrains a different spatial dimension in the

3690 query result columns.

3691 """

3692 butler = self.make_butler()

3693 registry = butler.registry

3694 self.load_data(butler, "base.yaml", "spatial.yaml")

3695 patch_regions = {

3696 (data_id["tract"], data_id["patch"]): data_id.region

3697 for data_id in registry.queryDataIds(["patch"]).expanded()

3698 }

3699 observation_regions = {

3700 (data_id["visit"], data_id["detector"]): data_id.region

3701 for data_id in registry.queryDataIds(["visit", "detector"]).expanded()

3702 }

3703 all_combos = {

3704 (patch_key, observation_key)

3705 for patch_key, observation_key in itertools.product(patch_regions, observation_regions)

3706 }

3707 overlapping_combos = {

3708 (patch_key, observation_key)

3709 for patch_key, observation_key in all_combos

3710 if not patch_regions[patch_key].isDisjointFrom(observation_regions[observation_key])

3711 }

3712 # Check a direct spatial join with no constraint first.

3713 self.assertEqual(

3714 {

3715 ((data_id["tract"], data_id["patch"]), (data_id["visit"], data_id["detector"]))

3716 for data_id in registry.queryDataIds(["patch", "visit", "detector"])

3717 },

3718 overlapping_combos,

3719 )

3720 overlaps_by_patch: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)

3721 overlaps_by_observation: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)

3722 for patch_key, observation_key in overlapping_combos:

3723 overlaps_by_patch[patch_key].add(observation_key)

3724 overlaps_by_observation[observation_key].add(patch_key)

3725 # Find patches and observations that overlap at least one of the other

3726 # but not all of the other.

3727 nontrivial_patch = next(

3728 iter(

3729 patch_key

3730 for patch_key, observation_keys in overlaps_by_patch.items()

3731 if observation_keys and observation_keys != observation_regions.keys()

3732 )

3733 )

3734 nontrivial_observation = next(

3735 iter(

3736 observation_key

3737 for observation_key, patch_keys in overlaps_by_observation.items()

3738 if patch_keys and patch_keys != patch_regions.keys()

3739 )

3740 )

3741 # Use the nontrivial patches and observations as constraints on the

3742 # other dimensions in various ways, first via a 'where' expression.

3743 # It's better in general to us 'bind' instead of f-strings, but these

3744 # all integers so there are no quoting concerns.

3745 self.assertEqual(

3746 {

3747 (data_id["visit"], data_id["detector"])

3748 for data_id in registry.queryDataIds(

3749 ["visit", "detector"],

3750 where=f"tract={nontrivial_patch[0]} AND patch={nontrivial_patch[1]}",

3751 skymap="SkyMap1",

3752 )

3753 },

3754 overlaps_by_patch[nontrivial_patch],

3755 )

3756 self.assertEqual(

3757 {

3758 (data_id["tract"], data_id["patch"])

3759 for data_id in registry.queryDataIds(

3760 ["patch"],

3761 where=f"visit={nontrivial_observation[0]} AND detector={nontrivial_observation[1]}",

3762 instrument="Cam1",

3763 )

3764 },

3765 overlaps_by_observation[nontrivial_observation],

3766 )

3767 # and then via the dataId argument.

3768 self.assertEqual(

3769 {

3770 (data_id["visit"], data_id["detector"])

3771 for data_id in registry.queryDataIds(

3772 ["visit", "detector"],

3773 dataId={

3774 "tract": nontrivial_patch[0],

3775 "patch": nontrivial_patch[1],

3776 },

3777 skymap="SkyMap1",

3778 )

3779 },

3780 overlaps_by_patch[nontrivial_patch],

3781 )

3782 self.assertEqual(

3783 {

3784 (data_id["tract"], data_id["patch"])

3785 for data_id in registry.queryDataIds(

3786 ["patch"],

3787 dataId={

3788 "visit": nontrivial_observation[0],

3789 "detector": nontrivial_observation[1],

3790 },

3791 instrument="Cam1",

3792 )

3793 },

3794 overlaps_by_observation[nontrivial_observation],

3795 )

3796

3797 def test_query_empty_collections(self) -> None:

3798 """Test for registry query methods with empty collections. The methods

3799 should return empty result set (or None when applicable) and provide

3800 "doomed" diagnostics.

3801 """

3802 butler = self.make_butler()

3803 registry = butler.registry

3804 self.load_data(butler, "base.yaml", "datasets.yaml")

3805

3806 # Tests for registry.findDataset()

3807 with self.assertRaises(NoDefaultCollectionError):

3808 registry.findDataset("bias", instrument="Cam1", detector=1)

3809 self.assertIsNotNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=...))

3810 self.assertIsNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=[]))

3811

3812 # Tests for registry.queryDatasets()

3813 with self.assertRaises(NoDefaultCollectionError):

3814 registry.queryDatasets("bias")

3815 self.assertTrue(list(registry.queryDatasets("bias", collections=...)))

3816

3817 result = registry.queryDatasets("bias", collections=[])

3818 self.assertEqual(len(list(result)), 0)

3819 messages = list(result.explain_no_results())

3820 self.assertTrue(messages)

3821 self.assertTrue(any("because collection list is empty" in message for message in messages))

3822

3823 # Tests for registry.queryDataIds()

3824 with self.assertRaises(NoDefaultCollectionError):

3825 registry.queryDataIds("detector", datasets="bias")

3826 self.assertTrue(list(registry.queryDataIds("detector", datasets="bias", collections=...)))

3827

3828 result = registry.queryDataIds("detector", datasets="bias", collections=[])

3829 self.assertEqual(len(list(result)), 0)

3830 messages = list(result.explain_no_results())

3831 self.assertTrue(messages)

3832 self.assertTrue(any("because collection list is empty" in message for message in messages))

3833

3834 # Tests for registry.queryDimensionRecords()

3835 with self.assertRaises(NoDefaultCollectionError):

3836 registry.queryDimensionRecords("detector", datasets="bias")

3837 self.assertTrue(list(registry.queryDimensionRecords("detector", datasets="bias", collections=...)))

3838

3839 result = registry.queryDimensionRecords("detector", datasets="bias", collections=[])

3840 self.assertEqual(len(list(result)), 0)

3841 messages = list(result.explain_no_results())

3842 self.assertTrue(messages)

3843 self.assertTrue(any("because collection list is empty" in message for message in messages))

3844

3845 def test_dataset_followup_spatial_joins(self) -> None:

3846 """Test queryDataIds(...).findRelatedDatasets(...) where a spatial join

3847 is involved.

3848 """

3849 butler = self.make_butler()

3850 registry = butler.registry

3851 self.load_data(butler, "base.yaml", "spatial.yaml")

3852 pvi_dataset_type = DatasetType(

3853 "pvi", {"visit", "detector"}, storageClass="StructuredDataDict", universe=registry.dimensions

3854 )

3855 registry.registerDatasetType(pvi_dataset_type)

3856 collection = "datasets"

3857 registry.registerRun(collection)

3858 (pvi1,) = registry.insertDatasets(

3859 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 1}], run=collection

3860 )

3861 (pvi2,) = registry.insertDatasets(

3862 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 2}], run=collection

3863 )

3864 (pvi3,) = registry.insertDatasets(

3865 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 3}], run=collection

3866 )

3867 self.assertEqual(

3868 set(

3869 registry.queryDataIds(["patch"], skymap="SkyMap1", tract=0)

3870 .expanded()

3871 .findRelatedDatasets("pvi", [collection])

3872 ),

3873 {

3874 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi1),

3875 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi2),

3876 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=1), pvi2),

3877 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi1),

3878 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi2),

3879 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi3),

3880 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=3), pvi2),

3881 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=4), pvi3),

3882 },

3883 )

3884

3885 def test_expanded_data_id_queries(self) -> None:

3886 """Tests for basic functionality of expanded() on queryDataIds and

3887 queryDatasets.

3888 """

3889 butler = self.make_butler()

3890 registry = butler.registry

3891 self.load_data(butler, "base.yaml", "spatial.yaml", "datasets.yaml")

3892

3893 result_obj = (

3894 registry.queryDataIds(["visit"], where="instrument = 'Cam1' and (visit.id = 1 or visit.id = 2)")

3895 .expanded()

3896 .order_by("visit.id")

3897 )

3898 self.assertTrue(result_obj.hasRecords())

3899 visits = list(result_obj)

3900 self.assertEqual(len(visits), 2)

3901

3902 self.assertEqual(visits[0]["visit"], 1)

3903 self.assertEqual(visits[1]["visit"], 2)

3904 self.assertEqual(visits[0].records["visit"].exposure_time, 60.0)

3905 self.assertEqual(visits[1].records["visit"].exposure_time, 45.0)

3906 # physical_filter is a "cacheable" dimension, so its records are loaded

3907 # from local cache rather than being part of the DB rows.

3908 self.assertEqual(visits[0].records["physical_filter"].band, "g")

3909 self.assertEqual(visits[1].records["physical_filter"].band, "r")

3910

3911 # Make sure that we can fetch nulls in dimension records

3912 registry.insertDimensionData(

3913 "detector",

3914 {

3915 "instrument": "Cam1",

3916 "id": 5,

3917 "raft": "Z",

3918 "name_in_raft": "z",

3919 "full_name": "Zz",

3920 "purpose": None,

3921 },

3922 )

3923 detectors = list(

3924 registry.queryDataIds("detector", dataId={"instrument": "Cam1", "detector": 5}).expanded()

3925 )

3926 self.assertIsNone(detectors[0].records["detector"].purpose)

3927

3928 datasets = list(

3929 registry.queryDatasets(

3930 "flat", collections="imported_g", where="instrument = 'Cam1' and detector <= 3"

3931 ).expanded()

3932 )

3933 datasets.sort(key=lambda ref: ref.dataId["detector"])

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

3935 self.assertEqual(datasets[0].id, uuid.UUID("60c8a65c-7290-4c38-b1de-e3b1cdcf872d"))

3936 self.assertEqual(datasets[1].id, uuid.UUID("84239e7f-c41f-46d5-97b9-a27976b98ceb"))

3937 # All of the dimensions for flat are "cached" dimensions.

3938 self.assertEqual(datasets[0].dataId.records["detector"].full_name, "Ab")

3939 self.assertEqual(datasets[1].dataId.records["detector"].full_name, "Ba")

3940 self.assertEqual(datasets[0].dataId.records["instrument"].visit_system, 1)

3941

3942 # None of the datasets in the test data include any uncached

3943 # dimensions, so we have to set one up.

3944 registry.registerDatasetType(DatasetType("test", ["visit"], "int", universe=registry.dimensions))

3945 registry.insertDatasets("test", [{"instrument": "Cam1", "visit": 1}], run="imported_g")

3946 ref = list(registry.queryDatasets("test", collections="imported_g").expanded())[0]

3947 self.assertEqual(ref.dataId.records["visit"].zenith_angle, 5.0)

3948 self.assertEqual(ref.dataId.records["physical_filter"].band, "g")

3949 self.assertEqual(

3950 ref.dataId.timespan,

3951 Timespan(

3952 begin=astropy.time.Time("2021-09-09 03:00:00.000000000", scale="tai"),

3953 end=astropy.time.Time("2021-09-09 03:01:00.000000000", scale="tai"),

3954 ),

3955 )

3956

3957 def test_collection_summary(self) -> None:

3958 """Test for collection summary methods."""

3959 butler = self.make_butler()

3960 registry = butler.registry

3961 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")

3962

3963 # Add one more dataset type, just for its existence to trigger a bug

3964 # in `associate` (DM-44311).

3965 test_dataset_type = DatasetType("test", ["tract", "patch"], "int", universe=registry.dimensions)

3966 registry.registerDatasetType(test_dataset_type)

3967

3968 # Check for what has been imported.

3969 summary = registry.getCollectionSummary("imported_g")

3970 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})

3971 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})

3972

3973 # Make a tagged collection and associate some datasets.

3974 tagged_coll = "tagged"

3975 registry.registerCollection(tagged_coll, CollectionType.TAGGED)

3976 refsets = registry.queryDatasets(..., collections=["imported_g"]).byParentDatasetType()

3977 for refs in refsets:

3978 registry.associate(tagged_coll, refs)

3979

3980 # Summary has to have the same dataset types.

3981 summary = registry.getCollectionSummary(tagged_coll)

3982 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})

3983 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})

3984

3985 # Remove all datasets from the tagged collection.

3986 refs = list(registry.queryDatasets(..., collections=[tagged_coll]))

3987 registry.disassociate(tagged_coll, refs)

3988

3989 # Summaries should not have changed.

3990 summary = registry.getCollectionSummary(tagged_coll)

3991 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})

3992 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})

3993

3994 # Cleanup summaries.

3995 registry.refresh_collection_summaries()

3996 summary = registry.getCollectionSummary(tagged_coll)

3997 self.assertFalse(summary.dataset_types.names)

3998 # We do not clean governor summaries yet, but because how the query is

3999 # run, it returns empty governors when collection is missing from

4000 # summaries.

4001 self.assertFalse(summary.governors)

4002

4003 # Add dataset with different governor, this is to test that governors

4004 # are not actually cleaned.

4005 refs = registry.insertDatasets("test", [{"skymap": "SkyMap1", "tract": 0, "patch": 0}], "imported_g")

4006 registry.associate(tagged_coll, refs)

4007 summary = registry.getCollectionSummary(tagged_coll)

4008 self.assertEqual(summary.dataset_types.names, {"test"})

4009 # Note that instrument governor resurrects here, even though there are

4010 # no datasets left with that governor.

4011 self.assertEqual(summary.governors, {"instrument": {"Cam1"}, "skymap": {"SkyMap1"}})

4012

4013 def test_temp_table_config(self) -> None:

4014 config = self.makeRegistryConfig()

4015 config["temporary_tables"] = False

4016 self.assertEqual(config.areTemporaryTablesAllowed, False)

4017 butler = self.make_butler(config)

4018 if not isinstance(butler, DirectButler):

4019 raise unittest.SkipTest("Test only makes sense for registry with direct database connection.")

4020 self.assertEqual(butler._registry._db.supports_temporary_tables, False)

4021 with self.assertRaisesRegex(ReadOnlyDatabaseError, "temporary tables"):

4022 with butler._registry._db.temporary_table(...):

4023 pass

Coverage for python/lsst/daf/butler/registry/tests/_registry.py: 6%

1660 statements