diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5ab8d5994..2f461e75f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -400,36 +400,42 @@ class AsyncPostgresSaver(BasePostgresSaver): channel: str, cur: Any, ) -> list[Any]: - """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async).""" + """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async). + + Two queries instead of a recursive CTE — see sync version for rationale. + """ await cur.execute( - """ - WITH RECURSIVE chain(cid, depth) AS ( - SELECT parent_checkpoint_id, 0 - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s - UNION ALL - SELECT c.parent_checkpoint_id, ch.depth + 1 - FROM checkpoints c - JOIN chain ch ON c.checkpoint_id = ch.cid - WHERE ch.cid IS NOT NULL - ) - SELECT cw.type, cw.blob - FROM checkpoint_writes cw - JOIN chain ON cw.checkpoint_id = chain.cid - WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s - ORDER BY chain.depth DESC, cw.task_id, cw.idx - """, - ( - thread_id, - checkpoint_ns, - checkpoint_id, - thread_id, - checkpoint_ns, - channel, - ), + "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " + "WHERE thread_id = %s AND checkpoint_ns = %s", + (thread_id, checkpoint_ns), + ) + parent_map: dict[str, str | None] = { + row["checkpoint_id"]: row["parent_checkpoint_id"] + for row in await cur.fetchall() + } + ancestor_ids: list[str] = [] + cid: str | None = parent_map.get(checkpoint_id) + while cid is not None: + ancestor_ids.append(cid) + cid = parent_map.get(cid) + if not ancestor_ids: + return [] + await cur.execute( + "SELECT checkpoint_id, type, blob FROM checkpoint_writes " + "WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s " + " AND checkpoint_id = ANY(%s) " + "ORDER BY task_id, idx", + (thread_id, checkpoint_ns, channel, ancestor_ids), ) rows = await cur.fetchall() - return [self.serde.loads_typed((row["type"], row["blob"])) for row in rows] + writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list) + for row in rows: + writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"])) + result = [] + for cid in reversed(ancestor_ids): + for type_tag, blob in writes_by_cp.get(cid, []): + result.append(self.serde.loads_typed((type_tag, blob))) + return result async def aget_channel_writes( self, config: RunnableConfig, channel: str diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 66eacb082..2e679b88e 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -2,6 +2,7 @@ from __future__ import annotations import random import warnings +from collections import defaultdict from collections.abc import Sequence from importlib.metadata import version as get_version from typing import Any, cast @@ -222,37 +223,46 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): channel: str, cur: Any, ) -> list[Any]: - """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.""" + """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest. + + Two queries instead of a recursive CTE: + 1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, just IDs. + 2. Walk the ancestor chain in Python, then fetch writes with a plain ANY() filter. + """ cur.execute( - """ - WITH RECURSIVE chain(cid, depth) AS ( - SELECT parent_checkpoint_id, 0 - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s - UNION ALL - SELECT c.parent_checkpoint_id, ch.depth + 1 - FROM checkpoints c - JOIN chain ch ON c.checkpoint_id = ch.cid - WHERE ch.cid IS NOT NULL - ) - SELECT cw.type, cw.blob - FROM checkpoint_writes cw - JOIN chain ON cw.checkpoint_id = chain.cid - WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s - ORDER BY chain.depth DESC, cw.task_id, cw.idx - """, - ( - thread_id, - checkpoint_ns, - checkpoint_id, - thread_id, - checkpoint_ns, - channel, - ), + "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " + "WHERE thread_id = %s AND checkpoint_ns = %s", + (thread_id, checkpoint_ns), ) - return [ - self.serde.loads_typed((row["type"], row["blob"])) for row in cur.fetchall() - ] + parent_map: dict[str, str | None] = { + row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall() + } + # Walk newest→oldest starting from the current checkpoint's parent. + # Writes stored under checkpoint C produced the state *after* C, so we + # want ancestors of the current checkpoint (not the checkpoint itself). + ancestor_ids: list[str] = [] + cid: str | None = parent_map.get(checkpoint_id) + while cid is not None: + ancestor_ids.append(cid) + cid = parent_map.get(cid) + if not ancestor_ids: + return [] + cur.execute( + "SELECT checkpoint_id, type, blob FROM checkpoint_writes " + "WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s " + " AND checkpoint_id = ANY(%s) " + "ORDER BY task_id, idx", + (thread_id, checkpoint_ns, channel, ancestor_ids), + ) + writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list) + for row in cur.fetchall(): + writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"])) + # ancestor_ids is newest→oldest; replay oldest→newest + result = [] + for cid in reversed(ancestor_ids): + for type_tag, blob in writes_by_cp.get(cid, []): + result.append(self.serde.loads_typed((type_tag, blob))) + return result def _dump_blobs( self,