langgraph-checkpoint-postgres 2.0.2__tar.gz → 2.0.4__tar.gz

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 2.0.2
+Version: 2.0.4
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/langgraph/checkpoint/postgres/__init__.py

@@ -1,9 +1,9 @@
 import threading
 from contextlib import contextmanager
-from typing import Any, Iterator, Optional, Sequence, Union
+from typing import Any, Iterator, Optional, Sequence
 
 from langchain_core.runnables import RunnableConfig
-from psycopg import Connection, Cursor, Pipeline
+from psycopg import Capabilities, Connection, Cursor, Pipeline
 from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
@@ -17,21 +17,11 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
     get_checkpoint_id,
 )
+from langgraph.checkpoint.postgres import _internal
 from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
-Conn = Union[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)}")
+Conn = _internal.Conn  # For backward compatibility
 
 
 class PostgresSaver(BasePostgresSaver):
@@ -39,7 +29,7 @@ class PostgresSaver(BasePostgresSaver):
 
     def __init__(
         self,
-        conn: Conn,
+        conn: _internal.Conn,
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -52,6 +42,7 @@ class PostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = threading.Lock()
+        self.supports_pipeline = Capabilities().has_pipeline()
 
     @classmethod
     @contextmanager
@@ -72,9 +63,9 @@ class PostgresSaver(BasePostgresSaver):
         ) as conn:
             if pipeline:
                 with conn.pipeline() as pipe:
-                    yield PostgresSaver(conn, pipe)
+                    yield cls(conn, pipe)
             else:
-                yield PostgresSaver(conn)
+                yield cls(conn)
 
     def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
@@ -365,7 +356,14 @@ class PostgresSaver(BasePostgresSaver):
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
-        with _get_connection(self.conn) as conn:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline (bool): 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 _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
@@ -379,13 +377,20 @@ class PostgresSaver(BasePostgresSaver):
             elif pipeline:
                 # a connection not in pipeline mode can only be used by one
                 # thread/coroutine at a time, so we acquire a lock
-                with self.lock, conn.pipeline(), conn.cursor(
-                    binary=True, row_factory=dict_row
-                ) as cur:
-                    yield cur
+                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
 
 
-__all__ = ["PostgresSaver", "Conn"]
+__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]

langgraph_checkpoint_postgres-2.0.4/langgraph/checkpoint/postgres/_ainternal.py

@@ -0,0 +1,23 @@
+"""Shared async utility functions for the Postgres checkpoint & storage classes."""
+
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Union
+
+from psycopg import AsyncConnection
+from psycopg.rows import DictRow
+from psycopg_pool import AsyncConnectionPool
+
+Conn = Union[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)}")

langgraph_checkpoint_postgres-2.0.4/langgraph/checkpoint/postgres/_internal.py

@@ -0,0 +1,21 @@
+"""Shared utility functions for the Postgres checkpoint & storage classes."""
+
+from contextlib import contextmanager
+from typing import Iterator, Union
+
+from psycopg import Connection
+from psycopg.rows import DictRow
+from psycopg_pool import ConnectionPool
+
+Conn = Union[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)}")

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/langgraph/checkpoint/postgres/aio.py

@@ -1,9 +1,9 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Union
+from typing import Any, AsyncIterator, Iterator, Optional, Sequence
 
 from langchain_core.runnables import RunnableConfig
-from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
+from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
 from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
@@ -17,23 +17,11 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
     get_checkpoint_id,
 )
+from langgraph.checkpoint.postgres import _ainternal
 from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
-Conn = Union[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)}")
+Conn = _ainternal.Conn  # For backward compatibility
 
 
 class AsyncPostgresSaver(BasePostgresSaver):
@@ -41,7 +29,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
 
     def __init__(
         self,
-        conn: Conn,
+        conn: _ainternal.Conn,
         pipe: Optional[AsyncPipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -55,6 +43,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         self.pipe = pipe
         self.lock = asyncio.Lock()
         self.loop = asyncio.get_running_loop()
+        self.supports_pipeline = Capabilities().has_pipeline()
 
     @classmethod
     @asynccontextmanager
@@ -79,9 +68,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
         ) as conn:
             if pipeline:
                 async with conn.pipeline() as pipe:
-                    yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde)
+                    yield cls(conn=conn, pipe=pipe, serde=serde)
             else:
-                yield AsyncPostgresSaver(conn=conn, serde=serde)
+                yield cls(conn=conn, serde=serde)
 
     async def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
@@ -156,15 +145,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
                         value["pending_sends"],
                     ),
                     self._load_metadata(value["metadata"]),
-                    {
-                        "configurable": {
-                            "thread_id": value["thread_id"],
-                            "checkpoint_ns": value["checkpoint_ns"],
-                            "checkpoint_id": value["parent_checkpoint_id"],
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": value["thread_id"],
+                                "checkpoint_ns": value["checkpoint_ns"],
+                                "checkpoint_id": value["parent_checkpoint_id"],
+                            }
                         }
-                    }
-                    if value["parent_checkpoint_id"]
-                    else None,
+                        if value["parent_checkpoint_id"]
+                        else None
+                    ),
                     await asyncio.to_thread(self._load_writes, value["pending_writes"]),
                 )
 
@@ -215,15 +206,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
                         value["pending_sends"],
                     ),
                     self._load_metadata(value["metadata"]),
-                    {
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": value["parent_checkpoint_id"],
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": value["parent_checkpoint_id"],
+                            }
                         }
-                    }
-                    if value["parent_checkpoint_id"]
-                    else None,
+                        if value["parent_checkpoint_id"]
+                        else None
+                    ),
                     await asyncio.to_thread(self._load_writes, value["pending_writes"]),
                 )
 
@@ -323,7 +316,14 @@ class AsyncPostgresSaver(BasePostgresSaver):
     async def _cursor(
         self, *, pipeline: bool = False
     ) -> AsyncIterator[AsyncCursor[DictRow]]:
-        async with _get_connection(self.conn) as conn:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline (bool): 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 _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
@@ -337,10 +337,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
             elif pipeline:
                 # a connection not in pipeline mode can only be used by one
                 # thread/coroutine at a time, so we acquire a lock
-                async with self.lock, conn.pipeline(), conn.cursor(
-                    binary=True, row_factory=dict_row
-                ) as cur:
-                    yield cur
+                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
@@ -452,3 +459,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
         return asyncio.run_coroutine_threadsafe(
             self.aput_writes(config, writes, task_id), self.loop
         ).result()
+
+
+__all__ = ["AsyncPostgresSaver", "Conn"]

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/langgraph/checkpoint/postgres/base.py

@@ -84,7 +84,7 @@ select
             and cw.checkpoint_id = checkpoints.checkpoint_id
     ) as pending_writes,
     (
-        select array_agg(array[cw.type::bytea, cw.blob] order by cw.idx)
+        select array_agg(array[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
@@ -133,6 +133,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
     INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
+    supports_pipeline: bool
 
     def _load_checkpoint(
         self,

langgraph_checkpoint_postgres-2.0.4/langgraph/store/postgres/aio.py

@@ -0,0 +1,294 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Iterable,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+)
+
+import orjson
+from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
+from psycopg.errors import UndefinedTable
+from psycopg.rows import DictRow, dict_row
+from psycopg_pool import AsyncConnectionPool
+
+from langgraph.checkpoint.postgres import _ainternal
+from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
+from langgraph.store.base.batch import AsyncBatchedBaseStore
+from langgraph.store.postgres.base import (
+    BasePostgresStore,
+    PoolConfig,
+    Row,
+    _decode_ns_bytes,
+    _group_ops,
+    _row_to_item,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
+    __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
+
+    def __init__(
+        self,
+        conn: _ainternal.Conn,
+        *,
+        pipe: Optional[AsyncPipeline] = None,
+        deserializer: Optional[
+            Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
+        ] = None,
+    ) -> None:
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+        super().__init__()
+        self._deserializer = deserializer
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        grouped_ops, num_ops = _group_ops(ops)
+        results: list[Result] = [None] * num_ops
+
+        async with _ainternal.get_connection(self.conn) as conn:
+            if self.pipe:
+                async with self.pipe:
+                    await self._execute_batch(grouped_ops, results, conn)
+            else:
+                await self._execute_batch(grouped_ops, results, conn)
+
+        return results
+
+    async def _execute_batch(
+        self,
+        grouped_ops: dict,
+        results: list[Result],
+        conn: AsyncConnection[DictRow],
+    ) -> None:
+        async with self._cursor(conn, pipeline=True) as cur:
+            if GetOp in grouped_ops:
+                await self._batch_get_ops(
+                    cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),
+                    results,
+                    cur,
+                )
+
+            if SearchOp in grouped_ops:
+                await self._batch_search_ops(
+                    cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
+                    results,
+                    cur,
+                )
+
+            if ListNamespacesOp in grouped_ops:
+                await self._batch_list_namespaces_ops(
+                    cast(
+                        Sequence[tuple[int, ListNamespacesOp]],
+                        grouped_ops[ListNamespacesOp],
+                    ),
+                    results,
+                    cur,
+                )
+
+            if PutOp in grouped_ops:
+                await self._batch_put_ops(
+                    cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]),
+                    cur,
+                )
+
+    async def _batch_get_ops(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+        results: list[Result],
+        cur: AsyncCursor[DictRow],
+    ) -> None:
+        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
+            await cur.execute(query, params)
+            rows = cast(list[Row], await cur.fetchall())
+            key_to_row = {row["key"]: row for row in rows}
+            for idx, key in items:
+                row = key_to_row.get(key)
+                if row:
+                    results[idx] = _row_to_item(
+                        namespace, row, loader=self._deserializer
+                    )
+                else:
+                    results[idx] = None
+
+    async def _batch_put_ops(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+        cur: AsyncCursor[DictRow],
+    ) -> None:
+        queries = self._get_batch_PUT_queries(put_ops)
+        for query, params in queries:
+            await cur.execute(query, params)
+
+    async def _batch_search_ops(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+        results: list[Result],
+        cur: AsyncCursor[DictRow],
+    ) -> None:
+        queries = self._get_batch_search_queries(search_ops)
+        for (query, params), (idx, _) in zip(queries, search_ops):
+            await cur.execute(query, params)
+            rows = cast(list[Row], await cur.fetchall())
+            items = [
+                _row_to_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                )
+                for row in rows
+            ]
+            results[idx] = items
+
+    async def _batch_list_namespaces_ops(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+        results: list[Result],
+        cur: AsyncCursor[DictRow],
+    ) -> None:
+        queries = self._get_batch_list_namespaces_queries(list_ops)
+        for (query, params), (idx, _) in zip(queries, list_ops):
+            await cur.execute(query, params)
+            rows = cast(list[dict], await cur.fetchall())
+            namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
+            results[idx] = namespaces
+
+    @asynccontextmanager
+    async def _cursor(
+        self, conn: AsyncConnection[DictRow], *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[Any]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            conn: The database connection to use
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        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
+            async with conn.cursor(binary=True) as cur:
+                try:
+                    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) as cur:
+                    yield cur
+            else:
+                async with self.lock, conn.transaction(), conn.cursor(
+                    binary=True
+                ) as cur:
+                    yield cur
+        else:
+            async with conn.cursor(binary=True) as cur:
+                yield cur
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        pool_config: Optional[PoolConfig] = None,
+    ) -> AsyncIterator["AsyncPostgresStore"]:
+        """Create a new AsyncPostgresStore instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+            pipeline (bool): Whether to use AsyncPipeline (only for single connections)
+            pool_config (Optional[PoolConfig]): Configuration for the connection pool.
+                If provided, will create a connection pool and use it instead of a single connection.
+                This overrides the `pipeline` argument.
+
+        Returns:
+            AsyncPostgresStore: A new AsyncPostgresStore instance.
+        """
+        if pool_config is not None:
+            pc = pool_config.copy()
+            async with cast(
+                AsyncConnectionPool[AsyncConnection[DictRow]],
+                AsyncConnectionPool(
+                    conn_string,
+                    min_size=pc.pop("min_size", 1),
+                    max_size=pc.pop("max_size", None),
+                    kwargs={
+                        "autocommit": True,
+                        "prepare_threshold": 0,
+                        "row_factory": dict_row,
+                        **(pc.pop("kwargs", None) or {}),
+                    },
+                    **cast(dict, pc),
+                ),
+            ) as pool:
+                yield cls(conn=pool)
+        else:
+            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)
+                else:
+                    yield cls(conn=conn)
+
+    async def setup(self) -> None:
+        """Set up the store 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 the store is used.
+        """
+        async with _ainternal.get_connection(self.conn) as conn:
+            async with conn.cursor() as cur:
+                try:
+                    await cur.execute(
+                        "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
+                    )
+                    row = cast(dict, await cur.fetchone())
+                    if row is None:
+                        version = -1
+                    else:
+                        version = row["v"]
+                except UndefinedTable:
+                    version = -1
+                    # Create store_migrations table if it doesn't exist
+                    await cur.execute(
+                        """
+                        CREATE TABLE IF NOT EXISTS store_migrations (
+                            v INTEGER PRIMARY KEY
+                        )
+                        """
+                    )
+                for v, migration in enumerate(
+                    self.MIGRATIONS[version + 1 :], start=version + 1
+                ):
+                    await cur.execute(migration)
+                    await cur.execute(
+                        "INSERT INTO store_migrations (v) VALUES (%s)", (v,)
+                    )
+            if self.pipe:
+                await self.pipe.sync()

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/langgraph/store/postgres/base.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import threading
 from collections import defaultdict
 from contextlib import contextmanager
 from datetime import datetime
@@ -18,12 +19,15 @@ from typing import (
 )
 
 import orjson
-from psycopg import BaseConnection, Connection, Cursor
+from psycopg import Capabilities, Connection, Cursor, Pipeline
 from psycopg.errors import UndefinedTable
-from psycopg.rows import dict_row
+from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
+from psycopg_pool import ConnectionPool
 from typing_extensions import TypedDict
 
+from langgraph.checkpoint.postgres import _ainternal as _ainternal
+from langgraph.checkpoint.postgres import _internal as _pg_internal
 from langgraph.store.base import (
     BaseStore,
     GetOp,
@@ -56,7 +60,32 @@ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pa
 """,
 ]
 
-C = TypeVar("C", bound=BaseConnection)
+C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
+
+
+class PoolConfig(TypedDict, total=False):
+    """Connection pool settings for PostgreSQL connections.
+
+    Controls connection lifecycle and resource utilization:
+    - Small pools (1-5) suit low-concurrency workloads
+    - Larger pools handle concurrent requests but consume more resources
+    - Setting max_size prevents resource exhaustion under load
+    """
+
+    min_size: int
+    """Minimum number of connections maintained in the pool. Defaults to 1."""
+
+    max_size: Optional[int]
+    """Maximum number of connections allowed in the pool. None means unlimited."""
+
+    kwargs: dict
+    """Additional connection arguments passed to each connection in the pool.
+    
+    Default kwargs set automatically:
+    - autocommit: True
+    - prepare_threshold: 0
+    - row_factory: dict_row
+    """
 
 
 class BasePostgresStore(Generic[C]):
@@ -88,9 +117,14 @@ class BasePostgresStore(Generic[C]):
         self,
         put_ops: Sequence[tuple[int, PutOp]],
     ) -> list[tuple[str, Sequence]]:
+        # Last-write wins
+        dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
+        for _, op in put_ops:
+            dedupped_ops[(op.namespace, op.key)] = op
+
         inserts: list[PutOp] = []
         deletes: list[PutOp] = []
-        for _, op in put_ops:
+        for op in dedupped_ops.values():
             if op.value is None:
                 deletes.append(op)
             else:
@@ -219,13 +253,14 @@ class BasePostgresStore(Generic[C]):
         return queries
 
 
-class PostgresStore(BaseStore, BasePostgresStore[Connection]):
-    __slots__ = ("_deserializer",)
+class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
+    __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
 
     def __init__(
         self,
-        conn: Connection[Any],
+        conn: _pg_internal.Conn,
         *,
+        pipe: Optional[Pipeline] = None,
         deserializer: Optional[
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
@@ -233,26 +268,110 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
         super().__init__()
         self._deserializer = deserializer
         self.conn = conn
+        self.pipe = pipe
+        self.supports_pipeline = Capabilities().has_pipeline()
+        self.lock = threading.Lock()
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        pool_config: Optional[PoolConfig] = None,
+    ) -> Iterator["PostgresStore"]:
+        """Create a new PostgresStore instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+            pipeline (bool): whether to use Pipeline (only for single connections)
+            pool_config (Optional[PoolArgs]): Configuration for the connection pool.
+                If provided, will create a connection pool and use it instead of a single connection.
+                This overrides the `pipeline` argument.
+        Returns:
+            PostgresStore: A new PostgresStore instance.
+        """
+        if pool_config is not None:
+            pc = pool_config.copy()
+            with cast(
+                ConnectionPool[Connection[DictRow]],
+                ConnectionPool(
+                    conn_string,
+                    min_size=pc.pop("min_size", 1),
+                    max_size=pc.pop("max_size", None),
+                    kwargs={
+                        "autocommit": True,
+                        "prepare_threshold": 0,
+                        "row_factory": dict_row,
+                        **(pc.pop("kwargs", None) or {}),
+                    },
+                    **cast(dict, pc),
+                ),
+            ) as pool:
+                yield cls(conn=pool)
+        else:
+            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=pipe)
+                else:
+                    yield cls(conn)
+
+    @contextmanager
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline (bool): whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        with _pg_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:
+                    with self.lock, 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 batch(self, ops: Iterable[Op]) -> list[Result]:
         grouped_ops, num_ops = _group_ops(ops)
         results: list[Result] = [None] * num_ops
 
-        with self.conn.pipeline():
+        with self._cursor(pipeline=True) as cur:
             if GetOp in grouped_ops:
                 self._batch_get_ops(
-                    cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
-                )
-
-            if PutOp in grouped_ops:
-                self._batch_put_ops(
-                    cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
+                    cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur
                 )
 
             if SearchOp in grouped_ops:
                 self._batch_search_ops(
                     cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
                     results,
+                    cur,
                 )
 
             if ListNamespacesOp in grouped_ops:
@@ -262,25 +381,23 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
                         grouped_ops[ListNamespacesOp],
                     ),
                     results,
+                    cur,
+                )
+            if PutOp in grouped_ops:
+                self._batch_put_ops(
+                    cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), cur
                 )
 
         return results
 
-    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
-        return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
-
     def _batch_get_ops(
         self,
         get_ops: Sequence[tuple[int, GetOp]],
         results: list[Result],
+        cur: Cursor[DictRow],
     ) -> None:
-        cursors = []
         for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
-            cur = self.conn.cursor(binary=True)
             cur.execute(query, params)
-            cursors.append((cur, namespace, items))
-
-        for cur, namespace, items in cursors:
             rows = cast(list[Row], cur.fetchall())
             key_to_row = {row["key"]: row for row in rows}
             for idx, key in items:
@@ -295,70 +412,44 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
     def _batch_put_ops(
         self,
         put_ops: Sequence[tuple[int, PutOp]],
+        cur: Cursor[DictRow],
     ) -> None:
         queries = self._get_batch_PUT_queries(put_ops)
         for query, params in queries:
-            cur = self.conn.cursor(binary=True)
             cur.execute(query, params)
 
     def _batch_search_ops(
         self,
         search_ops: Sequence[tuple[int, SearchOp]],
         results: list[Result],
+        cur: Cursor[DictRow],
     ) -> None:
-        queries = self._get_batch_search_queries(search_ops)
-        cursors: list[tuple[Cursor[Any], int]] = []
-
-        for (query, params), (idx, _) in zip(queries, search_ops):
-            cur = self.conn.cursor(binary=True)
+        for (query, params), (idx, _) in zip(
+            self._get_batch_search_queries(search_ops), search_ops
+        ):
             cur.execute(query, params)
-            cursors.append((cur, idx))
-
-        for cur, idx in cursors:
             rows = cast(list[Row], cur.fetchall())
-            items = [
+            results[idx] = [
                 _row_to_item(
                     _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
                 )
                 for row in rows
             ]
-            results[idx] = items
 
     def _batch_list_namespaces_ops(
         self,
         list_ops: Sequence[tuple[int, ListNamespacesOp]],
         results: list[Result],
+        cur: Cursor[DictRow],
     ) -> None:
-        queries = self._get_batch_list_namespaces_queries(list_ops)
-        cursors: list[tuple[Cursor[Any], int]] = []
-        for (query, params), (idx, _) in zip(queries, list_ops):
-            cur = self.conn.cursor(binary=True)
+        for (query, params), (idx, _) in zip(
+            self._get_batch_list_namespaces_queries(list_ops), list_ops
+        ):
             cur.execute(query, params)
-            cursors.append((cur, idx))
+            results[idx] = [_decode_ns_bytes(row["truncated_prefix"]) for row in cur]
 
-        for cur, idx in cursors:
-            rows = cast(list[dict], cur.fetchall())
-            namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
-            results[idx] = namespaces
-
-    @classmethod
-    @contextmanager
-    def from_conn_string(
-        cls,
-        conn_string: str,
-    ) -> Iterator["PostgresStore"]:
-        """Create a new BasePostgresStore instance from a connection string.
-
-        Args:
-            conn_string (str): The Postgres connection info string.
-
-        Returns:
-            BasePostgresStore: A new BasePostgresStore instance.
-        """
-        with Connection.connect(
-            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
-        ) as conn:
-            yield cls(conn=conn)
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
 
     def setup(self) -> None:
         """Set up the store database.
@@ -367,7 +458,7 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time the store is used.
         """
-        with self.conn.cursor(binary=True) as cur:
+        with self._cursor() as cur:
             try:
                 cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
                 row = cast(dict, cur.fetchone())
@@ -376,9 +467,7 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
                 else:
                     version = row["v"]
             except UndefinedTable:
-                self.conn.rollback()
                 version = -1
-                # Create store_migrations table if it doesn't exist
                 cur.execute(
                     """
                     CREATE TABLE IF NOT EXISTS store_migrations (

{langgraph_checkpoint_postgres-2.0.2 → langgraph_checkpoint_postgres-2.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.2"
+version = "2.0.4"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"

langgraph_checkpoint_postgres-2.0.2/langgraph/store/postgres/aio.py

@@ -1,218 +0,0 @@
-import asyncio
-import logging
-from contextlib import asynccontextmanager
-from typing import (
-    Any,
-    AsyncIterator,
-    Callable,
-    Iterable,
-    Optional,
-    Sequence,
-    Union,
-    cast,
-)
-
-import orjson
-from psycopg import AsyncConnection, AsyncCursor
-from psycopg.errors import UndefinedTable
-from psycopg.rows import dict_row
-
-from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
-from langgraph.store.base.batch import AsyncBatchedBaseStore
-from langgraph.store.postgres.base import (
-    BasePostgresStore,
-    Row,
-    _decode_ns_bytes,
-    _group_ops,
-    _row_to_item,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]):
-    __slots__ = ("_deserializer",)
-
-    def __init__(
-        self,
-        conn: AsyncConnection[Any],
-        *,
-        deserializer: Optional[
-            Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
-        ] = None,
-    ) -> None:
-        super().__init__()
-        self._deserializer = deserializer
-        self.conn = conn
-        self.loop = asyncio.get_running_loop()
-
-    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
-        grouped_ops, num_ops = _group_ops(ops)
-        results: list[Result] = [None] * num_ops
-
-        async with self.conn.pipeline():
-            tasks = []
-
-            if GetOp in grouped_ops:
-                tasks.append(
-                    self._batch_get_ops(
-                        cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
-                    )
-                )
-
-            if PutOp in grouped_ops:
-                tasks.append(
-                    self._batch_put_ops(
-                        cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
-                    )
-                )
-
-            if SearchOp in grouped_ops:
-                tasks.append(
-                    self._batch_search_ops(
-                        cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
-                        results,
-                    )
-                )
-
-            if ListNamespacesOp in grouped_ops:
-                tasks.append(
-                    self._batch_list_namespaces_ops(
-                        cast(
-                            Sequence[tuple[int, ListNamespacesOp]],
-                            grouped_ops[ListNamespacesOp],
-                        ),
-                        results,
-                    )
-                )
-
-            await asyncio.gather(*tasks)
-
-        return results
-
-    def batch(self, ops: Iterable[Op]) -> list[Result]:
-        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
-
-    async def _batch_get_ops(
-        self,
-        get_ops: Sequence[tuple[int, GetOp]],
-        results: list[Result],
-    ) -> None:
-        cursors = []
-        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
-            cur = self.conn.cursor(binary=True)
-            await cur.execute(query, params)
-            cursors.append((cur, namespace, items))
-
-        for cur, namespace, items in cursors:
-            rows = cast(list[Row], await cur.fetchall())
-            key_to_row = {row["key"]: row for row in rows}
-            for idx, key in items:
-                row = key_to_row.get(key)
-                if row:
-                    results[idx] = _row_to_item(
-                        namespace, row, loader=self._deserializer
-                    )
-                else:
-                    results[idx] = None
-
-    async def _batch_put_ops(
-        self,
-        put_ops: Sequence[tuple[int, PutOp]],
-    ) -> None:
-        queries = self._get_batch_PUT_queries(put_ops)
-        for query, params in queries:
-            cur = self.conn.cursor(binary=True)
-            await cur.execute(query, params)
-
-    async def _batch_search_ops(
-        self,
-        search_ops: Sequence[tuple[int, SearchOp]],
-        results: list[Result],
-    ) -> None:
-        queries = self._get_batch_search_queries(search_ops)
-        cursors: list[tuple[AsyncCursor[Any], int]] = []
-
-        for (query, params), (idx, _) in zip(queries, search_ops):
-            cur = self.conn.cursor(binary=True)
-            await cur.execute(query, params)
-            cursors.append((cur, idx))
-
-        for cur, idx in cursors:
-            rows = cast(list[Row], await cur.fetchall())
-            items = [
-                _row_to_item(
-                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
-                )
-                for row in rows
-            ]
-            results[idx] = items
-
-    async def _batch_list_namespaces_ops(
-        self,
-        list_ops: Sequence[tuple[int, ListNamespacesOp]],
-        results: list[Result],
-    ) -> None:
-        queries = self._get_batch_list_namespaces_queries(list_ops)
-        cursors: list[tuple[AsyncCursor[Any], int]] = []
-        for (query, params), (idx, _) in zip(queries, list_ops):
-            cur = self.conn.cursor(binary=True)
-            await cur.execute(query, params)
-            cursors.append((cur, idx))
-
-        for cur, idx in cursors:
-            rows = cast(list[dict], await cur.fetchall())
-            namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
-            results[idx] = namespaces
-
-    @classmethod
-    @asynccontextmanager
-    async def from_conn_string(
-        cls,
-        conn_string: str,
-    ) -> AsyncIterator["AsyncPostgresStore"]:
-        """Create a new AsyncPostgresStore instance from a connection string.
-
-        Args:
-            conn_string (str): The Postgres connection info string.
-
-        Returns:
-            AsyncPostgresStore: A new AsyncPostgresStore instance.
-        """
-        async with await AsyncConnection.connect(
-            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
-        ) as conn:
-            yield cls(conn=conn)
-
-    async def setup(self) -> None:
-        """Set up the store 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 the store is used.
-        """
-        async with self.conn.cursor() as cur:
-            try:
-                await cur.execute(
-                    "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
-                )
-                row = cast(dict, await cur.fetchone())
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                # Create store_migrations table if it doesn't exist
-                await cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS store_migrations (
-                        v INTEGER PRIMARY KEY
-                    )
-                    """
-                )
-            for v, migration in enumerate(
-                self.MIGRATIONS[version + 1 :], start=version + 1
-            ):
-                await cur.execute(migration)
-                await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))