feat: bound statements with a default command_timeout

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 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
 
 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=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
 

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:
@@ -34,8 +36,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 +73,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 = DEFAULT_COMMAND_TIMEOUT_SECONDS,
 ) -> 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:

docs/adr/0003-statements-get-a-default-bound.md

@@ -0,0 +1,24 @@
+# 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`.
+
+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.

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]

tests/test_command_timeout.py

@@ -0,0 +1,137 @@
+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
+
+
+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

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)