tenant-engine/src/tenant_engine/postgres_store.py

from __future__ import annotations

from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import replace
from datetime import UTC, datetime
from pathlib import Path
from threading import RLock
from typing import Any

try:  # Optional runtime dependency; SQLite/development installs stay small.
    import psycopg
    from psycopg.errors import UniqueViolation
    from psycopg.rows import dict_row
    from psycopg.types.json import Jsonb
    from psycopg_pool import ConnectionPool, PoolTimeout
except ImportError:  # pragma: no cover - exercised by the deployment guard
    psycopg = None
    ConnectionPool = None
    PoolTimeout = None

from tenant_engine.audit_core import envelope_for, new_event_id
from tenant_engine.domain import (
    CapabilityRole,
    PlanAssignment,
    RoleGrant,
    Tenant,
    TenantLifecycle,
    TenantRetiredError,
)
from tenant_engine.guardrail import (
    DEFAULT_REGISTRY,
    GuardrailChange,
    LimitValue,
    dump_limit,
    load_limit,
)
from tenant_engine.store import (
    AuthorizationRecord,
    DomainEvent,
    GrantNotFoundError,
    IdempotencyConflictError,
    OutboxRow,
    StoreUnavailableError,
    TenantAlreadyExistsError,
    TenantNotFoundError,
    VersionConflictError,
    guard_guardrail_write,
)


class RotatingConnectionPool:
    """A bounded pool whose DSN is read from a projected file per checkout.

    OpenBao can replace the file atomically when a database lease rotates.
    The next checkout swaps pools, while max_lifetime also bounds how long a
    still-valid connection can retain an old credential.
    """

    def __init__(self, dsn_file: str, *, min_size: int = 1, max_size: int = 4) -> None:
        if ConnectionPool is None:
            raise RuntimeError("install tenant-engine[postgres] for PostgreSQL support")
        self._dsn_file = Path(dsn_file)
        self._min_size = min_size
        self._max_size = max_size
        self._lock = RLock()
        self._dsn = ""
        self._pool = None

    @contextmanager
    def connection(self) -> Iterator[Any]:
        try:
            dsn = self._dsn_file.read_text(encoding="utf-8").strip()
            if not dsn:
                raise OSError("database URL file is empty")
            old_pool = None
            with self._lock:
                if self._pool is None or dsn != self._dsn:
                    old_pool = self._pool
                    self._pool = ConnectionPool(
                        conninfo=dsn,
                        min_size=self._min_size,
                        max_size=self._max_size,
                        timeout=5,
                        max_lifetime=300,
                        kwargs={"row_factory": dict_row},
                        open=True,
                    )
                    self._dsn = dsn
                pool = self._pool
            if old_pool is not None:
                old_pool.close()
            with pool.connection() as connection:
                yield connection
        except OSError as exc:
            raise StoreUnavailableError("database credential unavailable") from exc
        except (psycopg.Error, PoolTimeout) as exc:
            raise StoreUnavailableError("PostgreSQL unavailable") from exc

    def close(self) -> None:
        with self._lock:
            if self._pool is not None:
                self._pool.close()
                self._pool = None


class PostgresTenantStore:
    """PostgreSQL implementation of the complete TenantStore protocol."""

    backend_name = "postgresql"

    def __init__(self, dsn_file: str, *, min_pool_size: int = 1, max_pool_size: int = 4) -> None:
        self._pool = RotatingConnectionPool(
            dsn_file, min_size=min_pool_size, max_size=max_pool_size
        )

    def close(self) -> None:
        self._pool.close()

    def ping(self) -> None:
        with self._pool.connection() as conn:
            conn.execute("SELECT 1").fetchone()

    def create_tenant(self, tenant: Tenant, *, authz: dict[str, Any] | None = None) -> None:
        try:
            with self._pool.connection() as conn, conn.transaction():
                conn.execute(
                    """INSERT INTO tenants
                       (tenant_id, identifier, grouping_name, display_name, contact_email,
                        lifecycle, version, created_at, updated_at, retired_at, reactivated_at)
                       VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
                    (
                        tenant.tenant_id,
                        tenant.identifier,
                        tenant.grouping,
                        tenant.display_name,
                        tenant.contact_email,
                        tenant.lifecycle.value,
                        tenant.version,
                        tenant.created_at,
                        tenant.updated_at,
                        tenant.retired_at,
                        tenant.reactivated_at,
                    ),
                )
                self._emit(
                    conn,
                    "tenant_created",
                    tenant.tenant_id,
                    {"identifier": tenant.identifier, "grouping": tenant.grouping, **(authz or {})},
                )
        except StoreUnavailableError as exc:
            if isinstance(exc.__cause__, UniqueViolation):
                raise TenantAlreadyExistsError(tenant.identifier) from exc
            raise

    def get_tenant(self, tenant_id: str) -> Tenant:
        with self._pool.connection() as conn:
            row = conn.execute(
                "SELECT * FROM tenants WHERE tenant_id = %s OR identifier = %s",
                (tenant_id, tenant_id),
            ).fetchone()
        if row is None:
            raise TenantNotFoundError(tenant_id)
        return _tenant(row)

    def mutate_tenant(
        self,
        *,
        tenant_id: str,
        expected_version: int,
        mutate: Callable[[Tenant], Tenant],
        event_type: str,
        evidence: dict[str, Any],
        idempotency_key: str,
        request_fingerprint: str,
    ) -> tuple[Tenant, bool]:
        with self._pool.connection() as conn, conn.transaction():
            row = self._locked_tenant(conn, tenant_id)
            receipt = conn.execute(
                """SELECT request_fingerprint, result FROM idempotency_receipts
                   WHERE tenant_id = %s AND idempotency_key = %s""",
                (row["tenant_id"], idempotency_key),
            ).fetchone()
            if receipt is not None:
                if receipt["request_fingerprint"] != request_fingerprint:
                    raise IdempotencyConflictError(idempotency_key)
                return _tenant(receipt["result"]), True

            current = _tenant(row)
            if current.version != expected_version:
                raise VersionConflictError(expected=expected_version, actual=current.version)
            updated = mutate(current)
            conn.execute(
                """UPDATE tenants SET grouping_name = %s, display_name = %s,
                   contact_email = %s, lifecycle = %s, version = %s, updated_at = %s,
                   retired_at = %s, reactivated_at = %s WHERE tenant_id = %s""",
                (
                    updated.grouping,
                    updated.display_name,
                    updated.contact_email,
                    updated.lifecycle.value,
                    updated.version,
                    updated.updated_at,
                    updated.retired_at,
                    updated.reactivated_at,
                    updated.tenant_id,
                ),
            )
            self._record_receipt(conn, updated, idempotency_key, request_fingerprint)
            self._emit(
                conn, event_type, updated.tenant_id, {**evidence, "version": updated.version}
            )
            return updated, False

    def grant_role(self, grant: RoleGrant, *, authz: dict[str, Any] | None = None) -> None:
        tenant = self.get_tenant(grant.tenant_id)
        self._require_active(tenant, "grant a role")
        with self._pool.connection() as conn, conn.transaction():
            conn.execute(
                """INSERT INTO grants
                   (grant_id, tenant_id, role, grant_reason, plan_id, granted_by,
                    granted_at, correlation_id, revoked_at)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                   ON CONFLICT (grant_id) DO UPDATE SET
                     tenant_id = EXCLUDED.tenant_id, role = EXCLUDED.role,
                     grant_reason = EXCLUDED.grant_reason, plan_id = EXCLUDED.plan_id,
                     granted_by = EXCLUDED.granted_by, granted_at = EXCLUDED.granted_at,
                     correlation_id = EXCLUDED.correlation_id,
                     revoked_at = EXCLUDED.revoked_at""",
                (
                    grant.grant_id,
                    tenant.tenant_id,
                    grant.role.value,
                    grant.grant_reason,
                    grant.plan_id,
                    grant.granted_by,
                    grant.granted_at,
                    grant.correlation_id,
                    grant.revoked_at,
                ),
            )
            self._emit(
                conn,
                "role_granted",
                tenant.tenant_id,
                {
                    "grant_id": grant.grant_id,
                    "role": grant.role.value,
                    "grant_reason": grant.grant_reason,
                    "correlation_id": grant.correlation_id,
                    **(authz or {}),
                },
            )

    def revoke_role(
        self, *, tenant_id: str, grant_id: str, at: datetime, authz: dict[str, Any] | None = None
    ) -> RoleGrant:
        tenant = self.get_tenant(tenant_id)
        with self._pool.connection() as conn, conn.transaction():
            row = conn.execute(
                "SELECT * FROM grants WHERE tenant_id = %s AND grant_id = %s FOR UPDATE",
                (tenant.tenant_id, grant_id),
            ).fetchone()
            if row is None:
                raise GrantNotFoundError(grant_id)
            grant = _grant(row).revoke(at=at)
            conn.execute("UPDATE grants SET revoked_at = %s WHERE grant_id = %s", (at, grant_id))
            self._emit(
                conn,
                "role_revoked",
                tenant.tenant_id,
                {"grant_id": grant_id, "role": grant.role.value, **(authz or {})},
            )
            return grant

    def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]:
        tenant = self.get_tenant(tenant_id)
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT role FROM grants WHERE tenant_id = %s AND revoked_at IS NULL",
                (tenant.tenant_id,),
            ).fetchall()
        return frozenset(CapabilityRole(row["role"]) for row in rows)

    def assign_plan(
        self, assignment: PlanAssignment, *, authz: dict[str, Any] | None = None
    ) -> None:
        tenant = self.get_tenant(assignment.tenant_id)
        self._require_active(tenant, "assign a plan")
        with self._pool.connection() as conn, conn.transaction():
            conn.execute(
                """INSERT INTO plans (tenant_id, plan_id, assigned_at) VALUES (%s, %s, %s)
                   ON CONFLICT (tenant_id) DO UPDATE SET
                     plan_id = EXCLUDED.plan_id, assigned_at = EXCLUDED.assigned_at""",
                (tenant.tenant_id, assignment.plan_id, assignment.assigned_at),
            )
            self._emit(
                conn,
                "plan_assigned",
                tenant.tenant_id,
                {"plan_id": assignment.plan_id, **(authz or {})},
            )

    def events_for(self, tenant_id: str) -> list[DomainEvent]:
        tenant = self.get_tenant(tenant_id)
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM events WHERE tenant_id = %s ORDER BY seq",
                (tenant.tenant_id,),
            ).fetchall()
        return [
            DomainEvent(row["event_type"], row["tenant_id"], row["at"], row["payload"])
            for row in rows
        ]

    def record_authorization(self, record: AuthorizationRecord) -> None:
        with self._pool.connection() as conn, conn.transaction():
            conn.execute(
                """INSERT INTO authz_records
                   (action, tenant_id, actor, allowed, source, reason, at,
                    decision_id, request_digest, effect, stance)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
                (
                    record.action,
                    record.tenant_id,
                    record.actor,
                    record.allowed,
                    record.source,
                    record.reason,
                    record.at,
                    record.decision_id,
                    record.request_digest,
                    record.effect,
                    record.stance,
                ),
            )

    def authorization_records(self, tenant_id: str) -> list[AuthorizationRecord]:
        keys = [tenant_id]
        try:
            keys.append(self.get_tenant(tenant_id).tenant_id)
        except TenantNotFoundError:
            pass
        placeholders = ",".join(["%s"] * len(keys))
        with self._pool.connection() as conn:
            rows = conn.execute(
                f"SELECT * FROM authz_records WHERE tenant_id IN ({placeholders}) ORDER BY seq",
                keys,
            ).fetchall()
        return [_authz_row(row) for row in rows]

    def pending_outbox(self) -> list[OutboxRow]:
        with self._pool.connection() as conn:
            rows = conn.execute(
                """SELECT * FROM audit_outbox
                   WHERE delivered_at IS NULL AND dead_at IS NULL
                   ORDER BY created_at"""
            ).fetchall()
        return [_outbox_row(row) for row in rows]

    def mark_outbox(self, event_id: str, *, status: str, detail: str = "") -> None:
        with self._pool.connection() as conn, conn.transaction():
            if status in {"delivered", "duplicate"}:
                conn.execute(
                    """UPDATE audit_outbox SET delivered_at = %s, last_error = NULL
                       WHERE event_id = %s""",
                    (datetime.now(UTC), event_id),
                )
            elif status == "dead":
                conn.execute(
                    "UPDATE audit_outbox SET dead_at = %s, last_error = %s WHERE event_id = %s",
                    (datetime.now(UTC), detail, event_id),
                )
            else:
                conn.execute(
                    """UPDATE audit_outbox SET attempts = attempts + 1, last_error = %s
                       WHERE event_id = %s""",
                    (detail, event_id),
                )

    def guardrail_overrides(self, tenant_id: str) -> dict[str, LimitValue]:
        tenant = self.get_tenant(tenant_id)
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM guardrail_overrides WHERE tenant_id = %s", (tenant.tenant_id,)
            ).fetchall()
        return {row["limit_key"]: load_limit(row) for row in rows}

    def guardrail_changes(self, tenant_id: str) -> list[GuardrailChange]:
        tenant = self.get_tenant(tenant_id)
        with self._pool.connection() as conn:
            rows = conn.execute(
                """SELECT * FROM guardrail_changes WHERE tenant_id = %s
                   ORDER BY changed_at, change_id""",
                (tenant.tenant_id,),
            ).fetchall()
        return [_change(row) for row in rows]

    def set_guardrail_override(
        self,
        *,
        tenant_id: str,
        expected_version: int,
        limit_key: str,
        value: LimitValue | None,
        change_id: str,
        changed_by: str,
        reason: str,
        correlation_id: str,
        idempotency_key: str,
        request_fingerprint: str,
        at: datetime,
        authz: dict[str, Any] | None = None,
    ) -> tuple[Tenant, GuardrailChange | None, bool]:
        DEFAULT_REGISTRY.get(limit_key)
        with self._pool.connection() as conn, conn.transaction():
            row = self._locked_tenant(conn, tenant_id)
            receipt = conn.execute(
                """SELECT request_fingerprint, result FROM idempotency_receipts
                   WHERE tenant_id = %s AND idempotency_key = %s""",
                (row["tenant_id"], idempotency_key),
            ).fetchone()
            if receipt is not None:
                if receipt["request_fingerprint"] != request_fingerprint:
                    raise IdempotencyConflictError(idempotency_key)
                prior = conn.execute(
                    "SELECT * FROM guardrail_changes WHERE change_id = %s", (change_id,)
                ).fetchone()
                return _tenant(receipt["result"]), _change(prior) if prior else None, True

            current = _tenant(row)
            if current.version != expected_version:
                raise VersionConflictError(expected=expected_version, actual=current.version)
            override_rows = conn.execute(
                "SELECT * FROM guardrail_overrides WHERE tenant_id = %s",
                (current.tenant_id,),
            ).fetchall()
            overrides = {item["limit_key"]: load_limit(item) for item in override_rows}
            guard_guardrail_write(
                tenant=current, overrides=overrides, limit_key=limit_key, value=value
            )
            change = GuardrailChange(
                change_id=change_id,
                tenant_id=current.tenant_id,
                limit_key=limit_key,
                previous=overrides.get(limit_key),
                current=value,
                changed_by=changed_by,
                reason=reason,
                correlation_id=correlation_id,
                changed_at=at,
            )
            if value is None:
                conn.execute(
                    "DELETE FROM guardrail_overrides WHERE tenant_id = %s AND limit_key = %s",
                    (current.tenant_id, limit_key),
                )
            else:
                dumped = dump_limit(value)
                conn.execute(
                    """INSERT INTO guardrail_overrides
                       (tenant_id, limit_key, kind, amount, currency, period)
                       VALUES (%s, %s, %s, %s, %s, %s)
                       ON CONFLICT (tenant_id, limit_key) DO UPDATE SET
                         kind = EXCLUDED.kind, amount = EXCLUDED.amount,
                         currency = EXCLUDED.currency, period = EXCLUDED.period""",
                    (
                        current.tenant_id,
                        limit_key,
                        dumped["kind"],
                        dumped["amount"],
                        dumped["currency"],
                        dumped["period"],
                    ),
                )
            conn.execute(
                """INSERT INTO guardrail_changes
                   (change_id, tenant_id, limit_key, previous, current, changed_by,
                    reason, correlation_id, changed_at)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)""",
                (
                    change.change_id,
                    change.tenant_id,
                    change.limit_key,
                    Jsonb(dump_limit(change.previous)) if change.previous else None,
                    Jsonb(dump_limit(change.current)) if change.current else None,
                    change.changed_by,
                    change.reason,
                    change.correlation_id,
                    change.changed_at,
                ),
            )
            updated = replace(current, version=current.version + 1, updated_at=at)
            conn.execute(
                "UPDATE tenants SET version = %s, updated_at = %s WHERE tenant_id = %s",
                (updated.version, updated.updated_at, updated.tenant_id),
            )
            self._record_receipt(conn, updated, idempotency_key, request_fingerprint)
            self._emit(
                conn,
                "guardrail_changed",
                updated.tenant_id,
                {
                    "limit_key": limit_key,
                    "change_id": change_id,
                    "cleared": value is None,
                    "correlation_id": correlation_id,
                    "version": updated.version,
                    **(authz or {}),
                },
            )
            return updated, change, False

    @staticmethod
    def _require_active(tenant: Tenant, what: str) -> None:
        if tenant.lifecycle is not TenantLifecycle.ACTIVE:
            raise TenantRetiredError(f"cannot {what} on a retired tenant")

    @staticmethod
    def _locked_tenant(conn: Any, tenant_id: str) -> Mapping[str, Any]:
        row = conn.execute(
            """SELECT * FROM tenants WHERE tenant_id = %s OR identifier = %s
               FOR UPDATE""",
            (tenant_id, tenant_id),
        ).fetchone()
        if row is None:
            raise TenantNotFoundError(tenant_id)
        return row

    @staticmethod
    def _record_receipt(
        conn: Any, tenant: Tenant, idempotency_key: str, request_fingerprint: str
    ) -> None:
        conn.execute(
            """INSERT INTO idempotency_receipts
               (tenant_id, idempotency_key, request_fingerprint, result, recorded_at)
               VALUES (%s, %s, %s, %s, %s)""",
            (
                tenant.tenant_id,
                idempotency_key,
                request_fingerprint,
                Jsonb(_row(tenant)),
                datetime.now(UTC),
            ),
        )

    @staticmethod
    def _emit(conn: Any, event_type: str, tenant_id: str, payload: dict[str, Any]) -> None:
        event_id = str(payload.get("event_id") or new_event_id())
        payload = {**payload, "event_id": event_id}
        at = datetime.now(UTC)
        conn.execute(
            "INSERT INTO events (event_type, tenant_id, at, payload) VALUES (%s, %s, %s, %s)",
            (event_type, tenant_id, at, Jsonb(payload)),
        )
        envelope = envelope_for(
            event_id=event_id,
            event_type=event_type,
            tenant_id=tenant_id,
            observed_at=at.isoformat(),
            payload=payload,
        )
        conn.execute(
            """INSERT INTO audit_outbox (event_id, tenant_id, envelope, created_at)
               VALUES (%s, %s, %s, %s)""",
            (event_id, tenant_id, Jsonb(envelope), at),
        )


def _tenant(row: Mapping[str, Any]) -> Tenant:
    return Tenant(
        tenant_id=row["tenant_id"],
        identifier=row["identifier"],
        grouping=row["grouping_name"],
        display_name=row["display_name"],
        contact_email=row["contact_email"],
        lifecycle=TenantLifecycle(row["lifecycle"]),
        version=int(row["version"]),
        created_at=_dt(row["created_at"]),
        updated_at=_dt(row["updated_at"]),
        retired_at=_dt(row["retired_at"]),
        reactivated_at=_dt(row["reactivated_at"]),
    )


def _row(tenant: Tenant) -> dict[str, Any]:
    return {
        "tenant_id": tenant.tenant_id,
        "identifier": tenant.identifier,
        "grouping_name": tenant.grouping,
        "display_name": tenant.display_name,
        "contact_email": tenant.contact_email,
        "lifecycle": tenant.lifecycle.value,
        "version": tenant.version,
        "created_at": tenant.created_at.isoformat() if tenant.created_at else None,
        "updated_at": tenant.updated_at.isoformat() if tenant.updated_at else None,
        "retired_at": tenant.retired_at.isoformat() if tenant.retired_at else None,
        "reactivated_at": (tenant.reactivated_at.isoformat() if tenant.reactivated_at else None),
    }


def _dt(value: datetime | str | None) -> datetime | None:
    if isinstance(value, datetime) or value is None:
        return value
    return datetime.fromisoformat(value)


def _grant(row: Mapping[str, Any]) -> RoleGrant:
    return RoleGrant(
        row["grant_id"],
        row["tenant_id"],
        CapabilityRole(row["role"]),
        row["grant_reason"],
        row["plan_id"],
        row["granted_by"],
        _dt(row["granted_at"]),
        row["correlation_id"],
        _dt(row["revoked_at"]),
    )


def _change(row: Mapping[str, Any]) -> GuardrailChange:
    return GuardrailChange(
        change_id=row["change_id"],
        tenant_id=row["tenant_id"],
        limit_key=row["limit_key"],
        previous=load_limit(row["previous"]) if row["previous"] else None,
        current=load_limit(row["current"]) if row["current"] else None,
        changed_by=row["changed_by"],
        reason=row["reason"],
        correlation_id=row["correlation_id"],
        changed_at=_dt(row["changed_at"]),
    )


def _authz_row(row: Mapping[str, Any]) -> AuthorizationRecord:
    return AuthorizationRecord(
        action=row["action"],
        tenant_id=row["tenant_id"],
        actor=row["actor"],
        allowed=bool(row["allowed"]),
        source=row["source"],
        reason=row["reason"],
        at=_dt(row["at"]),
        decision_id=row["decision_id"],
        request_digest=row["request_digest"],
        effect=row["effect"],
        stance=row["stance"],
    )


def _outbox_row(row: Mapping[str, Any]) -> OutboxRow:
    return OutboxRow(
        event_id=row["event_id"],
        tenant_id=row["tenant_id"],
        envelope=row["envelope"],
        created_at=_dt(row["created_at"]),
        attempts=row["attempts"],
        last_error=row["last_error"],
        delivered_at=_dt(row["delivered_at"]) if row["delivered_at"] else None,
        dead_at=_dt(row["dead_at"]) if row["dead_at"] else None,
    )