refactor: explicit DB session propagation across backend paths (#38559)

Co-authored-by: WH-2099 <wh2099@pm.me>

api/controllers/openapi/workspaces.py

@@ -15,9 +15,11 @@ from itertools import starmap
 from urllib import parse
 
 from flask_restx import Resource
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, NotFound
 
 from configs import dify_config
+from controllers.common.session import with_session
 from controllers.openapi import openapi_ns
 from controllers.openapi._contract import accepts, returns
 from controllers.openapi._errors import MemberLicenseExceeded, MemberLimitExceeded
@@ -35,7 +37,6 @@ from controllers.openapi._models import (
 )
 from controllers.openapi.auth.composition import auth_router
 from controllers.openapi.auth.data import AuthData
-from extensions.ext_database import db
 from libs.oauth_bearer import Scope, TokenType
 from models import Account, Tenant, TenantAccountJoin
 from models.account import TenantAccountRole, TenantStatus
@@ -64,15 +65,15 @@ def _member_response(account: Account) -> MemberResponse:
     )
 
 
-def _load_tenant(workspace_id: str) -> Tenant:
-    tenant = TenantService.get_tenant_by_id(workspace_id, session=db.session())
+def _load_tenant(session: Session, workspace_id: str) -> Tenant:
+    tenant = TenantService.get_tenant_by_id(workspace_id, session=session)
     if tenant is None or tenant.status != TenantStatus.NORMAL:
         raise NotFound("workspace not found")
     return tenant
 
 
-def _load_account(account_id: object) -> Account:
-    account = AccountService.get_account_by_id(str(account_id), session=db.session()) if account_id else None
+def _load_account(session: Session, account_id: object) -> Account:
+    account = AccountService.get_account_by_id(str(account_id), session=session) if account_id else None
     if account is None:
         raise RuntimeError("authenticated account_id has no Account row")
     return account
@@ -94,8 +95,9 @@ def _check_member_invite_quota(tenant_id: str) -> None:
 class WorkspacesApi(Resource):
     @auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
     @returns(200, WorkspaceListResponse, description="Workspace list")
-    def get(self, *, auth_data: AuthData):
-        rows = TenantService.get_workspaces_for_account(str(auth_data.account_id), session=db.session())
+    @with_session(write=False)
+    def get(self, session: Session, *, auth_data: AuthData):
+        rows = TenantService.get_workspaces_for_account(str(auth_data.account_id), session=session)
 
         return WorkspaceListResponse(workspaces=list(starmap(_workspace_summary, rows)))
 
@@ -104,8 +106,9 @@ class WorkspacesApi(Resource):
 class WorkspaceByIdApi(Resource):
     @auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
     @returns(200, WorkspaceDetailResponse, description="Workspace detail")
-    def get(self, workspace_id: str, *, auth_data: AuthData):
-        row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session())
+    @with_session(write=False)
+    def get(self, session: Session, workspace_id: str, *, auth_data: AuthData):
+        row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=session)
         # 404 (not 403) on non-member so workspace IDs don't leak across tenants.
         if row is None:
             raise NotFound("workspace not found")
@@ -125,15 +128,16 @@ class WorkspaceSwitchApi(Resource):
 
     @auth_router.guard_workspace(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
     @returns(200, WorkspaceDetailResponse, description="Workspace detail")
-    def post(self, workspace_id: str, *, auth_data: AuthData):
-        account = _load_account(auth_data.account_id)
+    @with_session
+    def post(self, session: Session, workspace_id: str, *, auth_data: AuthData):
+        account = _load_account(session, auth_data.account_id)
 
         try:
-            TenantService.switch_tenant(account, workspace_id, session=db.session())
+            TenantService.switch_tenant(account, workspace_id, session=session)
         except AccountNotLinkTenantError:
             raise NotFound("workspace not found")
 
-        row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session())
+        row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=session)
         if row is None:
             raise NotFound("workspace not found")
         tenant, membership = row
@@ -151,9 +155,10 @@ class WorkspaceMembersApi(Resource):
     @auth_router.guard_workspace(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
     @returns(200, MemberListResponse, description="Member list")
     @accepts(query=MemberListQuery)
-    def get(self, workspace_id: str, *, auth_data: AuthData, query: MemberListQuery):
-        tenant = _load_tenant(workspace_id)
-        members = TenantService.get_tenant_members(tenant, session=db.session())
+    @with_session(write=False)
+    def get(self, session: Session, workspace_id: str, *, auth_data: AuthData, query: MemberListQuery):
+        tenant = _load_tenant(session, workspace_id)
+        members = TenantService.get_tenant_members(tenant, session=session)
         total = len(members)
         start = (query.page - 1) * query.limit
         page_items = members[start : start + query.limit]
@@ -172,9 +177,10 @@ class WorkspaceMembersApi(Resource):
     )
     @returns(201, MemberInviteResponse, description="Member invited")
     @accepts(body=MemberInvitePayload)
-    def post(self, workspace_id: str, *, auth_data: AuthData, body: MemberInvitePayload):
-        inviter = _load_account(auth_data.account_id)
-        tenant = _load_tenant(workspace_id)
+    @with_session
+    def post(self, session: Session, workspace_id: str, *, auth_data: AuthData, body: MemberInvitePayload):
+        inviter = _load_account(session, auth_data.account_id)
+        tenant = _load_tenant(session, workspace_id)
 
         _check_member_invite_quota(str(tenant.id))
 
@@ -185,7 +191,7 @@ class WorkspaceMembersApi(Resource):
                 language=None,
                 role=body.role,
                 inviter=inviter,
-                session=db.session(),
+                session=session,
             )
         except AccountAlreadyInTenantError as exc:
             raise BadRequest(str(exc))
@@ -197,7 +203,7 @@ class WorkspaceMembersApi(Resource):
             raise BadRequest(str(exc))
 
         normalized_email = body.email.lower()
-        member = AccountService.get_account_by_email_with_case_fallback(normalized_email, session=db.session())
+        member = AccountService.get_account_by_email_with_case_fallback(normalized_email, session=session)
         if member is None:
             # invite_new_member just created or fetched this account.
             raise RuntimeError("invited member missing from DB after invite")
@@ -229,15 +235,16 @@ class WorkspaceMemberApi(Resource):
         allowed_roles=frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}),
     )
     @returns(200, MemberActionResponse, description="Member removed")
-    def delete(self, workspace_id: str, member_id: str, *, auth_data: AuthData):
-        operator = _load_account(auth_data.account_id)
-        tenant = _load_tenant(workspace_id)
-        member = AccountService.get_account_by_id(member_id, session=db.session())
+    @with_session
+    def delete(self, session: Session, workspace_id: str, member_id: str, *, auth_data: AuthData):
+        operator = _load_account(session, auth_data.account_id)
+        tenant = _load_tenant(session, workspace_id)
+        member = AccountService.get_account_by_id(member_id, session=session)
         if member is None:
             raise NotFound("member not found")
 
         try:
-            TenantService.remove_member_from_tenant(tenant, member, operator, session=db.session())
+            TenantService.remove_member_from_tenant(tenant, member, operator, session=session)
         except CannotOperateSelfError as exc:
             raise BadRequest(str(exc))
         except NoPermissionError as exc:
@@ -254,15 +261,24 @@ class WorkspaceMemberApi(Resource):
     )
     @returns(200, MemberActionResponse, description="Role updated")
     @accepts(body=MemberRoleUpdatePayload)
-    def patch(self, workspace_id: str, member_id: str, *, auth_data: AuthData, body: MemberRoleUpdatePayload):
-        operator = _load_account(auth_data.account_id)
-        tenant = _load_tenant(workspace_id)
-        member = AccountService.get_account_by_id(member_id, session=db.session())
+    @with_session
+    def patch(
+        self,
+        session: Session,
+        workspace_id: str,
+        member_id: str,
+        *,
+        auth_data: AuthData,
+        body: MemberRoleUpdatePayload,
+    ):
+        operator = _load_account(session, auth_data.account_id)
+        tenant = _load_tenant(session, workspace_id)
+        member = AccountService.get_account_by_id(member_id, session=session)
         if member is None:
             raise NotFound("member not found")
 
         try:
-            TenantService.update_member_role(tenant, member, body.role, operator, session=db.session())
+            TenantService.update_member_role(tenant, member, body.role, operator, session=session)
         except CannotOperateSelfError as exc:
             raise BadRequest(str(exc))
         except NoPermissionError as exc: