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 <[email protected]> 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
