This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9004b834746 Fix FAB ignoring create_metadata_engine from local
settings (#71674)
9004b834746 is described below
commit 9004b8347461e69dd95570584dab13a5cc4db7ee
Author: rjgoyln <[email protected]>
AuthorDate: Thu Sep 10 05:37:17 2026 +0800
Fix FAB ignoring create_metadata_engine from local settings (#71674)
Deployments that authenticate to the metadata database with short-lived
credentials override create_metadata_engine in airflow_local_settings.py so
a
do_connect handler can mint a token per physical connection.
Flask-SQLAlchemy
built a second engine straight from the connection URI, so those handlers
never
ran and `airflow users create` — the Helm chart's create-user job —
connected
without a password while the migrate job and the API server succeeded.
---
.../fab/auth_manager/cli_commands/utils.py | 9 ++-
.../providers/fab/auth_manager/models/db.py | 4 +-
providers/fab/src/airflow/providers/fab/www/app.py | 4 +-
providers/fab/src/airflow/providers/fab/www/db.py | 50 +++++++++++++++
.../fab/auth_manager/cli_commands/test_utils.py | 15 +++++
providers/fab/tests/unit/fab/www/test_app.py | 38 +++++++++++
providers/fab/tests/unit/fab/www/test_db.py | 75 ++++++++++++++++++++++
7 files changed, 188 insertions(+), 7 deletions(-)
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/cli_commands/utils.py
b/providers/fab/src/airflow/providers/fab/auth_manager/cli_commands/utils.py
index a6a58c41525..2a3fb195146 100644
--- a/providers/fab/src/airflow/providers/fab/auth_manager/cli_commands/utils.py
+++ b/providers/fab/src/airflow/providers/fab/auth_manager/cli_commands/utils.py
@@ -25,12 +25,13 @@ from os.path import isabs
from typing import TYPE_CHECKING
from flask import Flask
-from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.engine import make_url
import airflow
+from airflow import settings
from airflow.exceptions import AirflowConfigException
from airflow.providers.common.compat.sdk import conf
+from airflow.providers.fab.www.db import AirflowSQLAlchemy
from airflow.providers.fab.www.extensions.init_appbuilder import
init_appbuilder
from airflow.providers.fab.www.extensions.init_session import
init_airflow_session_interface
from airflow.providers.fab.www.extensions.init_views import init_plugins
@@ -65,6 +66,8 @@ def get_application_builder() -> Generator[AirflowAppBuilder,
None, None]:
)
flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
- db = SQLAlchemy(flask_app)
+ db = AirflowSQLAlchemy(flask_app)
yield _return_appbuilder(flask_app, db)
- db.engine.dispose(close=True)
+ # The shared metadata engine outlives this helper — only a locally
built one may be disposed.
+ if db.engine is not settings.engine:
+ db.engine.dispose(close=True)
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 10b83612339..04eb86cc633 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
@@ -43,14 +43,14 @@ def _release_metadata_locks_if_supported(manager:
BaseDBManager) -> None:
def _get_flask_db(sql_database_uri):
from flask import Flask
- from flask_sqlalchemy import SQLAlchemy
+ from airflow.providers.fab.www.db import AirflowSQLAlchemy
from airflow.providers.fab.www.session import
AirflowDatabaseSessionInterface
flask_app = Flask(__name__)
flask_app.config["SQLALCHEMY_DATABASE_URI"] = sql_database_uri
flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
- db = SQLAlchemy(flask_app)
+ db = AirflowSQLAlchemy(flask_app)
AirflowDatabaseSessionInterface(app=flask_app, client=db, table="session",
key_prefix="")
return db, flask_app
diff --git a/providers/fab/src/airflow/providers/fab/www/app.py
b/providers/fab/src/airflow/providers/fab/www/app.py
index f28cd29c0db..cf7658b69c9 100644
--- a/providers/fab/src/airflow/providers/fab/www/app.py
+++ b/providers/fab/src/airflow/providers/fab/www/app.py
@@ -21,7 +21,6 @@ from datetime import timedelta
from os.path import isabs
from flask import Flask
-from flask_sqlalchemy import SQLAlchemy
from flask_wtf.csrf import CSRFProtect
from sqlalchemy.engine.url import make_url
@@ -31,6 +30,7 @@ from airflow.exceptions import AirflowConfigException
from airflow.logging_config import configure_logging
from airflow.providers.common.compat.sdk import conf
from airflow.providers.fab.version_compat import AIRFLOW_V_3_1_8_PLUS
+from airflow.providers.fab.www.db import AirflowSQLAlchemy
from airflow.providers.fab.www.extensions.init_appbuilder import
init_appbuilder
from airflow.providers.fab.www.extensions.init_jinja_globals import
init_jinja_globals
from airflow.providers.fab.www.extensions.init_manifest_files import
configure_manifest_files
@@ -104,7 +104,7 @@ def create_app(enable_plugins: bool):
csrf.init_app(flask_app)
- db = SQLAlchemy(flask_app)
+ db = AirflowSQLAlchemy(flask_app)
if settings.Session is None:
raise RuntimeError("Session not configured. Call configure_orm()
first.")
db.session = settings.Session
diff --git a/providers/fab/src/airflow/providers/fab/www/db.py
b/providers/fab/src/airflow/providers/fab/www/db.py
new file mode 100644
index 00000000000..d73e7d0c030
--- /dev/null
+++ b/providers/fab/src/airflow/providers/fab/www/db.py
@@ -0,0 +1,50 @@
+# 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
+
+from typing import TYPE_CHECKING, Any
+
+from flask_sqlalchemy import SQLAlchemy
+from sqlalchemy.engine import make_url
+
+from airflow import settings
+
+if TYPE_CHECKING:
+ from flask import Flask
+ from sqlalchemy.engine import Engine
+
+
+def _points_at_metadata_db(options: dict[str, Any]) -> bool:
+ url = options.get("url")
+ conn = settings.SQL_ALCHEMY_CONN
+ return url is not None and conn is not None and make_url(url) ==
make_url(conn)
+
+
+class AirflowSQLAlchemy(SQLAlchemy):
+ """``Flask-SQLAlchemy`` extension bound to Airflow's metadata engine."""
+
+ def _make_engine(self, bind_key: str | None, options: dict[str, Any], app:
Flask) -> Engine:
+ # A second engine built straight from the connection URI would skip
the ``do_connect``
+ # handlers that a ``settings.create_metadata_engine`` override
installs in
+ # ``airflow_local_settings.py`` to mint short-lived credentials per
connection. Only the
+ # default bind aimed at the metadata database is Airflow's to take
over — secondary binds
+ # and a ``webserver_config.py`` that points
``SQLALCHEMY_DATABASE_URI`` at another database
+ # keep the engine Flask-SQLAlchemy builds for them.
+ engine = settings.engine
+ if bind_key is None and engine is not None and
_points_at_metadata_db(options):
+ return engine
+ return super()._make_engine(bind_key, options, app)
diff --git
a/providers/fab/tests/unit/fab/auth_manager/cli_commands/test_utils.py
b/providers/fab/tests/unit/fab/auth_manager/cli_commands/test_utils.py
index 0acf228957c..4ec4e46b301 100644
--- a/providers/fab/tests/unit/fab/auth_manager/cli_commands/test_utils.py
+++ b/providers/fab/tests/unit/fab/auth_manager/cli_commands/test_utils.py
@@ -17,10 +17,12 @@
from __future__ import annotations
import os
+from unittest import mock
import pytest
import airflow
+from airflow import settings
from airflow.exceptions import AirflowConfigException
from airflow.providers.common.compat.sdk import conf
from airflow.providers.fab.auth_manager.cli_commands.utils import
get_application_builder
@@ -78,3 +80,16 @@ class TestCliUtils:
flask_app = appbuilder.app
# Ensure that the correct session interface is set (for 'database'
auth backend)
assert isinstance(flask_app.session_interface,
AirflowDatabaseSessionInterface)
+
+ def test_flask_app_binds_to_the_metadata_engine(self):
+ """A second engine would bypass a ``create_metadata_engine`` override
in local settings."""
+ with get_application_builder() as appbuilder:
+ db = appbuilder.app.extensions["sqlalchemy"]
+ assert db.engine is settings.engine
+
+ @mock.patch("airflow.settings.engine.dispose", autospec=True)
+ def test_metadata_engine_is_not_disposed(self, mock_dispose):
+ with get_application_builder():
+ pass
+
+ mock_dispose.assert_not_called()
diff --git a/providers/fab/tests/unit/fab/www/test_app.py
b/providers/fab/tests/unit/fab/www/test_app.py
new file mode 100644
index 00000000000..6ee442c3d84
--- /dev/null
+++ b/providers/fab/tests/unit/fab/www/test_app.py
@@ -0,0 +1,38 @@
+# 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 pytest
+
+from airflow import settings
+from airflow.providers.fab.www.app import create_app
+
+from tests_common.test_utils.config import conf_vars
+
+
[email protected]_test
+class TestCreateApp:
+ def test_flask_app_binds_to_the_metadata_engine(self):
+ """A second engine would bypass a ``create_metadata_engine`` override
in local settings."""
+ with conf_vars(
+ {("core", "auth_manager"):
"airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager"}
+ ):
+ flask_app = create_app(enable_plugins=False)
+
+ with flask_app.app_context():
+ db = flask_app.extensions["sqlalchemy"]
+ assert db.engine is settings.engine
diff --git a/providers/fab/tests/unit/fab/www/test_db.py
b/providers/fab/tests/unit/fab/www/test_db.py
new file mode 100644
index 00000000000..f28da79e89c
--- /dev/null
+++ b/providers/fab/tests/unit/fab/www/test_db.py
@@ -0,0 +1,75 @@
+# 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
+
+from unittest import mock
+
+import pytest
+from flask import Flask
+from sqlalchemy import create_engine
+
+from airflow.providers.fab.www.db import AirflowSQLAlchemy
+
+
[email protected]
+def flask_app():
+ app = Flask(__name__)
+ app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite://"
+ app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
+ return app
+
+
[email protected]
+def metadata_engine():
+ engine = create_engine("sqlite://")
+ with (
+ mock.patch("airflow.settings.engine", engine),
+ mock.patch("airflow.settings.SQL_ALCHEMY_CONN", "sqlite://"),
+ ):
+ yield engine
+
+
+class TestAirflowSQLAlchemy:
+ def test_binds_to_the_airflow_metadata_engine(self, flask_app,
metadata_engine):
+ db = AirflowSQLAlchemy(flask_app)
+
+ with flask_app.app_context():
+ assert db.engine is metadata_engine
+
+ def test_builds_its_own_engine_for_secondary_binds(self, flask_app,
metadata_engine):
+ flask_app.config["SQLALCHEMY_BINDS"] = {"other": "sqlite://"}
+ db = AirflowSQLAlchemy(flask_app)
+
+ with flask_app.app_context():
+ assert db.engines[None] is metadata_engine
+ assert db.engines["other"] is not metadata_engine
+
+ def test_builds_its_own_engine_for_another_database(self, flask_app,
metadata_engine):
+ """``webserver_config.py`` may aim ``SQLALCHEMY_DATABASE_URI`` at a
database of its own."""
+ flask_app.config["SQLALCHEMY_DATABASE_URI"] =
"sqlite:////var/lib/airflow/fab.db"
+ db = AirflowSQLAlchemy(flask_app)
+
+ with flask_app.app_context():
+ assert db.engine is not metadata_engine
+
+ @mock.patch("airflow.settings.engine", None)
+ def test_builds_its_own_engine_when_the_orm_is_not_configured(self,
flask_app):
+ db = AirflowSQLAlchemy(flask_app)
+
+ with flask_app.app_context():
+ assert db.engine is not None
+ assert str(db.engine.url) == "sqlite://"
