DataMate/runtime/datamate-python/app/module/kg_graphrag/kb_access.py

"""知识库访问权限校验。

在执行 GraphRAG 检索前，调用 Java rag-indexer-service 的
GET /knowledge-base/{id} 端点验证当前用户是否有权访问该知识库。

Java 侧实现参考：KnowledgeBaseService.getKnowledgeBaseWithAccessCheck()
- 查找 KB 是否存在
- 校验 createdBy == currentUserId（管理员跳过）
- 不满足则抛出 sys.0005 (INSUFFICIENT_PERMISSIONS)
"""

from __future__ import annotations

import httpx

from app.core.logging import get_logger

logger = get_logger(__name__)


class KnowledgeBaseAccessValidator:
    """通过 Java 后端校验用户是否有权访问指定知识库。"""

    def __init__(
        self,
        *,
        base_url: str = "http://datamate-backend:8080/api",
        timeout: float = 10.0,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    @classmethod
    def from_settings(cls) -> KnowledgeBaseAccessValidator:
        from app.core.config import settings

        return cls(base_url=settings.datamate_backend_base_url)

    def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._base_url,
                timeout=self._timeout,
            )
        return self._client

    async def check_access(
        self,
        knowledge_base_id: str,
        user_id: str,
        *,
        collection_name: str | None = None,
    ) -> bool:
        """校验用户是否有权访问指定知识库。

        调用 Java 后端 GET /knowledge-base/{id}，该端点内部执行
        owner 校验（createdBy == currentUserId，管理员跳过）。

        当 *collection_name* 不为 None 时，还会校验请求中的
        collection_name 与该知识库实际的 name 是否一致，防止
        用户提交合法 KB ID 但篡改 collection_name 来访问
        其他知识库的 Milvus 数据。

        Returns:
            True — 用户有权访问且 collection_name 匹配
            False — 无权访问、collection_name 不匹配或校验失败
        """
        try:
            client = self._get_client()
            resp = await client.get(
                f"/api/knowledge-base/{knowledge_base_id}",
                headers={"X-User-Id": user_id},
            )
            if resp.status_code == 200:
                body = resp.json()
                # Java 全局包装: {"code": 200, "data": {...}}
                # code != 200 说明业务层拒绝（如权限不足）
                code = body.get("code", resp.status_code)
                if code != 200:
                    logger.warning(
                        "KB access denied: kb_id=%s, user=%s, biz_code=%s, msg=%s",
                        knowledge_base_id, user_id, code, body.get("message", ""),
                    )
                    return False

                # 校验 collection_name 与 KB 实际名称的绑定关系
                if collection_name is not None:
                    data = body.get("data") or {}
                    actual_name = data.get("name") if isinstance(data, dict) else None
                    if actual_name != collection_name:
                        logger.warning(
                            "KB collection_name mismatch: kb_id=%s, "
                            "expected=%s, actual=%s, user=%s",
                            knowledge_base_id, collection_name,
                            actual_name, user_id,
                        )
                        return False

                return True
            # HTTP 4xx/5xx
            logger.warning(
                "KB access check returned HTTP %d: kb_id=%s, user=%s",
                resp.status_code, knowledge_base_id, user_id,
            )
            return False
        except Exception:
            # 网络异常时 fail-close：拒绝访问，防止绕过权限
            logger.exception(
                "KB access check failed (fail-close): kb_id=%s, user=%s",
                knowledge_base_id, user_id,
            )
            return False

    async def close(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None