refactor(api): drop global db.session from RBAC init task (#45736)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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

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"]]