dify/api/services/knowledge_fs/control_space_management.py

"""Operator-only registration, backfill, repair, and dry-run helpers."""

from __future__ import annotations

from collections import Counter
from collections.abc import Iterable
from typing import NamedTuple

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

from models.knowledge_fs import (
    KnowledgeFSAuthorizationRevision,
    KnowledgeFSControlSpace,
    KnowledgeFSControlSpacePermission,
    KnowledgeFSControlSpacePermissionRole,
    KnowledgeFSControlSpaceState,
)
from repositories.sqlalchemy_knowledge_fs_control_space_repository import (
    SQLAlchemyKnowledgeFSControlSpaceRepository,
)
from services.knowledge_fs.control_space_commands import KnowledgeFSControlSpaceIntentConflictError
from services.knowledge_fs.control_space_lifecycle import KnowledgeFSControlSpaceLifecycleService


class KnowledgeFSControlSpaceRegistration(NamedTuple):
    tenant_id: str
    owner_account_id: str
    provisioning_key: str
    knowledge_space_id: str
    knowledge_space_revision: int


class KnowledgeFSControlSpaceDryRunReport(NamedTuple):
    total: int
    by_state: dict[str, int]
    irreversible_deletions: int


class KnowledgeFSControlSpaceBackfillReport(NamedTuple):
    candidates: int
    registered: int
    replays: int


class KnowledgeFSControlSpaceManagementService:
    def __init__(self, session_maker: sessionmaker[Session]):
        self._session_maker = session_maker

    def dry_run(self, *, tenant_id: str | None = None) -> KnowledgeFSControlSpaceDryRunReport:
        with self._session_maker() as session:
            if tenant_id is None:
                from sqlalchemy import select

                control_spaces = tuple(session.scalars(select(KnowledgeFSControlSpace)))
            else:
                control_spaces = SQLAlchemyKnowledgeFSControlSpaceRepository(session).list_for_tenant(
                    tenant_id=tenant_id
                )
        counts = Counter(control_space.state.value for control_space in control_spaces)
        return KnowledgeFSControlSpaceDryRunReport(
            total=len(control_spaces),
            by_state=dict(sorted(counts.items())),
            irreversible_deletions=sum(
                control_space.deletion_irreversible_at is not None for control_space in control_spaces
            ),
        )

    def register(self, registration: KnowledgeFSControlSpaceRegistration) -> tuple[KnowledgeFSControlSpace, bool]:
        with self._session_maker.begin() as session:
            repository = SQLAlchemyKnowledgeFSControlSpaceRepository(session)
            by_key = repository.find_by_provisioning_key(provisioning_key=registration.provisioning_key)
            by_space = repository.find_by_knowledge_space_id(
                tenant_id=registration.tenant_id,
                knowledge_space_id=registration.knowledge_space_id,
            )
            existing = by_key or by_space
            if existing is not None:
                if (
                    existing.tenant_id != registration.tenant_id
                    or existing.owner_account_id != registration.owner_account_id
                    or existing.provisioning_key != registration.provisioning_key
                    or existing.knowledge_space_id != registration.knowledge_space_id
                ):
                    raise KnowledgeFSControlSpaceIntentConflictError(
                        "KnowledgeFS registration conflicts with an existing local identity"
                    )
                self._ensure_owner_permission(session, control_space=existing)
                return existing, True
            control_space = repository.add(
                KnowledgeFSControlSpace(
                    tenant_id=registration.tenant_id,
                    owner_account_id=registration.owner_account_id,
                    provisioning_key=registration.provisioning_key,
                    knowledge_space_id=registration.knowledge_space_id,
                    knowledge_space_revision=registration.knowledge_space_revision,
                    state=KnowledgeFSControlSpaceState.ACTIVE,
                    lifecycle_operation_id="management-register",
                )
            )
            session.add(
                KnowledgeFSAuthorizationRevision(
                    tenant_id=registration.tenant_id,
                    control_space_id=control_space.id,
                )
            )
            self._ensure_owner_permission(session, control_space=control_space)
            return control_space, False

    @staticmethod
    def _ensure_owner_permission(session: Session, *, control_space: KnowledgeFSControlSpace) -> None:
        permission = session.scalar(
            sa.select(KnowledgeFSControlSpacePermission).where(
                KnowledgeFSControlSpacePermission.tenant_id == control_space.tenant_id,
                KnowledgeFSControlSpacePermission.control_space_id == control_space.id,
                KnowledgeFSControlSpacePermission.account_id == control_space.owner_account_id,
            )
        )
        if permission is None:
            session.add(
                KnowledgeFSControlSpacePermission(
                    tenant_id=control_space.tenant_id,
                    control_space_id=control_space.id,
                    account_id=control_space.owner_account_id,
                    role=KnowledgeFSControlSpacePermissionRole.OWNER,
                    granted_by_account_id=control_space.owner_account_id,
                )
            )

    def backfill(
        self,
        registrations: Iterable[KnowledgeFSControlSpaceRegistration],
        *,
        apply: bool,
    ) -> KnowledgeFSControlSpaceBackfillReport:
        candidates = tuple(registrations)
        if not apply:
            return KnowledgeFSControlSpaceBackfillReport(len(candidates), 0, 0)
        registered = 0
        replays = 0
        for registration in candidates:
            _, replayed = self.register(registration)
            registered += not replayed
            replays += replayed
        return KnowledgeFSControlSpaceBackfillReport(len(candidates), registered, replays)

    def repair_registration(
        self,
        *,
        tenant_id: str,
        control_space_id: str,
        expected_resource_version: int,
        knowledge_space_id: str,
        knowledge_space_revision: int,
    ) -> KnowledgeFSControlSpace:
        with self._session_maker.begin() as session:
            repository = SQLAlchemyKnowledgeFSControlSpaceRepository(session)
            return KnowledgeFSControlSpaceLifecycleService(repository).transition(
                tenant_id=tenant_id,
                control_space_id=control_space_id,
                expected_resource_version=expected_resource_version,
                new_state=KnowledgeFSControlSpaceState.ACTIVE,
                lifecycle_operation_id="management-repair",
                knowledge_space_id=knowledge_space_id,
                knowledge_space_revision=knowledge_space_revision,
            )


__all__ = [
    "KnowledgeFSControlSpaceBackfillReport",
    "KnowledgeFSControlSpaceDryRunReport",
    "KnowledgeFSControlSpaceManagementService",
    "KnowledgeFSControlSpaceRegistration",
]