From ca2b22cd59ec0478d33c601fc7a234d6b08c206c Mon Sep 17 00:00:00 2001
From: Tim Jenness <tjenness@lsst.org>
Date: Tue, 25 Nov 2025 16:01:29 -0700
Subject: [PATCH 1/4] Use proper ruff check name in precommit

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ac1b671..7860463 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -12,7 +12,7 @@ repos:
     # Ruff version.
     rev: v0.13.3
     hooks:
-      - id: ruff
+      - id: ruff-check
         args: [--fix]
       - id: ruff-format
   - repo: https://github.com/numpy/numpydoc

From 4bb3ad6217053a3af9449a1480e2e122befe5bd0 Mon Sep 17 00:00:00 2001
From: Tim Jenness <tjenness@lsst.org>
Date: Tue, 25 Nov 2025 16:25:34 -0700
Subject: [PATCH 2/4] Free butler resources when no longer needed

---
 .../script/rewrite_sqlite_registry.py            | 16 ++++++++++------
 .../daf/butler_migrate/script/update_day_obs.py  | 14 +++++++++++++-
 tests/test_dimensions_json.py                    |  6 ++++--
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py b/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
index ba51d14..1953d22 100644
--- a/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
+++ b/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
@@ -64,14 +64,18 @@ def rewrite_sqlite_registry(source: str) -> None:
     """
     # Create the source butler early so we can ask it questions
     # without assuming things.
-    source_butler = Butler.from_config(source, writeable=False)
-    assert isinstance(source_butler, DirectButler)
-    assert isinstance(source_butler._registry, SqlRegistry), "Expecting SqlRegistry instance"
+    with Butler.from_config(source, writeable=False) as source_butler:
+        assert isinstance(source_butler, DirectButler)
+        assert isinstance(source_butler._registry, SqlRegistry), "Expecting SqlRegistry instance"
 
-    # Check that we are really working with a SQLite database.
-    if not isinstance(source_butler._registry._db, SqliteDatabase):
-        raise RuntimeError("This command can only be used on SQLite registries.")
+        # Check that we are really working with a SQLite database.
+        if not isinstance(source_butler._registry._db, SqliteDatabase):
+            raise RuntimeError("This command can only be used on SQLite registries.")
 
+        _rewrite_sqlite_registry(source_butler)
+
+
+def _rewrite_sqlite_registry(source_butler: DirectButler) -> None:
     # The source butler knows where its config came from.
     source_config_uri = source_butler._config.configFile
     assert source_config_uri is not None, "Config file must not be None"
diff --git a/python/lsst/daf/butler_migrate/script/update_day_obs.py b/python/lsst/daf/butler_migrate/script/update_day_obs.py
index a4e0114..fd74ed6 100644
--- a/python/lsst/daf/butler_migrate/script/update_day_obs.py
+++ b/python/lsst/daf/butler_migrate/script/update_day_obs.py
@@ -70,8 +70,20 @@ def update_day_obs(repo: str, instrument: str) -> None:
     class registered with the instrument record.
     """
     # Connect to the butler.
-    butler = Butler.from_config(repo, writeable=True)
+    with Butler.from_config(repo, writeable=True) as butler:
+        _update_day_obs(butler, instrument)
 
+
+def _update_day_obs(butler: Butler, instrument: str) -> None:
+    """Update day obs by taking Butler rather than repo string.
+
+    Parameters
+    ----------
+    butler : `lsst.daf.butler.Butler`
+        Butler repository to update.
+    instrument : `str`
+        Name of instrument to use to update records.
+    """
     # Need the instrument class, since that is used to calculate day_obs.
     # Do not depend directly on pipe_base but to load this class pipe_base
     # will have to be available.
diff --git a/tests/test_dimensions_json.py b/tests/test_dimensions_json.py
index e7fd4f3..be5f230 100644
--- a/tests/test_dimensions_json.py
+++ b/tests/test_dimensions_json.py
@@ -229,7 +229,8 @@ def test_upgrade_v2(self) -> None:
         versions = db.manager_versions(_NAMESPACE)
         self.assertEqual(versions[_MANAGER], (_NAMESPACE, "0", _revision_id(0)))
 
-        butler = Butler(butler_root, writeable=True)  # type: ignore[abstract]
+        butler = Butler.from_config(butler_root, writeable=True)
+        self.enterContext(butler)
         assert isinstance(butler, DirectButler), "Only DirectButler is supported"
         butler.import_(filename=os.path.join(TESTDIR, "data", "records.yaml"), without_datastore=True)
 
@@ -273,7 +274,8 @@ def test_upgrade_v2(self) -> None:
         versions = db.manager_versions()
         self.assertEqual(versions[_MANAGER], (_NAMESPACE, "2", _revision_id(2)))
 
-        butler = Butler(butler_root, writeable=False)  # type: ignore[abstract]
+        butler = Butler.from_config(butler_root, writeable=False)
+        self.enterContext(butler)
 
         # Check records for v2 attributes.
         records = list(butler.registry.queryDimensionRecords("instrument"))

From 4bc2846bf00b2871d898426eeae574486a9ff76b Mon Sep 17 00:00:00 2001
From: Tim Jenness <tjenness@lsst.org>
Date: Tue, 25 Nov 2025 17:27:35 -0700
Subject: [PATCH 3/4] Add Database.close() and allow use in context manager

---
 python/lsst/daf/butler_migrate/database.py    | 47 ++++++++-----
 .../butler_migrate/script/migrate_current.py  | 51 +++++++-------
 .../script/migrate_downgrade.py               | 41 ++++++------
 .../script/migrate_dump_schema.py             |  4 +-
 .../script/migrate_set_namespace.py           | 40 +++++------
 .../butler_migrate/script/migrate_stamp.py    | 67 +++++++++----------
 .../butler_migrate/script/migrate_upgrade.py  | 43 ++++++------
 tests/test_database.py                        | 25 +++----
 tests/test_dimensions_json.py                 | 28 ++++----
 9 files changed, 178 insertions(+), 168 deletions(-)

diff --git a/python/lsst/daf/butler_migrate/database.py b/python/lsst/daf/butler_migrate/database.py
index 8f43b43..ae7574a 100644
--- a/python/lsst/daf/butler_migrate/database.py
+++ b/python/lsst/daf/butler_migrate/database.py
@@ -24,8 +24,8 @@
 import json
 import logging
 from collections.abc import Iterable, Iterator, Mapping
-from contextlib import contextmanager
-from typing import cast
+from contextlib import AbstractContextManager, contextmanager
+from typing import Any, Literal, Self, cast
 
 import sqlalchemy
 from alembic.runtime.migration import MigrationContext
@@ -43,7 +43,7 @@ class RevisionConsistencyError(Exception):
     """
 
 
-class Database:
+class Database(AbstractContextManager):
     """Class implementing methods for database access needed for migrations.
 
     Parameters
@@ -69,6 +69,7 @@ class Database:
     def __init__(self, db_url: sqlalchemy.engine.url.URL, schema: str | None = None):
         self._db_url = db_url
         self._schema = schema
+        self._engine: sqlalchemy.engine.Engine | None = None
 
     @classmethod
     def from_repo(cls, repo: str) -> Database:
@@ -100,13 +101,35 @@ def schema(self) -> str | None:
         """Schema (namespace) name (`str`)."""
         return self._schema
 
+    @property
+    def engine(self) -> sqlalchemy.engine.Engine:
+        """Cached sqlalchemy Engine."""
+        if self._engine is None:
+            self._engine = sqlalchemy.engine.create_engine(self._db_url)
+        return self._engine
+
     @contextmanager
     def connect(self) -> Iterator[sqlalchemy.engine.Connection]:
         """Context manager for database connection."""
-        engine = sqlalchemy.engine.create_engine(self._db_url)
-        with engine.connect() as connection:
+        with self.engine.connect() as connection:
             yield connection
 
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Literal[False]:
+        try:
+            self.close()
+        except Exception:
+            _LOG.exception("An exception occurred during Database.close()")
+        return False
+
+    def close(self) -> None:
+        """Cleanup connection pool."""
+        if self._engine is not None:
+            self._engine.dispose()
+            self._engine = None
+
     def dimensions_namespace(self) -> str | None:
         """Return dimensions namespace from a stored configuration.
 
@@ -115,8 +138,6 @@ def dimensions_namespace(self) -> str | None:
         namespace: `str` or `None`
             Dimensions namespace or `None` if not defined.
         """
-        engine = sqlalchemy.engine.create_engine(self._db_url)
-
         meta = sqlalchemy.schema.MetaData(schema=self._schema)
         table = sqlalchemy.schema.Table(
             "butler_attributes",
@@ -126,7 +147,7 @@ def dimensions_namespace(self) -> str | None:
         )
 
         sql = sqlalchemy.sql.select(table.columns.value).where(table.columns.name == self.dimensions_json_key)
-        with engine.connect() as connection:
+        with self.engine.connect() as connection:
             result = connection.execute(sql)
             row = result.fetchone()
             if row is None:
@@ -150,8 +171,6 @@ def manager_versions(self, namespace: str | None = None) -> Mapping[str, tuple[s
             tuple consisting of manager class name (including package/module),
             version string in X.Y.Z format, and revision ID string/hash.
         """
-        engine = sqlalchemy.engine.create_engine(self._db_url)
-
         meta = sqlalchemy.schema.MetaData(schema=self._schema)
         table = sqlalchemy.schema.Table(
             "butler_attributes",
@@ -164,7 +183,7 @@ def manager_versions(self, namespace: str | None = None) -> Mapping[str, tuple[s
         managers: dict[str, str] = {}
         versions: dict[str, str] = {}
         sql = sqlalchemy.sql.select(table.columns.name, table.columns.value)
-        with engine.connect() as connection:
+        with self.engine.connect() as connection:
             result = connection.execute(sql)
             for name, value in result:
                 if name.startswith("config:registry.managers."):
@@ -211,8 +230,7 @@ def alembic_revisions(self) -> list[str]:
             Returned list is empty if alembic version table does not exist or
             is empty.
         """
-        engine = sqlalchemy.engine.create_engine(self._db_url)
-        with engine.connect() as connection:
+        with self.engine.connect() as connection:
             ctx = MigrationContext.configure(
                 connection=connection, opts={"version_table_schema": self._schema}
             )
@@ -286,8 +304,7 @@ def dump_schema(self, tables: list[str] | None) -> None:
             List of the tables, if missing or empty then schema for all tables
             is printed.
         """
-        engine = sqlalchemy.engine.create_engine(self._db_url)
-        inspector = sqlalchemy.inspect(engine)
+        inspector = sqlalchemy.inspect(self.engine)
         table_names = sorted(inspector.get_table_names(schema=self._schema))
         for table in table_names:
             if tables and table not in tables:
diff --git a/python/lsst/daf/butler_migrate/script/migrate_current.py b/python/lsst/daf/butler_migrate/script/migrate_current.py
index 64f5277..a41c55c 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_current.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_current.py
@@ -51,32 +51,31 @@ def migrate_current(repo: str, mig_path: str, verbose: bool, butler: bool, names
         Dimensions namespace to use when "namespace" key is not present in
         ``config:dimensions.json``.
     """
-    db = database.Database.from_repo(repo)
+    with database.Database.from_repo(repo) as db:
+        if namespace is None and db.dimensions_namespace() is None:
+            raise ValueError(
+                "The `--namespace` option is required when namespace is missing from"
+                " stored dimensions configuration"
+            )
 
-    if namespace is None and db.dimensions_namespace() is None:
-        raise ValueError(
-            "The `--namespace` option is required when namespace is missing from"
-            " stored dimensions configuration"
-        )
-
-    cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
-    if butler:
-        # Print current versions defined in butler.
-        script_info = scripts.Scripts(cfg)
-        heads = script_info.head_revisions()
-        manager_versions = db.manager_versions(namespace)
-        if manager_versions:
-            for manager, (klass, version, rev_id) in sorted(manager_versions.items()):
-                head = " (head)" if rev_id in heads else ""
-                print(f"{manager}: {klass} {version} -> {rev_id}{head}")
+        cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
+        if butler:
+            # Print current versions defined in butler.
+            script_info = scripts.Scripts(cfg)
+            heads = script_info.head_revisions()
+            manager_versions = db.manager_versions(namespace)
+            if manager_versions:
+                for manager, (klass, version, rev_id) in sorted(manager_versions.items()):
+                    head = " (head)" if rev_id in heads else ""
+                    print(f"{manager}: {klass} {version} -> {rev_id}{head}")
+            else:
+                print("No manager versions defined in butler_attributes table.")
         else:
-            print("No manager versions defined in butler_attributes table.")
-    else:
-        # Revisions from alembic.
-        command.current(cfg, verbose=verbose)
+            # Revisions from alembic.
+            command.current(cfg, verbose=verbose)
 
-    # Complain if alembic_version table is there but does not match manager
-    # versions.
-    if db.alembic_revisions():
-        script_info = scripts.Scripts(cfg)
-        db.validate_revisions(namespace, script_info.base_revisions())
+        # Complain if alembic_version table is there but does not match manager
+        # versions.
+        if db.alembic_revisions():
+            script_info = scripts.Scripts(cfg)
+            db.validate_revisions(namespace, script_info.base_revisions())
diff --git a/python/lsst/daf/butler_migrate/script/migrate_downgrade.py b/python/lsst/daf/butler_migrate/script/migrate_downgrade.py
index 5dd73c2..3bfcd94 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_downgrade.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_downgrade.py
@@ -54,28 +54,29 @@ def migrate_downgrade(
         Dimensions namespace to use when "namespace" key is not present in
         ``config:dimensions.json``.
     """
-    db = database.Database.from_repo(repo)
+    with database.Database.from_repo(repo) as db:
+        if namespace is None and db.dimensions_namespace() is None:
+            raise ValueError(
+                "The `--namespace` option is required when namespace is missing from"
+                " stored dimensions configuration"
+            )
 
-    if namespace is None and db.dimensions_namespace() is None:
-        raise ValueError(
-            "The `--namespace` option is required when namespace is missing from"
-            " stored dimensions configuration"
-        )
+        # Check that alembic versions exist in database, we do not support
+        # migrations from empty state.
+        if not db.alembic_revisions():
+            raise ValueError(
+                "Alembic version table does not exist, you may need to run `butler migrate stamp` first."
+            )
 
-    # Check that alembic versions exist in database, we do not support
-    # migrations from empty state.
-    if not db.alembic_revisions():
-        raise ValueError(
-            "Alembic version table does not exist, you may need to run `butler migrate stamp` first."
+        one_shot_arg: str | None = None
+        if one_shot_tree:
+            one_shot_arg = one_shot_tree
+        cfg = config.MigAlembicConfig.from_mig_path(
+            mig_path, repository=repo, db=db, one_shot_tree=one_shot_arg
         )
 
-    one_shot_arg: str | None = None
-    if one_shot_tree:
-        one_shot_arg = one_shot_tree
-    cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db, one_shot_tree=one_shot_arg)
-
-    # check that alembic versions are consistent with butler
-    script_info = scripts.Scripts(cfg)
-    db.validate_revisions(namespace, script_info.base_revisions())
+        # check that alembic versions are consistent with butler
+        script_info = scripts.Scripts(cfg)
+        db.validate_revisions(namespace, script_info.base_revisions())
 
-    command.downgrade(cfg, revision, sql=sql)
+        command.downgrade(cfg, revision, sql=sql)
diff --git a/python/lsst/daf/butler_migrate/script/migrate_dump_schema.py b/python/lsst/daf/butler_migrate/script/migrate_dump_schema.py
index 317fcbb..e25af0a 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_dump_schema.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_dump_schema.py
@@ -39,5 +39,5 @@ def migrate_dump_schema(repo: str, table: list[str]) -> None:
     table : `list`
         List of the tables, if empty then schema for all tables is printed.
     """
-    db = database.Database.from_repo(repo)
-    db.dump_schema(table)
+    with database.Database.from_repo(repo) as db:
+        db.dump_schema(table)
diff --git a/python/lsst/daf/butler_migrate/script/migrate_set_namespace.py b/python/lsst/daf/butler_migrate/script/migrate_set_namespace.py
index 29d196a..459f293 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_set_namespace.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_set_namespace.py
@@ -44,27 +44,27 @@ def migrate_set_namespace(repo: str, namespace: str | None, update: bool) -> Non
     update : `bool`
         Allows update of the existing namespace.
     """
-    db = database.Database.from_repo(repo)
-    db_namespace = db.dimensions_namespace()
+    with database.Database.from_repo(repo) as db:
+        db_namespace = db.dimensions_namespace()
 
-    if not namespace:
-        # Print current value
-        if not db_namespace:
-            print("No namespace defined in dimensions configuration.")
-        else:
-            print("Current dimensions namespace:", db_namespace)
+        if not namespace:
+            # Print current value
+            if not db_namespace:
+                print("No namespace defined in dimensions configuration.")
+            else:
+                print("Current dimensions namespace:", db_namespace)
 
-    else:
-        if db_namespace and not update:
-            raise ValueError(
-                f"Namespace is already defined ({db_namespace}), use --update option to replace it."
-            )
+        else:
+            if db_namespace and not update:
+                raise ValueError(
+                    f"Namespace is already defined ({db_namespace}), use --update option to replace it."
+                )
 
-        def update_namespace(config: dict) -> dict:
-            """Update namespace attribute"""
-            config["namespace"] = namespace
-            return config
+            def update_namespace(config: dict) -> dict:
+                """Update namespace attribute"""
+                config["namespace"] = namespace
+                return config
 
-        with db.connect() as connection:
-            attributes = butler_attributes.ButlerAttributes(connection, db.schema)
-            attributes.update_dimensions_json(update_namespace)
+            with db.connect() as connection:
+                attributes = butler_attributes.ButlerAttributes(connection, db.schema)
+                attributes.update_dimensions_json(update_namespace)
diff --git a/python/lsst/daf/butler_migrate/script/migrate_stamp.py b/python/lsst/daf/butler_migrate/script/migrate_stamp.py
index 884c788..5951acf 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_stamp.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_stamp.py
@@ -54,42 +54,41 @@ def migrate_stamp(
     manager : `str`, Optional
         Name of the manager to stamp, if `None` then all managers are stamped.
     """
-    db = database.Database.from_repo(repo)
+    with database.Database.from_repo(repo) as db:
+        if namespace is None and db.dimensions_namespace() is None:
+            raise ValueError(
+                "The `--namespace` option is required when namespace is missing from"
+                " stored dimensions configuration"
+            )
 
-    if namespace is None and db.dimensions_namespace() is None:
-        raise ValueError(
-            "The `--namespace` option is required when namespace is missing from"
-            " stored dimensions configuration"
-        )
+        manager_versions = db.manager_versions(namespace)
 
-    manager_versions = db.manager_versions(namespace)
+        revisions: dict[str, str] = {}
+        for mgr_name, (klass, version, rev_id) in manager_versions.items():
+            _LOG.debug("found revision (%s, %s, %s) -> %s", mgr_name, klass, version, rev_id)
+            revisions[mgr_name] = rev_id
 
-    revisions: dict[str, str] = {}
-    for mgr_name, (klass, version, rev_id) in manager_versions.items():
-        _LOG.debug("found revision (%s, %s, %s) -> %s", mgr_name, klass, version, rev_id)
-        revisions[mgr_name] = rev_id
+        cfg: config.MigAlembicConfig | None = None
+        if manager:
+            if manager in revisions:
+                revisions = {manager: revisions[manager]}
+            else:
+                # If specified manager not in the database, it may mean that an
+                # initial "tree-root" revision needs to be added to alembic
+                # table, if that manager is defined in the migration trees.
+                cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
+                script_info = scripts.Scripts(cfg)
+                base_revision = revision.rev_id(manager)
+                if base_revision not in script_info.base_revisions():
+                    raise ValueError(f"Unknown manager name {manager} (not in the database or migrations)")
+                revisions = {manager: base_revision}
 
-    cfg: config.MigAlembicConfig | None = None
-    if manager:
-        if manager in revisions:
-            revisions = {manager: revisions[manager]}
+        if dry_run:
+            print("Will store these revisions in alembic version table:")
+            for manager, rev_id in revisions.items():
+                print(f"  {manager}: {rev_id}")
         else:
-            # If specified manager not in the database, it may mean that an
-            # initial "tree-root" revision needs to be added to alembic
-            # table, if that manager is defined in the migration trees.
-            cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
-            script_info = scripts.Scripts(cfg)
-            base_revision = revision.rev_id(manager)
-            if base_revision not in script_info.base_revisions():
-                raise ValueError(f"Unknown manager name {manager} (not in the database or migrations)")
-            revisions = {manager: base_revision}
-
-    if dry_run:
-        print("Will store these revisions in alembic version table:")
-        for manager, rev_id in revisions.items():
-            print(f"  {manager}: {rev_id}")
-    else:
-        if cfg is None:
-            cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
-        for rev in revisions.values():
-            command.stamp(cfg, rev, purge=purge)
+            if cfg is None:
+                cfg = config.MigAlembicConfig.from_mig_path(mig_path, repository=repo, db=db)
+            for rev in revisions.values():
+                command.stamp(cfg, rev, purge=purge)
diff --git a/python/lsst/daf/butler_migrate/script/migrate_upgrade.py b/python/lsst/daf/butler_migrate/script/migrate_upgrade.py
index d048970..1f1eba0 100644
--- a/python/lsst/daf/butler_migrate/script/migrate_upgrade.py
+++ b/python/lsst/daf/butler_migrate/script/migrate_upgrade.py
@@ -62,30 +62,29 @@ def migrate_upgrade(
     options : `dict` [ `str`, `str` ] or `None`
         Options to select.
     """
-    db = database.Database.from_repo(repo)
+    with database.Database.from_repo(repo) as db:
+        if namespace is None and db.dimensions_namespace() is None:
+            raise ValueError(
+                "The `--namespace` option is required when namespace is missing from"
+                " stored dimensions configuration"
+            )
 
-    if namespace is None and db.dimensions_namespace() is None:
-        raise ValueError(
-            "The `--namespace` option is required when namespace is missing from"
-            " stored dimensions configuration"
-        )
+        # Check that alembic versions exist in database, we do not support
+        # migrations from empty state.
+        if not db.alembic_revisions():
+            raise ValueError(
+                "Alembic version table does not exist, you may need to run `butler migrate stamp` first."
+            )
 
-    # Check that alembic versions exist in database, we do not support
-    # migrations from empty state.
-    if not db.alembic_revisions():
-        raise ValueError(
-            "Alembic version table does not exist, you may need to run `butler migrate stamp` first."
+        one_shot_arg: str | None = None
+        if one_shot_tree:
+            one_shot_arg = one_shot_tree
+        cfg = config.MigAlembicConfig.from_mig_path(
+            mig_path, repository=repo, db=db, one_shot_tree=one_shot_arg, migration_options=options
         )
 
-    one_shot_arg: str | None = None
-    if one_shot_tree:
-        one_shot_arg = one_shot_tree
-    cfg = config.MigAlembicConfig.from_mig_path(
-        mig_path, repository=repo, db=db, one_shot_tree=one_shot_arg, migration_options=options
-    )
-
-    # check that alembic versions are consistent with butler
-    script_info = scripts.Scripts(cfg)
-    db.validate_revisions(namespace, script_info.base_revisions())
+        # check that alembic versions are consistent with butler
+        script_info = scripts.Scripts(cfg)
+        db.validate_revisions(namespace, script_info.base_revisions())
 
-    command.upgrade(cfg, revision, sql=sql)
+        command.upgrade(cfg, revision, sql=sql)
diff --git a/tests/test_database.py b/tests/test_database.py
index 91b3b20..13ad65b 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -92,6 +92,7 @@ def make_revision_tables(
             for query in queries:
                 conn.execute(sqlalchemy.text(query))
         yield db_url
+        engine.dispose()
 
 
 class DatabaseTestCase(unittest.TestCase):
@@ -99,8 +100,7 @@ class DatabaseTestCase(unittest.TestCase):
 
     def test_manager_versions(self) -> None:
         """Test for manager_versions() method"""
-        with make_revision_tables() as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables() as db_url, database.Database(db_url) as db:
             manager_versions = db.manager_versions()
             self.assertEqual(
                 manager_versions,
@@ -116,8 +116,7 @@ def test_manager_versions(self) -> None:
 
     def test_alembic_revisions(self) -> None:
         """Test for alembic_revisions() method"""
-        with make_revision_tables() as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables() as db_url, database.Database(db_url) as db:
             alembic_revisions = db.alembic_revisions()
             self.assertCountEqual(
                 alembic_revisions,
@@ -129,40 +128,34 @@ def test_alembic_revisions(self) -> None:
 
     def test_validate_revisions(self) -> None:
         """Test for validate_revisions() method"""
-        with make_revision_tables() as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables() as db_url, database.Database(db_url) as db:
             db.validate_revisions()
 
-        with make_revision_tables(make_alembic=False) as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables(make_alembic=False) as db_url, database.Database(db_url) as db:
             with self.assertRaisesRegex(
                 database.RevisionConsistencyError, "alembic_version table does not exist or is empty"
             ):
                 db.validate_revisions()
 
-        with make_revision_tables(fill_alembic=False) as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables(fill_alembic=False) as db_url, database.Database(db_url) as db:
             with self.assertRaisesRegex(
                 database.RevisionConsistencyError, "alembic_version table does not exist or is empty"
             ):
                 db.validate_revisions()
 
-        with make_revision_tables(make_butler=False) as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables(make_butler=False) as db_url, database.Database(db_url) as db:
             with self.assertRaisesRegex(
                 database.RevisionConsistencyError, "butler_attributes table does not exist"
             ):
                 db.validate_revisions()
 
-        with make_revision_tables(fill_butler=False) as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables(fill_butler=False) as db_url, database.Database(db_url) as db:
             with self.assertRaisesRegex(
                 database.RevisionConsistencyError, "butler_attributes table is empty"
             ):
                 db.validate_revisions()
 
-        with make_revision_tables(broken_alembic=True) as db_url:
-            db = database.Database(db_url)
+        with make_revision_tables(broken_alembic=True) as db_url, database.Database(db_url) as db:
             with self.assertRaisesRegex(
                 database.RevisionConsistencyError, "Butler and alembic revisions are inconsistent"
             ):
diff --git a/tests/test_dimensions_json.py b/tests/test_dimensions_json.py
index be5f230..087ca7a 100644
--- a/tests/test_dimensions_json.py
+++ b/tests/test_dimensions_json.py
@@ -109,6 +109,7 @@ def _upgrade_one(self, start_version: int) -> None:
         """Test version upgrade from N to N+1."""
         butler_root = self.make_butler(start_version)
         db = database.Database.from_repo(butler_root)
+        self.enterContext(db)
 
         versions = db.manager_versions(_NAMESPACE)
         self.assertEqual(versions[_MANAGER], (_NAMESPACE, str(start_version), _revision_id(start_version)))
@@ -140,6 +141,7 @@ def _downgrade_one(self, start_version: int) -> None:
         """Test version downgrade from N to N-1."""
         butler_root = self.make_butler(start_version)
         db = database.Database.from_repo(butler_root)
+        self.enterContext(db)
 
         versions = db.manager_versions(_NAMESPACE)
         self.assertEqual(versions[_MANAGER], (_NAMESPACE, str(start_version), _revision_id(start_version)))
@@ -184,6 +186,7 @@ def test_upgrade_v1(self) -> None:
         """
         butler_root = self.make_butler(0)
         db = database.Database.from_repo(butler_root)
+        self.enterContext(db)
 
         self.assertIsNone(db.dimensions_namespace())
         versions = db.manager_versions(_NAMESPACE)
@@ -224,26 +227,24 @@ def test_upgrade_v2(self) -> None:
         """
         butler_root = self.make_butler(0)
         db = database.Database.from_repo(butler_root)
+        self.enterContext(db)
 
         self.assertIsNone(db.dimensions_namespace())
         versions = db.manager_versions(_NAMESPACE)
         self.assertEqual(versions[_MANAGER], (_NAMESPACE, "0", _revision_id(0)))
 
-        butler = Butler.from_config(butler_root, writeable=True)
-        self.enterContext(butler)
-        assert isinstance(butler, DirectButler), "Only DirectButler is supported"
-        butler.import_(filename=os.path.join(TESTDIR, "data", "records.yaml"), without_datastore=True)
-
-        # Check records for v0 attributes.
-        records = list(butler.registry.queryDimensionRecords("visit"))
-        for record in records:
-            self.assertEqual(record.visit_system, 0)
+        with Butler.from_config(butler_root, writeable=True) as butler:
+            assert isinstance(butler, DirectButler), "Only DirectButler is supported"
+            butler.import_(filename=os.path.join(TESTDIR, "data", "records.yaml"), without_datastore=True)
 
-        records = list(butler.registry.queryDimensionRecords("visit_definition"))
-        for record in records:
-            self.assertEqual(record.visit_system, 0)
+            # Check records for v0 attributes.
+            records = list(butler.registry.queryDimensionRecords("visit"))
+            for record in records:
+                self.assertEqual(record.visit_system, 0)
 
-        del butler
+            records = list(butler.registry.queryDimensionRecords("visit_definition"))
+            for record in records:
+                self.assertEqual(record.visit_system, 0)
 
         # Upgrade to v1. We could upgrade to v2 in one step but I want to check
         # different arguments at each step.
@@ -311,6 +312,7 @@ def test_upgrade_v2(self) -> None:
     def test_validate_dimensions_json(self) -> None:
         butler_root = self.make_butler(0)
         db = database.Database.from_repo(butler_root)
+        self.enterContext(db)
         universe = 5
         with db.connect() as connection:
             attribs = butler_attributes.ButlerAttributes(connection, schema=db.schema)

From a7249ae2a3882413b49f5af1c801e210225c5b92 Mon Sep 17 00:00:00 2001
From: Tim Jenness <tjenness@lsst.org>
Date: Mon, 1 Dec 2025 16:39:41 -0700
Subject: [PATCH 4/4] Add more asserts to aid mypy

These became necessary once the code was split into two functions.
---
 .../lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py b/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
index 1953d22..93fa1b5 100644
--- a/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
+++ b/python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
@@ -151,6 +151,10 @@ def _rewrite_sqlite_registry(source_butler: DirectButler) -> None:
         # Now need to move this registry to the original location
         # and move the existing registry to a backup.
 
+        assert isinstance(source_butler, DirectButler)
+        assert isinstance(source_butler._registry, SqlRegistry), "Expecting SqlRegistry instance"
+        assert isinstance(source_butler._registry._db, SqliteDatabase), "Expecting SqliteDatabase instance"
+
         # Relocate the source registry first
         assert source_butler._registry._db.filename is not None, "Expecting non-None filename from registry"
         source_registry_uri = ResourcePath(source_butler._registry._db.filename)