Connection & Registry¶

@dataclass(slots=True)
class Session:
    """A unit of work: one connection scope, one identity map, one settings set.

    `settings` are *session settings* — Postgres settings (GUCs) that Ferro
    applies to every operation this session runs, so queries are scoped without
    a per-query `where`:

        async with engines.session(settings={"myapp.tenant_id": "acme"}):
            # Postgres sees myapp.tenant_id = 'acme' for every statement this
            # session sends, which is what a row-level-security policy reads
            # back — inside a transaction() block or, as here, outside one.
            invoices = await Invoice.where(lambda invoice: invoice.paid).all()

    They are validated when the `Session` is constructed and again when it is
    entered, which is when the effective set — this session's settings merged
    over anything it inherits from an enclosing session — is snapshotted. See
    `EngineManager.session` for the full contract, including what an operation
    outside a transaction sends and how inherited settings behave off Postgres.

    A value not known until partway through the session's life — resolved
    from an auth chain, say — is set with `set_config` rather than at open;
    see there for the mid-transaction and identity-map-eviction contract.
    """

    connection_name: str | None = None
    session_id: str | None = None
    settings: Mapping[str, str] | None = None
    effective_settings: dict[str, str] = field(
        default_factory=dict, init=False, repr=False, compare=False
    )
    _token: Any = field(default=None, repr=False, compare=False)
    _enter_context: contextvars.Context | None = field(
        default=None, repr=False, compare=False
    )
    _enter_task: asyncio.Task[Any] | None = field(
        default=None, repr=False, compare=False
    )
    _close_lock: asyncio.Lock | None = field(default=None, repr=False, compare=False)

    def __post_init__(self) -> None:
        # Eager validation: a bad key or value raises where it was written, not
        # on the first query that would have carried it. `settings` is left
        # exactly as passed and checked again at enter, which is the moment its
        # values are actually read.
        _validated_settings(self.settings)

    async def __aenter__(self) -> "Session":
        # Re-validate rather than trusting the construction-time check: the
        # caller may have reassigned `settings` or mutated their mapping since,
        # and what gets applied has to be what was checked.
        declared = _validated_settings(self.settings)
        self.effective_settings = self._merge_ambient_settings(declared)
        self.session_id, resolved_name = _core_open_session(
            self.connection_name,
            list(self.effective_settings.items()),
            list(declared.items()),
        )
        if self.connection_name is None:
            self.connection_name = resolved_name
        self._token = _CURRENT_SESSION.set(self)
        self._enter_context = contextvars.copy_context()
        self._enter_task = asyncio.current_task()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        try:
            await self.close()
        except Exception as close_exc:
            if exc_type is None:
                raise
            if exc is not None:
                raise exc from close_exc
            raise close_exc

    async def close(self) -> None:
        """Close this session and release its runtime state.

        Under ``settings_delivery="connection"`` this is also where the
        connection the session pinned goes back to the pool — after Ferro
        resets **exactly the settings keys this session set on it**, and
        nothing else. (Never ``RESET ALL``: that would also wipe the
        ``search_path`` the pool installs on every connection.) A session
        under the default ``transaction`` delivery has nothing to release and
        this costs no round trip.

        Safe to call from a different asyncio context than ``__aenter__``.
        Repeated calls are no-ops.

        Raises:
            RuntimeError: If the ambient session in this asyncio context does not
                match this handle (same-context lifecycle misuse), or if
                session-scoped transactions are still open, or if resetting a
                pinned connection failed — in which case the session is closed
                regardless and the connection was discarded rather than
                returned to the pool.
        """
        if self._close_lock is None:
            self._close_lock = asyncio.Lock()

        async with self._close_lock:
            if self.session_id is None and self._token is None:
                return

            self._assert_close_allowed()

            releasing = None
            if self.session_id is not None:
                # Rust takes the session out of its registry *here*, on the
                # call itself, and rejects the close outright — before any
                # awaitable exists — when transactions are still open. So an
                # exception from this line leaves the handle exactly as it
                # was: still open, still usable.
                releasing = _core_close_session(self.session_id)
                self.session_id = None

            try:
                if releasing is not None:
                    # Awaits a round trip only under `connection` settings
                    # delivery, where the pinned connection's settings are
                    # reset before it goes back to the pool.
                    await releasing
            finally:
                # Past that point the session is gone from the runtime
                # whatever happens, so the ambient session is restored
                # whatever happens: a failed reset must not leave the rest of
                # this context scoped to a session that no longer exists.
                self._detach_ambient()

    def _detach_ambient(self) -> None:
        """Stop being the ambient session for this asyncio context."""
        if self._token is None:
            self._enter_context = None
            self._enter_task = None
            return
        token = self._token
        self._token = None
        self._enter_context = None
        self._enter_task = None
        self._restore_ambient_session(token)

    def _assert_close_allowed(self) -> None:
        if self._token is None:
            return
        if asyncio.current_task() is not self._enter_task:
            return
        ambient = _CURRENT_SESSION.get()
        if ambient is self:
            return
        entered_ambient = (
            self._enter_context.get(_CURRENT_SESSION)
            if self._enter_context is not None
            else None
        )
        if entered_ambient is self and ambient is not None:
            raise RuntimeError(_SESSION_CLOSE_AMBIENT_MISMATCH)

    def _restore_ambient_session(self, token: Any) -> None:
        ambient = _CURRENT_SESSION.get()
        if ambient is not self:
            try:
                _CURRENT_SESSION.reset(token)
            except ValueError:
                return
            raise RuntimeError(_SESSION_CLOSE_AMBIENT_MISMATCH)
        try:
            _CURRENT_SESSION.reset(token)
        except ValueError:
            return

    def _merge_ambient_settings(self, declared: dict[str, str]) -> dict[str, str]:
        """Snapshot the settings this session runs with, at enter.

        A session opened inside another one starts from the outer session's
        effective settings and overrides them key by key, so helper code that
        opens its own session stays scoped:

            async with engines.session(settings={"myapp.tenant_id": "acme"}):
                async with engines.session(settings={"myapp.role": "auditor"}):
                    # myapp.tenant_id == 'acme' and myapp.role == 'auditor'
                    ...

        The merge is a snapshot: nothing propagates between the two sessions
        afterwards, and closing the inner one leaves the outer exactly as it was.

        Args:
            declared: This session's own validated settings.
        """
        declared = dict(declared)
        ambient = _CURRENT_SESSION.get()
        inherited = getattr(ambient, "effective_settings", None) if ambient else None
        if not inherited:
            return declared
        return {**inherited, **declared}

    async def set_config(self, key: str, value: str) -> None:
        """Mutate this session's settings while it is open.

        Everything you can `settings=` at open, you can also set later, for the
        case where the value isn't known until partway through a request — an
        auth chain that resolves the tenant only after checking a token, say:

            async with engines.session() as session:      # no settings yet
                tenant = await resolve_tenant_from_auth_header(request)
                await session.set_config("myapp.tenant_id", tenant)
                # every query from here on is scoped to `tenant`
                invoices = await Invoice.where(lambda invoice: invoice.paid).all()

        Deep in a call stack, reach the session through `ferro.current_session()`
        rather than threading it through every function signature:

            async def resolve_tenant_and_scope(request) -> None:
                tenant = await resolve_tenant_from_auth_header(request)
                await ferro.current_session().set_config("myapp.tenant_id", tenant)

        **When the new value takes effect, precisely.** Any operation or
        `transaction()` STARTED after `set_config` returns sees the new value
        — this holds for any task, not just the one that called `set_config`,
        because the change is committed before `set_config`'s `await`
        returns. An operation or transaction already in flight *when*
        `set_config` runs keeps running under the scope it started with —
        that is the existing per-operation-atomicity guarantee, working as
        intended, not something `set_config` reaches back into. The one place
        that needs (and gets) special handling is a `transaction()` block
        already open in the SAME task that calls `set_config` — there, the
        very next statement in that same transaction sees the new value
        immediately, rather than waiting for the transaction to end and a new
        one to begin (there is no race to resolve here: the two are
        sequential code in one task, by construction):

            async with engines.session() as session:
                async with transaction():
                    await session.set_config("myapp.tenant_id", "acme")
                    # this SELECT, in the SAME transaction, already sees it
                    rows = await Invoice.where(lambda invoice: invoice.paid).all()

        `set_config` also evicts this session's identity map on a real change:
        an instance `get()`-ed under the old scope is never handed back under
        the new one — the next `get()` for that primary key refetches instead.
        A call that doesn't actually change the value (the key already holds
        it) is a true no-op: no database round trip, no identity-map eviction.

        Validation matches `engines.session(settings=...)` exactly (see there
        for the full rationale): `key` and `value` must be `str`, and `key` must
        be a dotted custom setting name.

        Args:
            key: Dotted custom Postgres setting name (e.g. `"myapp.tenant_id"`).
            value: The setting's new value.

        Raises:
            RuntimeError: This session is not open (`set_config` mutates a live
                session, so there has to be one), or this session's connection
                is not Postgres — `set_config` is a declaration, exactly like
                `settings=` at open, so it has to be honourable rather than a
                silent no-op on a backend with no GUCs. On the rare failure
                while delivering the change into an already-open transaction,
                this session's recorded settings and identity map have
                already committed the change regardless (see the Ferro source
                for the exact sequencing) — the affected transaction cannot go
                on to leak an old-scope read, because Postgres aborts a
                transaction after a failed statement until it is rolled back.
            TypeError: `key` or `value` is not a `str`.
            ValueError: `key` is not a dotted custom setting name.
        """
        if self.session_id is None:
            raise RuntimeError(_SET_CONFIG_NOT_OPEN_MESSAGE)
        if self.connection_name is None:
            # `connection_name` is resolved by `__aenter__` alongside
            # `session_id` (see there), so a live `session_id` should make
            # this unreachable. Raised rather than asserted (`assert` strips
            # under `-O`) because this is a Ferro lifecycle invariant, not an
            # ordinary user mistake.
            raise RuntimeError(_SET_CONFIG_NOT_OPEN_MESSAGE)

        validated = _validated_settings({key: value})
        ((validated_key, validated_value),) = validated.items()

        # Only the one validated pair is sent — Rust merges it against the
        # session's own last-committed settings, inside a lock that
        # serializes the read and the write together, so two sibling tasks
        # setting different keys concurrently can never lose one to a
        # stale-mirror race (see `operations::set_session_config`). The
        # mirror below is always assigned from what Rust actually committed,
        # never computed here — both when this call changed something and
        # when it was a no-op.
        committed = await _core_set_session_config(
            self.session_id,
            self.connection_name,
            validated_key,
            validated_value,
        )
        self.effective_settings = dict(committed)

    def query(self, model_cls):
        from .query import Query

        return Query(model_cls, session=self)

class PoolConfig(BaseModel):
    """Connection pool settings for a named Ferro connection.

    ``settings_delivery`` chooses how sessions on this connection get their
    session settings (the Postgres GUCs row-level-security policies read) onto
    the database. The two modes send different SQL for the very same session:

        async with engines.session(settings={"myapp.tenant_id": "acme"}):
            invoices = await Invoice.where(lambda invoice: invoice.paid).all()

    ``"transaction"`` (the default) wraps the query in a transaction of its own
    and scopes the value to it::

        BEGIN
        SELECT set_config($1, $2, true)    -- 'myapp.tenant_id', 'acme'
        SELECT "invoice".* FROM "invoice" WHERE "paid"
        COMMIT

    ``"connection"`` sets the value once, on a connection it keeps for the
    session, and then sends your statements bare::

        SELECT set_config($1, $2, false), set_config($3, $4, false)
        -- 'myapp.tenant_id', 'acme', 'ferro.pinned_keys', 'myapp.tenant_id'
        SELECT "invoice".* FROM "invoice" WHERE "paid"
        ...                                -- every later query, no wrap
        SELECT set_config($1, NULL, false), set_config($2, NULL, false)
        -- at session close: exactly the keys above, reset

    The difference that matters is ``true`` vs ``false`` — Postgres'
    ``is_local`` flag. ``true`` makes the value die with the transaction, which
    is why the default is safe even behind a transaction pooler like PgBouncer,
    where consecutive statements can land on different backends. ``false``
    makes it live for the whole database session, which is only safe if that
    session belongs to one Ferro session and nobody else — hence the pinned
    connection.

    So ``settings_delivery="connection"`` is a promise about your deployment:
    **this pool talks to Postgres directly.** Ferro never guesses it. A pooler
    is invisible to its clients, and guessing wrong means one tenant's value
    answering another tenant's query.

    What you get for the promise: no per-operation wrap at all, so a
    settings-bearing session costs one extra round trip for its whole life
    instead of about two per operation outside ``transaction()``.

    What it costs, and there are four costs worth knowing before you opt in.

    **One connection per scoped session.** A settings-bearing session holds a
    pool connection from its first operation until it closes, so **no more
    settings-bearing sessions can run at once than the pool has
    connections**. Number ``max_connections`` + 1's first query waits for a
    connection, exactly like any other pool checkout. Size the pool for peak
    concurrent scoped sessions.

    **One session, one connection, therefore one thing at a time.** Everything
    a pinned session does is serialized: two sibling tasks sharing the session
    run one after the other, and while a ``transaction()`` block is open a
    sibling task's operation waits for it rather than running inside it.
    That is the honest meaning of a session that owns a single connection —
    the same fact as the cap above, seen from inside one session. Operations
    in the same task *inside* a ``transaction()`` block are unaffected; they
    already own the block.

    **A marker check on release while anyone is pinned.** The pool verifies
    that a connection coming back is not still carrying a session's settings.
    While no session is pinned this costs nothing at all. While *any* session
    on the pool is pinned, every connection release on that pool performs one
    marker check — including releases by settings-less sessions and by
    sessionless operations.

    **Schema changes wait for scoped sessions.** ``connect(migrate_updates=…)``
    and ``migrate()`` refresh the pool afterwards, and a refresh cannot finish
    until every connection comes back — including the pinned ones. Migrating
    while long-lived tenant sessions are open therefore blocks until they
    close; Ferro warns, naming how many it is waiting on. Run schema changes
    before opening tenant-scoped sessions.

    Sessions *without* settings never pin: same statements, same connections,
    no wrap, no serialization.

    Two smaller behaviours worth knowing:

    * A session that opens with no ``settings`` and gains them later via
      ``set_config`` has nothing pinned yet, so it pins on its **next
      operation** and applies the values then.
    * Closing a session resets its keys to the value the connection *started*
      with. For an ordinary custom setting that is the empty string, which is
      what fail-closed policies rely on. If an operator has set a default with
      ``ALTER ROLE ... SET myapp.tenant_id = ...`` or ``ALTER DATABASE ... SET
      ...``, resetting brings that value *back* rather than clearing it — so
      never configure a tenancy key as a role or database default.
    """

    model_config = ConfigDict(frozen=True)

    max_connections: int = PydanticField(default=5, ge=1)
    min_connections: int = PydanticField(default=0, ge=0)
    settings_delivery: Literal["transaction", "connection"] = "transaction"

    @model_validator(mode="after")
    def validate(self) -> "PoolConfig":
        if self.min_connections > self.max_connections:
            raise ValueError("min_connections cannot exceed max_connections")
        return self