(airflow) 12/12: Fix bug that makes `AirflowSecurityManagerV2` leave transactions in the `idle in transaction` state (#39935)

[ut...@apache.org]

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

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

commit 4fedbb1e5963a4655241cf8335b118025cd7a730
Author: Alejandro Rivera <al...@msn.com>
AuthorDate: Mon Jun 3 14:44:54 2024 -0600

    Fix bug that makes `AirflowSecurityManagerV2` leave transactions in the `idle in transaction` state (#39935)
    
    (cherry picked from commit b5bb039811784bbf52b7be55ff0df8d2b8b48509)
---
 airflow/www/security_manager.py    | 10 ++++------
 tests/www/test_security_manager.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/airflow/www/security_manager.py b/airflow/www/security_manager.py
index 316daace7d..cce3882132 100644
--- a/airflow/www/security_manager.py
+++ b/airflow/www/security_manager.py
@@ -64,7 +64,6 @@ from airflow.security.permissions import (
     RESOURCE_XCOM,
 )
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.www.extensions.init_auth_manager import get_auth_manager
 from airflow.www.utils import CustomSQLAInterface
 
@@ -77,8 +76,6 @@ EXISTING_ROLES = {
 }
 
 if TYPE_CHECKING:
-    from sqlalchemy.orm import Session
-
     from airflow.auth.managers.models.base_user import BaseUser
 
 
@@ -167,9 +164,8 @@ class AirflowSecurityManagerV2(LoggingMixin):
             )(baseview.blueprint)
 
     @cached_property
-    @provide_session
     def _auth_manager_is_authorized_map(
-        self, session: Session = NEW_SESSION
+        self,
     ) -> dict[str, Callable[[str, str | None, BaseUser | None], bool]]:
         """
         Return the map associating a FAB resource name to the corresponding auth manager is_authorized_ API.
@@ -179,6 +175,8 @@ class AirflowSecurityManagerV2(LoggingMixin):
         auth_manager = get_auth_manager()
         methods = get_method_from_fab_action_map()
 
+        session = self.appbuilder.session
+
         def get_connection_id(resource_pk):
             if not resource_pk:
                 return None
@@ -226,7 +224,7 @@ class AirflowSecurityManagerV2(LoggingMixin):
                 return None
             variable = session.scalar(select(Variable).where(Variable.id == resource_pk).limit(1))
             if not variable:
-                raise AirflowException("Connection not found")
+                raise AirflowException("Variable not found")
             return variable.key
 
         return {
diff --git a/tests/www/test_security_manager.py b/tests/www/test_security_manager.py
index 2051ccef09..ff66864188 100644
--- a/tests/www/test_security_manager.py
+++ b/tests/www/test_security_manager.py
@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import json
 from unittest import mock
 from unittest.mock import Mock
 
@@ -135,3 +136,33 @@ class TestAirflowSecurityManagerV2:
         if len(auth_manager_methods) > 1 and not expected:
             for method_name in auth_manager_methods:
                 getattr(auth_manager, method_name).assert_called()
+
+    @mock.patch("airflow.utils.session.create_session")
+    @mock.patch("airflow.www.security_manager.get_auth_manager")
+    def test_manager_does_not_create_extra_db_sessions(
+        self,
+        _,
+        mock_create_session,
+        security_manager,
+    ):
+        """
+        Test that the Security Manager doesn't create extra DB sessions and
+        instead uses the session already available through the appbuilder
+        object that is attached to it.
+        """
+        with mock.patch.object(security_manager.appbuilder, "session") as mock_appbuilder_session:
+            action_name = ACTION_CAN_READ
+            resource_pk = "PK"
+            user = Mock()
+            for func in security_manager._auth_manager_is_authorized_map.values():
+                try:
+                    func(action_name, resource_pk, user)
+                except json.JSONDecodeError:
+                    # The resource-retrieving function expects a "composite"
+                    # PK as a JSON string. Provide a mocked one.
+                    func(action_name, "[1, 1, 1, 1]", user)
+                mock_create_session.assert_not_called()
+
+        # The Security Manager's `appbuilder.session` object should have been
+        # put to use by many of the functions tested above.
+        assert len(mock_appbuilder_session.method_calls) > 0