mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-14 12:15:25 +02:00
Elior Nataf LackritzandGitHub
93a5a28008
Reading a DeltaChannel at an older checkpoint (the state at a past checkpoint, each entry of `get_state_history`, a fork or replay from one) walked `checkpoints` down from the newest row, so every newer row was read before the target's chain. A checkpoint's ancestors have smaller ids, the order `get_tuple` already relies on to find the latest checkpoint, so the walk now starts at the target: the first page is `checkpoint_id <= target`, later pages stay `< cursor`. The cursor is never NULL anymore, so the `(cursor IS NULL OR checkpoint_id < cursor)` predicate goes too. Once Postgres caches a generic plan for the prepared statement (`from_conn_string` sets `prepare_threshold=0`), that predicate is a filter rather than an index condition, and each later page re-read every newer row. `get_delta_channel_history` on a 6,000-checkpoint thread with a snapshot every 200 checkpoints, median of three runs of 60 reads each, same results before and after: | target | main | this PR | |---|---|---| | newest | 3.1 ms | 2.7 ms | | 1,000 back | 4.5 ms | 2.9 ms | | 3,000 back | 8.9 ms | 2.8 ms | | 5,000 back | 15.7 ms | 2.8 ms | The target is now always the first row read, so the pagination tests from #8556 can no longer leave it off the first page; they are renamed for what they still cover, a walk that continues across pages. ## Test plan - [x] `test_{sync,async}_walk_reads_nothing_newer_than_the_target` record the rows the walk reads; both fail on main, which also reads the three newer checkpoints. - [x] `checkpoint-postgres` suite on Postgres 16: 281 passed, 3 skipped. - [x] `langgraph` DeltaChannel tests against this saver (`-k "delta and not redis"`): 1,154 passed, 1 skipped.
666 lines
26 KiB
Python
666 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
import warnings
|
|
from collections.abc import Mapping, Sequence
|
|
from importlib.metadata import version as get_version
|
|
from typing import Any, TypedDict, cast
|
|
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.checkpoint.base import (
|
|
WRITES_IDX_MAP,
|
|
BaseCheckpointSaver,
|
|
ChannelVersions,
|
|
DeltaChannelHistory,
|
|
PendingWrite,
|
|
get_checkpoint_id,
|
|
writes_sort_key,
|
|
)
|
|
from langgraph.checkpoint.serde.types import TASKS
|
|
from psycopg.types.json import Jsonb
|
|
|
|
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
|
|
# constant — exposing this as a kwarg is left as a follow-up.
|
|
_DELTA_PAGE_SIZE = 1024
|
|
|
|
MetadataInput = dict[str, Any] | None
|
|
|
|
try:
|
|
major, minor = get_version("langgraph").split(".")[:2]
|
|
if int(major) == 0 and int(minor) < 5:
|
|
warnings.warn(
|
|
"You're using incompatible versions of langgraph and checkpoint-postgres. Please upgrade langgraph to avoid unexpected behavior.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
except Exception:
|
|
# skip version check if running from source
|
|
pass
|
|
|
|
"""
|
|
To add a new migration, add a new string to the MIGRATIONS list.
|
|
The position of the migration in the list is the version number.
|
|
"""
|
|
MIGRATIONS = [
|
|
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
|
|
v INTEGER PRIMARY KEY
|
|
);""",
|
|
"""CREATE TABLE IF NOT EXISTS checkpoints (
|
|
thread_id TEXT NOT NULL,
|
|
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
|
checkpoint_id TEXT NOT NULL,
|
|
parent_checkpoint_id TEXT,
|
|
type TEXT,
|
|
checkpoint JSONB NOT NULL,
|
|
metadata JSONB NOT NULL DEFAULT '{}',
|
|
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
|
|
);""",
|
|
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
|
|
thread_id TEXT NOT NULL,
|
|
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
|
channel TEXT NOT NULL,
|
|
version TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
blob BYTEA,
|
|
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
|
|
);""",
|
|
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
|
|
thread_id TEXT NOT NULL,
|
|
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
|
checkpoint_id TEXT NOT NULL,
|
|
task_id TEXT NOT NULL,
|
|
idx INTEGER NOT NULL,
|
|
channel TEXT NOT NULL,
|
|
type TEXT,
|
|
blob BYTEA NOT NULL,
|
|
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
|
);""",
|
|
"ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
|
|
# NOTE: this is a no-op migration to ensure that the versions in the migrations table are correct.
|
|
# This is necessary due to an empty migration previously added to the list.
|
|
"SELECT 1;",
|
|
"""
|
|
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
|
|
""",
|
|
"""
|
|
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
|
|
""",
|
|
"""
|
|
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
|
|
""",
|
|
"""ALTER TABLE checkpoint_writes ADD COLUMN IF NOT EXISTS task_path TEXT NOT NULL DEFAULT '';""",
|
|
]
|
|
|
|
SELECT_SQL = """
|
|
select
|
|
thread_id,
|
|
checkpoint,
|
|
checkpoint_ns,
|
|
checkpoint_id,
|
|
parent_checkpoint_id,
|
|
metadata,
|
|
(
|
|
select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
|
|
from jsonb_each_text(checkpoint -> 'channel_versions')
|
|
inner join checkpoint_blobs bl
|
|
on bl.thread_id = checkpoints.thread_id
|
|
and bl.checkpoint_ns = checkpoints.checkpoint_ns
|
|
and bl.channel = jsonb_each_text.key
|
|
and bl.version = jsonb_each_text.value
|
|
) as channel_values,
|
|
(
|
|
select
|
|
array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob, convert_to(cw.task_path, 'UTF8'), cw.idx::text::bytea])
|
|
from checkpoint_writes cw
|
|
where cw.thread_id = checkpoints.thread_id
|
|
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
|
and cw.checkpoint_id = checkpoints.checkpoint_id
|
|
) as pending_writes
|
|
from checkpoints """
|
|
|
|
SELECT_PENDING_SENDS_SQL = f"""
|
|
select
|
|
checkpoint_id,
|
|
array_agg(array[type::bytea, blob] order by task_path, task_id, idx) as sends
|
|
from checkpoint_writes
|
|
where thread_id = %s
|
|
and checkpoint_id = any(%s)
|
|
and channel = '{TASKS}'
|
|
group by checkpoint_id
|
|
"""
|
|
|
|
UPSERT_CHECKPOINT_BLOBS_SQL = """
|
|
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING
|
|
"""
|
|
|
|
UPSERT_CHECKPOINTS_SQL = """
|
|
INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id)
|
|
DO UPDATE SET
|
|
checkpoint = EXCLUDED.checkpoint,
|
|
metadata = EXCLUDED.metadata;
|
|
"""
|
|
|
|
UPSERT_CHECKPOINT_WRITES_SQL = """
|
|
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
|
|
channel = EXCLUDED.channel,
|
|
type = EXCLUDED.type,
|
|
blob = EXCLUDED.blob;
|
|
"""
|
|
|
|
INSERT_CHECKPOINT_WRITES_SQL = """
|
|
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
|
|
"""
|
|
|
|
|
|
class _DeltaStage2Row(TypedDict, total=False):
|
|
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
|
|
|
|
_kind: str # "w" or "b"
|
|
checkpoint_id: str | None # "w" rows only
|
|
channel: str | None # set on both "w" and "b" rows
|
|
type: str | None
|
|
blob: bytes | None
|
|
task_id: str | None # "w" rows only
|
|
task_path: str | None # "w" rows only
|
|
idx: int | None # "w" rows only
|
|
version: str | None # "b" rows only
|
|
|
|
|
|
# Delta history is rebuilt in two queries; `_build_delta_stage1_sql` and
|
|
# `_build_delta_stage2_sql` document their shapes.
|
|
|
|
|
|
def _build_delta_stage1_sql(
|
|
channels: Sequence[str], *, paged: bool, include_cursor: bool = False
|
|
) -> str:
|
|
"""Build stage 1 SQL with K parallel version lookups + seed probes.
|
|
|
|
For channels=["messages", "files"] (with `paged=True, include_cursor=True`)
|
|
the result is::
|
|
|
|
SELECT checkpoint_id, parent_checkpoint_id,
|
|
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
|
EXISTS (SELECT 1 FROM checkpoint_blobs b0
|
|
WHERE b0.thread_id = checkpoints.thread_id
|
|
AND b0.checkpoint_ns = checkpoints.checkpoint_ns
|
|
AND b0.channel = %s
|
|
AND b0.version = checkpoint -> 'channel_versions' ->> %s
|
|
AND b0.type <> 'empty') AS hb_0,
|
|
checkpoint -> 'channel_values' -> %s AS inline_0,
|
|
checkpoint -> 'channel_versions' ->> %s AS ver_1,
|
|
EXISTS (...) AS hb_1,
|
|
checkpoint -> 'channel_values' -> %s AS inline_1
|
|
FROM checkpoints
|
|
WHERE thread_id = %s AND checkpoint_ns = %s
|
|
AND checkpoint_id <= %s
|
|
ORDER BY checkpoint_id DESC
|
|
LIMIT %s
|
|
|
|
A stored value for a channel lives in one of two places, because `put`
|
|
splits them:
|
|
|
|
* **blob** — non-primitive values (and `_DeltaSnapshot`) are moved to
|
|
`checkpoint_blobs`. `hb_i` ("has blob") probes for one. The probe hits
|
|
that table's primary key `(thread_id, checkpoint_ns, channel, version)`
|
|
exactly, so it is an index lookup per row per channel.
|
|
* **inline** — `None`, `str`, `int`, `float` and `bool` stay in the
|
|
checkpoint's own `channel_values` and get no blob row at all. `inline_i`
|
|
returns that value.
|
|
|
|
Testing only for a key in `channel_values` (the previous approach) missed
|
|
blob-stored plain values, since `put` leaves an inline marker there for
|
|
`_DeltaSnapshot` but not for a plain value — which is what a thread
|
|
migrated from a pre-delta channel type leaves behind. Probing only the
|
|
blobs table would conversely miss inline primitives. Both are needed, and
|
|
the caller treats "either present" as the seed.
|
|
|
|
`hb_i` also disambiguates the two: for a `_DeltaSnapshot`, `inline_i` is the
|
|
literal `true` marker rather than the value, so a blob must win over an
|
|
inline reading whenever one exists. That ordering is what makes a genuine
|
|
inline `true` (a bool channel) distinguishable from the marker.
|
|
|
|
The `type <> 'empty'` predicate mirrors the check stage 2 already applies
|
|
when resolving the seed blob. `put` does not currently produce `empty` rows
|
|
on this path — `blob_versions` is filtered to keys present in
|
|
`channel_values`, so `_dump_blobs`' empty branch is unreachable from it —
|
|
but without the predicate the two stages could disagree: stage 1 would
|
|
terminate the walk on a row stage 2 then discards, yielding no seed *and* a
|
|
truncated write chain, which is the failure this function exists to avoid.
|
|
|
|
Channel names are passed as `%s` parameters (safe from SQL injection).
|
|
Only the column aliases `ver_i` / `hb_i` / `inline_i` and the subquery alias
|
|
`b{i}` are interpolated into the SQL string (i is bounded by len(channels)
|
|
and uses safe identifiers).
|
|
|
|
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
|
|
cursor, page_size]` when `paged=True` — four per channel: the version
|
|
lookup, the blob's channel, the version the blob must match, and the inline
|
|
lookup.
|
|
|
|
Pages run newest-first from the target down. The first page passes the
|
|
target as the cursor with `include_cursor=True`, so it opens with the
|
|
target's own row, whose parent starts the walk; each later page continues
|
|
below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6
|
|
is time-ordered, which `get_tuple` also relies on to find the latest
|
|
checkpoint), so no row newer than the target is part of its chain.
|
|
|
|
When `paged=False`, the WHERE has no cursor predicate and there's no
|
|
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
|
"""
|
|
cols = []
|
|
for i in range(len(channels)):
|
|
cols.append(
|
|
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, "
|
|
f"EXISTS (SELECT 1 FROM checkpoint_blobs b{i} "
|
|
f"WHERE b{i}.thread_id = checkpoints.thread_id "
|
|
f"AND b{i}.checkpoint_ns = checkpoints.checkpoint_ns "
|
|
f"AND b{i}.channel = %s "
|
|
f"AND b{i}.version = checkpoint -> 'channel_versions' ->> %s "
|
|
f"AND b{i}.type <> 'empty') AS hb_{i}, "
|
|
f"checkpoint -> 'channel_values' -> %s AS inline_{i}"
|
|
)
|
|
sql = (
|
|
"SELECT checkpoint_id, parent_checkpoint_id, "
|
|
+ ", ".join(cols)
|
|
+ " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
|
|
)
|
|
if paged:
|
|
sql += (
|
|
f" AND checkpoint_id {'<=' if include_cursor else '<'} %s"
|
|
" ORDER BY checkpoint_id DESC LIMIT %s"
|
|
)
|
|
return sql
|
|
|
|
|
|
def _build_delta_stage2_sql(
|
|
*,
|
|
channels_with_chain: Sequence[str],
|
|
channels_with_seed: Sequence[str],
|
|
) -> str:
|
|
"""Build stage 2 SQL as a per-channel UNION ALL.
|
|
|
|
For each channel with a non-empty chain, emit one branch reading
|
|
`checkpoint_writes` for that specific channel + chain_cids. For each
|
|
channel with a seed_version, emit one branch reading `checkpoint_blobs`
|
|
for that channel + version. This avoids the over-fetch of the prior
|
|
`channel = ANY(channels) AND checkpoint_id = ANY(union)` form when
|
|
channels have different chain depths.
|
|
|
|
The caller must pass parameters in matching order:
|
|
|
|
for ch in channels_with_chain:
|
|
params += [thread_id, checkpoint_ns, ch, chain_cids[ch]]
|
|
for ch in channels_with_seed:
|
|
params += [thread_id, checkpoint_ns, ch, seed_version[ch]]
|
|
|
|
Returns an empty SQL string if both channel lists are empty (caller
|
|
must skip executing in that case).
|
|
"""
|
|
branches: list[str] = []
|
|
for _ in channels_with_chain:
|
|
branches.append(
|
|
"SELECT 'w'::text AS _kind, "
|
|
"checkpoint_id, channel, "
|
|
"type, blob, task_id, task_path, idx, NULL::text AS version "
|
|
"FROM checkpoint_writes "
|
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
|
"AND checkpoint_id = ANY(%s)"
|
|
)
|
|
for _ in channels_with_seed:
|
|
branches.append(
|
|
"SELECT 'b'::text AS _kind, NULL::text AS checkpoint_id, channel, "
|
|
"type, blob, NULL::text AS task_id, NULL::text AS task_path, "
|
|
"NULL::int AS idx, version "
|
|
"FROM checkpoint_blobs "
|
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
|
"AND version = %s"
|
|
)
|
|
return " UNION ALL ".join(branches)
|
|
|
|
|
|
# Stage 1 rows are dicts keyed by the per-channel aliases
|
|
# `_build_delta_stage1_sql` emits, so there is no static TypedDict.
|
|
|
|
|
|
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
SELECT_SQL = SELECT_SQL
|
|
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
|
|
MIGRATIONS = MIGRATIONS
|
|
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
|
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
|
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
|
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
|
|
|
supports_pipeline: bool
|
|
|
|
def _migrate_pending_sends(
|
|
self,
|
|
pending_sends: list[tuple[bytes, bytes]],
|
|
checkpoint: dict[str, Any],
|
|
channel_values: list[tuple[bytes, bytes, bytes]],
|
|
) -> None:
|
|
if not pending_sends:
|
|
return
|
|
# add to values
|
|
enc, blob = self.serde.dumps_typed(
|
|
[self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends],
|
|
)
|
|
channel_values.append((TASKS.encode(), enc.encode(), blob))
|
|
# add to versions
|
|
checkpoint["channel_versions"][TASKS] = (
|
|
max(checkpoint["channel_versions"].values())
|
|
if checkpoint["channel_versions"]
|
|
else self.get_next_version(None, None)
|
|
)
|
|
|
|
def _load_blobs(
|
|
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
|
) -> dict[str, Any]:
|
|
if not blob_values:
|
|
return {}
|
|
return {
|
|
k.decode(): self.serde.loads_typed((t.decode(), v))
|
|
for k, t, v in blob_values
|
|
if t.decode() != "empty"
|
|
}
|
|
|
|
@staticmethod
|
|
def _ingest_stage1_page(
|
|
stage1_rows: Sequence[Mapping[str, Any]],
|
|
channels: Sequence[str],
|
|
parent_of: dict[str, str | None],
|
|
ver_by_i_by_cid: list[dict[str, str | None]],
|
|
hb_by_i_by_cid: list[dict[str, bool]],
|
|
inline_by_i_by_cid: list[dict[str, Any]],
|
|
) -> str | None:
|
|
"""Fold one stage-1 page into the running walk-state mappings.
|
|
|
|
Returns the oldest checkpoint_id seen on this page (smallest, since
|
|
pages come back DESC). Caller uses it as the cursor for the next
|
|
page (`AND checkpoint_id < cursor`).
|
|
"""
|
|
oldest: str | None = None
|
|
for r in stage1_rows:
|
|
cid = cast(str, r["checkpoint_id"])
|
|
parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
|
|
for i in range(len(channels)):
|
|
ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
|
|
hb_by_i_by_cid[i][cid] = bool(r.get(f"hb_{i}"))
|
|
inline_by_i_by_cid[i][cid] = r.get(f"inline_{i}")
|
|
# Rows are DESC; the last one is the smallest cid in the page.
|
|
oldest = cid
|
|
return oldest
|
|
|
|
@staticmethod
|
|
def _try_advance_walks(
|
|
target_id: str,
|
|
channels: Sequence[str],
|
|
parent_of: Mapping[str, str | None],
|
|
ver_by_i_by_cid: Sequence[Mapping[str, str | None]],
|
|
hb_by_i_by_cid: Sequence[Mapping[str, bool]],
|
|
inline_by_i_by_cid: Sequence[Mapping[str, Any]],
|
|
chain_by_ch: dict[str, list[str]],
|
|
seed_ver_by_ch: dict[str, str | None],
|
|
seed_inline_by_ch: dict[str, Any],
|
|
walk_cursor_by_ch: dict[str, str | None],
|
|
seeded: set[str],
|
|
) -> None:
|
|
"""Advance each not-yet-seeded channel's walk as far as possible.
|
|
|
|
Uses the partial `parent_of` map accumulated so far. A walk stops
|
|
either because:
|
|
(a) it found a stored value for its channel — a blob or an inline
|
|
primitive (channel becomes seeded),
|
|
(b) it reached a real root (parent_of[cid] is None — fully
|
|
materialized at this point),
|
|
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
|
a later page; the cursor stays put), or
|
|
(d) the target's own row isn't in `parent_of` (the target doesn't
|
|
exist, so the walk never starts).
|
|
|
|
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
|
`walk_cursor_by_ch`, and `seeded` in place.
|
|
"""
|
|
for i, ch in enumerate(channels):
|
|
if ch in seeded:
|
|
continue
|
|
# The first page opens with the target's row, so it's missing only
|
|
# when the target doesn't exist; a `None` cursor would read as
|
|
# "target is a root".
|
|
if ch not in walk_cursor_by_ch:
|
|
if target_id not in parent_of:
|
|
continue
|
|
walk_cursor_by_ch[ch] = parent_of[target_id]
|
|
cur_cid = walk_cursor_by_ch[ch]
|
|
ch_chain = chain_by_ch[ch]
|
|
hb_i = hb_by_i_by_cid[i]
|
|
inline_i = inline_by_i_by_cid[i]
|
|
ver_i = ver_by_i_by_cid[i]
|
|
while cur_cid is not None:
|
|
if cur_cid not in parent_of:
|
|
# Need more pages to continue this walk.
|
|
break
|
|
ch_chain.append(cur_cid)
|
|
has_blob = hb_i.get(cur_cid, False)
|
|
inline = inline_i.get(cur_cid)
|
|
if has_blob or inline is not None:
|
|
# A blob wins: for a `_DeltaSnapshot` the inline reading is
|
|
# the `true` marker, not the value.
|
|
seed_ver_by_ch[ch] = ver_i.get(cur_cid)
|
|
if not has_blob:
|
|
seed_inline_by_ch[ch] = inline
|
|
seeded.add(ch)
|
|
cur_cid = None
|
|
break
|
|
cur_cid = parent_of[cur_cid]
|
|
walk_cursor_by_ch[ch] = cur_cid
|
|
|
|
def _build_delta_channels_writes_history(
|
|
self,
|
|
*,
|
|
channels: Sequence[str],
|
|
chain_by_ch: Mapping[str, list[str]],
|
|
seed_ver_by_ch: Mapping[str, str | None],
|
|
seed_inline_by_ch: Mapping[str, Any],
|
|
stage2_rows: Sequence[_DeltaStage2Row],
|
|
) -> dict[str, DeltaChannelHistory]:
|
|
"""Demux stage 2 rows per channel; produce per-channel histories.
|
|
|
|
stage2_rows carry `channel` on every row. We build per-channel
|
|
`writes_by_cid` and per-channel `seed_blob` dicts, then assemble
|
|
a `DeltaChannelHistory` per requested channel.
|
|
|
|
A seed comes from the blobs table when the walk found one there, and
|
|
otherwise from `seed_inline_by_ch` — `put` keeps `None`, `str`, `int`,
|
|
`float` and `bool` values in the checkpoint's own `channel_values` with
|
|
no blob row, so those never appear in `stage2_rows`.
|
|
|
|
The `seed` key is omitted when the walk reached root without finding a
|
|
stored value, or when the seed blob is sentinel "empty" — in both cases
|
|
the consumer treats absence as "start empty".
|
|
"""
|
|
# writes_by_ch_by_cid[channel][cid] = list of
|
|
# (type, blob, task_id, idx, task_path)
|
|
writes_by_ch_by_cid: dict[
|
|
str, dict[str, list[tuple[str, bytes, str, int, str]]]
|
|
] = {ch: {} for ch in channels}
|
|
# seed_blob_by_ver[(channel, version)] = (type, blob)
|
|
seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {}
|
|
|
|
for r in stage2_rows:
|
|
ch = cast(str, r["channel"])
|
|
kind = r["_kind"]
|
|
if kind == "w":
|
|
cid = cast(str, r["checkpoint_id"])
|
|
writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
|
|
cast(
|
|
"tuple[str, bytes, str, int, str]",
|
|
(
|
|
r["type"],
|
|
r["blob"],
|
|
r["task_id"],
|
|
r["idx"],
|
|
r["task_path"],
|
|
),
|
|
)
|
|
)
|
|
else: # kind == "b"
|
|
ver = cast(str, r["version"])
|
|
seed_blob_by_ver[(ch, ver)] = cast(
|
|
"tuple[str, bytes]", (r["type"], r["blob"])
|
|
)
|
|
|
|
# Sort writes per (channel, cid) newest-first
|
|
for cid_map in writes_by_ch_by_cid.values():
|
|
for ws in cid_map.values():
|
|
ws.sort(key=lambda w: writes_sort_key(w[4], w[2], w[3]), reverse=True)
|
|
|
|
result: dict[str, DeltaChannelHistory] = {}
|
|
for ch in channels:
|
|
chain_cids = chain_by_ch.get(ch, [])
|
|
seed_version = seed_ver_by_ch.get(ch)
|
|
|
|
collected: list[PendingWrite] = []
|
|
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
|
for cid in chain_cids:
|
|
for type_tag, write_blob, task_id, _idx, _path in cid_writes.get(
|
|
cid, []
|
|
):
|
|
val = self.serde.loads_typed((type_tag, write_blob))
|
|
collected.append((task_id, ch, val))
|
|
collected.reverse()
|
|
|
|
entry: DeltaChannelHistory = {"writes": collected}
|
|
if seed_version is not None:
|
|
blob = seed_blob_by_ver.get((ch, seed_version))
|
|
if blob is not None and blob[0] != "empty":
|
|
entry["seed"] = self.serde.loads_typed(blob)
|
|
elif ch in seed_inline_by_ch:
|
|
# Inline primitive: stored in the checkpoint, not the blobs
|
|
# table, so stage 2 never returned a row for it.
|
|
entry["seed"] = seed_inline_by_ch[ch]
|
|
result[ch] = entry
|
|
return result
|
|
|
|
def _dump_blobs(
|
|
self,
|
|
thread_id: str,
|
|
checkpoint_ns: str,
|
|
values: dict[str, Any],
|
|
versions: ChannelVersions,
|
|
) -> list[tuple[str, str, str, str, str, bytes | None]]:
|
|
if not versions:
|
|
return []
|
|
|
|
return [
|
|
(
|
|
thread_id,
|
|
checkpoint_ns,
|
|
k,
|
|
cast(str, ver),
|
|
*(
|
|
self.serde.dumps_typed(values[k])
|
|
if k in values
|
|
else ("empty", None)
|
|
),
|
|
)
|
|
for k, ver in versions.items()
|
|
]
|
|
|
|
def _load_writes(
|
|
self, writes: list[tuple[bytes, bytes, bytes, bytes, bytes, bytes]] | None
|
|
) -> list[tuple[str, str, Any]]:
|
|
return [
|
|
(tid.decode(), channel.decode(), self.serde.loads_typed((t.decode(), v)))
|
|
for tid, channel, t, v, _, _ in sorted(
|
|
writes or [],
|
|
key=lambda w: writes_sort_key(w[4].decode(), w[0].decode(), int(w[5])),
|
|
)
|
|
]
|
|
|
|
def _dump_writes(
|
|
self,
|
|
thread_id: str,
|
|
checkpoint_ns: str,
|
|
checkpoint_id: str,
|
|
task_id: str,
|
|
task_path: str,
|
|
writes: Sequence[tuple[str, Any]],
|
|
) -> list[tuple[str, str, str, str, str, int, str, str, bytes]]:
|
|
return [
|
|
(
|
|
thread_id,
|
|
checkpoint_ns,
|
|
checkpoint_id,
|
|
task_id,
|
|
task_path,
|
|
WRITES_IDX_MAP.get(channel, idx),
|
|
channel,
|
|
*self.serde.dumps_typed(value),
|
|
)
|
|
for idx, (channel, value) in enumerate(writes)
|
|
]
|
|
|
|
def get_next_version(self, current: str | None, channel: None) -> str:
|
|
if current is None:
|
|
current_v = 0
|
|
elif isinstance(current, int):
|
|
current_v = current
|
|
else:
|
|
current_v = int(current.split(".")[0])
|
|
next_v = current_v + 1
|
|
next_h = random.random()
|
|
return f"{next_v:032}.{next_h:016}"
|
|
|
|
def _search_where(
|
|
self,
|
|
config: RunnableConfig | None,
|
|
filter: MetadataInput,
|
|
before: RunnableConfig | None = None,
|
|
) -> tuple[str, list[Any]]:
|
|
"""Return WHERE clause predicates for alist() given config, filter, before.
|
|
|
|
This method returns a tuple of a string and a tuple of values. The string
|
|
is the parametered WHERE clause predicate (including the WHERE keyword):
|
|
"WHERE column1 = $1 AND column2 IS $2". The list of values contains the
|
|
values for each of the corresponding parameters.
|
|
"""
|
|
wheres = []
|
|
param_values = []
|
|
|
|
# construct predicate for config filter
|
|
if config:
|
|
wheres.append("thread_id = %s ")
|
|
param_values.append(config["configurable"]["thread_id"])
|
|
checkpoint_ns = config["configurable"].get("checkpoint_ns")
|
|
if checkpoint_ns is not None:
|
|
wheres.append("checkpoint_ns = %s")
|
|
param_values.append(checkpoint_ns)
|
|
|
|
if checkpoint_id := get_checkpoint_id(config):
|
|
wheres.append("checkpoint_id = %s ")
|
|
param_values.append(checkpoint_id)
|
|
|
|
# construct predicate for metadata filter
|
|
if filter:
|
|
wheres.append("metadata @> %s ")
|
|
param_values.append(Jsonb(filter))
|
|
|
|
# construct predicate for `before`
|
|
if before is not None:
|
|
wheres.append("checkpoint_id < %s ")
|
|
param_values.append(get_checkpoint_id(before))
|
|
|
|
return (
|
|
"WHERE " + " AND ".join(wheres) if wheres else "",
|
|
param_values,
|
|
)
|