langgraph-checkpoint-postgres 3.1.0a1__tar.gz → 3.1.0a3__tar.gz

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-checkpoint-postgres
-Version: 3.1.0a1
+Version: 3.1.0a3
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Project-URL: Source, https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres
 Project-URL: Twitter, https://x.com/langchain_oss
@@ -9,7 +9,7 @@ Project-URL: Reddit, https://www.reddit.com/r/LangChain/
 License-Expression: MIT
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: langgraph-checkpoint<5.0.0,>=4.1.0a1
+Requires-Dist: langgraph-checkpoint<5.0.0,>=4.1.0a3
 Requires-Dist: orjson>=3.11.5
 Requires-Dist: psycopg-pool>=3.2.0
 Requires-Dist: psycopg>=3.2.0

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/langgraph/checkpoint/postgres/__init__.py

@@ -19,6 +19,7 @@ from langgraph.checkpoint.base import (
     get_serializable_checkpoint_metadata,
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.types import _DeltaSnapshot
 from psycopg import Capabilities, Connection, Cursor, Pipeline
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
@@ -26,9 +27,11 @@ from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.postgres import _internal
 from langgraph.checkpoint.postgres.base import (
-    SELECT_DELTA_COMBINED_SQL,
+    SELECT_DELTA_STAGE1_SQL,
+    SELECT_DELTA_STAGE2_SQL,
     BasePostgresSaver,
-    _DeltaCombinedRow,
+    _DeltaStage1Row,
+    _DeltaStage2Row,
 )
 from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
 
@@ -308,7 +311,12 @@ class PostgresSaver(BasePostgresSaver):
         # others are stored in blobs table
         blob_values = {}
         for k, v in checkpoint["channel_values"].items():
-            if v is None or isinstance(v, (str, int, float, bool)):
+            if v is DELTA_SENTINEL:
+                copy["channel_values"].pop(k)
+            elif isinstance(v, _DeltaSnapshot):
+                blob_values[k] = copy["channel_values"].pop(k)
+                copy["channel_values"][k] = True
+            elif v is None or isinstance(v, (str, int, float, bool)):
                 pass
             else:
                 blob_values[k] = copy["channel_values"].pop(k)
@@ -441,41 +449,49 @@ class PostgresSaver(BasePostgresSaver):
     ) -> _ChannelWritesHistory:
         """Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
 
-        One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
-        from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
-        single roundtrip; the ancestor walk runs in Python.
+        Two-stage query: stage 1 scans checkpoint metadata to walk the parent
+        chain and locate the nearest snapshot; stage 2 fetches only the
+        chain-limited writes and single seed blob.
         """
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         checkpoint_id = get_checkpoint_id(config)
         if checkpoint_id is None:
-            # Caller didn't specify a target — resolve to the latest
-            # checkpoint on the thread. `get_tuple` without `checkpoint_id`
-            # returns the newest; its config carries the resolved id.
             target = self.get_tuple(config)
             if target is None:
                 return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
             checkpoint_id = target.config["configurable"]["checkpoint_id"]
+
         with self._cursor() as cur:
             cur.execute(
-                SELECT_DELTA_COMBINED_SQL,
+                SELECT_DELTA_STAGE1_SQL,
+                (channel, channel, thread_id, checkpoint_ns),
+            )
+            stage1_rows = cur.fetchall()
+        chain_cids, seed_version = self._walk_stage1(
+            cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
+        )
+        seed_versions = [seed_version] if seed_version else []
+        with self._cursor() as cur:
+            cur.execute(
+                SELECT_DELTA_STAGE2_SQL,
                 (
-                    channel,
-                    thread_id,
-                    checkpoint_ns,
                     thread_id,
                     checkpoint_ns,
                     channel,
+                    chain_cids,
                     thread_id,
                     checkpoint_ns,
                     channel,
+                    seed_versions,
                 ),
             )
-            rows = cur.fetchall()
+            stage2_rows = cur.fetchall()
         return self._build_delta_channel_writes_history(
             channel=channel,
-            target_id=checkpoint_id,
-            rows=cast("list[_DeltaCombinedRow]", rows),
+            chain_cids=chain_cids,
+            seed_version=seed_version,
+            stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
         )
 
     def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/langgraph/checkpoint/postgres/aio.py

@@ -19,6 +19,7 @@ from langgraph.checkpoint.base import (
     get_serializable_checkpoint_metadata,
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.types import _DeltaSnapshot
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
@@ -26,9 +27,11 @@ from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.postgres import _ainternal
 from langgraph.checkpoint.postgres.base import (
-    SELECT_DELTA_COMBINED_SQL,
+    SELECT_DELTA_STAGE1_SQL,
+    SELECT_DELTA_STAGE2_SQL,
     BasePostgresSaver,
-    _DeltaCombinedRow,
+    _DeltaStage1Row,
+    _DeltaStage2Row,
 )
 from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
 
@@ -267,7 +270,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
         # others are stored in blobs table
         blob_values = {}
         for k, v in checkpoint["channel_values"].items():
-            if v is None or isinstance(v, (str, int, float, bool)):
+            if v is DELTA_SENTINEL:
+                copy["channel_values"].pop(k)
+            elif isinstance(v, _DeltaSnapshot):
+                blob_values[k] = copy["channel_values"].pop(k)
+                copy["channel_values"][k] = True
+            elif v is None or isinstance(v, (str, int, float, bool)):
                 pass
             else:
                 blob_values[k] = copy["channel_values"].pop(k)
@@ -402,10 +410,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
     ) -> _ChannelWritesHistory:
         """Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
 
-        One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
-        from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
-        single roundtrip; rows are assembled by the shared pure helper on
-        `BasePostgresSaver`.
+        Two-stage query: stage 1 scans checkpoint metadata to walk the parent
+        chain and locate the nearest snapshot; stage 2 fetches only the
+        chain-limited writes and single seed blob.
         """
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -415,26 +422,37 @@ class AsyncPostgresSaver(BasePostgresSaver):
             if target is None:
                 return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
             checkpoint_id = target.config["configurable"]["checkpoint_id"]
+
         async with self._cursor() as cur:
             await cur.execute(
-                SELECT_DELTA_COMBINED_SQL,
+                SELECT_DELTA_STAGE1_SQL,
+                (channel, channel, thread_id, checkpoint_ns),
+            )
+            stage1_rows = await cur.fetchall()
+        chain_cids, seed_version = self._walk_stage1(
+            cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
+        )
+        seed_versions = [seed_version] if seed_version else []
+        async with self._cursor() as cur:
+            await cur.execute(
+                SELECT_DELTA_STAGE2_SQL,
                 (
-                    channel,
-                    thread_id,
-                    checkpoint_ns,
                     thread_id,
                     checkpoint_ns,
                     channel,
+                    chain_cids,
                     thread_id,
                     checkpoint_ns,
                     channel,
+                    seed_versions,
                 ),
             )
-            rows = await cur.fetchall()
+            stage2_rows = await cur.fetchall()
         return self._build_delta_channel_writes_history(
             channel=channel,
-            target_id=checkpoint_id,
-            rows=cast("list[_DeltaCombinedRow]", rows),
+            chain_cids=chain_cids,
+            seed_version=seed_version,
+            stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
         )
 
     async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/langgraph/checkpoint/postgres/base.py

@@ -156,62 +156,62 @@ INSERT_CHECKPOINT_WRITES_SQL = """
 """
 
 
-class _DeltaCombinedRow(TypedDict, total=False):
-    """One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables).
+class _DeltaStage2Row(TypedDict, total=False):
+    """One row from `SELECT_DELTA_STAGE2_SQL` (a UNION ALL of writes and blobs)."""
 
-    Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are
-    relevant for that kind; irrelevant columns are NULL and typed as `None`.
-    """
-
-    _kind: str  # always present: "p", "w", or "b"
-    # checkpoint row ("p")
-    checkpoint_id: str | None
-    parent_checkpoint_id: str | None
-    ver: str | None
-    # write / blob rows ("w", "b")
+    _kind: str  # "w" or "b"
+    checkpoint_id: str | None  # "w" rows only
     type: str | None
     blob: bytes | None
-    # write row only ("w")
-    task_id: str | None
-    idx: int | None
-    # blob row only ("b")
-    version: str | None
+    task_id: str | None  # "w" rows only
+    idx: int | None  # "w" rows only
+    version: str | None  # "b" rows only
 
 
-# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints,
-# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs
-# in Python in `_build_delta_channel_writes_history`.
+# Two-stage DeltaChannel reconstruction.  Stage 1 scans checkpoint
+# metadata (no blob bytes) to walk the parent chain and locate the
+# nearest snapshot marker.  Stage 2 fetches only the chain-limited
+# writes and the single seed snapshot blob.
 #
-# 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,
+# Parameter order:
+#   stage1: (channel, channel, thread_id, checkpoint_ns)
+#   stage2: (thread_id, checkpoint_ns, channel, chain_cids[],
+#            thread_id, checkpoint_ns, channel, seed_versions[])
+
+SELECT_DELTA_STAGE1_SQL = """
+    SELECT checkpoint_id,
            parent_checkpoint_id,
            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
+           (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS has_snapshot
     FROM checkpoints
     WHERE thread_id = %s AND checkpoint_ns = %s
-    UNION ALL
-    SELECT 'w',
-           checkpoint_id, NULL, NULL,
-           type, blob, task_id, idx, NULL
+"""
+
+SELECT_DELTA_STAGE2_SQL = """
+    SELECT 'w'::text AS _kind,
+           checkpoint_id,
+           type, blob, task_id, idx, NULL::text AS version
     FROM checkpoint_writes
     WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
+      AND checkpoint_id = ANY(%s)
     UNION ALL
-    SELECT 'b',
-           NULL, NULL, NULL,
+    SELECT 'b', NULL,
            type, blob, NULL, NULL, version
     FROM checkpoint_blobs
     WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
+      AND version = ANY(%s)
 """
 
 
+class _DeltaStage1Row(TypedDict):
+    """One row from `SELECT_DELTA_STAGE1_SQL`."""
+
+    checkpoint_id: str
+    parent_checkpoint_id: str | None
+    ver: str | None
+    has_snapshot: bool
+
+
 class BasePostgresSaver(BaseCheckpointSaver[str]):
     SELECT_SQL = SELECT_SQL
     SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
@@ -254,38 +254,59 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
             if t.decode() != "empty"
         }
 
+    @staticmethod
+    def _walk_stage1(
+        stage1_rows: Sequence[_DeltaStage1Row],
+        target_id: str,
+    ) -> tuple[list[str], str | None]:
+        """Walk the parent chain from stage 1 metadata rows.
+
+        Returns (chain_cids, seed_version):
+          chain_cids: ancestor checkpoint IDs from target's parent down to
+                      the seed (or root), in newest-first order.
+          seed_version: the channel blob version at the nearest ancestor
+                        with has_snapshot=True, or None if pure delta.
+        """
+        parent_of: dict[str, str | None] = {}
+        ver_of: dict[str, str | None] = {}
+        snapshot_of: dict[str, bool] = {}
+        for r in stage1_rows:
+            cid = r["checkpoint_id"]
+            parent_of[cid] = r["parent_checkpoint_id"]
+            ver_of[cid] = r["ver"]
+            snapshot_of[cid] = r["has_snapshot"]
+
+        chain_cids: list[str] = []
+        seed_version: str | None = None
+        cur_cid: str | None = parent_of.get(target_id)
+        while cur_cid is not None:
+            chain_cids.append(cur_cid)
+            if snapshot_of.get(cur_cid, False):
+                seed_version = ver_of.get(cur_cid)
+                break
+            cur_cid = parent_of.get(cur_cid)
+        return chain_cids, seed_version
+
     def _build_delta_channel_writes_history(
         self,
         *,
         channel: str,
-        target_id: str,
-        rows: Sequence[_DeltaCombinedRow],
+        chain_cids: list[str],
+        seed_version: str | None,
+        stage2_rows: Sequence[_DeltaStage2Row],
     ) -> _ChannelWritesHistory:
-        """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 `SELECT_DELTA_COMBINED_SQL`
-        and feed the tagged rows here.
+        """Reconstruct delta channel history from two-stage query results.
 
-        Walk is newest → oldest from the target's parent. A non-sentinel
-        blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
-        walk and is returned as the seed so replay starts from it.
-
-        Writes stored at `target_id` itself are pending writes for the next
-        step and are excluded — the walk begins at the target's parent.
+        chain_cids are in newest-first order (target's parent first).
+        stage2_rows contain only writes for chain_cids and the single
+        seed blob at seed_version.
         """
-        parent_of: dict[str, str | None] = {}
-        ver_of: dict[str, str | None] = {}
         writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
-        blob_by_ver: dict[str, tuple[str, bytes]] = {}
+        seed_blob: tuple[str, bytes] | None = None
 
-        for r in rows:
+        for r in stage2_rows:
             kind = r["_kind"]
-            if kind == "p":
-                cid = cast(str, r["checkpoint_id"])
-                parent_of[cid] = r["parent_checkpoint_id"]
-                ver_of[cid] = r["ver"]
-            elif kind == "w":
+            if kind == "w":
                 cid = cast(str, r["checkpoint_id"])
                 writes_by_cid.setdefault(cid, []).append(
                     cast(
@@ -294,42 +315,26 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
                     )
                 )
             else:  # kind == "b"
-                blob_by_ver[cast(str, r["version"])] = cast(
-                    "tuple[str, bytes]", (r["type"], r["blob"])
-                )
+                seed_blob = cast("tuple[str, bytes]", (r["type"], r["blob"]))
 
-        # newest write first per ancestor (task_id DESC, idx DESC)
         for ws in writes_by_cid.values():
             ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
 
-        ancestors: list[str] = []
-        cur_cid: str | None = parent_of.get(target_id)
-        while cur_cid is not None:
-            ancestors.append(cur_cid)
-            cur_cid = parent_of.get(cur_cid)
-        if not ancestors:
+        if not chain_cids:
             return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
 
-        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.
+        collected: list[PendingWrite] = []
+        for cid in chain_cids:
             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:
-                        collected.reverse()
-                        return _ChannelWritesHistory(seed=blob_value, writes=collected)
-
-        collected.reverse()  # oldest → newest
-        return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
+
+        seed: Any = DELTA_SENTINEL
+        if seed_blob is not None and seed_blob[0] != "empty":
+            seed = self.serde.loads_typed(seed_blob)
+
+        collected.reverse()
+        return _ChannelWritesHistory(seed=seed, writes=collected)
 
     def _dump_blobs(
         self,

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-checkpoint-postgres"
-version = "3.1.0a1"
+version = "3.1.0a3"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 requires-python = ">=3.10"
@@ -12,7 +12,7 @@ readme = "README.md"
 license = "MIT"
 license-files = ['LICENSE']
 dependencies = [
-  "langgraph-checkpoint>=4.1.0a1,<5.0.0",
+  "langgraph-checkpoint>=4.1.0a3,<5.0.0",
   "orjson>=3.11.5",
   "psycopg>=3.2.0",
   "psycopg-pool>=3.2.0",

{langgraph_checkpoint_postgres-3.1.0a1 → langgraph_checkpoint_postgres-3.1.0a3}/uv.lock

@@ -259,7 +259,7 @@ wheels = [
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "4.1.0a1"
+version = "4.1.0a3"
 source = { editable = "../checkpoint" }
 dependencies = [
     { name = "langchain-core" },
@@ -307,7 +307,7 @@ test = [
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "3.1.0a1"
+version = "3.1.0a3"
 source = { editable = "." }
 dependencies = [
     { name = "langgraph-checkpoint" },