diff --git a/README.md b/README.md index 5b18524..7253fa6 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 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 diff --git a/db_retry/connections.py b/db_retry/connections.py index 8111dc2..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: @@ -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: 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..51d084b --- /dev/null +++ b/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. 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] diff --git a/tests/test_command_timeout.py b/tests/test_command_timeout.py new file mode 100644 index 0000000..deb4e24 --- /dev/null +++ b/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 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)