diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a794671..ccc6602 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,4 +32,27 @@ jobs: run: uv sync --extra test - name: Run tests - run: uv run pytest + run: uv run pytest --ignore=tests/test_strategy_postgres.py + + test-postgres: + name: Postgres Integration Tests + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v6.0.2 + + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + with: + enable-cache: true + + - name: Set up Python 3.12 + uses: actions/setup-python@v6.2.0 + with: + python-version: "3.12" + + - name: Install dependencies + run: uv sync --extra test --extra test-postgres + + - name: Run Postgres integration tests + run: uv run pytest tests/test_strategy_postgres.py -v diff --git a/pyproject.toml b/pyproject.toml index af40030..78fdeef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,11 @@ test = [ "langgraph-checkpoint-sqlite>=3.1.0", "aiosqlite>=0.20", ] +test-postgres = [ + "langgraph-checkpoint-postgres>=3.1.0", + "psycopg[binary]>=3.1", + "testcontainers[postgres]>=4.0", +] benchmark = [ "langgraph-checkpoint-sqlite>=3.0.3", "langgraph-checkpoint-postgres>=3.0.5", diff --git a/tests/test_policy_resolver.py b/tests/test_policy_resolver.py index 26c31bd..0639d86 100644 --- a/tests/test_policy_resolver.py +++ b/tests/test_policy_resolver.py @@ -128,3 +128,53 @@ def resolver(tid): assert "exempt" not in result.deleted_thread_ids assert "expired" in result.deleted_thread_ids + + +def test_fresh_policy_objects_per_thread_applied_correctly(): + tuples = [ + make_checkpoint_tuple("strict", None, iso_ts(-70)), + make_checkpoint_tuple("lenient", None, iso_ts(-70)), + ] + cp = MockCP(tuples) + + def resolver(tid): + if tid == "strict": + return TTLPolicy(idle_ttl_seconds=60) + return TTLPolicy(idle_ttl_seconds=3600) + + sweeper = Sweeper( + cp, + TTLPolicy(idle_ttl_seconds=9999), + policy_resolver=resolver, + safe_delete=False, + _strategy=list_strategy(tuples), + ) + result = sweeper.sweep() + + assert "strict" in result.deleted_thread_ids + assert "lenient" not in result.deleted_thread_ids + + +def test_many_threads_fresh_policies_all_evaluated_correctly(): + thread_ids = [f"t{i}" for i in range(50)] + tuples = [make_checkpoint_tuple(tid, None, iso_ts(-70)) for tid in thread_ids] + cp = MockCP(tuples) + + strict_ids = {tid for tid in thread_ids if int(tid[1:]) % 2 == 0} + + def resolver(tid): + if tid in strict_ids: + return TTLPolicy(idle_ttl_seconds=60) + return TTLPolicy(idle_ttl_seconds=3600) + + sweeper = Sweeper( + cp, + TTLPolicy(idle_ttl_seconds=9999), + policy_resolver=resolver, + safe_delete=False, + _strategy=list_strategy(tuples), + ) + result = sweeper.sweep() + + deleted = set(result.deleted_thread_ids) + assert deleted == strict_ids diff --git a/tests/test_strategy_postgres.py b/tests/test_strategy_postgres.py index 429af00..74f1c52 100644 --- a/tests/test_strategy_postgres.py +++ b/tests/test_strategy_postgres.py @@ -1,13 +1,18 @@ """Integration tests for PostgresSaver / AsyncPostgresSaver strategies. -Requires Docker and testcontainers: - pip install testcontainers[postgres] +Requires the test-postgres extra: + uv sync --extra test --extra test-postgres """ import time +from typing import TypedDict import pytest from langgraph_ephemeral_checkpointer import Sweeper, TTLPolicy +from langgraph_ephemeral_checkpointer._coordination import ( + AsyncAdvisoryLock, + SyncAdvisoryLock, +) from langgraph_ephemeral_checkpointer._strategies import detect from langgraph_ephemeral_checkpointer._strategies.postgres import ( AsyncPostgresStrategy, @@ -16,8 +21,6 @@ testcontainers = pytest.importorskip("testcontainers") -from typing import TypedDict # noqa: E402 - from langgraph.graph import END, START, StateGraph # noqa: E402 from testcontainers.postgres import PostgresContainer # noqa: E402 # pyrefly: ignore[missing-import] @@ -25,6 +28,7 @@ class _State(TypedDict): x: int + def _build_graph(saver): builder = StateGraph(_State) # pyrefly: ignore[bad-specialization] builder.add_node("inc", lambda s: {"x": s["x"] + 1}) # pyrefly: ignore[bad-argument-type] @@ -32,27 +36,100 @@ def _build_graph(saver): builder.add_edge("inc", END) return builder.compile(checkpointer=saver) + +def _row_count(saver, table: str, thread_id: str) -> int: + with saver._cursor() as cur: + cur.execute(f"SELECT COUNT(*) AS c FROM {table} WHERE thread_id = %s", (thread_id,)) + return cur.fetchone()["c"] + + +async def _arow_count(saver, table: str, thread_id: str) -> int: + async with saver._cursor() as cur: + await cur.execute(f"SELECT COUNT(*) AS c FROM {table} WHERE thread_id = %s", (thread_id,)) + return (await cur.fetchone())["c"] + + @pytest.fixture(scope="module") def postgres_dsn(): with PostgresContainer("postgres:16-alpine") as pg: - yield pg.get_connection_url().replace("psycopg2", "psycopg") + yield pg.get_connection_url().replace("postgresql+psycopg2", "postgresql") + def test_detect_returns_postgres_strategy(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() - strategy = detect(saver) - assert isinstance(strategy, PostgresStrategy) + assert isinstance(detect(saver), PostgresStrategy) + + +def test_detect_returns_async_postgres_strategy(postgres_dsn): + import asyncio + + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async def _inner(): + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + assert isinstance(detect(saver), AsyncPostgresStrategy) + + asyncio.run(_inner()) + def test_postgres_collect_empty(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() - strategy = PostgresStrategy(saver) - threads, _ = strategy.collect(None) + threads, cursor = PostgresStrategy(saver).collect(None) assert isinstance(threads, dict) + assert cursor is None or isinstance(cursor, str) + + +def test_postgres_collect_returns_thread_timestamps(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "collect_sync"}}) + + threads, cursor = PostgresStrategy(saver).collect(None) + + assert "collect_sync" in threads + ts = threads["collect_sync"] + assert ts.earliest_id <= ts.latest_id + assert cursor == ts.latest_id + + +@pytest.mark.asyncio +async def test_async_postgres_collect_returns_thread_timestamps(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "collect_async"}}) + + threads, cursor = await AsyncPostgresStrategy(saver).acollect(None) + + assert "collect_async" in threads + ts = threads["collect_async"] + assert ts.earliest_id <= ts.latest_id + assert cursor == ts.latest_id + + +def test_postgres_sweep_keeps_active_thread(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "active_sync"}}) + + result = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60)).sweep() + assert "active_sync" not in result.deleted_thread_ids + def test_postgres_sweep_deletes_expired(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver @@ -60,43 +137,186 @@ def test_postgres_sweep_deletes_expired(postgres_dsn): with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() graph = _build_graph(saver) - graph.invoke({"x": 0}, {"configurable": {"thread_id": "t1"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "expired_sync"}}) sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) - - result = sweeper.sweep() - assert "t1" not in result.deleted_thread_ids + assert "expired_sync" not in sweeper.sweep().deleted_thread_ids time.sleep(1.1) - result = sweeper.sweep() - assert "t1" in result.deleted_thread_ids + assert "expired_sync" in sweeper.sweep().deleted_thread_ids -def test_detect_returns_async_postgres_strategy(postgres_dsn): + +@pytest.mark.asyncio +async def test_async_postgres_sweep_deletes_expired(postgres_dsn): from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver - async def _inner(): - async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: - await saver.setup() - strategy = detect(saver) - assert isinstance(strategy, AsyncPostgresStrategy) + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "expired_async"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + assert "expired_async" not in (await sweeper.asweep()).deleted_thread_ids + + time.sleep(1.1) + assert "expired_async" in (await sweeper.asweep()).deleted_thread_ids + + +def test_postgres_batch_delete_clears_all_tables(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "all_tables_sync"}}) + + assert _row_count(saver, "checkpoints", "all_tables_sync") > 0 + assert _row_count(saver, "checkpoint_blobs", "all_tables_sync") > 0 + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = sweeper.sweep() + + assert "all_tables_sync" in result.deleted_thread_ids + assert _row_count(saver, "checkpoints", "all_tables_sync") == 0 + assert _row_count(saver, "checkpoint_writes", "all_tables_sync") == 0 + assert _row_count(saver, "checkpoint_blobs", "all_tables_sync") == 0 - import asyncio - asyncio.run(_inner()) @pytest.mark.asyncio -async def test_async_postgres_sweep(postgres_dsn): +async def test_async_postgres_batch_delete_clears_all_tables(postgres_dsn): from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: await saver.setup() graph = _build_graph(saver) - await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "t1"}}) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "all_tables_async"}}) - sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + assert await _arow_count(saver, "checkpoints", "all_tables_async") > 0 + assert await _arow_count(saver, "checkpoint_blobs", "all_tables_async") > 0 + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) result = await sweeper.asweep() - assert "t1" not in result.deleted_thread_ids + + assert "all_tables_async" in result.deleted_thread_ids + assert await _arow_count(saver, "checkpoints", "all_tables_async") == 0 + assert await _arow_count(saver, "checkpoint_writes", "all_tables_async") == 0 + assert await _arow_count(saver, "checkpoint_blobs", "all_tables_async") == 0 + + +def test_postgres_sweep_no_cross_contamination(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "survivor"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "victim"}}) time.sleep(1.1) - result = await sweeper.asweep() - assert "t1" in result.deleted_thread_ids + graph.invoke({"x": 1}, {"configurable": {"thread_id": "survivor"}}) + + result = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)).sweep() + + assert "victim" in result.deleted_thread_ids + assert "survivor" not in result.deleted_thread_ids + assert _row_count(saver, "checkpoints", "survivor") > 0 + assert _row_count(saver, "checkpoint_blobs", "survivor") > 0 + + +def test_postgres_hard_age_ttl(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "hard_age"}}) + + sweeper = Sweeper(saver, TTLPolicy(hard_age_ttl_seconds=1)) + assert "hard_age" not in sweeper.sweep().deleted_thread_ids + + time.sleep(1.1) + assert "hard_age" in sweeper.sweep().deleted_thread_ids + + +def test_advisory_lock_type_for_postgres(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + assert isinstance(sweeper._advisory_lock, SyncAdvisoryLock) + + +@pytest.mark.asyncio +async def test_advisory_lock_type_for_async_postgres(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + assert isinstance(sweeper._advisory_lock, AsyncAdvisoryLock) + + +def test_postgres_advisory_lock_second_sweeper_skips(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as s1: + with PostgresSaver.from_conn_string(postgres_dsn) as s2: + s1.setup() + + sw1 = Sweeper(s1, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + sw2 = Sweeper(s2, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + + assert sw1._advisory_lock.try_acquire() + try: + result = sw2.sweep() + assert result.deleted_thread_ids == [] + assert result.active_thread_count == 0 + finally: + sw1._advisory_lock.release() + + assert sw2.sweep() is not None + + +@pytest.mark.asyncio +async def test_async_postgres_advisory_lock_second_sweeper_skips(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as s1: + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as s2: + await s1.setup() + + sw1 = Sweeper(s1, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + sw2 = Sweeper(s2, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + + assert await sw1._advisory_lock.atry_acquire() + try: + result = await sw2.asweep() + assert result.deleted_thread_ids == [] + assert result.active_thread_count == 0 + finally: + await sw1._advisory_lock.arelease() + + +def test_postgres_incremental_collect(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "cursor_t1"}}) + + strategy = PostgresStrategy(saver) + threads1, cursor1 = strategy.collect(None) + assert "cursor_t1" in threads1 + assert cursor1 is not None + + threads2, _ = strategy.collect(cursor1) + assert "cursor_t1" not in threads2 + + graph.invoke({"x": 0}, {"configurable": {"thread_id": "cursor_t2"}}) + threads3, _ = strategy.collect(cursor1) + assert "cursor_t2" in threads3 + assert "cursor_t1" not in threads3 diff --git a/tests/test_strategy_sqlite.py b/tests/test_strategy_sqlite.py index 4e19168..61079c6 100644 --- a/tests/test_strategy_sqlite.py +++ b/tests/test_strategy_sqlite.py @@ -112,3 +112,66 @@ async def test_async_sweep_deletes_expired(): time.sleep(1.1) result = await sweeper.asweep() assert "t1" in result.deleted_thread_ids + + +def test_sqlite_batch_delete_clears_writes_table(): + with SqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "t_writes"}}) + graph.invoke({"x": 1}, {"configurable": {"thread_id": "t_writes"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = sweeper.sweep() + + assert "t_writes" in result.deleted_thread_ids + assert saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("t_writes",) + ).fetchone()[0] == 0 + assert saver.conn.execute( + "SELECT COUNT(*) FROM writes WHERE thread_id = ?", ("t_writes",) + ).fetchone()[0] == 0 + + +@pytest.mark.asyncio +async def test_async_sqlite_batch_delete_clears_writes_table(): + pytest.importorskip("aiosqlite") + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "t_writes_async"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = await sweeper.asweep() + + assert "t_writes_async" in result.deleted_thread_ids + + async with saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("t_writes_async",) + ) as cur: + assert (await cur.fetchone())[0] == 0 + + async with saver.conn.execute( + "SELECT COUNT(*) FROM writes WHERE thread_id = ?", ("t_writes_async",) + ) as cur: + assert (await cur.fetchone())[0] == 0 + + +def test_sqlite_sweep_no_cross_contamination(): + with SqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "old"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "new"}}) + + time.sleep(1.1) + graph.invoke({"x": 1}, {"configurable": {"thread_id": "new"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + result = sweeper.sweep() + + assert "old" in result.deleted_thread_ids + assert "new" not in result.deleted_thread_ids + assert saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("new",) + ).fetchone()[0] > 0 diff --git a/tests/test_sweeper.py b/tests/test_sweeper.py index 06f90a8..03a6727 100644 --- a/tests/test_sweeper.py +++ b/tests/test_sweeper.py @@ -161,22 +161,42 @@ async def test_start_twice_raises(): await sweeper.stop() @pytest.mark.asyncio -async def test_loop_fails_on_exception(): - """An exception in asweep() fails the background task.""" - async def failing_asweep(self, *, dry_run=False): - raise RuntimeError("sweep failed") +async def test_loop_continues_after_exception(): + first_done = asyncio.Event() + + async def maybe_fail(self, *, dry_run=False): + first_done.set() + raise RuntimeError("transient failure") cp = MockCheckpointer([]) sweeper = Sweeper(cp, TTLPolicy(idle_ttl_seconds=60), _strategy=list_strategy([])) - with patch.object(Sweeper, "asweep", failing_asweep): + with patch.object(Sweeper, "asweep", maybe_fail): await sweeper.start(interval_seconds=600) - await asyncio.sleep(0.05) + await asyncio.wait_for(first_done.wait(), timeout=2.0) assert sweeper._task is not None - assert sweeper._task.done() - with pytest.raises(RuntimeError, match="sweep failed"): - sweeper._task.result() + assert not sweeper._task.done() + await sweeper.stop() + + +@pytest.mark.asyncio +async def test_loop_logs_exception(caplog): + import logging + + async def failing_asweep(self, *, dry_run=False): + raise RuntimeError("database timeout") + + cp = MockCheckpointer([]) + sweeper = Sweeper(cp, TTLPolicy(idle_ttl_seconds=60), _strategy=list_strategy([])) + + with patch.object(Sweeper, "asweep", failing_asweep): + with caplog.at_level(logging.ERROR, logger="langgraph_ephemeral_checkpointer.sweeper"): + await sweeper.start(interval_seconds=600) + await asyncio.sleep(0.05) + await sweeper.stop() + + assert any("Sweep cycle failed" in r.message for r in caplog.records) @pytest.mark.asyncio async def test_stop_before_start_is_noop():