langgraph-checkpoint-postgres 3.1.2__tar.gz → 3.1.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/.gitignore +6 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/PKG-INFO +3 -3
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/__init__.py +12 -9
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/aio.py +23 -4
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/base.py +65 -66
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/shallow.py +1 -1
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/store/postgres/base.py +4 -2
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/pyproject.toml +2 -2
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_async_store.py +64 -0
- langgraph_checkpoint_postgres-3.1.3/tests/test_delta_pagination.py +211 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_store.py +49 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/uv.lock +219 -196
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/LICENSE +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/Makefile +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/README.md +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/_ainternal.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/_internal.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/checkpoint/postgres/py.typed +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/store/postgres/__init__.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/store/postgres/aio.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/langgraph/store/postgres/py.typed +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/__init__.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/compose-postgres.yml +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/conftest.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/embed_test_utils.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_async.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_conformance_delta.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_delta_plain_value_seed.py +0 -0
- {langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.1.3}/tests/test_sync.py +0 -0
|
@@ -76,7 +76,12 @@ __pypackages__/
|
|
|
76
76
|
|
|
77
77
|
# Environments
|
|
78
78
|
.env
|
|
79
|
+
.env.*
|
|
79
80
|
.envrc
|
|
81
|
+
*.crt
|
|
82
|
+
*.key
|
|
83
|
+
*.pem
|
|
84
|
+
credentials.json
|
|
80
85
|
.venv
|
|
81
86
|
.venvs
|
|
82
87
|
env/
|
|
@@ -98,6 +103,7 @@ dmypy.json
|
|
|
98
103
|
|
|
99
104
|
.vercel
|
|
100
105
|
.turbo
|
|
106
|
+
node_modules/
|
|
101
107
|
.editorconfig
|
|
102
108
|
.scratch
|
|
103
109
|
.worktrees/
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: langgraph-checkpoint-postgres
|
|
3
|
-
Version: 3.1.
|
|
3
|
+
Version: 3.1.3
|
|
4
4
|
Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
|
|
5
5
|
Project-URL: Source, https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres
|
|
6
6
|
Project-URL: Twitter, https://x.com/langchain_oss
|
|
@@ -9,7 +9,7 @@ Project-URL: Reddit, https://www.reddit.com/r/LangChain/
|
|
|
9
9
|
License-Expression: MIT
|
|
10
10
|
License-File: LICENSE
|
|
11
11
|
Requires-Python: >=3.10
|
|
12
|
-
Requires-Dist: langgraph-checkpoint<5.0.0,>=4.
|
|
12
|
+
Requires-Dist: langgraph-checkpoint<5.0.0,>=4.3.0
|
|
13
13
|
Requires-Dist: orjson>=3.11.5
|
|
14
14
|
Requires-Dist: psycopg-pool>=3.2.0
|
|
15
15
|
Requires-Dist: psycopg>=3.2.0
|
|
@@ -448,11 +448,13 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
448
448
|
|
|
449
449
|
Two-stage query, both stages cover ALL requested channels:
|
|
450
450
|
|
|
451
|
-
* Stage 1 (paged): dynamic SELECT over `checkpoints` with
|
|
452
|
-
|
|
453
|
-
|
|
454
|
-
|
|
455
|
-
|
|
451
|
+
* Stage 1 (paged): dynamic SELECT over `checkpoints` with three
|
|
452
|
+
columns per channel: its version, an `EXISTS` probe for a stored
|
|
453
|
+
blob at that version, and its inline value. Pages newest-first by
|
|
454
|
+
`checkpoint_id`, starting at the target; page size is
|
|
455
|
+
`_DELTA_PAGE_SIZE`.
|
|
456
|
+
Stops paging when every channel has found its seed or a page comes
|
|
457
|
+
back short.
|
|
456
458
|
|
|
457
459
|
* Stage 2 (per-channel UNION ALL): one branch per channel reading
|
|
458
460
|
`checkpoint_writes` filtered to that channel's specific
|
|
@@ -467,7 +469,7 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
467
469
|
thread_id = config["configurable"]["thread_id"]
|
|
468
470
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
|
469
471
|
checkpoint_id = get_checkpoint_id(config)
|
|
470
|
-
if checkpoint_id
|
|
472
|
+
if not checkpoint_id:
|
|
471
473
|
target = self.get_tuple(config)
|
|
472
474
|
if target is None:
|
|
473
475
|
return {ch: {"writes": []} for ch in channels}
|
|
@@ -475,7 +477,7 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
475
477
|
|
|
476
478
|
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
|
|
477
479
|
# Python after each page. Stops as soon as every channel has its seed.
|
|
478
|
-
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
|
480
|
+
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
|
479
481
|
parent_of: dict[str, str | None] = {}
|
|
480
482
|
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
|
481
483
|
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
|
@@ -485,7 +487,7 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
485
487
|
seed_inline_by_ch: dict[str, Any] = {}
|
|
486
488
|
walk_cursor_by_ch: dict[str, str | None] = {}
|
|
487
489
|
seeded: set[str] = set()
|
|
488
|
-
cursor: str | None =
|
|
490
|
+
cursor: str | None = checkpoint_id
|
|
489
491
|
|
|
490
492
|
with self._cursor() as cur:
|
|
491
493
|
while True:
|
|
@@ -494,7 +496,7 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
494
496
|
# ver_i, blob channel, blob version, inline_i
|
|
495
497
|
stage1_params.extend([ch, ch, ch, ch])
|
|
496
498
|
stage1_params.extend(
|
|
497
|
-
[thread_id, checkpoint_ns, cursor,
|
|
499
|
+
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
|
498
500
|
)
|
|
499
501
|
cur.execute(stage1_sql, stage1_params)
|
|
500
502
|
page = cur.fetchall()
|
|
@@ -526,6 +528,7 @@ class PostgresSaver(BasePostgresSaver):
|
|
|
526
528
|
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
|
527
529
|
break
|
|
528
530
|
cursor = oldest
|
|
531
|
+
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
|
529
532
|
|
|
530
533
|
# Stage 2: per-channel UNION ALL — one writes branch per channel
|
|
531
534
|
# with non-empty chain, plus one blob branch per seeded channel.
|
|
@@ -417,13 +417,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|
|
417
417
|
thread_id = config["configurable"]["thread_id"]
|
|
418
418
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
|
419
419
|
checkpoint_id = get_checkpoint_id(config)
|
|
420
|
-
if checkpoint_id
|
|
420
|
+
if not checkpoint_id:
|
|
421
421
|
target = await self.aget_tuple(config)
|
|
422
422
|
if target is None:
|
|
423
423
|
return {ch: {"writes": []} for ch in channels}
|
|
424
424
|
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
|
425
425
|
|
|
426
|
-
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
|
426
|
+
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
|
427
427
|
parent_of: dict[str, str | None] = {}
|
|
428
428
|
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
|
429
429
|
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
|
@@ -433,7 +433,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|
|
433
433
|
seed_inline_by_ch: dict[str, Any] = {}
|
|
434
434
|
walk_cursor_by_ch: dict[str, str | None] = {}
|
|
435
435
|
seeded: set[str] = set()
|
|
436
|
-
cursor: str | None =
|
|
436
|
+
cursor: str | None = checkpoint_id
|
|
437
437
|
|
|
438
438
|
async with self._cursor() as cur:
|
|
439
439
|
while True:
|
|
@@ -442,7 +442,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|
|
442
442
|
# ver_i, blob channel, blob version, inline_i
|
|
443
443
|
stage1_params.extend([ch, ch, ch, ch])
|
|
444
444
|
stage1_params.extend(
|
|
445
|
-
[thread_id, checkpoint_ns, cursor,
|
|
445
|
+
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
|
446
446
|
)
|
|
447
447
|
await cur.execute(stage1_sql, stage1_params)
|
|
448
448
|
page = await cur.fetchall()
|
|
@@ -472,6 +472,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|
|
472
472
|
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
|
473
473
|
break
|
|
474
474
|
cursor = oldest
|
|
475
|
+
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
|
475
476
|
|
|
476
477
|
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
|
477
478
|
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
|
@@ -687,5 +688,23 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|
|
687
688
|
self.adelete_thread(thread_id), self.loop
|
|
688
689
|
).result()
|
|
689
690
|
|
|
691
|
+
def get_delta_channel_history(
|
|
692
|
+
self, *, config: RunnableConfig, channels: Sequence[str]
|
|
693
|
+
) -> Mapping[str, DeltaChannelHistory]:
|
|
694
|
+
"""Sync bridge to `aget_delta_channel_history`, guarded like `get_tuple`."""
|
|
695
|
+
try:
|
|
696
|
+
if asyncio.get_running_loop() is self.loop:
|
|
697
|
+
raise asyncio.InvalidStateError(
|
|
698
|
+
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
|
|
699
|
+
"different thread. From the main thread, use the async interface. "
|
|
700
|
+
"For example, use `await checkpointer.aget_delta_channel_history(...)`."
|
|
701
|
+
)
|
|
702
|
+
except RuntimeError:
|
|
703
|
+
pass
|
|
704
|
+
return asyncio.run_coroutine_threadsafe(
|
|
705
|
+
self.aget_delta_channel_history(config=config, channels=channels),
|
|
706
|
+
self.loop,
|
|
707
|
+
).result()
|
|
708
|
+
|
|
690
709
|
|
|
691
710
|
__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]
|
|
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
|
|
|
14
14
|
DeltaChannelHistory,
|
|
15
15
|
PendingWrite,
|
|
16
16
|
get_checkpoint_id,
|
|
17
|
+
writes_sort_key,
|
|
17
18
|
)
|
|
18
19
|
from langgraph.checkpoint.serde.types import TASKS
|
|
19
20
|
from psycopg.types.json import Jsonb
|
|
@@ -109,7 +110,7 @@ select
|
|
|
109
110
|
) as channel_values,
|
|
110
111
|
(
|
|
111
112
|
select
|
|
112
|
-
array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob
|
|
113
|
+
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])
|
|
113
114
|
from checkpoint_writes cw
|
|
114
115
|
where cw.thread_id = checkpoints.thread_id
|
|
115
116
|
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
|
@@ -168,40 +169,22 @@ class _DeltaStage2Row(TypedDict, total=False):
|
|
|
168
169
|
type: str | None
|
|
169
170
|
blob: bytes | None
|
|
170
171
|
task_id: str | None # "w" rows only
|
|
172
|
+
task_path: str | None # "w" rows only
|
|
171
173
|
idx: int | None # "w" rows only
|
|
172
174
|
version: str | None # "b" rows only
|
|
173
175
|
|
|
174
176
|
|
|
175
|
-
#
|
|
176
|
-
#
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
# Stage 2 fetches all writes and the seed blobs for ALL channels in a
|
|
183
|
-
# single roundtrip via `channel = ANY(%s)` and chain/seed-version
|
|
184
|
-
# filtering.
|
|
185
|
-
#
|
|
186
|
-
# Empirical comparison vs an alternative "ship full channel_versions /
|
|
187
|
-
# channel_values JSONB and let Python pick" form (1000 checkpoints,
|
|
188
|
-
# 8 total channels in graph, 3 delta channels requested):
|
|
189
|
-
#
|
|
190
|
-
# Postgres execution: A=0.24ms vs B=0.38ms (both negligible)
|
|
191
|
-
# End-to-end latency: A=6.83ms vs B=2.28ms (B is 3.0x faster)
|
|
192
|
-
# Wire payload: A=836KB vs B=330KB (61% smaller)
|
|
193
|
-
# Buffer hits: identical (167 blocks)
|
|
194
|
-
#
|
|
195
|
-
# B (this dynamic-columns design) wins because it avoids JSONB
|
|
196
|
-
# serialization on the wire and JSONB-to-dict deserialization in
|
|
197
|
-
# psycopg. Even at K=8 (8 delta channels = 16 dynamic columns), B
|
|
198
|
-
# still beats A end-to-end (4.2ms vs 6.8ms).
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
|
177
|
+
# Delta history is rebuilt in two queries; `_build_delta_stage1_sql` and
|
|
178
|
+
# `_build_delta_stage2_sql` document their shapes.
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
def _build_delta_stage1_sql(
|
|
182
|
+
channels: Sequence[str], *, paged: bool, include_cursor: bool = False
|
|
183
|
+
) -> str:
|
|
202
184
|
"""Build stage 1 SQL with K parallel version lookups + seed probes.
|
|
203
185
|
|
|
204
|
-
For channels=["messages", "files"] (with `paged=True`)
|
|
186
|
+
For channels=["messages", "files"] (with `paged=True, include_cursor=True`)
|
|
187
|
+
the result is::
|
|
205
188
|
|
|
206
189
|
SELECT checkpoint_id, parent_checkpoint_id,
|
|
207
190
|
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
|
@@ -217,7 +200,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
|
|
217
200
|
checkpoint -> 'channel_values' -> %s AS inline_1
|
|
218
201
|
FROM checkpoints
|
|
219
202
|
WHERE thread_id = %s AND checkpoint_ns = %s
|
|
220
|
-
AND
|
|
203
|
+
AND checkpoint_id <= %s
|
|
221
204
|
ORDER BY checkpoint_id DESC
|
|
222
205
|
LIMIT %s
|
|
223
206
|
|
|
@@ -258,9 +241,16 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
|
|
258
241
|
and uses safe identifiers).
|
|
259
242
|
|
|
260
243
|
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
|
|
261
|
-
cursor,
|
|
262
|
-
|
|
263
|
-
|
|
244
|
+
cursor, page_size]` when `paged=True` — four per channel: the version
|
|
245
|
+
lookup, the blob's channel, the version the blob must match, and the inline
|
|
246
|
+
lookup.
|
|
247
|
+
|
|
248
|
+
Pages run newest-first from the target down. The first page passes the
|
|
249
|
+
target as the cursor with `include_cursor=True`, so it opens with the
|
|
250
|
+
target's own row, whose parent starts the walk; each later page continues
|
|
251
|
+
below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6
|
|
252
|
+
is time-ordered, which `get_tuple` also relies on to find the latest
|
|
253
|
+
checkpoint), so no row newer than the target is part of its chain.
|
|
264
254
|
|
|
265
255
|
When `paged=False`, the WHERE has no cursor predicate and there's no
|
|
266
256
|
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
|
@@ -284,7 +274,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
|
|
284
274
|
)
|
|
285
275
|
if paged:
|
|
286
276
|
sql += (
|
|
287
|
-
" AND
|
|
277
|
+
f" AND checkpoint_id {'<=' if include_cursor else '<'} %s"
|
|
288
278
|
" ORDER BY checkpoint_id DESC LIMIT %s"
|
|
289
279
|
)
|
|
290
280
|
return sql
|
|
@@ -319,7 +309,7 @@ def _build_delta_stage2_sql(
|
|
|
319
309
|
branches.append(
|
|
320
310
|
"SELECT 'w'::text AS _kind, "
|
|
321
311
|
"checkpoint_id, channel, "
|
|
322
|
-
"type, blob, task_id, idx, NULL::text AS version "
|
|
312
|
+
"type, blob, task_id, task_path, idx, NULL::text AS version "
|
|
323
313
|
"FROM checkpoint_writes "
|
|
324
314
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
|
325
315
|
"AND checkpoint_id = ANY(%s)"
|
|
@@ -327,7 +317,8 @@ def _build_delta_stage2_sql(
|
|
|
327
317
|
for _ in channels_with_seed:
|
|
328
318
|
branches.append(
|
|
329
319
|
"SELECT 'b'::text AS _kind, NULL::text AS checkpoint_id, channel, "
|
|
330
|
-
"type, blob, NULL::text AS task_id, NULL::
|
|
320
|
+
"type, blob, NULL::text AS task_id, NULL::text AS task_path, "
|
|
321
|
+
"NULL::int AS idx, version "
|
|
331
322
|
"FROM checkpoint_blobs "
|
|
332
323
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
|
333
324
|
"AND version = %s"
|
|
@@ -335,10 +326,8 @@ def _build_delta_stage2_sql(
|
|
|
335
326
|
return " UNION ALL ".join(branches)
|
|
336
327
|
|
|
337
328
|
|
|
338
|
-
# Stage 1 rows are
|
|
339
|
-
#
|
|
340
|
-
# list to map indices back to channel names — no static TypedDict here.
|
|
341
|
-
# `dict[str, Any]` is the practical signature.
|
|
329
|
+
# Stage 1 rows are dicts keyed by the per-channel aliases
|
|
330
|
+
# `_build_delta_stage1_sql` emits, so there is no static TypedDict.
|
|
342
331
|
|
|
343
332
|
|
|
344
333
|
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
@@ -431,9 +420,11 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
431
420
|
(a) it found a stored value for its channel — a blob or an inline
|
|
432
421
|
primitive (channel becomes seeded),
|
|
433
422
|
(b) it reached a real root (parent_of[cid] is None — fully
|
|
434
|
-
materialized at this point),
|
|
423
|
+
materialized at this point),
|
|
435
424
|
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
|
436
|
-
a later page; the cursor stays put)
|
|
425
|
+
a later page; the cursor stays put), or
|
|
426
|
+
(d) the target's own row isn't in `parent_of` (the target doesn't
|
|
427
|
+
exist, so the walk never starts).
|
|
437
428
|
|
|
438
429
|
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
|
439
430
|
`walk_cursor_by_ch`, and `seeded` in place.
|
|
@@ -441,9 +432,13 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
441
432
|
for i, ch in enumerate(channels):
|
|
442
433
|
if ch in seeded:
|
|
443
434
|
continue
|
|
444
|
-
#
|
|
435
|
+
# The first page opens with the target's row, so it's missing only
|
|
436
|
+
# when the target doesn't exist; a `None` cursor would read as
|
|
437
|
+
# "target is a root".
|
|
445
438
|
if ch not in walk_cursor_by_ch:
|
|
446
|
-
|
|
439
|
+
if target_id not in parent_of:
|
|
440
|
+
continue
|
|
441
|
+
walk_cursor_by_ch[ch] = parent_of[target_id]
|
|
447
442
|
cur_cid = walk_cursor_by_ch[ch]
|
|
448
443
|
ch_chain = chain_by_ch[ch]
|
|
449
444
|
hb_i = hb_by_i_by_cid[i]
|
|
@@ -492,10 +487,11 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
492
487
|
stored value, or when the seed blob is sentinel "empty" — in both cases
|
|
493
488
|
the consumer treats absence as "start empty".
|
|
494
489
|
"""
|
|
495
|
-
# writes_by_ch_by_cid[channel][cid] = list of
|
|
496
|
-
|
|
497
|
-
|
|
498
|
-
|
|
490
|
+
# writes_by_ch_by_cid[channel][cid] = list of
|
|
491
|
+
# (type, blob, task_id, idx, task_path)
|
|
492
|
+
writes_by_ch_by_cid: dict[
|
|
493
|
+
str, dict[str, list[tuple[str, bytes, str, int, str]]]
|
|
494
|
+
] = {ch: {} for ch in channels}
|
|
499
495
|
# seed_blob_by_ver[(channel, version)] = (type, blob)
|
|
500
496
|
seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {}
|
|
501
497
|
|
|
@@ -506,8 +502,14 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
506
502
|
cid = cast(str, r["checkpoint_id"])
|
|
507
503
|
writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
|
|
508
504
|
cast(
|
|
509
|
-
"tuple[str, bytes, str, int]",
|
|
510
|
-
(
|
|
505
|
+
"tuple[str, bytes, str, int, str]",
|
|
506
|
+
(
|
|
507
|
+
r["type"],
|
|
508
|
+
r["blob"],
|
|
509
|
+
r["task_id"],
|
|
510
|
+
r["idx"],
|
|
511
|
+
r["task_path"],
|
|
512
|
+
),
|
|
511
513
|
)
|
|
512
514
|
)
|
|
513
515
|
else: # kind == "b"
|
|
@@ -516,10 +518,10 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
516
518
|
"tuple[str, bytes]", (r["type"], r["blob"])
|
|
517
519
|
)
|
|
518
520
|
|
|
519
|
-
# Sort writes per (channel, cid) newest-first
|
|
521
|
+
# Sort writes per (channel, cid) newest-first
|
|
520
522
|
for cid_map in writes_by_ch_by_cid.values():
|
|
521
523
|
for ws in cid_map.values():
|
|
522
|
-
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
|
|
524
|
+
ws.sort(key=lambda w: writes_sort_key(w[4], w[2], w[3]), reverse=True)
|
|
523
525
|
|
|
524
526
|
result: dict[str, DeltaChannelHistory] = {}
|
|
525
527
|
for ch in channels:
|
|
@@ -529,7 +531,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
529
531
|
collected: list[PendingWrite] = []
|
|
530
532
|
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
|
531
533
|
for cid in chain_cids:
|
|
532
|
-
for type_tag, write_blob, task_id, _idx in cid_writes.get(
|
|
534
|
+
for type_tag, write_blob, task_id, _idx, _path in cid_writes.get(
|
|
535
|
+
cid, []
|
|
536
|
+
):
|
|
533
537
|
val = self.serde.loads_typed((type_tag, write_blob))
|
|
534
538
|
collected.append((task_id, ch, val))
|
|
535
539
|
collected.reverse()
|
|
@@ -572,20 +576,15 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|
|
572
576
|
]
|
|
573
577
|
|
|
574
578
|
def _load_writes(
|
|
575
|
-
self, writes: list[tuple[bytes, bytes, bytes, bytes]]
|
|
579
|
+
self, writes: list[tuple[bytes, bytes, bytes, bytes, bytes, bytes]] | None
|
|
576
580
|
) -> list[tuple[str, str, Any]]:
|
|
577
|
-
return
|
|
578
|
-
|
|
579
|
-
|
|
580
|
-
|
|
581
|
-
|
|
582
|
-
|
|
583
|
-
|
|
584
|
-
for tid, channel, t, v in writes
|
|
585
|
-
]
|
|
586
|
-
if writes
|
|
587
|
-
else []
|
|
588
|
-
)
|
|
581
|
+
return [
|
|
582
|
+
(tid.decode(), channel.decode(), self.serde.loads_typed((t.decode(), v)))
|
|
583
|
+
for tid, channel, t, v, _, _ in sorted(
|
|
584
|
+
writes or [],
|
|
585
|
+
key=lambda w: writes_sort_key(w[4].decode(), w[0].decode(), int(w[5])),
|
|
586
|
+
)
|
|
587
|
+
]
|
|
589
588
|
|
|
590
589
|
def _dump_writes(
|
|
591
590
|
self,
|
|
@@ -97,7 +97,7 @@ select
|
|
|
97
97
|
) as channel_values,
|
|
98
98
|
(
|
|
99
99
|
select
|
|
100
|
-
array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob
|
|
100
|
+
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])
|
|
101
101
|
from checkpoint_writes cw
|
|
102
102
|
where cw.thread_id = checkpoints.thread_id
|
|
103
103
|
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
|
@@ -7,7 +7,7 @@ import logging
|
|
|
7
7
|
import re
|
|
8
8
|
import threading
|
|
9
9
|
from collections import defaultdict
|
|
10
|
-
from collections.abc import Callable, Iterable, Iterator, Sequence
|
|
10
|
+
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
|
11
11
|
from contextlib import contextmanager
|
|
12
12
|
from datetime import datetime
|
|
13
13
|
from typing import (
|
|
@@ -36,6 +36,7 @@ from langgraph.store.base import (
|
|
|
36
36
|
ensure_embeddings,
|
|
37
37
|
get_text_at_path,
|
|
38
38
|
tokenize_path,
|
|
39
|
+
validate_op_namespace,
|
|
39
40
|
)
|
|
40
41
|
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
|
41
42
|
from psycopg.rows import DictRow, dict_row
|
|
@@ -354,7 +355,7 @@ class BasePostgresStore(Generic[C]):
|
|
|
354
355
|
(
|
|
355
356
|
_namespace_to_text(op.namespace),
|
|
356
357
|
op.key,
|
|
357
|
-
Jsonb(cast(
|
|
358
|
+
Jsonb(dict(cast(Mapping[str, Any], op.value))),
|
|
358
359
|
)
|
|
359
360
|
)
|
|
360
361
|
if op.ttl is not None:
|
|
@@ -1386,6 +1387,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
|
|
|
1386
1387
|
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
|
|
1387
1388
|
tot = 0
|
|
1388
1389
|
for idx, op in enumerate(ops):
|
|
1390
|
+
validate_op_namespace(op)
|
|
1389
1391
|
grouped_ops[type(op)].append((idx, op))
|
|
1390
1392
|
tot += 1
|
|
1391
1393
|
return grouped_ops, tot
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "langgraph-checkpoint-postgres"
|
|
7
|
-
version = "3.1.
|
|
7
|
+
version = "3.1.3"
|
|
8
8
|
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
|
9
9
|
authors = []
|
|
10
10
|
requires-python = ">=3.10"
|
|
@@ -12,7 +12,7 @@ readme = "README.md"
|
|
|
12
12
|
license = "MIT"
|
|
13
13
|
license-files = ['LICENSE']
|
|
14
14
|
dependencies = [
|
|
15
|
-
"langgraph-checkpoint>=4.
|
|
15
|
+
"langgraph-checkpoint>=4.3.0,<5.0.0",
|
|
16
16
|
"orjson>=3.11.5",
|
|
17
17
|
"psycopg>=3.2.0",
|
|
18
18
|
"psycopg-pool>=3.2.0",
|
|
@@ -13,8 +13,10 @@ import pytest
|
|
|
13
13
|
from langchain_core.embeddings import Embeddings
|
|
14
14
|
from langgraph.store.base import (
|
|
15
15
|
GetOp,
|
|
16
|
+
InvalidNamespaceError,
|
|
16
17
|
Item,
|
|
17
18
|
ListNamespacesOp,
|
|
19
|
+
MatchCondition,
|
|
18
20
|
PutOp,
|
|
19
21
|
SearchOp,
|
|
20
22
|
)
|
|
@@ -871,3 +873,65 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None
|
|
|
871
873
|
page2 = await store.asearch(ns, limit=2, offset=2)
|
|
872
874
|
assert [i.key for i in page1] == ["a", "b"]
|
|
873
875
|
assert [i.key for i in page2] == ["c"]
|
|
876
|
+
|
|
877
|
+
|
|
878
|
+
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
|
879
|
+
async def test_abatch_rejects_invalid_namespace_labels(
|
|
880
|
+
store: AsyncPostgresStore, namespace: tuple
|
|
881
|
+
) -> None:
|
|
882
|
+
await store.aput(("foo", "bar"), "key", {"original": True})
|
|
883
|
+
|
|
884
|
+
for op in (
|
|
885
|
+
GetOp(namespace, "key"),
|
|
886
|
+
GetOp(namespace, "key", refresh_ttl=True),
|
|
887
|
+
PutOp(namespace, "key", {"changed": True}),
|
|
888
|
+
PutOp(namespace, "key", None),
|
|
889
|
+
SearchOp(namespace),
|
|
890
|
+
ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
|
891
|
+
ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
|
892
|
+
):
|
|
893
|
+
with pytest.raises(InvalidNamespaceError):
|
|
894
|
+
await store.abatch([op])
|
|
895
|
+
|
|
896
|
+
item = await store.aget(("foo", "bar"), "key")
|
|
897
|
+
assert item is not None and item.value == {"original": True}
|
|
898
|
+
|
|
899
|
+
|
|
900
|
+
async def test_invalid_namespace_only_fails_its_own_call(
|
|
901
|
+
store: AsyncPostgresStore,
|
|
902
|
+
) -> None:
|
|
903
|
+
"""Concurrent calls share one `abatch`, which fails every op if it raises.
|
|
904
|
+
|
|
905
|
+
Labels are checked before an op is queued, so one caller's bad label cannot
|
|
906
|
+
fail another caller's request.
|
|
907
|
+
"""
|
|
908
|
+
await store.aput(("foo", "bar"), "key", {"original": True})
|
|
909
|
+
|
|
910
|
+
valid, invalid = await asyncio.gather(
|
|
911
|
+
store.aget(("foo", "bar"), "key"),
|
|
912
|
+
store.aget(("foo.bar",), "key"),
|
|
913
|
+
return_exceptions=True,
|
|
914
|
+
)
|
|
915
|
+
|
|
916
|
+
assert isinstance(valid, Item) and valid.value == {"original": True}
|
|
917
|
+
assert isinstance(invalid, InvalidNamespaceError)
|
|
918
|
+
|
|
919
|
+
|
|
920
|
+
async def test_sync_methods_reject_invalid_namespace_labels(
|
|
921
|
+
store: AsyncPostgresStore,
|
|
922
|
+
) -> None:
|
|
923
|
+
"""The sync wrappers run off the event loop thread and must validate too."""
|
|
924
|
+
await store.aput(("foo", "bar"), "key", {"original": True})
|
|
925
|
+
|
|
926
|
+
for call in (
|
|
927
|
+
lambda: store.get(("foo.bar",), "key"),
|
|
928
|
+
lambda: store.search(("foo.bar",)),
|
|
929
|
+
lambda: store.delete(("foo.bar",), "key"),
|
|
930
|
+
lambda: store.list_namespaces(prefix=("foo.bar",)),
|
|
931
|
+
lambda: store.batch([GetOp(("foo.bar",), "key")]),
|
|
932
|
+
):
|
|
933
|
+
with pytest.raises(InvalidNamespaceError):
|
|
934
|
+
await asyncio.to_thread(call)
|
|
935
|
+
|
|
936
|
+
item = await store.aget(("foo", "bar"), "key")
|
|
937
|
+
assert item is not None and item.value == {"original": True}
|