From e0205d63f62f2cc1be68b4288b644280ecfb26cd Mon Sep 17 00:00:00 2001
From: Jim Bosch <jbosch@astro.princeton.edu>
Date: Thu, 17 Oct 2024 10:59:12 -0400
Subject: [PATCH] Test and fix calibration dataset type registration bug.

Without this fix, calibration dataset types that were registered in
data repositories where the only existing dataset types with those
dimensions were not calibrations would incorrectly be registered as
non-calibrations.
---
 .../datasets/byDimensions/_manager.py         |  2 ++
 .../registry/datasets/byDimensions/tables.py  | 18 ++++++++++++
 tests/test_simpleButler.py                    | 29 +++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py b/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
index 9236cc636f..7d0a9c5e0b 100644
--- a/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
+++ b/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
@@ -327,6 +327,8 @@ def register_dataset_type(self, dataset_type: DatasetType) -> bool:
                     dataset_type.dimensions, dimensions_key, dataset_type.isCalibration()
                 )
                 dynamic_tables.create(self._db, type(self._collections))
+            elif dataset_type.isCalibration() and dynamic_tables.calibs_name is None:
+                dynamic_tables.add_calibs(self._db, type(self._collections))
             row, inserted = self._db.sync(
                 self._static.dataset_type,
                 keys={"name": dataset_type.name},
diff --git a/python/lsst/daf/butler/registry/datasets/byDimensions/tables.py b/python/lsst/daf/butler/registry/datasets/byDimensions/tables.py
index 7528253d63..029997ba64 100644
--- a/python/lsst/daf/butler/registry/datasets/byDimensions/tables.py
+++ b/python/lsst/daf/butler/registry/datasets/byDimensions/tables.py
@@ -531,6 +531,24 @@ def create(self, db: Database, collections: type[CollectionManager]) -> None:
                 makeCalibTableSpec(self._dimensions, collections, db.getTimespanRepresentation()),
             )
 
+    def add_calibs(self, db: Database, collections: type[CollectionManager]) -> None:
+        """Create a calibs table for a dataset type whose dimensions already
+        have a tags table.
+
+        Parameters
+        ----------
+        db : `Database`
+            Database interface.
+        collections : `type` [ `CollectionManager` ]
+            Manager class for collections; used to create foreign key columns
+            for collections.
+        """
+        self.calibs_name = makeCalibTableName(self.dimensions_key)
+        self._calibs_table = db.ensureTableExists(
+            self.calibs_name,
+            makeCalibTableSpec(self._dimensions, collections, db.getTimespanRepresentation()),
+        )
+
     def tags(self, db: Database, collections: type[CollectionManager]) -> sqlalchemy.Table:
         """Return the "tags" table that associates datasets with data IDs in
         TAGGED and RUN collections.
diff --git a/tests/test_simpleButler.py b/tests/test_simpleButler.py
index 6de7b81dfb..595b08957f 100644
--- a/tests/test_simpleButler.py
+++ b/tests/test_simpleButler.py
@@ -823,6 +823,35 @@ def test_clone(self):
         self.assertCountEqual(clone5.registry.defaults.collections, ["imported_r"])
         self.assertEqual(clone5.run, "imported_r")
 
+    def test_calibration_dataset_type_registration(self) -> None:
+        # Register two dataset types that should share the same tags table,
+        # but only one is a calibration and hence needs a calibs table.
+        butler1 = self.makeButler(writeable=True)
+        a = DatasetType("a", ["instrument"], universe=butler1.dimensions, storageClass="StructuredDataDict")
+        b = DatasetType(
+            "b",
+            ["instrument"],
+            universe=butler1.dimensions,
+            storageClass="StructuredDataDict",
+            isCalibration=True,
+        )
+        butler1.registry.registerDatasetType(a)
+        butler1.registry.registerDatasetType(b)
+        self.assertEqual(butler1.get_dataset_type("a"), a)
+        self.assertEqual(butler1.get_dataset_type("b"), b)
+        butler1.registry.refresh()
+        self.assertEqual(butler1.get_dataset_type("a"), a)
+        self.assertEqual(butler1.get_dataset_type("b"), b)
+        # Register them in the opposite order in a new repo.
+        butler2 = self.makeButler(writeable=True)
+        butler2.registry.registerDatasetType(b)
+        butler2.registry.registerDatasetType(a)
+        self.assertEqual(butler2.get_dataset_type("a"), a)
+        self.assertEqual(butler2.get_dataset_type("b"), b)
+        butler2.registry.refresh()
+        self.assertEqual(butler2.get_dataset_type("a"), a)
+        self.assertEqual(butler2.get_dataset_type("b"), b)
+
 
 class DirectSimpleButlerTestCase(SimpleButlerTests, unittest.TestCase):
     """Run tests against DirectButler implementation."""