Revert "Allow downgrading to 2.11 from 3.x (apache#54231)" (apache#54367)

This reverts commit b9cdc3d.

Author: ashb

airflow-core/src/airflow/cli/commands/db_command.py

@@ -50,39 +50,29 @@ def resetdb(args):
     db.resetdb(skip_init=args.skip_init)
 
 
-def _get_version_revision(version: str, revision_heads_map: dict[str, str] | None = None) -> str | None:
+def _get_version_revision(
+    version: str, recursion_limit: int = 10, revision_heads_map: dict[str, str] | None = None
+) -> str | None:
     """
-    Search for the revision of the given version in revision_heads_map.
+    Recursively 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:
-        wanted = tuple(map(int, version.split(".")))
+        major, minor, patch = map(int, version.split("."))
     except ValueError:
         return None
-
-    # 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:
+    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
         return None
+    return _get_version_revision(new_version, recursion_limit)
 
 
 def run_db_migrate_command(args, command, revision_heads_map: dict[str, str]):

airflow-core/src/airflow/utils/db.py

@@ -1189,21 +1189,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):
-        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"
+        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"
             )
-        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.")

airflow-core/src/airflow/utils/db_manager.py

@@ -216,6 +216,12 @@ def upgradedb(self, session):
             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:

airflow-core/tests/unit/cli/commands/test_db_command.py

@@ -658,19 +658,3 @@ def test_confirm_in_drop_archived_records_command(self, mock_drop_archived_recor
         )
         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

airflow-core/tests/unit/utils/test_db.py

@@ -33,8 +33,10 @@
 from sqlalchemy import Column, Integer, MetaData, Table, select
 
 from airflow import settings
+from airflow.exceptions import AirflowException
 from airflow.models import Base as airflow_base
 from airflow.utils.db import (
+    _REVISION_HEADS_MAP,
     AutocommitEngineForMySQL,
     LazySelectSequence,
     _get_alembic_config,
@@ -378,6 +380,18 @@ def scalar(self, stmt):
 
         assert bool(lss) is False
 
+    @conf_vars({("core", "unit_test_mode"): "False"})
+    def test_downgrade_raises_if_lower_than_v3_0_0_and_no_ab_user(self, mocker):
+        mock_inspect = mocker.patch("airflow.utils.db.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"])
+
 
 class TestAutocommitEngineForMySQL:
     """Test the AutocommitEngineForMySQL context manager."""

devel-common/src/tests_common/test_utils/db.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import json
-import os
 from tempfile import gettempdir
 from typing import TYPE_CHECKING
 
@@ -102,23 +101,9 @@ def initial_db_init():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
     db.resetdb()
-    # If we are testing core we want to downgrand and upgrade to ensure things migraitons are set correctly,
-    # but we can skip that for providers
-
     if AIRFLOW_V_3_0_PLUS:
-        if os.getenv("TEST_GROUP") != "providers":
-            db.downgrade(to_revision="5f2621c13b39")
-            db.upgradedb(to_revision="head")
-        else:
-            try:
-                from airflow.providers.fab.auth_manager.models.db import FABDBManager
-            except ModuleNotFoundError:
-                pass
-            else:
-                # If we loaded it (i.e. the provider exists) create the DB
-                with create_session() as session:
-                    FABDBManager(session).create_db_from_orm()
-                    session.commit()
+        db.downgrade(to_revision="5f2621c13b39")
+        db.upgradedb(to_revision="head")
     else:
         from flask import Flask
 

providers/fab/src/airflow/providers/fab/auth_manager/models/db.py

@@ -59,11 +59,6 @@ def create_db_from_orm(self):
         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:

providers/fab/tests/unit/fab/db_manager/test_fab_db_manager.py

@@ -22,7 +22,7 @@
 from sqlalchemy import Table
 
 from airflow.exceptions import AirflowException
-from airflow.utils.db import initdb
+from airflow.utils.db import downgrade, initdb
 from airflow.utils.db_manager import RunDBManager
 
 from tests_common.test_utils.config import conf_vars
@@ -57,12 +57,27 @@ def test_defining_table_same_name_as_airflow_table_name_raises(self):
             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, session):
+    def test_init_db_calls_rundbmanager(self, mock_initdb, mock_upgrade_db, mock_downgrade_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"}
@@ -78,7 +93,9 @@ def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session)
         # upgradedb
         ext_db.upgradedb(session=session)
         fabdb_manager.upgradedb.assert_called_once()
-        # drop_tables
+        # 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)