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.

    Every method taking a `tenant_id` accepts either the tenant's internal
    id (caller-chosen at creation, e.g. from an admin tool) or its profile
    identifier (e.g. "tenant:friendly:binky") -- external callers like
    key-cape and flex-auth only ever have the identifier (it's the IAM
    Profile `tenant` claim value), never tenant-engine's internal id. See
    `_resolve_tenant_id` for why this had to be added after real
    cross-service testing caught the gap (found integrating key-cape's
    tenant_roles claim, KEY-WP-0005-T02).
    """

    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._by_identifier: dict[str, str] = {}
        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)
        if tenant.identifier in self._by_identifier:
            raise TenantAlreadyExistsError(tenant.identifier)
        self._tenants[tenant.tenant_id] = tenant
        self._by_identifier[tenant.identifier] = tenant.tenant_id
        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:
        return self._tenants[self._resolve(tenant_id)]

    def grant_role(self, grant: RoleGrant) -> None:
        resolved = self._resolve(grant.tenant_id)
        self._grants[resolved][grant.grant_id] = grant
        self._emit(
            "role_granted",
            resolved,
            {
                "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:
        resolved = self._resolve(tenant_id)
        try:
            grant = self._grants[resolved][grant_id]
        except KeyError:
            raise GrantNotFoundError(grant_id) from None
        revoked = grant.revoke(at=at)
        self._grants[resolved][grant_id] = revoked
        self._emit(
            "role_revoked",
            resolved,
            {"grant_id": grant_id, "role": revoked.role.value},
        )
        return revoked

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

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

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

    def _resolve(self, tenant_id: str) -> str:
        """Resolve an internal tenant_id or a profile identifier to the

        canonical internal tenant_id every other private dict is keyed by.

        External callers (key-cape, flex-auth) only ever have the
        identifier (the IAM Profile `tenant` claim value) -- they have no
        way to know a tenant's internal tenant_id, which is caller-chosen
        at creation time and otherwise opaque. Found the hard way: a real
        cross-service check (key-cape's KEY-WP-0005-T02 against a live
        tenant-engine) returned tenant_not_found for a tenant that
        genuinely existed, because the caller only had the identifier.
        """
        resolved = self._by_identifier.get(tenant_id, tenant_id)
        if resolved not in self._tenants:
            raise TenantNotFoundError(tenant_id)
        return resolved

    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)
        )