Separate FAB migration from Core Airflow migration (apache#41437)

* Separate FAB migration from Core Airflow migration

This PR separates FAB migration from Airflow Core migration and provides
a way for apps to integrate into Airflow and run their migrations.

* minor fixes and use pre-commit on fab migrations

* support external db managers to drop tables when airflow db reset is run

* Add tests

* fix autogenerate and update function signatures

* Rename alembic directory to migrations

* add more tests

* Fix import

* Fix static check

* Fix tests

* fixup! Fix tests

* fixup! fixup! Fix tests

* Mark extra operator links test as non-db test

* update erd diagram

* import airflow settings at point of use

* skip import error for alembic env.py

* Remove type on typing

* Move code around

* fixup! Move code around

* Move import

* fixup! Move import

* Don't disable existing loggers

* Use abstract base and update tests

* Revert compat change

* Fix backcompat for fab provider

* Build fab provider for www test instead of installing from pypi

* fixup! Fix backcompat for fab provider

* Mark downgradedb method as abstract

* Ensure supports_table_dropping works and improve tests

* Skip tag check when preparing fab provider in CI

.github/workflows/basic-tests.yml

@@ -154,8 +154,11 @@ jobs:
         env:
           HATCH_ENV: "test"
         working-directory: ./clients/python
+      - name: "Prepare FAB provider packages: wheel"
+        run: >
+            breeze release-management prepare-provider-packages fab --package-format wheel --skip-tag-check
       - name: "Install Airflow with fab for webserver tests"
-        run: pip install ".[fab]"
+        run: pip install . dist/apache_airflow_providers_fab-*.whl
       - name: "Install Python client"
         run: pip install ./dist/apache_airflow_client-*.whl
       - name: "Initialize Airflow DB and start webserver"

.pre-commit-config.yaml

@@ -1295,7 +1295,8 @@ repos:
         language: python
         entry: ./scripts/ci/pre_commit/migration_reference.py
         pass_filenames: false
-        files: ^airflow/migrations/versions/.*\.py$|^docs/apache-airflow/migrations-ref\.rst$
+        files: |
+          ^airflow/migrations/versions/.*\.py$|^docs/apache-airflow/migrations-ref\.rst$|^airflow/providers/fab/alembic/versions/.*\.py$
         additional_dependencies: ['rich>=12.4.4']
       - id: update-er-diagram
         name: Update ER diagram

airflow/config_templates/config.yml

@@ -720,6 +720,14 @@ database:
       type: string
       example: ~
       default: "True"
+    external_db_managers:
+      description: |
+        List of DB managers to use to migrate external tables in airflow database. The managers must inherit
+        from BaseDBManager
+      version_added: 3.0.0
+      type: string
+      example: ~
+      default: "airflow.providers.fab.auth_manager.models.db.FABDBManager"
 logging:
   description: ~
   options:

airflow/migrations/env.py

@@ -32,8 +32,8 @@ def include_object(_, name, type_, *args):
     # Ignore the sqlite_sequence table, which is an internal SQLite construct
     if name == "sqlite_sequence":
         return False
-    # Ignore _anything_ to do with Celery, or FlaskSession's tables
-    if type_ == "table" and (name.startswith("celery_") or name == "session"):
+    # Only create migrations for objects that are in the target metadata
+    if type_ == "table" and name not in target_metadata.tables:
         return False
     else:
         return True
@@ -61,6 +61,9 @@ def include_object(_, name, type_, *args):
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
+# version table
+version_table = "alembic_version"
+
 
 def run_migrations_offline():
     """
@@ -82,6 +85,8 @@ def run_migrations_offline():
         compare_type=compare_type,
         compare_server_default=compare_server_default,
         render_as_batch=True,
+        include_object=include_object,
+        version_table=version_table,
     )
 
     with context.begin_transaction():
@@ -119,6 +124,7 @@ def process_revision_directives(context, revision, directives):
             include_object=include_object,
             render_as_batch=True,
             process_revision_directives=process_revision_directives,
+            version_table=version_table,
         )
 
         with context.begin_transaction():

airflow/providers/fab/alembic.ini

@@ -0,0 +1,133 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+# Use forward slashes (/) also on windows to provide an os agnostic path
+script_location = migrations
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "version_path_separator" below.
+# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
+
+# version path separator; As mentioned above, this is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
+# Valid values for version_path_separator are:
+#
+# version_path_separator = :
+# version_path_separator = ;
+# version_path_separator = space
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+sqlalchemy.url = scheme://localhost/airflow
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
+# hooks = ruff
+# ruff.type = exec
+# ruff.executable = %(here)s/.venv/bin/ruff
+# ruff.options = --fix REVISION_SCRIPT_FILENAME
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

airflow/providers/fab/auth_manager/models/__init__.py

@@ -23,6 +23,7 @@
 # Copyright 2013, Daniel Vaz Gaspar
 from typing import TYPE_CHECKING
 
+import packaging.version
 from flask import current_app, g
 from flask_appbuilder.models.sqla import Model
 from sqlalchemy import (
@@ -32,30 +33,42 @@
     ForeignKey,
     Index,
     Integer,
+    MetaData,
     String,
     Table,
     UniqueConstraint,
     event,
     func,
     select,
 )
-from sqlalchemy.orm import backref, declared_attr, relationship
+from sqlalchemy.orm import backref, declared_attr, registry, relationship
 
+from airflow import __version__ as airflow_version
 from airflow.auth.managers.models.base_user import BaseUser
-from airflow.models.base import Base
-
-"""
-Compatibility note: The models in this file are duplicated from Flask AppBuilder.
-"""
-# Use airflow metadata to create the tables
-Model.metadata = Base.metadata
+from airflow.models.base import _get_schema, naming_convention
 
 if TYPE_CHECKING:
     try:
         from sqlalchemy import Identity
     except Exception:
         Identity = None
 
+"""
+Compatibility note: The models in this file are duplicated from Flask AppBuilder.
+"""
+
+metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention)
+mapper_registry = registry(metadata=metadata)
+
+if packaging.version.parse(packaging.version.parse(airflow_version).base_version) >= packaging.version.parse(
+    "3.0.0"
+):
+    Model.metadata = metadata
+else:
+    from airflow.models.base import Base
+
+    Model.metadata = Base.metadata
+
 
 class Action(Model):
     """Represents permission actions such as `can_read`."""

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

@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import os
+
+import airflow
+from airflow.providers.fab.auth_manager.models import metadata
+from airflow.utils.db_manager import BaseDBManager
+
+PACKAGE_DIR = os.path.dirname(airflow.__file__)
+
+
+class FABDBManager(BaseDBManager):
+    """Manages FAB database."""
+
+    metadata = metadata
+    version_table_name = "fab_alembic_version"
+    migration_dir = os.path.join(PACKAGE_DIR, "providers/fab/migrations")
+    alembic_file = os.path.join(PACKAGE_DIR, "providers/fab/alembic.ini")
+    supports_table_dropping = True

airflow/providers/fab/migrations/README

@@ -0,0 +1 @@
+Generic single-database configuration.

airflow/providers/fab/migrations/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.