(airflow) branch main updated: Fix security manager inheritance in fab provider (#36538)

[vincbeck]

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 2093b6f3b9 Fix security manager inheritance in fab provider (#36538)
2093b6f3b9 is described below

commit 2093b6f3b94be9fae5d61042a9c280d9a835687b
Author: Vincent <97131062+vincb...@users.noreply.github.com>
AuthorDate: Wed Jan 3 13:02:36 2024 -0500

    Fix security manager inheritance in fab provider (#36538)
---
 .../providers/fab/auth_manager/fab_auth_manager.py | 15 +++--------
 .../fab/auth_manager/test_fab_auth_manager.py      | 30 ++++++++++++++++++++--
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/airflow/providers/fab/auth_manager/fab_auth_manager.py 
b/airflow/providers/fab/auth_manager/fab_auth_manager.py
index b6a1c563ce..dba6aca7ef 100644
--- a/airflow/providers/fab/auth_manager/fab_auth_manager.py
+++ b/airflow/providers/fab/auth_manager/fab_auth_manager.py
@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import warnings
 from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Container
@@ -43,7 +42,7 @@ from airflow.cli.cli_config import (
     GroupCommand,
 )
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, 
AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.models import DagModel
 from airflow.providers.fab.auth_manager.cli_commands.definition import (
     ROLES_COMMANDS,
@@ -334,20 +333,12 @@ class FabAuthManager(BaseAuthManager):
         from airflow.providers.fab.auth_manager.security_manager.override 
import (
             FabAirflowSecurityManagerOverride,
         )
-        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, AirflowSecurityManagerV2):
-                raise Exception(
-                    """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.",
-                    AirflowProviderDeprecationWarning,
+                raise Exception(
+                    """Your CUSTOM_SECURITY_MANAGER must extend 
FabAirflowSecurityManagerOverride."""
                 )
             return sm_from_config(self.appbuilder)
 
diff --git a/tests/providers/fab/auth_manager/test_fab_auth_manager.py 
b/tests/providers/fab/auth_manager/test_fab_auth_manager.py
index e4c5745536..72f5a76355 100644
--- a/tests/providers/fab/auth_manager/test_fab_auth_manager.py
+++ b/tests/providers/fab/auth_manager/test_fab_auth_manager.py
@@ -63,8 +63,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)
 
@@ -355,6 +359,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."):