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 0ade0fba989 providers-fab: close sessions after users and roles
collection reads (#72578)
0ade0fba989 is described below
commit 0ade0fba989cc39486024024feea785ae6ad389f
Author: MarthalaJagruthiReddy
<[email protected]>
AuthorDate: Wed Sep 9 07:43:36 2026 -0700
providers-fab: close sessions after users and roles collection reads
(#72578)
Use explicit non-scoped sessions for FAB users and roles collection queries
so sync FastAPI worker threads close their own PostgreSQL transactions.
---
.../fab/auth_manager/api_fastapi/services/roles.py | 21 +++++-----
.../fab/auth_manager/api_fastapi/services/users.py | 45 +++++++++++-----------
.../api_fastapi/services/test_roles.py | 22 ++++++-----
.../api_fastapi/services/test_users.py | 17 ++++----
4 files changed, 54 insertions(+), 51 deletions(-)
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/roles.py
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/roles.py
index 8d64a650da6..0cb782a7d0f 100644
---
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/roles.py
+++
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/roles.py
@@ -34,6 +34,7 @@ from
airflow.providers.fab.auth_manager.api_fastapi.datamodels.roles import (
from airflow.providers.fab.auth_manager.api_fastapi.sorting import
build_ordering
from airflow.providers.fab.auth_manager.models import Permission, Role
from airflow.providers.fab.www.utils import get_fab_auth_manager
+from airflow.utils.session import create_session
if TYPE_CHECKING:
from airflow.providers.fab.auth_manager.security_manager.override import
FabAirflowSecurityManagerOverride
@@ -87,20 +88,18 @@ class FABAuthManagerRoles:
@classmethod
def get_roles(cls, *, order_by: str, limit: int, offset: int) ->
RoleCollectionResponse:
- security_manager = get_fab_auth_manager().security_manager
- session = security_manager.session
-
- total_entries = session.scalars(select(func.count(Role.id))).one()
+ with create_session(scoped=False) as session:
+ total_entries = session.scalars(select(func.count(Role.id))).one()
- ordering = build_ordering(order_by, allowed={"name": Role.name,
"role_id": Role.id})
+ ordering = build_ordering(order_by, allowed={"name": Role.name,
"role_id": Role.id})
- stmt = select(Role).order_by(ordering).offset(offset).limit(limit)
- roles = session.scalars(stmt).unique().all()
+ stmt = select(Role).order_by(ordering).offset(offset).limit(limit)
+ roles = session.scalars(stmt).unique().all()
- return RoleCollectionResponse(
- roles=[RoleResponse.model_validate(r) for r in roles],
- total_entries=total_entries,
- )
+ return RoleCollectionResponse(
+ roles=[RoleResponse.model_validate(r) for r in roles],
+ total_entries=total_entries,
+ )
@classmethod
def delete_role(cls, name: str) -> None:
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/users.py
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/users.py
index e63577ba453..8958219322d 100644
---
a/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/users.py
+++
b/providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/services/users.py
@@ -31,6 +31,7 @@ from airflow.providers.fab.auth_manager.api_fastapi.sorting
import build_orderin
from airflow.providers.fab.auth_manager.models import User
from airflow.providers.fab.auth_manager.security_manager.override import
FabAirflowSecurityManagerOverride
from airflow.providers.fab.www.utils import get_fab_auth_manager
+from airflow.utils.session import create_session
class FABAuthManagerUsers:
@@ -66,31 +67,29 @@ class FABAuthManagerUsers:
@classmethod
def get_users(cls, *, order_by: str, limit: int, offset: int) ->
UserCollectionResponse:
"""Get users with pagination and ordering."""
- security_manager = get_fab_auth_manager().security_manager
- session = security_manager.session
-
- total_entries = session.scalars(select(func.count(User.id))).one()
-
- ordering = build_ordering(
- order_by,
- allowed={
- "id": User.id,
- "user_id": User.id,
- "first_name": User.first_name,
- "last_name": User.last_name,
- "username": User.username,
- "email": User.email,
- "active": User.active,
- },
- )
+ with create_session(scoped=False) as session:
+ total_entries = session.scalars(select(func.count(User.id))).one()
+
+ ordering = build_ordering(
+ order_by,
+ allowed={
+ "id": User.id,
+ "user_id": User.id,
+ "first_name": User.first_name,
+ "last_name": User.last_name,
+ "username": User.username,
+ "email": User.email,
+ "active": User.active,
+ },
+ )
- stmt = select(User).order_by(ordering).offset(offset).limit(limit)
- users = session.scalars(stmt).unique().all()
+ stmt = select(User).order_by(ordering).offset(offset).limit(limit)
+ users = session.scalars(stmt).unique().all()
- return UserCollectionResponse(
- users=[UserResponse.model_validate(u) for u in users],
- total_entries=total_entries,
- )
+ return UserCollectionResponse(
+ users=[UserResponse.model_validate(u) for u in users],
+ total_entries=total_entries,
+ )
@classmethod
def create_user(cls, body: UserBody) -> UserResponse:
diff --git
a/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_roles.py
b/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_roles.py
index 9bdd0235417..7942570618d 100644
---
a/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_roles.py
+++
b/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_roles.py
@@ -186,7 +186,8 @@ class TestRolesService:
# GET /roles
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.roles.build_ordering")
- def test_get_roles_happy_path(self, build_ordering, get_fab_auth_manager):
+
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.roles.create_session")
+ def test_get_roles_happy_path(self, create_session, build_ordering,
get_fab_auth_manager):
role1 = _make_role_obj("viewer", [("can_read", "DAG")])
role2 = _make_role_obj("admin", [("can_read", "DAG")])
fake_roles_result = _FakeScalarRoles([role1, role2])
@@ -196,10 +197,7 @@ class TestRolesService:
_FakeScalarCount(2),
fake_roles_result,
]
-
- fab_auth_manager = MagicMock()
- fab_auth_manager.security_manager = MagicMock(session=session)
- get_fab_auth_manager.return_value = fab_auth_manager
+ create_session.return_value.__enter__.return_value = session
build_ordering.return_value = column("name").desc()
@@ -215,19 +213,23 @@ class TestRolesService:
assert set(kwargs["allowed"].keys()) == {"name", "role_id"}
assert session.scalars.call_count == 2
+ create_session.assert_called_once_with(scoped=False)
+ create_session.return_value.__exit__.assert_called_once_with(None,
None, None)
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.roles.build_ordering")
- def test_get_roles_invalid_order_by_bubbles_400(self, build_ordering,
get_fab_auth_manager):
- session = MagicMock()
- fab_auth_manager = MagicMock()
- fab_auth_manager.security_manager = MagicMock(session=session)
- get_fab_auth_manager.return_value = fab_auth_manager
+
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.roles.create_session")
+ def test_get_roles_invalid_order_by_bubbles_400(
+ self, create_session, build_ordering, get_fab_auth_manager
+ ):
+ create_session.return_value.__enter__.return_value = MagicMock()
build_ordering.side_effect = HTTPException(status_code=400,
detail="disallowed")
with pytest.raises(HTTPException) as ex:
FABAuthManagerRoles.get_roles(order_by="nope", limit=10, offset=0)
assert ex.value.status_code == 400
+ create_session.assert_called_once_with(scoped=False)
+ create_session.return_value.__exit__.assert_called_once()
# DELETE /roles/{name}
diff --git
a/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_users.py
b/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_users.py
index 15511a3c2b8..84076a730e4 100644
---
a/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_users.py
+++
b/providers/fab/tests/unit/fab/auth_manager/api_fastapi/services/test_users.py
@@ -241,8 +241,9 @@ class TestUsersService:
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.users.build_ordering")
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.users.select")
+
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.users.create_session")
def test_get_users_success(
- self, mock_select, mock_build_ordering, get_fab_auth_manager,
fab_auth_manager, security_manager
+ self, mock_create_session, mock_select, mock_build_ordering,
get_fab_auth_manager
):
user1 = _make_user_obj(
username="alice", email="[email protected]", first_name="Alice",
last_name="Liddell"
@@ -252,9 +253,7 @@ class TestUsersService:
mock_session = MagicMock()
mock_session.scalars.return_value.one.return_value = 2
mock_session.scalars.return_value.unique.return_value.all.return_value
= [user1, user2]
- security_manager.session = mock_session
- fab_auth_manager.security_manager = security_manager
- get_fab_auth_manager.return_value = fab_auth_manager
+ mock_create_session.return_value.__enter__.return_value = mock_session
mock_build_ordering.return_value = "ordering"
@@ -264,11 +263,13 @@ class TestUsersService:
assert len(out.users) == 2
assert out.users[0].username == "alice"
assert out.users[1].username == "bob"
+ mock_create_session.assert_called_once_with(scoped=False)
+
mock_create_session.return_value.__exit__.assert_called_once_with(None, None,
None)
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.users.build_ordering")
- def test_get_users_invalid_order_by(
- self, mock_build_ordering, get_fab_auth_manager, fab_auth_manager,
security_manager
- ):
+
@patch("airflow.providers.fab.auth_manager.api_fastapi.services.users.create_session")
+ def test_get_users_invalid_order_by(self, mock_create_session,
mock_build_ordering, get_fab_auth_manager):
+ mock_create_session.return_value.__enter__.return_value = MagicMock()
mock_build_ordering.side_effect = HTTPException(
status_code=400,
detail="Ordering with 'invalid' is disallowed or the attribute
does not exist on the model",
@@ -279,6 +280,8 @@ class TestUsersService:
with pytest.raises(HTTPException) as ex:
FABAuthManagerUsers.get_users(order_by="invalid", limit=10,
offset=0)
assert ex.value.status_code == 400
+ mock_create_session.assert_called_once_with(scoped=False)
+ mock_create_session.return_value.__exit__.assert_called_once()
def test_update_user_success(self, get_fab_auth_manager, fab_auth_manager,
security_manager):
user_obj = _make_user_obj(
