Coverage for python/lsst/daf/butler/registry/tests/_registry.py: 6%
1684 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-13 01:31 -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 # Test dataset association query against a chained collection.
2544 # This is a regression test for DM-53179, as well as verification
2545 # that the flattenChains parameter has never had any effect.
2546 butler.collections.register("chain", CollectionType.CHAINED)
2547 butler.collections.redefine_chain("chain", [collection])
2548 expected_datasets = (
2549 DatasetAssociation(ref=bias2a, collection=collection, timespan=Timespan(begin=t2, end=t4)),
2550 DatasetAssociation(ref=bias3a, collection=collection, timespan=Timespan(begin=t1, end=t3)),
2551 DatasetAssociation(ref=bias2b, collection=collection, timespan=Timespan(begin=t4, end=None)),
2552 DatasetAssociation(ref=bias3b, collection=collection, timespan=Timespan(begin=t4, end=None)),
2553 )
2554 self.assertCountEqual(
2555 list(registry.queryDatasetAssociations("bias", collections=["chain"], flattenChains=False)),
2556 expected_datasets,
2557 )
2558 self.assertCountEqual(
2559 list(registry.queryDatasetAssociations("bias", collections=["chain"], flattenChains=True)),
2560 expected_datasets,
2561 )
2563 class Ambiguous:
2564 """Tag class to denote lookups that should be ambiguous."""
2566 pass
2568 def _assertLookup(
2569 detector: int, timespan: Timespan, expected: DatasetRef | type[Ambiguous] | None
2570 ) -> None:
2571 """Local function that asserts that a bias lookup returns the given
2572 expected result.
2573 """
2574 if expected is Ambiguous:
2575 with self.assertRaises((DatasetTypeError, LookupError)):
2576 registry.findDataset(
2577 "bias",
2578 collections=collection,
2579 instrument="Cam1",
2580 detector=detector,
2581 timespan=timespan,
2582 )
2583 else:
2584 self.assertEqual(
2585 expected,
2586 registry.findDataset(
2587 "bias",
2588 collections=collection,
2589 instrument="Cam1",
2590 detector=detector,
2591 timespan=timespan,
2592 ),
2593 )
2595 # Systematically test lookups against expected results.
2596 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)
2597 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)
2598 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)
2599 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)
2600 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=Ambiguous)
2601 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)
2602 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)
2603 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)
2604 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)
2605 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=Ambiguous)
2606 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)
2607 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)
2608 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)
2609 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=Ambiguous)
2610 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)
2611 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=bias2a)
2612 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=Ambiguous)
2613 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=Ambiguous)
2614 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=bias2b)
2615 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)
2616 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)
2617 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)
2618 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)
2619 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)
2620 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)
2621 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=Ambiguous)
2622 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)
2623 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)
2624 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)
2625 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)
2626 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=Ambiguous)
2627 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)
2628 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)
2629 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)
2630 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=Ambiguous)
2631 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)
2632 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)
2633 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=bias3b)
2634 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)
2635 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=bias3b)
2636 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)
2637 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)
2639 # Test lookups via temporal joins to exposures.
2640 self.assertEqual(
2641 set(
2642 registry.queryDataIds(
2643 ["exposure", "detector"], instrument="Cam1", detector=2
2644 ).findRelatedDatasets("bias", collections=[collection])
2645 ),
2646 {
2647 (registry.expandDataId(instrument="Cam1", exposure=1, detector=2), bias2a),
2648 (registry.expandDataId(instrument="Cam1", exposure=2, detector=2), bias2a),
2649 (registry.expandDataId(instrument="Cam1", exposure=3, detector=2), bias2b),
2650 },
2651 )
2652 self.assertEqual(
2653 set(
2654 registry.queryDataIds(
2655 ["exposure", "detector"], instrument="Cam1", detector=3
2656 ).findRelatedDatasets("bias", collections=[collection])
2657 ),
2658 {
2659 (registry.expandDataId(instrument="Cam1", exposure=0, detector=3), bias3a),
2660 (registry.expandDataId(instrument="Cam1", exposure=1, detector=3), bias3a),
2661 (registry.expandDataId(instrument="Cam1", exposure=3, detector=3), bias3b),
2662 },
2663 )
2665 # Decertify [t3, t5) for all data IDs, and do test lookups again.
2666 # This should truncate bias2a to [t2, t3), leave bias3a unchanged at
2667 # [t1, t3), and truncate bias2b and bias3b to [t5, ∞).
2668 registry.decertify(collection=collection, datasetType="bias", timespan=Timespan(t3, t5))
2669 _assertLookup(detector=2, timespan=Timespan(None, t1), expected=None)
2670 _assertLookup(detector=2, timespan=Timespan(None, t2), expected=None)
2671 _assertLookup(detector=2, timespan=Timespan(None, t3), expected=bias2a)
2672 _assertLookup(detector=2, timespan=Timespan(None, t4), expected=bias2a)
2673 _assertLookup(detector=2, timespan=Timespan(None, t5), expected=bias2a)
2674 _assertLookup(detector=2, timespan=Timespan(None, None), expected=Ambiguous)
2675 _assertLookup(detector=2, timespan=Timespan(t1, t2), expected=None)
2676 _assertLookup(detector=2, timespan=Timespan(t1, t3), expected=bias2a)
2677 _assertLookup(detector=2, timespan=Timespan(t1, t4), expected=bias2a)
2678 _assertLookup(detector=2, timespan=Timespan(t1, t5), expected=bias2a)
2679 _assertLookup(detector=2, timespan=Timespan(t1, None), expected=Ambiguous)
2680 _assertLookup(detector=2, timespan=Timespan(t2, t3), expected=bias2a)
2681 _assertLookup(detector=2, timespan=Timespan(t2, t4), expected=bias2a)
2682 _assertLookup(detector=2, timespan=Timespan(t2, t5), expected=bias2a)
2683 _assertLookup(detector=2, timespan=Timespan(t2, None), expected=Ambiguous)
2684 _assertLookup(detector=2, timespan=Timespan(t3, t4), expected=None)
2685 _assertLookup(detector=2, timespan=Timespan(t3, t5), expected=None)
2686 _assertLookup(detector=2, timespan=Timespan(t3, None), expected=bias2b)
2687 _assertLookup(detector=2, timespan=Timespan(t4, t5), expected=None)
2688 _assertLookup(detector=2, timespan=Timespan(t4, None), expected=bias2b)
2689 _assertLookup(detector=2, timespan=Timespan(t5, None), expected=bias2b)
2690 _assertLookup(detector=3, timespan=Timespan(None, t1), expected=None)
2691 _assertLookup(detector=3, timespan=Timespan(None, t2), expected=bias3a)
2692 _assertLookup(detector=3, timespan=Timespan(None, t3), expected=bias3a)
2693 _assertLookup(detector=3, timespan=Timespan(None, t4), expected=bias3a)
2694 _assertLookup(detector=3, timespan=Timespan(None, t5), expected=bias3a)
2695 _assertLookup(detector=3, timespan=Timespan(None, None), expected=Ambiguous)
2696 _assertLookup(detector=3, timespan=Timespan(t1, t2), expected=bias3a)
2697 _assertLookup(detector=3, timespan=Timespan(t1, t3), expected=bias3a)
2698 _assertLookup(detector=3, timespan=Timespan(t1, t4), expected=bias3a)
2699 _assertLookup(detector=3, timespan=Timespan(t1, t5), expected=bias3a)
2700 _assertLookup(detector=3, timespan=Timespan(t1, None), expected=Ambiguous)
2701 _assertLookup(detector=3, timespan=Timespan(t2, t3), expected=bias3a)
2702 _assertLookup(detector=3, timespan=Timespan(t2, t4), expected=bias3a)
2703 _assertLookup(detector=3, timespan=Timespan(t2, t5), expected=bias3a)
2704 _assertLookup(detector=3, timespan=Timespan(t2, None), expected=Ambiguous)
2705 _assertLookup(detector=3, timespan=Timespan(t3, t4), expected=None)
2706 _assertLookup(detector=3, timespan=Timespan(t3, t5), expected=None)
2707 _assertLookup(detector=3, timespan=Timespan(t3, None), expected=bias3b)
2708 _assertLookup(detector=3, timespan=Timespan(t4, t5), expected=None)
2709 _assertLookup(detector=3, timespan=Timespan(t4, None), expected=bias3b)
2710 _assertLookup(detector=3, timespan=Timespan(t5, None), expected=bias3b)
2712 # Decertify everything, this time with explicit data IDs, then check
2713 # that no lookups succeed.
2714 registry.decertify(
2715 collection,
2716 "bias",
2717 Timespan(None, None),
2718 dataIds=[
2719 dict(instrument="Cam1", detector=2),
2720 dict(instrument="Cam1", detector=3),
2721 ],
2722 )
2723 for detector in (2, 3):
2724 for timespan in allTimespans:
2725 _assertLookup(detector=detector, timespan=timespan, expected=None)
2726 # Certify bias2a and bias3a over (-∞, ∞), check that all lookups return
2727 # those.
2728 registry.certify(
2729 collection,
2730 [bias2a, bias3a],
2731 Timespan(None, None),
2732 )
2733 for timespan in allTimespans:
2734 _assertLookup(detector=2, timespan=timespan, expected=bias2a)
2735 _assertLookup(detector=3, timespan=timespan, expected=bias3a)
2736 # Decertify just bias2 over [t2, t4).
2737 # This should split a single certification row into two (and leave the
2738 # other existing row, for bias3a, alone).
2739 registry.decertify(
2740 collection, "bias", Timespan(t2, t4), dataIds=[dict(instrument="Cam1", detector=2)]
2741 )
2742 for timespan in allTimespans:
2743 _assertLookup(detector=3, timespan=timespan, expected=bias3a)
2744 overlapsBefore = timespan.overlaps(Timespan(None, t2))
2745 overlapsAfter = timespan.overlaps(Timespan(t4, None))
2746 if overlapsBefore and overlapsAfter:
2747 expected = Ambiguous
2748 elif overlapsBefore or overlapsAfter:
2749 expected = bias2a
2750 else:
2751 expected = None
2752 _assertLookup(detector=2, timespan=timespan, expected=expected)
2754 def testSkipCalibs(self):
2755 """Test how queries handle skipping of calibration collections."""
2756 butler = self.make_butler()
2757 registry = butler.registry
2758 self.load_data(butler, "base.yaml", "datasets.yaml")
2760 coll_calib = "Cam1/calibs/default"
2761 registry.registerCollection(coll_calib, type=CollectionType.CALIBRATION)
2763 # Add all biases to the calibration collection.
2764 # Without this, the logic that prunes dataset subqueries based on
2765 # datasetType-collection summary information will fire before the logic
2766 # we want to test below. This is a good thing (it avoids the dreaded
2767 # NotImplementedError a bit more often) everywhere but here.
2768 registry.certify(coll_calib, registry.queryDatasets("bias", collections=...), Timespan(None, None))
2770 coll_list = [coll_calib, "imported_g", "imported_r"]
2771 chain = "Cam1/chain"
2772 registry.registerCollection(chain, type=CollectionType.CHAINED)
2773 registry.setCollectionChain(chain, coll_list)
2775 # Lookup is ambiguous due to multiple datasets with the same data ID
2776 # in the calibration collection.
2777 with self.assertRaises(CalibrationLookupError):
2778 list(registry.queryDatasets("bias", collections=coll_list, findFirst=True))
2780 # chain will skip
2781 datasets = list(registry.queryDatasets("bias", collections=chain))
2782 self.assertGreater(len(datasets), 0)
2784 dataIds = list(registry.queryDataIds(["instrument", "detector"], datasets="bias", collections=chain))
2785 self.assertGreater(len(dataIds), 0)
2787 # glob will skip too
2788 datasets = list(registry.queryDatasets("bias", collections="*d*"))
2789 self.assertGreater(len(datasets), 0)
2791 # regular expression will skip too
2792 if self.supportsCollectionRegex:
2793 pattern = re.compile(".*")
2794 with self.assertWarns(FutureWarning):
2795 datasets = list(registry.queryDatasets("bias", collections=pattern))
2796 self.assertGreater(len(datasets), 0)
2798 # ellipsis should work as usual
2799 datasets = list(registry.queryDatasets("bias", collections=...))
2800 self.assertGreater(len(datasets), 0)
2802 # New query system correctly determines that this search is
2803 # ambiguous, because there are multiple datasets with the same
2804 # {instrument=Cam1, detector=2} data ID in the calibration
2805 # collection at the beginning of the chain.
2806 with self.assertRaises(CalibrationLookupError):
2807 datasets = list(registry.queryDatasets("bias", collections=chain, findFirst=True))
2809 def testIngestTimeQuery(self):
2810 butler = self.make_butler()
2811 registry = butler.registry
2812 dt0 = datetime.datetime.now(datetime.UTC)
2813 self.load_data(butler, "base.yaml", "datasets.yaml")
2814 dt1 = datetime.datetime.now(datetime.UTC)
2816 datasets = list(registry.queryDatasets(..., collections=...))
2817 len0 = len(datasets)
2818 self.assertGreater(len0, 0)
2820 for where in ("ingest_date > T'2000-01-01'", "T'2000-01-01' < ingest_date"):
2821 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2822 len1 = len(datasets)
2823 self.assertEqual(len0, len1)
2825 # no one will ever use this piece of software in 30 years
2826 for where in ("ingest_date > T'2050-01-01'", "T'2050-01-01' < ingest_date"):
2827 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2828 len2 = len(datasets)
2829 self.assertEqual(len2, 0)
2831 # Check more exact timing to make sure there is no 37 seconds offset
2832 # (after fixing DM-30124). SQLite time precision is 1 second, make
2833 # sure that we don't test with higher precision.
2834 tests = [
2835 # format: (timestamp, operator, expected_len)
2836 (dt0 - timedelta(seconds=1), ">", len0),
2837 (dt0 - timedelta(seconds=1), "<", 0),
2838 (dt1 + timedelta(seconds=1), "<", len0),
2839 (dt1 + timedelta(seconds=1), ">", 0),
2840 ]
2841 for dt, op, expect_len in tests:
2842 dt_str = dt.isoformat(sep=" ")
2844 where = f"ingest_date {op} T'{dt_str}'"
2845 datasets = list(registry.queryDatasets(..., collections=..., where=where))
2846 self.assertEqual(len(datasets), expect_len)
2848 # same with bind using datetime or astropy Time
2849 where = f"ingest_date {op} :ingest_time"
2850 datasets = list(
2851 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt})
2852 )
2853 self.assertEqual(len(datasets), expect_len)
2855 dt_astropy = astropy.time.Time(dt, format="datetime")
2856 datasets = list(
2857 registry.queryDatasets(..., collections=..., where=where, bind={"ingest_time": dt_astropy})
2858 )
2859 self.assertEqual(len(datasets), expect_len)
2861 def testTimespanQueries(self):
2862 """Test query expressions involving timespans."""
2863 butler = self.make_butler()
2864 registry = butler.registry
2865 self.load_data(butler, "ci_hsc-subset.yaml")
2866 # All exposures in the database; mapping from ID to timespan.
2867 visits = {record.id: record.timespan for record in registry.queryDimensionRecords("visit")}
2868 # Just those IDs, sorted (which is also temporal sorting, because HSC
2869 # exposure IDs are monotonically increasing).
2870 ids = sorted(visits.keys())
2871 self.assertEqual(len(ids), 11)
2872 # Pick some quasi-random indexes into `ids` to play with.
2873 i1 = int(len(ids) * 0.1)
2874 i2 = int(len(ids) * 0.3)
2875 i3 = int(len(ids) * 0.6)
2876 i4 = int(len(ids) * 0.8)
2877 # Extract some times from those: just before the beginning of i1 (which
2878 # should be after the end of the exposure before), exactly the
2879 # beginning of i2, just after the beginning of i3 (and before its end),
2880 # and the exact end of i4.
2881 t1 = visits[ids[i1]].begin - astropy.time.TimeDelta(1.0, format="sec")
2882 self.assertGreater(t1, visits[ids[i1 - 1]].end)
2883 t2 = visits[ids[i2]].begin
2884 t3 = visits[ids[i3]].begin + astropy.time.TimeDelta(1.0, format="sec")
2885 self.assertLess(t3, visits[ids[i3]].end)
2886 t4 = visits[ids[i4]].end
2887 # Make sure those are actually in order.
2888 self.assertEqual([t1, t2, t3, t4], sorted([t4, t3, t2, t1]))
2890 bind = {
2891 "t1": t1,
2892 "t2": t2,
2893 "t3": t3,
2894 "t4": t4,
2895 "ts23": Timespan(t2, t3),
2896 }
2898 def query(where):
2899 """Return results as a sorted, deduplicated list of visit IDs.
2901 Parameters
2902 ----------
2903 where : `str`
2904 The WHERE clause for the query.
2905 """
2906 return sorted(
2907 {
2908 dataId["visit"]
2909 for dataId in registry.queryDataIds("visit", instrument="HSC", bind=bind, where=where)
2910 }
2911 )
2913 # Try a bunch of timespan queries, mixing up the bounds themselves,
2914 # where they appear in the expression, and how we get the timespan into
2915 # the expression.
2917 # t1 is before the start of i1, so this should not include i1.
2918 self.assertEqual(ids[:i1], query("visit.timespan OVERLAPS (null, :t1)"))
2919 # t2 is exactly at the start of i2, but ends are exclusive, so these
2920 # should not include i2.
2921 self.assertEqual(ids[i1:i2], query("(:t1, :t2) OVERLAPS visit.timespan"))
2922 # t3 is in the middle of i3, so this should include i3.
2923 self.assertEqual(ids[i2 : i3 + 1], query("visit.timespan OVERLAPS :ts23"))
2924 # This one should not include t3 by the same reasoning.
2925 # t4 is exactly at the end of i4, so this should include i4.
2926 self.assertEqual(ids[i3 : i4 + 1], query(f"visit.timespan OVERLAPS (T'{t3.tai.isot}/tai', :t4)"))
2927 # i4's upper bound of t4 is exclusive so this should not include t4.
2928 self.assertEqual(ids[i4 + 1 :], query("visit.timespan OVERLAPS (:t4, NULL)"))
2930 # Now some timespan vs. time scalar queries.
2931 self.assertEqual(ids[i3 : i3 + 1], query("visit.timespan OVERLAPS :t3"))
2932 self.assertEqual(ids[i3 : i3 + 1], query(f"T'{t3.tai.isot}/tai' OVERLAPS visit.timespan"))
2934 # Empty timespans should not overlap anything.
2935 self.assertEqual([], query("visit.timespan OVERLAPS (:t3, :t2)"))
2937 # Make sure that expanded data IDs include the timespans.
2938 results = list(
2939 registry.queryDataIds(["visit"], dataId={"instrument": "HSC", "visit": ids[1]}).expanded()
2940 )
2941 self.assertEqual(len(results), 1)
2942 visit_timespan = visits[ids[1]]
2943 self.assertEqual(results[0].timespan, visit_timespan)
2944 visit_record = results[0].records["visit"]
2945 assert visit_record is not None
2946 self.assertEqual(visit_record.timespan, visit_timespan)
2947 day_obs_record = results[0].records["day_obs"]
2948 assert day_obs_record is not None
2949 self.assertEqual(day_obs_record.id, 20130617)
2950 self.assertEqual(
2951 day_obs_record.timespan,
2952 Timespan(
2953 astropy.time.Time("2013-06-17T00:00:00", scale="tai"),
2954 astropy.time.Time("2013-06-18T00:00:00", scale="tai"),
2955 ),
2956 )
2958 def testCollectionSummaries(self):
2959 """Test recording and retrieval of collection summaries."""
2960 self.maxDiff = None
2961 butler = self.make_butler()
2962 registry = butler.registry
2963 # Importing datasets from yaml should go through the code path where
2964 # we update collection summaries as we insert datasets.
2965 self.load_data(butler, "base.yaml", "datasets.yaml")
2966 flat = registry.getDatasetType("flat")
2967 expected1 = CollectionSummary()
2968 expected1.dataset_types.add(registry.getDatasetType("bias"))
2969 expected1.add_data_ids(
2970 flat, [DataCoordinate.standardize(instrument="Cam1", universe=registry.dimensions)]
2971 )
2972 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)
2973 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)
2974 # Create a chained collection with both of the imported runs; the
2975 # summary should be the same, because it's a union with itself.
2976 chain = "chain"
2977 registry.registerCollection(chain, CollectionType.CHAINED)
2978 registry.setCollectionChain(chain, ["imported_r", "imported_g"])
2979 self.assertEqual(registry.getCollectionSummary(chain), expected1)
2980 # Associate flats only into a tagged collection and a calibration
2981 # collection to check summaries of those.
2982 tag = "tag"
2983 registry.registerCollection(tag, CollectionType.TAGGED)
2984 registry.associate(tag, registry.queryDatasets(flat, collections="imported_g"))
2985 calibs = "calibs"
2986 registry.registerCollection(calibs, CollectionType.CALIBRATION)
2987 registry.certify(
2988 calibs, registry.queryDatasets(flat, collections="imported_g"), timespan=Timespan(None, None)
2989 )
2990 expected2 = expected1.copy()
2991 expected2.dataset_types.discard("bias")
2992 self.assertEqual(registry.getCollectionSummary(tag), expected2)
2993 self.assertEqual(registry.getCollectionSummary(calibs), expected2)
2994 # Explicitly calling SqlRegistry.refresh() should load those same
2995 # summaries, via a totally different code path.
2996 registry.refresh()
2997 self.assertEqual(registry.getCollectionSummary("imported_g"), expected1)
2998 self.assertEqual(registry.getCollectionSummary("imported_r"), expected1)
2999 self.assertEqual(registry.getCollectionSummary(tag), expected2)
3000 self.assertEqual(registry.getCollectionSummary(calibs), expected2)
3002 def testBindInQueryDatasets(self):
3003 """Test that the bind parameter is correctly forwarded in
3004 queryDatasets recursion.
3005 """
3006 butler = self.make_butler()
3007 registry = butler.registry
3008 # Importing datasets from yaml should go through the code path where
3009 # we update collection summaries as we insert datasets.
3010 self.load_data(butler, "base.yaml", "datasets.yaml")
3011 self.assertEqual(
3012 set(registry.queryDatasets("flat", band="r", collections=...)),
3013 set(
3014 registry.queryDatasets("flat", where="band=:my_band", bind={"my_band": "r"}, collections=...)
3015 ),
3016 )
3018 def testQueryIntRangeExpressions(self):
3019 """Test integer range expressions in ``where`` arguments.
3021 Note that our expressions use inclusive stop values, unlike Python's.
3022 """
3023 butler = self.make_butler()
3024 registry = butler.registry
3025 self.load_data(butler, "base.yaml")
3026 self.assertEqual(
3027 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..2)")),
3028 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 2]},
3029 )
3030 self.assertEqual(
3031 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (1..4:2)")),
3032 {registry.expandDataId(instrument="Cam1", detector=n) for n in [1, 3]},
3033 )
3034 self.assertEqual(
3035 set(registry.queryDataIds(["detector"], instrument="Cam1", where="detector IN (2..4:2)")),
3036 {registry.expandDataId(instrument="Cam1", detector=n) for n in [2, 4]},
3037 )
3039 def testQueryResultSummaries(self):
3040 """Test summary methods like `count`, `any`, and `explain_no_results`
3041 on `DataCoordinateQueryResults` and `DatasetQueryResults`.
3042 """
3043 butler = self.make_butler()
3044 registry = butler.registry
3045 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3046 # Default test dataset has two collections, each with both flats and
3047 # biases. Add a new collection with only biases.
3048 registry.registerCollection("biases", CollectionType.TAGGED)
3049 registry.associate("biases", registry.queryDatasets("bias", collections=["imported_g"]))
3050 # First query yields two results, and involves no postprocessing.
3051 query1 = registry.queryDataIds(["physical_filter"], band="r")
3052 self.assertTrue(query1.any(execute=False, exact=False))
3053 self.assertTrue(query1.any(execute=True, exact=False))
3054 self.assertTrue(query1.any(execute=True, exact=True))
3055 self.assertEqual(query1.count(exact=False), 2)
3056 self.assertEqual(query1.count(exact=True), 2)
3057 self.assertFalse(list(query1.explain_no_results()))
3058 # Second query should yield no results, which we should see when
3059 # we attempt to expand the data ID.
3060 query2 = registry.queryDataIds(["physical_filter"], band="h")
3061 # There's no execute=False, exact=False test here because the behavior
3062 # not something we want to guarantee in this case (and exact=False
3063 # says either answer is legal).
3064 self.assertFalse(query2.any(execute=True, exact=False))
3065 self.assertFalse(query2.any(execute=True, exact=True))
3066 self.assertEqual(query2.count(exact=False), 0)
3067 self.assertEqual(query2.count(exact=True), 0)
3068 # These queries yield no results due to various problems that can be
3069 # spotted prior to execution, yielding helpful diagnostics.
3070 base_query = registry.queryDataIds(["detector", "physical_filter"])
3071 queries_and_snippets = [
3072 (
3073 # Dataset type name doesn't match any existing dataset types.
3074 registry.queryDatasets("nonexistent", collections=...),
3075 ["nonexistent"],
3076 ),
3077 (
3078 # Dataset type object isn't registered.
3079 registry.queryDatasets(
3080 DatasetType(
3081 "nonexistent",
3082 dimensions=["instrument"],
3083 universe=registry.dimensions,
3084 storageClass="Image",
3085 ),
3086 collections=...,
3087 ),
3088 ["nonexistent"],
3089 ),
3090 (
3091 # No datasets of this type in this collection.
3092 registry.queryDatasets("flat", collections=["biases"]),
3093 ["flat", "biases"],
3094 ),
3095 (
3096 # No datasets of this type in this collection.
3097 base_query.findDatasets("flat", collections=["biases"]),
3098 ["flat", "biases"],
3099 ),
3100 (
3101 # No collections matching at all.
3102 registry.queryDatasets("flat", collections="potato*"),
3103 ["potato"],
3104 ),
3105 ]
3106 with self.assertRaises(MissingDatasetTypeError):
3107 # Dataset type name doesn't match any existing dataset types.
3108 list(registry.queryDataIds(["detector"], datasets=["nonexistent"], collections=...))
3109 with self.assertRaises(MissingDatasetTypeError):
3110 # Dataset type name doesn't match any existing dataset types.
3111 registry.queryDimensionRecords("detector", datasets=["nonexistent"], collections=...).any()
3112 with self.assertRaises(DatasetTypeExpressionError):
3113 # queryDimensionRecords does not allow dataset type wildcards.
3114 registry.queryDimensionRecords("detector", datasets=["f*"], collections=...).any()
3115 for query, snippets in queries_and_snippets:
3116 self.assertFalse(query.any(execute=False, exact=False))
3117 self.assertFalse(query.any(execute=True, exact=False))
3118 self.assertFalse(query.any(execute=True, exact=True))
3119 self.assertEqual(query.count(exact=False), 0)
3120 self.assertEqual(query.count(exact=True), 0)
3121 messages = list(query.explain_no_results())
3122 self.assertTrue(messages)
3123 # Want all expected snippets to appear in at least one message.
3124 self.assertTrue(
3125 any(
3126 all(snippet in message for snippet in snippets) for message in query.explain_no_results()
3127 ),
3128 messages,
3129 )
3131 # Wildcards on dataset types are not permitted in queryDataIds.
3132 with self.assertRaises(DatasetTypeExpressionError):
3133 registry.queryDataIds(["detector"], datasets=re.compile("^nonexistent$"), collections=...)
3135 # This query yields four overlaps in the database, but one is filtered
3136 # out in postprocessing. The count queries aren't accurate because
3137 # they don't account for duplication that happens due to an internal
3138 # join against commonSkyPix.
3139 query3 = registry.queryDataIds(["visit", "tract"], instrument="Cam1", skymap="SkyMap1")
3140 self.assertEqual(
3141 {
3142 DataCoordinate.standardize(
3143 instrument="Cam1",
3144 skymap="SkyMap1",
3145 visit=v,
3146 tract=t,
3147 universe=registry.dimensions,
3148 )
3149 for v, t in [(1, 0), (2, 0), (2, 1)]
3150 },
3151 set(query3),
3152 )
3153 self.assertTrue(query3.any(execute=False, exact=False))
3154 self.assertTrue(query3.any(execute=True, exact=False))
3155 self.assertTrue(query3.any(execute=True, exact=True))
3156 self.assertGreaterEqual(query3.count(exact=False), 3)
3157 self.assertGreaterEqual(query3.count(exact=True, discard=True), 3)
3158 self.assertFalse(list(query3.explain_no_results()))
3159 # This query yields overlaps in the database, but all are filtered
3160 # out in postprocessing. The count queries again aren't very useful.
3161 # We have to use `where=` here to avoid an optimization that
3162 # (currently) skips the spatial postprocess-filtering because it
3163 # recognizes that no spatial join is necessary. That's not ideal, but
3164 # fixing it is out of scope for this ticket.
3165 query4 = registry.queryDataIds(
3166 ["visit", "tract"],
3167 instrument="Cam1",
3168 skymap="SkyMap1",
3169 where="visit=1 AND detector=1 AND tract=0 AND patch=4",
3170 )
3171 self.assertFalse(set(query4))
3172 self.assertTrue(query4.any(execute=False, exact=False))
3173 self.assertTrue(query4.any(execute=True, exact=False))
3174 self.assertFalse(query4.any(execute=True, exact=True))
3175 self.assertGreaterEqual(query4.count(exact=False), 1)
3176 self.assertEqual(query4.count(exact=True, discard=True), 0)
3177 # This query should yield results from one dataset type but not the
3178 # other, which is not registered.
3179 query5 = registry.queryDatasets(["bias", "nonexistent"], collections=["biases"])
3180 self.assertTrue(set(query5))
3181 self.assertTrue(query5.any(execute=False, exact=False))
3182 self.assertTrue(query5.any(execute=True, exact=False))
3183 self.assertTrue(query5.any(execute=True, exact=True))
3184 self.assertGreaterEqual(query5.count(exact=False), 1)
3185 self.assertGreaterEqual(query5.count(exact=True), 1)
3186 # This query applies a selection that yields no results, fully in the
3187 # database. Explaining why it fails involves traversing the relation
3188 # tree and running a LIMIT 1 query at each level that has the potential
3189 # to remove rows.
3190 query6 = registry.queryDimensionRecords(
3191 "detector", where="detector.purpose = 'no-purpose'", instrument="Cam1"
3192 )
3193 self.assertEqual(query6.count(exact=True), 0)
3194 self.assertFalse(query6.any())
3196 def testQueryDataIdsExpressionError(self):
3197 """Test error checking of 'where' expressions in queryDataIds."""
3198 butler = self.make_butler()
3199 registry = butler.registry
3200 self.load_data(butler, "base.yaml")
3201 bind = {"time": astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")}
3202 # The diagnostics raised are slightly different between the old query
3203 # system (ValueError, first error string) and the new query system
3204 # (InvalidQueryError, second error string).
3205 with self.assertRaisesRegex(
3206 (LookupError, InvalidQueryError),
3207 r"(No dimension element with name 'foo' in 'foo\.bar'\.)|(Unrecognized identifier 'foo.bar')",
3208 ):
3209 list(registry.queryDataIds(["detector"], where="foo.bar = 12"))
3210 with self.assertRaisesRegex(
3211 (LookupError, InvalidQueryError),
3212 "(Dimension element name cannot be inferred in this context.)"
3213 "|(Unrecognized identifier 'timespan')",
3214 ):
3215 list(registry.queryDataIds(["detector"], where="timespan.end < :time", bind=bind))
3217 def testQueryDataIdsOrderBy(self):
3218 """Test order_by and limit on result returned by queryDataIds()."""
3219 butler = self.make_butler()
3220 registry = butler.registry
3221 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3223 def do_query(dimensions=("visit", "tract"), datasets=None, collections=None):
3224 return registry.queryDataIds(
3225 dimensions, datasets=datasets, collections=collections, instrument="Cam1", skymap="SkyMap1"
3226 )
3228 Test = namedtuple(
3229 "testQueryDataIdsOrderByTest",
3230 ("order_by", "keys", "result", "limit", "datasets", "collections"),
3231 defaults=(None, None, None),
3232 )
3234 test_data = [
3235 Test("tract,visit", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3236 Test("-tract,visit", "tract,visit", ((1, 2), (0, 1), (0, 2))),
3237 Test("tract,-visit", "tract,visit", ((0, 2), (0, 1), (1, 2))),
3238 Test("-tract,-visit", "tract,visit", ((1, 2), (0, 2), (0, 1))),
3239 Test("tract.id,visit.id", "tract,visit", ((0, 1),), limit=(1,)),
3240 Test("-tract,-visit", "tract,visit", ((1, 2),), limit=(1,)),
3241 Test("tract,visit.exposure_time", "tract,visit", ((0, 2), (0, 1), (1, 2))),
3242 Test("-tract,-visit.exposure_time", "tract,visit", ((1, 2), (0, 1), (0, 2))),
3243 Test("tract,-exposure_time", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3244 Test("tract,visit.name", "tract,visit", ((0, 1), (0, 2), (1, 2))),
3245 Test(
3246 "tract,-visit.timespan.begin,visit.timespan.end",
3247 "tract,visit",
3248 ((0, 2), (0, 1), (1, 2)),
3249 ),
3250 Test("visit.day_obs,exposure.day_obs", "visit,exposure", ()),
3251 Test("visit.timespan.begin,-exposure.timespan.begin", "visit,exposure", ()),
3252 Test(
3253 "tract,detector",
3254 "tract,detector",
3255 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3256 datasets="flat",
3257 collections="imported_r",
3258 ),
3259 Test(
3260 "tract,detector.full_name",
3261 "tract,detector",
3262 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3263 datasets="flat",
3264 collections="imported_r",
3265 ),
3266 Test(
3267 "tract,detector.raft,detector.name_in_raft",
3268 "tract,detector",
3269 ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), (1, 2), (1, 3), (1, 4)),
3270 datasets="flat",
3271 collections="imported_r",
3272 ),
3273 ]
3275 for test in test_data:
3276 with self.subTest(test=test):
3277 order_by = test.order_by.split(",")
3278 keys = test.keys.split(",")
3279 query = do_query(keys, test.datasets, test.collections).order_by(*order_by)
3280 if test.limit is not None:
3281 query = query.limit(*test.limit)
3282 dataIds = tuple(tuple(dataId[k] for k in keys) for dataId in query)
3283 self.assertEqual(dataIds, test.result)
3285 # and materialize
3286 query = do_query(keys).order_by(*order_by)
3287 if test.limit is not None:
3288 query = query.limit(*test.limit)
3290 # Test exceptions for errors in a name.
3291 # Many of these raise slightly different diagnostics in the old query
3292 # system (ValueError, first error string) than the new query system
3293 # (InvalidQueryError, second error string).
3294 for order_by in ("", "-"):
3295 with self.assertRaisesRegex((ValueError, InvalidQueryError), "Empty dimension name in ORDER BY"):
3296 list(do_query().order_by(order_by))
3298 for order_by in ("undimension.name", "-undimension.name"):
3299 with self.assertRaisesRegex(
3300 (ValueError, InvalidQueryError),
3301 "(Unknown dimension element 'undimension')|(Unrecognized identifier 'undimension.name')",
3302 ):
3303 list(do_query().order_by(order_by))
3305 for order_by in ("attract", "-attract"):
3306 with self.assertRaisesRegex(
3307 (ValueError, InvalidQueryError),
3308 "(Metadata 'attract' cannot be found in any dimension)|(Unrecognized identifier 'attract')",
3309 ):
3310 list(do_query().order_by(order_by))
3312 with self.assertRaisesRegex(
3313 (ValueError, InvalidQueryError),
3314 "(Metadata 'exposure_time' exists in more than one dimension)"
3315 "|(Ambiguous identifier 'exposure_time' matches multiple fields)",
3316 ):
3317 list(do_query(("exposure", "visit")).order_by("exposure_time"))
3319 with self.assertRaisesRegex(
3320 (ValueError, InvalidQueryError),
3321 r"(Timespan exists in more than one dimension element \(day_obs, exposure, visit\); "
3322 r"qualify timespan with specific dimension name\.)|"
3323 r"(Ambiguous identifier 'timespan' matches multiple fields)",
3324 ):
3325 list(do_query(("exposure", "visit")).order_by("timespan.begin"))
3327 with self.assertRaisesRegex(
3328 (ValueError, InvalidQueryError),
3329 "(Cannot find any temporal dimension element for 'timespan.begin')"
3330 "|(Unrecognized identifier 'timespan')",
3331 ):
3332 list(do_query("tract").order_by("timespan.begin"))
3334 with self.assertRaisesRegex(
3335 (ValueError, InvalidQueryError),
3336 "(Cannot use 'timespan.begin' with non-temporal element)"
3337 "|(Unrecognized field 'timespan' for tract)",
3338 ):
3339 list(do_query("tract").order_by("tract.timespan.begin"))
3341 with self.assertRaisesRegex(
3342 (ValueError, InvalidQueryError),
3343 "(Field 'name' does not exist in 'tract')|(Unrecognized field 'name' for tract.)",
3344 ):
3345 list(do_query("tract").order_by("tract.name"))
3347 with self.assertRaisesRegex(
3348 (ValueError, InvalidQueryError),
3349 r"(Unknown dimension element 'timestamp'; perhaps you meant 'timespan.begin'\?)"
3350 r"|(Unrecognized identifier 'timestamp.begin')",
3351 ):
3352 list(do_query("visit").order_by("timestamp.begin"))
3354 def testQueryDataIdsGovernorExceptions(self):
3355 """Test exceptions raised by queryDataIds() for incorrect governors."""
3356 butler = self.make_butler()
3357 registry = butler.registry
3358 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3360 def do_query(dimensions, dataId=None, where="", bind=None, **kwargs):
3361 return registry.queryDataIds(dimensions, dataId=dataId, where=where, bind=bind, **kwargs)
3363 Test = namedtuple(
3364 "testQueryDataIdExceptionsTest",
3365 ("dimensions", "dataId", "where", "bind", "kwargs", "exception", "count"),
3366 defaults=(None, None, None, {}, None, 0),
3367 )
3369 test_data = (
3370 Test("tract,visit", count=3),
3371 Test("tract,visit", kwargs={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),
3372 Test("tract,visit", kwargs={"instrument": "Cam2", "skymap": "SkyMap1"}, count=0),
3373 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap1"}, count=3),
3374 Test("tract,visit", dataId={"instrument": "Cam1", "skymap": "SkyMap2"}, count=0),
3375 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap1'", count=3),
3376 Test("tract,visit", where="instrument='Cam1' AND skymap='SkyMap5'", count=0),
3377 Test(
3378 "tract,visit",
3379 where="instrument=:cam AND skymap=:map",
3380 bind={"cam": "Cam1", "map": "SkyMap1"},
3381 count=3,
3382 ),
3383 Test(
3384 "tract,visit",
3385 where="instrument=:cam AND skymap=:map",
3386 bind={"cam": "Cam", "map": "SkyMap"},
3387 count=0,
3388 ),
3389 )
3391 for test in test_data:
3392 print(test)
3393 dimensions = test.dimensions.split(",")
3394 if test.exception:
3395 with self.assertRaises(test.exception):
3396 with ExitStack() as stack:
3397 if test.exception == DataIdValueError:
3398 stack.enter_context(self.assertWarns(FutureWarning))
3399 do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs).count()
3400 else:
3401 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3402 print(list(query))
3403 self.assertEqual(query.count(discard=True), test.count)
3405 # and materialize
3406 if test.exception:
3407 with self.assertRaises(test.exception):
3408 with ExitStack() as stack:
3409 if test.exception == DataIdValueError:
3410 stack.enter_context(self.assertWarns(FutureWarning))
3411 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3412 with query.materialize() as materialized:
3413 materialized.count(discard=True)
3414 else:
3415 query = do_query(dimensions, test.dataId, test.where, bind=test.bind, **test.kwargs)
3416 with query.materialize() as materialized:
3417 self.assertEqual(materialized.count(discard=True), test.count)
3419 def testQueryDimensionRecordsOrderBy(self):
3420 """Test order_by and limit on result returned by
3421 queryDimensionRecords().
3422 """
3423 butler = self.make_butler()
3424 registry = butler.registry
3425 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3427 def do_query(element, datasets=None, collections=None):
3428 return registry.queryDimensionRecords(
3429 element, instrument="Cam1", datasets=datasets, collections=collections
3430 )
3432 query = do_query("detector")
3433 self.assertEqual(len(list(query)), 4)
3435 Test = namedtuple(
3436 "testQueryDataIdsOrderByTest",
3437 ("element", "order_by", "result", "limit", "datasets", "collections"),
3438 defaults=(None, None, None),
3439 )
3441 test_data = [
3442 Test("detector", "detector", (1, 2, 3, 4)),
3443 Test("detector", "-detector", (4, 3, 2, 1)),
3444 Test("detector", "raft,-name_in_raft", (2, 1, 4, 3)),
3445 Test("detector", "-detector.purpose", (4,), limit=(1,)),
3446 Test("visit", "visit", (1, 2)),
3447 Test("visit", "-visit.id", (2, 1)),
3448 Test("visit", "zenith_angle", (1, 2)),
3449 Test("visit", "-visit.name", (2, 1)),
3450 Test("visit", "day_obs,-visit.timespan.begin", (2, 1)),
3451 ]
3453 def do_test(test: Test):
3454 order_by = test.order_by.split(",")
3455 query = do_query(test.element).order_by(*order_by)
3456 if test.limit is not None:
3457 query = query.limit(*test.limit)
3458 dataIds = tuple(rec.id for rec in query)
3459 self.assertEqual(dataIds, test.result)
3461 for test in test_data:
3462 do_test(test)
3464 # errors in a name
3465 for order_by in ("", "-"):
3466 with self.assertRaisesRegex(
3467 (ValueError, InvalidQueryError),
3468 "(Empty dimension name in ORDER BY)|(Unrecognized identifier)",
3469 ):
3470 list(do_query("detector").order_by(order_by))
3472 for order_by in ("undimension.name", "-undimension.name"):
3473 with self.assertRaisesRegex(
3474 (ValueError, InvalidQueryError),
3475 "(Element name mismatch: 'undimension')|(Unrecognized identifier)",
3476 ):
3477 list(do_query("detector").order_by(order_by))
3479 for order_by in ("attract", "-attract"):
3480 with self.assertRaisesRegex(
3481 (ValueError, InvalidQueryError),
3482 "(Field 'attract' does not exist in 'detector'.)|(Unrecognized identifier)",
3483 ):
3484 list(do_query("detector").order_by(order_by))
3486 for order_by in ("timestamp.begin", "-timestamp.begin"):
3487 with self.assertRaisesRegex(
3488 (ValueError, InvalidQueryError),
3489 r"(Element name mismatch: 'timestamp' instead of 'visit'; "
3490 r"perhaps you meant 'timespan.begin'\?)"
3491 r"|(Unrecognized identifier)",
3492 ):
3493 list(do_query("visit").order_by(order_by))
3495 def testQueryDimensionRecordsExceptions(self):
3496 """Test exceptions raised by queryDimensionRecords()."""
3497 butler = self.make_butler()
3498 registry = butler.registry
3499 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
3501 result = registry.queryDimensionRecords("detector")
3502 self.assertEqual(result.count(), 4)
3503 result = registry.queryDimensionRecords("detector", instrument="Cam1")
3504 self.assertEqual(result.count(), 4)
3505 result = registry.queryDimensionRecords("detector", dataId={"instrument": "Cam1"})
3506 self.assertEqual(result.count(), 4)
3508 # Test that values specified in kwargs override those specified in
3509 # dataId.
3510 result = registry.queryDimensionRecords(
3511 "detector", dataId={"instrument": "NotCam1"}, instrument="Cam1"
3512 )
3513 self.assertEqual(result.count(), 4)
3515 result = registry.queryDimensionRecords("detector", where="instrument='Cam1'")
3516 self.assertEqual(result.count(), 4)
3517 result = registry.queryDimensionRecords("detector", where="instrument=:instr", bind={"instr": "Cam1"})
3518 self.assertTrue(result.any())
3519 self.assertEqual(result.count(), 4)
3521 def testDatasetConstrainedDimensionRecordQueries(self):
3522 """Test that queryDimensionRecords works even when given a dataset
3523 constraint whose dimensions extend beyond the requested dimension
3524 element's.
3525 """
3526 butler = self.make_butler()
3527 registry = butler.registry
3528 self.load_data(butler, "base.yaml", "datasets.yaml")
3529 # Query for physical_filter dimension records, using a dataset that
3530 # has both physical_filter and dataset dimensions.
3531 records = registry.queryDimensionRecords(
3532 "physical_filter",
3533 datasets=["flat"],
3534 collections="imported_r",
3535 )
3536 self.assertEqual({record.name for record in records}, {"Cam1-R1", "Cam1-R2"})
3537 # Trying to constrain by all dataset types is an error.
3538 with self.assertRaises(TypeError):
3539 list(registry.queryDimensionRecords("physical_filter", datasets=..., collections="imported_r"))
3541 def testSkyPixDatasetQueries(self):
3542 """Test that we can build queries involving skypix dimensions as long
3543 as a dataset type that uses those dimensions is included.
3544 """
3545 butler = self.make_butler()
3546 registry = butler.registry
3547 self.load_data(butler, "base.yaml")
3548 dataset_type = DatasetType(
3549 "a", dimensions=["htm7", "instrument"], universe=registry.dimensions, storageClass="int"
3550 )
3551 registry.registerDatasetType(dataset_type)
3552 run = "r"
3553 registry.registerRun(run)
3554 # First try queries where there are no datasets; the concern is whether
3555 # we can even build and execute these queries without raising, even
3556 # when "doomed" query shortcuts are in play.
3557 self.assertFalse(
3558 list(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run))
3559 )
3560 self.assertFalse(list(registry.queryDatasets(dataset_type, collections=run)))
3561 # Now add a dataset and see that we can get it back.
3562 htm7 = registry.dimensions.skypix["htm"][7].pixelization
3563 data_id = registry.expandDataId(instrument="Cam1", htm7=htm7.universe()[0][0])
3564 (ref,) = registry.insertDatasets(dataset_type, [data_id], run=run)
3565 self.assertEqual(
3566 set(registry.queryDataIds(["htm7", "instrument"], datasets=dataset_type, collections=run)),
3567 {data_id},
3568 )
3569 self.assertEqual(set(registry.queryDatasets(dataset_type, collections=run)), {ref})
3571 def testDatasetIdFactory(self):
3572 """Simple test for DatasetIdFactory, mostly to catch potential changes
3573 in its API.
3574 """
3575 butler = self.make_butler()
3576 registry = butler.registry
3577 factory = DatasetIdFactory()
3578 dataset_type = DatasetType(
3579 "datasetType",
3580 dimensions=["detector", "instrument"],
3581 universe=registry.dimensions,
3582 storageClass="int",
3583 )
3584 run = "run"
3585 data_id = DataCoordinate.standardize(
3586 instrument="Cam1", detector=1, dimensions=dataset_type.dimensions
3587 )
3589 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.UNIQUE)
3590 self.assertIsInstance(datasetId, uuid.UUID)
3591 self.assertEqual(datasetId.version, 7)
3593 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE)
3594 self.assertIsInstance(datasetId, uuid.UUID)
3595 self.assertEqual(datasetId.version, 5)
3597 datasetId = factory.makeDatasetId(run, dataset_type, data_id, DatasetIdGenEnum.DATAID_TYPE_RUN)
3598 self.assertIsInstance(datasetId, uuid.UUID)
3599 self.assertEqual(datasetId.version, 5)
3601 def testExposureQueries(self):
3602 """Test query methods using arguments sourced from the exposure log
3603 service.
3605 The most complete test dataset currently available to daf_butler tests
3606 is ci_hsc-subset.yaml export , but that does not have 'exposure'
3607 dimension records. So in this test we need to translate queries that
3608 originally used the exposure dimension to use the (very similar) visit
3609 dimension instead.
3610 """
3611 butler = self.make_butler()
3612 registry = butler.registry
3613 self.load_data(butler, "ci_hsc-subset.yaml")
3614 self.assertEqual(
3615 [
3616 record.id
3617 for record in registry.queryDimensionRecords("visit", instrument="HSC")
3618 .order_by("visit")
3619 .limit(5)
3620 ],
3621 [903334, 903336, 903338, 903342, 903344],
3622 )
3623 self.assertEqual(
3624 [
3625 data_id["visit"]
3626 for data_id in registry.queryDataIds(["visit"], instrument="HSC").order_by("visit").limit(5)
3627 ],
3628 [903334, 903336, 903338, 903342, 903344],
3629 )
3630 self.assertEqual(
3631 [
3632 record.id
3633 for record in registry.queryDimensionRecords("detector", instrument="HSC")
3634 .order_by("full_name")
3635 .limit(5)
3636 ],
3637 [25, 24, 23, 22, 18],
3638 )
3639 self.assertEqual(
3640 [
3641 data_id["detector"]
3642 for data_id in registry.queryDataIds(["detector"], instrument="HSC")
3643 .order_by("full_name")
3644 .limit(5)
3645 ],
3646 [25, 24, 23, 22, 18],
3647 )
3649 def test_long_query_names(self) -> None:
3650 """Test that queries involving very long names are handled correctly.
3652 This is especially important for PostgreSQL, which truncates symbols
3653 longer than 64 chars, but it's worth testing for all DBs.
3654 """
3655 butler = self.make_butler()
3656 registry = butler.registry
3657 name = "abcd" * 17
3658 registry.registerDatasetType(
3659 DatasetType(
3660 name,
3661 dimensions=(),
3662 storageClass="Exposure",
3663 universe=registry.dimensions,
3664 )
3665 )
3666 # Need to search more than one collection actually containing a
3667 # matching dataset to avoid optimizations that sidestep bugs due to
3668 # truncation by making findFirst=True a no-op.
3669 run1 = "run1"
3670 registry.registerRun(run1)
3671 run2 = "run2"
3672 registry.registerRun(run2)
3673 (ref1,) = registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run1)
3674 registry.insertDatasets(name, [DataCoordinate.make_empty(registry.dimensions)], run2)
3675 self.assertEqual(
3676 set(registry.queryDatasets(name, collections=[run1, run2], findFirst=True)),
3677 {ref1},
3678 )
3680 def test_skypix_constraint_queries(self) -> None:
3681 """Test queries spatially constrained by a skypix data ID."""
3682 butler = self.make_butler()
3683 registry = butler.registry
3684 self.load_data(butler, "base.yaml", "spatial.yaml")
3685 patch_regions = {
3686 (data_id["tract"], data_id["patch"]): data_id.region
3687 for data_id in registry.queryDataIds(["patch"]).expanded()
3688 }
3689 skypix_dimension: SkyPixDimension = registry.dimensions["htm11"]
3690 # This check ensures the test doesn't become trivial due to a config
3691 # change; if it does, just pick a different HTML level.
3692 self.assertNotEqual(skypix_dimension, registry.dimensions.commonSkyPix)
3693 # Gather all skypix IDs that definitely overlap at least one of these
3694 # patches.
3695 relevant_skypix_ids = lsst.sphgeom.RangeSet()
3696 for patch_region in patch_regions.values():
3697 relevant_skypix_ids |= skypix_dimension.pixelization.interior(patch_region)
3698 # Look for a "nontrivial" skypix_id that overlaps at least one patch
3699 # and does not overlap at least one other patch.
3700 for skypix_id in itertools.chain.from_iterable(
3701 range(begin, end) for begin, end in relevant_skypix_ids
3702 ):
3703 skypix_region = skypix_dimension.pixelization.pixel(skypix_id)
3704 overlapping_patches = {
3705 patch_key
3706 for patch_key, patch_region in patch_regions.items()
3707 if not patch_region.isDisjointFrom(skypix_region)
3708 }
3709 if overlapping_patches and overlapping_patches != patch_regions.keys():
3710 break
3711 else:
3712 raise RuntimeError("Could not find usable skypix ID for this dimension configuration.")
3713 # Test that a three-way join that includes the common skypix system in
3714 # the dimensions doesn't generate redundant join terms in the query.
3715 with self.assertRaises(InvalidQueryError):
3716 set(
3717 registry.queryDataIds(
3718 ["tract", "visit", "htm7"], skymap="SkyMap1", instrument="Cam1"
3719 ).expanded()
3720 )
3722 def test_spatial_constraint_queries(self) -> None:
3723 """Test queries in which one spatial dimension in the constraint (data
3724 ID or ``where`` string) constrains a different spatial dimension in the
3725 query result columns.
3726 """
3727 butler = self.make_butler()
3728 registry = butler.registry
3729 self.load_data(butler, "base.yaml", "spatial.yaml")
3730 patch_regions = {
3731 (data_id["tract"], data_id["patch"]): data_id.region
3732 for data_id in registry.queryDataIds(["patch"]).expanded()
3733 }
3734 observation_regions = {
3735 (data_id["visit"], data_id["detector"]): data_id.region
3736 for data_id in registry.queryDataIds(["visit", "detector"]).expanded()
3737 }
3738 all_combos = {
3739 (patch_key, observation_key)
3740 for patch_key, observation_key in itertools.product(patch_regions, observation_regions)
3741 }
3742 overlapping_combos = {
3743 (patch_key, observation_key)
3744 for patch_key, observation_key in all_combos
3745 if not patch_regions[patch_key].isDisjointFrom(observation_regions[observation_key])
3746 }
3747 # Check a direct spatial join with no constraint first.
3748 self.assertEqual(
3749 {
3750 ((data_id["tract"], data_id["patch"]), (data_id["visit"], data_id["detector"]))
3751 for data_id in registry.queryDataIds(["patch", "visit", "detector"])
3752 },
3753 overlapping_combos,
3754 )
3755 overlaps_by_patch: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)
3756 overlaps_by_observation: defaultdict[tuple[int, int], set[tuple[str, str]]] = defaultdict(set)
3757 for patch_key, observation_key in overlapping_combos:
3758 overlaps_by_patch[patch_key].add(observation_key)
3759 overlaps_by_observation[observation_key].add(patch_key)
3760 # Find patches and observations that overlap at least one of the other
3761 # but not all of the other.
3762 nontrivial_patch = next(
3763 iter(
3764 patch_key
3765 for patch_key, observation_keys in overlaps_by_patch.items()
3766 if observation_keys and observation_keys != observation_regions.keys()
3767 )
3768 )
3769 nontrivial_observation = next(
3770 iter(
3771 observation_key
3772 for observation_key, patch_keys in overlaps_by_observation.items()
3773 if patch_keys and patch_keys != patch_regions.keys()
3774 )
3775 )
3776 # Use the nontrivial patches and observations as constraints on the
3777 # other dimensions in various ways, first via a 'where' expression.
3778 # It's better in general to us 'bind' instead of f-strings, but these
3779 # all integers so there are no quoting concerns.
3780 self.assertEqual(
3781 {
3782 (data_id["visit"], data_id["detector"])
3783 for data_id in registry.queryDataIds(
3784 ["visit", "detector"],
3785 where=f"tract={nontrivial_patch[0]} AND patch={nontrivial_patch[1]}",
3786 skymap="SkyMap1",
3787 )
3788 },
3789 overlaps_by_patch[nontrivial_patch],
3790 )
3791 self.assertEqual(
3792 {
3793 (data_id["tract"], data_id["patch"])
3794 for data_id in registry.queryDataIds(
3795 ["patch"],
3796 where=f"visit={nontrivial_observation[0]} AND detector={nontrivial_observation[1]}",
3797 instrument="Cam1",
3798 )
3799 },
3800 overlaps_by_observation[nontrivial_observation],
3801 )
3802 # and then via the dataId argument.
3803 self.assertEqual(
3804 {
3805 (data_id["visit"], data_id["detector"])
3806 for data_id in registry.queryDataIds(
3807 ["visit", "detector"],
3808 dataId={
3809 "tract": nontrivial_patch[0],
3810 "patch": nontrivial_patch[1],
3811 },
3812 skymap="SkyMap1",
3813 )
3814 },
3815 overlaps_by_patch[nontrivial_patch],
3816 )
3817 self.assertEqual(
3818 {
3819 (data_id["tract"], data_id["patch"])
3820 for data_id in registry.queryDataIds(
3821 ["patch"],
3822 dataId={
3823 "visit": nontrivial_observation[0],
3824 "detector": nontrivial_observation[1],
3825 },
3826 instrument="Cam1",
3827 )
3828 },
3829 overlaps_by_observation[nontrivial_observation],
3830 )
3832 def test_query_empty_collections(self) -> None:
3833 """Test for registry query methods with empty collections. The methods
3834 should return empty result set (or None when applicable) and provide
3835 "doomed" diagnostics.
3836 """
3837 butler = self.make_butler()
3838 registry = butler.registry
3839 self.load_data(butler, "base.yaml", "datasets.yaml")
3841 # Tests for registry.findDataset()
3842 with self.assertRaises(NoDefaultCollectionError):
3843 registry.findDataset("bias", instrument="Cam1", detector=1)
3844 self.assertIsNotNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=...))
3845 self.assertIsNone(registry.findDataset("bias", instrument="Cam1", detector=1, collections=[]))
3847 # Tests for registry.queryDatasets()
3848 with self.assertRaises(NoDefaultCollectionError):
3849 registry.queryDatasets("bias")
3850 self.assertTrue(list(registry.queryDatasets("bias", collections=...)))
3852 result = registry.queryDatasets("bias", collections=[])
3853 self.assertEqual(len(list(result)), 0)
3854 messages = list(result.explain_no_results())
3855 self.assertTrue(messages)
3856 self.assertTrue(any("because collection list is empty" in message for message in messages))
3858 # Tests for registry.queryDataIds()
3859 with self.assertRaises(NoDefaultCollectionError):
3860 registry.queryDataIds("detector", datasets="bias")
3861 self.assertTrue(list(registry.queryDataIds("detector", datasets="bias", collections=...)))
3863 result = registry.queryDataIds("detector", datasets="bias", collections=[])
3864 self.assertEqual(len(list(result)), 0)
3865 messages = list(result.explain_no_results())
3866 self.assertTrue(messages)
3867 self.assertTrue(any("because collection list is empty" in message for message in messages))
3869 # Tests for registry.queryDimensionRecords()
3870 with self.assertRaises(NoDefaultCollectionError):
3871 registry.queryDimensionRecords("detector", datasets="bias")
3872 self.assertTrue(list(registry.queryDimensionRecords("detector", datasets="bias", collections=...)))
3874 result = registry.queryDimensionRecords("detector", datasets="bias", collections=[])
3875 self.assertEqual(len(list(result)), 0)
3876 messages = list(result.explain_no_results())
3877 self.assertTrue(messages)
3878 self.assertTrue(any("because collection list is empty" in message for message in messages))
3880 def test_dataset_followup_spatial_joins(self) -> None:
3881 """Test queryDataIds(...).findRelatedDatasets(...) where a spatial join
3882 is involved.
3883 """
3884 butler = self.make_butler()
3885 registry = butler.registry
3886 self.load_data(butler, "base.yaml", "spatial.yaml")
3887 pvi_dataset_type = DatasetType(
3888 "pvi", {"visit", "detector"}, storageClass="StructuredDataDict", universe=registry.dimensions
3889 )
3890 registry.registerDatasetType(pvi_dataset_type)
3891 collection = "datasets"
3892 registry.registerRun(collection)
3893 (pvi1,) = registry.insertDatasets(
3894 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 1}], run=collection
3895 )
3896 (pvi2,) = registry.insertDatasets(
3897 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 2}], run=collection
3898 )
3899 (pvi3,) = registry.insertDatasets(
3900 pvi_dataset_type, [{"instrument": "Cam1", "visit": 1, "detector": 3}], run=collection
3901 )
3902 self.assertEqual(
3903 set(
3904 registry.queryDataIds(["patch"], skymap="SkyMap1", tract=0)
3905 .expanded()
3906 .findRelatedDatasets("pvi", [collection])
3907 ),
3908 {
3909 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi1),
3910 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=0), pvi2),
3911 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=1), pvi2),
3912 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi1),
3913 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi2),
3914 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=2), pvi3),
3915 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=3), pvi2),
3916 (registry.expandDataId(skymap="SkyMap1", tract=0, patch=4), pvi3),
3917 },
3918 )
3920 def test_expanded_data_id_queries(self) -> None:
3921 """Tests for basic functionality of expanded() on queryDataIds and
3922 queryDatasets.
3923 """
3924 butler = self.make_butler()
3925 registry = butler.registry
3926 self.load_data(butler, "base.yaml", "spatial.yaml", "datasets.yaml")
3928 result_obj = (
3929 registry.queryDataIds(["visit"], where="instrument = 'Cam1' and (visit.id = 1 or visit.id = 2)")
3930 .expanded()
3931 .order_by("visit.id")
3932 )
3933 self.assertTrue(result_obj.hasRecords())
3934 visits = list(result_obj)
3935 self.assertEqual(len(visits), 2)
3937 self.assertEqual(visits[0]["visit"], 1)
3938 self.assertEqual(visits[1]["visit"], 2)
3939 self.assertEqual(visits[0].records["visit"].exposure_time, 60.0)
3940 self.assertEqual(visits[1].records["visit"].exposure_time, 45.0)
3941 # physical_filter is a "cacheable" dimension, so its records are loaded
3942 # from local cache rather than being part of the DB rows.
3943 self.assertEqual(visits[0].records["physical_filter"].band, "g")
3944 self.assertEqual(visits[1].records["physical_filter"].band, "r")
3946 # Make sure that we can fetch nulls in dimension records
3947 registry.insertDimensionData(
3948 "detector",
3949 {
3950 "instrument": "Cam1",
3951 "id": 5,
3952 "raft": "Z",
3953 "name_in_raft": "z",
3954 "full_name": "Zz",
3955 "purpose": None,
3956 },
3957 )
3958 detectors = list(
3959 registry.queryDataIds("detector", dataId={"instrument": "Cam1", "detector": 5}).expanded()
3960 )
3961 self.assertIsNone(detectors[0].records["detector"].purpose)
3963 datasets_query = registry.queryDatasets(
3964 "flat", collections="imported_g", where="instrument = 'Cam1' and detector <= 3"
3965 ).expanded()
3966 datasets = list(datasets_query)
3967 datasets.sort(key=lambda ref: ref.dataId["detector"])
3968 self.assertEqual(len(datasets), 2)
3969 self.assertEqual(datasets[0].id, uuid.UUID("60c8a65c-7290-4c38-b1de-e3b1cdcf872d"))
3970 self.assertEqual(datasets[1].id, uuid.UUID("84239e7f-c41f-46d5-97b9-a27976b98ceb"))
3971 # All of the dimensions for flat are "cached" dimensions.
3972 self.assertEqual(datasets[0].dataId.records["detector"].full_name, "Ab")
3973 self.assertEqual(datasets[1].dataId.records["detector"].full_name, "Ba")
3974 self.assertEqual(datasets[0].dataId.records["instrument"].visit_system, 1)
3975 assert isinstance(datasets_query, ParentDatasetQueryResults)
3976 data_ids = list(datasets_query.dataIds)
3977 data_ids.sort(key=lambda data_id: data_id["detector"])
3978 self.assertEqual(len(data_ids), 2)
3979 self.assertEqual(data_ids[0].records["detector"].full_name, "Ab")
3980 self.assertEqual(data_ids[1].records["detector"].full_name, "Ba")
3981 self.assertEqual(data_ids[0].records["instrument"].visit_system, 1)
3983 # None of the datasets in the test data include any uncached
3984 # dimensions, so we have to set one up.
3985 registry.registerDatasetType(DatasetType("test", ["visit"], "int", universe=registry.dimensions))
3986 registry.insertDatasets("test", [{"instrument": "Cam1", "visit": 1}], run="imported_g")
3987 ref = list(registry.queryDatasets("test", collections="imported_g").expanded())[0]
3988 self.assertEqual(ref.dataId.records["visit"].zenith_angle, 5.0)
3989 self.assertEqual(ref.dataId.records["physical_filter"].band, "g")
3990 self.assertEqual(
3991 ref.dataId.timespan,
3992 Timespan(
3993 begin=astropy.time.Time("2021-09-09 03:00:00.000000000", scale="tai"),
3994 end=astropy.time.Time("2021-09-09 03:01:00.000000000", scale="tai"),
3995 ),
3996 )
3998 def test_collection_summary(self) -> None:
3999 """Test for collection summary methods."""
4000 butler = self.make_butler()
4001 registry = butler.registry
4002 self.load_data(butler, "base.yaml", "datasets.yaml", "spatial.yaml")
4004 # Add one more dataset type, just for its existence to trigger a bug
4005 # in `associate` (DM-44311).
4006 test_dataset_type = DatasetType("test", ["tract", "patch"], "int", universe=registry.dimensions)
4007 registry.registerDatasetType(test_dataset_type)
4009 # Check for what has been imported.
4010 summary = registry.getCollectionSummary("imported_g")
4011 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
4012 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
4014 # Make a tagged collection and associate some datasets.
4015 tagged_coll = "tagged"
4016 registry.registerCollection(tagged_coll, CollectionType.TAGGED)
4017 refsets = registry.queryDatasets(..., collections=["imported_g"]).byParentDatasetType()
4018 for refs in refsets:
4019 registry.associate(tagged_coll, refs)
4021 # Summary has to have the same dataset types.
4022 summary = registry.getCollectionSummary(tagged_coll)
4023 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
4024 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
4026 # Remove all datasets from the tagged collection.
4027 refs = list(registry.queryDatasets(..., collections=[tagged_coll]))
4028 registry.disassociate(tagged_coll, refs)
4030 # Summaries should not have changed.
4031 summary = registry.getCollectionSummary(tagged_coll)
4032 self.assertEqual(summary.dataset_types.names, {"bias", "flat"})
4033 self.assertEqual(summary.governors, {"instrument": {"Cam1"}})
4035 # Cleanup summaries.
4036 registry.refresh_collection_summaries()
4037 summary = registry.getCollectionSummary(tagged_coll)
4038 self.assertFalse(summary.dataset_types.names)
4039 # We do not clean governor summaries yet, but because how the query is
4040 # run, it returns empty governors when collection is missing from
4041 # summaries.
4042 self.assertFalse(summary.governors)
4044 # Add dataset with different governor, this is to test that governors
4045 # are not actually cleaned.
4046 refs = registry.insertDatasets("test", [{"skymap": "SkyMap1", "tract": 0, "patch": 0}], "imported_g")
4047 registry.associate(tagged_coll, refs)
4048 summary = registry.getCollectionSummary(tagged_coll)
4049 self.assertEqual(summary.dataset_types.names, {"test"})
4050 # Note that instrument governor resurrects here, even though there are
4051 # no datasets left with that governor.
4052 self.assertEqual(summary.governors, {"instrument": {"Cam1"}, "skymap": {"SkyMap1"}})
4054 def test_temp_table_config(self) -> None:
4055 config = self.makeRegistryConfig()
4056 config["temporary_tables"] = False
4057 self.assertEqual(config.areTemporaryTablesAllowed, False)
4058 butler = self.make_butler(config)
4059 if not isinstance(butler, DirectButler):
4060 raise unittest.SkipTest("Test only makes sense for registry with direct database connection.")
4061 self.assertEqual(butler._registry._db.supports_temporary_tables, False)
4062 with self.assertRaisesRegex(ReadOnlyDatabaseError, "temporary tables"):
4063 with butler._registry._db.temporary_table(...):
4064 pass