diff --git a/AGENTS.md b/AGENTS.md index 69ae7f8..f902a9c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,6 +17,10 @@ title. The words most easily got wrong here are the ones this package shares wit `just --list`, or read it. The one thing it does not say: `just test` is Docker-only, so without a Docker daemon, point `DB_DSN` at any reachable PostgreSQL and run `uv run pytest` directly. +`just test-dr` runs the failover drills in `tests_dr/` (primary killed, primary paused), each on a +fresh cluster from `docker-compose.dr.yml`: a primary, a replica that gets promoted, and a standby +that stays read-only. It needs the Docker socket, takes about a minute, and is not in CI. + ## Architecture Every module under `db_retry/` is named for what it does and is short enough to read whole. Read diff --git a/docker-compose.dr.yml b/docker-compose.dr.yml new file mode 100644 index 0000000..2db8a94 --- /dev/null +++ b/docker-compose.dr.yml @@ -0,0 +1,55 @@ +name: db-retry-dr + +x-standby: &standby + image: postgres:17 + user: postgres + entrypoint: ["bash", "/replica-entrypoint.sh"] + environment: + - PGPASSWORD=dr + volumes: + - ./tests_dr/replica-entrypoint.sh:/replica-entrypoint.sh:ro + depends_on: + pg-primary: + condition: service_healthy + healthcheck: + test: ["CMD-SHELL", "pg_isready -h 127.0.0.1 -U dr -d dr"] + interval: 1s + timeout: 3s + retries: 60 + +services: + dr: + build: + context: . + dockerfile: ./tests_dr/Dockerfile + user: root + volumes: + - .:/code + - /code/.venv + - /var/run/docker.sock:/var/run/docker.sock + depends_on: + pg-replica: + condition: service_healthy + pg-standby: + condition: service_healthy + environment: + - DB_DSN=postgresql://dr:dr@/placeholder?host=pg-primary:5432&host=pg-standby:5432&host=pg-replica:5432 + + pg-primary: + image: postgres:17 + container_name: db-retry-dr-primary + command: ["postgres", "-c", "wal_keep_size=256MB"] + environment: + - POSTGRES_USER=dr + - POSTGRES_PASSWORD=dr + - POSTGRES_DB=dr + volumes: + - ./tests_dr/primary-init.sh:/docker-entrypoint-initdb.d/primary-init.sh:ro + healthcheck: + test: ["CMD-SHELL", "pg_isready -h 127.0.0.1 -U dr -d dr"] + interval: 1s + timeout: 3s + retries: 30 + + pg-replica: *standby + pg-standby: *standby diff --git a/justfile b/justfile index e4e0cd6..051e539 100644 --- a/justfile +++ b/justfile @@ -9,6 +9,17 @@ sh: test *args: down && down docker compose run application uv run --no-sync pytest {{ args }} +# Primary/replica failover drills (kill, pause the primary), each on a fresh cluster. Not run in CI. +test-dr: + #!/usr/bin/env sh + set -eu + compose="docker compose -f docker-compose.dr.yml" + trap '$compose down --remove-orphans --volumes' EXIT + for scenario in kill pause; do + $compose down --remove-orphans --volumes + $compose run --rm --build dr uv run pytest tests_dr -s -k "$scenario" + done + test-ci: uv run --no-sync pytest --cov=. --cov-report term-missing --cov-report xml diff --git a/pyproject.toml b/pyproject.toml index cec9638..08d8695 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,7 @@ isort.lines-after-imports = 2 isort.no-lines-before = ["standard-library", "local-folder"] [tool.ruff.lint.per-file-ignores] -"tests/**" = ["S101"] +"{tests,tests_dr}/**" = ["S101"] [tool.pytest.ini_options] testpaths = ["tests"] @@ -91,6 +91,7 @@ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" [tool.coverage.run] +omit = ["tests_dr/*"] concurrency = ["thread", "greenlet"] [tool.coverage.report] diff --git a/tests_dr/Dockerfile b/tests_dr/Dockerfile new file mode 100644 index 0000000..26bd746 --- /dev/null +++ b/tests_dr/Dockerfile @@ -0,0 +1,14 @@ +FROM python:3.13-slim + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv + +ENV UV_PROJECT_ENVIRONMENT=/code/.venv \ + UV_NO_MANAGED_PYTHON=1 \ + UV_NO_CACHE=true \ + UV_LINK_MODE=copy + +WORKDIR /code + +COPY pyproject.toml uv.lock ./ + +RUN uv sync --all-extras --frozen --no-install-project diff --git a/tests_dr/__init__.py b/tests_dr/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests_dr/primary-init.sh b/tests_dr/primary-init.sh new file mode 100644 index 0000000..c7de9fe --- /dev/null +++ b/tests_dr/primary-init.sh @@ -0,0 +1,2 @@ +#!/bin/sh +echo "host replication all all scram-sha-256" >> "$PGDATA/pg_hba.conf" diff --git a/tests_dr/replica-entrypoint.sh b/tests_dr/replica-entrypoint.sh new file mode 100644 index 0000000..6a505d4 --- /dev/null +++ b/tests_dr/replica-entrypoint.sh @@ -0,0 +1,7 @@ +#!/bin/bash +set -eu +if [ ! -s "$PGDATA/PG_VERSION" ]; then + pg_basebackup -h pg-primary -U dr -D "$PGDATA" -R -X stream + chmod 0700 "$PGDATA" +fi +exec postgres diff --git a/tests_dr/test_failover.py b/tests_dr/test_failover.py new file mode 100644 index 0000000..866195a --- /dev/null +++ b/tests_dr/test_failover.py @@ -0,0 +1,182 @@ +"""Primary/replica failover drills. Run via ``just test-dr``, one fresh cluster per scenario.""" + +import asyncio +import collections +import http +import logging +import os +import socket +import time +import typing + +import asyncpg +import pytest +import sqlalchemy as sa +from sqlalchemy.ext import asyncio as sa_async + +from db_retry import Transaction, build_connection_factory, build_db_dsn, postgres_retry + + +DB_DSN: typing.Final = os.environ.get("DB_DSN", "") +DATABASE: typing.Final = "dr" +DRIVER: typing.Final = "postgresql+asyncpg" +DOCKER_SOCKET: typing.Final = "/var/run/docker.sock" +PRIMARY_CONTAINER: typing.Final = "db-retry-dr-primary" +PRIMARY_HOST: typing.Final = "pg-primary" +STANDBY_HOST: typing.Final = "pg-standby" +REPLICA_HOST: typing.Final = "pg-replica" +CONNECT_TIMEOUT_SECONDS: typing.Final = 2.0 +COMMAND_TIMEOUT_SECONDS: typing.Final = 3 +FRESH_CONNECTIONS: typing.Final = 10 +RECOVERY_BUDGET_SECONDS: typing.Final = 60.0 +TEST_BUDGET_SECONDS: typing.Final = 180.0 + +pytestmark = pytest.mark.skipif(not DB_DSN, reason="needs the docker-compose.dr.yml cluster") + + +def _host_dsn(host: str) -> str: + return sa.make_url(DB_DSN).set(host=host, database=DATABASE, query={}).render_as_string(hide_password=False) + + +async def _docker(method: str, path: str) -> None: + reader, writer = await asyncio.open_unix_connection(DOCKER_SOCKET) + writer.write(f"{method} {path} HTTP/1.1\r\nHost: docker\r\nConnection: close\r\nContent-Length: 0\r\n\r\n".encode()) + await writer.drain() + status_line = await reader.readline() + writer.close() + assert int(status_line.split(b" ", 2)[1]) < http.HTTPStatus.MULTIPLE_CHOICES, status_line.decode(errors="replace") + + +async def _kill_primary() -> None: + await _docker("POST", f"/containers/{PRIMARY_CONTAINER}/kill") + + +async def _pause_primary() -> None: + """Freeze the primary: its name resolves and the kernel accepts connects, but nothing ever answers.""" + await _docker("POST", f"/containers/{PRIMARY_CONTAINER}/pause") + + +async def _address(host: str) -> str: + infos = await asyncio.get_running_loop().getaddrinfo(host, None, family=socket.AF_INET, type=socket.SOCK_STREAM) + return str(infos[0][4][0]) + + +async def _wait_replica_caught_up() -> None: + primary = await asyncpg.connect(_host_dsn(PRIMARY_HOST)) + replica = await asyncpg.connect(_host_dsn(REPLICA_HOST)) + try: + target = await primary.fetchval("SELECT pg_current_wal_lsn()") + deadline = time.monotonic() + 30 + while await replica.fetchval("SELECT pg_last_wal_replay_lsn() < $1::pg_lsn", target): + assert time.monotonic() < deadline, "replica did not catch up" + await asyncio.sleep(0.1) + finally: + await primary.close() + await replica.close() + + +async def _promote_replica() -> None: + replica = await asyncpg.connect(_host_dsn(REPLICA_HOST)) + try: + assert await replica.fetchval("SELECT pg_promote()") + finally: + await replica.close() + + +async def _standby_in_recovery() -> bool: + standby = await asyncpg.connect(_host_dsn(STANDBY_HOST)) + try: + return bool(await standby.fetchval("SELECT pg_is_in_recovery()")) + finally: + await standby.close() + + +async def _write_until_accepted( + write: typing.Callable[[int], typing.Awaitable[None]], n: int, deadline: float +) -> collections.Counter[str]: + failures: collections.Counter[str] = collections.Counter() + while True: + try: + await write(n) + except Exception as exc: + if time.monotonic() > deadline: + raise + failures[type(exc).__name__] += 1 + await asyncio.sleep(0.5) + else: + return failures + + +async def _landed_hosts(url: sa.URL) -> collections.Counter[str]: + factory = build_connection_factory(url=url, timeout=CONNECT_TIMEOUT_SECONDS) + landed: collections.Counter[str] = collections.Counter() + for _ in range(FRESH_CONNECTIONS): + connection = await factory() + try: + landed[await connection.fetchval("SELECT host(inet_server_addr())")] += 1 + finally: + await connection.close() + return landed + + +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) + ) + + +@pytest.mark.parametrize("fail_primary", [_kill_primary, _pause_primary], ids=["kill", "pause"]) +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 = _engine(writes_url) + reads = _engine(reads_url) + sessions = sa_async.async_sessionmaker(writes, expire_on_commit=False) + attempts = 0 + + @postgres_retry + async def write(n: int) -> None: + nonlocal attempts + attempts += 1 + async with Transaction(session=sessions()) as transaction: + await transaction.session.execute(sa.text("INSERT INTO dr (n) VALUES (:n)"), {"n": n}) + await transaction.commit() + + @postgres_retry + async def read() -> int: + async with reads.connect() as connection: + return int(await connection.scalar(sa.text("SELECT count(*) FROM dr")) or 0) + + try: + async with asyncio.timeout(TEST_BUDGET_SECONDS): + async with writes.begin() as connection: + await connection.execute(sa.text("CREATE TABLE dr (n int PRIMARY KEY)")) + await write(0) + await _wait_replica_caught_up() + assert await read() == 1 + + failed_at = time.monotonic() + await fail_primary() + assert await read() == 1 + await _promote_replica() + + attempts = 0 + failures = await _write_until_accepted(write, 1, failed_at + RECOVERY_BUDGET_SECONDS) + recovered = time.monotonic() - failed_at + assert await read() >= 1 + + assert await _landed_hosts(writes_url) == {await _address(REPLICA_HOST): FRESH_CONNECTIONS} + assert await _standby_in_recovery() + finally: + await writes.dispose() + await reads.dispose() + + fallbacks = sum("Trying host by host" in r.getMessage() for r in caplog.records) + print( # noqa: T201 + f"\nwrite recovered {recovered:.1f}s, postgres_retry attempts {attempts}, " + f"failed outer calls {sum(failures.values())} {dict(failures)}, host-by-host fallbacks {fallbacks}" + )