This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-0-test by this push:
new 5c0781193e8 [v3-0-test] Allow downgrading to 2.11 from 3.x (#54371)
(#54399)
5c0781193e8 is described below
commit 5c0781193e86a78cef460ccf2b8492a023dd1fcc
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Wed Aug 13 10:01:09 2025 +0100
[v3-0-test] Allow downgrading to 2.11 from 3.x (#54371) (#54399)
---
.../src/airflow/cli/commands/db_command.py | 30 ++++++++++++++--------
airflow-core/src/airflow/utils/db.py | 28 ++++++++++----------
airflow-core/src/airflow/utils/db_manager.py | 6 -----
.../tests/unit/cli/commands/test_db_command.py | 16 ++++++++++++
airflow-core/tests/unit/utils/test_db.py | 14 ----------
airflow-core/tests/unit/utils/test_db_manager.py | 22 ++--------------
devel-common/src/tests_common/test_utils/db.py | 22 ++++++++++++++--
.../providers/fab/auth_manager/models/db.py | 5 ++++
8 files changed, 77 insertions(+), 66 deletions(-)
diff --git a/airflow-core/src/airflow/cli/commands/db_command.py
b/airflow-core/src/airflow/cli/commands/db_command.py
index 5a4f70cc604..f267f598b02 100644
--- a/airflow-core/src/airflow/cli/commands/db_command.py
+++ b/airflow-core/src/airflow/cli/commands/db_command.py
@@ -50,29 +50,39 @@ def resetdb(args):
db.resetdb(skip_init=args.skip_init)
-def _get_version_revision(
- version: str, recursion_limit: int = 10, revision_heads_map: dict[str,
str] | None = None
-) -> str | None:
+def _get_version_revision(version: str, revision_heads_map: dict[str, str] |
None = None) -> str | None:
"""
- Recursively search for the revision of the given version in
revision_heads_map.
+ Search for the revision of the given version in revision_heads_map.
This searches given revision_heads_map for the revision of the given
version, recursively
searching for the previous version if the given version is not found.
+
+ ``revision_heads_map`` must already be sorted in the dict in ascending
order for this function to work. No
+ checks are made that this is true
"""
if revision_heads_map is None:
revision_heads_map = _REVISION_HEADS_MAP
+ # Exact match found, we can just return it
if version in revision_heads_map:
return revision_heads_map[version]
+
try:
- major, minor, patch = map(int, version.split("."))
+ wanted = tuple(map(int, version.split(".")))
except ValueError:
return None
- new_version = f"{major}.{minor}.{patch - 1}"
- recursion_limit -= 1
- if recursion_limit <= 0:
- # Prevent infinite recursion as I can't imagine 10 successive versions
without migration
+
+ # Else, we walk backwards in the revision map until we find a version that
is < the target
+ for revision, head in reversed(revision_heads_map.items()):
+ try:
+ current = tuple(map(int, revision.split(".")))
+ except ValueError:
+ log.debug("Unable to parse HEAD revision", exc_info=True)
+ return None
+
+ if current < wanted:
+ return head
+ else:
return None
- return _get_version_revision(new_version, recursion_limit)
def run_db_migrate_command(args, command, revision_heads_map: dict[str, str]):
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index f091f0868b0..a0dcb5fdf03 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -1211,21 +1211,21 @@ def downgrade(*, to_revision, from_revision=None,
show_sql_only=False, session:
config = _get_alembic_config()
# Check if downgrade is less than 3.0.0 and requires that `ab_user` fab
table is present
if _revision_greater(config, _REVISION_HEADS_MAP["2.10.3"], to_revision):
- unitest_mode = conf.getboolean("core", "unit_test_mode")
- if unitest_mode:
- try:
- from airflow.providers.fab.auth_manager.models.db import
FABDBManager
-
- dbm = FABDBManager(session)
- dbm.initdb()
- except ImportError:
- log.warning("Import error occurred while importing
FABDBManager. Skipping the check.")
- return
- if not inspect(settings.engine).has_table("ab_user") and not
unitest_mode:
- raise AirflowException(
- "Downgrade to revision less than 3.0.0 requires that `ab_user`
table is present. "
- "Please add FabDBManager to [core] external_db_managers and
run fab migrations before proceeding"
+ try:
+ from airflow.providers.fab.auth_manager.models.db import
FABDBManager
+ except ImportError:
+ # Raise the error with a new message
+ raise RuntimeError(
+ "Import error occurred while importing FABDBManager. We need
that to exist before we can "
+ "downgrade to <3.0.0"
)
+ dbm = FABDBManager(session)
+ if hasattr(dbm, "reset_to_2_x"):
+ dbm.reset_to_2_x()
+ else:
+ # Older version before we added that function, it only has a
single migration so we can just
+ # created
+ dbm.create_db_from_orm()
with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
if show_sql_only:
log.warning("Generating sql scripts for manual migration.")
diff --git a/airflow-core/src/airflow/utils/db_manager.py
b/airflow-core/src/airflow/utils/db_manager.py
index 6483c446f01..a11c4cdad25 100644
--- a/airflow-core/src/airflow/utils/db_manager.py
+++ b/airflow-core/src/airflow/utils/db_manager.py
@@ -213,12 +213,6 @@ class RunDBManager(LoggingMixin):
m = manager(session)
m.upgradedb()
- def downgrade(self, session):
- """Downgrade the external database managers."""
- for manager in self._managers:
- m = manager(session)
- m.downgrade()
-
def drop_tables(self, session, connection):
"""Drop the external database managers."""
for manager in self._managers:
diff --git a/airflow-core/tests/unit/cli/commands/test_db_command.py
b/airflow-core/tests/unit/cli/commands/test_db_command.py
index e6d1b9bd6d9..b9f5e0f04d0 100644
--- a/airflow-core/tests/unit/cli/commands/test_db_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_db_command.py
@@ -624,3 +624,19 @@ class TestCLIDBClean:
)
db_command.drop_archived(args)
mock_drop_archived_records.assert_called_once_with(table_names=None,
needs_confirm=expected)
+
+
+def test_get_version_revision():
+ heads: dict[str, str] = {
+ "2.10.0": "22ed7efa9da2",
+ "2.10.3": "5f2621c13b39",
+ "3.0.0": "29ce7909c52b",
+ "3.0.3": "fe199e1abd77",
+ "3.1.0": "808787349f22",
+ }
+
+ assert db_command._get_version_revision("3.1.0", heads) == "808787349f22"
+ assert db_command._get_version_revision("3.1.1", heads) == "808787349f22"
+ assert db_command._get_version_revision("2.11.1", heads) == "5f2621c13b39"
+ assert db_command._get_version_revision("2.10.1", heads) == "22ed7efa9da2"
+ assert db_command._get_version_revision("2.0.0", heads) is None
diff --git a/airflow-core/tests/unit/utils/test_db.py
b/airflow-core/tests/unit/utils/test_db.py
index cf692750253..b17c523ce27 100644
--- a/airflow-core/tests/unit/utils/test_db.py
+++ b/airflow-core/tests/unit/utils/test_db.py
@@ -34,12 +34,10 @@ from alembic.runtime.environment import EnvironmentContext
from alembic.script import ScriptDirectory
from sqlalchemy import Column, Integer, MetaData, Table, select
-from airflow.exceptions import AirflowException
from airflow.models import Base as airflow_base
from airflow.providers.fab.auth_manager.models.db import FABDBManager
from airflow.settings import engine
from airflow.utils.db import (
- _REVISION_HEADS_MAP,
LazySelectSequence,
_get_alembic_config,
check_migrations,
@@ -293,15 +291,3 @@ class TestDb:
lss = LazySelectSequence.from_select(select(t.c.id), order_by=[],
session=MockSession())
assert bool(lss) is False
-
- @conf_vars({("core", "unit_test_mode"): "False"})
- @mock.patch("airflow.utils.db.inspect")
- def test_downgrade_raises_if_lower_than_v3_0_0_and_no_ab_user(self,
mock_inspect):
- mock_inspect.return_value.has_table.return_value = False
- msg = (
- "Downgrade to revision less than 3.0.0 requires that `ab_user`
table is present. "
- "Please add FabDBManager to [core] external_db_managers and run
fab migrations before "
- "proceeding"
- )
- with pytest.raises(AirflowException, match=re.escape(msg)):
- downgrade(to_revision=_REVISION_HEADS_MAP["2.7.0"])
diff --git a/airflow-core/tests/unit/utils/test_db_manager.py
b/airflow-core/tests/unit/utils/test_db_manager.py
index 075e87c0cbe..501fa02ec38 100644
--- a/airflow-core/tests/unit/utils/test_db_manager.py
+++ b/airflow-core/tests/unit/utils/test_db_manager.py
@@ -23,7 +23,7 @@ from sqlalchemy import Table
from airflow.exceptions import AirflowException
from airflow.models import Base
-from airflow.utils.db import downgrade, initdb
+from airflow.utils.db import initdb
from airflow.utils.db_manager import BaseDBManager, RunDBManager
from tests_common.test_utils.config import conf_vars
@@ -58,27 +58,12 @@ class TestRunDBManager:
run_db_manager.validate()
metadata._remove_table("dag_run", None)
- @mock.patch.object(RunDBManager, "downgrade")
@mock.patch.object(RunDBManager, "upgradedb")
@mock.patch.object(RunDBManager, "initdb")
- def test_init_db_calls_rundbmanager(self, mock_initdb, mock_upgrade_db,
mock_downgrade_db, session):
+ def test_init_db_calls_rundbmanager(self, mock_initdb, mock_upgrade_db,
session):
initdb(session=session)
mock_initdb.assert_called()
mock_initdb.assert_called_once_with(session)
- mock_downgrade_db.assert_not_called()
-
- @mock.patch.object(RunDBManager, "downgrade")
- @mock.patch.object(RunDBManager, "upgradedb")
- @mock.patch.object(RunDBManager, "initdb")
- @mock.patch("alembic.command")
- def test_downgrade_dont_call_rundbmanager(
- self, mock_alembic_command, mock_initdb, mock_upgrade_db,
mock_downgrade_db, session
- ):
- downgrade(to_revision="base")
- mock_alembic_command.downgrade.assert_called_once_with(mock.ANY,
revision="base", sql=False)
- mock_upgrade_db.assert_not_called()
- mock_initdb.assert_not_called()
- mock_downgrade_db.assert_not_called()
@conf_vars(
{("database", "external_db_managers"):
"airflow.providers.fab.auth_manager.models.db.FABDBManager"}
@@ -94,9 +79,6 @@ class TestRunDBManager:
# upgradedb
ext_db.upgradedb(session=session)
fabdb_manager.upgradedb.assert_called_once()
- # downgrade
- ext_db.downgrade(session=session)
- mock_fabdb_manager.return_value.downgrade.assert_called_once()
connection = mock.MagicMock()
ext_db.drop_tables(session, connection)
mock_fabdb_manager.return_value.drop_tables.assert_called_once_with(connection)
diff --git a/devel-common/src/tests_common/test_utils/db.py
b/devel-common/src/tests_common/test_utils/db.py
index 4ab88d1c75e..8b3b85c33b6 100644
--- a/devel-common/src/tests_common/test_utils/db.py
+++ b/devel-common/src/tests_common/test_utils/db.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import os
from typing import TYPE_CHECKING
from airflow.configuration import conf
@@ -92,9 +93,26 @@ def initial_db_init():
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
db.resetdb()
+
if AIRFLOW_V_3_0_PLUS:
- db.downgrade(to_revision="5f2621c13b39")
- db.upgradedb(to_revision="head")
+ try:
+ from airflow.providers.fab.auth_manager.models.db import
FABDBManager
+ except ModuleNotFoundError:
+ # Reasons it might fail: we're in a provider bundle without FAB,
or we're on a version of Python
+ # where FAB isn't yet supported
+ pass
+ else:
+ if os.getenv("TEST_GROUP") != "providers":
+ # If we loaded the provider, and we're running core (or
running via breeze where TEST_GROUP
+ # isn't specified) run the downgrade+upgrade to ensure
migrations are in sync with Model
+ # classes
+ db.downgrade(to_revision="5f2621c13b39")
+ db.upgradedb(to_revision="head")
+ else:
+ # Just create the tables so they are there
+ with create_session() as session:
+ FABDBManager(session).create_db_from_orm()
+ session.commit()
else:
from flask import Flask
diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/models/db.py
b/providers/fab/src/airflow/providers/fab/auth_manager/models/db.py
index 8ffb6cfeab8..1e055085ba1 100644
--- a/providers/fab/src/airflow/providers/fab/auth_manager/models/db.py
+++ b/providers/fab/src/airflow/providers/fab/auth_manager/models/db.py
@@ -58,6 +58,11 @@ class FABDBManager(BaseDBManager):
super().create_db_from_orm()
_get_flask_db(settings.SQL_ALCHEMY_CONN).create_all()
+ def reset_to_2_x(self):
+ self.create_db_from_orm()
+ # And ensure it's at the oldest version
+ self.downgrade(_REVISION_HEADS_MAP["1.4.0"])
+
def upgradedb(self, to_revision=None, from_revision=None,
show_sql_only=False):
"""Upgrade the database."""
if from_revision and not show_sql_only:
