From df7d74a5a7d98512239464580909c9bcb9bded8a Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 29 Oct 2026 15:50:23 +0300 Subject: [PATCH] feat!: add InboxCleaner and drop the Inbox cleanup methods InboxCleaner runs a whole cleanup run: one cutoff from the database clock, then batches in their own READ COMMITTED transactions until one deletes nothing. Inbox.cleanup_cutoff and Inbox.delete_expired_batch are removed; Inbox now only claims. --- AGENTS.md | 2 + CONTEXT.md | 4 + README.md | 50 +++--- docs/adr/0001-postgres-only.md | 13 ++ sqlalchemy_inbox/__init__.py | 3 +- sqlalchemy_inbox/cleaner.py | 78 +++++++++ sqlalchemy_inbox/inbox.py | 32 +--- tests/test_cleaner.py | 282 +++++++++++++++++++++++++++++++++ tests/test_cleanup.py | 204 ------------------------ tests/test_readme_examples.py | 38 ++--- 10 files changed, 421 insertions(+), 285 deletions(-) create mode 100644 docs/adr/0001-postgres-only.md create mode 100644 sqlalchemy_inbox/cleaner.py create mode 100644 tests/test_cleaner.py delete mode 100644 tests/test_cleanup.py diff --git a/AGENTS.md b/AGENTS.md index 42cc593..855a31f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -22,6 +22,8 @@ The things it does not say: - `Inbox` never commits, rolls back, begins, or opens a session or connection. Transaction ownership stays with the caller. Any new method must keep that. +- `InboxCleaner` is the only code that begins and commits, and only on connections it opens from + the caller's engine. It never takes a caller's session and never disposes the engine. - `Inbox.claim` requires an open transaction on its session. That is what turns a wrongly wired second session into an error instead of a marker that never commits. - Concurrency claims are proven only by tests on separate connections, never one shared test diff --git a/CONTEXT.md b/CONTEXT.md index 8357d64..2995ba2 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -36,3 +36,7 @@ _Avoid_: TTL (nothing expires on its own; cleanup deletes) **Cutoff**: The database timestamp, taken once per cleanup run, before which markers are deleted. _Avoid_: deadline, expiry time + +**Cleanup run**: +One pass that takes a cutoff and deletes the markers older than it, batch by batch. +_Avoid_: vacuum (Postgres VACUUM is something else), purge, sweep, expiry diff --git a/README.md b/README.md index c0ddf1d..282aa28 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,9 @@ transaction as the consumer's own writes. A redelivered message finds its marker and is skipped. A failed transaction rolls the marker back, so a retry processes the message again. -It works with SQLAlchemy's asyncio extension on PostgreSQL. It stores no payload, runs no worker, -and never begins, commits or rolls back a transaction for you. +It works with SQLAlchemy's asyncio extension on PostgreSQL. It stores no payload and runs no worker. +`Inbox` never begins, commits or rolls back a transaction for you. Only `InboxCleaner` does that, on +its own connections. ```bash pip install sqlalchemy-inbox @@ -92,31 +93,38 @@ or fall back to Kafka offsets. ## Cleaning up -Markers are kept until you delete them. Take one cutoff per cleanup run from the database clock, -then delete in batches, one transaction per batch, until a batch deletes nothing: +Markers are kept until you delete them. Run `InboxCleaner` from a cron job: ```python -@dataclasses.dataclass(kw_only=True, frozen=True, slots=True) -class CleanupInboxUseCase: - transaction: Transaction - inbox: Inbox +import asyncio +import datetime - async def __call__(self, *, retention: datetime.timedelta, batch_size: int) -> None: - async with self.transaction: - cutoff = await self.inbox.cleanup_cutoff(retention=retention) - while await self._delete_batch(cutoff=cutoff, limit=batch_size): - pass +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy_inbox import InboxCleaner - @postgres_retry - async def _delete_batch(self, *, cutoff: datetime.datetime, limit: int) -> int: - async with self.transaction: - deleted = await self.inbox.delete_expired_batch(cutoff=cutoff, limit=limit) - await self.transaction.commit() - return deleted + +async def main() -> None: + engine = create_async_engine(DATABASE_DSN) + try: + cleaner = InboxCleaner(engine=engine, table=INBOX_TABLE, retention=datetime.timedelta(days=7)) + await cleaner.cleanup() + finally: + await engine.dispose() + + +asyncio.run(main()) ``` -Batches that committed stay deleted if a later batch fails, and the error reaches your scheduler. -Cleanup skips markers that another transaction holds, so it never waits on a consumer. +A cleanup run takes one cutoff from the database clock, `now()` minus `retention`, and deletes the +markers older than it in batches of `batch_size` (1000 by default), one transaction per batch, until +a batch deletes nothing. `cleanup()` returns a `CleanupResult` with the `cutoff` and the +`deleted_count`. It logs each batch at DEBUG and the total at INFO to the `sqlalchemy_inbox` logger. + +Batches that committed stay deleted if a later batch fails or the task is cancelled, and the error +reaches your cron job. The cleaner doesn't retry: the next run deletes what is left. It skips markers +that another transaction holds, so it never waits on a consumer, and two runs at once share the +work. Each batch runs at READ COMMITTED whatever the engine's default. The engine stays yours to +dispose. ## Guarantees diff --git a/docs/adr/0001-postgres-only.md b/docs/adr/0001-postgres-only.md new file mode 100644 index 0000000..cf1dd45 --- /dev/null +++ b/docs/adr/0001-postgres-only.md @@ -0,0 +1,13 @@ +# PostgreSQL only + +The package supports PostgreSQL and no other database. Its value is the guarantees in the README, and +those come from how Postgres behaves: `INSERT ... ON CONFLICT DO NOTHING` waits on an uncommitted +duplicate under READ COMMITTED, and cleanup relies on `FOR UPDATE SKIP LOCKED` and `DELETE ... +RETURNING`. Each extra dialect would need its own code paths, its own guarantees and its own CI matrix. + +## Considered Options + +- SQLite as well, for tests without a server. The insert syntax matches, but SQLite has one writer at a + time and no row locks, so it would test a different concurrency model from the one in production. +- MySQL. It has no `RETURNING`, rejects `LIMIT` in an `IN` subquery, and its `INSERT IGNORE` turns + unrelated errors into warnings, so both claim and cleanup would need separate implementations. diff --git a/sqlalchemy_inbox/__init__.py b/sqlalchemy_inbox/__init__.py index 40b7a2c..fee896c 100644 --- a/sqlalchemy_inbox/__init__.py +++ b/sqlalchemy_inbox/__init__.py @@ -1,5 +1,6 @@ +from sqlalchemy_inbox.cleaner import CleanupResult, InboxCleaner from sqlalchemy_inbox.inbox import Inbox, SessionOrConnection from sqlalchemy_inbox.schema import make_inbox_table -__all__ = ["Inbox", "SessionOrConnection", "make_inbox_table"] +__all__ = ["CleanupResult", "Inbox", "InboxCleaner", "SessionOrConnection", "make_inbox_table"] diff --git a/sqlalchemy_inbox/cleaner.py b/sqlalchemy_inbox/cleaner.py new file mode 100644 index 0000000..6bfc290 --- /dev/null +++ b/sqlalchemy_inbox/cleaner.py @@ -0,0 +1,78 @@ +import contextlib +import dataclasses +import datetime as dt +import logging +from collections.abc import AsyncGenerator + +from sqlalchemy import Table, delete, func, select, tuple_ +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine + +from sqlalchemy_inbox.schema import validate_table_identifiers + + +_LOGGER = logging.getLogger(__name__) + + +@dataclasses.dataclass(kw_only=True, frozen=True, slots=True) +class CleanupResult: + cutoff: dt.datetime + deleted_count: int + + +class InboxCleaner: + """Deletes markers older than *retention*, one transaction per batch. + + Opens its own connections from *engine* but never disposes it: the caller owns the engine. + """ + + def __init__(self, *, engine: AsyncEngine, table: Table, retention: dt.timedelta, batch_size: int = 1000) -> None: + if not isinstance(engine, AsyncEngine): + msg = "InboxCleaner requires an sqlalchemy.ext.asyncio.AsyncEngine" + raise TypeError(msg) + validate_table_identifiers(table.name) + if retention <= dt.timedelta(0): + msg = f"retention must be positive (got {retention})" + raise ValueError(msg) + if batch_size < 1: + msg = f"batch_size must be positive (got {batch_size})" + raise ValueError(msg) + self._engine = engine + self._table = table + self._retention = retention + self._batch_size = batch_size + + async def cleanup(self) -> CleanupResult: + """Take one cutoff from the database clock, then delete the markers older than it until a batch deletes nothing. + + Each batch commits on its own and skips markers other transactions hold. If a batch fails or the task + is cancelled, that batch rolls back, earlier batches stay deleted, and the error propagates. + """ + async with self._transaction() as connection: + cutoff = (await connection.execute(select(func.now() - self._retention))).scalar_one() + deleted_count = 0 + while deleted := await self._delete_batch(cutoff): + deleted_count += deleted + _LOGGER.debug("Inbox cleanup batch deleted %d markers (%d so far)", deleted, deleted_count) + _LOGGER.info("Inbox cleanup deleted %d markers older than %s", deleted_count, cutoff.isoformat()) + return CleanupResult(cutoff=cutoff, deleted_count=deleted_count) + + async def _delete_batch(self, cutoff: dt.datetime) -> int: + key = tuple_(self._table.c.scope, self._table.c.message_id) + expired = ( + select(self._table.c.scope, self._table.c.message_id) + .where(self._table.c.processed_at < cutoff) + .limit(self._batch_size) + .with_for_update(skip_locked=True) + ) + async with self._transaction() as connection: + result = await connection.execute( + delete(self._table).where(key.in_(expired)).returning(self._table.c.scope) + ) + return len(result.all()) + + @contextlib.asynccontextmanager + async def _transaction(self) -> AsyncGenerator[AsyncConnection]: + async with self._engine.connect() as connection: + await connection.execution_options(isolation_level="READ COMMITTED") + async with connection.begin(): + yield connection diff --git a/sqlalchemy_inbox/inbox.py b/sqlalchemy_inbox/inbox.py index dbccbac..eb537b8 100644 --- a/sqlalchemy_inbox/inbox.py +++ b/sqlalchemy_inbox/inbox.py @@ -1,6 +1,4 @@ -import datetime as dt - -from sqlalchemy import Table, delete, func, select, tuple_ +from sqlalchemy import Table from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession @@ -52,34 +50,6 @@ async def claim(self, *, scope: str, message_id: str) -> bool: result = await self._session.execute(statement) return result.first() is not None - async def cleanup_cutoff(self, *, retention: dt.timedelta) -> dt.datetime: - """Return the database's ``now()`` minus *retention*: take it once per cleanup run.""" - if retention <= dt.timedelta(0): - msg = f"retention must be positive (got {retention})" - raise ValueError(msg) - self._require_transaction("cleanup_cutoff") - result = await self._session.execute(select(func.now() - retention)) - return result.scalar_one() - - async def delete_expired_batch(self, *, cutoff: dt.datetime, limit: int) -> int: - """Delete up to *limit* markers older than *cutoff*, skipping rows other transactions hold. - - Returns the number deleted; a run is done when a batch returns 0. - """ - if limit < 1: - msg = f"limit must be positive (got {limit})" - raise ValueError(msg) - self._require_transaction("delete_expired_batch") - key = tuple_(self._table.c.scope, self._table.c.message_id) - expired = ( - select(self._table.c.scope, self._table.c.message_id) - .where(self._table.c.processed_at < cutoff) - .limit(limit) - .with_for_update(skip_locked=True) - ) - result = await self._session.execute(delete(self._table).where(key.in_(expired)).returning(self._table.c.scope)) - return len(result.all()) - def _require_transaction(self, method: str) -> None: if not self._session.in_transaction(): msg = f"Inbox.{method} requires an open transaction on its session" diff --git a/tests/test_cleaner.py b/tests/test_cleaner.py new file mode 100644 index 0000000..7ea7530 --- /dev/null +++ b/tests/test_cleaner.py @@ -0,0 +1,282 @@ +import asyncio +import contextlib +import datetime as dt +import logging +import typing +import uuid +from collections.abc import AsyncGenerator, AsyncIterator + +import pytest +from sqlalchemy import MetaData, Table, event, func, insert, select, text +from sqlalchemy.exc import DBAPIError +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession, async_sessionmaker +from sqlalchemy.pool import QueuePool + +from sqlalchemy_inbox import CleanupResult, Inbox, InboxCleaner, make_inbox_table +from tests.conftest import Tables, wait_until_lock_wait + + +HOUR = dt.timedelta(hours=1) +_ADVISORY_KEY = 7_331 + + +class _BatchError(Exception): + pass + + +async def _add_markers(engine: AsyncEngine, table: Table, *message_ids: str, age: dt.timedelta = 2 * HOUR) -> None: + async with engine.begin() as conn: + for message_id in message_ids: + await conn.execute(insert(table).values(scope="s", message_id=message_id, processed_at=func.now() - age)) + + +async def _message_ids(engine: AsyncEngine, table: Table) -> list[str]: + async with engine.connect() as conn: + return list((await conn.execute(select(table.c.message_id).order_by(table.c.message_id))).scalars()) + + +def _cleaner(engine: AsyncEngine, table: Table, *, batch_size: int = 100) -> InboxCleaner: + return InboxCleaner(engine=engine, table=table, retention=HOUR, batch_size=batch_size) + + +@contextlib.asynccontextmanager +async def _delete_trigger(engine: AsyncEngine, table: Table, body: str) -> AsyncGenerator[None]: + function = f"{table.name}_on_delete" + async with engine.begin() as conn: + await conn.execute( + text( + f"CREATE FUNCTION {function}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN {body} RETURN NULL; END $$" + ) + ) + await conn.execute( + text(f"CREATE TRIGGER {table.name}_on_delete AFTER DELETE ON {table.name} EXECUTE FUNCTION {function}()") + ) + yield + async with engine.begin() as conn: + await conn.execute(text(f"DROP FUNCTION {function}() CASCADE")) + + +@pytest.fixture +async def delete_gate(pg_engine: AsyncEngine, tables: Tables) -> AsyncIterator[AsyncConnection]: + """Make every cleanup batch wait after its delete until the test unlocks the yielded connection's lock.""" + async with ( + _delete_trigger(pg_engine, tables.inbox, f"PERFORM pg_advisory_xact_lock_shared({_ADVISORY_KEY});"), + pg_engine.connect() as holder, + ): + await holder.execute(select(func.pg_advisory_lock(_ADVISORY_KEY))) + try: + yield holder + finally: + await holder.execute(select(func.pg_advisory_unlock_all())) + + +async def _open_gate(holder: AsyncConnection) -> None: + await holder.execute(select(func.pg_advisory_unlock(_ADVISORY_KEY))) + + +async def _wait_for_blocked_delete(engine: AsyncEngine) -> None: + async with engine.connect() as conn: + for _ in range(200): + blocked = ( + await conn.execute( + text( + "SELECT count(*) FROM pg_stat_activity WHERE wait_event_type = 'Lock' " + "AND query LIKE 'DELETE%' AND pid <> pg_backend_pid()" + ) + ) + ).scalar_one() + await conn.commit() + if blocked: + return + await asyncio.sleep(0.025) # pragma: no cover - only when the first poll runs before the lock wait + msg = "the cleanup batch never waited on the advisory lock" # pragma: no cover - only when nothing blocked + raise AssertionError(msg) # pragma: no cover - only when nothing blocked + + +async def test_cleanup_deletes_only_markers_older_than_the_cutoff(pg_engine: AsyncEngine, tables: Tables) -> None: + await _add_markers(pg_engine, tables.inbox, "old-1", "old-2") + await _add_markers(pg_engine, tables.inbox, "recent", age=dt.timedelta(minutes=10)) + + result = await _cleaner(pg_engine, tables.inbox).cleanup() + + assert result.deleted_count == 2 + assert await _message_ids(pg_engine, tables.inbox) == ["recent"] + + +async def test_cutoff_comes_from_the_database_clock(pg_engine: AsyncEngine, tables: Tables) -> None: + async with pg_engine.connect() as conn: + before = (await conn.execute(select(func.now()))).scalar_one() + + result = await _cleaner(pg_engine, tables.inbox).cleanup() + + async with pg_engine.connect() as conn: + after = (await conn.execute(select(func.now()))).scalar_one() + assert before - HOUR <= result.cutoff <= after - HOUR + + +async def test_cleanup_deletes_in_batches_until_one_deletes_nothing( + pg_engine: AsyncEngine, tables: Tables, caplog: pytest.LogCaptureFixture +) -> None: + await _add_markers(pg_engine, tables.inbox, *(f"m-{index}" for index in range(5))) + + with caplog.at_level(logging.DEBUG, logger="sqlalchemy_inbox"): + result = await _cleaner(pg_engine, tables.inbox, batch_size=2).cleanup() + + assert result == CleanupResult(cutoff=result.cutoff, deleted_count=5) + assert await _message_ids(pg_engine, tables.inbox) == [] + levels = [record.levelno for record in caplog.records if record.name.startswith("sqlalchemy_inbox")] + assert levels == [logging.DEBUG] * 3 + [logging.INFO] + + +async def test_every_batch_uses_the_cutoff_taken_at_the_start(pg_engine: AsyncEngine, tables: Tables) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1", "m-2", "m-3") + cutoffs: list[dt.datetime] = [] + + @event.listens_for(pg_engine.sync_engine, "before_cursor_execute") + def record_cutoff( + _conn: object, _cursor: object, statement: str, parameters: tuple[object, ...], *_: object + ) -> None: + if statement.startswith("DELETE"): + cutoffs.extend(value for value in parameters if isinstance(value, dt.datetime)) + + result = await _cleaner(pg_engine, tables.inbox, batch_size=1).cleanup() + + assert cutoffs == [result.cutoff] * 4 + + +async def test_committed_batches_survive_a_failed_batch(pg_engine: AsyncEngine, tables: Tables) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1", "m-2", "m-3", "m-4") + deletes = 0 + + @event.listens_for(pg_engine.sync_engine, "before_cursor_execute") + def fail_second_delete(_conn: object, _cursor: object, statement: str, *_: object) -> None: + nonlocal deletes + if statement.startswith("DELETE"): + deletes += 1 + if deletes == 2: + raise _BatchError + + with pytest.raises(_BatchError): + await _cleaner(pg_engine, tables.inbox, batch_size=2).cleanup() + + assert len(await _message_ids(pg_engine, tables.inbox)) == 2 + + +async def test_cancelled_cleanup_rolls_back_the_batch_in_progress( + pg_engine: AsyncEngine, tables: Tables, delete_gate: AsyncConnection +) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1", "m-2") + cleanup = asyncio.create_task(_cleaner(pg_engine, tables.inbox).cleanup()) + await _wait_for_blocked_delete(pg_engine) + + cleanup.cancel() + with pytest.raises(asyncio.CancelledError): + await cleanup + await _open_gate(delete_gate) + + assert isinstance(pg_engine.pool, QueuePool) + assert pg_engine.pool.checkedout() == 1 + assert await _message_ids(pg_engine, tables.inbox) == ["m-1", "m-2"] + + +async def test_cleanup_skips_markers_locked_by_another_transaction( + pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] +) -> None: + await _add_markers(pg_engine, tables.inbox, "locked", "free") + + async with session_maker() as holder, holder.begin(): + await holder.execute(select(tables.inbox).where(tables.inbox.c.message_id == "locked").with_for_update()) + result = await asyncio.wait_for(_cleaner(pg_engine, tables.inbox).cleanup(), timeout=5) + + assert result.deleted_count == 1 + assert await _message_ids(pg_engine, tables.inbox) == ["locked"] + + +async def test_concurrent_runs_delete_every_marker_once(pg_engine: AsyncEngine, tables: Tables) -> None: + await _add_markers(pg_engine, tables.inbox, *(f"m-{index:02}" for index in range(20))) + + results = await asyncio.gather(*(_cleaner(pg_engine, tables.inbox, batch_size=1).cleanup() for _ in range(3))) + + assert sum(result.deleted_count for result in results) == 20 + assert await _message_ids(pg_engine, tables.inbox) == [] + + +async def test_cleanup_runs_at_read_committed_whatever_the_engine_default( + pg_engine: AsyncEngine, tables: Tables +) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1") + async with _delete_trigger( + pg_engine, + tables.inbox, + "IF current_setting('transaction_isolation') <> 'read committed' THEN RAISE EXCEPTION 'isolation'; END IF;", + ): + repeatable_read = pg_engine.execution_options(isolation_level="REPEATABLE READ") + result = await _cleaner(repeatable_read, tables.inbox).cleanup() + + assert result.deleted_count == 1 + + +async def test_an_expired_message_is_processed_again( + pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] +) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1") + async with session_maker() as session, session.begin(): + assert not await Inbox(session, tables.inbox).claim(scope="s", message_id="m-1") + + await _cleaner(pg_engine, tables.inbox).cleanup() + + async with session_maker() as session, session.begin(): + assert await Inbox(session, tables.inbox).claim(scope="s", message_id="m-1") + + +async def test_claim_during_uncommitted_cleanup_waits_then_claims( + pg_engine: AsyncEngine, + tables: Tables, + session_maker: async_sessionmaker[AsyncSession], + delete_gate: AsyncConnection, +) -> None: + await _add_markers(pg_engine, tables.inbox, "m-1") + cleanup = asyncio.create_task(_cleaner(pg_engine, tables.inbox).cleanup()) + await _wait_for_blocked_delete(pg_engine) + + async with session_maker() as consumer, consumer.begin(): + pid = (await consumer.execute(text("SELECT pg_backend_pid()"))).scalar_one() + competing = asyncio.create_task(Inbox(consumer, tables.inbox).claim(scope="s", message_id="m-1")) + await wait_until_lock_wait(pg_engine, pid) + assert not competing.done() + + await _open_gate(delete_gate) + assert (await cleanup).deleted_count == 1 + assert await competing + + assert await _message_ids(pg_engine, tables.inbox) == ["m-1"] + + +async def test_cleanup_errors_reach_the_caller(pg_engine: AsyncEngine) -> None: + missing = make_inbox_table(MetaData(), table_name=f"missing_{uuid.uuid4().hex[:12]}") + + with pytest.raises(DBAPIError): + await _cleaner(pg_engine, missing).cleanup() + + +@pytest.mark.parametrize( + ("overrides", "error", "match"), + [ + ({"retention": dt.timedelta(0)}, ValueError, "retention"), + ({"batch_size": 0}, ValueError, "batch_size"), + ({"table": Table("t" * 50, MetaData())}, ValueError, "too long"), + ({"engine": object()}, TypeError, "AsyncEngine"), + ], +) +async def test_constructor_rejects_invalid_arguments( + pg_engine: AsyncEngine, tables: Tables, overrides: dict[str, typing.Any], error: type[Exception], match: str +) -> None: + arguments: dict[str, typing.Any] = {"engine": pg_engine, "table": tables.inbox, "retention": HOUR} + + with pytest.raises(error, match=match): + InboxCleaner(**(arguments | overrides)) + + +async def test_constructor_takes_only_named_arguments(pg_engine: AsyncEngine, tables: Tables) -> None: + with pytest.raises(TypeError): + InboxCleaner(pg_engine, tables.inbox, HOUR) # ty: ignore[missing-argument, too-many-positional-arguments] diff --git a/tests/test_cleanup.py b/tests/test_cleanup.py deleted file mode 100644 index f6dc28f..0000000 --- a/tests/test_cleanup.py +++ /dev/null @@ -1,204 +0,0 @@ -import asyncio -import datetime as dt -import uuid - -import pytest -from sqlalchemy import MetaData, Table, func, insert, select, text -from sqlalchemy.exc import DBAPIError -from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker - -from sqlalchemy_inbox import Inbox, make_inbox_table -from tests.conftest import Tables, wait_until_lock_wait - - -HOUR = dt.timedelta(hours=1) - - -class _BatchError(Exception): - pass - - -async def _add_marker(engine: AsyncEngine, table: Table, message_id: str, age: dt.timedelta) -> None: - async with engine.begin() as conn: - await conn.execute(insert(table).values(scope="s", message_id=message_id, processed_at=func.now() - age)) - - -async def _message_ids(engine: AsyncEngine, table: Table) -> list[str]: - async with engine.connect() as conn: - return list((await conn.execute(select(table.c.message_id).order_by(table.c.message_id))).scalars()) - - -async def _cutoff( - session_maker: async_sessionmaker[AsyncSession], table: Table, retention: dt.timedelta -) -> dt.datetime: - async with session_maker() as session, session.begin(): - return await Inbox(session, table).cleanup_cutoff(retention=retention) - - -async def _delete_batch( - session_maker: async_sessionmaker[AsyncSession], table: Table, cutoff: dt.datetime, limit: int -) -> int: - async with session_maker() as session, session.begin(): - return await Inbox(session, table).delete_expired_batch(cutoff=cutoff, limit=limit) - - -async def test_cutoff_comes_from_the_database_clock( - tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - async with session_maker() as session, session.begin(): - cutoff = await Inbox(session, tables.inbox).cleanup_cutoff(retention=HOUR) - database_now = (await session.execute(select(func.now()))).scalar_one() - - assert cutoff == database_now - HOUR - - -async def test_only_markers_older_than_the_cutoff_are_deleted( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - await _add_marker(pg_engine, tables.inbox, "old", 2 * HOUR) - await _add_marker(pg_engine, tables.inbox, "recent", dt.timedelta(minutes=10)) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - - assert await _delete_batch(session_maker, tables.inbox, cutoff, 100) == 1 - assert await _message_ids(pg_engine, tables.inbox) == ["recent"] - - -async def test_each_batch_respects_its_limit( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - for index in range(5): - await _add_marker(pg_engine, tables.inbox, f"m-{index}", 2 * HOUR) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - - deleted = [await _delete_batch(session_maker, tables.inbox, cutoff, 2) for _ in range(4)] - - assert deleted == [2, 2, 1, 0] - - -async def test_a_run_keeps_its_cutoff_for_markers_added_later( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - await _add_marker(pg_engine, tables.inbox, "old", 2 * HOUR) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - async with pg_engine.begin() as conn: - await conn.execute( - insert(tables.inbox).values( - scope="s", message_id="after-cutoff", processed_at=cutoff + dt.timedelta(seconds=1) - ) - ) - - assert await _delete_batch(session_maker, tables.inbox, cutoff, 100) == 1 - assert await _message_ids(pg_engine, tables.inbox) == ["after-cutoff"] - - -async def test_committed_batches_survive_a_failed_batch( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - for index in range(4): - await _add_marker(pg_engine, tables.inbox, f"m-{index}", 2 * HOUR) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - - assert await _delete_batch(session_maker, tables.inbox, cutoff, 2) == 2 - - async def fail_after_delete(session: AsyncSession) -> None: - async with session.begin(): - assert await Inbox(session, tables.inbox).delete_expired_batch(cutoff=cutoff, limit=2) == 2 - raise _BatchError - - async with session_maker() as session: - with pytest.raises(_BatchError): - await fail_after_delete(session) - - assert len(await _message_ids(pg_engine, tables.inbox)) == 2 - assert await _delete_batch(session_maker, tables.inbox, cutoff, 2) == 2 - assert await _message_ids(pg_engine, tables.inbox) == [] - - -async def test_cleanup_skips_markers_locked_by_another_transaction( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - await _add_marker(pg_engine, tables.inbox, "locked", 2 * HOUR) - await _add_marker(pg_engine, tables.inbox, "free", 2 * HOUR) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - - async with session_maker() as holder, holder.begin(): - await holder.execute(select(tables.inbox).where(tables.inbox.c.message_id == "locked").with_for_update()) - deleted = await asyncio.wait_for(_delete_batch(session_maker, tables.inbox, cutoff, 100), timeout=5) - - assert deleted == 1 - assert await _message_ids(pg_engine, tables.inbox) == ["locked"] - - -async def test_cleanup_errors_reach_the_caller( - session_maker: async_sessionmaker[AsyncSession], -) -> None: - missing = make_inbox_table(MetaData(), table_name=f"missing_{uuid.uuid4().hex[:12]}") - async with session_maker() as session: - with pytest.raises(DBAPIError): - async with session.begin(): - await Inbox(session, missing).delete_expired_batch(cutoff=dt.datetime.now(tz=dt.UTC), limit=10) - - -async def test_an_expired_message_is_processed_again( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - await _add_marker(pg_engine, tables.inbox, "m-1", 2 * HOUR) - async with session_maker() as session, session.begin(): - assert not await Inbox(session, tables.inbox).claim(scope="s", message_id="m-1") - - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - assert await _delete_batch(session_maker, tables.inbox, cutoff, 100) == 1 - - async with session_maker() as session, session.begin(): - assert await Inbox(session, tables.inbox).claim(scope="s", message_id="m-1") - - -async def test_claim_during_uncommitted_cleanup_waits_then_claims( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - await _add_marker(pg_engine, tables.inbox, "m-1", 2 * HOUR) - cutoff = await _cutoff(session_maker, tables.inbox, HOUR) - - async with session_maker() as cleaner, session_maker() as consumer: - await cleaner.begin() - assert await Inbox(cleaner, tables.inbox).delete_expired_batch(cutoff=cutoff, limit=100) == 1 - - await consumer.begin() - pid = (await consumer.execute(text("SELECT pg_backend_pid()"))).scalar_one() - competing = asyncio.create_task(Inbox(consumer, tables.inbox).claim(scope="s", message_id="m-1")) - await wait_until_lock_wait(pg_engine, pid) - assert not competing.done() - - await cleaner.commit() - assert await competing - await consumer.commit() - - assert await _message_ids(pg_engine, tables.inbox) == ["m-1"] - - -async def test_cutoff_rejects_a_non_positive_retention( - tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - async with session_maker() as session, session.begin(): - with pytest.raises(ValueError, match="retention"): - await Inbox(session, tables.inbox).cleanup_cutoff(retention=dt.timedelta(0)) - - -async def test_delete_rejects_a_non_positive_limit( - tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - async with session_maker() as session, session.begin(): - with pytest.raises(ValueError, match="limit"): - await Inbox(session, tables.inbox).delete_expired_batch(cutoff=dt.datetime.now(tz=dt.UTC), limit=0) - - -async def test_cleanup_without_an_open_transaction_raises( - tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: - async with session_maker() as session: - inbox = Inbox(session, tables.inbox) - with pytest.raises(RuntimeError, match="open transaction"): - await inbox.cleanup_cutoff(retention=HOUR) - with pytest.raises(RuntimeError, match="open transaction"): - await inbox.delete_expired_batch(cutoff=dt.datetime.now(tz=dt.UTC), limit=10) - assert not session.in_transaction() diff --git a/tests/test_readme_examples.py b/tests/test_readme_examples.py index 6531864..0c31ca1 100644 --- a/tests/test_readme_examples.py +++ b/tests/test_readme_examples.py @@ -3,10 +3,10 @@ from db_retry import Transaction, postgres_retry from sqlalchemy import Table, func, insert, select, text -from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine -from sqlalchemy_inbox import Inbox -from tests.conftest import Tables, count_rows +from sqlalchemy_inbox import Inbox, InboxCleaner +from tests.conftest import PG_DSN, Tables, count_rows @dataclasses.dataclass(kw_only=True, frozen=True, slots=True) @@ -83,28 +83,7 @@ async def test_readme_processing_use_case_skips_a_duplicate_and_handles_the_next assert await count_rows(pg_engine, tables.business) == 2 -@dataclasses.dataclass(kw_only=True, frozen=True, slots=True) -class CleanupInboxUseCase: - transaction: Transaction - inbox: Inbox - - async def __call__(self, *, retention: dt.timedelta, batch_size: int) -> None: - async with self.transaction: - cutoff = await self.inbox.cleanup_cutoff(retention=retention) - while await self._delete_batch(cutoff=cutoff, limit=batch_size): - pass - - @postgres_retry - async def _delete_batch(self, *, cutoff: dt.datetime, limit: int) -> int: - async with self.transaction: - deleted = await self.inbox.delete_expired_batch(cutoff=cutoff, limit=limit) - await self.transaction.commit() - return deleted - - -async def test_readme_cleanup_use_case_deletes_every_expired_marker( - pg_engine: AsyncEngine, tables: Tables, session_maker: async_sessionmaker[AsyncSession] -) -> None: +async def test_readme_cleanup_cron_job_deletes_every_expired_marker(pg_engine: AsyncEngine, tables: Tables) -> None: async with pg_engine.begin() as conn: for index in range(5): await conn.execute( @@ -114,9 +93,12 @@ async def test_readme_cleanup_use_case_deletes_every_expired_marker( ) await conn.execute(insert(tables.inbox).values(scope="s", message_id="recent")) - session = session_maker() - use_case = CleanupInboxUseCase(transaction=Transaction(session=session), inbox=Inbox(session, tables.inbox)) - await use_case(retention=dt.timedelta(hours=1), batch_size=2) + engine = create_async_engine(PG_DSN) + try: + cleaner = InboxCleaner(engine=engine, table=tables.inbox, retention=dt.timedelta(hours=1)) + await cleaner.cleanup() + finally: + await engine.dispose() async with pg_engine.connect() as conn: remaining = list((await conn.execute(select(tables.inbox.c.message_id))).scalars())