diff --git a/providers/src/airflow/providers/fab/auth_manager/cli_commands/utils.py b/providers/src/airflow/providers/fab/auth_manager/cli_commands/utils.py index 03bfd526d09c4..81b2ab025351b 100644 --- a/providers/src/airflow/providers/fab/auth_manager/cli_commands/utils.py +++ b/providers/src/airflow/providers/fab/auth_manager/cli_commands/utils.py @@ -27,7 +27,10 @@ import airflow from airflow.configuration import conf +from airflow.exceptions import AirflowConfigException +from airflow.www.app import isabs, make_url from airflow.www.extensions.init_appbuilder import init_appbuilder +from airflow.www.extensions.init_session import init_airflow_session_interface from airflow.www.extensions.init_views import init_plugins if TYPE_CHECKING: @@ -39,6 +42,7 @@ def _return_appbuilder(app: Flask) -> AirflowAppBuilder: """Return an appbuilder instance for the given app.""" init_appbuilder(app) init_plugins(app) + init_airflow_session_interface(app) return app.appbuilder # type: ignore[attr-defined] @@ -50,4 +54,12 @@ def get_application_builder() -> Generator[AirflowAppBuilder, None, None]: with flask_app.app_context(): # Enable customizations in webserver_config.py to be applied via Flask.current_app. flask_app.config.from_pyfile(webserver_config, silent=True) + flask_app.config["SQLALCHEMY_DATABASE_URI"] = conf.get("database", "SQL_ALCHEMY_CONN") + url = make_url(flask_app.config["SQLALCHEMY_DATABASE_URI"]) + if url.drivername == "sqlite" and url.database and not isabs(url.database): + raise AirflowConfigException( + f'Cannot use relative path: `{conf.get("database", "SQL_ALCHEMY_CONN")}` to connect to sqlite. ' + "Please use absolute path such as `sqlite:////tmp/airflow.db`." + ) + flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False yield _return_appbuilder(flask_app) diff --git a/providers/src/airflow/providers/fab/auth_manager/security_manager/override.py b/providers/src/airflow/providers/fab/auth_manager/security_manager/override.py index 82c8ab0bdcd41..292f1e1f2af87 100644 --- a/providers/src/airflow/providers/fab/auth_manager/security_manager/override.py +++ b/providers/src/airflow/providers/fab/auth_manager/security_manager/override.py @@ -572,6 +572,7 @@ def reset_user_sessions(self, user: User) -> None: session_details = interface.serializer.loads(want_bytes(s.data)) if session_details.get("_user_id") == user.id: session.delete(s) + session.commit() else: self._cli_safe_flash( "Since you are using `securecookie` session backend mechanism, we cannot prevent " diff --git a/providers/tests/fab/auth_manager/cli_commands/test_utils.py b/providers/tests/fab/auth_manager/cli_commands/test_utils.py index cacf614a82fd0..e7f25185b5c98 100644 --- a/providers/tests/fab/auth_manager/cli_commands/test_utils.py +++ b/providers/tests/fab/auth_manager/cli_commands/test_utils.py @@ -16,19 +16,69 @@ # under the License. from __future__ import annotations +import os + import pytest +import airflow +from airflow.configuration import conf +from airflow.exceptions import AirflowConfigException +from airflow.www.extensions.init_appbuilder import AirflowAppBuilder +from airflow.www.session import AirflowDatabaseSessionInterface + from tests_common.test_utils.compat import ignore_provider_compatibility_error +from tests_common.test_utils.config import conf_vars with ignore_provider_compatibility_error("2.9.0+", __file__): from airflow.providers.fab.auth_manager.cli_commands.utils import get_application_builder -from airflow.www.extensions.init_appbuilder import AirflowAppBuilder - pytestmark = pytest.mark.db_test +@pytest.fixture +def flask_app(): + """Fixture to set up the Flask app with the necessary configuration.""" + # Get the webserver config file path + webserver_config = conf.get_mandatory_value("webserver", "config_file") + + with get_application_builder() as appbuilder: + flask_app = appbuilder.app + + # Load webserver configuration + flask_app.config.from_pyfile(webserver_config, silent=True) + + yield flask_app + + class TestCliUtils: def test_get_application_builder(self): + """Test that get_application_builder returns an AirflowAppBuilder instance.""" with get_application_builder() as appbuilder: assert isinstance(appbuilder, AirflowAppBuilder) + + def test_sqlalchemy_uri_configured(self, flask_app): + """Test that the SQLALCHEMY_DATABASE_URI is correctly set in the Flask app.""" + sqlalchemy_uri = conf.get("database", "SQL_ALCHEMY_CONN") + + # Assert that the SQLAlchemy URI is correctly set + assert sqlalchemy_uri == flask_app.config["SQLALCHEMY_DATABASE_URI"] + + def test_relative_path_sqlite_raises_exception(self): + """Test that a relative SQLite path raises an AirflowConfigException.""" + # Directly simulate the configuration for relative SQLite path + with conf_vars({("database", "SQL_ALCHEMY_CONN"): "sqlite://relative/path"}): + with pytest.raises(AirflowConfigException, match="Cannot use relative path"): + with get_application_builder(): + pass + + def test_static_folder_exists(self, flask_app): + """Test that the static folder is correctly configured in the Flask app.""" + static_folder = os.path.join(os.path.dirname(airflow.__file__), "www", "static") + assert flask_app.static_folder == static_folder + + def test_database_auth_backend_in_session(self, flask_app): + """Test that the database is used for session management when AUTH_BACKEND is set to 'database'.""" + with get_application_builder() as appbuilder: + flask_app = appbuilder.app + # Ensure that the correct session interface is set (for 'database' auth backend) + assert isinstance(flask_app.session_interface, AirflowDatabaseSessionInterface)