From eb2be7ff816d2ff3b0e574a9e98b6f38b608fcdd Mon Sep 17 00:00:00 2001 From: Paul Y Hui Date: Fri, 6 Nov 2026 13:21:20 +0000 Subject: [PATCH] refactor(api): drop global db.session from RBAC init task (#45736) Co-authored-by: Claude Opus 5.5 --- ...initialize_created_app_rbac_access_task.py | 22 +++-- ...initialize_created_app_rbac_access_task.py | 89 +++++++++++++++++++ 2 files changed, 103 insertions(+), 8 deletions(-) diff --git a/api/tasks/initialize_created_app_rbac_access_task.py b/api/tasks/initialize_created_app_rbac_access_task.py index caba953503e..f1c984a8ccd 100644 --- a/api/tasks/initialize_created_app_rbac_access_task.py +++ b/api/tasks/initialize_created_app_rbac_access_task.py @@ -8,7 +8,7 @@ from celery import shared_task from sqlalchemy import select from configs import dify_config -from extensions.ext_database import db +from core.db.session_factory import session_factory from models import Agent, App, Dataset, TenantAccountJoin, TenantAccountRole from services.enterprise import rbac_service as enterprise_rbac_service from services.rbac import contracts as rbac_contracts @@ -23,12 +23,13 @@ APP_RBAC_MEMBER_APPEND_BATCH_SIZE = 500 def _owner_account_id(tenant_id: str) -> str | None: - return db.session().scalar( - select(TenantAccountJoin.account_id) - .where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.role == TenantAccountRole.OWNER) - .order_by(TenantAccountJoin.id.asc()) - .limit(1) - ) + with session_factory.create_session() as session: + return session.scalar( + select(TenantAccountJoin.account_id) + .where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.role == TenantAccountRole.OWNER) + .order_by(TenantAccountJoin.id.asc()) + .limit(1) + ) @dataclass(frozen=True) @@ -39,6 +40,10 @@ class _WhitelistResourceKind[ItemT]: append_members: Callable[[str, str | None, Sequence[ItemT]], None] def iter_id_batches(self, tenant_id: str, batch_size: int) -> Iterator[list[str]]: + """Yield resource ids in ascending batches, closing each batch's session before yielding. + + Callers call the RBAC service between batches, so no database connection is held across a yield. + """ last_id: str | None = None while True: stmt = ( @@ -49,7 +54,8 @@ class _WhitelistResourceKind[ItemT]: ) if last_id: stmt = stmt.where(self.model.id > last_id) - ids = list(db.session().scalars(stmt).all()) + with session_factory.create_session() as session: + ids = list(session.scalars(stmt).all()) if not ids: return yield ids diff --git a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py index b7ac59dae43..2a67ac345b9 100644 --- a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py +++ b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py @@ -5,6 +5,8 @@ from typing import Protocol, cast import pytest from flask import Flask +from sqlalchemy import Engine +from sqlalchemy.pool import QueuePool import tasks.initialize_created_app_rbac_access_task as task_module from enums import DeploymentEdition @@ -12,6 +14,7 @@ from enums.account import TenantAccountRole from extensions.ext_application_services import build_application_services from extensions.ext_redis import redis_client from models import TenantAccountJoin +from models.model import App, AppMode from services.rbac import contracts as dto from tasks.initialize_created_app_rbac_access_task import ( initialize_created_app_rbac_access_task, @@ -159,3 +162,89 @@ def test_sync_joined_workspace_member_rbac_access_task_appends_auto_included_res } for request in (lookup, *appends): assert (request.tenant_id, request.account_id) == ("tenant-1", "actor-1") + + +def _seed_owner_lookup_memberships(domain: RBACDomain) -> None: + with domain.sessions.begin() as session: + for membership_id, tenant_id, account_id, role in ( + ("0", "other", "foreign-owner", TenantAccountRole.OWNER), + ("1", "tenant-1", "member-0", TenantAccountRole.NORMAL), + ("2", "tenant-1", "owner-1", TenantAccountRole.OWNER), + ("3", "tenant-1", "owner-2", TenantAccountRole.OWNER), + ): + membership = TenantAccountJoin(tenant_id=tenant_id, account_id=account_id, role=role) + membership.id = membership_id + session.add(membership) + + +def test_sync_joined_workspace_member_rbac_access_task_acts_as_the_workspace_owner( + monkeypatch: pytest.MonkeyPatch, rbac_domain: RBACDomain +) -> None: + _seed_owner_lookup_memberships(rbac_domain) + resource = dto.ResourceWhitelistConfigResource(resource_type=dto.RBACResourceType.APP, resource_id="app-1") + + def resource_batches(tenant_id: str, batch_size: int) -> Iterator[list[dto.ResourceWhitelistConfigResource]]: + assert tenant_id == "tenant-1" + assert batch_size == task_module.APP_RBAC_RESOURCE_CONFIG_BATCH_SIZE + yield [resource] + + configs = dto.ResourceWhitelistConfigsResponse( + data=[ + dto.ResourceWhitelistConfigItem( + resource_type=resource.resource_type, + resource_id=resource.resource_id, + automatic_include_workspace_members=True, + ) + ] + ) + rbac_domain.transport.responses["/rbac/whitelist/configs"] = configs.model_dump(mode="json") + monkeypatch.setattr(task_module, "_iter_resource_config_batches", resource_batches) + + sync_joined_workspace_member_rbac_access_task.run("tenant-1", "member-1") + + lookup, append = rbac_domain.transport.requests + assert (lookup.endpoint, append.endpoint) == ("/rbac/whitelist/configs", "/rbac/apps/whitelist/members/batch") + for request in (lookup, append): + assert (request.tenant_id, request.account_id) == ("tenant-1", "owner-1") + + +def test_sync_joined_workspace_member_rbac_access_task_skips_a_workspace_without_owner( + rbac_domain: RBACDomain, +) -> None: + _seed_members(rbac_domain, ("acct-1",)) + + sync_joined_workspace_member_rbac_access_task.run("tenant-1", "member-1") + + assert rbac_domain.transport.requests == [] + + +def test_resource_id_batches_are_tenant_scoped_and_hold_no_connection_between_batches( + rbac_domain: RBACDomain, sqlite_engine: Engine +) -> None: + with rbac_domain.sessions.begin() as session: + for app_id, tenant_id in ( + ("app-3", "tenant-1"), + ("app-0", "other"), + ("app-1", "tenant-1"), + ("app-2", "tenant-1"), + ): + app = App( + tenant_id=tenant_id, + name=app_id, + mode=AppMode.CHAT, + created_by="actor-1", + enable_api=True, + enable_site=True, + ) + app.id = app_id + session.add(app) + pool = sqlite_engine.pool + assert isinstance(pool, QueuePool) + app_kind = task_module._WHITELIST_RESOURCE_KIND_BY_TYPE[dto.RBACResourceType.APP] + + batches: list[list[str]] = [] + for batch in app_kind.iter_id_batches("tenant-1", 2): + assert pool.checkedout() == 0 + batches.append(batch) + + assert batches == [["app-1", "app-2"], ["app-3"]]