Skip to content

Commit b5bb039

Browse files
Fix bug that makes AirflowSecurityManagerV2 leave transactions in the idle in transaction state (#39935)
1 parent 651a6d6 commit b5bb039

2 files changed

Lines changed: 35 additions & 6 deletions

File tree

‎airflow/www/security_manager.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@
6464
RESOURCE_XCOM,
6565
)
6666
from airflow.utils.log.logging_mixin import LoggingMixin
67-
from airflow.utils.session import NEW_SESSION, provide_session
6867
from airflow.www.extensions.init_auth_manager import get_auth_manager
6968
from airflow.www.utils import CustomSQLAInterface
7069

@@ -77,8 +76,6 @@
7776
}
7877

7978
if TYPE_CHECKING:
80-
from sqlalchemy.orm import Session
81-
8279
from airflow.auth.managers.models.base_user import BaseUser
8380

8481

@@ -167,9 +164,8 @@ def add_limit_view(self, baseview):
167164
)(baseview.blueprint)
168165

169166
@cached_property
170-
@provide_session
171167
def _auth_manager_is_authorized_map(
172-
self, session: Session = NEW_SESSION
168+
self,
173169
) -> dict[str, Callable[[str, str | None, BaseUser | None], bool]]:
174170
"""
175171
Return the map associating a FAB resource name to the corresponding auth manager is_authorized_ API.
@@ -179,6 +175,8 @@ def _auth_manager_is_authorized_map(
179175
auth_manager = get_auth_manager()
180176
methods = get_method_from_fab_action_map()
181177

178+
session = self.appbuilder.session
179+
182180
def get_connection_id(resource_pk):
183181
if not resource_pk:
184182
return None
@@ -226,7 +224,7 @@ def get_variable_key(resource_pk):
226224
return None
227225
variable = session.scalar(select(Variable).where(Variable.id == resource_pk).limit(1))
228226
if not variable:
229-
raise AirflowException("Connection not found")
227+
raise AirflowException("Variable not found")
230228
return variable.key
231229

232230
return {

‎tests/www/test_security_manager.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
# under the License.
1818
from __future__ import annotations
1919

20+
import json
2021
from unittest import mock
2122
from unittest.mock import Mock
2223

@@ -135,3 +136,33 @@ def test_has_access(
135136
if len(auth_manager_methods) > 1 and not expected:
136137
for method_name in auth_manager_methods:
137138
getattr(auth_manager, method_name).assert_called()
139+
140+
@mock.patch("airflow.utils.session.create_session")
141+
@mock.patch("airflow.www.security_manager.get_auth_manager")
142+
def test_manager_does_not_create_extra_db_sessions(
143+
self,
144+
_,
145+
mock_create_session,
146+
security_manager,
147+
):
148+
"""
149+
Test that the Security Manager doesn't create extra DB sessions and
150+
instead uses the session already available through the appbuilder
151+
object that is attached to it.
152+
"""
153+
with mock.patch.object(security_manager.appbuilder, "session") as mock_appbuilder_session:
154+
action_name = ACTION_CAN_READ
155+
resource_pk = "PK"
156+
user = Mock()
157+
for func in security_manager._auth_manager_is_authorized_map.values():
158+
try:
159+
func(action_name, resource_pk, user)
160+
except json.JSONDecodeError:
161+
# The resource-retrieving function expects a "composite"
162+
# PK as a JSON string. Provide a mocked one.
163+
func(action_name, "[1, 1, 1, 1]", user)
164+
mock_create_session.assert_not_called()
165+
166+
# The Security Manager's `appbuilder.session` object should have been
167+
# put to use by many of the functions tested above.
168+
assert len(mock_appbuilder_session.method_calls) > 0

0 commit comments

Comments
 (0)