Merge pull request #316 from lsst/tickets/DM-25355

DM-25355: Add support for specifying schema versions

python/lsst/daf/butler/registry/_registry.py

@@ -70,6 +70,7 @@
 from ._exceptions import ConflictingDefinitionError, InconsistentDataIdError, OrphanedRecordError
 from .wildcards import CategorizedWildcard, CollectionQuery, CollectionSearch, Ellipsis
 from .interfaces import ChainedCollectionRecord, RunRecord
+from .versions import ButlerVersionsManager
 
 if TYPE_CHECKING:
     from ..butlerConfig import ButlerConfig
@@ -138,9 +139,11 @@ def fromConfig(cls, config: Union[ButlerConfig, RegistryConfig, Config, str], cr
         collections = doImport(config["managers", "collections"])
         datasets = doImport(config["managers", "datasets"])
         datastoreBridges = doImport(config["managers", "datastores"])
+        versions = ButlerVersionsManager.fromConfig(config.get("schema_versions"))
+
         return cls(database, universe, dimensions=dimensions, attributes=attributes, opaque=opaque,
                    collections=collections, datasets=datasets, datastoreBridges=datastoreBridges,
-                   create=create)
+                   versions=versions, writeable=writeable, create=create)
 
     def __init__(self, database: Database, universe: DimensionUniverse, *,
                  attributes: Type[ButlerAttributeManager],
@@ -149,6 +152,8 @@ def __init__(self, database: Database, universe: DimensionUniverse, *,
                  collections: Type[CollectionManager],
                  datasets: Type[DatasetRecordStorageManager],
                  datastoreBridges: Type[DatastoreRegistryBridgeManager],
+                 versions: ButlerVersionsManager,
+                 writeable: bool = True,
                  create: bool = False):
         self._db = database
         self.storageClasses = StorageClassFactory()
@@ -164,6 +169,15 @@ def __init__(self, database: Database, universe: DimensionUniverse, *,
                                                                  opaque=self._opaque,
                                                                  datasets=datasets,
                                                                  universe=self.dimensions)
+            context.addInitializer(lambda db: versions.storeVersions(self._attributes))
+
+        # This call does not do anything right now as we do not have a way to
+        # split tables between sub-schemas yet.
+        versions.checkVersionDigests()
+        if not create:
+            # verify that configured versions are compatible with schema
+            versions.checkStoredVersions(self._attributes, writeable)
+
         self._collections.refresh()
         self._datasets.refresh(universe=self._dimensions.universe)
 

python/lsst/daf/butler/registry/interfaces/_database.py

@@ -31,6 +31,7 @@
 from contextlib import contextmanager
 from typing import (
     Any,
+    Callable,
     Dict,
     Iterable,
     Iterator,
@@ -100,6 +101,7 @@ def __init__(self, db: Database):
         self._foreignKeys: List[Tuple[sqlalchemy.schema.Table, sqlalchemy.schema.ForeignKeyConstraint]] = []
         self._inspector = sqlalchemy.engine.reflection.Inspector(self._db._connection)
         self._tableNames = frozenset(self._inspector.get_table_names(schema=self._db.namespace))
+        self._initializers: List[Callable[[Database], None]] = []
 
     def addTable(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table:
         """Add a new table to the schema, returning its sqlalchemy
@@ -141,6 +143,20 @@ def addTableTuple(self, specs: Tuple[ddl.TableSpec, ...]) -> Tuple[sqlalchemy.sc
         return specs._make(self.addTable(name, spec)                     # type: ignore
                            for name, spec in zip(specs._fields, specs))  # type: ignore
 
+    def addInitializer(self, initializer: Callable[[Database], None]) -> None:
+        """Add a method that does one-time initialization of a database.
+
+        Initialization can mean anything that changes state of a database
+        and needs to be done exactly once after database schema was created.
+        An example for that could be population of schema attributes.
+
+        Parameters
+        ----------
+        initializer : callable
+            Method of a single argument which is a `Database` instance.
+        """
+        self._initializers.append(initializer)
+
 
 class Database(ABC):
     """An abstract interface that represents a particular database engine's
@@ -387,7 +403,12 @@ def declareStaticTables(self, *, create: bool) -> Iterator[StaticTablesContext]:
                 with warnings.catch_warnings():
                     warnings.simplefilter("ignore", category=sqlalchemy.exc.SADeprecationWarning)
                     self._metadata.create_all(self._connection)
+                # call all initializer methods sequentially
+                for init in context._initializers:
+                    init(self)
         except BaseException:
+            # TODO: this is potentially dangerous if we run it on
+            # pre-existing schema.
             self._metadata.drop_all(self._connection)
             self._metadata = None
             raise

python/lsst/daf/butler/registry/tests/_registry.py

@@ -1020,14 +1020,17 @@ def testAbstractFilterQuery(self):
     def testAttributeManager(self):
         """Test basic functionality of attribute manager.
         """
+        # number of attributes with schema versions in a fresh database
+        VERSION_COUNT = 0
+
         registry = self.makeRegistry()
         attributes = registry._attributes
 
         # check what get() returns for non-existing key
         self.assertIsNone(attributes.get("attr"))
         self.assertEqual(attributes.get("attr", ""), "")
         self.assertEqual(attributes.get("attr", "Value"), "Value")
-        self.assertEqual(len(list(attributes.items())), 0)
+        self.assertEqual(len(list(attributes.items())), VERSION_COUNT)
 
         # cannot store empty key or value
         with self.assertRaises(ValueError):
@@ -1037,20 +1040,20 @@ def testAttributeManager(self):
 
         # set value of non-existing key
         attributes.set("attr", "value")
-        self.assertEqual(len(list(attributes.items())), 1)
+        self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)
         self.assertEqual(attributes.get("attr"), "value")
 
         # update value of existing key
         with self.assertRaises(ButlerAttributeExistsError):
             attributes.set("attr", "value2")
 
         attributes.set("attr", "value2", force=True)
-        self.assertEqual(len(list(attributes.items())), 1)
+        self.assertEqual(len(list(attributes.items())), VERSION_COUNT + 1)
         self.assertEqual(attributes.get("attr"), "value2")
 
         # delete existing key
         self.assertTrue(attributes.delete("attr"))
-        self.assertEqual(len(list(attributes.items())), 0)
+        self.assertEqual(len(list(attributes.items())), VERSION_COUNT)
 
         # delete non-existing key
         self.assertFalse(attributes.delete("non-attr"))
@@ -1063,5 +1066,6 @@ def testAttributeManager(self):
         ]
         for key, value in data:
             attributes.set(key, value)
-        items = list(attributes.items())
-        self.assertCountEqual(items, data)
+        items = dict(attributes.items())
+        for key, value in data:
+            self.assertEqual(items[key], value)