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