From 8563948e70030ee693ae668fe013bebdcefeae8f Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Thu, 19 Feb 2026 18:12:59 -0800 Subject: [PATCH] feat(checkpoint-postgres): add extended methods and conformance tests Implements adelete_for_runs, acopy_thread, aprune (and sync counterparts) for AsyncPostgresSaver and PostgresSaver. Adds conformance test harness using langgraph-checkpoint-conformance as an editable dev dependency. Co-Authored-By: Claude Opus 4.6 --- .../langgraph/checkpoint/postgres/__init__.py | 136 +++++++++++++ .../langgraph/checkpoint/postgres/aio.py | 184 ++++++++++++++++++ libs/checkpoint-postgres/pyproject.toml | 2 + .../tests/test_conformance.py | 56 ++++++ libs/checkpoint-postgres/uv.lock | 31 +++ libs/langgraph/uv.lock | 2 + libs/prebuilt/uv.lock | 2 + 7 files changed, 413 insertions(+) create mode 100644 libs/checkpoint-postgres/tests/test_conformance.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 600127333..f7b2db41c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -390,6 +390,142 @@ class PostgresSaver(BasePostgresSaver): (str(thread_id),), ) + def delete_for_runs(self, run_ids: Sequence[str]) -> None: + """Delete all checkpoints and writes for the given run IDs. + + Args: + run_ids: The run IDs whose checkpoints should be deleted. + """ + if not run_ids: + return + with self._cursor(pipeline=True) as cur: + cur.execute( + """DELETE FROM checkpoint_writes + WHERE (thread_id, checkpoint_ns, checkpoint_id) IN ( + SELECT thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE metadata->>'run_id' = ANY(%s) + )""", + (list(run_ids),), + ) + cur.execute( + """DELETE FROM checkpoint_blobs + WHERE (thread_id, checkpoint_ns) IN ( + SELECT DISTINCT thread_id, checkpoint_ns + FROM checkpoints + WHERE metadata->>'run_id' = ANY(%s) + ) + AND NOT EXISTS ( + SELECT 1 FROM checkpoints c2 + WHERE c2.thread_id = checkpoint_blobs.thread_id + AND c2.checkpoint_ns = checkpoint_blobs.checkpoint_ns + AND (c2.metadata->>'run_id' IS NULL OR c2.metadata->>'run_id' != ALL(%s)) + AND c2.checkpoint->'channel_versions' ? checkpoint_blobs.channel + )""", + (list(run_ids), list(run_ids)), + ) + cur.execute( + "DELETE FROM checkpoints WHERE metadata->>'run_id' = ANY(%s)", + (list(run_ids),), + ) + + def copy_thread(self, source_thread_id: str, target_thread_id: str) -> None: + """Copy all checkpoints and writes from source thread to target thread. + + Args: + source_thread_id: The thread ID to copy from. + target_thread_id: The thread ID to copy to. + """ + with self._cursor(pipeline=True) as cur: + cur.execute( + """INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata) + SELECT %s, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata + FROM checkpoints + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + cur.execute( + """INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob) + SELECT %s, checkpoint_ns, channel, version, type, blob + FROM checkpoint_blobs + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + cur.execute( + """INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) + SELECT %s, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob + FROM checkpoint_writes + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + + def prune( + self, + thread_ids: Sequence[str], + *, + strategy: str = "keep_latest", + ) -> None: + """Prune checkpoints for given threads. + + Args: + thread_ids: The thread IDs to prune. + strategy: `"keep_latest"` keeps only the latest checkpoint per + thread+namespace. `"delete"` removes everything. + """ + if not thread_ids: + return + if strategy == "delete": + with self._cursor(pipeline=True) as cur: + cur.execute( + "DELETE FROM checkpoints WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + cur.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + cur.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + elif strategy == "keep_latest": + with self._cursor(pipeline=True) as cur: + cur.execute( + """DELETE FROM checkpoints + WHERE thread_id = ANY(%s) + AND (thread_id, checkpoint_ns, checkpoint_id) NOT IN ( + SELECT DISTINCT ON (thread_id, checkpoint_ns) + thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE thread_id = ANY(%s) + ORDER BY thread_id, checkpoint_ns, checkpoint_id DESC + )""", + (list(thread_ids), list(thread_ids)), + ) + cur.execute( + """DELETE FROM checkpoint_writes + WHERE thread_id = ANY(%s) + AND (thread_id, checkpoint_ns, checkpoint_id) NOT IN ( + SELECT thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE thread_id = ANY(%s) + )""", + (list(thread_ids), list(thread_ids)), + ) + cur.execute( + """DELETE FROM checkpoint_blobs + WHERE thread_id = ANY(%s) + AND NOT EXISTS ( + SELECT 1 FROM checkpoints c + WHERE c.thread_id = checkpoint_blobs.thread_id + AND c.checkpoint_ns = checkpoint_blobs.checkpoint_ns + )""", + (list(thread_ids),), + ) + @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: """Create a database cursor as a context manager. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5cd533f8d..40059613f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -349,6 +349,152 @@ class AsyncPostgresSaver(BasePostgresSaver): (str(thread_id),), ) + async def adelete_for_runs(self, run_ids: Sequence[str]) -> None: + """Delete all checkpoints and writes for the given run IDs. + + Args: + run_ids: The run IDs whose checkpoints should be deleted. + """ + if not run_ids: + return + async with self._cursor(pipeline=True) as cur: + # Delete writes associated with checkpoints that have matching run_ids + await cur.execute( + """DELETE FROM checkpoint_writes + WHERE (thread_id, checkpoint_ns, checkpoint_id) IN ( + SELECT thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE metadata->>'run_id' = ANY(%s) + )""", + (list(run_ids),), + ) + # Delete blobs associated with checkpoints that have matching run_ids + # We need to delete blobs for channels referenced by these checkpoints + await cur.execute( + """DELETE FROM checkpoint_blobs + WHERE (thread_id, checkpoint_ns) IN ( + SELECT DISTINCT thread_id, checkpoint_ns + FROM checkpoints + WHERE metadata->>'run_id' = ANY(%s) + ) + AND NOT EXISTS ( + SELECT 1 FROM checkpoints c2 + WHERE c2.thread_id = checkpoint_blobs.thread_id + AND c2.checkpoint_ns = checkpoint_blobs.checkpoint_ns + AND (c2.metadata->>'run_id' IS NULL OR c2.metadata->>'run_id' != ALL(%s)) + AND c2.checkpoint->'channel_versions' ? checkpoint_blobs.channel + )""", + (list(run_ids), list(run_ids)), + ) + # Delete the checkpoints themselves + await cur.execute( + "DELETE FROM checkpoints WHERE metadata->>'run_id' = ANY(%s)", + (list(run_ids),), + ) + + async def acopy_thread(self, source_thread_id: str, target_thread_id: str) -> None: + """Copy all checkpoints and writes from source thread to target thread. + + Args: + source_thread_id: The thread ID to copy from. + target_thread_id: The thread ID to copy to. + """ + async with self._cursor(pipeline=True) as cur: + # Copy checkpoints + await cur.execute( + """INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata) + SELECT %s, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata + FROM checkpoints + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + # Copy blobs + await cur.execute( + """INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob) + SELECT %s, checkpoint_ns, channel, version, type, blob + FROM checkpoint_blobs + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + # Copy writes + await cur.execute( + """INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) + SELECT %s, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob + FROM checkpoint_writes + WHERE thread_id = %s + ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING""", + (target_thread_id, source_thread_id), + ) + + async def aprune( + self, + thread_ids: Sequence[str], + *, + strategy: str = "keep_latest", + ) -> None: + """Prune checkpoints for given threads. + + Args: + thread_ids: The thread IDs to prune. + strategy: `"keep_latest"` keeps only the latest checkpoint per + thread+namespace. `"delete"` removes everything. + """ + if not thread_ids: + return + if strategy == "delete": + async with self._cursor(pipeline=True) as cur: + await cur.execute( + "DELETE FROM checkpoints WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + await cur.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + await cur.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)", + (list(thread_ids),), + ) + elif strategy == "keep_latest": + async with self._cursor(pipeline=True) as cur: + # Delete non-latest checkpoints + await cur.execute( + """DELETE FROM checkpoints + WHERE thread_id = ANY(%s) + AND (thread_id, checkpoint_ns, checkpoint_id) NOT IN ( + SELECT DISTINCT ON (thread_id, checkpoint_ns) + thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE thread_id = ANY(%s) + ORDER BY thread_id, checkpoint_ns, checkpoint_id DESC + )""", + (list(thread_ids), list(thread_ids)), + ) + # Delete writes for removed checkpoints (keep only writes for remaining checkpoints) + await cur.execute( + """DELETE FROM checkpoint_writes + WHERE thread_id = ANY(%s) + AND (thread_id, checkpoint_ns, checkpoint_id) NOT IN ( + SELECT thread_id, checkpoint_ns, checkpoint_id + FROM checkpoints + WHERE thread_id = ANY(%s) + )""", + (list(thread_ids), list(thread_ids)), + ) + # Clean up orphaned blobs + await cur.execute( + """DELETE FROM checkpoint_blobs + WHERE thread_id = ANY(%s) + AND NOT EXISTS ( + SELECT 1 FROM checkpoints c + WHERE c.thread_id = checkpoint_blobs.thread_id + AND c.checkpoint_ns = checkpoint_blobs.checkpoint_ns + )""", + (list(thread_ids),), + ) + @asynccontextmanager async def _cursor( self, *, pipeline: bool = False @@ -578,5 +724,43 @@ class AsyncPostgresSaver(BasePostgresSaver): self.adelete_thread(thread_id), self.loop ).result() + def delete_for_runs(self, run_ids: Sequence[str]) -> None: + """Delete all checkpoints and writes for the given run IDs. + + Args: + run_ids: The run IDs whose checkpoints should be deleted. + """ + return asyncio.run_coroutine_threadsafe( + self.adelete_for_runs(run_ids), self.loop + ).result() + + def copy_thread(self, source_thread_id: str, target_thread_id: str) -> None: + """Copy all checkpoints and writes from source thread to target thread. + + Args: + source_thread_id: The thread ID to copy from. + target_thread_id: The thread ID to copy to. + """ + return asyncio.run_coroutine_threadsafe( + self.acopy_thread(source_thread_id, target_thread_id), self.loop + ).result() + + def prune( + self, + thread_ids: Sequence[str], + *, + strategy: str = "keep_latest", + ) -> None: + """Prune checkpoints for given threads. + + Args: + thread_ids: The thread IDs to prune. + strategy: `"keep_latest"` keeps only the latest checkpoint per + thread+namespace. `"delete"` removes everything. + """ + return asyncio.run_coroutine_threadsafe( + self.aprune(thread_ids, strategy=strategy), self.loop + ).result() + __all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index f9e1ddc17..6898487a0 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -32,6 +32,7 @@ test = [ "pytest-mock", "psycopg[binary]", "langgraph-checkpoint", + "langgraph-checkpoint-conformance", "pytest-watcher", ] lint = [ @@ -49,6 +50,7 @@ default-groups = ['dev'] [tool.uv.sources] langgraph-checkpoint = { path = "../checkpoint", editable = true } +langgraph-checkpoint-conformance = { path = "../checkpoint-conformance", editable = true } [tool.hatch.build.targets.wheel] include = ["langgraph"] diff --git a/libs/checkpoint-postgres/tests/test_conformance.py b/libs/checkpoint-postgres/tests/test_conformance.py new file mode 100644 index 000000000..4f195edfd --- /dev/null +++ b/libs/checkpoint-postgres/tests/test_conformance.py @@ -0,0 +1,56 @@ +"""Conformance tests for AsyncPostgresSaver.""" +# mypy: disable-error-code="import-untyped" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from uuid import uuid4 + +import pytest +from langgraph.checkpoint.conformance import checkpointer_test, validate +from langgraph.checkpoint.conformance.report import ProgressCallbacks +from psycopg import AsyncConnection +from psycopg.rows import dict_row + +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from tests.conftest import DEFAULT_POSTGRES_URI + + +async def pg_lifespan() -> AsyncGenerator[None, None]: + """No-op lifespan; databases are created per-checkpointer instance.""" + yield + + +@checkpointer_test(name="AsyncPostgresSaver", lifespan=pg_lifespan) +async def postgres_checkpointer() -> AsyncGenerator[AsyncPostgresSaver, None]: + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI + database, + autocommit=True, + prepare_threshold=0, + row_factory=dict_row, + ) as conn: + saver = AsyncPostgresSaver(conn) + await saver.setup() + yield saver + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@pytest.mark.asyncio +async def test_full_conformance() -> None: + """AsyncPostgresSaver passes ALL conformance tests.""" + report = await validate( + postgres_checkpointer, + progress=ProgressCallbacks.verbose(), + ) + report.print_report() + assert report.passed_all(), f"Conformance failed: {report.to_dict()}" diff --git a/libs/checkpoint-postgres/uv.lock b/libs/checkpoint-postgres/uv.lock index 9923453e2..eaaf89174 100644 --- a/libs/checkpoint-postgres/uv.lock +++ b/libs/checkpoint-postgres/uv.lock @@ -304,6 +304,33 @@ test = [ { name = "redis" }, ] +[[package]] +name = "langgraph-checkpoint-conformance" +version = "0.0.1" +source = { editable = "../checkpoint-conformance" } +dependencies = [ + { name = "langgraph-checkpoint" }, +] + +[package.metadata] +requires-dist = [{ name = "langgraph-checkpoint", specifier = ">=2.0.0" }] + +[package.metadata.requires-dev] +dev = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "ruff" }, + { name = "ty" }, +] +lint = [ + { name = "ruff" }, + { name = "ty" }, +] +test = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, +] + [[package]] name = "langgraph-checkpoint-postgres" version = "3.0.4" @@ -320,6 +347,7 @@ dev = [ { name = "anyio" }, { name = "codespell" }, { name = "langgraph-checkpoint" }, + { name = "langgraph-checkpoint-conformance" }, { name = "mypy" }, { name = "psycopg", extra = ["binary"] }, { name = "pytest" }, @@ -336,6 +364,7 @@ lint = [ test = [ { name = "anyio" }, { name = "langgraph-checkpoint" }, + { name = "langgraph-checkpoint-conformance" }, { name = "psycopg", extra = ["binary"] }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -356,6 +385,7 @@ dev = [ { name = "anyio" }, { name = "codespell" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "mypy" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, @@ -372,6 +402,7 @@ lint = [ test = [ { name = "anyio" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, { name = "pytest-asyncio" }, diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index a5186b591..a9c2a44b5 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -1617,6 +1617,7 @@ dev = [ { name = "anyio" }, { name = "codespell" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "mypy" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, @@ -1633,6 +1634,7 @@ lint = [ test = [ { name = "anyio" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, { name = "pytest-asyncio" }, diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index 109c99461..63405b684 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -421,6 +421,7 @@ dev = [ { name = "anyio" }, { name = "codespell" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "mypy" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, @@ -437,6 +438,7 @@ lint = [ test = [ { name = "anyio" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, + { name = "langgraph-checkpoint-conformance", editable = "../checkpoint-conformance" }, { name = "psycopg", extras = ["binary"] }, { name = "pytest" }, { name = "pytest-asyncio" },