From cb254ba8554adec490c523215133228d6722b59a Mon Sep 17 00:00:00 2001 From: Alexey Shalaev <75322386+AlexeyShalaev@users.noreply.github.com> Date: Mon, 7 Sep 2026 08:52:58 +0300 Subject: [PATCH] feat: stop sending jit and search_path as startup parameters `BasePostgresConfig.jit` defaults to None and is put into asyncpg `server_settings` only when set explicitly. `db_schema` no longer goes as a `search_path` startup parameter: `create_async_session_manager` hands it to the manager, which attaches a `begin` listener running `set_config('search_path', ..., true)` -- SET LOCAL -- as the first statement of every transaction. That is the one scope PgBouncer in transaction mode honours: it refused every connection on the default `jit` (`unsupported startup parameter: jit`) and, with `ignore_startup_parameters`, silently dropped the schema. New: `AsyncSessionManager(search_path=...)`, `AsyncSessionManagerBuilder.with_search_path()`, `session.manager.attach_search_path(engine, search_path)`. Two defaults change for direct-to-PostgreSQL users: JIT follows the server setting unless `jit` is set, and the schema is per transaction rather than per connection (a statement under AUTOCOMMIT is not covered). The configuration guide's PgBouncer section is rewritten around the measurements in the issue; the agents page follows. --- README.md | 4 +- docs/agents.md | 85 +++++++--- docs/guide/advanced.md | 31 ++-- docs/guide/configuration.md | 120 +++++++++++--- docs/index.md | 2 +- sqlalchemy_foundation_kit/config/postgres.py | 6 +- .../contrib/settings/postgres.py | 17 +- sqlalchemy_foundation_kit/session/builder.py | 18 +++ .../session/factories.py | 15 +- sqlalchemy_foundation_kit/session/manager.py | 41 ++++- tests/integration/conftest.py | 26 +++ tests/integration/test_session_integration.py | 153 +++++++++++++++++- tests/integration/test_uow_integration.py | 98 +++++++++++ tests/unit/contrib/settings/test_postgres.py | 7 +- tests/unit/session/test_builder.py | 34 ++++ tests/unit/session/test_factories.py | 87 ++++++++-- tests/unit/session/test_manager.py | 59 +++++++ 17 files changed, 707 insertions(+), 96 deletions(-) diff --git a/README.md b/README.md index 919dd90..f0c5535 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ Only `sqlalchemy[asyncio]`, `pydantic` and `asyncpg` are required by default — ✅ **Single dependency** — All foundation pieces in one place ✅ **Unit of Work pattern** — Transactional consistency with automatic commit/rollback ✅ **Connection pool management** — `AsyncSessionManager` with metrics and health checks -✅ **pgbouncer compatible** — Custom connection class for transaction mode +✅ **PgBouncer compatible** — A startup packet PgBouncer accepts, `search_path` per transaction, unique statement names ✅ **Observability built-in** — Prometheus metrics + OpenTelemetry tracing ✅ **Type-safe configuration** — Pydantic settings with validation ✅ **Base ORM models** — Pre-configured `Base` with naming conventions and mixins @@ -147,7 +147,7 @@ class PostgresConfig: application_name: str = "my-service" db_schema: str | None = None use_orjson_serialization: bool = True - jit: str | None = "off" + jit: str | None = None def to_dsn(self) -> str: return f"postgresql+asyncpg://{self.connection.user}@{self.connection.host}:{self.connection.port}/{self.connection.database}" diff --git a/docs/agents.md b/docs/agents.md index fc3097c..9bb3467 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -11,7 +11,7 @@ | Install | `pip install sqlalchemy-foundation-kit` · extras: `settings`, `metrics`, `orjson`, `dishka`, `dependency-injector`, `telemetry`, `all` | | Async | the whole library. `AsyncEngine`, `AsyncSession`, `asyncpg` | | Sync | none. There is no sync mirror and no sync entry point | -| Version | everything below was read from the source this site was built from. Three calls described here do not work on 0.2.0 — see [fixed since 0.2.0](#fixed-since-020) | +| Version | everything below was read from the source this site was built from. Some calls described here do not work on 0.2.0 or 0.2.1 — see [fixed since 0.2.0](#fixed-since-020) | | Source | | ## How to read this page @@ -27,8 +27,8 @@ docstrings in the source. Top to bottom before writing code. [Rules that hold or break the code](#rules-that-hold-or-break-the-code) is the section correctness lives in, and it is followed by the short list of calls that -are [fixed since 0.2.0](#fixed-since-020) — if the installed version is 0.2.0, those -raise before they reach the database. Every name used below is in the public API; if you +are [fixed since 0.2.0](#fixed-since-020) — on the version each row names, those fail +before they reach the database. Every name used below is in the public API; if you need something not listed here, fetch the page the [documentation map](#documentation-map) points at rather than guessing a method that sounds plausible. @@ -137,8 +137,9 @@ asyncio.run(main()) ``` With the `settings` extra, `create_async_session_manager(config)` builds the same manager -from a `PostgresSettingsProtocol` and fills in the pgbouncer-safe defaults -(`AsyncCConnection`, `application_name`, `search_path`, `jit`, statement caches) for you. +from a `PostgresSettingsProtocol` and fills in the PgBouncer-safe defaults for you: +`AsyncCConnection`, both statement caches at 0, `application_name` as the only startup +parameter, and `db_schema` applied per transaction as `search_path` (rule 16). ## The API @@ -162,8 +163,11 @@ from the submodule named beside it further down. `AsyncSessionManager(url, echo=False, poolclass="null", session_class=None, expire_on_commit=False, connect_args=None, isolation_level=None, pool_settings=None, -use_orjson=False, metrics=None, on_engine_created=None, dispose_timeout=30.0, **kwargs)` — -`**kwargs` reach `create_async_engine`. +use_orjson=False, metrics=None, on_engine_created=None, dispose_timeout=30.0, +search_path=None, **kwargs)` — `**kwargs` reach `create_async_engine`. `search_path` +attaches a `begin` listener to the engine that runs `set_config('search_path', …, true)` — +`SET LOCAL` — as the first statement of every transaction, autobegun ones included; +nothing is attached when it is `None`. | Member | What it does | |---|---| @@ -182,21 +186,23 @@ After `aclose()`, `get_session()` and `get_transaction()` raise `.with_session_class(cls)`, `.with_expire_on_commit(bool)`, `.with_connect_args(**kw)`, `.with_isolation_level(str)`, `.with_metrics(m)`, `.with_callbacks(on_engine_created=fn)`, `.with_json_serialization(orjson=True)`, `.with_extra_kwargs(**kw)`, -`.with_dispose_timeout(float)`, `.build()`. A builder is reusable: `build()` does not -consume it. +`.with_dispose_timeout(float)`, `.with_search_path(str)`, `.build()`. A builder is +reusable: `build()` does not consume it. `create_async_session_manager(postgres_config, application_name=None, metrics=None, on_engine_created=None, connection_class=None, extra_server_settings=None, extra_connect_args=None, **kwargs)` takes a `PostgresSettingsProtocol` and returns a -manager. It defaults `connection_class` to `AsyncCConnection`, sets `server_settings` from -`application_name`, `jit` and `db_schema` (as `search_path`), and passes both statement -cache sizes through. Your keys in `extra_server_settings` / `extra_connect_args` win over -the library's. - -`attach_metrics(engine, metrics)` — in `sqlalchemy_foundation_kit.session.manager`, not -re-exported — wires the pool listeners onto an engine you built yourself. The manager -calls it when `metrics` is passed. A raising metrics callback is logged and swallowed, -never propagated. +manager. It defaults `connection_class` to `AsyncCConnection`, sends `application_name` — +and `jit`, only when the config sets it — as `server_settings`, hands `db_schema` to the +manager as `search_path`, and passes both statement cache sizes through. Your keys in +`extra_server_settings` / `extra_connect_args` win over the library's; `server_settings` +are startup parameters, so read rule 16 before adding one. + +`attach_metrics(engine, metrics)` and `attach_search_path(engine, search_path)` — in +`sqlalchemy_foundation_kit.session.manager`, not re-exported — wire the pool listeners, +respectively the per-transaction `search_path`, onto an engine you built yourself. The +manager calls them when `metrics` / `search_path` is passed. A raising metrics callback is +logged and swallowed, never propagated. ### The unit of work @@ -293,7 +299,7 @@ With the `settings` extra, `sqlalchemy_foundation_kit.contrib.settings` implemen | `ConnectionSettings` | `host="localhost"`, `port=5432` (1–65535), `user="postgres"`, `password: SecretStr` **required**, `database: str` **required** | | `PoolSettings` | `kind="async_adapted_queue"`, `size=10`, `max_overflow=20`, `pre_ping=True`, `recycle=3600`, `timeout=30.0`. `kind="static"` with `max_overflow > 0` raises | | `QuerySettings` | `echo=False`, `statement_cache_size=0`, `prepared_statement_cache_size=0`, `isolation_level=None` | -| `BasePostgresConfig` | `connection` **required**, `pool`, `query`, `application_name: str` **required**, `db_schema=None`, `use_orjson_serialization=True`, `jit="off"`, `metrics_enabled=False`. `to_dsn(driver="asyncpg", mask_password=False)`; `__repr__` prints the masked DSN | +| `BasePostgresConfig` | `connection` **required**, `pool`, `query`, `application_name: str` **required**, `db_schema=None`, `use_orjson_serialization=True`, `jit=None`, `metrics_enabled=False`. `to_dsn(driver="asyncpg", mask_password=False)`; `__repr__` prints the masked DSN. `db_schema` is a `search_path` value applied per transaction (rule 16); `jit` is a startup parameter and is sent only when set | | `BasePostgresMigrationsConfig` | `postgres: BasePostgresConfig`, with `env_nested_delimiter="__"` and `extra="ignore"` | `BasePostgresConfig` declares no `model_config` of its own, so it reads environment @@ -391,17 +397,34 @@ decorator. metrics backend is logged at exception level and discarded, and the query proceeds. 15. **Close the manager.** `await manager.aclose()` disposes the engine under a shield and a `dispose_timeout` (30s). Skipping it leaks connections; calling it twice is fine. +16. **Through PgBouncer in transaction mode, only `SET LOCAL` and the server's own settings + hold.** asyncpg `server_settings` are startup parameters. PgBouncer forwards the ones it + tracks (`client_encoding`, `datestyle`, `timezone`, `standard_conforming_strings`, + `application_name`, plus `track_extra_parameters`), refuses the connection on any other + (`unsupported startup parameter: jit`), and with `ignore_startup_parameters` drops them + silently. So `jit` defaults to `None` and is sent only when set, and `db_schema` is + applied with `SET LOCAL search_path` as the first statement of every transaction — each + `transaction()`, `query()`, `managed_session()`, `get_session()`, `get_transaction()` + and `engine.connect()` block, and the transaction after a `commit()` in the same + session. Not covered: a statement under `isolation_level="AUTOCOMMIT"`, which begins no + transaction. To have the schema without the round-trip use `ALTER ROLE … SET + search_path`, `ALTER DATABASE … SET`, or schema-qualified metadata. A plain `SET` on a + connection leaks to the next client through a pooler; never issue one in a `connect` + listener. Statement caches stay at 0 and `AsyncCConnection` stays, for PgBouncer before + 1.22 (`max_prepared_statements=0`); 1.22+ tracks prepared statements itself. ### Fixed since 0.2.0 -Three published entry points raise before they reach the database on 0.2.0. They work on -current versions; the workaround column is what to do if the installed version is 0.2.0. +Three published entry points raise before they reach the database on 0.2.0, and on 0.2.1 +the PgBouncer-safe defaults cannot connect through PgBouncer. All work on current versions; +the workaround column is what to do on the version the row names. -| Call | What 0.2.0 does | Workaround on 0.2.0 | +| Call | What it does on that version | Workaround there | |---|---|---| -| `manager.get_transaction()`, with or without `isolation_level` | `TypeError: Session.__init__() got an unexpected keyword argument 'execution_options'` — the argument was passed to the session factory unconditionally | `async with manager.get_session() as s, s.begin():`, or the unit of work | -| `uow.transaction(isolation_level=…)`, `uow.managed_session(isolation_level=…)`, `uow.query(isolation_level=…)` | `InvalidRequestError: This connection has already initialized a SQLAlchemy Transaction()… isolation_level may not be altered` — the level was applied after the connection had autobegun | set the level on the engine: `AsyncSessionManager(..., isolation_level="SERIALIZABLE")` or `QuerySettings(isolation_level=...)` | -| `import sqlalchemy_foundation_kit.contrib.di` (or `.contrib.dependency_injector`) without its extra | `AttributeError: 'NoneType' object has no attribute 'APP'` (resp. `'DeclarativeContainer'`) instead of the intended `ImportError` | install the extra; the message is not the one the code meant to give you | +| `manager.get_transaction()`, with or without `isolation_level` (0.2.0) | `TypeError: Session.__init__() got an unexpected keyword argument 'execution_options'` — the argument was passed to the session factory unconditionally | `async with manager.get_session() as s, s.begin():`, or the unit of work | +| `uow.transaction(isolation_level=…)`, `uow.managed_session(isolation_level=…)`, `uow.query(isolation_level=…)` (0.2.0) | `InvalidRequestError: This connection has already initialized a SQLAlchemy Transaction()… isolation_level may not be altered` — the level was applied after the connection had autobegun | set the level on the engine: `AsyncSessionManager(..., isolation_level="SERIALIZABLE")` or `QuerySettings(isolation_level=...)` | +| `import sqlalchemy_foundation_kit.contrib.di` (or `.contrib.dependency_injector`) without its extra (0.2.0) | `AttributeError: 'NoneType' object has no attribute 'APP'` (resp. `'DeclarativeContainer'`) instead of the intended `ImportError` | install the extra; the message is not the one the code meant to give you | +| `create_async_session_manager(config)` through PgBouncer in transaction mode (0.2.1 and earlier) | `jit="off"` was the default and `db_schema` went as `search_path`, both as startup parameters: a default PgBouncer refuses every connection with `ProtocolViolationError: unsupported startup parameter: jit`; with `ignore_startup_parameters=jit,search_path` it connects and silently drops both, so every query lands in `public` | `jit=None`, `db_schema=None`, and `ALTER ROLE … SET search_path` on the server | `IsolationLevel` itself was always fine — `READ_UNCOMMITTED`, `READ_COMMITTED`, `REPEATABLE_READ`, `SERIALIZABLE`, whose values are the PostgreSQL spellings with spaces — @@ -498,6 +521,17 @@ manager = AsyncSessionManager("postgresql+asyncpg://…") manager = AsyncSessionManager("postgresql+asyncpg://…", poolclass="async_adapted_queue") ``` +```python +# WRONG — both are startup parameters; PgBouncer in transaction mode refuses the connection, +# or with ignore_startup_parameters drops them and every query lands in public +config = BasePostgresConfig(..., jit="off") # "for pgbouncer" +manager = create_async_session_manager(config, extra_server_settings={"search_path": "app"}) + +# RIGHT — send nothing PgBouncer will not carry; the schema is applied per transaction +config = BasePostgresConfig(..., db_schema="app") # jit stays None +manager = create_async_session_manager(config) +``` + ## Errors The library defines no exception classes of its own. It raises the standard ones and lets @@ -513,6 +547,7 @@ SQLAlchemy's through untouched. | `TypeError` | a value orjson cannot serialize; `manager.get_transaction()` on 0.2.0 (see above) | | `sqlalchemy.exc.IllegalStateChangeError` | one session driven by two tasks at once (rule 6) | | `sqlalchemy.exc.InvalidRequestError` | `isolation_level` on a unit-of-work method on 0.2.0 (see above) | +| `asyncpg.exceptions.ProtocolViolationError` | `unsupported startup parameter: …` — a `server_settings` key PgBouncer does not track (rule 16); on 0.2.1 the default `jit` did this on every connection | | `sqlalchemy.exc.IntegrityError`, `OperationalError`, … | the database refused the statement. Inside `tx.savepoint()` these are re-raised with the surrounding transaction still usable; anywhere else they roll the whole block back | ## Documentation map diff --git a/docs/guide/advanced.md b/docs/guide/advanced.md index 57b2d4e..aea1886 100644 --- a/docs/guide/advanced.md +++ b/docs/guide/advanced.md @@ -605,30 +605,35 @@ async def wait_for_database(session_manager: AsyncSessionManager) -> None: await wait_for_database(session_manager) ``` -### pgbouncer Compatibility +### PgBouncer Compatibility -For pgbouncer transaction mode, disable prepared statements: +`create_async_session_manager` is safe behind PgBouncer in transaction mode as it stands: +both statement caches default to 0 and the connection class is `AsyncCConnection`, whose +statement names are unique per connection (both for PgBouncer before 1.22); `jit` is not +sent unless you set it; and `db_schema` is applied with `SET LOCAL` at the start of every +transaction rather than as a startup parameter. The one thing left to size is the pool: +PgBouncer owns the server connections, so `size` is how many client connections this +process may hold open. ```python from sqlalchemy_foundation_kit.contrib.settings import ( BasePostgresConfig, - QuerySettings, + ConnectionSettings, + PoolSettings, ) config = BasePostgresConfig( - connection=..., - pool=PoolSettings( - size=20, # pgbouncer manages actual connections - max_overflow=0, # No overflow with pgbouncer - ), - query=QuerySettings( - statement_cache_size=0, # Required for pgbouncer - prepared_statement_cache_size=0, # Required for pgbouncer - ), - jit="off", # Required for pgbouncer + connection=ConnectionSettings(host="pgbouncer", port=6432, ...), + pool=PoolSettings(size=20, max_overflow=0), + application_name="my-service", + db_schema="app", ) ``` +Do not set `jit="off"` for PgBouncer. It travels as a startup parameter, which PgBouncer in +transaction mode rejects (`unsupported startup parameter: jit`), so every connection fails. +The measurements and the reasoning are in [Configuration → PgBouncer](configuration.md#pgbouncer). + ## Custom Types ### PydanticJSONB diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index 9e018b9..e92988c 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -54,7 +54,7 @@ class Settings(BaseSettings): application_name="my-service", db_schema=None, use_orjson_serialization=True, - jit="off", + jit=None, metrics_enabled=True, ) @@ -103,7 +103,7 @@ POSTGRES__QUERY__ISOLATION_LEVEL="READ COMMITTED" POSTGRES__APPLICATION_NAME=my-service POSTGRES__DB_SCHEMA=public POSTGRES__USE_ORJSON_SERIALIZATION=true -POSTGRES__JIT=off +POSTGRES__JIT=off # a startup parameter: direct connections only, see PgBouncer below POSTGRES__METRICS_ENABLED=true ``` @@ -284,9 +284,9 @@ QuerySettings( | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `application_name` | `str` | **required** | Application identifier in PostgreSQL logs | -| `db_schema` | `str \| None` | `None` | Default PostgreSQL schema name | +| `db_schema` | `str \| None` | `None` | `search_path` applied to every transaction with `SET LOCAL` — a schema, or a comma-separated list | | `use_orjson_serialization` | `bool` | `True` | Use `orjson` for JSON serialization (requires `[orjson]` extra) | -| `jit` | `"off" \| "on" \| None` | `"off"` | PostgreSQL JIT compilation setting | +| `jit` | `"off" \| "on" \| None` | `None` | PostgreSQL JIT setting, sent as a startup parameter only when set | | `metrics_enabled` | `bool` | `False` | Enable Prometheus metrics (requires `[metrics]` extra) | **Application Name:** @@ -317,7 +317,23 @@ BasePostgresConfig( ) ``` -Sets the default `search_path` for all connections. +The value is a `search_path`, so `"tenant_123, public"` is valid too. It is applied at the +start of every transaction with `SET LOCAL` semantics — `set_config('search_path', …, true)` +as the first statement after `BEGIN` — which is the one scope that survives a +transaction-mode pooler. Every `transaction()`, `query()`, `managed_session()`, +`get_session()`, `get_transaction()` and raw `engine.connect()` block sees it, including +the transaction that follows a `commit()` in the same session. Two things it is not: a +startup parameter, which PgBouncer rejects or drops (see [PgBouncer](#pgbouncer)), and a +session-level `SET`, which leaks between clients through a pooler. A statement run with +`isolation_level="AUTOCOMMIT"` begins no transaction and therefore runs with the server's +default `search_path`. + +If you would rather not pay one round-trip per transaction, set it on the server instead — +`ALTER ROLE app_user SET search_path = tenant_123` or `ALTER DATABASE mydb SET search_path = +tenant_123` — and leave `db_schema` unset; or qualify the schema in your metadata +(`__table_args__ = {"schema": "tenant_123"}`) and need no `search_path` at all. Outside +`create_async_session_manager` the same mechanism is `AsyncSessionManager(..., +search_path="tenant_123")` or `AsyncSessionManagerBuilder(url).with_search_path("tenant_123")`. **orjson Serialization:** @@ -338,18 +354,17 @@ Automatically used by `PydanticJSONB` type for better performance. **JIT (Just-In-Time Compilation):** -PostgreSQL 11+ includes JIT compilation for complex queries. Disable when using pgbouncer in transaction mode: +`jit` is `None` by default: nothing is sent and PostgreSQL's own setting applies (`on` +since PostgreSQL 12). Set `"off"` or `"on"` only to override the server for this +application, and know that the value travels as a startup parameter — fine on a direct +connection, rejected by PgBouncer in transaction mode unless it is listed in +`track_extra_parameters` (see [PgBouncer](#pgbouncer)). Behind a pooler, set it on the +server instead: `ALTER ROLE app_user SET jit = off`. ```python -# pgbouncer transaction mode +# Direct PostgreSQL connection: override the server setting for this application BasePostgresConfig( - jit="off", # Required for pgbouncer - # ... -) - -# Direct PostgreSQL connection -BasePostgresConfig( - jit="on", # Can improve performance for complex queries + jit="off", # ... ) ``` @@ -400,7 +415,7 @@ class MyPostgresConfig: application_name: str = "my-service" db_schema: str | None = None use_orjson_serialization: bool = True - jit: str | None = "off" + jit: str | None = None metrics_enabled: bool = False def to_dsn(self, driver: str | None = "asyncpg", mask_password: bool = False) -> str: @@ -410,6 +425,69 @@ class MyPostgresConfig: return f"{scheme}://{conn.user}:{password}@{conn.host}:{conn.port}/{conn.database}" ``` +## PgBouncer + +Measured on PostgreSQL 17 and PgBouncer 1.25 in transaction mode, PgBouncer at its +defaults unless the row says otherwise: twenty clients, twenty distinct statements each, +run twice, so every cache fills. The scripts are on +[issue #23](https://github.com/bedrock-python/sqlalchemy-foundation-kit/issues/23). + +| Client | Through | Result | +|---|---|---| +| plain SQLAlchemy + asyncpg, driver defaults | PostgreSQL directly | 800 of 800 ok | +| plain SQLAlchemy + asyncpg, driver defaults | PgBouncer, defaults (`max_prepared_statements=200`) | 800 of 800 ok | +| plain SQLAlchemy + asyncpg, driver defaults | PgBouncer, `max_prepared_statements=0` (the default before 1.22) | 485 of 800 ok; 315 × `prepared statement "__asyncpg_stmt_…" does not exist` | +| `create_async_session_manager(config)`, 0.2.1 | PgBouncer, defaults | 0 of 800: every connection refused with `unsupported startup parameter: jit` | +| `create_async_session_manager(config)`, 0.2.1 | PgBouncer, `ignore_startup_parameters=jit,search_path` | 800 of 800 ok — and `SHOW jit` is `on`, `SHOW search_path` is `"$user", public`: both parameters silently dropped | +| `create_async_session_manager(config)`, current | PgBouncer, defaults | 800 of 800 ok; `SHOW search_path` is the configured `db_schema` | + +Three things follow, and the library is built around each. + +**Startup parameters do not pass through.** asyncpg's `server_settings` are parameters of +the PostgreSQL startup packet. PgBouncer forwards only the ones it tracks — +`client_encoding`, `datestyle`, `timezone`, `standard_conforming_strings`, +`application_name`, plus whatever you list in `track_extra_parameters` (1.18+) — and +refuses the connection on any other. `ignore_startup_parameters` makes the connection +succeed by throwing the parameter away, which is worse: a `search_path` the library +accepted and the database never saw. So `create_async_session_manager` sends +`application_name` and nothing else by default; `jit` goes only when you set it, and +`db_schema` never goes as a startup parameter. + +**Session state does not survive a transaction.** In transaction mode a client owns a +server connection from `BEGIN` to `COMMIT` and not a moment longer. A plain `SET` outside +a transaction lands on whichever server connection was free and is visible to the next +client handed that connection — the last row of the lab: client A ran +`SET search_path TO leaked`, client B read `search_path = 'leaked'`. The only per-connection +state that behaves is state scoped to the transaction (`SET LOCAL`) or set on the server +for the role or database. That is why `db_schema` is applied as `SET LOCAL search_path` at +the start of every transaction, and why the alternatives are `ALTER ROLE … SET`, +`ALTER DATABASE … SET`, `track_extra_parameters`, or schema-qualified metadata — never a +`SET` in a `connect` listener. + +**Prepared statements are fine on a current PgBouncer, and still not on an old one.** +Since 1.22 PgBouncer tracks protocol-level prepared statements itself +(`max_prepared_statements=200` by default), so the driver's default statement cache runs +clean. Before 1.22, or with `max_prepared_statements=0`, a cached statement name is unknown +to the server connection the next transaction lands on and the classic failure comes back. +`QuerySettings` defaults both caches to 0 and `create_async_session_manager` uses +`AsyncCConnection`, whose statement names are unique per connection, so the library is +safe on either; on 1.22+ you may raise the caches to win the round-trip back. + +```python +# Behind PgBouncer in transaction mode this is all it takes +config = BasePostgresConfig( + connection=ConnectionSettings(host="pgbouncer", port=6432, ...), + application_name="my-service", + db_schema="app", # SET LOCAL at the start of every transaction: arrives + # jit stays None # nothing sent: the connection is accepted +) +manager = create_async_session_manager(config) +``` + +What not to do: `jit="off"` "for pgbouncer" (it is the one setting that makes the +connection impossible), `extra_server_settings={"search_path": ...}` (rejected, or dropped), +and a `SET search_path` in an `on_engine_created` connect listener (leaks between clients). + ## Configuration Best Practices ### 1. Use Environment Variables @@ -467,15 +545,11 @@ PoolSettings(size=5, max_overflow=10) PoolSettings(size=2, max_overflow=0) ``` -### 5. Disable Caching for pgbouncer +### 5. Behind PgBouncer, Send Nothing It Will Not Carry -```python -# pgbouncer transaction mode -QuerySettings( - statement_cache_size=0, # Required - prepared_statement_cache_size=0, # Required -) -``` +The defaults already do the right thing — statement caches at 0, `jit` unset, `db_schema` +applied per transaction. Do not put `jit`, `search_path` or any other untracked parameter +into `extra_server_settings`; see [PgBouncer](#pgbouncer). ### 6. Enable Metrics in Production diff --git a/docs/index.md b/docs/index.md index 773f07e..2f66dfd 100644 --- a/docs/index.md +++ b/docs/index.md @@ -18,7 +18,7 @@ Only `sqlalchemy[asyncio]`, `pydantic` and `asyncpg` are required by default — ✅ **Single dependency** — All foundation pieces in one place ✅ **Unit of Work pattern** — Transactional consistency with automatic commit/rollback ✅ **Connection pool management** — `AsyncSessionManager` with metrics and health checks -✅ **pgbouncer compatible** — Custom connection class for transaction mode +✅ **PgBouncer compatible** — A startup packet PgBouncer accepts, `search_path` per transaction, unique statement names ✅ **Observability built-in** — Prometheus metrics + OpenTelemetry tracing ✅ **Type-safe configuration** — Pydantic settings with validation ✅ **Base ORM models** — Pre-configured `Base` with naming conventions and mixins diff --git a/sqlalchemy_foundation_kit/config/postgres.py b/sqlalchemy_foundation_kit/config/postgres.py index 81db857..caf21ca 100644 --- a/sqlalchemy_foundation_kit/config/postgres.py +++ b/sqlalchemy_foundation_kit/config/postgres.py @@ -131,9 +131,9 @@ class PostgresSettingsProtocol(Protocol): pool: Connection pool settings. query: Query execution and transaction settings. application_name: Application identifier for connections. - db_schema: Optional PostgreSQL schema name. + db_schema: Optional PostgreSQL ``search_path``, applied to every transaction. use_orjson_serialization: Enable orjson for JSON operations. - jit: JIT compilation setting (PgBouncer compatibility). + jit: JIT compilation setting, sent as a startup parameter only when not ``None``. Examples: Implementing the protocol: @@ -144,7 +144,7 @@ class PostgresSettingsProtocol(Protocol): ... application_name: str = "my-app" ... db_schema: str | None = None ... use_orjson_serialization: bool = True - ... jit: str | None = "off" + ... jit: str | None = None ... ... def to_dsn(self) -> str: ... return f"postgresql://{self.connection.user}@{self.connection.host}..." diff --git a/sqlalchemy_foundation_kit/contrib/settings/postgres.py b/sqlalchemy_foundation_kit/contrib/settings/postgres.py index 532594a..f92673c 100644 --- a/sqlalchemy_foundation_kit/contrib/settings/postgres.py +++ b/sqlalchemy_foundation_kit/contrib/settings/postgres.py @@ -98,9 +98,12 @@ class BasePostgresConfig(BaseSettings): pool: Connection pool configuration (size, overflow, timeouts). query: Query execution settings (echo, caching, isolation level). application_name: Application name for connection identification. - db_schema: Optional PostgreSQL schema name. + db_schema: Optional PostgreSQL ``search_path`` — a schema, or a comma-separated + list — applied to every transaction with ``SET LOCAL`` semantics. use_orjson_serialization: Use orjson for JSON serialization (requires orjson). - jit: JIT compilation setting (off/on) for PgBouncer compatibility. + jit: JIT compilation setting (off/on), sent as a startup parameter only when set. + ``None`` (the default) sends nothing and leaves the server's own setting; a + transaction-mode PgBouncer rejects the parameter unless it tracks it. metrics_enabled: Enable connection pool metrics collection. Examples: @@ -125,12 +128,18 @@ class BasePostgresConfig(BaseSettings): # Top-level settings application_name: str = Field(description="Application name for PostgreSQL") - db_schema: str | None = Field(default=None, description="PostgreSQL schema name") + db_schema: str | None = Field( + default=None, + description="PostgreSQL search_path (a schema or a comma-separated list), applied to every transaction", + ) use_orjson_serialization: bool = Field( default=True, description="Use orjson for JSON serialization (requires orjson installed)", ) - jit: PostgresJit | None = Field(default="off", description="JIT setting (off/on)") + jit: PostgresJit | None = Field( + default=None, + description="JIT setting (off/on), sent as a startup parameter only when set", + ) metrics_enabled: bool = Field(default=False, description="Enable PostgreSQL metrics") def __repr__(self) -> str: diff --git a/sqlalchemy_foundation_kit/session/builder.py b/sqlalchemy_foundation_kit/session/builder.py index 42df709..80cefbd 100644 --- a/sqlalchemy_foundation_kit/session/builder.py +++ b/sqlalchemy_foundation_kit/session/builder.py @@ -77,6 +77,7 @@ def __init__(self, url: str) -> None: self._metrics: PostgresMetricsProtocol | None = None self._on_engine_created: Callable[[AsyncEngine], None] | None = None self._dispose_timeout: float | None = None + self._search_path: str | None = None self._extra_kwargs: dict[str, object] = {} def with_echo(self, echo: bool = True) -> AsyncSessionManagerBuilder[SessionT]: @@ -258,6 +259,22 @@ def with_dispose_timeout(self, timeout: float) -> AsyncSessionManagerBuilder[Ses self._dispose_timeout = timeout return self + def with_search_path(self, search_path: str) -> AsyncSessionManagerBuilder[SessionT]: + """Apply a PostgreSQL ``search_path`` to every transaction. + + Issued with ``SET LOCAL`` semantics at the start of each transaction — the one + scope a transaction-mode pooler such as PgBouncer honours. See + :func:`~sqlalchemy_foundation_kit.session.manager.attach_search_path`. + + Args: + search_path: A schema, or a comma-separated list such as ``"tenant_7, public"``. + + Returns: + Self for method chaining. + """ + self._search_path = search_path + return self + def build(self) -> AsyncSessionManager[SessionT]: """Build AsyncSessionManager instance with configured parameters. @@ -284,6 +301,7 @@ def build(self) -> AsyncSessionManager[SessionT]: "use_orjson": self._use_orjson, "metrics": self._metrics, "on_engine_created": self._on_engine_created, + "search_path": self._search_path, } if self._dispose_timeout is not None: kwargs["dispose_timeout"] = self._dispose_timeout diff --git a/sqlalchemy_foundation_kit/session/factories.py b/sqlalchemy_foundation_kit/session/factories.py index 8158c4a..589d742 100644 --- a/sqlalchemy_foundation_kit/session/factories.py +++ b/sqlalchemy_foundation_kit/session/factories.py @@ -31,6 +31,13 @@ def create_async_session_manager( ) -> AsyncSessionManager[AsyncSession]: """Create async session manager with PostgreSQL-specific configuration. + The startup packet carries ``application_name`` and, only when the config sets it, + ``jit``; nothing else, because a transaction-mode pooler such as PgBouncer rejects + startup parameters it does not track (``unsupported startup parameter: jit``) or, + with ``ignore_startup_parameters``, silently drops them. ``db_schema`` therefore + travels as the manager's ``search_path`` and is applied with ``SET LOCAL`` semantics + at the start of every transaction, which is the one scope such a pooler honours. + Args: postgres_config: PostgreSQL configuration implementing PostgresSettingsProtocol. application_name: Optional custom application name. If None, uses postgres_config.application_name. @@ -41,7 +48,8 @@ def create_async_session_manager( which provides pgbouncer transaction-mode compatibility. extra_server_settings: Additional PostgreSQL ``server_settings`` to merge with defaults (e.g., ``{"statement_timeout": "30000", "timezone": "UTC"}``). User-provided keys - override library defaults. + override library defaults. These are startup parameters: through PgBouncer only + the ones it tracks arrive (``track_extra_parameters``). extra_connect_args: Additional asyncpg ``connect_args`` to merge with defaults (e.g., ``{"command_timeout": 60}``). User-provided keys override library defaults. **kwargs: Additional keyword arguments passed to AsyncSessionManager. @@ -73,11 +81,11 @@ def create_async_session_manager( """ app_name = application_name or postgres_config.application_name - # Build server settings with optional overrides + # Build server settings with optional overrides. These are startup parameters, so + # search_path is deliberately not among them -- see the docstring. server_settings: dict[str, str] = { "application_name": app_name, **({"jit": postgres_config.jit} if postgres_config.jit is not None else {}), - **({"search_path": postgres_config.db_schema} if postgres_config.db_schema is not None else {}), **(extra_server_settings or {}), } @@ -100,5 +108,6 @@ def create_async_session_manager( use_orjson=postgres_config.use_orjson_serialization, metrics=metrics, on_engine_created=on_engine_created, + search_path=postgres_config.db_schema, **kwargs, ) diff --git a/sqlalchemy_foundation_kit/session/manager.py b/sqlalchemy_foundation_kit/session/manager.py index 9372a66..097008f 100644 --- a/sqlalchemy_foundation_kit/session/manager.py +++ b/sqlalchemy_foundation_kit/session/manager.py @@ -10,7 +10,7 @@ from types import TracebackType from typing import TYPE_CHECKING, Any, Generic, cast -from sqlalchemy import event +from sqlalchemy import event, text from sqlalchemy.exc import TimeoutError as SATimeoutError from sqlalchemy.ext.asyncio import ( AsyncEngine, @@ -23,6 +23,8 @@ from ..base import build_engine_kwargs, resolve_pool_class if TYPE_CHECKING: + from sqlalchemy.engine import Connection + from ..config import PoolSettingsProtocol from ..protocols import PostgresMetricsProtocol @@ -100,6 +102,35 @@ def on_error(exception_context: Any) -> None: event.listen(engine.sync_engine, "handle_error", on_error) +def attach_search_path(engine: AsyncEngine, search_path: str) -> None: + """Apply ``search_path`` to every transaction the engine begins. + + Registers a ``begin`` listener that issues ``set_config('search_path', …, true)`` — + the function form of ``SET LOCAL``, with the value as a bind parameter — as the first + statement of each transaction, so the setting lives exactly as long as the transaction + does. That is the one scope a transaction-mode pooler such as PgBouncer honours: a + startup parameter is rejected or dropped before it reaches PostgreSQL, and a plain + ``SET`` on a server connection leaks to whichever client is handed it next. + + Under the asyncpg adapter ``BEGIN`` is sent lazily with the first statement, so the + ``set_config`` call lands inside the transaction rather than ahead of it. Every + transaction is covered — a session's autobegin, ``session.begin()``, a raw + ``engine.connect()``, the one that follows a ``commit()`` — but a statement run with + ``isolation_level="AUTOCOMMIT"`` begins none and runs with the server's default. + + Args: + engine: SQLAlchemy ``AsyncEngine`` to attach the listener to. + search_path: Value for ``search_path`` — a schema, or a comma-separated list such + as ``"tenant_7, public"``. + """ + statement = text("SELECT set_config('search_path', :search_path, true)") + + def on_begin(conn: Connection) -> None: + conn.execute(statement, {"search_path": search_path}) + + event.listen(engine.sync_engine, "begin", on_begin) + + class AsyncSessionManager(Generic[SessionT]): """Manages async database sessions with configurable connection pooling. @@ -124,6 +155,7 @@ def __init__( metrics: PostgresMetricsProtocol | None = None, on_engine_created: Callable[[AsyncEngine], None] | None = None, dispose_timeout: float = DEFAULT_DISPOSE_TIMEOUT_SECONDS, + search_path: str | None = None, **kwargs: object, ) -> None: """Initialize session manager with direct configuration. @@ -145,6 +177,10 @@ def __init__( dispose_timeout: Maximum seconds to wait for engine disposal in :meth:`aclose` (default: 30.0). Lower this in tests or short-lived environments; raise it if you have long-running transactions that need more time to settle. + search_path: PostgreSQL ``search_path`` applied to every transaction with + ``SET LOCAL`` semantics, which is what survives a transaction-mode pooler + (default: None — nothing is sent, the server default applies). See + :func:`attach_search_path`. **kwargs: Additional keyword arguments for ``create_async_engine``. """ self._closed = False @@ -174,6 +210,9 @@ def __init__( if metrics: attach_metrics(self._engine, metrics) + if search_path is not None: + attach_search_path(self._engine, search_path) + if on_engine_created is not None: on_engine_created(self._engine) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 2d85e13..c51577c 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -28,6 +28,8 @@ from sqlalchemy_foundation_kit.session.manager import AsyncSessionManager from tests.integration.models import Base +SEARCH_PATH_SCHEMA = "app" + def asyncpg_url(container: PostgresContainer) -> str: """Rewrite the container's connection URL onto the asyncpg driver.""" @@ -115,3 +117,27 @@ async def session_manager( ) yield manager await manager.aclose() + + +@pytest_asyncio.fixture +async def schema_session_manager( + postgres_container: PostgresContainer, + db_engine: AsyncEngine, +) -> AsyncGenerator[AsyncSessionManager[AsyncSession], None]: + """An AsyncSessionManager whose transactions run with ``search_path = app``. + + The schema holds one table, ``app.schema_probe``, so a test can prove an unqualified + name resolves there; it starts empty for every test. + """ + async with db_engine.begin() as conn: + await conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {SEARCH_PATH_SCHEMA}")) + await conn.execute(text(f"CREATE TABLE IF NOT EXISTS {SEARCH_PATH_SCHEMA}.schema_probe (v int)")) + await conn.execute(text(f"TRUNCATE TABLE {SEARCH_PATH_SCHEMA}.schema_probe")) + + manager: AsyncSessionManager[AsyncSession] = AsyncSessionManager( + asyncpg_url(postgres_container), + poolclass="async_adapted_queue", + search_path=SEARCH_PATH_SCHEMA, + ) + yield manager + await manager.aclose() diff --git a/tests/integration/test_session_integration.py b/tests/integration/test_session_integration.py index 2c694ac..f693b71 100644 --- a/tests/integration/test_session_integration.py +++ b/tests/integration/test_session_integration.py @@ -6,9 +6,10 @@ import pytest from sqlalchemy import select, text -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession from sqlalchemy_foundation_kit.session.manager import AsyncSessionManager +from tests.integration.conftest import SEARCH_PATH_SCHEMA, PostgresContainer from tests.integration.models import Status, TestUser # ============================================================================ @@ -427,3 +428,153 @@ async def test__get_transaction__no_isolation_level__leaves_the_server_default( # Assert assert actual == "read committed" + + +# ============================================================================ +# search_path Tests (issue #23 scenario) +# ============================================================================ + + +async def _search_path(session: AsyncSession) -> str: + return (await session.execute(text("SHOW search_path"))).scalar_one() # type: ignore[no-any-return] + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_session__search_path__applied_to_the_autobegun_transaction( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act + async with schema_session_manager.get_session() as session: + actual = await _search_path(session) + + # Assert + assert actual == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_transaction__search_path__applied_to_the_transaction( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act + async with schema_session_manager.get_transaction() as session: + actual = await _search_path(session) + + # Assert + assert actual == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_transaction__search_path_and_isolation_level__both_applied( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act + async with schema_session_manager.get_transaction(isolation_level="SERIALIZABLE") as session: + search_path = await _search_path(session) + isolation = (await session.execute(text("SHOW transaction_isolation"))).scalar_one() + + # Assert + assert search_path == SEARCH_PATH_SCHEMA + assert isolation == "serializable" + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__engine_connect__search_path__applied_outside_any_session( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act + async with schema_session_manager.engine.connect() as conn: + actual = (await conn.execute(text("SHOW search_path"))).scalar_one() + + # Assert + assert actual == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_session__search_path__reapplied_to_the_transaction_after_a_commit( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act: SET LOCAL ends with the transaction; the next one has to set it again + async with schema_session_manager.get_session() as session: + before = await _search_path(session) + await session.commit() + after = await _search_path(session) + + # Assert + assert before == SEARCH_PATH_SCHEMA + assert after == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_transaction__search_path__unqualified_table_resolves_in_the_schema( + schema_session_manager: AsyncSessionManager[AsyncSession], + db_engine: AsyncEngine, +) -> None: + # Act + async with schema_session_manager.get_transaction() as session: + await session.execute(text("INSERT INTO schema_probe (v) VALUES (1)")) + + # Assert + async with db_engine.connect() as conn: + count = (await conn.execute(text("SELECT count(*) FROM app.schema_probe"))).scalar_one() + assert count == 1 + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__get_session__no_search_path__leaves_the_server_default( + session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Act: the control case -- nothing is attached when no schema is configured + async with session_manager.get_session() as session: + actual = await _search_path(session) + + # Assert + assert actual == '"$user", public' + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__create_async_session_manager__db_schema__applied_per_transaction( + postgres_container: PostgresContainer, + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> None: + # Arrange: the reporter's config -- host, credentials, an application name and db_schema + pytest.importorskip("pydantic_settings") + from pydantic import SecretStr + + from sqlalchemy_foundation_kit.contrib.settings.postgres import BasePostgresConfig, ConnectionSettings + from sqlalchemy_foundation_kit.session.factories import create_async_session_manager + + config = BasePostgresConfig( + connection=ConnectionSettings( + host=postgres_container.get_container_host_ip(), + port=int(postgres_container.get_exposed_port(postgres_container.port)), + user=postgres_container.username, + password=SecretStr(postgres_container.password), + database=postgres_container.dbname, + ), + application_name="probe", + db_schema=SEARCH_PATH_SCHEMA, + use_orjson_serialization=False, + ) + manager = create_async_session_manager(config) + + # Act + try: + async with manager.get_session() as session: + jit = (await session.execute(text("SHOW jit"))).scalar_one() + search_path = await _search_path(session) + application_name = (await session.execute(text("SHOW application_name"))).scalar_one() + finally: + await manager.aclose() + + # Assert: the schema arrives, and jit is whatever the server says -- nothing was sent for it + assert search_path == SEARCH_PATH_SCHEMA + assert application_name == "probe" + assert jit == "on" diff --git a/tests/integration/test_uow_integration.py b/tests/integration/test_uow_integration.py index fe80ac3..abbfe6c 100644 --- a/tests/integration/test_uow_integration.py +++ b/tests/integration/test_uow_integration.py @@ -7,8 +7,10 @@ from sqlalchemy.exc import DBAPIError, IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy_foundation_kit.session.manager import AsyncSessionManager from sqlalchemy_foundation_kit.uow.enums import IsolationLevel from sqlalchemy_foundation_kit.uow.sqlalchemy import AsyncSQLAlchemyUnitOfWork, AsyncSQLAlchemyUowTransaction +from tests.integration.conftest import SEARCH_PATH_SCHEMA from tests.integration.models import TestUser @@ -288,3 +290,99 @@ async def test__transaction__invalid_isolation_level__raises_value_error( with pytest.raises(ValueError, match="Invalid isolation level"): async with uow.transaction(isolation_level="NOT A LEVEL"): pass + + +# ============================================================================ +# search_path Tests (issue #23 scenario) +# ============================================================================ + + +@pytest.fixture +def schema_uow( + schema_session_manager: AsyncSessionManager[AsyncSession], +) -> AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction]: + return AsyncSQLAlchemyUnitOfWork( + schema_session_manager.session_maker, transaction_factory=AsyncSQLAlchemyUowTransaction + ) + + +async def _search_path(session: AsyncSession) -> str: + return (await session.execute(text("SHOW search_path"))).scalar_one() # type: ignore[no-any-return] + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__transaction__search_path__applied_to_the_transaction( + schema_uow: AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction], +) -> None: + # Act + async with schema_uow.transaction() as tx: + actual = await _search_path(tx.session) + + # Assert + assert actual == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__query__search_path__applied_to_the_transaction( + schema_uow: AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction], +) -> None: + # Act + async with schema_uow.query() as qx: + actual = await _search_path(qx.session) + + # Assert + assert actual == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__transaction__search_path_and_isolation_level__both_applied( + schema_uow: AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction], +) -> None: + # Act + async with schema_uow.transaction(isolation_level=IsolationLevel.SERIALIZABLE) as tx: + search_path = await _search_path(tx.session) + isolation = await _transaction_isolation(tx.session) + + # Assert + assert search_path == SEARCH_PATH_SCHEMA + assert isolation == "serializable" + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__transaction__search_path__held_inside_and_after_a_savepoint( + schema_uow: AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction], +) -> None: + # Act + async with schema_uow.transaction() as tx: + async with tx.savepoint(): + inside = await _search_path(tx.session) + after = await _search_path(tx.session) + + # Assert + assert inside == SEARCH_PATH_SCHEMA + assert after == SEARCH_PATH_SCHEMA + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__managed_session__search_path__reapplied_after_the_caller_commits( + schema_uow: AsyncSQLAlchemyUnitOfWork[AsyncSQLAlchemyUowTransaction], +) -> None: + # Act: the commit ends the transaction the schema was set on; the next statement + # autobegins a new one, which has to see the schema again + async with schema_uow.managed_session() as (_tx, session): + before = await _search_path(session) + await session.execute(text("INSERT INTO schema_probe (v) VALUES (1)")) + await session.commit() + after = await _search_path(session) + count = (await session.execute(text("SELECT count(*) FROM schema_probe"))).scalar_one() + await session.rollback() + + # Assert + assert before == SEARCH_PATH_SCHEMA + assert after == SEARCH_PATH_SCHEMA + assert count == 1 diff --git a/tests/unit/contrib/settings/test_postgres.py b/tests/unit/contrib/settings/test_postgres.py index 2065531..3ae1c2e 100644 --- a/tests/unit/contrib/settings/test_postgres.py +++ b/tests/unit/contrib/settings/test_postgres.py @@ -413,7 +413,7 @@ def test__base_postgres_config__repr__masks_password() -> None: @pytest.mark.unit -def test__base_postgres_config__default_jit_off__succeeds() -> None: +def test__base_postgres_config__default_jit__is_none() -> None: # Arrange & Act config = BasePostgresConfig( connection=ConnectionSettings( @@ -423,8 +423,9 @@ def test__base_postgres_config__default_jit_off__succeeds() -> None: application_name="test", ) - # Assert - assert config.jit == "off" + # Assert: nothing is sent, so the server's own setting applies -- and a transaction-mode + # PgBouncer, which rejects the parameter, accepts the connection + assert config.jit is None @pytest.mark.unit diff --git a/tests/unit/session/test_builder.py b/tests/unit/session/test_builder.py index ad15cfb..d89bf74 100644 --- a/tests/unit/session/test_builder.py +++ b/tests/unit/session/test_builder.py @@ -37,6 +37,7 @@ def test__async_session_manager_builder__init__sets_defaults() -> None: assert builder._metrics is None assert builder._on_engine_created is None assert builder._dispose_timeout is None + assert builder._search_path is None assert builder._extra_kwargs == {} @@ -104,6 +105,7 @@ def test__async_session_manager_builder__with_echo__default_true() -> None: ("with_json_serialization", (True,), {}), ("with_extra_kwargs", (), {"custom": "value"}), ("with_dispose_timeout", (30.0,), {}), + ("with_search_path", ("app",), {}), ], ) def test__async_session_manager_builder__methods__return_self(method_name: str, args: tuple, kwargs: dict) -> None: @@ -410,6 +412,22 @@ def test__async_session_manager_builder__with_dispose_timeout__sets_timeout() -> assert builder._dispose_timeout == 60.0 +# ============================================================================ +# AsyncSessionManagerBuilder - with_search_path Tests +# ============================================================================ + + +def test__async_session_manager_builder__with_search_path__sets_value() -> None: + # Arrange + builder = AsyncSessionManagerBuilder("postgresql://localhost/test") + + # Act + builder.with_search_path("tenant_7, public") + + # Assert + assert builder._search_path == "tenant_7, public" + + # ============================================================================ # AsyncSessionManagerBuilder - Method Chaining Tests # ============================================================================ @@ -451,6 +469,7 @@ def test__async_session_manager_builder__method_chaining__all_methods() -> None: .with_json_serialization(True) .with_extra_kwargs(custom="value") .with_dispose_timeout(45.0) + .with_search_path("app") ) # Assert @@ -467,6 +486,7 @@ def test__async_session_manager_builder__method_chaining__all_methods() -> None: assert builder._use_orjson is True assert builder._extra_kwargs["custom"] == "value" assert builder._dispose_timeout == 45.0 + assert builder._search_path == "app" # ============================================================================ @@ -602,6 +622,18 @@ def test__async_session_manager_builder__build__excludes_none_dispose_timeout() assert "dispose_timeout" not in call_kwargs +def test__async_session_manager_builder__build__no_search_path__passes_none() -> None: + # Arrange + builder = AsyncSessionManagerBuilder("postgresql://localhost/test") + + # Act + with patch("sqlalchemy_foundation_kit.session.builder.AsyncSessionManager") as mock_manager_class: + builder.build() + + # Assert + assert mock_manager_class.call_args[1]["search_path"] is None + + def test__async_session_manager_builder__build__passes_extra_kwargs() -> None: # Arrange builder = AsyncSessionManagerBuilder("postgresql://localhost/test").with_extra_kwargs( @@ -638,6 +670,7 @@ def test__async_session_manager_builder__build__all_params_combined() -> None: .with_json_serialization(True) .with_extra_kwargs(custom="value") .with_dispose_timeout(60.0) + .with_search_path("app") ) # Act @@ -659,6 +692,7 @@ def test__async_session_manager_builder__build__all_params_combined() -> None: assert call_kwargs["use_orjson"] is True assert call_kwargs["custom"] == "value" assert call_kwargs["dispose_timeout"] == 60.0 + assert call_kwargs["search_path"] == "app" # ============================================================================ diff --git a/tests/unit/session/test_factories.py b/tests/unit/session/test_factories.py index a8d4c12..5cafc75 100644 --- a/tests/unit/session/test_factories.py +++ b/tests/unit/session/test_factories.py @@ -2,11 +2,17 @@ from __future__ import annotations +from typing import Any from unittest.mock import Mock, patch import asyncpg import pytest +pytest.importorskip("pydantic_settings") + +from pydantic import SecretStr + +from sqlalchemy_foundation_kit.contrib.settings.postgres import BasePostgresConfig, ConnectionSettings from sqlalchemy_foundation_kit.session.connection import AsyncCConnection from sqlalchemy_foundation_kit.session.factories import create_async_session_manager @@ -160,16 +166,9 @@ def test__create_async_session_manager__jit__handled_correctly( assert "jit" not in server_settings -@pytest.mark.parametrize( - "db_schema,should_be_in_settings", - [ - ("public", True), - ("custom_schema", True), - (None, False), - ], -) -def test__create_async_session_manager__db_schema__handled_correctly( - db_schema: str | None, should_be_in_settings: bool +@pytest.mark.parametrize("db_schema", ["public", "custom_schema", None]) +def test__create_async_session_manager__db_schema__search_path_on_the_manager_not_in_startup( + db_schema: str | None, ) -> None: # Arrange mock_config = Mock() @@ -189,13 +188,11 @@ def test__create_async_session_manager__db_schema__handled_correctly( with patch("sqlalchemy_foundation_kit.session.factories.AsyncSessionManager") as mock_manager_class: create_async_session_manager(mock_config) - # Assert + # Assert: a startup parameter never reaches PostgreSQL through a transaction-mode pooler, + # so the schema goes to the manager, which applies it per transaction call_kwargs = mock_manager_class.call_args[1] - server_settings = call_kwargs["connect_args"]["server_settings"] - if should_be_in_settings: - assert server_settings["search_path"] == db_schema - else: - assert "search_path" not in server_settings + assert "search_path" not in call_kwargs["connect_args"]["server_settings"] + assert call_kwargs["search_path"] == db_schema def test__create_async_session_manager__extra_server_settings__merges_with_defaults() -> None: @@ -562,5 +559,61 @@ class CustomConn(asyncpg.Connection): server_settings = connect_args["server_settings"] assert server_settings["application_name"] == "override-app" assert server_settings["jit"] == "off" - assert server_settings["search_path"] == "custom_schema" + assert "search_path" not in server_settings assert server_settings["timezone"] == "UTC" + assert call_kwargs["search_path"] == "custom_schema" + + +# ============================================================================ +# create_async_session_manager - Startup Parameters (issue #23 scenario) +# ============================================================================ + + +def _postgres_config(**overrides: Any) -> BasePostgresConfig: + """The issue #23 shape: host, credentials and an application name, nothing else.""" + return BasePostgresConfig( + connection=ConnectionSettings(password=SecretStr("secret"), database="app"), + application_name="pgbouncer-lab", + **overrides, + ) + + +def test__create_async_session_manager__default_config__startup_carries_only_application_name() -> None: + # Arrange + config = _postgres_config() + + # Act + with patch("sqlalchemy_foundation_kit.session.factories.AsyncSessionManager") as mock_manager_class: + create_async_session_manager(config) + + # Assert: PgBouncer in transaction mode rejects any startup parameter it does not track + call_kwargs = mock_manager_class.call_args[1] + assert call_kwargs["connect_args"]["server_settings"] == {"application_name": "pgbouncer-lab"} + assert call_kwargs["search_path"] is None + + +def test__create_async_session_manager__explicit_jit_off__is_still_a_startup_parameter() -> None: + # Arrange + config = _postgres_config(jit="off") + + # Act + with patch("sqlalchemy_foundation_kit.session.factories.AsyncSessionManager") as mock_manager_class: + create_async_session_manager(config) + + # Assert + server_settings = mock_manager_class.call_args[1]["connect_args"]["server_settings"] + assert server_settings["jit"] == "off" + + +def test__create_async_session_manager__db_schema__becomes_the_manager_search_path() -> None: + # Arrange + config = _postgres_config(db_schema="app") + + # Act + with patch("sqlalchemy_foundation_kit.session.factories.AsyncSessionManager") as mock_manager_class: + create_async_session_manager(config) + + # Assert + call_kwargs = mock_manager_class.call_args[1] + assert call_kwargs["connect_args"]["server_settings"] == {"application_name": "pgbouncer-lab"} + assert call_kwargs["search_path"] == "app" diff --git a/tests/unit/session/test_manager.py b/tests/unit/session/test_manager.py index 83e5ff6..a0760d8 100644 --- a/tests/unit/session/test_manager.py +++ b/tests/unit/session/test_manager.py @@ -13,6 +13,7 @@ AsyncSessionManager, _safe_metric_call, attach_metrics, + attach_search_path, ) # ============================================================================ @@ -294,6 +295,44 @@ def test__attach_metrics__pool_without_size__uses_zero() -> None: assert call_args["pool_overflow"] == 0 +# ============================================================================ +# attach_search_path Tests +# ============================================================================ + + +def test__attach_search_path__registers_begin_listener() -> None: + # Arrange + mock_engine = Mock() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: + attach_search_path(mock_engine, "app") + + # Assert + mock_event.listen.assert_called_once() + target, identifier, _on_begin = mock_event.listen.call_args[0] + assert target == mock_engine.sync_engine + assert identifier == "begin" + + +def test__attach_search_path__on_begin__sets_search_path_locally_with_a_bind_parameter() -> None: + # Arrange + mock_engine = Mock() + mock_conn = Mock() + with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: + attach_search_path(mock_engine, "tenant_7, public") + on_begin = mock_event.listen.call_args[0][2] + + # Act + on_begin(mock_conn) + + # Assert + mock_conn.execute.assert_called_once() + statement, params = mock_conn.execute.call_args[0] + assert str(statement) == "SELECT set_config('search_path', :search_path, true)" + assert params == {"search_path": "tenant_7, public"} + + # ============================================================================ # AsyncSessionManager - Initialization Tests # ============================================================================ @@ -363,6 +402,26 @@ def test__async_session_manager__init__no_metrics__does_not_attach() -> None: mock_attach.assert_not_called() +def test__async_session_manager__init__search_path__attaches_it() -> None: + # Arrange & Act + with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine"): + with patch("sqlalchemy_foundation_kit.session.manager.attach_search_path") as mock_attach: + manager = AsyncSessionManager("postgresql://localhost/test", search_path="app") + + # Assert + mock_attach.assert_called_once_with(manager._engine, "app") + + +def test__async_session_manager__init__no_search_path__does_not_attach() -> None: + # Arrange & Act + with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine"): + with patch("sqlalchemy_foundation_kit.session.manager.attach_search_path") as mock_attach: + AsyncSessionManager("postgresql://localhost/test") + + # Assert + mock_attach.assert_not_called() + + def test__async_session_manager__init__calls_on_engine_created() -> None: # Arrange mock_callback = Mock()