chore: import upstream snapshot with attribution

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -0,0 +1,595 @@
+from __future__ import annotations
+
+import threading
+from collections import defaultdict
+from collections.abc import Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from typing import Any, cast
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    DeltaChannelHistory,
+    get_checkpoint_id,
+    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
+from psycopg_pool import ConnectionPool
+
+from langgraph.checkpoint.postgres import _internal
+from langgraph.checkpoint.postgres.base import (
+    _DELTA_PAGE_SIZE,
+    BasePostgresSaver,
+    _build_delta_stage1_sql,
+    _build_delta_stage2_sql,
+    _DeltaStage2Row,
+)
+from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
+
+Conn = _internal.Conn  # For backward compatibility
+
+
+class PostgresSaver(BasePostgresSaver):
+    """Checkpointer that stores checkpoints in a Postgres database."""
+
+    lock: threading.Lock
+
+    def __init__(
+        self,
+        conn: _internal.Conn,
+        pipe: Pipeline | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        if isinstance(conn, ConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single Connection, not ConnectionPool."
+            )
+
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = threading.Lock()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls, conn_string: str, *, pipeline: bool = False
+    ) -> Iterator[PostgresSaver]:
+        """Create a new PostgresSaver instance from a connection string.
+
+        Args:
+            conn_string: The Postgres connection info string.
+            pipeline: whether to use Pipeline
+
+        Returns:
+            PostgresSaver: A new PostgresSaver instance.
+        """
+        with Connection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                with conn.pipeline() as pipe:
+                    yield cls(conn, pipe)
+            else:
+                yield cls(conn)
+
+    def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        with self._cursor() as cur:
+            cur.execute(self.MIGRATIONS[0])
+            results = cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = results.fetchone()
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+                strict=False,
+            ):
+                cur.execute(migration)
+                cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,))
+        if self.pipe:
+            self.pipe.sync()
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config: The config to use for listing the checkpoints.
+            filter: Additional filtering criteria for metadata.
+            before: If provided, only checkpoints before the specified checkpoint ID are returned.
+            limit: The maximum number of checkpoints to return.
+
+        Yields:
+            An iterator of checkpoint tuples.
+
+        Examples:
+            >>> from langgraph.checkpoint.postgres import PostgresSaver
+            >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     config = {"configurable": {"thread_id": "1"}}
+            >>>     checkpoints = list(memory.list(config, limit=2))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), CheckpointTuple(...)]
+
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     checkpoints = list(memory.list(config, before=before))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), ...]
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        params = list(args)
+        if limit is not None:
+            query += " LIMIT %s"
+            params.append(int(limit))
+        # if we change this to use .stream() we need to make sure to close the cursor
+        with self._cursor() as cur:
+            cur.execute(query, params)
+            values = cur.fetchall()
+            if not values:
+                return
+            # migrate pending sends if necessary
+            if to_migrate := [
+                v
+                for v in values
+                if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
+            ]:
+                cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (
+                        values[0]["thread_id"],
+                        [v["parent_checkpoint_id"] for v in to_migrate],
+                    ),
+                )
+                grouped_by_parent = defaultdict(list)
+                for value in to_migrate:
+                    grouped_by_parent[value["parent_checkpoint_id"]].append(value)
+                for sends in cur:
+                    for value in grouped_by_parent[sends["checkpoint_id"]]:
+                        if value["channel_values"] is None:
+                            value["channel_values"] = []
+                        self._migrate_pending_sends(
+                            sends["sends"],
+                            value["checkpoint"],
+                            value["channel_values"],
+                        )
+            for value in values:
+                yield self._load_checkpoint_tuple(value)
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a `checkpoint_id` key, the checkpoint with
+        the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+
+        Examples:
+
+            Basic:
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+
+            With timestamp:
+
+            >>> config = {
+            ...    "configurable": {
+            ...        "thread_id": "1",
+            ...        "checkpoint_ns": "",
+            ...        "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+            ...    }
+            ... }
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+        """  # noqa
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
+        else:
+            args = (thread_id, checkpoint_ns)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
+
+        with self._cursor() as cur:
+            cur.execute(
+                self.SELECT_SQL + where,
+                args,
+            )
+            value = cur.fetchone()
+            if value is None:
+                return None
+
+            # migrate pending sends if necessary
+            if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
+                cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (thread_id, [value["parent_checkpoint_id"]]),
+                )
+                if sends := cur.fetchone():
+                    if value["channel_values"] is None:
+                        value["channel_values"] = []
+                    self._migrate_pending_sends(
+                        sends["sends"],
+                        value["checkpoint"],
+                        value["channel_values"],
+                    )
+
+            return self._load_checkpoint_tuple(value)
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Examples:
+
+            >>> from langgraph.checkpoint.postgres import PostgresSaver
+            >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
+            >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
+            >>> print(saved_config)
+            {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop("checkpoint_id", None)
+        copy = checkpoint.copy()
+        copy["channel_values"] = copy["channel_values"].copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        # inline primitive values in checkpoint table
+        # others are stored in blobs table
+        blob_values = {}
+        for k, v in checkpoint["channel_values"].items():
+            if 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)
+
+        with self._cursor(pipeline=True) as cur:
+            if blob_versions := {
+                k: v for k, v in new_versions.items() if k in blob_values
+            }:
+                cur.executemany(
+                    self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                    self._dump_blobs(
+                        thread_id,
+                        checkpoint_ns,
+                        blob_values,
+                        blob_versions,
+                    ),
+                )
+            cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    Jsonb(copy),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
+                ),
+            )
+        return next_config
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the Postgres database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store.
+            task_id: Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        with self._cursor(pipeline=True) as cur:
+            cur.executemany(
+                query,
+                self._dump_writes(
+                    config["configurable"]["thread_id"],
+                    config["configurable"]["checkpoint_ns"],
+                    config["configurable"]["checkpoint_id"],
+                    task_id,
+                    task_path,
+                    writes,
+                ),
+            )
+
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id: The thread ID to delete.
+
+        Returns:
+            None
+        """
+        with self._cursor(pipeline=True) as cur:
+            cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            cur.execute(
+                "DELETE FROM checkpoint_blobs WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            cur.execute(
+                "DELETE FROM checkpoint_writes WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+
+    @contextmanager
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the PostgresSaver instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        with self.lock, _internal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                if self.supports_pipeline:
+                    with (
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+                else:
+                    # Use connection's transaction context manager when pipeline mode not supported
+                    with (
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+            else:
+                with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                    yield cur
+
+    def get_delta_channel_history(
+        self, *, config: RunnableConfig, channels: Sequence[str]
+    ) -> Mapping[str, DeltaChannelHistory]:
+        """Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
+
+        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 2 (per-channel UNION ALL): one branch per channel reading
+          `checkpoint_writes` filtered to that channel's specific
+          `chain_cids`, plus one branch per channel that has a seed reading
+          `checkpoint_blobs` for that channel + version. Avoids the
+          over-fetch of a single `channel = ANY(channels)` filter when
+          channels have different chain depths.
+        """
+        if not channels:
+            return {}
+        channels = list(channels)
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_id = get_checkpoint_id(config)
+        if checkpoint_id is None:
+            target = self.get_tuple(config)
+            if target is None:
+                return {ch: {"writes": []} for ch in channels}
+            checkpoint_id = target.config["configurable"]["checkpoint_id"]
+
+        # 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)
+        parent_of: dict[str, str | None] = {}
+        ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
+        hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
+        chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
+        seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
+        walk_cursor_by_ch: dict[str, str | None] = {}
+        seeded: set[str] = set()
+        cursor: str | None = None
+
+        with self._cursor() as cur:
+            while True:
+                stage1_params: list[Any] = []
+                for ch in channels:
+                    stage1_params.extend([ch, ch])
+                stage1_params.extend(
+                    [thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
+                )
+                cur.execute(stage1_sql, stage1_params)
+                page = cur.fetchall()
+                if not page:
+                    break
+                oldest = self._ingest_stage1_page(
+                    cast("list[Mapping[str, Any]]", page),
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hs_by_i_by_cid,
+                )
+                self._try_advance_walks(
+                    checkpoint_id,
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hs_by_i_by_cid,
+                    chain_by_ch,
+                    seed_ver_by_ch,
+                    walk_cursor_by_ch,
+                    seeded,
+                )
+                # Stop if every channel is seeded, or the page was short
+                # (chain exhausted — no more rows to fetch).
+                if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
+                    break
+                cursor = oldest
+
+        # Stage 2: per-channel UNION ALL — one writes branch per channel
+        # with non-empty chain, plus one blob branch per seeded channel.
+        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]
+        stage2_sql = _build_delta_stage2_sql(
+            channels_with_chain=channels_with_chain,
+            channels_with_seed=channels_with_seed,
+        )
+
+        if stage2_sql:
+            stage2_params: list[Any] = []
+            for ch in channels_with_chain:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]])
+            for ch in channels_with_seed:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
+            with self._cursor() as cur:
+                cur.execute(stage2_sql, stage2_params)
+                stage2_rows = cur.fetchall()
+        else:
+            stage2_rows = []
+
+        return self._build_delta_channels_writes_history(
+            channels=channels,
+            chain_by_ch=chain_by_ch,
+            seed_ver_by_ch=seed_ver_by_ch,
+            stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
+        )
+
+    def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
+        """
+        Convert a database row into a CheckpointTuple object.
+
+        Args:
+            value: A row from the database containing checkpoint data.
+
+        Returns:
+            CheckpointTuple: A structured representation of the checkpoint,
+            including its configuration, metadata, parent checkpoint (if any),
+            and pending writes.
+        """
+        return CheckpointTuple(
+            {
+                "configurable": {
+                    "thread_id": value["thread_id"],
+                    "checkpoint_ns": value["checkpoint_ns"],
+                    "checkpoint_id": value["checkpoint_id"],
+                }
+            },
+            {
+                **value["checkpoint"],
+                "channel_values": {
+                    **(value["checkpoint"].get("channel_values") or {}),
+                    **self._load_blobs(value["channel_values"]),
+                },
+            },
+            value["metadata"],
+            (
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["parent_checkpoint_id"],
+                    }
+                }
+                if value["parent_checkpoint_id"]
+                else None
+            ),
+            self._load_writes(value["pending_writes"]),
+        )
+
+
+__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"]

libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py

@@ -0,0 +1,23 @@
+"""Shared async utility functions for the Postgres checkpoint & storage classes."""
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+from psycopg import AsyncConnection
+from psycopg.rows import DictRow
+from psycopg_pool import AsyncConnectionPool
+
+Conn = AsyncConnection[DictRow] | AsyncConnectionPool[AsyncConnection[DictRow]]
+
+
+@asynccontextmanager
+async def get_connection(
+    conn: Conn,
+) -> AsyncIterator[AsyncConnection[DictRow]]:
+    if isinstance(conn, AsyncConnection):
+        yield conn
+    elif isinstance(conn, AsyncConnectionPool):
+        async with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")

libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py

@@ -0,0 +1,21 @@
+"""Shared utility functions for the Postgres checkpoint & storage classes."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+from psycopg import Connection
+from psycopg.rows import DictRow
+from psycopg_pool import ConnectionPool
+
+Conn = Connection[DictRow] | ConnectionPool[Connection[DictRow]]
+
+
+@contextmanager
+def get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
+    if isinstance(conn, Connection):
+        yield conn
+    elif isinstance(conn, ConnectionPool):
+        with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -0,0 +1,684 @@
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from contextlib import asynccontextmanager
+from typing import Any, cast
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    DeltaChannelHistory,
+    get_checkpoint_id,
+    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
+from psycopg_pool import AsyncConnectionPool
+
+from langgraph.checkpoint.postgres import _ainternal
+from langgraph.checkpoint.postgres.base import (
+    _DELTA_PAGE_SIZE,
+    BasePostgresSaver,
+    _build_delta_stage1_sql,
+    _build_delta_stage2_sql,
+    _DeltaStage2Row,
+)
+from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
+
+Conn = _ainternal.Conn  # For backward compatibility
+
+
+class AsyncPostgresSaver(BasePostgresSaver):
+    """Asynchronous checkpointer that stores checkpoints in a Postgres database."""
+
+    lock: asyncio.Lock
+
+    def __init__(
+        self,
+        conn: _ainternal.Conn,
+        pipe: AsyncPipeline | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        serde: SerializerProtocol | None = None,
+    ) -> AsyncIterator[AsyncPostgresSaver]:
+        """Create a new AsyncPostgresSaver instance from a connection string.
+
+        Args:
+            conn_string: The Postgres connection info string.
+            pipeline: whether to use AsyncPipeline
+
+        Returns:
+            AsyncPostgresSaver: A new AsyncPostgresSaver instance.
+        """
+        async with await AsyncConnection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                async with conn.pipeline() as pipe:
+                    yield cls(conn=conn, pipe=pipe, serde=serde)
+            else:
+                yield cls(conn=conn, serde=serde)
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        async with self._cursor() as cur:
+            await cur.execute(self.MIGRATIONS[0])
+            results = await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await results.fetchone()
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+                strict=False,
+            ):
+                await cur.execute(migration)
+                await cur.execute(
+                    "INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
+                )
+        if self.pipe:
+            await self.pipe.sync()
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        """List checkpoints from the database asynchronously.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config: Base configuration for filtering checkpoints.
+            filter: Additional filtering criteria for metadata.
+            before: If provided, only checkpoints before the specified checkpoint ID are returned.
+            limit: Maximum number of checkpoints to return.
+
+        Yields:
+            An asynchronous iterator of matching checkpoint tuples.
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        params = list(args)
+        if limit is not None:
+            query += " LIMIT %s"
+            params.append(int(limit))
+        # if we change this to use .stream() we need to make sure to close the cursor
+        async with self._cursor() as cur:
+            await cur.execute(query, params, binary=True)
+            values = await cur.fetchall()
+            if not values:
+                return
+            # migrate pending sends if necessary
+            if to_migrate := [
+                v
+                for v in values
+                if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
+            ]:
+                await cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (
+                        values[0]["thread_id"],
+                        [v["parent_checkpoint_id"] for v in to_migrate],
+                    ),
+                )
+                grouped_by_parent = defaultdict(list)
+                for value in to_migrate:
+                    grouped_by_parent[value["parent_checkpoint_id"]].append(value)
+                async for sends in cur:
+                    for value in grouped_by_parent[sends["checkpoint_id"]]:
+                        if value["channel_values"] is None:
+                            value["channel_values"] = []
+                        self._migrate_pending_sends(
+                            sends["sends"],
+                            value["checkpoint"],
+                            value["channel_values"],
+                        )
+            for value in values:
+                yield await self._load_checkpoint_tuple(value)
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database asynchronously.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a `checkpoint_id` key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
+        else:
+            args = (thread_id, checkpoint_ns)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
+
+        async with self._cursor() as cur:
+            await cur.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+            value = await cur.fetchone()
+            if value is None:
+                return None
+
+            # migrate pending sends if necessary
+            if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
+                await cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (thread_id, [value["parent_checkpoint_id"]]),
+                )
+                if sends := await cur.fetchone():
+                    if value["channel_values"] is None:
+                        value["channel_values"] = []
+                    self._migrate_pending_sends(
+                        sends["sends"],
+                        value["checkpoint"],
+                        value["channel_values"],
+                    )
+
+            return await self._load_checkpoint_tuple(value)
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database asynchronously.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop("checkpoint_id", None)
+
+        copy = checkpoint.copy()
+        copy["channel_values"] = copy["channel_values"].copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        # inline primitive values in checkpoint table
+        # others are stored in blobs table
+        blob_values = {}
+        for k, v in checkpoint["channel_values"].items():
+            if 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)
+
+        async with self._cursor(pipeline=True) as cur:
+            if blob_versions := {
+                k: v for k, v in new_versions.items() if k in blob_values
+            }:
+                await cur.executemany(
+                    self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                    await asyncio.to_thread(
+                        self._dump_blobs,
+                        thread_id,
+                        checkpoint_ns,
+                        blob_values,
+                        blob_versions,
+                    ),
+                )
+            await cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    Jsonb(copy),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
+                ),
+            )
+        return next_config
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            task_path,
+            writes,
+        )
+        async with self._cursor(pipeline=True) as cur:
+            await cur.executemany(query, params)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id: The thread ID to delete.
+
+        Returns:
+            None
+        """
+        async with self._cursor(pipeline=True) as cur:
+            await cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_blobs WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_writes WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+
+    @asynccontextmanager
+    async def _cursor(
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        async with self.lock, _ainternal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        await self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                if self.supports_pipeline:
+                    async with (
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+                else:
+                    # Use connection's transaction context manager when pipeline mode not supported
+                    async with (
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+            else:
+                async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                    yield cur
+
+    async def aget_delta_channel_history(
+        self, *, config: RunnableConfig, channels: Sequence[str]
+    ) -> Mapping[str, DeltaChannelHistory]:
+        """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
+
+        See `PostgresSaver.get_delta_channel_history` for design notes; this is
+        the async equivalent with internal stage-1 paging and per-channel
+        UNION ALL stage-2.
+        """
+        if not channels:
+            return {}
+        channels = list(channels)
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_id = get_checkpoint_id(config)
+        if checkpoint_id is None:
+            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)
+        parent_of: dict[str, str | None] = {}
+        ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
+        hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
+        chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
+        seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
+        walk_cursor_by_ch: dict[str, str | None] = {}
+        seeded: set[str] = set()
+        cursor: str | None = None
+
+        async with self._cursor() as cur:
+            while True:
+                stage1_params: list[Any] = []
+                for ch in channels:
+                    stage1_params.extend([ch, ch])
+                stage1_params.extend(
+                    [thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
+                )
+                await cur.execute(stage1_sql, stage1_params)
+                page = await cur.fetchall()
+                if not page:
+                    break
+                oldest = self._ingest_stage1_page(
+                    cast("list[Mapping[str, Any]]", page),
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hs_by_i_by_cid,
+                )
+                self._try_advance_walks(
+                    checkpoint_id,
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hs_by_i_by_cid,
+                    chain_by_ch,
+                    seed_ver_by_ch,
+                    walk_cursor_by_ch,
+                    seeded,
+                )
+                if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
+                    break
+                cursor = oldest
+
+        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]
+        stage2_sql = _build_delta_stage2_sql(
+            channels_with_chain=channels_with_chain,
+            channels_with_seed=channels_with_seed,
+        )
+
+        if stage2_sql:
+            stage2_params: list[Any] = []
+            for ch in channels_with_chain:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]])
+            for ch in channels_with_seed:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
+            async with self._cursor() as cur:
+                await cur.execute(stage2_sql, stage2_params)
+                stage2_rows = await cur.fetchall()
+        else:
+            stage2_rows = []
+
+        return self._build_delta_channels_writes_history(
+            channels=channels,
+            chain_by_ch=chain_by_ch,
+            seed_ver_by_ch=seed_ver_by_ch,
+            stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
+        )
+
+    async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
+        """
+        Convert a database row into a CheckpointTuple object.
+
+        Args:
+            value: A row from the database containing checkpoint data.
+
+        Returns:
+            CheckpointTuple: A structured representation of the checkpoint,
+            including its configuration, metadata, parent checkpoint (if any),
+            and pending writes.
+        """
+        return CheckpointTuple(
+            {
+                "configurable": {
+                    "thread_id": value["thread_id"],
+                    "checkpoint_ns": value["checkpoint_ns"],
+                    "checkpoint_id": value["checkpoint_id"],
+                }
+            },
+            {
+                **value["checkpoint"],
+                "channel_values": {
+                    **(value["checkpoint"].get("channel_values") or {}),
+                    **self._load_blobs(value["channel_values"]),
+                },
+            },
+            value["metadata"],
+            (
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["parent_checkpoint_id"],
+                    }
+                }
+                if value["parent_checkpoint_id"]
+                else None
+            ),
+            await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+        )
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config: Base configuration for filtering checkpoints.
+            filter: Additional filtering criteria for metadata.
+            before: If provided, only checkpoints before the specified checkpoint ID are returned.
+            limit: Maximum number of checkpoints to return.
+
+        Yields:
+            An iterator of matching checkpoint tuples.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            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 `checkpointer.alist(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_),  # type: ignore[arg-type]  # noqa: F821
+                    self.loop,
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a `checkpoint_id` key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            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_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+            task_path: Path of the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id, task_path), self.loop
+        ).result()
+
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id: The thread ID to delete.
+
+        Returns:
+            None
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            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_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.adelete_thread(thread_id), self.loop
+        ).result()
+
+
+__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -0,0 +1,596 @@
+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,
+)
+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] order by cw.task_id, cw.idx)
+        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
+    idx: int | None  # "w" rows only
+    version: str | None  # "b" rows only
+
+
+# Multi-channel two-stage DeltaChannel reconstruction.
+#
+# Stage 1 scans checkpoint metadata (no blob bytes) and emits one row per
+# checkpoint with K parallel JSONB key lookups (one column pair per
+# requested delta channel: ver_i / hs_i).  No subqueries, no aggregation.
+# Python walks the parent chain once across all channels.
+#
+# Stage 2 fetches all writes and the seed blobs for ALL channels in a
+# single roundtrip via `channel = ANY(%s)` and chain/seed-version
+# filtering.
+#
+# Empirical comparison vs an alternative "ship full channel_versions /
+# channel_values JSONB and let Python pick" form (1000 checkpoints,
+# 8 total channels in graph, 3 delta channels requested):
+#
+#   Postgres execution:    A=0.24ms vs B=0.38ms   (both negligible)
+#   End-to-end latency:    A=6.83ms vs B=2.28ms   (B is 3.0x faster)
+#   Wire payload:          A=836KB  vs B=330KB    (61% smaller)
+#   Buffer hits:           identical (167 blocks)
+#
+# B (this dynamic-columns design) wins because it avoids JSONB
+# serialization on the wire and JSONB-to-dict deserialization in
+# psycopg.  Even at K=8 (8 delta channels = 16 dynamic columns), B
+# still beats A end-to-end (4.2ms vs 6.8ms).
+
+
+def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
+    """Build stage 1 SQL with 2K parallel JSONB key lookups.
+
+    For channels=["messages", "files"] (with `paged=True`) the result is::
+
+        SELECT checkpoint_id, parent_checkpoint_id,
+               checkpoint -> 'channel_versions' ->> %s AS ver_0,
+               (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0,
+               checkpoint -> 'channel_versions' ->> %s AS ver_1,
+               (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1
+        FROM checkpoints
+        WHERE thread_id = %s AND checkpoint_ns = %s
+          AND (%s::text IS NULL OR checkpoint_id < %s)
+        ORDER BY checkpoint_id DESC
+        LIMIT %s
+
+    Channel names are passed as `%s` parameters (safe from SQL injection).
+    Only the column aliases `ver_i` / `hs_i` are interpolated into the
+    SQL string (i is bounded by len(channels) and uses safe identifiers).
+
+    Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ...,
+    thread_id, ns, cursor, cursor, page_size]` when `paged=True`.
+
+    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"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}"
+        )
+    sql = (
+        "SELECT checkpoint_id, parent_checkpoint_id, "
+        + ", ".join(cols)
+        + " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
+    )
+    if paged:
+        sql += (
+            " AND (%s::text IS NULL OR checkpoint_id < %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, 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::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 dynamic-shape dicts: {checkpoint_id, parent_checkpoint_id,
+# ver_0, hs_0, ver_1, hs_1, ...}.  Walking is parameterized by the channel
+# list to map indices back to channel names — no static TypedDict here.
+# `dict[str, Any]` is the practical signature.
+
+
+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]],
+        hs_by_i_by_cid: list[dict[str, bool]],
+    ) -> 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}"))
+                hs_by_i_by_cid[i][cid] = bool(r.get(f"hs_{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]],
+        hs_by_i_by_cid: Sequence[Mapping[str, bool]],
+        chain_by_ch: dict[str, list[str]],
+        seed_ver_by_ch: dict[str, str | None],
+        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 snapshot for its channel (channel becomes seeded),
+          (b) it reached a real root (parent_of[cid] is None — fully
+              materialized at this point), or
+          (c) the next ancestor cid isn't in `parent_of` yet (waiting for
+              a later page; the cursor stays put).
+
+        Mutates `chain_by_ch`, `seed_ver_by_ch`, `walk_cursor_by_ch`, and
+        `seeded` in place.
+        """
+        for i, ch in enumerate(channels):
+            if ch in seeded:
+                continue
+            # First-time entry: cursor starts at the target's parent.
+            if ch not in walk_cursor_by_ch:
+                walk_cursor_by_ch[ch] = parent_of.get(target_id)
+            cur_cid = walk_cursor_by_ch[ch]
+            ch_chain = chain_by_ch[ch]
+            hs_i = hs_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)
+                if hs_i.get(cur_cid, False):
+                    seed_ver_by_ch[ch] = ver_i.get(cur_cid)
+                    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],
+        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. The `seed` key is omitted
+        when the walk reached root with no snapshot found, 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)
+        writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
+            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]",
+                        (r["type"], r["blob"], r["task_id"], r["idx"]),
+                    )
+                )
+            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 by (task_id, idx)
+        for cid_map in writes_by_ch_by_cid.values():
+            for ws in cid_map.values():
+                ws.sort(key=lambda w: (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 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)
+            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]]
+    ) -> list[tuple[str, str, Any]]:
+        return (
+            [
+                (
+                    tid.decode(),
+                    channel.decode(),
+                    self.serde.loads_typed((t.decode(), v)),
+                )
+                for tid, channel, t, v in writes
+            ]
+            if writes
+            else []
+        )
+
+    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,
+        )

libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py

@@ -0,0 +1,967 @@
+import asyncio
+import threading
+import warnings
+from collections.abc import AsyncIterator, Iterator, Sequence
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_serializable_checkpoint_metadata,
+)
+from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.types import TASKS
+from psycopg import (
+    AsyncConnection,
+    AsyncCursor,
+    AsyncPipeline,
+    Capabilities,
+    Connection,
+    Cursor,
+    Pipeline,
+)
+from psycopg.rows import DictRow, dict_row
+from psycopg.types.json import Jsonb
+from psycopg_pool import AsyncConnectionPool, ConnectionPool
+
+from langgraph.checkpoint.postgres import _ainternal, _internal
+from langgraph.checkpoint.postgres.base import BasePostgresSaver
+
+"""
+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 '',
+    type TEXT,
+    checkpoint JSONB NOT NULL,
+    metadata JSONB NOT NULL DEFAULT '{}',
+    PRIMARY KEY (thread_id, checkpoint_ns)
+);""",
+    """CREATE TABLE IF NOT EXISTS checkpoint_blobs (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    channel TEXT NOT NULL,
+    type TEXT NOT NULL,
+    blob BYTEA,
+    PRIMARY KEY (thread_id, checkpoint_ns, channel)
+);""",
+    """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)
+);""",
+    """
+    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 = f"""
+select
+    thread_id,
+    checkpoint,
+    checkpoint_ns,
+    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
+    ) as channel_values,
+    (
+        select
+        array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.checkpoint_id = (checkpoint->>'id')
+    ) as pending_writes,
+    (
+        select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.channel = '{TASKS}'
+    ) as pending_sends
+from checkpoints """
+
+UPSERT_CHECKPOINT_BLOBS_SQL = """
+    INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, type, blob)
+    VALUES (%s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, channel) DO UPDATE SET
+        type = EXCLUDED.type,
+        blob = EXCLUDED.blob;
+"""
+
+UPSERT_CHECKPOINTS_SQL = """
+    INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint, metadata)
+    VALUES (%s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns)
+    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
+"""
+
+
+def _dump_blobs(
+    serde: SerializerProtocol,
+    thread_id: str,
+    checkpoint_ns: str,
+    values: dict[str, Any],
+    versions: ChannelVersions,
+) -> list[tuple[str, str, str, str, bytes | None]]:
+    if not versions:
+        return []
+
+    return [
+        (
+            thread_id,
+            checkpoint_ns,
+            k,
+            *(serde.dumps_typed(values[k]) if k in values else ("empty", None)),
+        )
+        for k in versions
+    ]
+
+
+class ShallowPostgresSaver(BasePostgresSaver):
+    """A checkpoint saver that uses Postgres to store checkpoints.
+
+    This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
+    It is meant to be a light-weight drop-in replacement for the PostgresSaver that
+    supports most of the LangGraph persistence functionality with the exception of time travel.
+    """
+
+    SELECT_SQL = SELECT_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
+
+    lock: threading.Lock
+
+    def __init__(
+        self,
+        conn: _internal.Conn,
+        pipe: Pipeline | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        warnings.warn(
+            "ShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
+            "Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., durability='exit')`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(serde=serde)
+        if isinstance(conn, ConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single Connection, not ConnectionPool."
+            )
+
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = threading.Lock()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls, conn_string: str, *, pipeline: bool = False
+    ) -> Iterator["ShallowPostgresSaver"]:
+        """Create a new ShallowPostgresSaver instance from a connection string.
+
+        Args:
+            conn_string: The Postgres connection info string.
+            pipeline: whether to use Pipeline
+
+        Returns:
+            ShallowPostgresSaver: A new ShallowPostgresSaver instance.
+        """
+        with Connection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                with conn.pipeline() as pipe:
+                    yield cls(conn, pipe)
+            else:
+                yield cls(conn)
+
+    def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        with self._cursor() as cur:
+            cur.execute(self.MIGRATIONS[0])
+            results = cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = results.fetchone()
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+                strict=False,
+            ):
+                cur.execute(migration)
+                cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,))
+        if self.pipe:
+            self.pipe.sync()
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. For ShallowPostgresSaver, this method returns a list with
+        ONLY the most recent checkpoint.
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where
+        params = list(args)
+        if limit is not None:
+            query += " LIMIT %s"
+            params.append(int(limit))
+        with self._cursor() as cur:
+            cur.execute(query, params, binary=True)
+            for value in cur:
+                checkpoint: Checkpoint = {
+                    **value["checkpoint"],
+                    "channel_values": self._load_blobs(value["channel_values"]),
+                    "pending_sends": [
+                        self.serde.loads_typed((t.decode(), v))
+                        for t, v in value["pending_sends"]
+                    ]
+                    if value["pending_sends"]
+                    else [],
+                }
+                yield CheckpointTuple(
+                    config={
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": checkpoint["id"],
+                        }
+                    },
+                    checkpoint=checkpoint,
+                    metadata=value["metadata"],
+                    pending_writes=self._load_writes(value["pending_writes"]),
+                )
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config (matching the thread ID in the config).
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+
+        Examples:
+
+            Basic:
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+
+            With timestamp:
+
+            >>> config = {
+            ...    "configurable": {
+            ...        "thread_id": "1",
+            ...        "checkpoint_ns": "",
+            ...        "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+            ...    }
+            ... }
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+        """  # noqa
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        args = (thread_id, checkpoint_ns)
+        where = "WHERE thread_id = %s AND checkpoint_ns = %s"
+
+        with self._cursor() as cur:
+            cur.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+
+            for value in cur:
+                checkpoint: Checkpoint = {
+                    **value["checkpoint"],
+                    "channel_values": self._load_blobs(value["channel_values"]),
+                    "pending_sends": [
+                        self.serde.loads_typed((t.decode(), v))
+                        for t, v in value["pending_sends"]
+                    ]
+                    if value["pending_sends"]
+                    else [],
+                }
+                return CheckpointTuple(
+                    config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint["id"],
+                        }
+                    },
+                    checkpoint=checkpoint,
+                    metadata=value["metadata"],
+                    pending_writes=self._load_writes(value["pending_writes"]),
+                )
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config. For ShallowPostgresSaver, this method saves ONLY the most recent
+        checkpoint and overwrites a previous checkpoint, if it exists.
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Examples:
+
+            >>> from langgraph.checkpoint.postgres import ShallowPostgresSaver
+            >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+            >>> with ShallowPostgresSaver.from_conn_string(DB_URI) as memory:
+            >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
+            >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
+            >>> print(saved_config)
+            {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        with self._cursor(pipeline=True) as cur:
+            cur.execute(
+                """DELETE FROM checkpoint_writes
+                WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    configurable.get("checkpoint_id", ""),
+                ),
+            )
+            cur.executemany(
+                self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                _dump_blobs(
+                    self.serde,
+                    thread_id,
+                    checkpoint_ns,
+                    copy.pop("channel_values"),  # type: ignore[misc]
+                    new_versions,
+                ),
+            )
+            cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    Jsonb(copy),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
+                ),
+            )
+        return next_config
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the Postgres database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store.
+            task_id: Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        with self._cursor(pipeline=True) as cur:
+            cur.executemany(
+                query,
+                self._dump_writes(
+                    config["configurable"]["thread_id"],
+                    config["configurable"]["checkpoint_ns"],
+                    config["configurable"]["checkpoint_id"],
+                    task_id,
+                    task_path,
+                    writes,
+                ),
+            )
+
+    @contextmanager
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the ShallowPostgresSaver instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        with _internal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                if self.supports_pipeline:
+                    with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+                else:
+                    # Use connection's transaction context manager when pipeline mode not supported
+                    with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+            else:
+                with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
+                    yield cur
+
+
+class AsyncShallowPostgresSaver(BasePostgresSaver):
+    """A checkpoint saver that uses Postgres to store checkpoints asynchronously.
+
+    This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
+    It is meant to be a light-weight drop-in replacement for the AsyncPostgresSaver that
+    supports most of the LangGraph persistence functionality with the exception of time travel.
+    """
+
+    SELECT_SQL = SELECT_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
+    lock: asyncio.Lock
+
+    def __init__(
+        self,
+        conn: _ainternal.Conn,
+        pipe: AsyncPipeline | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        warnings.warn(
+            "AsyncShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
+            "Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., durability='exit')`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(serde=serde)
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        serde: SerializerProtocol | None = None,
+    ) -> AsyncIterator["AsyncShallowPostgresSaver"]:
+        """Create a new AsyncShallowPostgresSaver instance from a connection string.
+
+        Args:
+            conn_string: The Postgres connection info string.
+            pipeline: whether to use AsyncPipeline
+
+        Returns:
+            AsyncShallowPostgresSaver: A new AsyncShallowPostgresSaver instance.
+        """
+        async with await AsyncConnection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                async with conn.pipeline() as pipe:
+                    yield cls(conn=conn, pipe=pipe, serde=serde)
+            else:
+                yield cls(conn=conn, serde=serde)
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        async with self._cursor() as cur:
+            await cur.execute(self.MIGRATIONS[0])
+            results = await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await results.fetchone()
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+                strict=False,
+            ):
+                await cur.execute(migration)
+                await cur.execute(
+                    "INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
+                )
+        if self.pipe:
+            await self.pipe.sync()
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        """List checkpoints from the database asynchronously.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. For ShallowPostgresSaver, this method returns a list with
+        ONLY the most recent checkpoint.
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where
+        params = list(args)
+        if limit is not None:
+            query += " LIMIT %s"
+            params.append(int(limit))
+        async with self._cursor() as cur:
+            await cur.execute(query, params, binary=True)
+            async for value in cur:
+                checkpoint: Checkpoint = {
+                    **value["checkpoint"],
+                    "channel_values": self._load_blobs(value["channel_values"]),
+                    "pending_sends": [
+                        self.serde.loads_typed((t.decode(), v))
+                        for t, v in value["pending_sends"]
+                    ]
+                    if value["pending_sends"]
+                    else [],
+                }
+                yield CheckpointTuple(
+                    config={
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": checkpoint["id"],
+                        }
+                    },
+                    checkpoint=checkpoint,
+                    metadata=value["metadata"],
+                    pending_writes=await asyncio.to_thread(
+                        self._load_writes, value["pending_writes"]
+                    ),
+                )
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database asynchronously.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config (matching the thread ID in the config).
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        args = (thread_id, checkpoint_ns)
+        where = "WHERE thread_id = %s AND checkpoint_ns = %s"
+
+        async with self._cursor() as cur:
+            await cur.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+
+            async for value in cur:
+                checkpoint: Checkpoint = {
+                    **value["checkpoint"],
+                    "channel_values": self._load_blobs(value["channel_values"]),
+                    "pending_sends": [
+                        self.serde.loads_typed((t.decode(), v))
+                        for t, v in value["pending_sends"]
+                    ]
+                    if value["pending_sends"]
+                    else [],
+                }
+                return CheckpointTuple(
+                    config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint["id"],
+                        }
+                    },
+                    checkpoint=checkpoint,
+                    metadata=value["metadata"],
+                    pending_writes=await asyncio.to_thread(
+                        self._load_writes, value["pending_writes"]
+                    ),
+                )
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database asynchronously.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
+        checkpoint and overwrites a previous checkpoint, if it exists.
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        async with self._cursor(pipeline=True) as cur:
+            await cur.execute(
+                """DELETE FROM checkpoint_writes
+                WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    configurable.get("checkpoint_id", ""),
+                ),
+            )
+            await cur.executemany(
+                self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                _dump_blobs(
+                    self.serde,
+                    thread_id,
+                    checkpoint_ns,
+                    copy.pop("channel_values"),  # type: ignore[misc]
+                    new_versions,
+                ),
+            )
+            await cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    Jsonb(copy),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
+                ),
+            )
+        return next_config
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            task_path,
+            writes,
+        )
+        async with self._cursor(pipeline=True) as cur:
+            await cur.executemany(query, params)
+
+    @asynccontextmanager
+    async def _cursor(
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        async with _ainternal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        await self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                if self.supports_pipeline:
+                    async with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+                else:
+                    # Use connection's transaction context manager when pipeline mode not supported
+                    async with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+            else:
+                async with (
+                    self.lock,
+                    conn.cursor(binary=True, row_factory=dict_row) as cur,
+                ):
+                    yield cur
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. For ShallowPostgresSaver, this method returns a list with
+        ONLY the most recent checkpoint.
+        """
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_),  # type: ignore[arg-type]  # noqa: F821
+                    self.loop,
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config (matching the thread ID in the config).
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncShallowPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface."
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
+        checkpoint and overwrites a previous checkpoint, if it exists.
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+            task_path: Path of the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id, task_path), self.loop
+        ).result()