WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 17:57:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(checkpoint-postgres): single-roundtrip CTE query + write-ordering fix

Browse Source 
Replaces the three sequential SELECT roundtrips in _get_channel_writes_history
(checkpoints, checkpoint_writes, checkpoint_blobs) with one combined
UNION ALL query tagged by a _kind discriminator column. Both sync and async
paths now do one execute + one fetchall regardless of pipeline mode.

_build_delta_channel_writes_history is updated to accept the single tagged
rows list and dispatch on _kind while building its lookup dicts; the three
old SQL constants are removed.

Also fixes write-collection ordering in _build_delta_channel_writes_history:
the seed-terminator blob check previously fired before collecting that
ancestor's writes, silently dropping the transition writes needed to
reconstruct the child's state. Writes are now collected first.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-30 14:49:05 -04:00
co-authored by Claude Sonnet 4.6
parent 88e06f16ca
commit 54ee1b5b1e
3 changed files with 95 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+21 -16
View File
@@ -26,9 +26,7 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -442,10 +440,10 @@ class PostgresSaver(BasePostgresSaver):
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`) each filtered by `(thread_id, checkpoint_ns)` and
the per-table key. Plain SELECTs let the planner pick straight index
scans; rationale + benchmark in `notes/delta_channel_query_bench.md`.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip. Rationale + benchmark in
`notes/delta_channel_query_bench.md`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -459,18 +457,25 @@ class PostgresSaver(BasePostgresSaver):
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
with self._cursor() as cur:
cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
parents_rows = cur.fetchall()
cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
writes_rows = cur.fetchall()
cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
blobs_rows = cur.fetchall()
cur.execute(
SELECT_DELTA_COMBINED_SQL,
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
thread_id,
checkpoint_ns,
channel,
),
)
rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
rows=rows,
)
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+18 -18
View File
@@ -26,9 +26,7 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -403,8 +401,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`); rows assembled by the shared pure helper on
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; rows assembled by the shared pure helper on
`BasePostgresSaver`. Rationale + benchmark in
`notes/delta_channel_query_bench.md`.
"""
@@ -418,23 +417,24 @@ class AsyncPostgresSaver(BasePostgresSaver):
checkpoint_id = target.config["configurable"]["checkpoint_id"]
async with self._cursor() as cur:
await cur.execute(
SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns)
SELECT_DELTA_COMBINED_SQL,
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
thread_id,
checkpoint_ns,
channel,
),
)
parents_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel)
)
writes_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel)
)
blobs_rows = await cur.fetchall()
rows = await cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
rows=rows,
)
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+56 -60
View File
@@ -155,26 +155,38 @@ INSERT_CHECKPOINT_WRITES_SQL = """
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
# DeltaChannel reconstruction: three plain indexed SELECTs per channel.
# DeltaChannel reconstruction: one combined CTE+UNION ALL query per channel.
# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE
# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are
# 3x-100x faster in the realistic depth range and the Python walk is O(n).
SELECT_DELTA_PARENTS_SQL = """
SELECT checkpoint_id,
# The three plain SELECTs are further collapsed into one UNION ALL query so
# that only one roundtrip is needed per channel reconstruction.
#
# Parameter order: (channel, thread_id, checkpoint_ns,
# thread_id, checkpoint_ns, channel,
# thread_id, checkpoint_ns, channel)
SELECT_DELTA_COMBINED_SQL = """
SELECT 'p'::text AS _kind,
checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver
checkpoint -> 'channel_versions' ->> %s AS ver,
NULL::text AS type,
NULL::bytea AS blob,
NULL::text AS task_id,
NULL::int AS idx,
NULL::text AS version
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
"""
SELECT_DELTA_WRITES_SQL = """
SELECT checkpoint_id, type, blob, task_id, idx
UNION ALL
SELECT 'w',
checkpoint_id, NULL, NULL,
type, blob, task_id, idx, NULL
FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
SELECT_DELTA_BLOBS_SQL = """
SELECT version, type, blob
UNION ALL
SELECT 'b',
NULL, NULL, NULL,
type, blob, NULL, NULL, version
FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
@@ -229,15 +241,13 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
*,
channel: str,
target_id: str,
parents_rows: Sequence[Any],
writes_rows: Sequence[Any],
blobs_rows: Sequence[Any],
rows: Sequence[Any],
) -> _ChannelWritesHistory:
"""Reconstruct one delta channel's history from rows of the three SELECTs.
"""Reconstruct one delta channel's history from the combined UNION ALL rows.
Pure data transform shared by sync (`PostgresSaver`) and async
(`AsyncPostgresSaver`); both paths run the queries themselves and
feed the rows here.
(`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL`
and feed the tagged rows here.
Walk is newest → oldest from the target's parent. A non-sentinel
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
@@ -248,10 +258,27 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
"""
parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {}
for r in parents_rows:
cid = r["checkpoint_id"]
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
blob_by_ver: dict[str, tuple[str, bytes]] = {}
for r in rows:
kind = r["_kind"]
if kind == "p":
cid = r["checkpoint_id"]
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
elif kind == "w":
cid = r["checkpoint_id"]
writes_by_cid.setdefault(cid, []).append(
(r["type"], r["blob"], r["task_id"], r["idx"])
)
else: # kind == "b"
blob_by_ver[r["version"]] = (r["type"], r["blob"])
# Sort writes within each checkpoint (task_id DESC, idx DESC) to match
# the prior CTE ordering — newest write first per ancestor.
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
ancestors: list[str] = []
cid = parent_of.get(target_id)
@@ -260,55 +287,24 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
cid = parent_of.get(cid)
if not ancestors:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
ancestor_set = set(ancestors)
# Group writes by ancestor cid; sort within (task_id DESC, idx DESC)
# to match the prior CTE ordering — newest write first per ancestor.
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
for r in writes_rows:
cid = r["checkpoint_id"]
if cid not in ancestor_set:
continue
writes_by_cid.setdefault(cid, []).append(
(r["type"], r["blob"], r["task_id"], r["idx"])
)
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
blob_by_ver: dict[str, tuple[str, bytes]] = {
r["version"]: (r["type"], r["blob"]) for r in blobs_rows
}
collected: list[PendingWrite] = [] # newest first; reversed at the end
for cid in ancestors:
# Collect writes first — they encode the transition FROM this
# ancestor's state to its child's and must be included even if
# this ancestor is also the seed checkpoint.
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, channel, val))
# Then check seed terminator.
ver = ver_of.get(cid)
if ver is not None:
seed_blob = blob_by_ver.get(ver)
if seed_blob is not None and seed_blob[0] != "empty":
blob_value = self.serde.loads_typed(seed_blob)
if blob_value is not DELTA_SENTINEL:
if isinstance(blob_value, _DeltaSnapshot):
# Step-based snapshot: collect this ancestor's
# pending_writes first (they encode the NEXT step's
# transition, not subsumed by the snapshot blob).
for (
type_tag,
write_blob,
task_id,
_idx,
) in writes_by_cid.get(cid, []):
val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, channel, val))
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
# Pre-delta blob: subsumes this ancestor's writes.
collected.reverse()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, channel, val))
collected.reverse() # oldest → newest
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)