Merge pull request #33 from gianchub/engine-fix

Engine Disposal

Author: gianchub

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## [1.1.1]
+
+12/12/2025
+
+- Add `dispose_engines` flag to `Comparer` constructor.
+- Add `dispose` method to `Comparer` class.
+- Add tests for engine disposal.
+- Update README.md with new information.
+
 ## [1.1.0]
 
 10/12/2025

Makefile

@@ -1,5 +1,5 @@
 .PHONY: ruff-fix ruff-check ruff-format ruff-format-check lint format test install-tox-uv test-sqlalchemy14
-.PHONY: ty install-reqs docker-test-db-run build publish-test publish bump-version
+.PHONY: ty install-reqs update-reqs docker-test-db-run build publish-test publish bump-version
 
 # Misc
 
@@ -41,6 +41,9 @@ ty:
 
 # requirements
 
+update-reqs:
+	uv lock --upgrade
+
 install-reqs:
 	uv pip install -U -e ."[dev,lint,test]"
 

README.md

@@ -42,7 +42,7 @@ result.dump_result('comparison_result.json')
 result.dump_errors('comparison_errors.json')
 ```
 
-You can also create a comparer directly from database URIs:
+You can create a comparer directly from database URIs, using the `from_params` classmethod:
 
 ```python
 from sqlalchemydiff.comparer import Comparer
@@ -56,23 +56,36 @@ comparer = Comparer.from_params(
 result = comparer.compare()
 ```
 
+
+> [!NOTE]
+> When using the `from_params` classmethod, the engines will be disposed after the comparison is complete, to avoi leaving pooled connections open.
+> If instead you supply your own engines, **manage their lifecycle as needed**.
+> You can still pass a flag, `dispose_engines=True`, to the constructor to dispose the engines after the comparison is complete.
+
+
+### Aliases
+
 You can use meaningful aliases for the results:
 
 ```python
 result = comparer.compare(one_alias='production', two_alias='staging')
 ```
 
+## Inspectors
+
 The built-in inspectors includes: **tables**, **columns**, **primary keys**, **foreign keys**, **indexes**, **unique constraints**, **check constraints**, and **enums**.
 
-### To ignore specific inspectors:
+### Ignoring inspectors:
+
+To ignore specific inspectors, you can pass a list of inspector keys to the `compare` method.
 
-For example, ignore enums and check constraints inspectors
+For example, to ignore enums and check constraints inspectors:
 
 ```python
 result = comparer.compare(ignore_inspectors=['enums', 'check_constraints'])
 ```
 
-## Custom Inspectors
+### Custom Inspectors
 
 You can create your own custom inspectors to compare specific aspects of your database schemas.
 
@@ -140,12 +153,11 @@ class MyCustomInspector(BaseInspector, DiffMixin):
         return hasattr(inspector, 'get_something')
 ```
 
-### Important Notes
-
-- The `key` attribute must be unique, non-empty and must not start or end with whitespace
-- Use the `DiffMixin` helper methods (`_listdiff`, `_dictdiff`, `_itemsdiff`) for consistent comparison logic
+> [!IMPORTANT]
+> - The `key` attribute must be unique, non-empty and must not start or end with whitespace
+> - Use the `DiffMixin` helper methods (`_listdiff`, `_dictdiff`, `_itemsdiff`) for consistent comparison logic
 
-### Example: A Custom Sequences Inspector
+#### Example: A Custom Sequences Inspector
 
 This is a working example of a inspector that compares sequences.
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlalchemy-diff"
-version = "1.1.0"
+version = "1.1.1"
 authors = [
     { name = "Fabrizio Romano", email = "gianchub@gmail.com" },
     { name = "Mark McArdle", email = "m.mc4rdle@gmail.com" },

src/sqlalchemydiff/comparer.py

@@ -88,16 +88,23 @@ class Comparer:
 
     Simply call the `compare` method to get the result.
 
+    You can pass a flag, `dispose_engines`, to the constructor to dispose the engines after
+    the comparison is complete. If you use the `from_params` classmethod, the engines will
+    automatically be disposed after the comparison is complete.
+
     You can customise how certain aspects of the comparison are performed by setting your own
     classes for the `ignore_spec_factory` and `compare_result_class` attributes.
     """
 
     ignore_spec_factory_class = IgnoreSpecFactory
     compare_result_class = CompareResult
 
-    def __init__(self, db_one_engine: Engine, db_two_engine: Engine):
+    def __init__(
+        self, db_one_engine: Engine, db_two_engine: Engine, *, dispose_engines: bool = False
+    ):
         self.db_one_engine = db_one_engine
         self.db_two_engine = db_two_engine
+        self.__dispose_engines = dispose_engines
 
     @classmethod
     def from_params(
@@ -112,7 +119,7 @@ def from_params(
         db_one_engine = DBConnectionFactory.create_engine(db_one_uri, **db_one_params)
         db_two_engine = DBConnectionFactory.create_engine(db_two_uri, **db_two_params)
 
-        return cls(db_one_engine, db_two_engine)
+        return cls(db_one_engine, db_two_engine, dispose_engines=True)
 
     def compare(
         self,
@@ -125,18 +132,23 @@ def compare(
 
         filtered_inspectors = self._filter_inspectors(set(ignore_inspectors or set()))
 
-        result = {}
-        with self.db_one_engine.begin(), self.db_two_engine.begin():
-            for key, inspector_class in filtered_inspectors:
-                inspector = inspector_class(one_alias=one_alias, two_alias=two_alias)
+        try:
+            result = {}
+            with self.db_one_engine.begin(), self.db_two_engine.begin():
+                for key, inspector_class in filtered_inspectors:
+                    inspector = inspector_class(one_alias=one_alias, two_alias=two_alias)
 
-                db_one_info = self._get_db_info(ignore_specs, inspector, self.db_one_engine)
-                db_two_info = self._get_db_info(ignore_specs, inspector, self.db_two_engine)
+                    db_one_info = self._get_db_info(ignore_specs, inspector, self.db_one_engine)
+                    db_two_info = self._get_db_info(ignore_specs, inspector, self.db_two_engine)
 
-                if db_one_info is not None and db_two_info is not None:
-                    result[key] = inspector.diff(db_one_info, db_two_info)
+                    if db_one_info is not None and db_two_info is not None:
+                        result[key] = inspector.diff(db_one_info, db_two_info)
 
-        return self.compare_result_class(result, one_alias=one_alias, two_alias=two_alias)
+            return self.compare_result_class(result, one_alias=one_alias, two_alias=two_alias)
+
+        finally:
+            if self.__dispose_engines:
+                self._dispose_engines()
 
     def _filter_inspectors(
         self, ignore_inspectors: set[str] | None
@@ -157,3 +169,8 @@ def _get_db_info(
             return inspector.inspect(engine, ignore_specs)
         except InspectorNotSupported as e:
             logger.warning({"engine": engine, "inspector": inspector.key, "error": e.message})
+
+    def _dispose_engines(self) -> None:
+        """Dispose engines to close any pooled connections."""
+        self.db_one_engine.dispose()
+        self.db_two_engine.dispose()

tests/base.py

@@ -27,11 +27,19 @@ def db_uri_two(self, make_postgres_uri, db_name_two):
 
     @pytest.fixture
     def db_engine_one(self, db_uri_one):
-        return get_engine(db_uri_one)
+        engine = get_engine(db_uri_one)
+        try:
+            yield engine
+        finally:
+            engine.dispose()
 
     @pytest.fixture
     def db_engine_two(self, db_uri_two):
-        return get_engine(db_uri_two)
+        engine = get_engine(db_uri_two)
+        try:
+            yield engine
+        finally:
+            engine.dispose()
 
     @pytest.fixture
     def setup_db_one(self, db_uri_one, db_engine_one):

tests/test_comparer.py

@@ -1,6 +1,7 @@
 import json
 import logging
-from unittest.mock import patch
+from contextlib import nullcontext
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -135,8 +136,8 @@ def test_compare_transaction(
             result = comparer.compare()
             assert result.result == compare_result
             assert result.errors == compare_errors
-            mock_begin_one.assert_called_once_with()
-            mock_begin_two.assert_called_once_with()
+            mock_begin_one.assert_called_once()
+            mock_begin_two.assert_called_once()
 
     @pytest.mark.usefixtures("setup_db_one", "setup_db_two")
     def test_dump_result(
@@ -158,6 +159,39 @@ def test_dump_result(
             assert json.load(f) == compare_errors
 
 
+class TestComparerEngineDisposal(BaseTest):
+    @pytest.fixture
+    def engines(self):
+        engine_one = MagicMock()
+        engine_two = MagicMock()
+        engine_one.begin.side_effect = lambda: nullcontext()
+        engine_two.begin.side_effect = lambda: nullcontext()
+        return engine_one, engine_two
+
+    def test_from_params_disposes_engines(self, monkeypatch, engines):
+        engine_one, engine_two = engines
+
+        monkeypatch.setattr(
+            "sqlalchemydiff.connection.DBConnectionFactory.create_engine",
+            MagicMock(side_effect=[engine_one, engine_two]),
+        )
+
+        comparer = Comparer.from_params("postgresql://db_one", "postgresql://db_two")
+        comparer.compare()
+
+        engine_one.dispose.assert_called_once()
+        engine_two.dispose.assert_called_once()
+
+    def test_does_not_dispose_passed_engines(self, engines):
+        engine_one, engine_two = engines
+
+        comparer = Comparer(engine_one, engine_two)
+        comparer.compare()
+
+        engine_one.dispose.assert_not_called()
+        engine_two.dispose.assert_not_called()
+
+
 @pytest.mark.is_sqlalchemy_1_4
 class TestComparerV14(BaseTest):
     @pytest.mark.usefixtures("setup_db_one", "setup_db_two")
@@ -171,11 +205,19 @@ def test_compare(self, db_engine_one, db_engine_two, compare_result_v14, compare
 class TestComparerSqlite(BaseTest):
     @pytest.fixture
     def sqlite_db_engine_one(self):
-        return get_engine("sqlite:///:memory:")
+        engine = get_engine("sqlite:///:memory:")
+        try:
+            yield engine
+        finally:
+            engine.dispose()
 
     @pytest.fixture
     def sqlite_db_engine_two(self):
-        return get_engine("sqlite:///:memory:")
+        engine = get_engine("sqlite:///:memory:")
+        try:
+            yield engine
+        finally:
+            engine.dispose()
 
     @pytest.fixture()
     def setup_db_one(self, sqlite_db_engine_one):