(airflow) 02/02: Fix security manager inheritance in fab provider (#36538)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch v2-8-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 676e91037572502c4c73e3a56c92820feef94ed1
Author: Vincent <97...@users.noreply.github.com>
AuthorDate: Wed Jan 3 13:02:36 2024 -0500

    Fix security manager inheritance in fab provider (#36538)
    
    (cherry picked from commit 2093b6f3b94be9fae5d61042a9c280d9a835687b)
---
 airflow/auth/managers/fab/fab_auth_manager.py    | 14 +++++------
 tests/auth/managers/fab/test_fab_auth_manager.py | 30 ++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py
index 5224509039..e05abc97d3 100644
--- a/airflow/auth/managers/fab/fab_auth_manager.py
+++ b/airflow/auth/managers/fab/fab_auth_manager.py
@@ -222,10 +222,10 @@ class FabAuthManager(BaseAuthManager):
         entity (e.g. DAG runs).
         2. ``dag_access`` is provided which means the user wants to access a sub entity of the DAG
         (e.g. DAG runs).
-            a. If ``method`` is GET, then check the user has READ permissions on the DAG and the sub entity.
-            b. Else, check the user has EDIT permissions on the DAG and ``method`` on the sub entity.
 
-            However, if no specific DAG is targeted, just check the sub entity.
+            a. If ``method`` is GET, then check the user has READ permissions on the DAG and the sub entity.
+            b. Else, check the user has EDIT permissions on the DAG and ``method`` on the sub entity. However,
+                if no specific DAG is targeted, just check the sub entity.
 
         :param method: The method to authorize.
         :param access_entity: The dag access entity.
@@ -335,19 +335,19 @@ class FabAuthManager(BaseAuthManager):
     def security_manager(self) -> FabAirflowSecurityManagerOverride:
         """Return the security manager specific to FAB."""
         from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride
-        from airflow.www.security import AirflowSecurityManager
+        from airflow.www.security_manager import AirflowSecurityManagerV2
 
         sm_from_config = self.appbuilder.get_app.config.get("SECURITY_MANAGER_CLASS")
         if sm_from_config:
-            if not issubclass(sm_from_config, AirflowSecurityManager):
+            if not issubclass(sm_from_config, AirflowSecurityManagerV2):
                 raise Exception(
-                    """Your CUSTOM_SECURITY_MANAGER must extend FabAirflowSecurityManagerOverride,
+                    """Your CUSTOM_SECURITY_MANAGER must extend AirflowSecurityManagerV2,
                      not FAB's own security manager."""
                 )
             if not issubclass(sm_from_config, FabAirflowSecurityManagerOverride):
                 warnings.warn(
                     "Please make your custom security manager inherit from "
-                    "FabAirflowSecurityManagerOverride instead of AirflowSecurityManager.",
+                    "FabAirflowSecurityManagerOverride instead of FabAirflowSecurityManagerOverride.",
                     DeprecationWarning,
                 )
             return sm_from_config(self.appbuilder)
diff --git a/tests/auth/managers/fab/test_fab_auth_manager.py b/tests/auth/managers/fab/test_fab_auth_manager.py
index 74d20ab13b..361e805d71 100644
--- a/tests/auth/managers/fab/test_fab_auth_manager.py
+++ b/tests/auth/managers/fab/test_fab_auth_manager.py
@@ -65,8 +65,12 @@ def auth_manager():
 
 
 @pytest.fixture
-def auth_manager_with_appbuilder():
-    flask_app = Flask(__name__)
+def flask_app():
+    return Flask(__name__)
+
+
+@pytest.fixture
+def auth_manager_with_appbuilder(flask_app):
     appbuilder = init_appbuilder(flask_app)
     return FabAuthManager(appbuilder)
 
@@ -357,6 +361,28 @@ class TestFabAuthManager:
     def test_security_manager_return_fab_security_manager_override(self, auth_manager_with_appbuilder):
         assert isinstance(auth_manager_with_appbuilder.security_manager, FabAirflowSecurityManagerOverride)
 
+    @pytest.mark.db_test
+    def test_security_manager_return_custom_provided(self, flask_app, auth_manager_with_appbuilder):
+        class TestSecurityManager(FabAirflowSecurityManagerOverride):
+            pass
+
+        flask_app.config["SECURITY_MANAGER_CLASS"] = TestSecurityManager
+        assert isinstance(auth_manager_with_appbuilder.security_manager, TestSecurityManager)
+
+    @pytest.mark.db_test
+    def test_security_manager_wrong_inheritance_raise_exception(
+        self, flask_app, auth_manager_with_appbuilder
+    ):
+        class TestSecurityManager:
+            pass
+
+        flask_app.config["SECURITY_MANAGER_CLASS"] = TestSecurityManager
+
+        with pytest.raises(
+            Exception, match="Your CUSTOM_SECURITY_MANAGER must extend FabAirflowSecurityManagerOverride."
+        ):
+            auth_manager_with_appbuilder.security_manager
+
     @pytest.mark.db_test
     def test_get_url_login_when_auth_view_not_defined(self, auth_manager_with_appbuilder):
         with pytest.raises(AirflowException, match="`auth_view` not defined in the security manager."):