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