pgbulkhead 0.1.0__py3-none-any.whl

pgbulkhead/__init__.py

@@ -0,0 +1,71 @@
+"""pgbulkhead: tenant isolation for multi-tenant MCP servers on Postgres.
+
+One database, many tenants and workspaces: Postgres enforces the walls, and no agent can talk
+its way through them.
+
+  * :class:`TenancyModel` names the walls (a tenant column, an optional workspace column) and
+    the settings that carry them; :func:`policy_sql` renders its row-level security policies
+    for a migration.
+  * :class:`SecurityContext` is the verified identity of one request, and
+    :func:`scoped_session` binds it for the length of a transaction, so every query runs
+    inside its walls and a pooled connection never carries it to the next request.
+  * :class:`SecuredDatabase` holds one server's engine, pool settings and per-workspace
+    admission, so one agent's fan-out queues behind its own workspace's slots, and refuses
+    an overloaded call with :class:`Overloaded` before it runs anything.
+
+Isolation holds against query bugs and against identity taken from untrusted input. It does
+not hold against arbitrary SQL executed as the app role, which can set any context setting.
+The README's threat model says what each guarantee rests on.
+"""
+
+from .admission import (
+    AdmittedSlot,
+    Overloaded,
+    PoolExhausted,
+    RefusalText,
+    WorkspaceBusy,
+    WorkspaceLimiter,
+    admitted_session,
+)
+from .context import ClaimNames, SecurityContext, context_from_claims
+from .database import SecuredDatabase
+from .model import GucBypass, RoleBypass, TenancyModel
+from .policy import policy_sql, policy_statements
+from .pool import CONNECTIONS_PER_CALL, PoolConfigError, PoolSettings, create_pooled_engine
+from .session import (
+    CLEANUP_TIMEOUT_SECONDS,
+    bind_context,
+    check_context,
+    scoped_session,
+    system_session,
+    transaction,
+)
+
+__all__ = [
+    "CLEANUP_TIMEOUT_SECONDS",
+    "CONNECTIONS_PER_CALL",
+    "AdmittedSlot",
+    "ClaimNames",
+    "GucBypass",
+    "Overloaded",
+    "PoolConfigError",
+    "PoolExhausted",
+    "PoolSettings",
+    "RefusalText",
+    "RoleBypass",
+    "SecuredDatabase",
+    "SecurityContext",
+    "TenancyModel",
+    "WorkspaceBusy",
+    "WorkspaceLimiter",
+    "admitted_session",
+    "bind_context",
+    "check_context",
+    "context_from_claims",
+    "create_pooled_engine",
+    "policy_sql",
+    "policy_statements",
+    "scoped_session",
+    "system_session",
+    "transaction",
+]

pgbulkhead/admission.py

@@ -0,0 +1,235 @@
+"""Per-workspace admission: no one workspace can take every pooled connection on a pod.
+
+Every call takes a pooled connection for its session, and the pool is shared by every
+workspace the pod serves. Without a limit, one client that fires a burst of calls holds the
+whole pool, and every other workspace's call waits behind it. :class:`WorkspaceLimiter` admits
+at most ``limit`` calls per workspace at once, *before* the call checks a connection out, so a
+burst queues behind its own workspace's slots while other workspaces keep theirs.
+
+The key is the verified ``(tenant_id, workspace_id)`` of the call's :class:`SecurityContext`,
+which a server builds from a verified identity alone; nothing a caller passes can choose it.
+
+A call over the limit waits up to ``wait_seconds`` for a slot, then is refused with
+:class:`WorkspaceBusy`. An admitted call whose checkout then waits past the pool's timeout is
+refused with :class:`PoolExhausted`. Both are :class:`Overloaded`: the call ran nothing, so
+retrying it is safe, and each server maps it onto its own error surface as ``rate_limited``.
+
+State is bounded by the calls in flight: a workspace's entry exists only while some call holds
+or waits for one of its slots, and is dropped when the last one leaves.
+
+Admit only the call's own session (:func:`admitted_session`). A write a call commits apart
+from its session opens a second session while the call already holds a slot; admitting that
+one too could deadlock a workspace whose every slot is held by a call waiting for its second
+session. A request that runs several sessions one after another holds one slot across them
+all (:class:`AdmittedSlot`, from :meth:`SecuredDatabase.admitted
+<pgbulkhead.SecuredDatabase.admitted>`).
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Callable, Mapping
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+import anyio
+from sqlalchemy.exc import TimeoutError as PoolTimeout
+
+from .context import SecurityContext
+from .pool import PoolSettings
+
+__all__ = [
+    "AdmittedSlot",
+    "Overloaded",
+    "RefusalText",
+    "PoolExhausted",
+    "WorkspaceBusy",
+    "WorkspaceLimiter",
+    "admitted_session",
+]
+
+
+class Overloaded(Exception):
+    """The call was refused before it ran anything, because the server is at capacity.
+
+    ``message`` and ``next_step`` are text a server can return to the agent as they are: they
+    say nothing ran, so a retry is safe. Map it onto your error surface as a rate limit. Replace
+    the text with ``refusals`` on :class:`~pgbulkhead.SecuredDatabase` or
+    :func:`admitted_session`.
+    """
+
+    message = "The server is at capacity right now. Nothing was run."
+    next_step = "Wait a few seconds, then retry the call."
+
+    def __init__(self) -> None:
+        super().__init__(self.message)
+
+
+class WorkspaceBusy(Overloaded):
+    """A call waited its bound for one of its workspace's slots and did not get one."""
+
+    message = (
+        "This workspace already has as many requests running as the server allows at once, "
+        "and this one waited too long for its turn. Nothing was run."
+    )
+    next_step = "Wait a few seconds, then retry the call; send fewer calls at once."
+
+
+class PoolExhausted(Overloaded):
+    """An admitted call found no free database connection within the pool's timeout."""
+
+    message = "The server has no free database connection right now. Nothing was run."
+    next_step = "Wait a few seconds, then retry the call."
+
+
+@dataclass(frozen=True, slots=True)
+class RefusalText:
+    """Text to return in place of an :class:`Overloaded` refusal's defaults."""
+
+    message: str
+    next_step: str
+
+
+#: Replacement text per refusal class. A refusal takes its own class's entry, else the entry of
+#: its nearest listed base class, else keeps its defaults.
+Refusals = Mapping[type[Overloaded], RefusalText]
+
+
+def _apply(exc: Overloaded, refusals: Refusals | None) -> Overloaded:
+    """``exc`` carrying the text ``refusals`` gives its class, if any."""
+    if refusals:
+        for cls in type(exc).__mro__:
+            text = refusals.get(cls)
+            if text is not None:
+                exc.message = text.message
+                exc.next_step = text.next_step
+                exc.args = (text.message,)
+                break
+    return exc
+
+
+@dataclass(slots=True)
+class _Slots:
+    semaphore: anyio.Semaphore
+    # The calls holding or waiting for one of these slots; the entry is dropped at zero.
+    users: int = field(default=0)
+
+
+class WorkspaceLimiter:
+    """At most ``limit`` concurrent admitted calls per workspace, waiting ``wait_seconds``."""
+
+    def __init__(self, *, limit: int, wait_seconds: float) -> None:
+        if limit < 1:
+            raise ValueError("limit must be at least 1")
+        self.limit = limit
+        self.wait_seconds = wait_seconds
+        self._slots: dict[tuple[str, str | None], _Slots] = {}
+
+    @classmethod
+    def from_settings(cls, settings: PoolSettings) -> WorkspaceLimiter:
+        """The limiter ``settings`` sized together with the pool."""
+        return cls(
+            limit=settings.workspace_concurrency, wait_seconds=settings.workspace_wait_seconds
+        )
+
+    def tracked(self) -> int:
+        """How many workspaces have an entry: those with a call holding or waiting for a slot."""
+        return len(self._slots)
+
+    @asynccontextmanager
+    async def admit(self, ctx: SecurityContext) -> AsyncGenerator[None, None]:
+        """Hold one of ``ctx``'s workspace's slots for the body, or raise :class:`WorkspaceBusy`.
+
+        The lookup and the count change run with no await between them, so on one event loop
+        no other call can see the entry between its creation and its first user.
+        """
+        key = (ctx.tenant_id, ctx.workspace_id)
+        slots = self._slots.get(key)
+        if slots is None:
+            slots = self._slots[key] = _Slots(anyio.Semaphore(self.limit))
+        slots.users += 1
+        try:
+            try:
+                with anyio.fail_after(self.wait_seconds):
+                    await slots.semaphore.acquire()
+            except TimeoutError:
+                raise WorkspaceBusy from None
+            try:
+                yield
+            finally:
+                slots.semaphore.release()
+        finally:
+            slots.users -= 1
+            if slots.users == 0:
+                del self._slots[key]
+
+
+@asynccontextmanager
+async def _refusing_checkout_timeout[T](
+    session: AbstractAsyncContextManager[T], refusals: Refusals | None = None
+) -> AsyncGenerator[T, None]:
+    """Enter ``session``, turning a pool timeout at its own checkout into :class:`PoolExhausted`.
+
+    A pool timeout raised after the session is open belongs to the body (a second session it
+    opened, say, which may follow a committed write), so it propagates unchanged.
+    """
+    entered = False
+    try:
+        async with session as opened:
+            entered = True
+            yield opened
+    except PoolTimeout:
+        if entered:
+            raise
+        raise _apply(PoolExhausted(), refusals) from None
+
+
+class AdmittedSlot:
+    """One of a workspace's slots, held by a request that opens several sessions in turn.
+
+    Each :meth:`session` runs under the held slot, so the request uses one connection at a
+    time and never waits on its own workspace's slots. A checkout that times out raises the
+    pool's own ``TimeoutError``, not :class:`PoolExhausted`: an earlier session under the slot
+    may have committed, so the request is not safe to retry as a whole.
+    """
+
+    def __init__(
+        self,
+        ctx: SecurityContext,
+        open_session: Callable[[SecurityContext], AbstractAsyncContextManager[Any]],
+    ) -> None:
+        self.ctx = ctx
+        self._open_session = open_session
+
+    def session(self) -> AbstractAsyncContextManager[Any]:
+        """A session for the slot's context, under the held slot."""
+        return self._open_session(self.ctx)
+
+
+@asynccontextmanager
+async def admitted_session[T](
+    limiter: WorkspaceLimiter,
+    ctx: SecurityContext,
+    open_session: Callable[[SecurityContext], AbstractAsyncContextManager[T]],
+    *,
+    refusals: Refusals | None = None,
+) -> AsyncGenerator[T, None]:
+    """Admit the call for ``ctx``'s workspace, then open its session with ``open_session``.
+
+    Raises :class:`WorkspaceBusy` when no slot frees within the limiter's wait, and
+    :class:`PoolExhausted` when the session's own checkout times out. A pool timeout raised
+    after the session is open belongs to the body, so it propagates unchanged.
+
+    The slot is released by the same exit that returns the session's connection, so a call
+    left by cancellation gives its slot back.
+    """
+    try:
+        async with (
+            limiter.admit(ctx),
+            _refusing_checkout_timeout(open_session(ctx), refusals) as session,
+        ):
+            yield session
+    except WorkspaceBusy as exc:
+        # The same object, re-raised: a refusal from the body is this server's too.
+        _apply(exc, refusals)
+        raise

pgbulkhead/context.py

@@ -0,0 +1,114 @@
+"""The verified identity a session binds.
+
+A :class:`SecurityContext` is built **only** from something your server has verified: the
+claims of a token it checked, or headers set by a gateway that checked one and that a caller
+cannot reach past. Never from a tool argument or a request body. An agent writes those, and an
+agent writes what its prompt steers it to, so a tenant taken from one is a tenant anyone who can
+put text in front of the agent can choose. The body may *ask* for a workspace; the verified
+identity decides whether it gets one.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, overload
+
+__all__ = ["ClaimNames", "SecurityContext", "context_from_claims"]
+
+
+@dataclass(frozen=True, slots=True)
+class SecurityContext:
+    """The identity and grants of one request.
+
+    - ``tenant_id``: the hard wall. Never crossed.
+    - ``workspace_id``: the active workspace, the soft wall inside the tenant, and the only one
+      writes go to. ``None`` for a single-key model, which has no workspace.
+    - ``allowed_workspaces``: the workspaces the request may **read**, from grants in the
+      verified identity. A session reads the active workspace whether or not it is listed. A
+      read grant never implies a write.
+    - ``scopes``: capability scopes, checked with :meth:`require_scope`.
+    """
+
+    tenant_id: str
+    workspace_id: str | None = None
+    allowed_workspaces: frozenset[str] = frozenset()
+    scopes: frozenset[str] = frozenset()
+    user_id: str | None = None
+    correlation_id: str | None = None
+
+    def readable_workspaces(self) -> frozenset[str]:
+        """The workspaces a session may read: the grants plus the active one."""
+        if self.workspace_id is None:
+            return self.allowed_workspaces
+        return self.allowed_workspaces | {self.workspace_id}
+
+    def require_scope(self, scope: str) -> None:
+        """Raise :class:`PermissionError` unless the context carries ``scope``."""
+        if scope not in self.scopes:
+            raise PermissionError(f"missing required scope: {scope}")
+
+
+@dataclass(frozen=True, slots=True)
+class ClaimNames:
+    """The claim each context field is read from. ``workspace`` is ``None`` for a single-key
+    model, whose context has no workspace."""
+
+    tenant: str = "tenant_id"
+    workspace: str | None = "workspace_id"
+    allowed_workspaces: str = "allowed_workspaces"
+    scopes: str = "scopes"
+    user_id: str = "user_id"
+    correlation_id: str = "correlation_id"
+
+
+@overload
+def context_from_claims(
+    claims: Mapping[str, object], names: ClaimNames | None = None
+) -> SecurityContext: ...
+@overload
+def context_from_claims[C: SecurityContext](
+    claims: Mapping[str, object], names: ClaimNames | None = None, *, cls: type[C]
+) -> C: ...
+def context_from_claims(
+    claims: Mapping[str, object],
+    names: ClaimNames | None = None,
+    *,
+    cls: type[Any] = SecurityContext,
+) -> Any:
+    """A context from **verified** token claims; ``cls`` builds a subclass.
+
+    Fails closed: a missing or empty tenant (or workspace, when ``names.workspace`` is set)
+    raises :class:`ValueError` rather than defaulting to a wider scope. The active workspace is
+    always readable, so it is added to ``allowed_workspaces``.
+    """
+    names = names or ClaimNames()
+    tenant = claims.get(names.tenant)
+    if not isinstance(tenant, str) or not tenant:
+        raise ValueError(f"security context requires a verified {names.tenant} claim")
+    workspace: str | None = None
+    if names.workspace is not None:
+        value = claims.get(names.workspace)
+        if not isinstance(value, str) or not value:
+            raise ValueError(f"security context requires a verified {names.workspace} claim")
+        workspace = value
+    allowed = _str_set(claims.get(names.allowed_workspaces))
+    return cls(
+        tenant_id=tenant,
+        workspace_id=workspace,
+        allowed_workspaces=allowed | {workspace} if workspace else allowed,
+        scopes=_str_set(claims.get(names.scopes)),
+        user_id=_opt_str(claims.get(names.user_id)),
+        correlation_id=_opt_str(claims.get(names.correlation_id)),
+    )
+
+
+def _opt_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None
+
+
+def _str_set(value: object) -> frozenset[str]:
+    """A loosely typed claim (a list, tuple or set, or nothing) as a set of strings."""
+    if isinstance(value, list | tuple | set | frozenset):
+        return frozenset(str(x) for x in value)
+    return frozenset()

pgbulkhead/database.py

@@ -0,0 +1,144 @@
+"""One server's database: its tenancy model, pool settings, pooled engine and workspace limiter.
+
+A server declares one :class:`SecuredDatabase` per process::
+
+    db = SecuredDatabase(model=MODEL, url_env="NOTES_DATABASE_URL", prefix="NOTES")
+
+and calls :meth:`SecuredDatabase.check` once before it serves. The object reads and validates
+the settings once, builds the engine lazily on the first session (importing the server needs no
+database), and admits every :meth:`~SecuredDatabase.session` against the workspace's
+concurrency limit before it checks a connection out.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import AsyncGenerator, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+from .admission import (
+    AdmittedSlot,
+    Refusals,
+    WorkspaceBusy,
+    WorkspaceLimiter,
+    _apply,
+    admitted_session,
+)
+from .context import SecurityContext
+from .model import TenancyModel
+from .pool import PoolSettings, create_pooled_engine
+from .session import scoped_session, system_session
+
+__all__ = ["SecuredDatabase"]
+
+
+class SecuredDatabase:
+    """The model, settings, engine and limiter one server process uses for its own database.
+
+    ``model`` is how its tables are isolated. ``url_env`` names the env var holding the async
+    Postgres URL, for a role that is neither a superuser nor ``BYPASSRLS``; ``url`` instead
+    supplies it with a callable. ``prefix`` is the env namespace of the pool settings
+    (:class:`PoolSettings`), and ``defaults`` the server's own when its traffic needs more than
+    the class defaults. ``refusals`` replaces the text of the overload errors its sessions raise.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: TenancyModel,
+        prefix: str,
+        url_env: str | None = None,
+        url: Callable[[], str] | None = None,
+        defaults: PoolSettings | None = None,
+        refusals: Refusals | None = None,
+    ) -> None:
+        if url_env is None and url is None:
+            raise ValueError("SecuredDatabase needs url_env or url")
+        self.model = model
+        self.url_env = url_env
+        self.prefix = prefix
+        self.refusals = refusals
+        self._defaults = defaults
+        self._url = url or self._url_from_env
+        self._settings: PoolSettings | None = None
+        self._limiter: WorkspaceLimiter | None = None
+        self._engine: AsyncEngine | None = None
+
+    def _url_from_env(self) -> str:
+        assert self.url_env is not None
+        url = os.environ.get(self.url_env, "").strip()
+        if not url:
+            raise RuntimeError(
+                f"{self.url_env} is not set — the async Postgres URL for this server's app "
+                "role (NOSUPERUSER NOBYPASSRLS) is required."
+            )
+        return url
+
+    @property
+    def settings(self) -> PoolSettings:
+        """The pool and admission settings, read and validated on first use, then fixed."""
+        if self._settings is None:
+            self._settings = PoolSettings.from_env(self.prefix, defaults=self._defaults)
+        return self._settings
+
+    @property
+    def limiter(self) -> WorkspaceLimiter:
+        """The process-wide per-workspace limiter, sized with the pool."""
+        if self._limiter is None:
+            self._limiter = WorkspaceLimiter.from_settings(self.settings)
+        return self._limiter
+
+    def check(self) -> None:
+        """Read and validate the settings and build the limiter; call it at startup.
+
+        Raises :class:`~pgbulkhead.PoolConfigError` for settings that cannot hold, so a bad
+        value fails the rollout rather than a user's call. Opens no connection.
+        """
+        _ = self.limiter
+
+    def get_engine(self) -> AsyncEngine:
+        """The pooled engine, built on first use."""
+        if self._engine is None:
+            self._engine = create_pooled_engine(self._url(), self.settings)
+        return self._engine
+
+    async def dispose(self) -> None:
+        """Dispose the engine and its pool; the next session builds a new one."""
+        if self._engine is not None:
+            engine, self._engine = self._engine, None
+            await engine.dispose()
+
+    def session(self, ctx: SecurityContext) -> AbstractAsyncContextManager[AsyncConnection]:
+        """The call's own scoped session, admitted against ``ctx``'s workspace limit first.
+
+        Raises :class:`~pgbulkhead.Overloaded`, having run nothing, when no slot frees in time
+        or, once admitted, no connection does.
+        """
+        return admitted_session(self.limiter, ctx, self.unadmitted_session, refusals=self.refusals)
+
+    def unadmitted_session(
+        self, ctx: SecurityContext
+    ) -> AbstractAsyncContextManager[AsyncConnection]:
+        """A scoped session with no admission, for one opened while the request holds a slot
+        (a write committed apart from its session, say). Admitting that one too could deadlock
+        a workspace whose every slot waits on a second session."""
+        return scoped_session(self.get_engine(), self.model, ctx)
+
+    @asynccontextmanager
+    async def admitted(self, ctx: SecurityContext) -> AsyncGenerator[AdmittedSlot, None]:
+        """Hold one of ``ctx``'s workspace slots across several sessions opened in turn."""
+        try:
+            async with self.limiter.admit(ctx):
+                yield AdmittedSlot(ctx, self.unadmitted_session)
+        except WorkspaceBusy as exc:
+            _apply(exc, self.refusals)
+            raise
+
+    def system_session(
+        self, tenant_id: str | None = None
+    ) -> AbstractAsyncContextManager[AsyncConnection]:
+        """A session through the model's :class:`~pgbulkhead.GucBypass`; never admitted, since
+        it belongs to no one workspace. See :func:`pgbulkhead.system_session`."""
+        return system_session(self.get_engine(), self.model, tenant_id)