tenant-engine/src/tenant_engine/store.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Protocol

from tenant_engine.domain import CapabilityRole, PlanAssignment, RoleGrant, Tenant


class TenantNotFoundError(KeyError):
    pass


class TenantAlreadyExistsError(ValueError):
    pass


class GrantNotFoundError(KeyError):
    pass


class StoreUnavailableError(RuntimeError):
    """The store could not answer -- callers on a privileged decision path

    (flex-auth's live lookup) must treat this as deny, never as "zero roles".
    """


@dataclass(frozen=True, slots=True)
class DomainEvent:
    """Boundary contract's Audit Correlation Contract, in event form."""

    event_type: str
    tenant_id: str
    at: datetime
    payload: dict[str, Any]


class TenantStore(Protocol):
    """Swappable persistence seam -- domain/ and api/ depend on this, not a backend."""

    def create_tenant(self, tenant: Tenant) -> None: ...

    def get_tenant(self, tenant_id: str) -> Tenant: ...

    def grant_role(self, grant: RoleGrant) -> None: ...

    def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant: ...

    def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]: ...

    def assign_plan(self, assignment: PlanAssignment) -> None: ...

    def events(self) -> list[DomainEvent]: ...


class InMemoryTenantStore:
    def __init__(self) -> None:
        self._tenants: dict[str, Tenant] = {}
        self._grants: dict[str, dict[str, RoleGrant]] = {}
        self._plans: dict[str, PlanAssignment] = {}
        self._events: list[DomainEvent] = []

    def create_tenant(self, tenant: Tenant) -> None:
        if tenant.tenant_id in self._tenants:
            raise TenantAlreadyExistsError(tenant.tenant_id)
        self._tenants[tenant.tenant_id] = tenant
        self._grants[tenant.tenant_id] = {}
        self._emit(
            "tenant_created",
            tenant.tenant_id,
            {"identifier": tenant.identifier, "grouping": tenant.grouping},
        )

    def get_tenant(self, tenant_id: str) -> Tenant:
        try:
            return self._tenants[tenant_id]
        except KeyError:
            raise TenantNotFoundError(tenant_id) from None

    def grant_role(self, grant: RoleGrant) -> None:
        self.get_tenant(grant.tenant_id)
        self._grants[grant.tenant_id][grant.grant_id] = grant
        self._emit(
            "role_granted",
            grant.tenant_id,
            {
                "grant_id": grant.grant_id,
                "role": grant.role.value,
                "grant_reason": grant.grant_reason,
                "correlation_id": grant.correlation_id,
            },
        )

    def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant:
        self.get_tenant(tenant_id)
        try:
            grant = self._grants[tenant_id][grant_id]
        except KeyError:
            raise GrantNotFoundError(grant_id) from None
        revoked = grant.revoke(at=at)
        self._grants[tenant_id][grant_id] = revoked
        self._emit(
            "role_revoked",
            tenant_id,
            {"grant_id": grant_id, "role": revoked.role.value},
        )
        return revoked

    def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]:
        self.get_tenant(tenant_id)
        return frozenset(
            grant.role for grant in self._grants.get(tenant_id, {}).values() if grant.active
        )

    def assign_plan(self, assignment: PlanAssignment) -> None:
        self.get_tenant(assignment.tenant_id)
        self._plans[assignment.tenant_id] = assignment
        self._emit("plan_assigned", assignment.tenant_id, {"plan_id": assignment.plan_id})

    def events(self) -> list[DomainEvent]:
        return list(self._events)

    def _emit(self, event_type: str, tenant_id: str, payload: dict[str, Any]) -> None:
        self._events.append(
            DomainEvent(event_type=event_type, tenant_id=tenant_id, at=datetime.now(UTC), payload=payload)
        )