Coverage for python/lsst/daf/butler/registry/tests/_registry.py: 6%
1679 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 01:36 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
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 ..queries import ParentDatasetQueryResults
90from ..sql_registry import SqlRegistry
92_T = TypeVar("_T")
95class RegistryTests(ABC):
96 """Generic tests for the `SqlRegistry` class that can be subclassed to
97 generate tests for different configurations.
98 """
100 collectionsManager: str | None = None
101 """Name of the collections manager class, if subclass provides value for
102 this member then it overrides name specified in default configuration
103 (`str`).
104 """
106 datasetsManager: str | dict[str, str] | None = None
107 """Name or configuration dictionary of the datasets manager class, if
108 subclass provides value for this member then it overrides name specified
109 in default configuration (`str` or `dict`).
110 """
112 supportsCollectionRegex: bool = False
113 """True if the registry class being tested supports regex searches for
114 collections."""
116 def makeRegistryConfig(self) -> RegistryConfig:
117 """Create RegistryConfig used to create a registry.
119 This method should be called by a subclass from `makeRegistry`.
120 Returned instance will be pre-configured based on the values of class
121 members, and default-configured for all other parameters. Subclasses
122 that need default configuration should just instantiate
123 `RegistryConfig` directly.
124 """
125 config = RegistryConfig()
126 if self.collectionsManager:
127 config["managers", "collections"] = self.collectionsManager
128 if self.datasetsManager:
129 config["managers", "datasets"] = self.datasetsManager
130 return config
132 @abstractmethod
133 def make_butler(self, registry_config: RegistryConfig | None = None) -> Butler:
134 """Return the butler to be tested.
136 Parameters
137 ----------
138 registry_config : `RegistryConfig`, optional
139 Registry configuration used when instantiating the Butler.
141 Returns
142 -------
143 butler : `~lsst.daf.butler.Butler`
144 The butler with a registry to be tested.
145 """
146 raise NotImplementedError()
148 def load_data(self, butler: Butler, *filenames: str) -> None:
149 """Load registry test data from
150 ``resource://lsst.daf.butler/tests/registry_data/<filename>``,
151 which should be a YAML import/export file.
153 Parameters
154 ----------
155 butler : `Butler`
156 The butler to load into.
157 *filenames : `str`
158 The names of the files to load.
159 """
160 for filename in filenames:
161 butler.import_(
162 filename=f"resource://lsst.daf.butler/tests/registry_data/{filename}", without_datastore=True
163 )
165 def checkQueryResults(self, results, expected):
166 """Check that a query results object contains expected values.
168 Parameters
169 ----------
170 results : `DataCoordinateQueryResults` or `DatasetQueryResults`
171 A lazy-evaluation query results object.
172 expected : `list`
173 A list of `DataCoordinate` o `DatasetRef` objects that should be
174 equal to results of the query, aside from ordering.
175 """
176 self.assertCountEqual(list(results), expected)
177 self.assertEqual(results.count(), len(expected))
178 if expected:
179 self.assertTrue(results.any())
180 else:
181 self.assertFalse(results.any())
183 def testOpaque(self):
184 """Tests for `SqlRegistry.registerOpaqueTable`,
185 `SqlRegistry.insertOpaqueData`, `SqlRegistry.fetchOpaqueData`, and
186 `SqlRegistry.deleteOpaqueData`.
187 """
188 butler = self.make_butler()
189 registry = butler._registry
190 table = "opaque_table_for_testing"
191 registry.registerOpaqueTable(
192 table,
193 spec=ddl.TableSpec(
194 fields=[
195 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, primaryKey=True),
196 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=16, nullable=False),
197 ddl.FieldSpec("count", dtype=sqlalchemy.SmallInteger, nullable=True),
198 ],
199 ),
200 )
201 rows = [
202 {"id": 1, "name": "one", "count": None},
203 {"id": 2, "name": "two", "count": 5},
204 {"id": 3, "name": "three", "count": 6},
205 ]
206 registry.insertOpaqueData(table, *rows)
207 self.assertCountEqual(rows, list(registry.fetchOpaqueData(table)))
208 self.assertEqual(rows[0:1], list(registry.fetchOpaqueData(table, id=1)))
209 self.assertEqual(rows[1:2], list(registry.fetchOpaqueData(table, name="two")))
210 self.assertEqual(rows[0:1], list(registry.fetchOpaqueData(table, id=(1, 3), name=("one", "two"))))
211 self.assertEqual(rows, list(registry.fetchOpaqueData(table, id=(1, 2, 3))))
212 # Test very long IN clause which exceeds sqlite limit on number of
213 # parameters. SQLite says the limit is 32k but it looks like it is
214 # much higher.
215 self.assertEqual(rows, list(registry.fetchOpaqueData(table, id=list(range(300_000)))))
216 # Two IN clauses, each longer than 1k batch size, first with
217 # duplicates, second has matching elements in different batches (after
218 # sorting).
219 self.assertEqual(
220 rows[0:2],
221 list(
222 registry.fetchOpaqueData(
223 table,
224 id=list(range(1000)) + list(range(100, 0, -1)),
225 name=["one"] + [f"q{i}" for i in range(2200)] + ["two"],
226 )
227 ),
228 )
229 self.assertEqual([], list(registry.fetchOpaqueData(table, id=1, name="two")))
230 registry.deleteOpaqueData(table, id=3)
231 self.assertCountEqual(rows[:2], list(registry.fetchOpaqueData(table)))
232 registry.deleteOpaqueData(table)
233 self.assertEqual([], list(registry.fetchOpaqueData(table)))
235 def testDatasetType(self):
236 """Tests for `SqlRegistry.registerDatasetType` and
237 `SqlRegistry.getDatasetType`.
238 """
239 butler = self.make_butler()
240 registry = butler.registry
241 # Check valid insert
242 datasetTypeName = "test"
243 storageClass = StorageClass("testDatasetType")
244 registry.storageClasses.registerStorageClass(storageClass)
245 dimensions = registry.dimensions.conform(("instrument", "visit"))
246 differentDimensions = registry.dimensions.conform(("instrument", "patch"))
247 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass)
248 # Inserting for the first time should return True
249 self.assertTrue(registry.registerDatasetType(inDatasetType))
250 outDatasetType1 = registry.getDatasetType(datasetTypeName)
251 self.assertEqual(outDatasetType1, inDatasetType)
253 # Re-inserting should work
254 self.assertFalse(registry.registerDatasetType(inDatasetType))
255 # Except when they are not identical
256 with self.assertRaises(ConflictingDefinitionError):
257 nonIdenticalDatasetType = DatasetType(datasetTypeName, differentDimensions, storageClass)
258 registry.registerDatasetType(nonIdenticalDatasetType)
260 # Template can be None
261 datasetTypeName = "testNoneTemplate"
262 storageClass = StorageClass("testDatasetType2")
263 registry.storageClasses.registerStorageClass(storageClass)
264 dimensions = registry.dimensions.conform(("instrument", "visit"))
265 inDatasetType = DatasetType(datasetTypeName, dimensions, storageClass)
266 registry.registerDatasetType(inDatasetType)
267 outDatasetType2 = registry.getDatasetType(datasetTypeName)
268 self.assertEqual(outDatasetType2, inDatasetType)
270 allTypes = set(registry.queryDatasetTypes())
271 self.assertEqual(allTypes, {outDatasetType1, outDatasetType2})
273 # Test some basic queryDatasetTypes functionality
274 missing: list[str] = []
275 types = registry.queryDatasetTypes(["te*", "notarealdatasettype"], missing=missing)
276 self.assertCountEqual([dt.name for dt in types], ["test", "testNoneTemplate"])
277 self.assertEqual(missing, ["notarealdatasettype"])
279 # Trying to register a dataset type with different universe version or
280 # namespace will raise.
281 wrong_universes = (DimensionUniverse(version=-1), DimensionUniverse(namespace="🔭"))
282 for universe in wrong_universes:
283 storageClass = StorageClass("testDatasetType")
284 dataset_type = DatasetType(
285 "wrong_universe", ("instrument", "visit"), storageClass, universe=universe
286 )
287 with self.assertRaisesRegex(ValueError, "Incompatible dimension universe versions"):
288 registry.registerDatasetType(dataset_type)
290 def testDatasetTypeCache(self):
291 """Test for dataset type cache update logic after a cache miss."""
292 butler1 = self.make_butler()
293 butler2 = butler1.clone()
294 self.load_data(butler1, "base.yaml")
296 # Trigger full cache load.
297 butler2.get_dataset_type("flat")
298 # Have an external process register a dataset type.
299 butler1.registry.registerDatasetType(
300 DatasetType("test_type", ["instrument"], "int", universe=butler1.dimensions)
301 )
302 # Try to read the new dataset type -- this is a cache miss that
303 # triggers fetch of a single dataset type.
304 dt = butler2.get_dataset_type("test_type")
305 self.assertEqual(dt.name, "test_type")
306 self.assertEqual(list(dt.dimensions.names), ["instrument"])
307 # Read it again -- this time it should pull from the cache.
308 dt = butler2.get_dataset_type("test_type")
309 self.assertEqual(dt.name, "test_type")
310 self.assertEqual(list(dt.dimensions.names), ["instrument"])
311 # Do a query that uses the dataset type's tags table.
312 self.assertEqual(
313 butler2.query_datasets("test_type", collections="*", find_first=False, explain=False), []
314 )
316 def testDimensions(self):
317 """Tests for `SqlRegistry.insertDimensionData`,
318 `SqlRegistry.syncDimensionData`, and `SqlRegistry.expandDataId`.
319 """
320 butler = self.make_butler()
321 registry = butler.registry
322 dimensionName = "instrument"
323 dimension = registry.dimensions[dimensionName]
324 dimensionValue = {
325 "name": "DummyCam",
326 "visit_max": 10,
327 "visit_system": 0,
328 "exposure_max": 10,
329 "detector_max": 2,
330 "class_name": "lsst.pipe.base.Instrument",
331 }
332 registry.insertDimensionData(dimensionName, dimensionValue)
333 # Inserting the same value twice should fail
334 with self.assertRaises(sqlalchemy.exc.IntegrityError):
335 registry.insertDimensionData(dimensionName, dimensionValue)
336 # expandDataId should retrieve the record we just inserted
337 self.assertEqual(
338 registry.expandDataId(instrument="DummyCam", dimensions=dimension.minimal_group)
339 .records[dimensionName]
340 .toDict(),
341 dimensionValue,
342 )
343 # expandDataId should raise if there is no record with the given ID.
344 with self.assertRaises(DataIdValueError):
345 registry.expandDataId({"instrument": "Unknown"}, dimensions=dimension.minimal_group)
346 # band doesn't have a table; insert should fail.
347 with self.assertRaises(TypeError):
348 registry.insertDimensionData("band", {"band": "i"})
349 dimensionName2 = "physical_filter"
350 dimension2 = registry.dimensions[dimensionName2]
351 dimensionValue2 = {"name": "DummyCam_i", "band": "i"}
352 # Missing required dependency ("instrument") should fail
353 with self.assertRaises(KeyError):
354 registry.insertDimensionData(dimensionName2, dimensionValue2)
355 # Adding required dependency should fix the failure
356 dimensionValue2["instrument"] = "DummyCam"
357 registry.insertDimensionData(dimensionName2, dimensionValue2)
358 # expandDataId should retrieve the record we just inserted.
359 self.assertEqual(
360 registry.expandDataId(
361 instrument="DummyCam", physical_filter="DummyCam_i", dimensions=dimension2.minimal_group
362 )
363 .records[dimensionName2]
364 .toDict(),
365 dimensionValue2,
366 )
367 # Use syncDimensionData to insert a new record successfully.
368 dimensionName3 = "detector"
369 dimensionValue3 = {
370 "instrument": "DummyCam",
371 "id": 1,
372 "full_name": "one",
373 "name_in_raft": "zero",
374 "purpose": "SCIENCE",
375 }
376 self.assertTrue(registry.syncDimensionData(dimensionName3, dimensionValue3))
377 # Sync that again. Note that one field ("raft") is NULL, and that
378 # should be okay.
379 self.assertFalse(registry.syncDimensionData(dimensionName3, dimensionValue3))
380 # Now try that sync with the same primary key but a different value.
381 # This should fail.
382 with self.assertRaises(ConflictingDefinitionError):
383 registry.syncDimensionData(
384 dimensionName3,
385 {
386 "instrument": "DummyCam",
387 "id": 1,
388 "full_name": "one",
389 "name_in_raft": "four",
390 "purpose": "SCIENCE",
391 },
392 )
394 @unittest.skipIf(np is None, "numpy not available.")
395 def testNumpyDataId(self):
396 """Test that we can use a numpy int in a dataId."""
397 butler = self.make_butler()
398 registry = butler.registry
399 dimensionEntries = [
400 ("instrument", {"instrument": "DummyCam"}),
401 ("physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "R"}),
402 ("day_obs", {"instrument": "DummyCam", "id": 20250101}),
403 # Using an np.int64 here fails unless Records.fromDict is also
404 # patched to look for numbers.Integral
405 (
406 "visit",
407 {
408 "instrument": "DummyCam",
409 "id": 42,
410 "name": "fortytwo",
411 "physical_filter": "d-r",
412 "day_obs": 20250101,
413 },
414 ),
415 ]
416 for args in dimensionEntries:
417 registry.insertDimensionData(*args)
419 # Try a normal integer and something that looks like an int but
420 # is not.
421 for visit_id in (42, np.int64(42)):
422 with self.subTest(visit_id=visit_id, id_type=type(visit_id).__name__):
423 expanded = registry.expandDataId({"instrument": "DummyCam", "visit": visit_id})
424 self.assertEqual(expanded["visit"], int(visit_id))
425 self.assertIsInstance(expanded["visit"], int)
427 def testDataIdRelationships(self):
428 """Test that `SqlRegistry.expandDataId` raises an exception when the
429 given keys are inconsistent.
430 """
431 butler = self.make_butler()
432 self.load_data(butler, "base.yaml")
433 registry = butler.registry
434 # Insert a few more dimension records for the next test.
435 registry.insertDimensionData(
436 "day_obs",
437 {"instrument": "Cam1", "id": 20250101},
438 )
439 registry.insertDimensionData(
440 "group",
441 {"instrument": "Cam1", "name": "group1"},
442 )
443 registry.insertDimensionData(
444 "exposure",
445 {
446 "instrument": "Cam1",
447 "id": 1,
448 "obs_id": "one",
449 "physical_filter": "Cam1-G",
450 "group": "group1",
451 "day_obs": 20250101,
452 },
453 )
454 registry.insertDimensionData(
455 "group",
456 {"instrument": "Cam1", "name": "group2"},
457 )
458 registry.insertDimensionData(
459 "exposure",
460 {
461 "instrument": "Cam1",
462 "id": 2,
463 "obs_id": "two",
464 "physical_filter": "Cam1-G",
465 "group": "group2",
466 "day_obs": 20250101,
467 },
468 )
469 registry.insertDimensionData(
470 "visit_system",
471 {"instrument": "Cam1", "id": 0, "name": "one-to-one"},
472 )
473 registry.insertDimensionData(
474 "visit",
475 {"instrument": "Cam1", "id": 1, "name": "one", "physical_filter": "Cam1-G", "day_obs": 20250101},
476 )
477 registry.insertDimensionData(
478 "visit_definition",
479 {"instrument": "Cam1", "visit": 1, "exposure": 1},
480 )
481 with self.assertRaises(InconsistentDataIdError):
482 registry.expandDataId(
483 {"instrument": "Cam1", "visit": 1, "exposure": 2},
484 )
486 def testDataset(self):
487 """Basic tests for `SqlRegistry.insertDatasets`,
488 `SqlRegistry.getDataset`, and `SqlRegistry.removeDatasets`.
489 """
490 butler = self.make_butler()
491 registry = butler.registry
492 self.load_data(butler, "base.yaml")
493 run = "tésτ"
494 registry.registerRun(run)
495 datasetType = registry.getDatasetType("bias")
496 dataId = {"instrument": "Cam1", "detector": 2}
497 (ref,) = registry.insertDatasets(datasetType, dataIds=[dataId], run=run)
498 outRef = registry.getDataset(ref.id)
499 self.assertIsNotNone(ref.id)
500 self.assertEqual(ref, outRef)
501 with self.assertRaises(ConflictingDefinitionError):
502 registry.insertDatasets(datasetType, dataIds=[dataId], run=run)
503 registry.removeDatasets([ref])
504 self.assertIsNone(registry.findDataset(datasetType, dataId, collections=[run]))
506 def test_fetch_run_dataset_ids(self):
507 butler = self.make_butler()
508 registry = butler._registry
509 self.load_data(butler, "base.yaml", "datasets.yaml")
510 dataset_ids = registry._fetch_run_dataset_ids("imported_r")
511 self.assertEqual(len(dataset_ids), 7)
512 refs = butler.query_all_datasets("imported_r")
513 self.assertCountEqual(dataset_ids, [ref.id for ref in refs])
515 def testFindDataset(self):
516 """Tests for `SqlRegistry.findDataset`."""
517 butler = self.make_butler()
518 registry = butler.registry
519 self.load_data(butler, "base.yaml")
520 run = "tésτ"
521 datasetType = registry.getDatasetType("bias")
522 dataId = {"instrument": "Cam1", "detector": 4}
523 registry.registerRun(run)
524 (inputRef,) = registry.insertDatasets(datasetType, dataIds=[dataId], run=run)
525 outputRef = registry.findDataset(datasetType, dataId, collections=[run])
526 self.assertEqual(outputRef, inputRef)
527 # Check that retrieval with invalid dataId raises
528 with self.assertRaises(LookupError):
529 dataId = {"instrument": "Cam1"} # no detector
530 registry.findDataset(datasetType, dataId, collections=run)
531 # Check that different dataIds match to different datasets
532 dataId1 = {"instrument": "Cam1", "detector": 1}
533 (inputRef1,) = registry.insertDatasets(datasetType, dataIds=[dataId1], run=run)
534 dataId2 = {"instrument": "Cam1", "detector": 2}
535 (inputRef2,) = registry.insertDatasets(datasetType, dataIds=[dataId2], run=run)
536 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef1)
537 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef2)
538 self.assertNotEqual(registry.findDataset(datasetType, dataId1, collections=run), inputRef2)
539 self.assertNotEqual(registry.findDataset(datasetType, dataId2, collections=run), inputRef1)
540 # Check that requesting a non-existing dataId returns None
541 nonExistingDataId = {"instrument": "Cam1", "detector": 3}
542 self.assertIsNone(registry.findDataset(datasetType, nonExistingDataId, collections=run))
543 # Search more than one collection, in which two have the right
544 # dataset type and another does not.
545 registry.registerRun("empty")
546 self.load_data(butler, "datasets.yaml")
547 bias1 = registry.findDataset("bias", instrument="Cam1", detector=2, collections=["imported_g"])
548 self.assertIsNotNone(bias1)
549 bias2 = registry.findDataset("bias", instrument="Cam1", detector=2, collections=["imported_r"])
550 self.assertIsNotNone(bias2)
551 self.assertEqual(
552 bias1,
553 registry.findDataset(
554 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_g", "imported_r"]
555 ),
556 )
557 self.assertEqual(
558 bias2,
559 registry.findDataset(
560 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_r", "imported_g"]
561 ),
562 )
563 # If the input data ID was an expanded DataCoordinate with records,
564 # then the output ref has records, too.
565 expanded_id = registry.expandDataId({"instrument": "Cam1", "detector": 2})
566 expanded_ref = registry.findDataset("bias", expanded_id, collections=["imported_r"])
567 self.assertTrue(expanded_ref.dataId.hasRecords())
568 # Search more than one collection, with one of them a CALIBRATION
569 # collection.
570 registry.registerCollection("Cam1/calib", CollectionType.CALIBRATION)
571 timespan = Timespan(
572 begin=astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai"),
573 end=astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai"),
574 )
575 registry.certify("Cam1/calib", [bias2], timespan=timespan)
576 self.assertEqual(
577 bias1,
578 registry.findDataset(
579 "bias",
580 instrument="Cam1",
581 detector=2,
582 collections=["empty", "imported_g", "Cam1/calib"],
583 timespan=timespan,
584 ),
585 )
586 self.assertEqual(
587 bias1,
588 registry.findDataset(
589 "bias",
590 instrument="Cam1",
591 detector=2,
592 # Calibration dataset type, with no calibration collection, but
593 # a timespan was provided.
594 collections=["imported_g"],
595 timespan=timespan,
596 ),
597 )
598 self.assertEqual(
599 bias2,
600 registry.findDataset(
601 "bias",
602 instrument="Cam1",
603 detector=2,
604 collections=["empty", "Cam1/calib", "imported_g"],
605 timespan=timespan,
606 ),
607 )
608 # If we try to search those same collections without a timespan, it
609 # should still work, since the CALIBRATION collection is ignored.
610 self.assertEqual(
611 bias1,
612 registry.findDataset(
613 "bias", instrument="Cam1", detector=2, collections=["empty", "imported_g", "Cam1/calib"]
614 ),
615 )
616 self.assertEqual(
617 bias1,
618 registry.findDataset(
619 "bias", instrument="Cam1", detector=2, collections=["empty", "Cam1/calib", "imported_g"]
620 ),
621 )
622 self.assertIsNone(
623 registry.findDataset("bias", instrument="Cam1", detector=2, collections=["Cam1/calib"])
624 )
625 # Test non-calibration dataset type.
626 registry.registerDatasetType(
627 DatasetType("noncalibration", ["instrument", "detector"], "int", universe=butler.dimensions)
628 )
629 (non_calibration_ref,) = registry.insertDatasets("noncalibration", dataIds=[dataId2], run=run)
630 self.assertIsNone(
631 registry.findDataset("noncalibration", instrument="Cam1", detector=2, collections=["imported_g"])
632 )
633 self.assertEqual(
634 non_calibration_ref,
635 registry.findDataset("noncalibration", instrument="Cam1", detector=2, collections=[run]),
636 )
637 # Timespan parameter is ignored for non-calibration dataset types.
638 self.assertIsNone(
639 registry.findDataset(
640 "noncalibration", instrument="Cam1", detector=2, collections=["imported_g"], timespan=timespan
641 )
642 )
643 self.assertEqual(
644 non_calibration_ref,
645 registry.findDataset(
646 "noncalibration", instrument="Cam1", detector=2, collections=[run], timespan=timespan
647 ),
648 )
649 self.assertEqual(
650 non_calibration_ref,
651 registry.findDataset(
652 "noncalibration",
653 instrument="Cam1",
654 detector=2,
655 collections=["Cam1/calib", run],
656 timespan=timespan,
657 ),
658 )
659 # Add a dataset type whose dimension group involves an "implied"
660 # dimension. ("physical_filter" implies "band".)
661 registry.registerDatasetType(
662 DatasetType(
663 "dt_with_implied",
664 [
665 "instrument",
666 "physical_filter",
667 ],
668 "int",
669 universe=butler.dimensions,
670 )
671 )
672 data_id = {"instrument": "Cam1", "physical_filter": "Cam1-G"}
673 (implied_ref,) = registry.insertDatasets("dt_with_implied", dataIds=[data_id], run=run)
674 found_ref = registry.findDataset("dt_with_implied", data_id, collections=[run])
675 self.assertEqual(implied_ref, found_ref)
676 # The "full" data ID with implied values is looked up, even though we
677 # provided only the "required" values.
678 self.assertTrue(found_ref.dataId.hasFull())
679 # The search ignores excess data ID values beyond the 'required' set.
680 # This is not the correct band value for this physical_filter, but
681 # the mismatch is ignored.
682 self.assertEqual(
683 implied_ref,
684 registry.findDataset(
685 "dt_with_implied",
686 {"instrument": "Cam1", "physical_filter": "Cam1-G", "band": "r"},
687 collections=[run],
688 ),
689 )
690 # Correct band value, wrong physical_filter.
691 self.assertIsNone(
692 registry.findDataset(
693 "dt_with_implied",
694 {"instrument": "Cam1", "physical_filter": "Cam1-R1", "band": "g"},
695 collections=[run],
696 ),
697 )
699 def testRemoveDatasetTypeSuccess(self):
700 """Test that SqlRegistry.removeDatasetType works when there are no
701 datasets of that type present.
702 """
703 butler = self.make_butler()
704 registry = butler.registry
705 self.load_data(butler, "base.yaml")
706 registry.removeDatasetType("flat")
707 with self.assertRaises(MissingDatasetTypeError):
708 registry.getDatasetType("flat")
710 def testRemoveDatasetTypeFailure(self):
711 """Test that SqlRegistry.removeDatasetType raises when there are
712 datasets of that type present or if the dataset type is for a
713 component.
714 """
715 butler = self.make_butler()
716 registry = butler.registry
717 self.load_data(butler, "base.yaml", "datasets.yaml")
718 with self.assertRaises(OrphanedRecordError):
719 registry.removeDatasetType("flat")
720 with self.assertRaises(DatasetTypeError):
721 registry.removeDatasetType(DatasetType.nameWithComponent("flat", "image"))
723 def testImportDatasetsUUID(self):
724 """Test for `SqlRegistry._importDatasets` with UUID dataset ID."""
725 if isinstance(self.datasetsManager, str):
726 if not self.datasetsManager.endswith(".ByDimensionsDatasetRecordStorageManagerUUID"):
727 self.skipTest(f"Unexpected dataset manager {self.datasetsManager}")
728 elif isinstance(self.datasetsManager, dict) and not self.datasetsManager["cls"].endswith(
729 ".ByDimensionsDatasetRecordStorageManagerUUID"
730 ):
731 self.skipTest(f"Unexpected dataset manager {self.datasetsManager['cls']}")
733 butler = self.make_butler()
734 registry = butler.registry
735 self.load_data(butler, "base.yaml")
736 for run in range(6):
737 registry.registerRun(f"run{run}")
738 datasetTypeBias = registry.getDatasetType("bias")
739 datasetTypeFlat = registry.getDatasetType("flat")
740 dataIdBias1 = {"instrument": "Cam1", "detector": 1}
741 dataIdBias2 = {"instrument": "Cam1", "detector": 2}
742 dataIdFlat1 = {"instrument": "Cam1", "detector": 1, "physical_filter": "Cam1-G", "band": "g"}
744 ref = DatasetRef(datasetTypeBias, dataIdBias1, run="run0")
745 (ref1,) = registry._importDatasets([ref], assume_new=True)
746 # UUID is used without change
747 self.assertEqual(ref.id, ref1.id)
749 # Inserting this ref with assume_new=True should fail, since this
750 # dataset exists.
751 with self.assertRaises(ConflictingDefinitionError):
752 registry._importDatasets([ref], assume_new=True)
754 # All different failure modes
755 refs = (
756 # Importing same DatasetRef with different dataset ID is an error
757 DatasetRef(datasetTypeBias, dataIdBias1, run="run0"),
758 # Same DatasetId but different DataId
759 DatasetRef(datasetTypeBias, dataIdBias2, id=ref1.id, run="run0"),
760 DatasetRef(datasetTypeFlat, dataIdFlat1, id=ref1.id, run="run0"),
761 # Same DatasetRef and DatasetId but different run
762 DatasetRef(datasetTypeBias, dataIdBias1, id=ref1.id, run="run1"),
763 )
764 for ref in refs:
765 with self.assertRaises(ConflictingDefinitionError):
766 registry._importDatasets([ref])
768 # Test for non-unique IDs, they can be re-imported multiple times.
769 for run, idGenMode in ((2, DatasetIdGenEnum.DATAID_TYPE), (4, DatasetIdGenEnum.DATAID_TYPE_RUN)):
770 with self.subTest(idGenMode=idGenMode):
771 # Make dataset ref with reproducible dataset ID.
772 ref = DatasetRef(datasetTypeBias, dataIdBias1, run=f"run{run}", id_generation_mode=idGenMode)
773 (ref1,) = registry._importDatasets([ref])
774 self.assertIsInstance(ref1.id, uuid.UUID)
775 self.assertEqual(ref1.id.version, 5)
776 self.assertEqual(ref1.id, ref.id)
778 # Importing it again is OK
779 (ref2,) = registry._importDatasets([ref1])
780 self.assertEqual(ref2.id, ref1.id)
782 # Cannot import to different run with the same ID
783 ref = DatasetRef(datasetTypeBias, dataIdBias1, id=ref1.id, run=f"run{run + 1}")
784 with self.assertRaises(ConflictingDefinitionError):
785 registry._importDatasets([ref])
787 ref = DatasetRef(
788 datasetTypeBias, dataIdBias1, run=f"run{run + 1}", id_generation_mode=idGenMode
789 )
790 if idGenMode is DatasetIdGenEnum.DATAID_TYPE:
791 # Cannot import same DATAID_TYPE ref into a new run
792 with self.assertRaises(ConflictingDefinitionError):
793 (ref2,) = registry._importDatasets([ref])
794 else:
795 # DATAID_TYPE_RUN ref can be imported into a new run
796 (ref2,) = registry._importDatasets([ref])
798 def testComponentLookups(self):
799 """Test searching for component datasets via their parents.
801 Components can no longer be found by registry. This test checks
802 that this now fails.
803 """
804 butler = self.make_butler()
805 registry = butler.registry
806 self.load_data(butler, "base.yaml", "datasets.yaml")
807 # Test getting the child dataset type (which does still exist in the
808 # Registry), and check for consistency with
809 # DatasetRef.makeComponentRef.
810 collection = "imported_g"
811 parentType = registry.getDatasetType("bias")
812 childType = registry.getDatasetType("bias.wcs")
813 parentRefResolved = registry.findDataset(
814 parentType, collections=collection, instrument="Cam1", detector=1
815 )
816 self.assertIsInstance(parentRefResolved, DatasetRef)
817 self.assertEqual(childType, parentRefResolved.makeComponentRef("wcs").datasetType)
818 # Search for a single dataset with findDataset.
819 with self.assertRaises(DatasetTypeError):
820 registry.findDataset("bias.wcs", collections=collection, dataId=parentRefResolved.dataId)
822 def testCollections(self):
823 """Tests for registry methods that manage collections."""
824 butler = self.make_butler()
825 registry = butler.registry
826 other_registry = butler.clone().registry
827 self.load_data(butler, "base.yaml", "datasets.yaml")
828 run1 = "imported_g"
829 run2 = "imported_r"
830 # Test setting a collection docstring after it has been created.
831 registry.setCollectionDocumentation(run1, "doc for run1")
832 self.assertEqual(registry.getCollectionDocumentation(run1), "doc for run1")
833 registry.setCollectionDocumentation(run1, None)
834 self.assertIsNone(registry.getCollectionDocumentation(run1))
835 datasetType = "bias"
836 # Find some datasets via their run's collection.
837 dataId1 = {"instrument": "Cam1", "detector": 1}
838 ref1 = registry.findDataset(datasetType, dataId1, collections=run1)
839 self.assertIsNotNone(ref1)
840 dataId2 = {"instrument": "Cam1", "detector": 2}
841 ref2 = registry.findDataset(datasetType, dataId2, collections=run1)
842 self.assertIsNotNone(ref2)
843 # Associate those into a new collection, then look for them there.
844 tag1 = "tag1"
845 registry.registerCollection(tag1, type=CollectionType.TAGGED, doc="doc for tag1")
846 # Check that we can query for old and new collections by type.
847 self.assertEqual(set(registry.queryCollections(collectionTypes=CollectionType.RUN)), {run1, run2})
848 self.assertEqual(
849 set(registry.queryCollections(collectionTypes={CollectionType.TAGGED, CollectionType.RUN})),
850 {tag1, run1, run2},
851 )
852 self.assertEqual(registry.getCollectionDocumentation(tag1), "doc for tag1")
853 registry.associate(tag1, [ref1, ref2])
854 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)
855 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)
856 # Disassociate one and verify that we can't it there anymore...
857 registry.disassociate(tag1, [ref1])
858 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=tag1))
859 # ...but we can still find ref2 in tag1, and ref1 in the run.
860 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=run1), ref1)
861 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)
862 collections = set(registry.queryCollections())
863 self.assertEqual(collections, {run1, run2, tag1})
864 # Associate both refs into tag1 again; ref2 is already there, but that
865 # should be a harmless no-op.
866 registry.associate(tag1, [ref1, ref2])
867 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)
868 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)
869 # Get a different dataset (from a different run) that has the same
870 # dataset type and data ID as ref2.
871 ref2b = registry.findDataset(datasetType, dataId2, collections=run2)
872 self.assertNotEqual(ref2, ref2b)
873 # Attempting to associate that into tag1 should be an error.
874 with self.assertRaises(ConflictingDefinitionError):
875 registry.associate(tag1, [ref2b])
876 # That error shouldn't have messed up what we had before.
877 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)
878 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)
879 # Attempt to associate the conflicting dataset again, this time with
880 # a dataset that isn't in the collection and won't cause a conflict.
881 # Should also fail without modifying anything.
882 dataId3 = {"instrument": "Cam1", "detector": 3}
883 ref3 = registry.findDataset(datasetType, dataId3, collections=run1)
884 with self.assertRaises(ConflictingDefinitionError):
885 registry.associate(tag1, [ref3, ref2b])
886 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=tag1), ref1)
887 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=tag1), ref2)
888 self.assertIsNone(registry.findDataset(datasetType, dataId3, collections=tag1))
889 # Register a chained collection that searches [tag1, run2]
890 chain1 = "chain1"
891 registry.registerCollection(chain1, type=CollectionType.CHAINED)
892 self.assertIs(registry.getCollectionType(chain1), CollectionType.CHAINED)
893 # Chained collection exists, but has no collections in it.
894 self.assertFalse(registry.getCollectionChain(chain1))
895 # If we query for all collections, we should get the chained collection
896 # if we don't ask to flatten it (i.e. yield only its children) or if we
897 # explicitly ask to include it too.
898 self.assertEqual(set(registry.queryCollections(flattenChains=False)), {tag1, run1, run2, chain1})
899 self.assertEqual(set(registry.queryCollections(flattenChains=True)), {tag1, run1, run2})
900 self.assertEqual(
901 set(registry.queryCollections(flattenChains=True, includeChains=True)), {tag1, run1, run2, chain1}
902 )
903 # Attempt to set its child collections to something circular; that
904 # should fail.
905 with self.assertRaises(ValueError):
906 registry.setCollectionChain(chain1, [tag1, chain1])
907 # Add the child collections.
908 registry.setCollectionChain(chain1, [tag1, run2])
909 self.assertEqual(list(registry.getCollectionChain(chain1)), [tag1, run2])
910 self.assertEqual(registry.getCollectionParentChains(tag1), {chain1})
911 self.assertEqual(registry.getCollectionParentChains(run2), {chain1})
912 # Refresh the other registry that points to the same repo, and make
913 # sure it can see the things we've done (note that this does require
914 # an explicit refresh(); that's the documented behavior, because
915 # caching is ~impossible otherwise).
916 if other_registry is not None:
917 other_registry.refresh()
918 self.assertEqual(list(other_registry.getCollectionChain(chain1)), [tag1, run2])
919 self.assertEqual(other_registry.getCollectionParentChains(tag1), {chain1})
920 self.assertEqual(other_registry.getCollectionParentChains(run2), {chain1})
921 # Searching for dataId1 or dataId2 in the chain should return ref1 and
922 # ref2, because both are in tag1.
923 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain1), ref1)
924 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain1), ref2)
925 # Now disassociate ref2 from tag1. The search (for bias) with
926 # dataId2 in chain1 should then:
927 # 1. not find it in tag1
928 # 2. find a different dataset in run2
929 registry.disassociate(tag1, [ref2])
930 ref2b = registry.findDataset(datasetType, dataId2, collections=chain1)
931 self.assertNotEqual(ref2b, ref2)
932 self.assertEqual(ref2b, registry.findDataset(datasetType, dataId2, collections=run2))
933 # Define a new chain so we can test recursive chains.
934 chain2 = "chain2"
935 registry.registerCollection(chain2, type=CollectionType.CHAINED)
936 registry.setCollectionChain(chain2, [run2, chain1])
937 self.assertEqual(registry.getCollectionParentChains(chain1), {chain2})
938 self.assertEqual(registry.getCollectionParentChains(run2), {chain1, chain2})
940 if self.supportsCollectionRegex:
941 # Query for collections matching a regex.
942 with self.assertWarns(FutureWarning):
943 self.assertCountEqual(
944 list(registry.queryCollections(re.compile("imported_."), flattenChains=False)),
945 ["imported_r", "imported_g"],
946 )
947 # Query for collections matching a regex or an explicit str.
948 with self.assertWarns(FutureWarning):
949 self.assertCountEqual(
950 list(
951 registry.queryCollections([re.compile("imported_."), "chain1"], flattenChains=False)
952 ),
953 ["imported_r", "imported_g", "chain1"],
954 )
955 # Same queries as the regex ones above, but using globs instead of
956 # regex.
957 self.assertCountEqual(
958 list(registry.queryCollections("imported_*", flattenChains=False)),
959 ["imported_r", "imported_g"],
960 )
961 # Query for collections matching a regex or an explicit str.
962 self.assertCountEqual(
963 list(registry.queryCollections(["imported_*", "chain1"], flattenChains=False)),
964 ["imported_r", "imported_g", "chain1"],
965 )
966 # Query for collection matching chain names, by flattening it should
967 # only return non-chain names.
968 self.assertCountEqual(list(registry.queryCollections("chain?", flattenChains=True)), [tag1, run2])
969 # Query for collection matching chain name, by flattening and
970 # asking to include chains it should return everything.
971 self.assertCountEqual(
972 list(registry.queryCollections("chain*", flattenChains=True, includeChains=True)),
973 [tag1, run2, chain1, chain2],
974 )
975 # Order of children in chained collections is preserved.
976 self.assertEqual(list(registry.queryCollections("chain1", flattenChains=True)), [tag1, run2])
977 self.assertEqual(list(registry.queryCollections("cha*2", flattenChains=True)), [run2, tag1])
978 self.assertEqual(
979 list(registry.queryCollections("chain1", flattenChains=True, includeChains=True)),
980 [chain1, tag1, run2],
981 )
982 self.assertEqual(
983 list(registry.queryCollections("chain2", flattenChains=True, includeChains=True)),
984 [chain2, run2, chain1, tag1],
985 )
987 # Search for bias with dataId1 should find it via tag1 in chain2,
988 # recursing, because is not in run1.
989 self.assertIsNone(registry.findDataset(datasetType, dataId1, collections=run2))
990 self.assertEqual(registry.findDataset(datasetType, dataId1, collections=chain2), ref1)
991 # Search for bias with dataId2 should find it in run2 (ref2b).
992 self.assertEqual(registry.findDataset(datasetType, dataId2, collections=chain2), ref2b)
993 # Search for a flat that is in run2. That should not be found
994 # at the front of chain2, because of the restriction to bias
995 # on run2 there, but it should be found in at the end of chain1.
996 dataId4 = {"instrument": "Cam1", "detector": 3, "physical_filter": "Cam1-R2"}
997 ref4 = registry.findDataset("flat", dataId4, collections=run2)
998 self.assertIsNotNone(ref4)
999 self.assertEqual(ref4, registry.findDataset("flat", dataId4, collections=chain2))
1000 # Deleting a collection that's part of a CHAINED collection is not
1001 # allowed, and is exception-safe.
1002 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1003 registry.removeCollection(run2)
1004 self.assertEqual(registry.getCollectionType(run2), CollectionType.RUN)
1005 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1006 registry.removeCollection(chain1)
1007 self.assertEqual(registry.getCollectionType(chain1), CollectionType.CHAINED)
1008 # Actually remove chain2, test that it's gone by asking for its type.
1009 registry.removeCollection(chain2)
1010 with self.assertRaises(MissingCollectionError):
1011 registry.getCollectionType(chain2)
1012 # Actually remove run2 and chain1, which should work now.
1013 registry.removeCollection(chain1)
1014 registry.removeCollection(run2)
1015 with self.assertRaises(MissingCollectionError):
1016 registry.getCollectionType(run2)
1017 with self.assertRaises(MissingCollectionError):
1018 registry.getCollectionType(chain1)
1019 # Remove tag1 as well, just to test that we can remove TAGGED
1020 # collections.
1021 registry.removeCollection(tag1)
1022 with self.assertRaises(MissingCollectionError):
1023 registry.getCollectionType(tag1)
1025 def test_collection_clearing(self) -> None:
1026 """Test that we can delete TAGGED and CALIBRATION collections without
1027 manually removing all associated datasets first.
1028 """
1029 butler = self.make_butler()
1030 self.load_data(butler, "base.yaml", "datasets.yaml")
1032 # This brings in datasets of two different types, with the same
1033 # dimension group.
1034 original_datasets = tuple(butler.query_all_datasets("imported_r", instrument="Cam1", detector=2))
1035 self.assertEqual(len(original_datasets), 2)
1037 # Test tagged collections.
1038 butler.collections.register("tag1", CollectionType.TAGGED)
1039 butler.collections.register("tag2", CollectionType.TAGGED)
1040 butler.registry.associate("tag1", original_datasets)
1041 butler.registry.associate("tag2", original_datasets)
1042 butler.collections.x_remove("tag1")
1043 with self.assertRaises(MissingCollectionError):
1044 butler.collections.get_info("tag1")
1045 # Make sure there was no collateral damage -- tag2 should still be
1046 # intact.
1047 self.assertEqual(set(butler.query_all_datasets("tag2")), set(original_datasets))
1049 # Test calibration collections.
1050 butler.collections.register("calib1", CollectionType.CALIBRATION)
1051 butler.collections.register("calib2", CollectionType.CALIBRATION)
1052 butler.registry.certify("calib1", original_datasets, Timespan(None, None))
1053 butler.registry.certify("calib2", original_datasets, Timespan(None, None))
1054 butler.collections.x_remove("calib1")
1055 with self.assertRaises(MissingCollectionError):
1056 butler.collections.get_info("calib1")
1057 # Make sure there was no collateral damage -- calib2 should still be
1058 # intact.
1059 self.assertEqual(set(butler.query_all_datasets("calib2")), set(original_datasets))
1061 def testCollectionChainCaching(self):
1062 butler = self.make_butler()
1063 registry = butler.registry
1064 with registry.caching_context():
1065 registry.registerCollection("a")
1066 registry.registerCollection("chain", CollectionType.CHAINED)
1067 # There used to be a caching bug (DM-43750) that would throw an
1068 # exception if you modified a collection chain for a collection
1069 # that was already in the cache.
1070 registry.setCollectionChain("chain", ["a"])
1071 self.assertEqual(list(registry.getCollectionChain("chain")), ["a"])
1073 def testCollectionChainFlatten(self):
1074 """Test that `SqlRegistry.setCollectionChain` obeys its 'flatten'
1075 option.
1076 """
1077 butler = self.make_butler()
1078 registry = butler.registry
1079 registry.registerCollection("inner", CollectionType.CHAINED)
1080 registry.registerCollection("innermost", CollectionType.RUN)
1081 registry.setCollectionChain("inner", ["innermost"])
1082 registry.registerCollection("outer", CollectionType.CHAINED)
1083 registry.setCollectionChain("outer", ["inner"], flatten=False)
1084 self.assertEqual(list(registry.getCollectionChain("outer")), ["inner"])
1085 registry.setCollectionChain("outer", ["inner"], flatten=True)
1086 self.assertEqual(list(registry.getCollectionChain("outer")), ["innermost"])
1088 def testCollectionChainPrependConcurrency(self):
1089 """Verify that locking via database row locks is working as
1090 expected.
1091 """
1093 def blocked_thread_func(butler: Butler):
1094 # This call will become blocked after it has decided on positions
1095 # for the new children in the collection chain, but before
1096 # inserting them.
1097 butler.collections.prepend_chain("chain", ["a"])
1099 def unblocked_thread_func(butler: Butler):
1100 butler.collections.prepend_chain("chain", ["b"])
1102 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)
1104 # blocked_thread_func should have finished first, inserting "a".
1105 # unblocked_thread_func should have finished second, prepending "b".
1106 self.assertEqual(("b", "a"), registry.getCollectionChain("chain"))
1108 def testCollectionChainReplaceConcurrency(self):
1109 """Verify that locking via database row locks is working as
1110 expected.
1111 """
1113 def blocked_thread_func(butler: Butler):
1114 # This call will become blocked after deleting children, but before
1115 # inserting new ones.
1116 butler.collections.redefine_chain("chain", ["a"])
1118 def unblocked_thread_func(butler: Butler):
1119 butler.collections.redefine_chain("chain", ["b"])
1121 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)
1123 # blocked_thread_func should have finished first.
1124 # unblocked_thread_func should have finished second, overwriting the
1125 # chain with "b".
1126 self.assertEqual(("b",), registry.getCollectionChain("chain"))
1128 def testCollectionChainRemoveConcurrency(self):
1129 def blocked_thread_func(butler: Butler):
1130 # This call will become blocked after taking the lock, but before
1131 # deleting the children.
1132 butler.collections.remove_from_chain("chain", ["b"])
1134 def unblocked_thread_func(butler: Butler):
1135 butler.collections.redefine_chain("chain", ["b", "a"])
1137 registry = self._do_collection_concurrency_test(blocked_thread_func, unblocked_thread_func)
1139 # blocked_thread_func should have finished first, removing "b".
1140 # unblocked_thread_func should have finished second, putting "b" back.
1141 self.assertEqual(("b", "a"), registry.getCollectionChain("chain"))
1143 def _do_collection_concurrency_test(
1144 self, blocked_thread_func: Callable[[Butler]], unblocked_thread_func: Callable[[Butler]]
1145 ) -> SqlRegistry:
1146 # This function:
1147 # 1. Sets up two registries pointing at the same database.
1148 # 2. Start running 'blocked_thread_func' in a background thread,
1149 # arranging for it to become blocked during a critical section in
1150 # the collections manager.
1151 # 3. Wait for 'blocked_thread_func' to reach the critical section
1152 # 4. Start running 'unblocked_thread_func'.
1153 # 5. Allow both functions to run to completion.
1155 # Set up two registries pointing to the same DB
1156 butler1 = self.make_butler()
1157 butler2 = butler1.clone()
1158 registry1 = butler1._registry
1159 assert isinstance(registry1, SqlRegistry)
1160 registry2 = butler2._registry
1162 with contextlib.suppress(AttributeError):
1163 if ":memory:" in str(registry2._db):
1164 raise unittest.SkipTest("Testing concurrency requires two connections to the same DB.")
1166 registry1.registerCollection("chain", CollectionType.CHAINED)
1167 for collection in ["a", "b"]:
1168 registry1.registerCollection(collection)
1170 # Arrange for registry1 to block during its critical section, allowing
1171 # us to detect this and control when it becomes unblocked.
1172 enter_barrier = Barrier(2, timeout=60)
1173 exit_barrier = Barrier(2, timeout=60)
1175 def wait_for_barrier():
1176 enter_barrier.wait()
1177 exit_barrier.wait()
1179 registry1._managers.collections._block_for_concurrency_test = wait_for_barrier
1181 with ThreadPoolExecutor(max_workers=1) as exec1:
1182 with ThreadPoolExecutor(max_workers=1) as exec2:
1183 future1 = exec1.submit(blocked_thread_func, butler1)
1184 enter_barrier.wait()
1186 # At this point registry 1 has entered the critical section and
1187 # is waiting for us to release it. Start the other thread.
1188 future2 = exec2.submit(unblocked_thread_func, butler2)
1189 # thread2 should block inside a database call, but we have no
1190 # way to detect when it is in this state.
1191 time.sleep(0.200)
1193 # Let the threads run to completion.
1194 exit_barrier.wait()
1195 future1.result()
1196 future2.result()
1198 return registry1
1200 def testBasicTransaction(self):
1201 """Test that all operations within a single transaction block are
1202 rolled back if an exception propagates out of the block.
1203 """
1204 butler = self.make_butler()
1205 registry = butler.registry
1206 storageClass = StorageClass("testDatasetType")
1207 registry.storageClasses.registerStorageClass(storageClass)
1208 with registry.transaction():
1209 registry.insertDimensionData("instrument", {"name": "Cam1", "class_name": "A"})
1210 with self.assertRaises(ValueError):
1211 with registry.transaction():
1212 registry.insertDimensionData("instrument", {"name": "Cam2"})
1213 raise ValueError("Oops, something went wrong")
1214 # Cam1 should exist
1215 self.assertEqual(registry.expandDataId(instrument="Cam1").records["instrument"].class_name, "A")
1216 # But Cam2 and Cam3 should both not exist
1217 with self.assertRaises(DataIdValueError):
1218 registry.expandDataId(instrument="Cam2")
1219 with self.assertRaises(DataIdValueError):
1220 registry.expandDataId(instrument="Cam3")
1222 def testNestedTransaction(self):
1223 """Test that operations within a transaction block are not rolled back
1224 if an exception propagates out of an inner transaction block and is
1225 then caught.
1226 """
1227 butler = self.make_butler()
1228 registry = butler.registry
1229 dimension = registry.dimensions["instrument"]
1230 dataId1 = {"instrument": "DummyCam"}
1231 dataId2 = {"instrument": "DummyCam2"}
1232 checkpointReached = False
1233 with registry.transaction():
1234 # This should be added and (ultimately) committed.
1235 registry.insertDimensionData(dimension, dataId1)
1236 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1237 with registry.transaction(savepoint=True):
1238 # This does not conflict, and should succeed (but not
1239 # be committed).
1240 registry.insertDimensionData(dimension, dataId2)
1241 checkpointReached = True
1242 # This should conflict and raise, triggering a rollback
1243 # of the previous insertion within the same transaction
1244 # context, but not the original insertion in the outer
1245 # block.
1246 registry.insertDimensionData(dimension, dataId1)
1247 self.assertTrue(checkpointReached)
1248 self.assertIsNotNone(registry.expandDataId(dataId1, dimensions=dimension.minimal_group))
1249 with self.assertRaises(DataIdValueError):
1250 registry.expandDataId(dataId2, dimensions=dimension.minimal_group)
1252 def testInstrumentDimensions(self):
1253 """Test queries involving only instrument dimensions, with no joins to
1254 skymap.
1255 """
1256 butler = self.make_butler()
1257 registry = butler.registry
1259 # need a bunch of dimensions and datasets for test
1260 registry.insertDimensionData(
1261 "instrument", dict(name="DummyCam", visit_max=25, exposure_max=300, detector_max=6)
1262 )
1263 registry.insertDimensionData("day_obs", dict(instrument="DummyCam", id=20250101))
1264 registry.insertDimensionData(
1265 "physical_filter",
1266 dict(instrument="DummyCam", name="dummy_r", band="r"),
1267 dict(instrument="DummyCam", name="dummy_i", band="i"),
1268 )
1269 registry.insertDimensionData(
1270 "detector", *[dict(instrument="DummyCam", id=i, full_name=str(i)) for i in range(1, 6)]
1271 )
1272 registry.insertDimensionData(
1273 "visit",
1274 dict(instrument="DummyCam", id=10, name="ten", physical_filter="dummy_i", day_obs=20250101),
1275 dict(instrument="DummyCam", id=11, name="eleven", physical_filter="dummy_r", day_obs=20250101),
1276 dict(instrument="DummyCam", id=20, name="twelve", physical_filter="dummy_r", day_obs=20250101),
1277 )
1278 registry.insertDimensionData(
1279 "group",
1280 dict(instrument="DummyCam", name="ten"),
1281 dict(instrument="DummyCam", name="eleven"),
1282 dict(instrument="DummyCam", name="twelve"),
1283 )
1284 for i in range(1, 6):
1285 registry.insertDimensionData(
1286 "visit_detector_region",
1287 dict(instrument="DummyCam", visit=10, detector=i),
1288 dict(instrument="DummyCam", visit=11, detector=i),
1289 dict(instrument="DummyCam", visit=20, detector=i),
1290 )
1291 registry.insertDimensionData(
1292 "exposure",
1293 dict(
1294 instrument="DummyCam",
1295 id=100,
1296 obs_id="100",
1297 physical_filter="dummy_i",
1298 group="ten",
1299 day_obs=20250101,
1300 ),
1301 dict(
1302 instrument="DummyCam",
1303 id=101,
1304 obs_id="101",
1305 physical_filter="dummy_i",
1306 group="ten",
1307 day_obs=20250101,
1308 ),
1309 dict(
1310 instrument="DummyCam",
1311 id=110,
1312 obs_id="110",
1313 physical_filter="dummy_r",
1314 group="eleven",
1315 day_obs=20250101,
1316 ),
1317 dict(
1318 instrument="DummyCam",
1319 id=111,
1320 obs_id="111",
1321 physical_filter="dummy_r",
1322 group="eleven",
1323 day_obs=20250101,
1324 ),
1325 dict(
1326 instrument="DummyCam",
1327 id=200,
1328 obs_id="200",
1329 physical_filter="dummy_r",
1330 group="twelve",
1331 day_obs=20250101,
1332 ),
1333 dict(
1334 instrument="DummyCam",
1335 id=201,
1336 obs_id="201",
1337 physical_filter="dummy_r",
1338 group="twelve",
1339 day_obs=20250101,
1340 ),
1341 )
1342 registry.insertDimensionData(
1343 "visit_definition",
1344 dict(instrument="DummyCam", exposure=100, visit=10),
1345 dict(instrument="DummyCam", exposure=101, visit=10),
1346 dict(instrument="DummyCam", exposure=110, visit=11),
1347 dict(instrument="DummyCam", exposure=111, visit=11),
1348 dict(instrument="DummyCam", exposure=200, visit=20),
1349 dict(instrument="DummyCam", exposure=201, visit=20),
1350 )
1351 # dataset types
1352 run1 = "test1_r"
1353 run2 = "test2_r"
1354 tagged2 = "test2_t"
1355 registry.registerRun(run1)
1356 registry.registerRun(run2)
1357 registry.registerCollection(tagged2)
1358 storageClass = StorageClass("testDataset")
1359 registry.storageClasses.registerStorageClass(storageClass)
1360 rawType = DatasetType(
1361 name="RAW",
1362 dimensions=registry.dimensions.conform(("instrument", "exposure", "detector")),
1363 storageClass=storageClass,
1364 )
1365 registry.registerDatasetType(rawType)
1366 calexpType = DatasetType(
1367 name="CALEXP",
1368 dimensions=registry.dimensions.conform(("instrument", "visit", "detector")),
1369 storageClass=storageClass,
1370 )
1371 registry.registerDatasetType(calexpType)
1373 # add pre-existing datasets
1374 for exposure in (100, 101, 110, 111):
1375 for detector in (1, 2, 3):
1376 # note that only 3 of 5 detectors have datasets
1377 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector)
1378 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run1)
1379 # exposures 100 and 101 appear in both run1 and tagged2.
1380 # 100 has different datasets in the different collections
1381 # 101 has the same dataset in both collections.
1382 if exposure == 100:
1383 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run2)
1384 if exposure in (100, 101):
1385 registry.associate(tagged2, [ref])
1386 # Add pre-existing datasets to tagged2.
1387 for exposure in (200, 201):
1388 for detector in (3, 4, 5):
1389 # note that only 3 of 5 detectors have datasets
1390 dataId = dict(instrument="DummyCam", exposure=exposure, detector=detector)
1391 (ref,) = registry.insertDatasets(rawType, dataIds=[dataId], run=run2)
1392 registry.associate(tagged2, [ref])
1394 dimensions = registry.dimensions.conform(rawType.dimensions.required | calexpType.dimensions.required)
1395 # Test that single dim string works as well as list of str
1396 rows = registry.queryDataIds("visit", datasets=rawType, collections=run1).expanded().toSet()
1397 rowsI = registry.queryDataIds(["visit"], datasets=rawType, collections=run1).expanded().toSet()
1398 self.assertEqual(rows, rowsI)
1399 # with empty expression
1400 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=run1).expanded().toSet()
1401 self.assertEqual(len(rows), 4 * 3) # 4 exposures times 3 detectors
1402 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 110, 111))
1403 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 11))
1404 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))
1406 # second collection
1407 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=tagged2).toSet()
1408 self.assertEqual(len(rows), 4 * 3) # 4 exposures times 3 detectors
1409 for dataId in rows:
1410 self.assertCountEqual(dataId.dimensions.required, ("instrument", "detector", "exposure", "visit"))
1411 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 200, 201))
1412 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 20))
1413 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3, 4, 5))
1415 # with two input datasets
1416 rows = registry.queryDataIds(dimensions, datasets=rawType, collections=[run1, tagged2]).toSet()
1417 self.assertEqual(len(set(rows)), 6 * 3) # 6 exposures times 3 detectors; set needed to de-dupe
1418 for dataId in rows:
1419 self.assertCountEqual(dataId.dimensions.required, ("instrument", "detector", "exposure", "visit"))
1420 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101, 110, 111, 200, 201))
1421 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10, 11, 20))
1422 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3, 4, 5))
1424 # limit to single visit
1425 rows = registry.queryDataIds(
1426 dimensions, datasets=rawType, collections=run1, where="visit = 10", instrument="DummyCam"
1427 ).toSet()
1428 self.assertEqual(len(rows), 2 * 3) # 2 exposures times 3 detectors
1429 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101))
1430 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10,))
1431 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))
1433 # more limiting expression, using link names instead of Table.column
1434 rows = registry.queryDataIds(
1435 dimensions,
1436 datasets=rawType,
1437 collections=run1,
1438 where="visit = 10 and detector > 1 and 'DummyCam'=instrument",
1439 ).toSet()
1440 self.assertEqual(len(rows), 2 * 2) # 2 exposures times 2 detectors
1441 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (100, 101))
1442 self.assertCountEqual({dataId["visit"] for dataId in rows}, (10,))
1443 self.assertCountEqual({dataId["detector"] for dataId in rows}, (2, 3))
1445 # queryDataIds with only one of `datasets` and `collections` is an
1446 # error.
1447 with self.assertRaises(CollectionError):
1448 registry.queryDataIds(dimensions, datasets=rawType)
1449 with self.assertRaises(ArgumentError):
1450 registry.queryDataIds(dimensions, collections=run1)
1452 # expression excludes everything
1453 rows = registry.queryDataIds(
1454 dimensions, datasets=rawType, collections=run1, where="visit > 1000", instrument="DummyCam"
1455 ).toSet()
1456 self.assertEqual(len(rows), 0)
1458 # Selecting by physical_filter, this is not in the dimensions, but it
1459 # is a part of the full expression so it should work too.
1460 rows = registry.queryDataIds(
1461 dimensions,
1462 datasets=rawType,
1463 collections=run1,
1464 where="physical_filter = 'dummy_r'",
1465 instrument="DummyCam",
1466 ).toSet()
1467 self.assertEqual(len(rows), 2 * 3) # 2 exposures times 3 detectors
1468 self.assertCountEqual({dataId["exposure"] for dataId in rows}, (110, 111))
1469 self.assertCountEqual({dataId["visit"] for dataId in rows}, (11,))
1470 self.assertCountEqual({dataId["detector"] for dataId in rows}, (1, 2, 3))
1472 def testSkyMapDimensions(self):
1473 """Tests involving only skymap dimensions, no joins to instrument."""
1474 butler = self.make_butler()
1475 registry = butler.registry
1477 # need a bunch of dimensions and datasets for test, we want
1478 # "band" in the test so also have to add physical_filter
1479 # dimensions
1480 registry.insertDimensionData("instrument", dict(instrument="DummyCam"))
1481 registry.insertDimensionData(
1482 "physical_filter",
1483 dict(instrument="DummyCam", name="dummy_r", band="r"),
1484 dict(instrument="DummyCam", name="dummy_i", band="i"),
1485 )
1486 registry.insertDimensionData("skymap", dict(name="DummyMap", hash=b"sha!"))
1487 for tract in range(10):
1488 registry.insertDimensionData("tract", dict(skymap="DummyMap", id=tract))
1489 registry.insertDimensionData(
1490 "patch",
1491 *[dict(skymap="DummyMap", tract=tract, id=patch, cell_x=0, cell_y=0) for patch in range(10)],
1492 )
1494 # dataset types
1495 run = "tésτ"
1496 registry.registerRun(run)
1497 storageClass = StorageClass("testDataset")
1498 registry.storageClasses.registerStorageClass(storageClass)
1499 calexpType = DatasetType(
1500 name="deepCoadd_calexp",
1501 dimensions=registry.dimensions.conform(("skymap", "tract", "patch", "band")),
1502 storageClass=storageClass,
1503 )
1504 registry.registerDatasetType(calexpType)
1505 mergeType = DatasetType(
1506 name="deepCoadd_mergeDet",
1507 dimensions=registry.dimensions.conform(("skymap", "tract", "patch")),
1508 storageClass=storageClass,
1509 )
1510 registry.registerDatasetType(mergeType)
1511 measType = DatasetType(
1512 name="deepCoadd_meas",
1513 dimensions=registry.dimensions.conform(("skymap", "tract", "patch", "band")),
1514 storageClass=storageClass,
1515 )
1516 registry.registerDatasetType(measType)
1518 dimensions = registry.dimensions.conform(
1519 calexpType.dimensions.required | mergeType.dimensions.required | measType.dimensions.required
1520 )
1522 # add pre-existing datasets
1523 for tract in (1, 3, 5):
1524 for patch in (2, 4, 6, 7):
1525 dataId = dict(skymap="DummyMap", tract=tract, patch=patch)
1526 registry.insertDatasets(mergeType, dataIds=[dataId], run=run)
1527 for aFilter in ("i", "r"):
1528 dataId = dict(skymap="DummyMap", tract=tract, patch=patch, band=aFilter)
1529 registry.insertDatasets(calexpType, dataIds=[dataId], run=run)
1531 # with empty expression
1532 rows = registry.queryDataIds(dimensions, datasets=[calexpType, mergeType], collections=run).toSet()
1533 self.assertEqual(len(rows), 3 * 4 * 2) # 4 tracts x 4 patches x 2 filters
1534 for dataId in rows:
1535 self.assertCountEqual(dataId.dimensions.required, ("skymap", "tract", "patch", "band"))
1536 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 3, 5))
1537 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 4, 6, 7))
1538 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i", "r"))
1540 # limit to 2 tracts and 2 patches
1541 rows = registry.queryDataIds(
1542 dimensions,
1543 datasets=[calexpType, mergeType],
1544 collections=run,
1545 where="tract IN (1, 5) AND patch IN (2, 7)",
1546 skymap="DummyMap",
1547 ).toSet()
1548 self.assertEqual(len(rows), 2 * 2 * 2) # 2 tracts x 2 patches x 2 filters
1549 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 5))
1550 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 7))
1551 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i", "r"))
1553 # limit to single filter
1554 rows = registry.queryDataIds(
1555 dimensions, datasets=[calexpType, mergeType], collections=run, where="band = 'i'"
1556 ).toSet()
1557 self.assertEqual(len(rows), 3 * 4 * 1) # 4 tracts x 4 patches x 2 filters
1558 self.assertCountEqual({dataId["tract"] for dataId in rows}, (1, 3, 5))
1559 self.assertCountEqual({dataId["patch"] for dataId in rows}, (2, 4, 6, 7))
1560 self.assertCountEqual({dataId["band"] for dataId in rows}, ("i",))
1562 def do_query():
1563 return registry.queryDataIds(
1564 dimensions, datasets=[calexpType, mergeType], collections=run, where="skymap = 'Mars'"
1565 ).toSet()
1567 self.assertEqual(len(do_query()), 0)
1569 def testSpatialJoin(self):
1570 """Test queries that involve spatial overlap joins."""
1571 butler = self.make_butler()
1572 registry = butler.registry
1573 self.load_data(butler, "base.yaml", "spatial.yaml")
1575 # Dictionary of spatial DatabaseDimensionElements, keyed by the name of
1576 # the TopologicalFamily they belong to. We'll relate all elements in
1577 # each family to all of the elements in each other family.
1578 families = defaultdict(set)
1579 # Dictionary of {element.name: {dataId: region}}.
1580 regions = {}
1581 for element in registry.dimensions.database_elements:
1582 if element.spatial is not None:
1583 families[element.spatial.name].add(element)
1584 regions[element.name] = {
1585 record.dataId: record.region for record in registry.queryDimensionRecords(element)
1586 }
1588 # If this check fails, it's not necessarily a problem - it may just be
1589 # a reasonable change to the default dimension definitions - but the
1590 # test below depends on there being more than one family to do anything
1591 # useful.
1592 self.assertEqual(len(families), 2)
1594 # Overlap DatabaseDimensionElements with each other.
1595 for family1, family2 in itertools.combinations(families, 2):
1596 for element1, element2 in itertools.product(families[family1], families[family2]):
1597 dimensions = element1.minimal_group | element2.minimal_group
1598 # Construct expected set of overlapping data IDs via a
1599 # brute-force comparison of the regions we've already fetched.
1600 expected = {
1601 DataCoordinate.standardize(
1602 {**dataId1.required, **dataId2.required}, dimensions=dimensions
1603 )
1604 for (dataId1, region1), (dataId2, region2) in itertools.product(
1605 regions[element1.name].items(), regions[element2.name].items()
1606 )
1607 if not region1.isDisjointFrom(region2)
1608 }
1609 self.assertGreater(len(expected), 2, msg="Test that we aren't just comparing empty sets.")
1610 queried = set(registry.queryDataIds(dimensions))
1611 self.assertEqual(expected, queried)
1613 # Overlap each DatabaseDimensionElement with the commonSkyPix system.
1614 commonSkyPix = registry.dimensions.commonSkyPix
1615 for elementName, these_regions in regions.items():
1616 dimensions = registry.dimensions[elementName].minimal_group | commonSkyPix.minimal_group
1617 expected = set()
1618 for dataId, region in these_regions.items():
1619 for begin, end in commonSkyPix.pixelization.envelope(region):
1620 expected.update(
1621 DataCoordinate.standardize(
1622 {commonSkyPix.name: index, **dataId.required}, dimensions=dimensions
1623 )
1624 for index in range(begin, end)
1625 )
1626 self.assertGreater(len(expected), 2, msg="Test that we aren't just comparing empty sets.")
1627 queried = set(registry.queryDataIds(dimensions))
1628 self.assertEqual(expected, queried)
1630 def testAbstractQuery(self):
1631 """Test that we can run a query that just lists the known
1632 bands. This is tricky because band is
1633 backed by a query against physical_filter.
1634 """
1635 butler = self.make_butler()
1636 registry = butler.registry
1637 registry.insertDimensionData("instrument", dict(name="DummyCam"))
1638 registry.insertDimensionData(
1639 "physical_filter",
1640 dict(instrument="DummyCam", name="dummy_i", band="i"),
1641 dict(instrument="DummyCam", name="dummy_i2", band="i"),
1642 dict(instrument="DummyCam", name="dummy_r", band="r"),
1643 )
1644 rows = registry.queryDataIds(["band"]).toSet()
1645 self.assertCountEqual(
1646 rows,
1647 [
1648 DataCoordinate.standardize(band="i", universe=registry.dimensions),
1649 DataCoordinate.standardize(band="r", universe=registry.dimensions),
1650 ],
1651 )
1653 def testAttributeManager(self):
1654 """Test basic functionality of attribute manager."""
1655 # number of attributes with schema versions in a fresh database,
1656 # 6 managers with 2 records per manager, plus config for dimensions
1657 VERSION_COUNT = 6 * 2 + 1
1659 butler = self.make_butler()
1660 registry = butler._registry
1661 attributes = registry._managers.attributes
1663 # check what get() returns for non-existing key
1664 self.assertIsNone(attributes.get("attr"))
1665 self.assertEqual(attributes.get("attr", ""), "")
1666 self.assertEqual(attributes.get("attr", "Value"), "Value")
1667 self.assertEqual(len(list(attributes.items())), VERSION_COUNT)
1669 # cannot store empty key or value
1670 with self.assertRaises(ValueError):
1671 attributes.set("", "value")
1672 with self.assertRaises(ValueError):
1673 attributes.set("attr", "")
1675 # set value of non-existing key
1676 attributes.set("attr", "value")
1677 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)
1678 self.assertEqual(attributes.get("attr"), "value")
1680 # update value of existing key
1681 with self.assertRaises(ButlerAttributeExistsError):
1682 attributes.set("attr", "value2")
1684 attributes.set("attr", "value2", force=True)
1685 self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)
1686 self.assertEqual(attributes.get("attr"), "value2")
1688 # delete existing key
1689 self.assertTrue(attributes.delete("attr"))
1690 self.assertEqual(len(list(attributes.items())), VERSION_COUNT)
1692 # delete non-existing key
1693 self.assertFalse(attributes.delete("non-attr"))
1695 # store bunch of keys and get the list back
1696 data = [
1697 ("version.core", "1.2.3"),
1698 ("version.dimensions", "3.2.1"),
1699 ("config.managers.opaque", "ByNameOpaqueTableStorageManager"),
1700 ]
1701 for key, value in data:
1702 attributes.set(key, value)
1703 items = dict(attributes.items())
1704 for key, value in data:
1705 self.assertEqual(items[key], value)
1707 def testQueryDatasetsDeduplication(self):
1708 """Test that the findFirst option to queryDatasets selects datasets
1709 from collections in the order given".
1710 """
1711 butler = self.make_butler()
1712 registry = butler.registry
1713 self.load_data(butler, "base.yaml", "datasets.yaml")
1714 self.assertCountEqual(
1715 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"])),
1716 [
1717 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),
1718 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g"),
1719 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g"),
1720 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r"),
1721 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r"),
1722 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),
1723 ],
1724 )
1725 self.assertCountEqual(
1726 list(registry.queryDatasets("bias", collections=["imported_g", "imported_r"], findFirst=True)),
1727 [
1728 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),
1729 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g"),
1730 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g"),
1731 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),
1732 ],
1733 )
1734 self.assertCountEqual(
1735 list(registry.queryDatasets("bias", collections=["imported_r", "imported_g"], findFirst=True)),
1736 [
1737 registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),
1738 registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r"),
1739 registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r"),
1740 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),
1741 ],
1742 )
1744 with self.assertRaises(TypeError):
1745 # Collection wildcards not allowed in find-first searches because
1746 # they do not guarantee the ordering of collections.
1747 registry.queryDatasets("bias", collections="imported_*", findFirst=True)
1749 def testQueryDatasetsExtraDimensions(self):
1750 butler = self.make_butler()
1751 registry = butler.registry
1752 self.load_data(butler, "base.yaml", "datasets.yaml")
1753 # Bias dataset type does not include physical filter. By adding
1754 # "physical_filter" to dimensions, we are effectively searching here
1755 # for bias datasets with an instrument that has a specific filter
1756 # available, even though that filter has nothing to do with the bias
1757 # datasets we are finding.
1758 self.assertEqual(
1759 0,
1760 registry.queryDatasets(
1761 "bias",
1762 collections=...,
1763 dimensions=["physical_filter"],
1764 dataId={
1765 "instrument": "Cam1",
1766 "band": "not_a_real_band",
1767 },
1768 ).count(),
1769 )
1770 self.assertEqual(
1771 6,
1772 len(
1773 set(
1774 registry.queryDatasets(
1775 "bias",
1776 collections=...,
1777 dimensions=["physical_filter"],
1778 dataId={
1779 "instrument": "Cam1",
1780 "band": "r",
1781 },
1782 )
1783 )
1784 ),
1785 )
1787 def testQueryResults(self):
1788 """Test querying for data IDs and then manipulating the QueryResults
1789 object returned to perform other queries.
1790 """
1791 butler = self.make_butler()
1792 registry = butler.registry
1793 self.load_data(butler, "base.yaml", "datasets.yaml")
1794 bias = registry.getDatasetType("bias")
1795 flat = registry.getDatasetType("flat")
1796 # Obtain expected results from methods other than those we're testing
1797 # here. That includes:
1798 # - the dimensions of the data IDs we want to query:
1799 expected_dimensions = registry.dimensions.conform(["detector", "physical_filter"])
1800 # - the dimensions of some other data IDs we'll extract from that:
1801 expected_subset_dimensions = registry.dimensions.conform(["detector"])
1802 # - the data IDs we expect to obtain from the first queries:
1803 expectedDataIds = DataCoordinateSet(
1804 {
1805 DataCoordinate.standardize(
1806 instrument="Cam1", detector=d, physical_filter=p, universe=registry.dimensions
1807 )
1808 for d, p in itertools.product({1, 2, 3}, {"Cam1-G", "Cam1-R1", "Cam1-R2"})
1809 },
1810 dimensions=expected_dimensions,
1811 hasFull=False,
1812 hasRecords=False,
1813 )
1814 # - the flat datasets we expect to find from those data IDs, in just
1815 # one collection (so deduplication is irrelevant):
1816 expectedFlats = [
1817 registry.findDataset(
1818 flat, instrument="Cam1", detector=1, physical_filter="Cam1-R1", collections="imported_r"
1819 ),
1820 registry.findDataset(
1821 flat, instrument="Cam1", detector=2, physical_filter="Cam1-R1", collections="imported_r"
1822 ),
1823 registry.findDataset(
1824 flat, instrument="Cam1", detector=3, physical_filter="Cam1-R2", collections="imported_r"
1825 ),
1826 ]
1827 # - the data IDs we expect to extract from that:
1828 expectedSubsetDataIds = expectedDataIds.subset(expected_subset_dimensions)
1829 # - the bias datasets we expect to find from those data IDs, after we
1830 # subset-out the physical_filter dimension, both with duplicates:
1831 expectedAllBiases = [
1832 registry.findDataset(bias, instrument="Cam1", detector=1, collections="imported_g"),
1833 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_g"),
1834 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_g"),
1835 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_r"),
1836 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_r"),
1837 ]
1838 # - ...and without duplicates:
1839 expectedDeduplicatedBiases = [
1840 registry.findDataset(bias, instrument="Cam1", detector=1, collections="imported_g"),
1841 registry.findDataset(bias, instrument="Cam1", detector=2, collections="imported_r"),
1842 registry.findDataset(bias, instrument="Cam1", detector=3, collections="imported_r"),
1843 ]
1844 # Test against those expected results, using a "lazy" query for the
1845 # data IDs (which re-executes that query each time we use it to do
1846 # something new).
1847 dataIds = registry.queryDataIds(
1848 ["detector", "physical_filter"],
1849 where="detector.purpose = 'SCIENCE'", # this rejects detector=4
1850 instrument="Cam1",
1851 )
1852 self.assertEqual(dataIds.dimensions, expected_dimensions)
1853 self.assertEqual(dataIds.toSet(), expectedDataIds)
1854 self.assertCountEqual(
1855 list(
1856 dataIds.findDatasets(
1857 flat,
1858 collections=["imported_r"],
1859 )
1860 ),
1861 expectedFlats,
1862 )
1863 subsetDataIds = dataIds.subset(expected_subset_dimensions, unique=True)
1864 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)
1865 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)
1866 self.assertCountEqual(
1867 list(subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=False)),
1868 expectedAllBiases,
1869 )
1870 self.assertCountEqual(
1871 list(subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)),
1872 expectedDeduplicatedBiases,
1873 )
1875 # Searching for a dataset with dimensions we had projected away
1876 # restores those dimensions.
1877 self.assertCountEqual(
1878 list(subsetDataIds.findDatasets("flat", collections=["imported_r"], findFirst=True)),
1879 expectedFlats,
1880 )
1882 # Use a named dataset type that does not exist and a dataset type
1883 # object that does not exist.
1884 unknown_type = DatasetType("not_known", dimensions=bias.dimensions, storageClass="Exposure")
1886 # Test both string name and dataset type object.
1887 test_type: str | DatasetType
1888 for test_type, test_type_name in (
1889 (unknown_type, unknown_type.name),
1890 (unknown_type.name, unknown_type.name),
1891 ):
1892 with self.assertRaisesRegex(DatasetTypeError, expected_regex=test_type_name):
1893 list(
1894 subsetDataIds.findDatasets(
1895 test_type, collections=["imported_r", "imported_g"], findFirst=True
1896 )
1897 )
1899 # Materialize the data ID subset query, but not the dataset queries.
1900 with subsetDataIds.materialize() as subsetDataIds:
1901 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)
1902 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)
1903 self.assertCountEqual(
1904 list(
1905 subsetDataIds.findDatasets(
1906 bias, collections=["imported_r", "imported_g"], findFirst=False
1907 )
1908 ),
1909 expectedAllBiases,
1910 )
1911 self.assertCountEqual(
1912 list(
1913 subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)
1914 ),
1915 expectedDeduplicatedBiases,
1916 )
1917 # Materialize the original query, but none of the follow-up queries.
1918 with dataIds.materialize() as dataIds:
1919 self.assertEqual(dataIds.dimensions, expected_dimensions)
1920 self.assertEqual(dataIds.toSet(), expectedDataIds)
1921 self.assertCountEqual(
1922 list(
1923 dataIds.findDatasets(
1924 flat,
1925 collections=["imported_r"],
1926 )
1927 ),
1928 expectedFlats,
1929 )
1930 subsetDataIds = dataIds.subset(expected_subset_dimensions, unique=True)
1931 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)
1932 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)
1933 self.assertCountEqual(
1934 list(
1935 subsetDataIds.findDatasets(
1936 bias, collections=["imported_r", "imported_g"], findFirst=False
1937 )
1938 ),
1939 expectedAllBiases,
1940 )
1941 self.assertCountEqual(
1942 list(
1943 subsetDataIds.findDatasets(bias, collections=["imported_r", "imported_g"], findFirst=True)
1944 ),
1945 expectedDeduplicatedBiases,
1946 )
1947 # Materialize the subset data ID query, but not the dataset
1948 # queries.
1949 with subsetDataIds.materialize() as subsetDataIds:
1950 self.assertEqual(subsetDataIds.dimensions, expected_subset_dimensions)
1951 self.assertEqual(subsetDataIds.toSet(), expectedSubsetDataIds)
1952 self.assertCountEqual(
1953 list(
1954 subsetDataIds.findDatasets(
1955 bias, collections=["imported_r", "imported_g"], findFirst=False
1956 )
1957 ),
1958 expectedAllBiases,
1959 )
1960 self.assertCountEqual(
1961 list(
1962 subsetDataIds.findDatasets(
1963 bias, collections=["imported_r", "imported_g"], findFirst=True
1964 )
1965 ),
1966 expectedDeduplicatedBiases,
1967 )
1969 def testStorageClassPropagation(self):
1970 """Test that queries for datasets respect the storage class passed in
1971 as part of a full dataset type.
1972 """
1973 butler = self.make_butler()
1974 registry = butler.registry
1975 self.load_data(butler, "base.yaml")
1976 dataset_type_in_registry = DatasetType(
1977 "tbl", dimensions=["instrument"], storageClass="Packages", universe=registry.dimensions
1978 )
1979 registry.registerDatasetType(dataset_type_in_registry)
1980 run = "run1"
1981 registry.registerRun(run)
1982 (inserted_ref,) = registry.insertDatasets(
1983 dataset_type_in_registry, [registry.expandDataId(instrument="Cam1")], run=run
1984 )
1985 self.assertEqual(inserted_ref.datasetType, dataset_type_in_registry)
1986 query_dataset_type = DatasetType(
1987 "tbl", dimensions=["instrument"], storageClass="StructuredDataDict", universe=registry.dimensions
1988 )
1989 self.assertNotEqual(dataset_type_in_registry, query_dataset_type)
1990 query_datasets_result = registry.queryDatasets(query_dataset_type, collections=[run])
1991 self.assertEqual(query_datasets_result.parentDatasetType, query_dataset_type) # type: ignore
1992 (query_datasets_ref,) = query_datasets_result
1993 self.assertEqual(query_datasets_ref.datasetType, query_dataset_type)
1994 query_data_ids_find_datasets_result = registry.queryDataIds(["instrument"]).findDatasets(
1995 query_dataset_type, collections=[run]
1996 )
1997 self.assertEqual(query_data_ids_find_datasets_result.parentDatasetType, query_dataset_type)
1998 (query_data_ids_find_datasets_ref,) = query_data_ids_find_datasets_result
1999 self.assertEqual(query_data_ids_find_datasets_ref.datasetType, query_dataset_type)
2000 query_dataset_types_result = registry.queryDatasetTypes(query_dataset_type)
2001 self.assertEqual(list(query_dataset_types_result), [query_dataset_type])
2002 find_dataset_ref = registry.findDataset(query_dataset_type, instrument="Cam1", collections=[run])
2003 self.assertEqual(find_dataset_ref.datasetType, query_dataset_type)
2005 def testEmptyDimensionsQueries(self):
2006 """Test Query and QueryResults objects in the case where there are no
2007 dimensions.
2008 """
2009 # Set up test data: one dataset type, two runs, one dataset in each.
2010 butler = self.make_butler()
2011 registry = butler.registry
2012 self.load_data(butler, "base.yaml")
2013 schema = DatasetType("schema", dimensions=registry.dimensions.empty, storageClass="Catalog")
2014 registry.registerDatasetType(schema)
2015 dataId = DataCoordinate.make_empty(registry.dimensions)
2016 run1 = "run1"
2017 run2 = "run2"
2018 registry.registerRun(run1)
2019 registry.registerRun(run2)
2020 (dataset1,) = registry.insertDatasets(schema, dataIds=[dataId], run=run1)
2021 (dataset2,) = registry.insertDatasets(schema, dataIds=[dataId], run=run2)
2022 # Query directly for both of the datasets, and each one, one at a time.
2023 self.checkQueryResults(
2024 registry.queryDatasets(schema, collections=[run1, run2], findFirst=False), [dataset1, dataset2]
2025 )
2026 self.checkQueryResults(
2027 registry.queryDatasets(schema, collections=[run1, run2], findFirst=True),
2028 [dataset1],
2029 )
2030 self.checkQueryResults(
2031 registry.queryDatasets(schema, collections=[run2, run1], findFirst=True),
2032 [dataset2],
2033 )
2034 # Query for data IDs with no dimensions.
2035 dataIds = registry.queryDataIds([])
2036 self.checkQueryResults(dataIds, [dataId])
2037 # Use queried data IDs to find the datasets.
2038 self.checkQueryResults(
2039 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),
2040 [dataset1, dataset2],
2041 )
2042 self.checkQueryResults(
2043 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2044 [dataset1],
2045 )
2046 self.checkQueryResults(
2047 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2048 [dataset2],
2049 )
2050 # Now materialize the data ID query results and repeat those tests.
2051 with dataIds.materialize() as dataIds:
2052 self.checkQueryResults(dataIds, [dataId])
2053 self.checkQueryResults(
2054 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2055 [dataset1],
2056 )
2057 self.checkQueryResults(
2058 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2059 [dataset2],
2060 )
2061 # Query for non-empty data IDs, then subset that to get the empty one.
2062 # Repeat the above tests starting from that.
2063 dataIds = registry.queryDataIds(["instrument"]).subset(registry.dimensions.empty, unique=True)
2064 self.checkQueryResults(dataIds, [dataId])
2065 self.checkQueryResults(
2066 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),
2067 [dataset1, dataset2],
2068 )
2069 self.checkQueryResults(
2070 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2071 [dataset1],
2072 )
2073 self.checkQueryResults(
2074 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2075 [dataset2],
2076 )
2077 with dataIds.materialize() as dataIds:
2078 self.checkQueryResults(dataIds, [dataId])
2079 self.checkQueryResults(
2080 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),
2081 [dataset1, dataset2],
2082 )
2083 self.checkQueryResults(
2084 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2085 [dataset1],
2086 )
2087 self.checkQueryResults(
2088 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2089 [dataset2],
2090 )
2091 # Query for non-empty data IDs, then materialize, then subset to get
2092 # the empty one. Repeat again.
2093 with registry.queryDataIds(["instrument"]).materialize() as nonEmptyDataIds:
2094 dataIds = nonEmptyDataIds.subset(registry.dimensions.empty, unique=True)
2095 self.checkQueryResults(dataIds, [dataId])
2096 self.checkQueryResults(
2097 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),
2098 [dataset1, dataset2],
2099 )
2100 self.checkQueryResults(
2101 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2102 [dataset1],
2103 )
2104 self.checkQueryResults(
2105 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2106 [dataset2],
2107 )
2108 with dataIds.materialize() as dataIds:
2109 self.checkQueryResults(dataIds, [dataId])
2110 self.checkQueryResults(
2111 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=False),
2112 [dataset1, dataset2],
2113 )
2114 self.checkQueryResults(
2115 dataIds.findDatasets(schema, collections=[run1, run2], findFirst=True),
2116 [dataset1],
2117 )
2118 self.checkQueryResults(
2119 dataIds.findDatasets(schema, collections=[run2, run1], findFirst=True),
2120 [dataset2],
2121 )
2122 # Query for non-empty data IDs with a constraint on an empty-data-ID
2123 # dataset that exists.
2124 dataIds = registry.queryDataIds(["instrument"], datasets="schema", collections=...)
2125 self.checkQueryResults(
2126 dataIds.subset(unique=True),
2127 [DataCoordinate.standardize(instrument="Cam1", universe=registry.dimensions)],
2128 )
2129 # Again query for non-empty data IDs with a constraint on empty-data-ID
2130 # datasets, but when the datasets don't exist. We delete the existing
2131 # dataset and query just that collection rather than creating a new
2132 # empty collection because this is a bit less likely for our build-time
2133 # logic to shortcut-out (via the collection summaries), and such a
2134 # shortcut would make this test a bit more trivial than we'd like.
2135 registry.removeDatasets([dataset2])
2136 dataIds = registry.queryDataIds(["instrument"], datasets="schema", collections=run2)
2137 self.checkQueryResults(dataIds, [])
2139 def testDimensionDataModifications(self):
2140 """Test that modifying dimension records via:
2141 syncDimensionData(..., update=True) and
2142 insertDimensionData(..., replace=True) works as expected, even in the
2143 presence of datasets using those dimensions and spatial overlap
2144 relationships.
2145 """
2147 def _unpack_range_set(ranges: lsst.sphgeom.RangeSet) -> Iterator[int]:
2148 """Unpack a sphgeom.RangeSet into the integers it contains."""
2149 for begin, end in ranges:
2150 yield from range(begin, end)
2152 def _range_set_hull(
2153 ranges: lsst.sphgeom.RangeSet,
2154 pixelization: lsst.sphgeom.HtmPixelization,
2155 ) -> lsst.sphgeom.ConvexPolygon:
2156 """Create a ConvexPolygon hull of the region defined by a set of
2157 HTM pixelization index ranges.
2158 """
2159 points = []
2160 for index in _unpack_range_set(ranges):
2161 points.extend(pixelization.triangle(index).getVertices())
2162 return lsst.sphgeom.ConvexPolygon(points)
2164 # Use HTM to set up an initial parent region (one arbitrary trixel)
2165 # and four child regions (the trixels within the parent at the next
2166 # level. We'll use the parent as a tract/visit region and the children
2167 # as its patch/visit_detector regions.
2168 butler = self.make_butler()
2169 registry = butler.registry
2170 htm6 = registry.dimensions.skypix["htm"][6].pixelization
2171 commonSkyPix = registry.dimensions.commonSkyPix.pixelization
2172 index = 12288
2173 child_ranges_small = lsst.sphgeom.RangeSet(index).scaled(4)
2174 assert htm6.universe().contains(child_ranges_small)
2175 child_regions_small = [htm6.triangle(i) for i in _unpack_range_set(child_ranges_small)]
2176 parent_region_small = lsst.sphgeom.ConvexPolygon(
2177 list(itertools.chain.from_iterable(c.getVertices() for c in child_regions_small))
2178 )
2179 assert all(parent_region_small.contains(c) for c in child_regions_small)
2180 # Make a larger version of each child region, defined to be the set of
2181 # htm6 trixels that overlap the original's bounding circle. Make a new
2182 # parent that's the convex hull of the new children.
2183 child_regions_large = [
2184 _range_set_hull(htm6.envelope(c.getBoundingCircle()), htm6) for c in child_regions_small
2185 ]
2186 assert all(
2187 large.contains(small)
2188 for large, small in zip(child_regions_large, child_regions_small, strict=True)
2189 )
2190 parent_region_large = lsst.sphgeom.ConvexPolygon(
2191 list(itertools.chain.from_iterable(c.getVertices() for c in child_regions_large))
2192 )
2193 assert all(parent_region_large.contains(c) for c in child_regions_large)
2194 assert parent_region_large.contains(parent_region_small)
2195 assert not parent_region_small.contains(parent_region_large)
2196 assert not all(parent_region_small.contains(c) for c in child_regions_large)
2197 # Find some commonSkyPix indices that overlap the large regions but not
2198 # overlap the small regions. We use commonSkyPix here to make sure the
2199 # real tests later involve what's in the database, not just post-query
2200 # filtering of regions.
2201 child_difference_indices = []
2202 for large, small in zip(child_regions_large, child_regions_small, strict=True):
2203 difference = list(_unpack_range_set(commonSkyPix.envelope(large) - commonSkyPix.envelope(small)))
2204 assert difference, "if this is empty, we can't test anything useful with these regions"
2205 assert all(
2206 not commonSkyPix.triangle(d).isDisjointFrom(large)
2207 and commonSkyPix.triangle(d).isDisjointFrom(small)
2208 for d in difference
2209 )
2210 child_difference_indices.append(difference)
2211 parent_difference_indices = list(
2212 _unpack_range_set(
2213 commonSkyPix.envelope(parent_region_large) - commonSkyPix.envelope(parent_region_small)
2214 )
2215 )
2216 assert parent_difference_indices, "if this is empty, we can't test anything useful with these regions"
2217 assert all(
2218 (
2219 not commonSkyPix.triangle(d).isDisjointFrom(parent_region_large)
2220 and commonSkyPix.triangle(d).isDisjointFrom(parent_region_small)
2221 )
2222 for d in parent_difference_indices
2223 )
2224 # Now that we've finally got those regions, we'll insert the large ones
2225 # as tract/patch dimension records.
2226 skymap_name = "testing_v1"
2227 registry.insertDimensionData(
2228 "skymap",
2229 {
2230 "name": skymap_name,
2231 "hash": bytes([42]),
2232 "tract_max": 1,
2233 "patch_nx_max": 2,
2234 "patch_ny_max": 2,
2235 },
2236 )
2237 registry.insertDimensionData("tract", {"skymap": skymap_name, "id": 0, "region": parent_region_large})
2238 registry.insertDimensionData(
2239 "patch",
2240 *[
2241 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}
2242 for n, c in enumerate(child_regions_large)
2243 ],
2244 )
2245 # Add at dataset that uses these dimensions to make sure that modifying
2246 # them doesn't disrupt foreign keys (need to make sure DB doesn't
2247 # implement insert with replace=True as delete-then-insert).
2248 dataset_type = DatasetType(
2249 "coadd",
2250 dimensions=["tract", "patch"],
2251 universe=registry.dimensions,
2252 storageClass="Exposure",
2253 )
2254 registry.registerDatasetType(dataset_type)
2255 registry.registerCollection("the_run", CollectionType.RUN)
2256 registry.insertDatasets(
2257 dataset_type,
2258 [{"skymap": skymap_name, "tract": 0, "patch": 2}],
2259 run="the_run",
2260 )
2261 # Query for tracts and patches that overlap some "difference" htm9
2262 # pixels; there should be overlaps, because the database has
2263 # the "large" suite of regions.
2264 self.assertEqual(
2265 {0},
2266 {
2267 data_id["tract"]
2268 for data_id in registry.queryDataIds(
2269 ["tract"],
2270 skymap=skymap_name,
2271 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},
2272 )
2273 },
2274 )
2275 for patch_id, patch_difference_indices in enumerate(child_difference_indices):
2276 self.assertIn(
2277 patch_id,
2278 {
2279 data_id["patch"]
2280 for data_id in registry.queryDataIds(
2281 ["patch"],
2282 skymap=skymap_name,
2283 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},
2284 )
2285 },
2286 )
2287 # Use sync to update the tract region and insert to update the regions
2288 # of the patches, to the "small" suite.
2289 updated = registry.syncDimensionData(
2290 "tract",
2291 {"skymap": skymap_name, "id": 0, "region": parent_region_small},
2292 update=True,
2293 )
2294 self.assertEqual(updated, {"region": parent_region_large})
2295 registry.insertDimensionData(
2296 "patch",
2297 *[
2298 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}
2299 for n, c in enumerate(child_regions_small)
2300 ],
2301 replace=True,
2302 )
2303 # Query again; there now should be no such overlaps, because the
2304 # database has the "small" suite of regions.
2305 self.assertFalse(
2306 set(
2307 registry.queryDataIds(
2308 ["tract"],
2309 skymap=skymap_name,
2310 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},
2311 )
2312 )
2313 )
2314 for patch_id, patch_difference_indices in enumerate(child_difference_indices):
2315 self.assertNotIn(
2316 patch_id,
2317 {
2318 data_id["patch"]
2319 for data_id in registry.queryDataIds(
2320 ["patch"],
2321 skymap=skymap_name,
2322 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},
2323 )
2324 },
2325 )
2326 # Update back to the large regions and query one more time.
2327 updated = registry.syncDimensionData(
2328 "tract",
2329 {"skymap": skymap_name, "id": 0, "region": parent_region_large},
2330 update=True,
2331 )
2332 self.assertEqual(updated, {"region": parent_region_small})
2333 registry.insertDimensionData(
2334 "patch",
2335 *[
2336 {"skymap": skymap_name, "tract": 0, "id": n, "cell_x": n % 2, "cell_y": n // 2, "region": c}
2337 for n, c in enumerate(child_regions_large)
2338 ],
2339 replace=True,
2340 )
2341 self.assertEqual(
2342 {0},
2343 {
2344 data_id["tract"]
2345 for data_id in registry.queryDataIds(
2346 ["tract"],
2347 skymap=skymap_name,
2348 dataId={registry.dimensions.commonSkyPix.name: parent_difference_indices[0]},
2349 )
2350 },
2351 )
2352 for patch_id, patch_difference_indices in enumerate(child_difference_indices):
2353 self.assertIn(
2354 patch_id,
2355 {
2356 data_id["patch"]
2357 for data_id in registry.queryDataIds(
2358 ["patch"],
2359 skymap=skymap_name,
2360 dataId={registry.dimensions.commonSkyPix.name: patch_difference_indices[0]},
2361 )
2362 },
2363 )
2365 def testCalibrationCollections(self):
2366 """Test operations on `~CollectionType.CALIBRATION` collections,
2367 including `SqlRegistry.certify`, `SqlRegistry.decertify`,
2368 `SqlRegistry.findDataset`, and
2369 `DataCoordinateQueryResults.findRelatedDatasets`.
2370 """
2371 # Setup - make a Registry, fill it with some datasets in
2372 # non-calibration collections.
2373 butler = self.make_butler()
2374 registry = butler.registry
2375 self.load_data(butler, "base.yaml", "datasets.yaml")
2376 # Set up some timestamps.
2377 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")
2378 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai")
2379 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai")
2380 t4 = astropy.time.Time("2020-01-01T04:00:00", format="isot", scale="tai")
2381 t5 = astropy.time.Time("2020-01-01T05:00:00", format="isot", scale="tai")
2382 allTimespans = [
2383 Timespan(a, b) for a, b in itertools.combinations([None, t1, t2, t3, t4, t5, None], r=2)
2384 ]
2385 # Insert some exposure records with timespans between each sequential
2386 # pair of those.
2387 registry.insertDimensionData(
2388 "day_obs", {"instrument": "Cam1", "id": 20200101, "timespan": Timespan(t1, t5)}
2389 )
2390 registry.insertDimensionData(
2391 "group",
2392 {"instrument": "Cam1", "name": "group0"},
2393 {"instrument": "Cam1", "name": "group1"},
2394 {"instrument": "Cam1", "name": "group2"},
2395 {"instrument": "Cam1", "name": "group3"},
2396 )
2397 registry.insertDimensionData(
2398 "exposure",
2399 {
2400 "instrument": "Cam1",
2401 "id": 0,
2402 "group": "group0",
2403 "obs_id": "zero",
2404 "physical_filter": "Cam1-G",
2405 "day_obs": 20200101,
2406 "timespan": Timespan(t1, t2),
2407 },
2408 {
2409 "instrument": "Cam1",
2410 "id": 1,
2411 "group": "group1",
2412 "obs_id": "one",
2413 "physical_filter": "Cam1-G",
2414 "day_obs": 20200101,
2415 "timespan": Timespan(t2, t3),
2416 },
2417 {
2418 "instrument": "Cam1",
2419 "id": 2,
2420 "group": "group2",
2421 "obs_id": "two",
2422 "physical_filter": "Cam1-G",
2423 "day_obs": 20200101,
2424 "timespan": Timespan(t3, t4),
2425 },
2426 {
2427 "instrument": "Cam1",
2428 "id": 3,
2429 "group": "group3",
2430 "obs_id": "three",
2431 "physical_filter": "Cam1-G",
2432 "day_obs": 20200101,
2433 "timespan": Timespan(t4, t5),
2434 },
2435 )
2436 # Get references to some datasets.
2437 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g")
2438 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g")
2439 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r")
2440 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r")
2441 # Register the main calibration collection we'll be working with.
2442 collection = "Cam1/calibs/default"
2443 registry.registerCollection(collection, type=CollectionType.CALIBRATION)
2444 # Cannot associate into a calibration collection (no timespan).
2445 with self.assertRaises(CollectionTypeError):
2446 registry.associate(collection, [bias2a])
2447 # Certify 2a dataset with [t2, t4) validity.
2448 registry.certify(collection, [bias2a], Timespan(begin=t2, end=t4))
2449 # Test that we can query for this dataset via the new collection, both
2450 # on its own and with a RUN collection.
2451 self.assertEqual(
2452 set(registry.queryDatasets("bias", findFirst=False, collections=collection)),
2453 {bias2a},
2454 )
2455 self.assertEqual(
2456 set(registry.queryDatasets("bias", findFirst=False, collections=[collection, "imported_r"])),
2457 {
2458 bias2a,
2459 bias2b,
2460 bias3b,
2461 registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),
2462 },
2463 )
2464 self.assertEqual(
2465 set(registry.queryDataIds("detector", datasets="bias", collections=collection)),
2466 {registry.expandDataId(instrument="Cam1", detector=2)},
2467 )
2468 self.assertEqual(
2469 set(registry.queryDataIds("detector", datasets="bias", collections=[collection, "imported_r"])),
2470 {
2471 registry.expandDataId(instrument="Cam1", detector=2),
2472 registry.expandDataId(instrument="Cam1", detector=3),
2473 registry.expandDataId(instrument="Cam1", detector=4),
2474 },
2475 )
2476 self.assertEqual(
2477 set(
2478 registry.queryDataIds(["exposure", "detector"]).findRelatedDatasets(
2479 "bias", findFirst=True, collections=[collection]
2480 )
2481 ),
2482 {
2483 (registry.expandDataId(instrument="Cam1", detector=2, exposure=1), bias2a),
2484 (registry.expandDataId(instrument="Cam1", detector=2, exposure=2), bias2a),
2485 },
2486 )
2488 # We should not be able to certify 2b with anything overlapping that
2489 # window.
2490 with self.assertRaises(ConflictingDefinitionError):
2491 registry.certify(collection, [bias2b], Timespan(begin=None, end=t3))
2492 with self.assertRaises(ConflictingDefinitionError):
2493 registry.certify(collection, [bias2b], Timespan(begin=None, end=t5))
2494 with self.assertRaises(ConflictingDefinitionError):
2495 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t3))
2496 with self.assertRaises(ConflictingDefinitionError):
2497 registry.certify(collection, [bias2b], Timespan(begin=t1, end=t5))
2498 with self.assertRaises(ConflictingDefinitionError):
2499 registry.certify(collection, [bias2b], Timespan(begin=t1, end=None))
2500 with self.assertRaises(ConflictingDefinitionError):
2501 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t3))
2502 with self.assertRaises(ConflictingDefinitionError):
2503 registry.certify(collection, [bias2b], Timespan(begin=t2, end=t5))
2504 with self.assertRaises(ConflictingDefinitionError):
2505 registry.certify(collection, [bias2b], Timespan(begin=t2, end=None))
2506 # We should be able to certify 3a with a range overlapping that window,
2507 # because it's for a different detector.
2508 # We'll certify 3a over [t1, t3).
2509 registry.certify(collection, [bias3a], Timespan(begin=t1, end=t3))
2510 # Now we'll certify 2b and 3b together over [t4, ∞).
2511 registry.certify(collection, [bias2b, bias3b], Timespan(begin=t4, end=None))
2513 # Fetch all associations and check that they are what we expect.
2514 self.assertCountEqual(
2515 list(
2516 registry.queryDatasetAssociations(
2517 "bias",
2518 collections=[collection, "imported_g", "imported_r"],
2519 )
2520 ),
2521 [
2522 DatasetAssociation(
2523 ref=registry.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g"),
2524 collection="imported_g",
2525 timespan=None,
2526 ),
2527 DatasetAssociation(
2528 ref=registry.findDataset("bias", instrument="Cam1", detector=4, collections="imported_r"),
2529 collection="imported_r",
2530 timespan=None,
2531 ),
2532 DatasetAssociation(ref=bias2a, collection="imported_g", timespan=None),
2533 DatasetAssociation(ref=bias3a, collection="imported_g", timespan=None),
2534 DatasetAssociation(ref=bias2b, collection="imported_r", timespan=None),
2535 DatasetAssociation(ref=bias3b, collection="imported_r", timespan=None),
2536 DatasetAssociation(ref=bias2a, collection=collection, timespan=Timespan(begin=t2, end=t4)),
2537 DatasetAssociation(ref=bias3a, collection=collection, timespan=Timespan(begin=t1, end=t3)),
2538 DatasetAssociation(ref=bias2b, collection=collection, timespan=Timespan(begin=t4, end=None)),
2539 DatasetAssociation(ref=bias3b, collection=collection, timespan=Timespan(begin=t4, end=None)),
2540 ],
2541 )
2543 class Ambiguous:
2544 """Tag class to denote lookups that should be ambiguous."""
2546 pass
2548 def _assertLookup(
2549 detector: int, timespan: Timespan, expected: DatasetRef | type[Ambiguous] | None
2550 ) -> None:
2551 """Local function that asserts that a bias lookup returns the given
2552 expected result.
2553 """
2554 if expected is Ambiguous:
2555 with self.assertRaises((DatasetTypeError, LookupError)):
2556 registry.findDataset(
2557 "bias",
2558 collections=collection,
2559 instrument="Cam1",
2560 detector=detector,
2561 timespan=timespan,
2562 )
2563 else:
2564 self.assertEqual(
2565 expected,
2566 registry.findDataset(
2567 "bias",
2568 collections=collection,
2569 instrument="Cam1",
2570 detector=detector,
2571 timespan=timespan,
2572 ),
2573 )
2575 # Systematically test lookups against expected results.
2576 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)
2577 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)
2578 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)
2579 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)
2580 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=Ambiguous)
2581 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)
2582 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)
2583 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)
2584 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)
2585 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=Ambiguous)
2586 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)
2587 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)
2588 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)
2589 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=Ambiguous)
2590 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)
2591 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=bias2a)
2592 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=Ambiguous)
2593 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=Ambiguous)
2594 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=bias2b)
2595 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)
2596 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)
2597 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)
2598 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)
2599 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)
2600 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)
2601 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=Ambiguous)
2602 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)
2603 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)
2604 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)
2605 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)
2606 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=Ambiguous)
2607 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)
2608 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)
2609 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)
2610 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=Ambiguous)
2611 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)
2612 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)
2613 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=bias3b)
2614 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)
2615 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=bias3b)
2616 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)
2617 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)
2619 # Test lookups via temporal joins to exposures.
2620 self.assertEqual(
2621 set(
2622 registry.queryDataIds(
2623 ["exposure", "detector"], instrument="Cam1", detector=2
2624 ).findRelatedDatasets("bias", collections=[collection])
2625 ),
2626 {
2627 (registry.expandDataId(instrument="Cam1", exposure=1, detector=2), bias2a),
2628 (registry.expandDataId(instrument="Cam1", exposure=2, detector=2), bias2a),
2629 (registry.expandDataId(instrument="Cam1", exposure=3, detector=2), bias2b),
2630 },
2631 )
2632 self.assertEqual(
2633 set(
2634 registry.queryDataIds(
2635 ["exposure", "detector"], instrument="Cam1", detector=3
2636 ).findRelatedDatasets("bias", collections=[collection])
2637 ),
2638 {
2639 (registry.expandDataId(instrument="Cam1", exposure=0, detector=3), bias3a),
2640 (registry.expandDataId(instrument="Cam1", exposure=1, detector=3), bias3a),
2641 (registry.expandDataId(instrument="Cam1", exposure=3, detector=3), bias3b),
2642 },
2643 )
2645 # Decertify [t3, t5) for all data IDs, and do test lookups again.
2646 # This should truncate bias2a to [t2, t3), leave bias3a unchanged at
2647 # [t1, t3), and truncate bias2b and bias3b to [t5, ∞).
2648 registry.decertify(collection=collection, datasetType="bias", timespan=Timespan(t3, t5))
2649 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)
2650 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)
2651 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)
2652 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)
2653 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=bias2a)
2654 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)
2655 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)
2656 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)
2657 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)
2658 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=bias2a)
2659 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)
2660 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)
2661 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)
2662 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=bias2a)
2663 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)
2664 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=None)
2665 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=None)
2666 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=bias2b)
2667 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=None)
2668 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)
2669 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)
2670 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)
2671 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)
2672 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)
2673 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)
2674 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=bias3a)
2675 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)
2676 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)
2677 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)
2678 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)
2679 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=bias3a)
2680 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)
2681 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)
2682 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)
2683 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=bias3a)
2684 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)
2685 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)
2686 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=None)
2687 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)
2688 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=None)
2689 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)
2690 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)
2692 # Decertify everything, this time with explicit data IDs, then check
2693 # that no lookups succeed.
2694 registry.decertify(
2695 collection,
2696 "bias",
2697 Timespan(None, None),
2698 dataIds=[
2699 dict(instrument="Cam1", detector=2),
2700 dict(instrument="Cam1", detector=3),
2701 ],
2702 )
2703 for detector in (2, 3):
2704 for timespan in allTimespans:
2705 _assertLookup(detector=detector, timespan=timespan, expected=None)
2706 # Certify bias2a and bias3a over (-∞, ∞), check that all lookups return
2707 # those.
2708 registry.certify(
2709 collection,
2710 [bias2a, bias3a],
2711 Timespan(None, None),
2712 )
2713 for timespan in allTimespans:
2714 _assertLookup(detector=2, timespan=timespan, expected=bias2a)
2715 _assertLookup(detector=3, timespan=timespan, expected=bias3a)
2716 # Decertify just bias2 over [t2, t4).
2717 # This should split a single certification row into two (and leave the
2718 # other existing row, for bias3a, alone).
2719 registry.decertify(
2720 collection, "bias", Timespan(t2, t4), dataIds=[dict(instrument="Cam1", detector=2)]
2721 )
2722 for timespan in allTimespans:
2723 _assertLookup(detector=3, timespan=timespan, expected=bias3a)
2724 overlapsBefore = timespan.overlaps(Timespan(None, t2))
2725 overlapsAfter = timespan.overlaps(Timespan(t4, None))
2726 if overlapsBefore and overlapsAfter:
2727 expected = Ambiguous
2728 elif overlapsBefore or overlapsAfter:
2729 expected = bias2a
2730 else:
2731 expected = None
2732 _assertLookup(detector=2, timespan=timespan, expected=expected)
2734 def testSkipCalibs(self):
2735 """Test how queries handle skipping of calibration collections."""
2736 butler = self.make_butler()
2737 registry = butler.registry
2738 self.load_data(butler, "base.yaml", "datasets.yaml")
2740 coll_calib = "Cam1/calibs/default"
2741 registry.registerCollection(coll_calib, type=CollectionType.CALIBRATION)
2743 # Add all biases to the calibration collection.
2744 # Without this, the logic that prunes dataset subqueries based on
2745 # datasetType-collection summary information will fire before the logic
2746 # we want to test below. This is a good thing (it avoids the dreaded
2747 # NotImplementedError a bit more often) everywhere but here.
2748 registry.certify(coll_calib, registry.queryDatasets("bias", collections=...), Timespan(None, None))
2750 coll_list = [coll_calib, "imported_g", "imported_r"]
2751 chain = "Cam1/chain"
2752 registry.registerCollection(chain, type=CollectionType.CHAINED)
2753 registry.setCollectionChain(chain, coll_list)
2755 # Lookup is ambiguous due to multiple datasets with the same data ID
2756 # in the calibration collection.
2757 with self.assertRaises(CalibrationLookupError):
2758 list(registry.queryDatasets("bias", collections=coll_list, findFirst=True))
2760 # chain will skip
2761 datasets = list(registry.queryDatasets("bias", collections=chain))
2762 self.assertGreater(len(datasets), 0)
2764 dataIds = list(registry.queryDataIds(["instrument", "detector"], datasets="bias", collections=chain))
2765 self.assertGreater(len(dataIds), 0)
2767 # glob will skip too
2768 datasets = list(registry.queryDatasets("bias", collections="*d*"))
2769 self.assertGreater(len(datasets), 0)
2771 # regular expression will skip too
2772 if self.supportsCollectionRegex:
2773 pattern = re.compile(".*")
2774 with self.assertWarns(FutureWarning):
2775 datasets = list(registry.queryDatasets("bias", collections=pattern))
2776 self.assertGreater(len(datasets), 0)
2778 # ellipsis should work as usual
2779 datasets = list(registry.queryDatasets("bias", collections=...))
2780 self.assertGreater(len(datasets), 0)
2782 # New query system correctly determines that this search is
2783 # ambiguous, because there are multiple datasets with the same
2784 # {instrument=Cam1, detector=2} data ID in the calibration
2785 # collection at the beginning of the chain.
2786 with self.assertRaises(CalibrationLookupError):
2787 datasets = list(registry.queryDatasets("bias", collections=chain, findFirst=True))
2789 def testIngestTimeQuery(self):
2790 butler = self.make_butler()
2791 registry = butler.registry
2792 dt0 = datetime.datetime.now(datetime.UTC)
2793 self.load_data(butler, "base.yaml", "datasets.yaml")
2794 dt1 = datetime.datetime.now(datetime.UTC)
2796 datasets = list(registry.queryDatasets(..., collections=...))
2797 len0 = len(datasets)
2798 self.assertGreater(len0, 0)
2800 for where in ("ingest_date > T'2000-01-01'", "T'2000-01-01' < ingest_date"):
2801 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2802 len1 = len(datasets)
2803 self.assertEqual(len0, len1)
2805 # no one will ever use this piece of software in 30 years
2806 for where in ("ingest_date > T'2050-01-01'", "T'2050-01-01' < ingest_date"):
2807 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2808 len2 = len(datasets)
2809 self.assertEqual(len2, 0)
2811 # Check more exact timing to make sure there is no 37 seconds offset
2812 # (after fixing DM-30124). SQLite time precision is 1 second, make
2813 # sure that we don't test with higher precision.
2814 tests = [
2815 # format: (timestamp, operator, expected_len)
2816 (dt0 - timedelta(seconds=1), ">", len0),
2817 (dt0 - timedelta(seconds=1), "<", 0),
2818 (dt1 + timedelta(seconds=1), "<", len0),
2819 (dt1 + timedelta(seconds=1), ">", 0),
2820 ]
2821 for dt, op, expect_len in tests:
2822 dt_str = dt.isoformat(sep=" ")
2824 where = f"ingest_date {op} T'{dt_str}'"
2825 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2826 self.assertEqual(len(datasets), expect_len)
2828 # same with bind using datetime or astropy Time
2829 where = f"ingest_date {op} :ingest_time"
2830 datasets = list(
2831 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt})
2832 )
2833 self.assertEqual(len(datasets), expect_len)
2835 dt_astropy = astropy.time.Time(dt, format="datetime")
2836 datasets = list(
2837 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt_astropy})
2838 )
2839 self.assertEqual(len(datasets), expect_len)
2841 def testTimespanQueries(self):
2842 """Test query expressions involving timespans."""
2843 butler = self.make_butler()
2844 registry = butler.registry
2845 self.load_data(butler, "ci_hsc-subset.yaml")
2846 # All exposures in the database; mapping from ID to timespan.
2847 visits = {record.id: record.timespan for record in registry.queryDimensionRecords("visit")}
2848 # Just those IDs, sorted (which is also temporal sorting, because HSC
2849 # exposure IDs are monotonically increasing).
2850 ids = sorted(visits.keys())
2851 self.assertEqual(len(ids), 11)
2852 # Pick some quasi-random indexes into `ids` to play with.
2853 i1 = int(len(ids) * 0.1)
2854 i2 = int(len(ids) * 0.3)
2855 i3 = int(len(ids) * 0.6)
2856 i4 = int(len(ids) * 0.8)
2857 # Extract some times from those: just before the beginning of i1 (which
2858 # should be after the end of the exposure before), exactly the
2859 # beginning of i2, just after the beginning of i3 (and before its end),
2860 # and the exact end of i4.
2861 t1 = visits[ids[i1]].begin - astropy.time.TimeDelta(1.0, format="sec")
2862 self.assertGreater(t1, visits[ids[i1 - 1]].end)
2863 t2 = visits[ids[i2]].begin
2864 t3 = visits[ids[i3]].begin + astropy.time.TimeDelta(1.0, format="sec")
2865 self.assertLess(t3, visits[ids[i3]].end)
2866 t4 = visits[ids[i4]].end
2867 # Make sure those are actually in order.
2868 self.assertEqual([t1, t2, t3, t4], sorted([t4, t3, t2, t1]))
2870 bind = {
2871 "t1": t1,
2872 "t2": t2,
2873 "t3": t3,
2874 "t4": t4,
2875 "ts23": Timespan(t2, t3),
2876 }
2878 def query(where):
2879 """Return results as a sorted, deduplicated list of visit IDs.
2881 Parameters
2882 ----------
2883 where : `str`
2884 The WHERE clause for the query.
2885 """
2886 return sorted(
2887 {
2888 dataId["visit"]
2889 for dataId in registry.queryDataIds("visit", instrument="HSC", bind=bind, where=where)
2890 }
2891 )
2893 # Try a bunch of timespan queries, mixing up the bounds themselves,
2894 # where they appear in the expression, and how we get the timespan into
2895 # the expression.
2897 # t1 is before the start of i1, so this should not include i1.
2898 self.assertEqual(ids[:i1], query("visit.timespan OVERLAPS (null, :t1)"))
2899 # t2 is exactly at the start of i2, but ends are exclusive, so these
2900 # should not include i2.
2901 self.assertEqual(ids[i1:i2], query("(:t1, :t2) OVERLAPS visit.timespan"))
2902 # t3 is in the middle of i3, so this should include i3.
2903 self.assertEqual(ids[i2 : i3 + 1], query("visit.timespan OVERLAPS :ts23"))
2904 # This one should not include t3 by the same reasoning.
2905 # t4 is exactly at the end of i4, so this should include i4.
2906 self.assertEqual(ids[i3 : i4 + 1], query(f"visit.timespan OVERLAPS (T'{t3.tai.isot}/tai', :t4)"))
2907 # i4's upper bound of t4 is exclusive so this should not include t4.
2908 self.assertEqual(ids[i4 + 1 :], query("visit.timespan OVERLAPS (:t4, NULL)"))
2910 # Now some timespan vs. time scalar queries.
2911 self.assertEqual(ids[i3 : i3 + 1], query("visit.timespan OVERLAPS :t3"))
2912 self.assertEqual(ids[i3 : i3 + 1], query(f"T'{t3.tai.isot}/tai' OVERLAPS visit.timespan"))
2914 # Empty timespans should not overlap anything.
2915 self.assertEqual([], query("visit.timespan OVERLAPS (:t3, :t2)"))
2917 # Make sure that expanded data IDs include the timespans.
2918 results = list(
2919 registry.queryDataIds(["visit"], dataId={"instrument": "HSC", "visit": ids[1]}).expanded()
2920 )
2921 self.assertEqual(len(results), 1)
2922 visit_timespan = visits[ids[1]]
2923 self.assertEqual(results[0].timespan, visit_timespan)
2924 visit_record = results[0].records["visit"]
2925 assert visit_record is not None
2926 self.assertEqual(visit_record.timespan, visit_timespan)
2927 day_obs_record = results[0].records["day_obs"]
2928 assert day_obs_record is not None
2929 self.assertEqual(day_obs_record.id, 20130617)
2930 self.assertEqual(
2931 day_obs_record.timespan,
2932 Timespan(
2933 astropy.time.Time("2013-06-17T00:00:00", scale="tai"),
2934 astropy.time.Time("2013-06-18T00:00:00", scale="tai"),
2935 ),
2936 )
2938 def testCollectionSummaries(self):
2939 """Test recording and retrieval of collection summaries."""
2940 self.maxDiff = None
2941 butler = self.make_butler()
2942 registry = butler.registry
2943 # Importing datasets from yaml should go through the code path where
2944 # we update collection summaries as we insert datasets.
2945 self.load_data(butler, "base.yaml", "datasets.yaml")
2946 flat = registry.getDatasetType("flat")
2947 expected1 = CollectionSummary()
2948 expected1.dataset_types.add(registry.getDatasetType("bias"))
2949 expected1.add_data_ids(
2950 flat, [DataCoordinate.standardize(instrument="Cam1", universe=registry.dimensions)]
2951 )
2952 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)
2953 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)
2954 # Create a chained collection with both of the imported runs; the
2955 # summary should be the same, because it's a union with itself.
2956 chain = "chain"
2957 registry.registerCollection(chain, CollectionType.CHAINED)
2958 registry.setCollectionChain(chain, ["imported_r", "imported_g"])
2959 self.assertEqual(registry.getCollectionSummary(chain), expected1)
2960 # Associate flats only into a tagged collection and a calibration
2961 # collection to check summaries of those.
2962 tag = "tag"
2963 registry.registerCollection(tag, CollectionType.TAGGED)
2964 registry.associate(tag, registry.queryDatasets(flat, collections="imported_g"))
2965 calibs = "calibs"
2966 registry.registerCollection(calibs, CollectionType.CALIBRATION)
2967 registry.certify(
2968 calibs, registry.queryDatasets(flat, collections="imported_g"), timespan=Timespan(None, None)
2969 )
2970 expected2 = expected1.copy()
2971 expected2.dataset_types.discard("bias")
2972 self.assertEqual(registry.getCollectionSummary(tag), expected2)
2973 self.assertEqual(registry.getCollectionSummary(calibs), expected2)
2974 # Explicitly calling SqlRegistry.refresh() should load those same
2975 # summaries, via a totally different code path.
2976 registry.refresh()
2977 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)
2978 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)
2979 self.assertEqual(registry.getCollectionSummary(tag), expected2)
2980 self.assertEqual(registry.getCollectionSummary(calibs), expected2)
2982 def testBindInQueryDatasets(self):
2983 """Test that the bind parameter is correctly forwarded in
2984 queryDatasets recursion.
2985 """
2986 butler = self.make_butler()
2987 registry = butler.registry
2988 # Importing datasets from yaml should go through the code path where
2989 # we update collection summaries as we insert datasets.
2990 self.load_data(butler, "base.yaml", "datasets.yaml")
2991 self.assertEqual(
2992 set(registry.queryDatasets("flat", band="r", collections=...)),
2993 set(
2994 registry.queryDatasets("flat", where="band=:my_band", bind={"my_band": "r"}, collections=...)
2995 ),
2996 )
2998 def testQueryIntRangeExpressions(self):
2999 """Test integer range expressions in ``where`` arguments.
3001 Note that our expressions use inclusive stop values, unlike Python's.
3002 """
3003 butler = self.make_butler()
3004 registry = butler.registry
3005 self.load_data(butler, "base.yaml")
3006 self.assertEqual(
3007 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..2)")),
3008 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 2]},
3009 )
3010 self.assertEqual(
3011 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..4:2)")),
3012 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 3]},
3013 )
3014 self.assertEqual(
3015 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (2..4:2)")),
3016 {registry.expandDataId(instrument="Cam1", detector=n) for n in [2, 4]},
3017 )
3019 def testQueryResultSummaries(self):
3020 """Test summary methods like `count`, `any`, and `explain_no_results`
3021 on `DataCoordinateQueryResults` and `DatasetQueryResults`.
3022 """
3023 butler = self.make_butler()
3024 registry = butler.registry
3025 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3026 # Default test dataset has two collections, each with both flats and
3027 # biases. Add a new collection with only biases.
3028 registry.registerCollection("biases", CollectionType.TAGGED)
3029 registry.associate("biases", registry.queryDatasets("bias", collections=["imported_g"]))
3030 # First query yields two results, and involves no postprocessing.
3031 query1 = registry.queryDataIds(["physical_filter"], band="r")
3032 self.assertTrue(query1.any(execute=False, exact=False))
3033 self.assertTrue(query1.any(execute=True, exact=False))
3034 self.assertTrue(query1.any(execute=True, exact=True))
3035 self.assertEqual(query1.count(exact=False), 2)
3036 self.assertEqual(query1.count(exact=True), 2)
3037 self.assertFalse(list(query1.explain_no_results()))
3038 # Second query should yield no results, which we should see when
3039 # we attempt to expand the data ID.
3040 query2 = registry.queryDataIds(["physical_filter"], band="h")
3041 # There's no execute=False, exact=False test here because the behavior
3042 # not something we want to guarantee in this case (and exact=False
3043 # says either answer is legal).
3044 self.assertFalse(query2.any(execute=True, exact=False))
3045 self.assertFalse(query2.any(execute=True, exact=True))
3046 self.assertEqual(query2.count(exact=False), 0)
3047 self.assertEqual(query2.count(exact=True), 0)
3048 # These queries yield no results due to various problems that can be
3049 # spotted prior to execution, yielding helpful diagnostics.
3050 base_query = registry.queryDataIds(["detector", "physical_filter"])
3051 queries_and_snippets = [
3052 (
3053 # Dataset type name doesn't match any existing dataset types.
3054 registry.queryDatasets("nonexistent", collections=...),
3055 ["nonexistent"],
3056 ),
3057 (
3058 # Dataset type object isn't registered.
3059 registry.queryDatasets(
3060 DatasetType(
3061 "nonexistent",
3062 dimensions=["instrument"],
3063 universe=registry.dimensions,
3064 storageClass="Image",
3065 ),
3066 collections=...,
3067 ),
3068 ["nonexistent"],
3069 ),
3070 (
3071 # No datasets of this type in this collection.
3072 registry.queryDatasets("flat", collections=["biases"]),
3073 ["flat", "biases"],
3074 ),
3075 (
3076 # No datasets of this type in this collection.
3077 base_query.findDatasets("flat", collections=["biases"]),
3078 ["flat", "biases"],
3079 ),
3080 (
3081 # No collections matching at all.
3082 registry.queryDatasets("flat", collections="potato*"),
3083 ["potato"],
3084 ),
3085 ]
3086 with self.assertRaises(MissingDatasetTypeError):
3087 # Dataset type name doesn't match any existing dataset types.
3088 list(registry.queryDataIds(["detector"], datasets=["nonexistent"], collections=...))
3089 with self.assertRaises(MissingDatasetTypeError):
3090 # Dataset type name doesn't match any existing dataset types.
3091 registry.queryDimensionRecords("detector", datasets=["nonexistent"], collections=...).any()
3092 with self.assertRaises(DatasetTypeExpressionError):
3093 # queryDimensionRecords does not allow dataset type wildcards.
3094 registry.queryDimensionRecords("detector", datasets=["f*"], collections=...).any()
3095 for query, snippets in queries_and_snippets:
3096 self.assertFalse(query.any(execute=False, exact=False))
3097 self.assertFalse(query.any(execute=True, exact=False))
3098 self.assertFalse(query.any(execute=True, exact=True))
3099 self.assertEqual(query.count(exact=False), 0)
3100 self.assertEqual(query.count(exact=True), 0)
3101 messages = list(query.explain_no_results())
3102 self.assertTrue(messages)
3103 # Want all expected snippets to appear in at least one message.
3104 self.assertTrue(
3105 any(
3106 all(snippet in message for snippet in snippets) for message in query.explain_no_results()
3107 ),
3108 messages,
3109 )
3111 # Wildcards on dataset types are not permitted in queryDataIds.
3112 with self.assertRaises(DatasetTypeExpressionError):
3113 registry.queryDataIds(["detector"], datasets=re.compile("^nonexistent$"), collections=...)
3115 # This query yields four overlaps in the database, but one is filtered
3116 # out in postprocessing. The count queries aren't accurate because
3117 # they don't account for duplication that happens due to an internal
3118 # join against commonSkyPix.
3119 query3 = registry.queryDataIds(["visit", "tract"], instrument="Cam1", skymap="SkyMap1")
3120 self.assertEqual(
3121 {
3122 DataCoordinate.standardize(
3123 instrument="Cam1",
3124 skymap="SkyMap1",
3125 visit=v,
3126 tract=t,
3127 universe=registry.dimensions,
3128 )
3129 for v, t in [(1, 0), (2, 0), (2, 1)]
3130 },
3131 set(query3),
3132 )
3133 self.assertTrue(query3.any(execute=False, exact=False))
3134 self.assertTrue(query3.any(execute=True, exact=False))
3135 self.assertTrue(query3.any(execute=True, exact=True))
3136 self.assertGreaterEqual(query3.count(exact=False), 3)
3137 self.assertGreaterEqual(query3.count(exact=True, discard=True), 3)
3138 self.assertFalse(list(query3.explain_no_results()))
3139 # This query yields overlaps in the database, but all are filtered
3140 # out in postprocessing. The count queries again aren't very useful.
3141 # We have to use `where=` here to avoid an optimization that
3142 # (currently) skips the spatial postprocess-filtering because it
3143 # recognizes that no spatial join is necessary. That's not ideal, but
3144 # fixing it is out of scope for this ticket.
3145 query4 = registry.queryDataIds(
3146 ["visit", "tract"],
3147 instrument="Cam1",
3148 skymap="SkyMap1",
3149 where="visit=1 AND detector=1 AND tract=0 AND patch=4",
3150 )
3151 self.assertFalse(set(query4))
3152 self.assertTrue(query4.any(execute=False, exact=False))
3153 self.assertTrue(query4.any(execute=True, exact=False))
3154 self.assertFalse(query4.any(execute=True, exact=True))
3155 self.assertGreaterEqual(query4.count(exact=False), 1)
3156 self.assertEqual(query4.count(exact=True, discard=True), 0)
3157 # This query should yield results from one dataset type but not the
3158 # other, which is not registered.
3159 query5 = registry.queryDatasets(["bias", "nonexistent"], collections=["biases"])
3160 self.assertTrue(set(query5))
3161 self.assertTrue(query5.any(execute=False, exact=False))
3162 self.assertTrue(query5.any(execute=True, exact=False))
3163 self.assertTrue(query5.any(execute=True, exact=True))
3164 self.assertGreaterEqual(query5.count(exact=False), 1)
3165 self.assertGreaterEqual(query5.count(exact=True), 1)
3166 # This query applies a selection that yields no results, fully in the
3167 # database. Explaining why it fails involves traversing the relation
3168 # tree and running a LIMIT 1 query at each level that has the potential
3169 # to remove rows.
3170 query6 = registry.queryDimensionRecords(
3171 "detector", where="detector.purpose = 'no-purpose'", instrument="Cam1"
3172 )
3173 self.assertEqual(query6.count(exact=True), 0)
3174 self.assertFalse(query6.any())
3176 def testQueryDataIdsExpressionError(self):
3177 """Test error checking of 'where' expressions in queryDataIds."""
3178 butler = self.make_butler()
3179 registry = butler.registry
3180 self.load_data(butler, "base.yaml")
3181 bind = {"time": astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")}
3182 # The diagnostics raised are slightly different between the old query
3183 # system (ValueError, first error string) and the new query system
3184 # (InvalidQueryError, second error string).
3185 with self.assertRaisesRegex(
3186 (LookupError, InvalidQueryError),
3187 r"(No dimension element with name 'foo' in 'foo\.bar'\.)|(Unrecognized identifier 'foo.bar')",
3188 ):
3189 list(registry.queryDataIds(["detector"], where="foo.bar = 12"))
3190 with self.assertRaisesRegex(
3191 (LookupError, InvalidQueryError),
3192 "(Dimension element name cannot be inferred in this context.)"
3193 "|(Unrecognized identifier 'timespan')",
3194 ):
3195 list(registry.queryDataIds(["detector"], where="timespan.end < :time", bind=bind))
3197 def testQueryDataIdsOrderBy(self):
3198 """Test order_by and limit on result returned by queryDataIds()."""
3199 butler = self.make_butler()
3200 registry = butler.registry
3201 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3203 def do_query(dimensions=("visit", "tract"), datasets=None, collections=None):
3204 return registry.queryDataIds(
3205 dimensions, datasets=datasets, collections=collections, instrument="Cam1", skymap="SkyMap1"
3206 )
3208 Test = namedtuple(
3209 "testQueryDataIdsOrderByTest",
3210 ("order_by", "keys", "result", "limit", "datasets", "collections"),
3211 defaults=(None, None, None),
3212 )
3214 test_data = [
3215 Test("tract,visit", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3216 Test("-tract,visit", "tract,visit", ((1, 2), (0, 1), (0, 2))),
3217 Test("tract,-visit", "tract,visit", ((0, 2), (0, 1), (1, 2))),
3218 Test("-tract,-visit", "tract,visit", ((1, 2), (0, 2), (0, 1))),
3219 Test("tract.id,visit.id", "tract,visit", ((0, 1),), limit=(1,)),
3220 Test("-tract,-visit", "tract,visit", ((1, 2),), limit=(1,)),
3221 Test("tract,visit.exposure_time", "tract,visit", ((0, 2), (0, 1), (1, 2))),
3222 Test("-tract,-visit.exposure_time", "tract,visit", ((1, 2), (0, 1), (0, 2))),
3223 Test("tract,-exposure_time", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3224 Test("tract,visit.name", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3225 Test(
3226 "tract,-visit.timespan.begin,visit.timespan.end",
3227 "tract,visit",
3228 ((0, 2), (0, 1), (1, 2)),
3229 ),
3230 Test("visit.day_obs,exposure.day_obs", "visit,exposure", ()),
3231 Test("visit.timespan.begin,-exposure.timespan.begin", "visit,exposure", ()),
3232 Test(
3233 "tract,detector",
3234 "tract,detector",
3235 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3236 datasets="flat",
3237 collections="imported_r",
3238 ),
3239 Test(
3240 "tract,detector.full_name",
3241 "tract,detector",
3242 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3243 datasets="flat",
3244 collections="imported_r",
3245 ),
3246 Test(
3247 "tract,detector.raft,detector.name_in_raft",
3248 "tract,detector",
3249 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3250 datasets="flat",
3251 collections="imported_r",
3252 ),
3253 ]
3255 for test in test_data:
3256 with self.subTest(test=test):
3257 order_by = test.order_by.split(",")
3258 keys = test.keys.split(",")
3259 query = do_query(keys, test.datasets, test.collections).order_by(*order_by)
3260 if test.limit is not None:
3261 query = query.limit(*test.limit)
3262 dataIds = tuple(tuple(dataId[k] for k in keys) for dataId in query)
3263 self.assertEqual(dataIds, test.result)
3265 # and materialize
3266 query = do_query(keys).order_by(*order_by)
3267 if test.limit is not None:
3268 query = query.limit(*test.limit)
3270 # Test exceptions for errors in a name.
3271 # Many of these raise slightly different diagnostics in the old query
3272 # system (ValueError, first error string) than the new query system
3273 # (InvalidQueryError, second error string).
3274 for order_by in ("", "-"):
3275 with self.assertRaisesRegex((ValueError, InvalidQueryError), "Empty dimension name in ORDER BY"):
3276 list(do_query().order_by(order_by))
3278 for order_by in ("undimension.name", "-undimension.name"):
3279 with self.assertRaisesRegex(
3280 (ValueError, InvalidQueryError),
3281 "(Unknown dimension element 'undimension')|(Unrecognized identifier 'undimension.name')",
3282 ):
3283 list(do_query().order_by(order_by))
3285 for order_by in ("attract", "-attract"):
3286 with self.assertRaisesRegex(
3287 (ValueError, InvalidQueryError),
3288 "(Metadata 'attract' cannot be found in any dimension)|(Unrecognized identifier 'attract')",
3289 ):
3290 list(do_query().order_by(order_by))
3292 with self.assertRaisesRegex(
3293 (ValueError, InvalidQueryError),
3294 "(Metadata 'exposure_time' exists in more than one dimension)"
3295 "|(Ambiguous identifier 'exposure_time' matches multiple fields)",
3296 ):
3297 list(do_query(("exposure", "visit")).order_by("exposure_time"))
3299 with self.assertRaisesRegex(
3300 (ValueError, InvalidQueryError),
3301 r"(Timespan exists in more than one dimension element \(day_obs, exposure, visit\); "
3302 r"qualify timespan with specific dimension name\.)|"
3303 r"(Ambiguous identifier 'timespan' matches multiple fields)",
3304 ):
3305 list(do_query(("exposure", "visit")).order_by("timespan.begin"))
3307 with self.assertRaisesRegex(
3308 (ValueError, InvalidQueryError),
3309 "(Cannot find any temporal dimension element for 'timespan.begin')"
3310 "|(Unrecognized identifier 'timespan')",
3311 ):
3312 list(do_query("tract").order_by("timespan.begin"))
3314 with self.assertRaisesRegex(
3315 (ValueError, InvalidQueryError),
3316 "(Cannot use 'timespan.begin' with non-temporal element)"
3317 "|(Unrecognized field 'timespan' for tract)",
3318 ):
3319 list(do_query("tract").order_by("tract.timespan.begin"))
3321 with self.assertRaisesRegex(
3322 (ValueError, InvalidQueryError),
3323 "(Field 'name' does not exist in 'tract')|(Unrecognized field 'name' for tract.)",
3324 ):
3325 list(do_query("tract").order_by("tract.name"))
3327 with self.assertRaisesRegex(
3328 (ValueError, InvalidQueryError),
3329 r"(Unknown dimension element 'timestamp'; perhaps you meant 'timespan.begin'\?)"
3330 r"|(Unrecognized identifier 'timestamp.begin')",
3331 ):
3332 list(do_query("visit").order_by("timestamp.begin"))
3334 def testQueryDataIdsGovernorExceptions(self):
3335 """Test exceptions raised by queryDataIds() for incorrect governors."""
3336 butler = self.make_butler()
3337 registry = butler.registry
3338 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3340 def do_query(dimensions, dataId=None, where="", bind=None, **kwargs):
3341 return registry.queryDataIds(dimensions, dataId=dataId, where=where, bind=bind, **kwargs)
3343 Test = namedtuple(
3344 "testQueryDataIdExceptionsTest",
3345 ("dimensions", "dataId", "where", "bind", "kwargs", "exception", "count"),
3346 defaults=(None, None, None, {}, None, 0),
3347 )
3349 test_data = (
3350 Test("tract,visit", count=3),
3351 Test("tract,visit", kwargs={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),
3352 Test("tract,visit", kwargs={"instrument": "Cam2", "skymap": "SkyMap1"}, count=0),
3353 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),
3354 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap2"}, count=0),
3355 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap1'", count=3),
3356 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap5'", count=0),
3357 Test(
3358 "tract,visit",
3359 where="instrument=:cam AND skymap=:map",
3360 bind={"cam": "Cam1", "map": "SkyMap1"},
3361 count=3,
3362 ),
3363 Test(
3364 "tract,visit",
3365 where="instrument=:cam AND skymap=:map",
3366 bind={"cam": "Cam", "map": "SkyMap"},
3367 count=0,
3368 ),
3369 )
3371 for test in test_data:
3372 print(test)
3373 dimensions = test.dimensions.split(",")
3374 if test.exception:
3375 with self.assertRaises(test.exception):
3376 with ExitStack() as stack:
3377 if test.exception == DataIdValueError:
3378 stack.enter_context(self.assertWarns(FutureWarning))
3379 do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs).count()
3380 else:
3381 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3382 print(list(query))
3383 self.assertEqual(query.count(discard=True), test.count)
3385 # and materialize
3386 if test.exception:
3387 with self.assertRaises(test.exception):
3388 with ExitStack() as stack:
3389 if test.exception == DataIdValueError:
3390 stack.enter_context(self.assertWarns(FutureWarning))
3391 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3392 with query.materialize() as materialized:
3393 materialized.count(discard=True)
3394 else:
3395 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3396 with query.materialize() as materialized:
3397 self.assertEqual(materialized.count(discard=True), test.count)
3399 def testQueryDimensionRecordsOrderBy(self):
3400 """Test order_by and limit on result returned by
3401 queryDimensionRecords().
3402 """
3403 butler = self.make_butler()
3404 registry = butler.registry
3405 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3407 def do_query(element, datasets=None, collections=None):
3408 return registry.queryDimensionRecords(
3409 element, instrument="Cam1", datasets=datasets, collections=collections
3410 )
3412 query = do_query("detector")
3413 self.assertEqual(len(list(query)), 4)
3415 Test = namedtuple(
3416 "testQueryDataIdsOrderByTest",
3417 ("element", "order_by", "result", "limit", "datasets", "collections"),
3418 defaults=(None, None, None),
3419 )
3421 test_data = [
3422 Test("detector", "detector", (1, 2, 3, 4)),
3423 Test("detector", "-detector", (4, 3, 2, 1)),
3424 Test("detector", "raft,-name_in_raft", (2, 1, 4, 3)),
3425 Test("detector", "-detector.purpose", (4,), limit=(1,)),
3426 Test("visit", "visit", (1, 2)),
3427 Test("visit", "-visit.id", (2, 1)),
3428 Test("visit", "zenith_angle", (1, 2)),
3429 Test("visit", "-visit.name", (2, 1)),
3430 Test("visit", "day_obs,-visit.timespan.begin", (2, 1)),
3431 ]
3433 def do_test(test: Test):
3434 order_by = test.order_by.split(",")
3435 query = do_query(test.element).order_by(*order_by)
3436 if test.limit is not None:
3437 query = query.limit(*test.limit)
3438 dataIds = tuple(rec.id for rec in query)
3439 self.assertEqual(dataIds, test.result)
3441 for test in test_data:
3442 do_test(test)
3444 # errors in a name
3445 for order_by in ("", "-"):
3446 with self.assertRaisesRegex(
3447 (ValueError, InvalidQueryError),
3448 "(Empty dimension name in ORDER BY)|(Unrecognized identifier)",
3449 ):
3450 list(do_query("detector").order_by(order_by))
3452 for order_by in ("undimension.name", "-undimension.name"):
3453 with self.assertRaisesRegex(
3454 (ValueError, InvalidQueryError),
3455 "(Element name mismatch: 'undimension')|(Unrecognized identifier)",
3456 ):
3457 list(do_query("detector").order_by(order_by))
3459 for order_by in ("attract", "-attract"):
3460 with self.assertRaisesRegex(
3461 (ValueError, InvalidQueryError),
3462 "(Field 'attract' does not exist in 'detector'.)|(Unrecognized identifier)",
3463 ):
3464 list(do_query("detector").order_by(order_by))
3466 for order_by in ("timestamp.begin", "-timestamp.begin"):
3467 with self.assertRaisesRegex(
3468 (ValueError, InvalidQueryError),
3469 r"(Element name mismatch: 'timestamp' instead of 'visit'; "
3470 r"perhaps you meant 'timespan.begin'\?)"
3471 r"|(Unrecognized identifier)",
3472 ):
3473 list(do_query("visit").order_by(order_by))
3475 def testQueryDimensionRecordsExceptions(self):
3476 """Test exceptions raised by queryDimensionRecords()."""
3477 butler = self.make_butler()
3478 registry = butler.registry
3479 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3481 result = registry.queryDimensionRecords("detector")
3482 self.assertEqual(result.count(), 4)
3483 result = registry.queryDimensionRecords("detector", instrument="Cam1")
3484 self.assertEqual(result.count(), 4)
3485 result = registry.queryDimensionRecords("detector", dataId={"instrument": "Cam1"})
3486 self.assertEqual(result.count(), 4)
3488 # Test that values specified in kwargs override those specified in
3489 # dataId.
3490 result = registry.queryDimensionRecords(
3491 "detector", dataId={"instrument": "NotCam1"}, instrument="Cam1"
3492 )
3493 self.assertEqual(result.count(), 4)
3495 result = registry.queryDimensionRecords("detector", where="instrument='Cam1'")
3496 self.assertEqual(result.count(), 4)
3497 result = registry.queryDimensionRecords("detector", where="instrument=:instr", bind={"instr": "Cam1"})
3498 self.assertTrue(result.any())
3499 self.assertEqual(result.count(), 4)
3501 def testDatasetConstrainedDimensionRecordQueries(self):
3502 """Test that queryDimensionRecords works even when given a dataset
3503 constraint whose dimensions extend beyond the requested dimension
3504 element's.
3505 """
3506 butler = self.make_butler()
3507 registry = butler.registry
3508 self.load_data(butler, "base.yaml", "datasets.yaml")
3509 # Query for physical_filter dimension records, using a dataset that
3510 # has both physical_filter and dataset dimensions.
3511 records = registry.queryDimensionRecords(
3512 "physical_filter",
3513 datasets=["flat"],
3514 collections="imported_r",
3515 )
3516 self.assertEqual({record.name for record in records}, {"Cam1-R1", "Cam1-R2"})
3517 # Trying to constrain by all dataset types is an error.
3518 with self.assertRaises(TypeError):
3519 list(registry.queryDimensionRecords("physical_filter", datasets=..., collections="imported_r"))
3521 def testSkyPixDatasetQueries(self):
3522 """Test that we can build queries involving skypix dimensions as long
3523 as a dataset type that uses those dimensions is included.
3524 """
3525 butler = self.make_butler()
3526 registry = butler.registry
3527 self.load_data(butler, "base.yaml")
3528 dataset_type = DatasetType(
3529 "a", dimensions=["htm7", "instrument"], universe=registry.dimensions, storageClass="int"
3530 )
3531 registry.registerDatasetType(dataset_type)
3532 run = "r"
3533 registry.registerRun(run)
3534 # First try queries where there are no datasets; the concern is whether
3535 # we can even build and execute these queries without raising, even
3536 # when "doomed" query shortcuts are in play.
3537 self.assertFalse(
3538 list(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run))
3539 )
3540 self.assertFalse(list(registry.queryDatasets(dataset_type, collections=run)))
3541 # Now add a dataset and see that we can get it back.
3542 htm7 = registry.dimensions.skypix["htm"][7].pixelization
3543 data_id = registry.expandDataId(instrument="Cam1", htm7=htm7.universe()[0][0])
3544 (ref,) = registry.insertDatasets(dataset_type, [data_id], run=run)
3545 self.assertEqual(
3546 set(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run)),
3547 {data_id},
3548 )
3549 self.assertEqual(set(registry.queryDatasets(dataset_type, collections=run)), {ref})
3551 def testDatasetIdFactory(self):
3552 """Simple test for DatasetIdFactory, mostly to catch potential changes
3553 in its API.
3554 """
3555 butler = self.make_butler()
3556 registry = butler.registry
3557 factory = DatasetIdFactory()
3558 dataset_type = DatasetType(
3559 "datasetType",
3560 dimensions=["detector", "instrument"],
3561 universe=registry.dimensions,
3562 storageClass="int",
3563 )
3564 run = "run"
3565 data_id = DataCoordinate.standardize(
3566 instrument="Cam1", detector=1, dimensions=dataset_type.dimensions
3567 )
3569 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.UNIQUE)
3570 self.assertIsInstance(datasetId, uuid.UUID)
3571 self.assertEqual(datasetId.version, 7)
3573 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE)
3574 self.assertIsInstance(datasetId, uuid.UUID)
3575 self.assertEqual(datasetId.version, 5)
3577 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE_RUN)
3578 self.assertIsInstance(datasetId, uuid.UUID)
3579 self.assertEqual(datasetId.version, 5)
3581 def testExposureQueries(self):
3582 """Test query methods using arguments sourced from the exposure log
3583 service.
3585 The most complete test dataset currently available to daf_butler tests
3586 is ci_hsc-subset.yaml export , but that does not have 'exposure'
3587 dimension records. So in this test we need to translate queries that
3588 originally used the exposure dimension to use the (very similar) visit
3589 dimension instead.
3590 """
3591 butler = self.make_butler()
3592 registry = butler.registry
3593 self.load_data(butler, "ci_hsc-subset.yaml")
3594 self.assertEqual(
3595 [
3596 record.id
3597 for record in registry.queryDimensionRecords("visit", instrument="HSC")
3598 .order_by("visit")
3599 .limit(5)
3600 ],
3601 [903334, 903336, 903338, 903342, 903344],
3602 )
3603 self.assertEqual(
3604 [
3605 data_id["visit"]
3606 for data_id in registry.queryDataIds(["visit"], instrument="HSC").order_by("visit").limit(5)
3607 ],
3608 [903334, 903336, 903338, 903342, 903344],
3609 )
3610 self.assertEqual(
3611 [
3612 record.id
3613 for record in registry.queryDimensionRecords("detector", instrument="HSC")
3614 .order_by("full_name")
3615 .limit(5)
3616 ],
3617 [25, 24, 23, 22, 18],
3618 )
3619 self.assertEqual(
3620 [
3621 data_id["detector"]
3622 for data_id in registry.queryDataIds(["detector"], instrument="HSC")
3623 .order_by("full_name")
3624 .limit(5)
3625 ],
3626 [25, 24, 23, 22, 18],
3627 )
3629 def test_long_query_names(self) -> None:
3630 """Test that queries involving very long names are handled correctly.
3632 This is especially important for PostgreSQL, which truncates symbols
3633 longer than 64 chars, but it's worth testing for all DBs.
3634 """
3635 butler = self.make_butler()
3636 registry = butler.registry
3637 name = "abcd" * 17
3638 registry.registerDatasetType(
3639 DatasetType(
3640 name,
3641 dimensions=(),
3642 storageClass="Exposure",
3643 universe=registry.dimensions,
3644 )
3645 )
3646 # Need to search more than one collection actually containing a
3647 # matching dataset to avoid optimizations that sidestep bugs due to
3648 # truncation by making findFirst=True a no-op.
3649 run1 = "run1"
3650 registry.registerRun(run1)
3651 run2 = "run2"
3652 registry.registerRun(run2)
3653 (ref1,) = registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run1)
3654 registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run2)
3655 self.assertEqual(
3656 set(registry.queryDatasets(name, collections=[run1, run2], findFirst=True)),
3657 {ref1},
3658 )
3660 def test_skypix_constraint_queries(self) -> None:
3661 """Test queries spatially constrained by a skypix data ID."""
3662 butler = self.make_butler()
3663 registry = butler.registry
3664 self.load_data(butler, "base.yaml", "spatial.yaml")
3665 patch_regions = {
3666 (data_id["tract"], data_id["patch"]): data_id.region
3667 for data_id in registry.queryDataIds(["patch"]).expanded()
3668 }
3669 skypix_dimension: SkyPixDimension = registry.dimensions["htm11"]
3670 # This check ensures the test doesn't become trivial due to a config
3671 # change; if it does, just pick a different HTML level.
3672 self.assertNotEqual(skypix_dimension, registry.dimensions.commonSkyPix)
3673 # Gather all skypix IDs that definitely overlap at least one of these
3674 # patches.
3675 relevant_skypix_ids = lsst.sphgeom.RangeSet()
3676 for patch_region in patch_regions.values():
3677 relevant_skypix_ids |= skypix_dimension.pixelization.interior(patch_region)
3678 # Look for a "nontrivial" skypix_id that overlaps at least one patch
3679 # and does not overlap at least one other patch.
3680 for skypix_id in itertools.chain.from_iterable(
3681 range(begin, end) for begin, end in relevant_skypix_ids
3682 ):
3683 skypix_region = skypix_dimension.pixelization.pixel(skypix_id)
3684 overlapping_patches = {
3685 patch_key
3686 for patch_key, patch_region in patch_regions.items()
3687 if not patch_region.isDisjointFrom(skypix_region)
3688 }
3689 if overlapping_patches and overlapping_patches != patch_regions.keys():
3690 break
3691 else:
3692 raise RuntimeError("Could not find usable skypix ID for this dimension configuration.")
3693 # Test that a three-way join that includes the common skypix system in
3694 # the dimensions doesn't generate redundant join terms in the query.
3695 with self.assertRaises(InvalidQueryError):
3696 set(
3697 registry.queryDataIds(
3698 ["tract", "visit", "htm7"], skymap="SkyMap1", instrument="Cam1"
3699 ).expanded()
3700 )
3702 def test_spatial_constraint_queries(self) -> None:
3703 """Test queries in which one spatial dimension in the constraint (data
3704 ID or ``where`` string) constrains a different spatial dimension in the
3705 query result columns.
3706 """
3707 butler = self.make_butler()
3708 registry = butler.registry
3709 self.load_data(butler, "base.yaml", "spatial.yaml")
3710 patch_regions = {
3711 (data_id["tract"], data_id["patch"]): data_id.region
3712 for data_id in registry.queryDataIds(["patch"]).expanded()
3713 }
3714 observation_regions = {
3715 (data_id["visit"], data_id["detector"]): data_id.region
3716 for data_id in registry.queryDataIds(["visit", "detector"]).expanded()
3717 }
3718 all_combos = {
3719 (patch_key, observation_key)
3720 for patch_key, observation_key in itertools.product(patch_regions, observation_regions)
3721 }
3722 overlapping_combos = {
3723 (patch_key, observation_key)
3724 for patch_key, observation_key in all_combos
3725 if not patch_regions[patch_key].isDisjointFrom(observation_regions[observation_key])
3726 }
3727 # Check a direct spatial join with no constraint first.
3728 self.assertEqual(
3729 {
3730 ((data_id["tract"], data_id["patch"]), (data_id["visit"], data_id["detector"]))
3731 for data_id in registry.queryDataIds(["patch", "visit", "detector"])
3732 },
3733 overlapping_combos,
3734 )
3735 overlaps_by_patch: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)
3736 overlaps_by_observation: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)
3737 for patch_key, observation_key in overlapping_combos:
3738 overlaps_by_patch[patch_key].add(observation_key)
3739 overlaps_by_observation[observation_key].add(patch_key)
3740 # Find patches and observations that overlap at least one of the other
3741 # but not all of the other.
3742 nontrivial_patch = next(
3743 iter(
3744 patch_key
3745 for patch_key, observation_keys in overlaps_by_patch.items()
3746 if observation_keys and observation_keys != observation_regions.keys()
3747 )
3748 )
3749 nontrivial_observation = next(
3750 iter(
3751 observation_key
3752 for observation_key, patch_keys in overlaps_by_observation.items()
3753 if patch_keys and patch_keys != patch_regions.keys()
3754 )
3755 )
3756 # Use the nontrivial patches and observations as constraints on the
3757 # other dimensions in various ways, first via a 'where' expression.
3758 # It's better in general to us 'bind' instead of f-strings, but these
3759 # all integers so there are no quoting concerns.
3760 self.assertEqual(
3761 {
3762 (data_id["visit"], data_id["detector"])
3763 for data_id in registry.queryDataIds(
3764 ["visit", "detector"],
3765 where=f"tract={nontrivial_patch[0]} AND patch={nontrivial_patch[1]}",
3766 skymap="SkyMap1",
3767 )
3768 },
3769 overlaps_by_patch[nontrivial_patch],
3770 )
3771 self.assertEqual(
3772 {
3773 (data_id["tract"], data_id["patch"])
3774 for data_id in registry.queryDataIds(
3775 ["patch"],
3776 where=f"visit={nontrivial_observation[0]} AND detector={nontrivial_observation[1]}",
3777 instrument="Cam1",
3778 )
3779 },
3780 overlaps_by_observation[nontrivial_observation],
3781 )
3782 # and then via the dataId argument.
3783 self.assertEqual(
3784 {
3785 (data_id["visit"], data_id["detector"])
3786 for data_id in registry.queryDataIds(
3787 ["visit", "detector"],
3788 dataId={
3789 "tract": nontrivial_patch[0],
3790 "patch": nontrivial_patch[1],
3791 },
3792 skymap="SkyMap1",
3793 )
3794 },
3795 overlaps_by_patch[nontrivial_patch],
3796 )
3797 self.assertEqual(
3798 {
3799 (data_id["tract"], data_id["patch"])
3800 for data_id in registry.queryDataIds(
3801 ["patch"],
3802 dataId={
3803 "visit": nontrivial_observation[0],
3804 "detector": nontrivial_observation[1],
3805 },
3806 instrument="Cam1",
3807 )
3808 },
3809 overlaps_by_observation[nontrivial_observation],
3810 )
3812 def test_query_empty_collections(self) -> None:
3813 """Test for registry query methods with empty collections. The methods
3814 should return empty result set (or None when applicable) and provide
3815 "doomed" diagnostics.
3816 """
3817 butler = self.make_butler()
3818 registry = butler.registry
3819 self.load_data(butler, "base.yaml", "datasets.yaml")
3821 # Tests for registry.findDataset()
3822 with self.assertRaises(NoDefaultCollectionError):
3823 registry.findDataset("bias", instrument="Cam1", detector=1)
3824 self.assertIsNotNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=...))
3825 self.assertIsNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=[]))
3827 # Tests for registry.queryDatasets()
3828 with self.assertRaises(NoDefaultCollectionError):
3829 registry.queryDatasets("bias")
3830 self.assertTrue(list(registry.queryDatasets("bias", collections=...)))
3832 result = registry.queryDatasets("bias", collections=[])
3833 self.assertEqual(len(list(result)), 0)
3834 messages = list(result.explain_no_results())
3835 self.assertTrue(messages)
3836 self.assertTrue(any("because collection list is empty" in message for message in messages))
3838 # Tests for registry.queryDataIds()
3839 with self.assertRaises(NoDefaultCollectionError):
3840 registry.queryDataIds("detector", datasets="bias")
3841 self.assertTrue(list(registry.queryDataIds("detector", datasets="bias", collections=...)))
3843 result = registry.queryDataIds("detector", datasets="bias", collections=[])
3844 self.assertEqual(len(list(result)), 0)
3845 messages = list(result.explain_no_results())
3846 self.assertTrue(messages)
3847 self.assertTrue(any("because collection list is empty" in message for message in messages))
3849 # Tests for registry.queryDimensionRecords()
3850 with self.assertRaises(NoDefaultCollectionError):
3851 registry.queryDimensionRecords("detector", datasets="bias")
3852 self.assertTrue(list(registry.queryDimensionRecords("detector", datasets="bias", collections=...)))
3854 result = registry.queryDimensionRecords("detector", datasets="bias", collections=[])
3855 self.assertEqual(len(list(result)), 0)
3856 messages = list(result.explain_no_results())
3857 self.assertTrue(messages)
3858 self.assertTrue(any("because collection list is empty" in message for message in messages))
3860 def test_dataset_followup_spatial_joins(self) -> None:
3861 """Test queryDataIds(...).findRelatedDatasets(...) where a spatial join
3862 is involved.
3863 """
3864 butler = self.make_butler()
3865 registry = butler.registry
3866 self.load_data(butler, "base.yaml", "spatial.yaml")
3867 pvi_dataset_type = DatasetType(
3868 "pvi", {"visit", "detector"}, storageClass="StructuredDataDict", universe=registry.dimensions
3869 )
3870 registry.registerDatasetType(pvi_dataset_type)
3871 collection = "datasets"
3872 registry.registerRun(collection)
3873 (pvi1,) = registry.insertDatasets(
3874 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 1}], run=collection
3875 )
3876 (pvi2,) = registry.insertDatasets(
3877 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 2}], run=collection
3878 )
3879 (pvi3,) = registry.insertDatasets(
3880 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 3}], run=collection
3881 )
3882 self.assertEqual(
3883 set(
3884 registry.queryDataIds(["patch"], skymap="SkyMap1", tract=0)
3885 .expanded()
3886 .findRelatedDatasets("pvi", [collection])
3887 ),
3888 {
3889 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi1),
3890 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi2),
3891 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=1), pvi2),
3892 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi1),
3893 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi2),
3894 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi3),
3895 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=3), pvi2),
3896 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=4), pvi3),
3897 },
3898 )
3900 def test_expanded_data_id_queries(self) -> None:
3901 """Tests for basic functionality of expanded() on queryDataIds and
3902 queryDatasets.
3903 """
3904 butler = self.make_butler()
3905 registry = butler.registry
3906 self.load_data(butler, "base.yaml", "spatial.yaml", "datasets.yaml")
3908 result_obj = (
3909 registry.queryDataIds(["visit"], where="instrument = 'Cam1' and (visit.id = 1 or visit.id = 2)")
3910 .expanded()
3911 .order_by("visit.id")
3912 )
3913 self.assertTrue(result_obj.hasRecords())
3914 visits = list(result_obj)
3915 self.assertEqual(len(visits), 2)
3917 self.assertEqual(visits[0]["visit"], 1)
3918 self.assertEqual(visits[1]["visit"], 2)
3919 self.assertEqual(visits[0].records["visit"].exposure_time, 60.0)
3920 self.assertEqual(visits[1].records["visit"].exposure_time, 45.0)
3921 # physical_filter is a "cacheable" dimension, so its records are loaded
3922 # from local cache rather than being part of the DB rows.
3923 self.assertEqual(visits[0].records["physical_filter"].band, "g")
3924 self.assertEqual(visits[1].records["physical_filter"].band, "r")
3926 # Make sure that we can fetch nulls in dimension records
3927 registry.insertDimensionData(
3928 "detector",
3929 {
3930 "instrument": "Cam1",
3931 "id": 5,
3932 "raft": "Z",
3933 "name_in_raft": "z",
3934 "full_name": "Zz",
3935 "purpose": None,
3936 },
3937 )
3938 detectors = list(
3939 registry.queryDataIds("detector", dataId={"instrument": "Cam1", "detector": 5}).expanded()
3940 )
3941 self.assertIsNone(detectors[0].records["detector"].purpose)
3943 datasets_query = registry.queryDatasets(
3944 "flat", collections="imported_g", where="instrument = 'Cam1' and detector <= 3"
3945 ).expanded()
3946 datasets = list(datasets_query)
3947 datasets.sort(key=lambda ref: ref.dataId["detector"])
3948 self.assertEqual(len(datasets), 2)
3949 self.assertEqual(datasets[0].id, uuid.UUID("60c8a65c-7290-4c38-b1de-e3b1cdcf872d"))
3950 self.assertEqual(datasets[1].id, uuid.UUID("84239e7f-c41f-46d5-97b9-a27976b98ceb"))
3951 # All of the dimensions for flat are "cached" dimensions.
3952 self.assertEqual(datasets[0].dataId.records["detector"].full_name, "Ab")
3953 self.assertEqual(datasets[1].dataId.records["detector"].full_name, "Ba")
3954 self.assertEqual(datasets[0].dataId.records["instrument"].visit_system, 1)
3955 assert isinstance(datasets_query, ParentDatasetQueryResults)
3956 data_ids = list(datasets_query.dataIds)
3957 data_ids.sort(key=lambda data_id: data_id["detector"])
3958 self.assertEqual(len(data_ids), 2)
3959 self.assertEqual(data_ids[0].records["detector"].full_name, "Ab")
3960 self.assertEqual(data_ids[1].records["detector"].full_name, "Ba")
3961 self.assertEqual(data_ids[0].records["instrument"].visit_system, 1)
3963 # None of the datasets in the test data include any uncached
3964 # dimensions, so we have to set one up.
3965 registry.registerDatasetType(DatasetType("test", ["visit"], "int", universe=registry.dimensions))
3966 registry.insertDatasets("test", [{"instrument": "Cam1", "visit": 1}], run="imported_g")
3967 ref = list(registry.queryDatasets("test", collections="imported_g").expanded())[0]
3968 self.assertEqual(ref.dataId.records["visit"].zenith_angle, 5.0)
3969 self.assertEqual(ref.dataId.records["physical_filter"].band, "g")
3970 self.assertEqual(
3971 ref.dataId.timespan,
3972 Timespan(
3973 begin=astropy.time.Time("2021-09-09 03:00:00.000000000", scale="tai"),
3974 end=astropy.time.Time("2021-09-09 03:01:00.000000000", scale="tai"),
3975 ),
3976 )
3978 def test_collection_summary(self) -> None:
3979 """Test for collection summary methods."""
3980 butler = self.make_butler()
3981 registry = butler.registry
3982 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3984 # Add one more dataset type, just for its existence to trigger a bug
3985 # in `associate` (DM-44311).
3986 test_dataset_type = DatasetType("test", ["tract", "patch"], "int", universe=registry.dimensions)
3987 registry.registerDatasetType(test_dataset_type)
3989 # Check for what has been imported.
3990 summary = registry.getCollectionSummary("imported_g")
3991 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
3992 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
3994 # Make a tagged collection and associate some datasets.
3995 tagged_coll = "tagged"
3996 registry.registerCollection(tagged_coll, CollectionType.TAGGED)
3997 refsets = registry.queryDatasets(..., collections=["imported_g"]).byParentDatasetType()
3998 for refs in refsets:
3999 registry.associate(tagged_coll, refs)
4001 # Summary has to have the same dataset types.
4002 summary = registry.getCollectionSummary(tagged_coll)
4003 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
4004 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
4006 # Remove all datasets from the tagged collection.
4007 refs = list(registry.queryDatasets(..., collections=[tagged_coll]))
4008 registry.disassociate(tagged_coll, refs)
4010 # Summaries should not have changed.
4011 summary = registry.getCollectionSummary(tagged_coll)
4012 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
4013 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
4015 # Cleanup summaries.
4016 registry.refresh_collection_summaries()
4017 summary = registry.getCollectionSummary(tagged_coll)
4018 self.assertFalse(summary.dataset_types.names)
4019 # We do not clean governor summaries yet, but because how the query is
4020 # run, it returns empty governors when collection is missing from
4021 # summaries.
4022 self.assertFalse(summary.governors)
4024 # Add dataset with different governor, this is to test that governors
4025 # are not actually cleaned.
4026 refs = registry.insertDatasets("test", [{"skymap": "SkyMap1", "tract": 0, "patch": 0}], "imported_g")
4027 registry.associate(tagged_coll, refs)
4028 summary = registry.getCollectionSummary(tagged_coll)
4029 self.assertEqual(summary.dataset_types.names, {"test"})
4030 # Note that instrument governor resurrects here, even though there are
4031 # no datasets left with that governor.
4032 self.assertEqual(summary.governors, {"instrument": {"Cam1"}, "skymap": {"SkyMap1"}})
4034 def test_temp_table_config(self) -> None:
4035 config = self.makeRegistryConfig()
4036 config["temporary_tables"] = False
4037 self.assertEqual(config.areTemporaryTablesAllowed, False)
4038 butler = self.make_butler(config)
4039 if not isinstance(butler, DirectButler):
4040 raise unittest.SkipTest("Test only makes sense for registry with direct database connection.")
4041 self.assertEqual(butler._registry._db.supports_temporary_tables, False)
4042 with self.assertRaisesRegex(ReadOnlyDatabaseError, "temporary tables"):
4043 with butler._registry._db.temporary_table(...):
4044 pass