Tell the session about SET, RESET and set_config

pg_mimic/__init__.py

@@ -17,7 +17,7 @@
 # describe is the machinery behind "column shape from a declared schema, without
 # executing" -- what TableSession answers Describe with, and what any session
 # declaring its own schema() should reach for rather than reimplement (#88).
-from . import catalog, describe, errors, middleware
+from . import catalog, describe, errors, middleware, settings_values
 from .arrays import ARRAY_OID
 from .auth import (
     AuthPlugin,
@@ -39,6 +39,7 @@
 from .results import ResultColumn
 from .server import PgServer
 from .session import BaseSession, Portal, Session, Statement, StaticStatement
+from .state import SettingValue
 
 # TableSession is headline public API, not a niche extension point: "serve these
 # tables" is the entry point most users want before they ever write a Session.
@@ -83,13 +84,15 @@
     "describe",
     "errors",
     "middleware",
+    "settings_values",
     "BaseSession",
     "Session",
     "TableSession",
     "Statement",
     "StaticStatement",
     "Portal",
     "ResultColumn",
+    "SettingValue",
     "ARRAY_OID",
     "BOOL",
     "BYTEA",

pg_mimic/middleware.py

@@ -27,6 +27,7 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
 import re
 from functools import partial
@@ -52,7 +53,7 @@
     PgError,
 )
 from .results import ResultColumn
-from .session import Statement, StaticStatement, statement_from_rows
+from .session import Session, Statement, StaticStatement, statement_from_rows
 from .typeinfo import is_typeinfo_query, typeinfo_statement
 from .types import TEXT
 
@@ -715,7 +716,7 @@ def on_execute() -> None:
 
 
 def _discard_statement(connection: Connection, sql: str, kind: str) -> Statement:
-    def on_execute() -> None:
+    async def on_execute() -> None:
         if kind != "ALL":
             return  # PLANS/SEQUENCES/TEMP: a mimic has none of those to throw away
         # Real Postgres refuses DISCARD ALL inside a transaction block (25001);
@@ -725,7 +726,7 @@ def on_execute() -> None:
         # _reset_all first, because it does what reset() cannot: it reports each
         # changed setting back to the client. reset() then clears the rest of the
         # DISCARD ALL surface (and the settings again, by then already empty).
-        _reset_all(connection)
+        await _reset_all(connection)
         connection.state.reset()
         # `DISCARD ALL` includes `UNLISTEN *` -- reset() cleared the set, and this
         # is what makes the server stop delivering to us.
@@ -841,7 +842,7 @@ def _assignment_statement(
     assignments: Sequence[tuple[str, str | None]],
     local: bool = False,
 ) -> Statement:
-    def on_execute() -> None:
+    async def on_execute() -> None:
         outside_a_transaction = local and connection.tx_status == b"I"
         if outside_a_transaction:
             # Real Postgres does nothing and warns while doing it -- the exact
@@ -860,26 +861,29 @@ def on_execute() -> None:
         if outside_a_transaction:
             return
         for name, value in assignments:
-            _apply_set_config(connection, name, value, local=local)
+            await _apply_set_config(connection, name, value, local=local)
 
     return StaticStatement(sql, None, [], on_execute)
 
 
 def _reset_all_statement(connection: Connection, sql: str) -> Statement:
-    def on_execute() -> None:
-        _reset_all(connection)
+    async def on_execute() -> None:
+        await _reset_all(connection)
 
     return StaticStatement(sql, None, [], on_execute)
 
 
-def _reset_all(connection: Connection) -> None:
-    """Drop every SET this connection has made. Reported settings are reported
-    again on the way out, at whatever value they revert to."""
-    names = list(connection.state.session_vars)
-    connection.state.session_vars.clear()
-    connection.state.committed_vars.clear()
-    for name in names:
-        _report_setting(connection, name)
+async def _reset_all(connection: Connection) -> None:
+    """Drop every SET this connection has made. Reported settings are reported again
+    on the way out, at whatever value they revert to.
+
+    One `_apply_set_config` per name rather than clearing the dicts, so a session
+    hears each reset as it would hear `RESET x` written out (#35). A session that
+    refuses one leaves the rest reset, which is what Postgres does with a failure
+    partway through -- there is no transaction around a RESET ALL.
+    """
+    for name in list(connection.state.session_vars):
+        await _apply_set_config(connection, name, None)
 
 
 # The GUC_REPORT settings: real Postgres sends a ParameterStatus whenever one of
@@ -1212,15 +1216,24 @@ def _substitute_session_functions(
     return expr, substitutions, _run_all(effects) if effects else None
 
 
-def _run_all(effects: list[Callable[[], None]]) -> Callable[[], None]:
-    def run() -> None:
+def _run_all(effects: list[Callable[[], Any]]) -> Callable[[], Awaitable[None]]:
+    """One effect out of many, awaiting the ones that need it.
+
+    A query may carry several -- `SELECT set_config('a','1',false), pg_notify(...)`
+    -- and they run in the order they were found. Portal.execute awaits what this
+    returns, so an effect reaching the session is no different here from one that
+    only writes connection state (#35)."""
+
+    async def run() -> None:
         for effect in effects:
-            effect()
+            result = effect()
+            if inspect.isawaitable(result):
+                await result
 
     return run
 
 
-def _apply_set_config(connection: Connection, key: str, value: str | None, local: bool = False) -> None:
+async def _apply_set_config(connection: Connection, key: str, value: str | None, local: bool = False) -> None:
     """Apply one setting change, from SET/RESET or from set_config(). None means
     "back to the built-in default", i.e. forget the override.
 
@@ -1250,15 +1263,62 @@ def _apply_set_config(connection: Connection, key: str, value: str | None, local
     # reads back `on` -- see _setting_value, which renders it again (#115). RESET
     # has no value to parse.
     stored = None if value is None else settings_values.parse(key, value)
+    undo = _record_setting(connection, key, stored, drop=value is None, local=local)
+    _report_setting(connection, key)
+    try:
+        await _tell_session(connection, key, value, stored)
+    except Exception:
+        # The session refused it. Put back what was there and say so again, so a
+        # client's ParameterStatus cache does not keep a value this connection no
+        # longer holds (#35).
+        undo()
+        _report_setting(connection, key)
+        raise
+
+
+def _record_setting(connection: Connection, key: str, stored: Any, *, drop: bool, local: bool) -> Callable[[], None]:
+    """Write one setting, and hand back what would put things as they were.
+
+    `SET LOCAL` writes only session_vars, which is what makes it local. The undo is
+    built before the write rather than reconstructed after, because only here knows
+    which dicts were touched.
+    """
+    targets = [connection.state.session_vars]
+    if not local:
+        targets.append(connection.state.committed_vars)
+    previous = [(target, key in target, target.get(key)) for target in targets]
+    was_known = key in connection.state.known_settings
+
     connection.state.known_settings.add(key)
-    for target in (connection.state.session_vars, connection.state.committed_vars):
-        if value is None:
+    for target in targets:
+        if drop:
             target.pop(key, None)
         else:
             target[key] = stored
-        if local:
-            break  # session_vars only
-    _report_setting(connection, key)
+
+    def undo() -> None:
+        for target, had, value in previous:
+            if had:
+                target[key] = value
+            else:
+                target.pop(key, None)
+        if not was_known:
+            connection.state.known_settings.discard(key)
+
+    return undo
+
+
+async def _tell_session(connection: Connection, key: str, raw: str | None, parsed: Any) -> None:
+    """Hand the change to the session, if there is one that takes it.
+
+    Both spellings go over, so a session need not know that this happens to record
+    before it tells -- see Session.set_parameter.
+
+    A bare BaseSession bypasses the middleware chain entirely and has no hook; a
+    Session has one that does nothing until it is overridden."""
+    session = getattr(connection, "session", None)
+    if isinstance(session, Session):
+        await session.set_parameter(key, raw, parsed)
 
 
 def _evaluate_select(

pg_mimic/session.py

@@ -26,7 +26,7 @@
 from typing import Any, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Sequence
 
 from .results import ResultColumn
-from .state import SessionState
+from .state import SessionState, SettingValue
 from .types import TEXT
 
 Row = tuple
@@ -302,6 +302,46 @@ async def describe(self, sql: str, param_oids: list[int | None]) -> list[ResultC
     async def query(self, sql: str, params: list[str | None]) -> RowSource:
         raise NotImplementedError
 
+    async def set_parameter(self, name: str, raw_value: str | None, parsed_value: SettingValue | None) -> None:
+        """Told about every SET, RESET and set_config() the middleware handled.
+
+        Both spellings, because a session wants different ones for different jobs:
+        `raw_value` is the text the client wrote, which is what to forward to a real
+        backend, and `parsed_value` is what it means, which is what to act on::
+
+            SET work_mem = '32MB'      raw='32MB'  parsed=32768   (its own unit, kB)
+            SET row_security = 'tr'    raw='tr'    parsed=True
+            SET client_encoding='utf8' raw='utf8'  parsed='UTF8'
+            SET app.tenant = 'acme'    raw='acme'  parsed='acme'  (a custom GUC has
+                                                                   no type to parse)
+            RESET work_mem             raw=None    parsed=None
+
+        Both are None for a RESET, which is the only thing None means here -- no
+        parameter parses to it.
+
+        Handed over rather than read off `self.state.session_vars`, which holds the
+        same parsed value by the time this runs. Depending on that would make the
+        order the connection happens to do things in part of this contract.
+
+        `pg_mimic.settings_values.parse` and `.render` are the same pair the
+        middleware used to produce it, for a session that needs to read a value this
+        did not hand it -- one a real backend reported back, say.
+
+        `name` is lowercased. Dotted custom GUCs (`app.tenant_id`) arrive here too,
+        which is the point: the alternative was re-parsing raw SQL out of `query()`
+        to discover it was connection boilerplate at all.
+
+        The connection has already recorded the change and will send any
+        ParameterStatus it owes. Raising `PgError` rejects the setting and undoes
+        that record, so a session may refuse one it cannot honour::
+
+            async def set_parameter(self, name, raw_value, parsed_value):
+                if name == "app.tenant_id" and not self.may_use(raw_value):
+                    raise PgError(INVALID_PARAMETER_VALUE, f"no such tenant: {raw_value}")
+
+        Does nothing by default.
+        """
+
     async def schema(self) -> dict | None:
         """Optional: describe your tables for information_schema/pg_catalog
         emulation. See pg_mimic.catalog."""