This is an automated email from the ASF dual-hosted git repository. kaxilnaik pushed a commit to branch v3-1-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit bee83f6a8458b396c5ce41f086b598e25858adca Author: Ephraim Anierobi <[email protected]> AuthorDate: Sat Sep 20 15:04:16 2025 +0100 Change the signature of BaseDBManager downgrade method (#55907) (cherry picked from commit bff6c9503b28806ef51d2c65ed1d464ae1ac0568) --- airflow-core/src/airflow/utils/db_manager.py | 2 +- airflow-core/tests/unit/utils/test_db_manager.py | 30 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/utils/db_manager.py b/airflow-core/src/airflow/utils/db_manager.py index 59f4921ce92..995cb81c803 100644 --- a/airflow-core/src/airflow/utils/db_manager.py +++ b/airflow-core/src/airflow/utils/db_manager.py @@ -131,7 +131,7 @@ class BaseDBManager(LoggingMixin): command.upgrade(config, revision=to_revision or "heads", sql=show_sql_only) self.log.info("Migrated the %s database", self.__class__.__name__) - def downgrade(self, to_version, from_version=None, show_sql_only=False): + def downgrade(self, to_revision, from_revision=None, show_sql_only=False): """Downgrade the database.""" raise NotImplementedError diff --git a/airflow-core/tests/unit/utils/test_db_manager.py b/airflow-core/tests/unit/utils/test_db_manager.py index e46ac4d4265..529dacf58be 100644 --- a/airflow-core/tests/unit/utils/test_db_manager.py +++ b/airflow-core/tests/unit/utils/test_db_manager.py @@ -34,6 +34,19 @@ class MockDBManager(BaseDBManager): supports_table_dropping = True +class CustomDBManager(BaseDBManager): + metadata = Base.metadata + version_table_name = "custom_alembic_version" + migration_dir = "custom_migration_dir" + alembic_file = "custom_alembic.ini" + + def downgrade(self, to_revision, from_revision=None, show_sql_only=False): + from alembic import command as alembic_command + + config = self.get_alembic_config() + alembic_command.downgrade(config, revision=to_revision, sql=show_sql_only) + + class TestBaseDBManager: @mock.patch.object(BaseDBManager, "get_alembic_config") @mock.patch.object(BaseDBManager, "get_current_revision") @@ -60,3 +73,20 @@ class TestBaseDBManager: def test_check_migration(self, mock_script_obj, mock_current_revision, session): manager = MockDBManager(session) manager.check_migration() # just ensure this can be called + + def test_custom_db_manager_downgrade_uses_revision_kwarg(self, session): + manager = CustomDBManager(session) + with ( + mock.patch.object(BaseDBManager, "get_alembic_config") as mock_config, + mock.patch("alembic.command.downgrade") as mock_alembic_downgrade, + ): + cfg = object() + mock_config.return_value = cfg + manager.downgrade(to_revision="abc123", show_sql_only=True) + mock_alembic_downgrade.assert_called_once_with(cfg, revision="abc123", sql=True) + + def test_custom_db_manager_downgrade_rejects_to_version_kwarg(self, session): + manager = CustomDBManager(session) + with pytest.raises(TypeError): + # Ensure the old kwarg name is not accepted anymore + manager.downgrade(to_version="1.2.3") # type: ignore[call-arg]
