From 55b41b973ad1ff88deaab69d067a0fe696c30897 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Wed, 28 Oct 2026 19:32:53 +0300 Subject: [PATCH 1/3] feat: opt-in command_timeout on build_connection_factory A statement on a pooled connection to a primary that stops answering hung until TCP keepalive gave up. The keyword bounds each statement; a command_timeout in the DSN query wins. Default stays None. --- README.md | 5 +- db_retry/connections.py | 8 ++- tests/test_command_timeout.py | 128 ++++++++++++++++++++++++++++++++++ tests_dr/test_failover.py | 10 +-- 4 files changed, 144 insertions(+), 7 deletions(-) create mode 100644 tests/test_command_timeout.py diff --git a/README.md b/README.md index 5b18524..e52e5bc 100644 --- a/README.md +++ b/README.md @@ -111,12 +111,15 @@ dsn = build_db_dsn(db_dsn=multi_host_dsn, database_name="production_database", d connection_factory = build_connection_factory( url=dsn, timeout=5.0, # 5 second connection timeout + command_timeout=30.0, # 30 second bound on each statement ) # On a connect timeout, the factory tries each host in turn engine = create_async_engine(dsn, async_creator=connection_factory) ``` +`timeout` bounds only the connect. Statements have no bound unless you set `command_timeout`, either as the keyword or as `command_timeout` in the DSN query (the DSN value wins). Without one, a statement on a pooled connection to a primary that stops answering, for example a frozen VM or a network partition that drops packets, waits until TCP keepalive gives up, which can take hours. With one, asyncpg raises `TimeoutError` and the next connection goes to the new primary. `postgres_retry` does not retry that `TimeoutError`, because the statement may have run. Set it whenever writes must survive a failover, and pick a value above your slowest statement. + ### 3. Transaction management Wrap a unit of work in a transaction that is rolled back and closed on exit: @@ -207,7 +210,7 @@ export DB_RETRY_RETRIES_NUMBER=5 - Retries also fire when the retriable `asyncpg` error is wrapped by [`advanced-alchemy`](https://github.com/litestar-org/advanced-alchemy)'s `wrap_sqlalchemy_exception()` (i.e. surfaced as `RepositoryError` / `IntegrityError`); the handler walks the `__cause__` / `__context__` chain. ### Connection utilities -- `build_connection_factory(url, timeout)` - Creates a connection factory for multi-host setups +- `build_connection_factory(url, timeout, *, command_timeout=None)` - Creates a connection factory for multi-host setups; `timeout` bounds the connect, `command_timeout` bounds each statement (off by default; a `command_timeout` in the DSN wins) - `build_db_dsn(db_dsn, database_name, use_replica=False, drivername="postgresql")` - Builds a DSN with the given database name and driver, and sets `target_session_attrs` to `read-write`, or to `prefer-standby` when `use_replica=True` - `is_dsn_multihost(db_dsn)` - Checks if a DSN contains multiple hosts diff --git a/db_retry/connections.py b/db_retry/connections.py index 8111dc2..e9530aa 100644 --- a/db_retry/connections.py +++ b/db_retry/connections.py @@ -34,8 +34,10 @@ def parse_connect_args(url: sqlalchemy.URL) -> tuple[dict[str, typing.Any], list return connect_args, [] -def build_connection_plan(url: sqlalchemy.URL) -> ConnectionPlan: +def build_connection_plan(url: sqlalchemy.URL, command_timeout: float | None = None) -> ConnectionPlan: connect_args, hosts_and_ports = parse_connect_args(url) + if command_timeout is not None: + connect_args.setdefault("command_timeout", command_timeout) raw_target_session_attrs: str | None = connect_args.pop("target_session_attrs", None) target_session_attrs: SessionAttribute | None = ( SessionAttribute(raw_target_session_attrs) if raw_target_session_attrs else None @@ -69,8 +71,10 @@ def _shuffled(failover: tuple[tuple[str, int], ...]) -> list[tuple[str, int]]: def build_connection_factory( url: sqlalchemy.URL, timeout: float, + *, + command_timeout: float | None = None, ) -> typing.Callable[[], typing.Awaitable["ConnectionType"]]: - plan: typing.Final = build_connection_plan(url) + plan: typing.Final = build_connection_plan(url, command_timeout) async def _connection_factory() -> "ConnectionType": if not plan.failover: diff --git a/tests/test_command_timeout.py b/tests/test_command_timeout.py new file mode 100644 index 0000000..37a6c7b --- /dev/null +++ b/tests/test_command_timeout.py @@ -0,0 +1,128 @@ +import asyncio +import contextlib +import os +import typing +from unittest import mock + +import pytest +import sqlalchemy +from sqlalchemy.ext import asyncio as sa_async + +from db_retry.connections import build_connection_factory + + +COMMAND_TIMEOUT_SECONDS: typing.Final = 0.5 +HANG_BUDGET_SECONDS: typing.Final = 5.0 + + +class FreezingProxy: + """TCP proxy whose freeze() silences the connections open so far while still forwarding new ones.""" + + def __init__(self, upstream_host: str, upstream_port: int) -> None: + self._upstream: typing.Final = (upstream_host, upstream_port) + self._pumps: list[asyncio.Task[None]] = [] + self._writers: list[asyncio.StreamWriter] = [] + self._server: asyncio.Server | None = None + + async def start(self) -> int: + self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0) + return int(self._server.sockets[0].getsockname()[1]) + + def freeze(self) -> None: + for pump in self._pumps: + pump.cancel() + self._pumps.clear() + + async def close(self) -> None: + self.freeze() + for writer in self._writers: + writer.close() + if self._server is not None: + self._server.close() + await self._server.wait_closed() + + async def _handle(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None: + upstream_reader, upstream_writer = await asyncio.open_connection(*self._upstream) + self._writers += [client_writer, upstream_writer] + self._pumps += [ + asyncio.create_task(self._pump(client_reader, upstream_writer)), + asyncio.create_task(self._pump(upstream_reader, client_writer)), + ] + + @staticmethod + async def _pump(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + with contextlib.suppress(ConnectionError): + while chunk := await reader.read(65536): + writer.write(chunk) + await writer.drain() + + +@pytest.fixture +async def proxied_url() -> typing.AsyncIterator[tuple[sqlalchemy.URL, FreezingProxy]]: + url: typing.Final = sqlalchemy.make_url(os.environ["DB_DSN"]) + proxy: typing.Final = FreezingProxy(url.host or "localhost", url.port or 5432) + port: typing.Final = await proxy.start() + try: + yield url.set(host="127.0.0.1", port=port), proxy + finally: + await proxy.close() + + +async def test_statement_on_frozen_connection_times_out( + proxied_url: tuple[sqlalchemy.URL, FreezingProxy], +) -> None: + url, proxy = proxied_url + engine: typing.Final = sa_async.create_async_engine( + url=url, + async_creator=build_connection_factory(url=url, timeout=1.0, command_timeout=COMMAND_TIMEOUT_SECONDS), + ) + try: + async with engine.connect() as connection: + await connection.execute(sqlalchemy.text("SELECT 1")) + proxy.freeze() + + async with asyncio.timeout(HANG_BUDGET_SECONDS): + with pytest.raises(TimeoutError): + async with engine.connect() as connection: + await connection.execute(sqlalchemy.text("SELECT 1")) + async with engine.connect() as connection: + assert await connection.scalar(sqlalchemy.text("SELECT 1")) == 1 + finally: + await engine.dispose() + + +@pytest.mark.parametrize( + "dsn", + [ + "postgresql+asyncpg://user:password@host/database", + "postgresql+asyncpg://user:password@/database?host=host1:5432&host=host2:5432", + ], +) +@pytest.mark.parametrize( + ("dsn_command_timeout", "keyword_command_timeout", "expected"), + [ + (None, None, None), + (None, 2.0, 2.0), + ("7", 2.0, "7"), + ("7", None, "7"), + ], +) +async def test_command_timeout_reaches_asyncpg( + monkeypatch: pytest.MonkeyPatch, + dsn: str, + dsn_command_timeout: str | None, + keyword_command_timeout: float | None, + expected: str | float | None, +) -> None: + connect: typing.Final = mock.AsyncMock() + monkeypatch.setattr("asyncpg.connect", connect) + url: typing.Final = sqlalchemy.make_url(dsn) + factory: typing.Final = build_connection_factory( + url=url.update_query_dict({"command_timeout": dsn_command_timeout}) if dsn_command_timeout else url, + timeout=1.0, + command_timeout=keyword_command_timeout, + ) + + await factory() + + assert connect.call_args.kwargs.get("command_timeout") == expected diff --git a/tests_dr/test_failover.py b/tests_dr/test_failover.py index 866195a..46f7326 100644 --- a/tests_dr/test_failover.py +++ b/tests_dr/test_failover.py @@ -121,7 +121,10 @@ async def _landed_hosts(url: sa.URL) -> collections.Counter[str]: def _engine(url: sa.URL) -> sa_async.AsyncEngine: return sa_async.create_async_engine( - url, async_creator=build_connection_factory(url=url, timeout=CONNECT_TIMEOUT_SECONDS) + url, + async_creator=build_connection_factory( + url=url, timeout=CONNECT_TIMEOUT_SECONDS, command_timeout=COMMAND_TIMEOUT_SECONDS + ), ) @@ -130,9 +133,8 @@ async def test_writes_reach_promoted_replica( fail_primary: typing.Callable[[], typing.Awaitable[None]], caplog: pytest.LogCaptureFixture ) -> None: caplog.set_level(logging.WARNING, logger="db_retry.connections") - dsn = f"{DB_DSN}&command_timeout={COMMAND_TIMEOUT_SECONDS}" - writes_url = build_db_dsn(dsn, DATABASE, drivername=DRIVER) - reads_url = build_db_dsn(dsn, DATABASE, use_replica=True, drivername=DRIVER) + writes_url = build_db_dsn(DB_DSN, DATABASE, drivername=DRIVER) + reads_url = build_db_dsn(DB_DSN, DATABASE, use_replica=True, drivername=DRIVER) writes = _engine(writes_url) reads = _engine(reads_url) sessions = sa_async.async_sessionmaker(writes, expire_on_commit=False) From 35e8d2a7f976c2e094f2ce21e54cdde505f646d7 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Wed, 28 Oct 2026 19:35:21 +0300 Subject: [PATCH 2/3] feat: default 60s command_timeout on build_connection_factory The bound is now on by default. None turns it off; a command_timeout in the DSN still wins. Adds ADR 0003. --- README.md | 4 ++-- db_retry/connections.py | 4 +++- .../0003-statements-get-a-default-bound.md | 19 +++++++++++++++++++ tests/test_command_timeout.py | 9 +++++++++ 4 files changed, 33 insertions(+), 3 deletions(-) create mode 100644 docs/adr/0003-statements-get-a-default-bound.md diff --git a/README.md b/README.md index e52e5bc..7253fa6 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,7 @@ connection_factory = build_connection_factory( engine = create_async_engine(dsn, async_creator=connection_factory) ``` -`timeout` bounds only the connect. Statements have no bound unless you set `command_timeout`, either as the keyword or as `command_timeout` in the DSN query (the DSN value wins). Without one, a statement on a pooled connection to a primary that stops answering, for example a frozen VM or a network partition that drops packets, waits until TCP keepalive gives up, which can take hours. With one, asyncpg raises `TimeoutError` and the next connection goes to the new primary. `postgres_retry` does not retry that `TimeoutError`, because the statement may have run. Set it whenever writes must survive a failover, and pick a value above your slowest statement. +`timeout` bounds the connect and `command_timeout` bounds each statement. `command_timeout` defaults to 60 seconds; a `command_timeout` in the DSN query wins over the keyword, and `command_timeout=None` turns the bound off. Without a bound, a statement on a pooled connection to a primary that stops answering, for example a frozen VM or a network partition that drops packets, waits until TCP keepalive gives up, which can take hours. With one, asyncpg raises `TimeoutError` and the next connection goes to the new primary. `postgres_retry` does not retry that `TimeoutError`, because the statement may have run. Raise the value, or set it to `None`, for engines that run statements longer than a minute, such as migrations or reports. ### 3. Transaction management @@ -210,7 +210,7 @@ export DB_RETRY_RETRIES_NUMBER=5 - Retries also fire when the retriable `asyncpg` error is wrapped by [`advanced-alchemy`](https://github.com/litestar-org/advanced-alchemy)'s `wrap_sqlalchemy_exception()` (i.e. surfaced as `RepositoryError` / `IntegrityError`); the handler walks the `__cause__` / `__context__` chain. ### Connection utilities -- `build_connection_factory(url, timeout, *, command_timeout=None)` - Creates a connection factory for multi-host setups; `timeout` bounds the connect, `command_timeout` bounds each statement (off by default; a `command_timeout` in the DSN wins) +- `build_connection_factory(url, timeout, *, command_timeout=60.0)` - Creates a connection factory for multi-host setups; `timeout` bounds the connect, `command_timeout` bounds each statement (`None` turns it off; a `command_timeout` in the DSN wins) - `build_db_dsn(db_dsn, database_name, use_replica=False, drivername="postgresql")` - Builds a DSN with the given database name and driver, and sets `target_session_attrs` to `read-write`, or to `prefer-standby` when `use_replica=True` - `is_dsn_multihost(db_dsn)` - Checks if a DSN contains multiple hosts diff --git a/db_retry/connections.py b/db_retry/connections.py index e9530aa..e12ccb8 100644 --- a/db_retry/connections.py +++ b/db_retry/connections.py @@ -17,6 +17,8 @@ logger = logging.getLogger(__name__) +DEFAULT_COMMAND_TIMEOUT_SECONDS: typing.Final = 60.0 + @dataclasses.dataclass(kw_only=True, frozen=True, slots=True) class ConnectionPlan: @@ -72,7 +74,7 @@ def build_connection_factory( url: sqlalchemy.URL, timeout: float, *, - command_timeout: float | None = None, + command_timeout: float | None = DEFAULT_COMMAND_TIMEOUT_SECONDS, ) -> typing.Callable[[], typing.Awaitable["ConnectionType"]]: plan: typing.Final = build_connection_plan(url, command_timeout) diff --git a/docs/adr/0003-statements-get-a-default-bound.md b/docs/adr/0003-statements-get-a-default-bound.md new file mode 100644 index 0000000..6f463f7 --- /dev/null +++ b/docs/adr/0003-statements-get-a-default-bound.md @@ -0,0 +1,19 @@ +# Statements get a default bound + +`build_connection_factory` passes asyncpg a `command_timeout` of 60 seconds unless the caller sets +the keyword or the DSN query carries its own `command_timeout`, which wins. `None` turns the bound +off. Before this, db-retry bounded only the connect, so a statement on a pooled connection to a +primary that stopped answering (a paused VM, a partition that drops packets) waited for TCP +keepalive, which can take hours, and failover never reached the caller (#85). + +An opt-in keyword defaulting to `None` was considered and declined. It changes nothing for existing +users, but the people it misses are the ones who never read the failover docs, and they are the ones +who hang. The sibling libraries put a default bound inside the library too: faststream-outbox and +faststream-redis-timers default theirs to the lease. db-retry has no lease to borrow, so 60 seconds +is a chosen number, high enough for ordinary OLTP statements and low enough that a frozen primary +costs a minute instead of hours. + +The cost is a behaviour change. A statement that runs longer than 60 seconds now raises +`TimeoutError`, and `postgres_retry` does not retry it, because the statement may have completed +(see [ADR 0002](0002-retriable-taxonomy-is-asyncpg-classes.md)). Engines that run migrations or +long reports must raise the keyword or pass `None`. diff --git a/tests/test_command_timeout.py b/tests/test_command_timeout.py index 37a6c7b..deb4e24 100644 --- a/tests/test_command_timeout.py +++ b/tests/test_command_timeout.py @@ -126,3 +126,12 @@ async def test_command_timeout_reaches_asyncpg( await factory() assert connect.call_args.kwargs.get("command_timeout") == expected + + +async def test_command_timeout_defaults_to_a_minute(monkeypatch: pytest.MonkeyPatch) -> None: + connect: typing.Final = mock.AsyncMock() + monkeypatch.setattr("asyncpg.connect", connect) + + await build_connection_factory(url=sqlalchemy.make_url("postgresql+asyncpg://user:password@host/db"), timeout=1.0)() + + assert connect.call_args.kwargs["command_timeout"] == 60.0 # noqa: PLR2004 From 10db90ad804d94dcc4cb3fa68bc38323c881db8c Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Wed, 28 Oct 2026 19:55:32 +0300 Subject: [PATCH 3/3] build: require asyncpg>=0.32 Before 0.32, close() after a command_timeout waits with no bound for a cancel a frozen server never answers, so the caller still hangs. --- docs/adr/0003-statements-get-a-default-bound.md | 5 +++++ pyproject.toml | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/adr/0003-statements-get-a-default-bound.md b/docs/adr/0003-statements-get-a-default-bound.md index 6f463f7..51d084b 100644 --- a/docs/adr/0003-statements-get-a-default-bound.md +++ b/docs/adr/0003-statements-get-a-default-bound.md @@ -17,3 +17,8 @@ The cost is a behaviour change. A statement that runs longer than 60 seconds now `TimeoutError`, and `postgres_retry` does not retry it, because the statement may have completed (see [ADR 0002](0002-retriable-taxonomy-is-asyncpg-classes.md)). Engines that run migrations or long reports must raise the keyword or pass `None`. + +The bound needs asyncpg 0.32 or newer, which is now the floor. On 0.30 and 0.31 the statement times +out, but SQLAlchemy then discards the connection through asyncpg's `close()`, which waits with no +bound for the server to acknowledge the cancel. A frozen primary never does, so the caller hangs +there instead. diff --git a/pyproject.toml b/pyproject.toml index 08d8695..5200164 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,8 @@ dependencies = [ # 2.0.20, not 2.0.18: advanced-alchemy requires it, so nothing here can exercise 2.0.18. "sqlalchemy[asyncio]>=2.0.20; python_version < '3.13'", "sqlalchemy[asyncio]>=2.0.31; python_version >= '3.13'", - "asyncpg>=0.28", + # 0.32: before it, close() after a command_timeout waits with no bound for a cancel a frozen server never answers. + "asyncpg>=0.32", ] [project.urls]