merge

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

@@ -1,6 +1,7 @@
 import threading
+from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
-from typing import Any, Iterator, Optional, Sequence
+from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import Capabilities, Connection, Cursor, Pipeline
@@ -21,6 +22,8 @@ from langgraph.checkpoint.postgres import _internal
 from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
+Conn = _internal.Conn  # For backward compatibility
+
 
 class PostgresSaver(BasePostgresSaver):
     lock: threading.Lock
@@ -61,9 +64,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.
@@ -376,19 +379,23 @@ class PostgresSaver(BasePostgresSaver):
                 # 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:
+                    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:
+                    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"]

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

@@ -1,7 +1,8 @@
 """Shared async utility functions for the Postgres checkpoint & storage classes."""
 
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, Union
+from typing import Union
 
 from psycopg import AsyncConnection
 from psycopg.rows import DictRow

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

@@ -1,7 +1,8 @@
 """Shared utility functions for the Postgres checkpoint & storage classes."""
 
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Iterator, Union
+from typing import Union
 
 from psycopg import Connection
 from psycopg.rows import DictRow

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

@@ -1,6 +1,7 @@
 import asyncio
+from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Iterator, Optional, Sequence
+from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -21,6 +22,8 @@ from langgraph.checkpoint.postgres import _ainternal
 from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
+Conn = _ainternal.Conn  # For backward compatibility
+
 
 class AsyncPostgresSaver(BasePostgresSaver):
     lock: asyncio.Lock
@@ -66,9 +69,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.
@@ -143,15 +146,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"]),
                 )
 
@@ -202,15 +207,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"]),
                 )
 
@@ -332,20 +339,25 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 # 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:
+                    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:
+                    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:
+                async with (
+                    self.lock,
+                    conn.cursor(binary=True, row_factory=dict_row) as cur,
+                ):
                     yield cur
 
     def list(
@@ -374,7 +386,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         while True:
             try:
                 yield asyncio.run_coroutine_threadsafe(
-                    anext(aiter_),
+                    anext(aiter_),  # noqa: F821
                     self.loop,
                 ).result()
             except StopAsyncIteration:
@@ -453,3 +465,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
         return asyncio.run_coroutine_threadsafe(
             self.aput_writes(config, writes, task_id), self.loop
         ).result()
+
+
+__all__ = ["AsyncPostgresSaver", "Conn"]

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

@@ -1,5 +1,6 @@
 import random
-from typing import Any, List, Optional, Sequence, Tuple, cast
+from collections.abc import Sequence
+from typing import Any, Optional, cast
 
 from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
@@ -249,7 +250,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         config: Optional[RunnableConfig],
         filter: MetadataInput,
         before: Optional[RunnableConfig] = None,
-    ) -> Tuple[str, List[Any]]:
+    ) -> 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

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

@@ -1,13 +1,11 @@
 import asyncio
 import logging
+from collections.abc import AsyncIterator, Iterable, Sequence
 from contextlib import asynccontextmanager
 from typing import (
     Any,
-    AsyncIterator,
     Callable,
-    Iterable,
     Optional,
-    Sequence,
     Union,
     cast,
 )
@@ -32,6 +30,7 @@ from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
     BasePostgresStore,
     EmbeddingConfig,
+    PoolConfig,
     Row,
     _decode_ns_bytes,
     _group_ops,
@@ -241,14 +240,18 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                 # 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:
+                    async with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
                 else:
-                    async with self.lock, conn.transaction(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    async with (
+                        self.lock,
+                        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:
@@ -264,9 +267,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         conn_string: str,
         *,
         pipeline: bool = False,
-        min_size: int = 1,
-        max_size: Optional[int] = None,
-        use_pool: bool = False,
+        pool_config: Optional[PoolConfig] = None,
         embedding: Optional[EmbeddingConfig] = None,
     ) -> AsyncIterator["AsyncPostgresStore"]:
         """Create a new AsyncPostgresStore instance from a connection string.
@@ -274,26 +275,29 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         Args:
             conn_string (str): The Postgres connection info string.
             pipeline (bool): Whether to use AsyncPipeline (only for single connections)
-            min_size (int): Minimum number of connections when using a pool
-            max_size (Optional[int]): Maximum number of connections when using a pool
-            use_pool (bool): Whether to use a connection pool
-            embedding (Optional[EmbeddingConfig]): Configuration for vector embeddings
+            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.
+            embedding (Optional[EmbeddingConfig]): The embedding config.
 
         Returns:
             AsyncPostgresStore: A new AsyncPostgresStore instance.
         """
-        if use_pool:
+        if pool_config is not None:
+            pc = pool_config.copy()
             async with cast(
                 AsyncConnectionPool[AsyncConnection[DictRow]],
                 AsyncConnectionPool(
                     conn_string,
-                    min_size=min_size,
-                    max_size=max_size,
+                    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, embedding=embedding)

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

@@ -3,20 +3,15 @@ import json
 import logging
 import threading
 from collections import defaultdict
+from collections.abc import Awaitable, Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import datetime
 from typing import (
     Any,
-    Awaitable,
     Callable,
     Generic,
-    Iterable,
-    Iterator,
-    List,
     Optional,
-    Sequence,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -28,6 +23,7 @@ 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
+from psycopg_pool import ConnectionPool
 from typing_extensions import TypedDict
 
 from langgraph.checkpoint.postgres import _ainternal as _ainternal
@@ -177,6 +173,31 @@ CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
 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]):
     MIGRATIONS = MIGRATIONS
     conn: C
@@ -459,6 +480,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         conn_string: str,
         *,
         pipeline: bool = False,
+        pool_config: Optional[PoolConfig] = None,
         embedding: Optional[EmbeddingConfig] = None,
     ) -> Iterator["PostgresStore"]:
         """Create a new PostgresStore instance from a connection string.
@@ -466,19 +488,41 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         Args:
             conn_string (str): The Postgres connection info string.
             pipeline (bool): whether to use Pipeline
+            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.
             embedding (Optional[EmbeddingConfig]): The embedding config.
 
         Returns:
             PostgresStore: A new PostgresStore 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=pipe, embedding=embedding)
-            else:
-                yield cls(conn, embedding=embedding)
+        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, embedding=embedding)
+        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, embedding=embedding)
+                else:
+                    yield cls(conn, embedding=embedding)
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -504,14 +548,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                 # 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:
+                    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:
+                    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:
@@ -617,7 +665,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         for (idx, _), (query, params) in zip(search_ops, queries):
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
-            items = [
+            results[idx] = [
                 _row_to_item(
                     _decode_ns_bytes(row["prefix"]),
                     row,
@@ -626,7 +674,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                 )
                 for row in rows
             ]
-            results[idx] = items
 
     def _batch_list_namespaces_ops(
         self,
@@ -726,7 +773,7 @@ def _row_to_item(
     row: Row,
     *,
     loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
-    cls: Union[Type[SearchItem], Type[Item]] = Item,
+    cls: Union[type[SearchItem], type[Item]] = Item,
 ) -> Union[Item, SearchItem]:
     """Convert a row from the database into an Item.
 
@@ -796,7 +843,7 @@ def _tokenize_path(path: str) -> list[str]:
         return []
 
     tokens = []
-    current: List[str] = []
+    current: list[str] = []
     i = 0
     while i < len(path):
         char = path[i]

libs/checkpoint-postgres/tests/conftest.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
 from psycopg import AsyncConnection

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,7 +1,7 @@
 # type: ignore
 import sys
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
 from conftest import DEFAULT_URI  # type: ignore
@@ -43,9 +43,8 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
                 yield store
         elif request.param == "pool":
             async with AsyncPostgresStore.from_conn_string(
-                conn_string, use_pool=True, max_size=10
+                conn_string, pool_config={"min_size": 1, "max_size": 10}
             ) as store:
-                await store.setup()
                 yield store
         else:  # default
             async with AsyncPostgresStore.from_conn_string(conn_string) as store:

libs/checkpoint-postgres/tests/test_store.py

@@ -7,7 +7,6 @@ import pytest
 from conftest import DEFAULT_URI  # type: ignore
 from langchain_core.embeddings import Embeddings
 from psycopg import Connection
-from psycopg_pool import ConnectionPool
 from utils import CharacterEmbeddings
 
 from langgraph.store.base import (
@@ -45,10 +44,9 @@ def store(request) -> PostgresStore:
             with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
                 yield store
         elif request.param == "pool":
-            with ConnectionPool(
-                conn_string, max_size=10, kwargs={"autocommit": True}
-            ) as pool:
-                store = PostgresStore(pool)
+            with PostgresStore.from_conn_string(
+                conn_string, pool_config={"min_size": 1, "max_size": 10}
+            ) as store:
                 yield store
         else:  # default
             with PostgresStore.from_conn_string(conn_string) as store:
@@ -566,7 +564,7 @@ def test_extract_text_by_path():
     assert _extract_text_by_path(nested_data, "empty_dict") == ["{}"]
 
     zeros = _extract_text_by_path(nested_data, "zeros[*]")
-    assert set(zeros) == {"0", "0.0", "0"}
+    assert set(zeros) == {"0", "0.0"}
 
     assert _extract_text_by_path(nested_data, "items[].value") == []
     assert _extract_text_by_path(nested_data, "items[abc].value") == []