Allow downgrading to 2.11 from 3.x

There were two things blocking this:

1) The revision heads map didn't have any 2.11.x versions in it, so the
   previous implementation of `_get_version_revision` was only looking within
   the same <major.minor> pathc version.

   We change it to rely on the fact that our pre-commit checks ensure this map
   is ordered, and iterate over the dictionary reversed, and when we find the
   first thing less than the target revision we use that (direct equal is
   handled already above)

2) The `ab_*` tables not existing were blocking the migration. Part of this is
   now fixable manually with apache#54227, but I have decided that since FAB was
   required and the only option in 2.x, so I have decided to just create the
   tables if they are missing

   In order to try and cope with possible future changes I create the tables
   at the latest version and then downgrade to the oldest known revision.

   This is all handled in a `reset_to_2_x()` method on the FABDBManager, with
   a fallback to just blindly create the tables from the ORM for versions of
   the provider that don't yet have that function.

Author: ashb

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

@@ -50,29 +50,35 @@ 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("."))
-    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
+
+    wanted = tuple(map(int, version.split(".")))
+    # 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]):

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):
-        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.")

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

@@ -658,3 +658,19 @@ 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

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

@@ -59,6 +59,11 @@ 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: