dify/api/services/knowledge_fs/service_api_authorization.py

"""Workspace Dataset API key authorization for KnowledgeFS Service API calls."""

from __future__ import annotations

from typing import NamedTuple

import sqlalchemy as sa
from sqlalchemy.orm import Session, sessionmaker

from models.enums import ApiTokenType
from models.knowledge_fs import (
    KnowledgeFSAuthorizationRevision,
    KnowledgeFSControlSpace,
    KnowledgeFSControlSpaceState,
    KnowledgeFSExternalAccessPolicy,
)
from models.model import ApiToken


class KnowledgeFSServiceApiAuthorizationError(RuntimeError):
    """A Dataset API key cannot access the requested KnowledgeFS control space."""


class KnowledgeFSServiceApiProfile(NamedTuple):
    tenant_id: str
    control_space_id: str
    api_token_id: str
    principal_id: str
    knowledge_space_id: str
    knowledge_space_revision: int
    membership_epoch: int
    space_acl_epoch: int
    external_access_epoch: int
    content_policy_revision: int


class KnowledgeFSServiceApiAuthorizationService:
    def __init__(self, session_maker: sessionmaker[Session]) -> None:
        self._session_maker = session_maker

    def authorize(
        self,
        *,
        api_token_id: str,
        tenant_id: str,
        control_space_id: str,
    ) -> KnowledgeFSServiceApiProfile:
        with self._session_maker() as session:
            row = session.execute(
                sa.select(
                    ApiToken,
                    KnowledgeFSControlSpace,
                    KnowledgeFSExternalAccessPolicy,
                    KnowledgeFSAuthorizationRevision,
                )
                .join(
                    KnowledgeFSControlSpace,
                    KnowledgeFSControlSpace.tenant_id == ApiToken.tenant_id,
                )
                .outerjoin(
                    KnowledgeFSExternalAccessPolicy,
                    sa.and_(
                        KnowledgeFSExternalAccessPolicy.tenant_id == KnowledgeFSControlSpace.tenant_id,
                        KnowledgeFSExternalAccessPolicy.control_space_id == KnowledgeFSControlSpace.id,
                    ),
                )
                .join(
                    KnowledgeFSAuthorizationRevision,
                    sa.and_(
                        KnowledgeFSAuthorizationRevision.tenant_id == KnowledgeFSControlSpace.tenant_id,
                        KnowledgeFSAuthorizationRevision.control_space_id == KnowledgeFSControlSpace.id,
                    ),
                )
                .where(
                    ApiToken.id == api_token_id,
                    ApiToken.tenant_id == tenant_id,
                    ApiToken.type == ApiTokenType.DATASET,
                    KnowledgeFSControlSpace.id == control_space_id,
                )
            ).one_or_none()
        if row is None:
            raise KnowledgeFSServiceApiAuthorizationError("Invalid Dataset API key or KnowledgeFS space")
        api_token, control_space, policy, revision = row._t
        if (
            control_space.state is not KnowledgeFSControlSpaceState.ACTIVE
            or control_space.knowledge_space_id is None
            or policy is None
            or not policy.service_api_enabled
        ):
            raise KnowledgeFSServiceApiAuthorizationError("KnowledgeFS Service API access is not enabled")
        return KnowledgeFSServiceApiProfile(
            tenant_id=tenant_id,
            control_space_id=control_space_id,
            api_token_id=api_token.id,
            principal_id=api_token.id,
            knowledge_space_id=control_space.knowledge_space_id,
            knowledge_space_revision=control_space.knowledge_space_revision,
            membership_epoch=revision.membership_epoch,
            space_acl_epoch=revision.space_acl_epoch,
            external_access_epoch=revision.external_access_epoch,
            content_policy_revision=revision.content_policy_revision,
        )


__all__ = [
    "KnowledgeFSServiceApiAuthorizationError",
    "KnowledgeFSServiceApiAuthorizationService",
    "KnowledgeFSServiceApiProfile",
]