chore: remove unnessary rbac call (#45857)

api/services/enterprise/base.py

@@ -1,7 +1,10 @@
+import atexit
 import logging
 import os
+import threading
 from collections.abc import Mapping
-from typing import Any
+from http.cookiejar import CookieJar, DefaultCookiePolicy
+from typing import Any, ClassVar
 
 import httpx
 
@@ -159,6 +162,34 @@ class EnterpriseRequest(BaseRequest):
     secret_key = os.environ.get("ENTERPRISE_API_SECRET_KEY", "ENTERPRISE_API_SECRET_KEY")
     secret_key_header = "Enterprise-Api-Secret-Key"
 
+    _rbac_client: ClassVar[httpx.Client | None] = None
+    _rbac_client_lock: ClassVar[threading.Lock] = threading.Lock()
+
+    @classmethod
+    def _close_rbac_clients(cls) -> None:
+        """Close this process's RBAC connections during shutdown."""
+        with cls._rbac_client_lock:
+            if cls._rbac_client is not None:
+                cls._rbac_client.close()
+                cls._rbac_client = None
+
+    @classmethod
+    def _get_rbac_client(cls) -> httpx.Client:
+        """Return a process-local client for fixed endpoint/proxy settings.
+
+        Initialize only in workers after process startup, never before forking.
+        Connections are shared; credentials and tenant headers remain request-local.
+        Response cookies are rejected so one account cannot affect another call.
+        The pool owns shutdown; callers must not close the returned client.
+        """
+        with cls._rbac_client_lock:
+            if cls._rbac_client is None:
+                cls._rbac_client = httpx.Client(
+                    mounts=cls._build_mounts(),
+                    cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=())),
+                )
+            return cls._rbac_client
+
     @classmethod
     def send_inner_rbac_request(
         cls,
@@ -194,7 +225,6 @@ class EnterpriseRequest(BaseRequest):
             raise ValueError("ENTERPRISE_RBAC_API_URL is required when RBAC_ENABLED=true")
 
         url = f"{cls.rbac_base_url}{endpoint}"
-        mounts = cls._build_mounts()
 
         try:
             traceparent = generate_traceparent_header()
@@ -204,19 +234,22 @@ class EnterpriseRequest(BaseRequest):
         except Exception:
             logger.debug("Failed to generate traceparent header", exc_info=True)
 
-        with httpx.Client(mounts=mounts) as client:
-            request_kwargs: dict[str, Any] = {
-                "json": json,
-                "params": params,
-                "headers": {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key, **inner_headers},
-                "timeout": timeout if timeout is not None else dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT,
-            }
-            response = client.request(method, url, **request_kwargs)
-            if not response.is_success:
-                cls._handle_error_response(response)
+        client = cls._get_rbac_client()
+        request_kwargs: dict[str, Any] = {
+            "json": json,
+            "params": params,
+            "headers": {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key, **inner_headers},
+            "timeout": timeout if timeout is not None else dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT,
+        }
+        response = client.request(method, url, **request_kwargs)
+        if not response.is_success:
+            cls._handle_error_response(response)
         return response.json()
 
 
+atexit.register(EnterpriseRequest._close_rbac_clients)
+
+
 class EnterprisePluginManagerRequest(BaseRequest):
     base_url = os.environ.get("ENTERPRISE_PLUGIN_MANAGER_API_URL", "ENTERPRISE_PLUGIN_MANAGER_API_URL")
     secret_key = os.environ.get("ENTERPRISE_PLUGIN_MANAGER_API_SECRET_KEY", "ENTERPRISE_PLUGIN_MANAGER_API_SECRET_KEY")

api/services/knowledge/dataset_service.py

@@ -133,15 +133,6 @@ class AutoDisableLogsDict(TypedDict):
 
 
 class DatasetService:
-    @staticmethod
-    def _can_manage_all_datasets(tenant_id: str, account_id: str) -> bool:
-        if not dify_config.RBAC_ENABLED:
-            return False
-
-        permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id)
-        workspace_permission_keys = getattr(getattr(permissions, "workspace", None), "permission_keys", []) or []
-        return "dataset.create_and_management" in workspace_permission_keys
-
     @staticmethod
     def get_datasets(
         page: int,
@@ -164,6 +155,9 @@ class DatasetService:
 
         Service API callers skip account ACLs while retaining the tenant boundary.
         A resource-scoped grant further restricts the query to its authorized IDs.
+        Account callers must supply the resolved RBAC resource scope when restricted.
+        This method preserves that scope and pagination without fetching another
+        permission snapshot; include_all does not bypass the RBAC resource filter.
         """
         query = select(Dataset).where(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc(), Dataset.id)
 
@@ -193,46 +187,34 @@ class DatasetService:
                     query = query.where(Dataset.id.in_(permitted_dataset_ids))
                 else:
                     return [], 0
-            else:
-                if dify_config.RBAC_ENABLED:
-                    can_manage_all_datasets = include_all and DatasetService._can_manage_all_datasets(
-                        str(tenant_id), str(user.id)
-                    )
-                    should_show_all_datasets = include_all and can_manage_all_datasets
-                else:
-                    should_show_all_datasets = user.current_role == TenantAccountRole.OWNER and include_all
-
+            elif not dify_config.RBAC_ENABLED:
+                should_show_all_datasets = user.current_role == TenantAccountRole.OWNER and include_all
                 if not should_show_all_datasets:
-                    if dify_config.RBAC_ENABLED:
-                        # RBAC mode: show all datasets.  Permission control is enforced
-                        # via permission_keys on each item and @rbac_permission_required decorators.
-                        pass
+                    # Keep legacy visibility rules when RBAC is disabled.
+                    if permitted_dataset_ids and len(permitted_dataset_ids) > 0:
+                        query = query.where(
+                            sa.or_(
+                                Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
+                                sa.and_(
+                                    Dataset.permission == DatasetPermissionEnum.ONLY_ME,
+                                    Dataset.maintainer == user.id,
+                                ),
+                                sa.and_(
+                                    Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM,
+                                    Dataset.id.in_(permitted_dataset_ids),
+                                ),
+                            )
+                        )
                     else:
-                        # Keep legacy visibility rules when RBAC is disabled.
-                        if permitted_dataset_ids and len(permitted_dataset_ids) > 0:
-                            query = query.where(
-                                sa.or_(
-                                    Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
-                                    sa.and_(
-                                        Dataset.permission == DatasetPermissionEnum.ONLY_ME,
-                                        Dataset.maintainer == user.id,
-                                    ),
-                                    sa.and_(
-                                        Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM,
-                                        Dataset.id.in_(permitted_dataset_ids),
-                                    ),
-                                )
-                            )
-                        else:
-                            query = query.where(
-                                sa.or_(
-                                    Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
-                                    sa.and_(
-                                        Dataset.permission == DatasetPermissionEnum.ONLY_ME,
-                                        Dataset.maintainer == user.id,
-                                    ),
-                                )
+                        query = query.where(
+                            sa.or_(
+                                Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
+                                sa.and_(
+                                    Dataset.permission == DatasetPermissionEnum.ONLY_ME,
+                                    Dataset.maintainer == user.id,
+                                ),
                             )
+                        )
         elif service_api_grant is None:
             if dify_config.RBAC_ENABLED:
                 # Without an account we cannot resolve RBAC resource visibility.

api/tests/unit_tests/services/enterprise/test_base.py

@@ -5,8 +5,10 @@ generated and included in HTTP requests made by EnterpriseRequest.
 """
 
 from collections.abc import Iterator
+from http.cookiejar import CookieJar
 from unittest.mock import MagicMock, patch
 
+import httpx
 import pytest
 
 from services.enterprise.base import EnterpriseRequest
@@ -64,3 +66,88 @@ class TestTraceparentPropagation:
             assert headers["traceparent"] == expected_traceparent
             assert headers["Content-Type"] == "application/json"
             assert headers["Enterprise-Api-Secret-Key"] == "test-secret-key"
+
+
+type PooledTransport = tuple[list[httpx.Request], list[httpx.Client]]
+
+
+class TestRBACConnectionPooling:
+    @pytest.fixture
+    def pooled_transport(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[PooledTransport]:
+        from services.enterprise import base
+
+        EnterpriseRequest._close_rbac_clients()
+        requests: list[httpx.Request] = []
+
+        def handle(request: httpx.Request) -> httpx.Response:
+            requests.append(request)
+            return httpx.Response(200, json={"ok": True}, headers={"set-cookie": "session=first; Path=/"})
+
+        real_client = httpx.Client
+        clients: list[httpx.Client] = []
+
+        def build(*, mounts: dict[str, httpx.BaseTransport] | None, cookies: CookieJar) -> httpx.Client:
+            client = real_client(transport=httpx.MockTransport(handle), mounts=mounts, cookies=cookies)
+            clients.append(client)
+            return client
+
+        monkeypatch.setattr(base.httpx, "Client", build)
+        monkeypatch.setattr(EnterpriseRequest, "rbac_base_url", "https://rbac.example.com")
+        monkeypatch.setattr(base, "generate_traceparent_header", lambda: None)
+        yield requests, clients
+        EnterpriseRequest._close_rbac_clients()
+        for client in clients:
+            client.close()
+
+    def test_reuses_client_without_sharing_identity_or_cookies(self, pooled_transport: PooledTransport) -> None:
+        requests, clients = pooled_transport
+        assert EnterpriseRequest.send_inner_rbac_request(
+            "GET", "/permissions", tenant_id="tenant-a", account_id="account-a", timeout=3
+        ) == {"ok": True}
+        EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-b", timeout=7)
+
+        assert len(clients) == 1
+        assert not clients[0].is_closed
+        assert requests[0].headers["X-Inner-Tenant-Id"] == "tenant-a"
+        assert requests[0].headers["X-Inner-Account-Id"] == "account-a"
+        assert requests[1].headers["X-Inner-Tenant-Id"] == "tenant-b"
+        assert "X-Inner-Account-Id" not in requests[1].headers
+        assert "cookie" not in requests[1].headers
+        assert requests[0].extensions["timeout"]["read"] == 3
+        assert requests[1].extensions["timeout"]["read"] == 7
+
+    def test_close_releases_shared_client(self, pooled_transport: PooledTransport) -> None:
+        client = EnterpriseRequest._get_rbac_client()
+        EnterpriseRequest._close_rbac_clients()
+        assert client.is_closed
+        assert EnterpriseRequest._rbac_client is None
+
+    def test_concurrent_calls_share_one_client(self, pooled_transport: PooledTransport) -> None:
+        from concurrent.futures import ThreadPoolExecutor
+
+        requests, clients = pooled_transport
+
+        def send(index: int) -> object:
+            return EnterpriseRequest.send_inner_rbac_request(
+                "GET", "/permissions", tenant_id=f"tenant-{index}", account_id=f"account-{index}"
+            )
+
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            assert list(executor.map(send, range(12))) == [{"ok": True}] * 12
+        assert len(clients) == 1
+        assert len(requests) == 12
+        for request in requests:
+            assert request.headers["X-Inner-Tenant-Id"].removeprefix("tenant-") == request.headers[
+                "X-Inner-Account-Id"
+            ].removeprefix("account-")
+            assert "cookie" not in request.headers
+
+    def test_error_does_not_close_shared_client(self, pooled_transport: PooledTransport) -> None:
+        from services.errors.enterprise import EnterpriseAPIForbiddenError
+
+        client = EnterpriseRequest._get_rbac_client()
+        with patch.object(client, "request", return_value=httpx.Response(403, json={"message": "denied"})):
+            with pytest.raises(EnterpriseAPIForbiddenError):
+                EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-a")
+        assert not client.is_closed
+        assert EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-a") == {"ok": True}

api/tests/unit_tests/services/knowledge/test_dataset_service.py

@@ -649,8 +649,9 @@ class TestDatasetServiceDatasetCases:
             assert total == 2
             assert {dataset.id for dataset in datasets} == {creator_one.id, creator_two.id}
 
+        @pytest.mark.parametrize("include_all", [False, True])
         def test_get_datasets_applies_rbac_resource_scope_and_maintainer_override(
-            self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session
+            self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session, include_all: bool
         ) -> None:
             user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL)
             accessible = _dataset_service_dataset_dataset(
@@ -667,7 +668,7 @@ class TestDatasetServiceDatasetCases:
             with (
                 _dataset_service_dataset_patch(
                     "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get",
-                    return_value=SimpleNamespace(workspace=SimpleNamespace(permission_keys=[])),
+                    side_effect=AssertionError("List queries must reuse the caller-resolved RBAC scope"),
                 ),
             ):
                 datasets, total = _dataset_service_dataset_DatasetService.get_datasets(
@@ -677,6 +678,7 @@ class TestDatasetServiceDatasetCases:
                     tenant_id="tenant-1",
                     user=user,
                     accessible_dataset_ids=[accessible.id],
+                    include_all=include_all,
                     tags=application_tags,
                 )