langgraph-checkpoint-postgres 3.1.2__tar.gz → 3.2.0__tar.gz

{langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.2.0}/.gitignore

@@ -76,7 +76,12 @@ __pypackages__/
 
 # Environments
 .env
+.env.*
 .envrc
+*.crt
+*.key
+*.pem
+credentials.json
 .venv
 .venvs
 env/
@@ -98,6 +103,7 @@ dmypy.json
 
 .vercel
 .turbo
+node_modules/
 .editorconfig
 .scratch
 .worktrees/

{langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: langgraph-checkpoint-postgres
-Version: 3.1.2
+Version: 3.2.0
 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.0
+Requires-Dist: langgraph-checkpoint<5.0.0,>=4.3.0
 Requires-Dist: orjson>=3.11.5
 Requires-Dist: psycopg-pool>=3.2.0
 Requires-Dist: psycopg>=3.2.0

{langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.2.0}/langgraph/checkpoint/postgres/__init__.py

@@ -27,10 +27,15 @@ from psycopg_pool import ConnectionPool
 from langgraph.checkpoint.postgres import _internal
 from langgraph.checkpoint.postgres.base import (
     _DELTA_PAGE_SIZE,
+    PRUNE_DELETE_THREADS_SQL,
+    PRUNE_STRATEGIES,
     BasePostgresSaver,
     _build_delta_stage1_sql,
     _build_delta_stage2_sql,
     _DeltaStage2Row,
+    _execute,
+    _keep_latest,
+    _refuse_delta_threads,
 )
 from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
 
@@ -401,6 +406,73 @@ class PostgresSaver(BasePostgresSaver):
                 (str(thread_id),),
             )
 
+    def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        `"delete"` removes every checkpoint, write and blob of the threads, in
+        one transaction.
+
+        `"keep_latest"` keeps the latest checkpoint of each namespace (the one
+        `get_tuple` loads, plus its parent when it predates v4 and has pending
+        sends), with its pending writes and the blobs it references, and
+        deletes the rest. A kept checkpoint whose parent it deletes is left with
+        no parent. A subgraph call's namespace is deleted with the root
+        checkpoint the call started from, except a `checkpointer=True`
+        subgraph's, which keeps its latest checkpoint. Each thread is pruned in
+        its own transaction, on one snapshot of it, so a run appending to the
+        thread meanwhile is safe; forking or replaying from a checkpoint being
+        deleted is not. A thread that fails because another transaction changed
+        it leaves the threads before it pruned, and retrying is safe.
+
+        `"keep_latest"` refuses threads that use a `DeltaChannel`, since it
+        rebuilds the latest value from older checkpoints: if any requested
+        thread does, nothing is pruned. A thread that starts using one before
+        its turn in the batch is refused then, after the threads before it.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+
+        Raises:
+            TypeError: If `thread_ids` is a single string.
+            ValueError: If `strategy` is unknown.
+            NotImplementedError: If `strategy` is `"keep_latest"` and a thread
+                uses a `DeltaChannel`.
+            psycopg.errors.SerializationFailure: If another transaction, such
+                as a second prune, changes a thread's rows while `"keep_latest"`
+                prunes it.
+            psycopg.errors.DeadlockDetected: If a second prune of a thread runs
+                at the same time.
+        """
+        if isinstance(thread_ids, str):
+            raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
+        if strategy not in PRUNE_STRATEGIES:
+            raise ValueError(f"Unknown prune strategy: {strategy!r}")
+        thread_ids = [str(t) for t in thread_ids]
+        if not thread_ids:
+            return
+        if strategy == "delete":
+            with self._cursor() as cur, cur.connection.transaction():
+                for sql in PRUNE_DELETE_THREADS_SQL:
+                    cur.execute(sql, (thread_ids,))
+            return
+        with self._cursor() as cur:
+            _execute(cur, _refuse_delta_threads(thread_ids))
+        for thread_id in thread_ids:
+            with self._cursor() as cur:
+                if self.pipe:
+                    # The transaction has to start with its isolation level, and
+                    # in pipeline mode it would take in what ran since the last
+                    # sync.
+                    self.pipe.sync()
+                with cur.connection.transaction():
+                    _execute(cur, _keep_latest(thread_id))
+
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
         """Create a database cursor as a context manager.
@@ -448,11 +520,13 @@ class PostgresSaver(BasePostgresSaver):
 
         Two-stage query, both stages cover ALL requested channels:
 
-        * Stage 1 (paged): dynamic SELECT over `checkpoints` with K parallel
-          JSONB key lookups (one column pair per channel) — no subquery, no
-          aggregation. Pages newest-first by `checkpoint_id` with a cursor;
-          page size is `_DELTA_PAGE_SIZE`. Stops paging when every channel
-          has found its seed or the chain is exhausted.
+        * Stage 1 (paged): dynamic SELECT over `checkpoints` with three
+          columns per channel: its version, an `EXISTS` probe for a stored
+          blob at that version, and its inline value. Pages newest-first by
+          `checkpoint_id`, starting at the target; page size is
+          `_DELTA_PAGE_SIZE`.
+          Stops paging when every channel has found its seed or a page comes
+          back short.
 
         * Stage 2 (per-channel UNION ALL): one branch per channel reading
           `checkpoint_writes` filtered to that channel's specific
@@ -467,7 +541,7 @@ class PostgresSaver(BasePostgresSaver):
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         checkpoint_id = get_checkpoint_id(config)
-        if checkpoint_id is None:
+        if not checkpoint_id:
             target = self.get_tuple(config)
             if target is None:
                 return {ch: {"writes": []} for ch in channels}
@@ -475,7 +549,7 @@ class PostgresSaver(BasePostgresSaver):
 
         # Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
         # Python after each page. Stops as soon as every channel has its seed.
-        stage1_sql = _build_delta_stage1_sql(channels, paged=True)
+        stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
         parent_of: dict[str, str | None] = {}
         ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
         hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
@@ -485,7 +559,7 @@ class PostgresSaver(BasePostgresSaver):
         seed_inline_by_ch: dict[str, Any] = {}
         walk_cursor_by_ch: dict[str, str | None] = {}
         seeded: set[str] = set()
-        cursor: str | None = None
+        cursor: str | None = checkpoint_id
 
         with self._cursor() as cur:
             while True:
@@ -494,7 +568,7 @@ class PostgresSaver(BasePostgresSaver):
                     # ver_i, blob channel, blob version, inline_i
                     stage1_params.extend([ch, ch, ch, ch])
                 stage1_params.extend(
-                    [thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
+                    [thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
                 )
                 cur.execute(stage1_sql, stage1_params)
                 page = cur.fetchall()
@@ -526,6 +600,7 @@ class PostgresSaver(BasePostgresSaver):
                 if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
                     break
                 cursor = oldest
+                stage1_sql = _build_delta_stage1_sql(channels, paged=True)
 
         # Stage 2: per-channel UNION ALL — one writes branch per channel
         # with non-empty chain, plus one blob branch per seeded channel.

{langgraph_checkpoint_postgres-3.1.2 → langgraph_checkpoint_postgres-3.2.0}/langgraph/checkpoint/postgres/aio.py

@@ -27,10 +27,15 @@ from psycopg_pool import AsyncConnectionPool
 from langgraph.checkpoint.postgres import _ainternal
 from langgraph.checkpoint.postgres.base import (
     _DELTA_PAGE_SIZE,
+    PRUNE_DELETE_THREADS_SQL,
+    PRUNE_STRATEGIES,
     BasePostgresSaver,
+    _aexecute,
     _build_delta_stage1_sql,
     _build_delta_stage2_sql,
     _DeltaStage2Row,
+    _keep_latest,
+    _refuse_delta_threads,
 )
 from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
 
@@ -360,6 +365,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 (str(thread_id),),
             )
 
+    async def aprune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        See `PostgresSaver.prune` for what each strategy keeps.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+
+        Raises:
+            TypeError: If `thread_ids` is a single string.
+            ValueError: If `strategy` is unknown.
+            NotImplementedError: If `strategy` is `"keep_latest"` and a thread
+                uses a `DeltaChannel`.
+            psycopg.errors.SerializationFailure: If another transaction, such
+                as a second prune, changes a thread's rows while `"keep_latest"`
+                prunes it.
+            psycopg.errors.DeadlockDetected: If a second prune of a thread runs
+                at the same time.
+        """
+        if isinstance(thread_ids, str):
+            raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
+        if strategy not in PRUNE_STRATEGIES:
+            raise ValueError(f"Unknown prune strategy: {strategy!r}")
+        thread_ids = [str(t) for t in thread_ids]
+        if not thread_ids:
+            return
+        if strategy == "delete":
+            async with self._cursor() as cur, cur.connection.transaction():
+                for sql in PRUNE_DELETE_THREADS_SQL:
+                    await cur.execute(sql, (thread_ids,))
+            return
+        async with self._cursor() as cur:
+            await _aexecute(cur, _refuse_delta_threads(thread_ids))
+        for thread_id in thread_ids:
+            async with self._cursor() as cur:
+                if self.pipe:
+                    # The transaction has to start with its isolation level, and
+                    # in pipeline mode it would take in what ran since the last
+                    # sync.
+                    await self.pipe.sync()
+                async with cur.connection.transaction():
+                    await _aexecute(cur, _keep_latest(thread_id))
+
     @asynccontextmanager
     async def _cursor(
         self, *, pipeline: bool = False
@@ -417,13 +471,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         checkpoint_id = get_checkpoint_id(config)
-        if checkpoint_id is None:
+        if not checkpoint_id:
             target = await self.aget_tuple(config)
             if target is None:
                 return {ch: {"writes": []} for ch in channels}
             checkpoint_id = target.config["configurable"]["checkpoint_id"]
 
-        stage1_sql = _build_delta_stage1_sql(channels, paged=True)
+        stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
         parent_of: dict[str, str | None] = {}
         ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
         hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
@@ -433,7 +487,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         seed_inline_by_ch: dict[str, Any] = {}
         walk_cursor_by_ch: dict[str, str | None] = {}
         seeded: set[str] = set()
-        cursor: str | None = None
+        cursor: str | None = checkpoint_id
 
         async with self._cursor() as cur:
             while True:
@@ -442,7 +496,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     # ver_i, blob channel, blob version, inline_i
                     stage1_params.extend([ch, ch, ch, ch])
                 stage1_params.extend(
-                    [thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
+                    [thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
                 )
                 await cur.execute(stage1_sql, stage1_params)
                 page = await cur.fetchall()
@@ -472,6 +526,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
                     break
                 cursor = oldest
+                stage1_sql = _build_delta_stage1_sql(channels, paged=True)
 
         channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
         channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
@@ -687,5 +742,50 @@ class AsyncPostgresSaver(BasePostgresSaver):
             self.adelete_thread(thread_id), self.loop
         ).result()
 
+    def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        See `PostgresSaver.prune` for what each strategy keeps.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+        """
+        try:
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aprune(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aprune(thread_ids, strategy=strategy), self.loop
+        ).result()
+
+    def get_delta_channel_history(
+        self, *, config: RunnableConfig, channels: Sequence[str]
+    ) -> Mapping[str, DeltaChannelHistory]:
+        """Sync bridge to `aget_delta_channel_history`, guarded like `get_tuple`."""
+        try:
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aget_delta_channel_history(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_delta_channel_history(config=config, channels=channels),
+            self.loop,
+        ).result()
+
 
 __all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]