From 98935e1ffd8ab82d5587b90fe154ecf92fa61121 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 25 Nov 2024 12:19:52 -0800 Subject: [PATCH] fix: Fix race condition in PostgresSaver (#2494) Signed-off-by: Tyler Ball Co-authored-by: Phoenix Logan Co-authored-by: Tyler Ball <2481463+tyler-ball@users.noreply.github.com> --- .../langgraph/checkpoint/postgres/__init__.py | 26 +- .../checkpoint/postgres/_ainternal.py | 23 + .../checkpoint/postgres/_internal.py | 21 + .../langgraph/checkpoint/postgres/aio.py | 65 ++- .../langgraph/store/postgres/aio.py | 254 +++++---- .../langgraph/store/postgres/base.py | 217 +++++--- libs/checkpoint-postgres/tests/conftest.py | 6 +- .../tests/test_async_store.py | 223 +++----- libs/checkpoint-postgres/tests/test_store.py | 501 +++++++----------- libs/langgraph/Makefile | 7 +- libs/langgraph/tests/conftest.py | 109 +++- libs/langgraph/tests/test_pregel.py | 115 ++-- libs/langgraph/tests/test_pregel_async.py | 90 +++- 13 files changed, 906 insertions(+), 751 deletions(-) create mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py create mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index b8138a945..2107b05a3 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -1,6 +1,6 @@ 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 Capabilities, Connection, Cursor, Pipeline @@ -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: @@ -73,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. @@ -373,7 +363,7 @@ class PostgresSaver(BasePostgresSaver): 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 _get_connection(self.conn) as conn: + 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 @@ -403,4 +393,4 @@ class PostgresSaver(BasePostgresSaver): yield cur -__all__ = ["PostgresSaver", "Conn"] +__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py new file mode 100644 index 000000000..a0b8b10f5 --- /dev/null +++ b/libs/checkpoint-postgres/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)}") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py new file mode 100644 index 000000000..b703262f2 --- /dev/null +++ b/libs/checkpoint-postgres/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)}") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5b67e4ca9..589520efc 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,6 +1,6 @@ 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, Capabilities @@ -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: @@ -80,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. @@ -157,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"]), ) @@ -216,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"]), ) @@ -331,7 +323,7 @@ class AsyncPostgresSaver(BasePostgresSaver): 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 _get_connection(self.conn) as conn: + 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 @@ -467,3 +459,6 @@ class AsyncPostgresSaver(BasePostgresSaver): return asyncio.run_coroutine_threadsafe( self.aput_writes(config, writes, task_id), self.loop ).result() + + +__all__ = ["AsyncPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index dda7321d0..578523052 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -13,14 +13,17 @@ from typing import ( ) import orjson -from psycopg import AsyncConnection, AsyncCursor +from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +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, @@ -30,81 +33,88 @@ from langgraph.store.postgres.base import ( logger = logging.getLogger(__name__) -class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]): - __slots__ = ("_deserializer",) +class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]): + __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline") def __init__( self, - conn: AsyncConnection[Any], + 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 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) + 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 - def batch(self, ops: Iterable[Op]) -> list[Result]: - return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result() + 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: - 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: @@ -119,26 +129,21 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio 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: - 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], + cur: AsyncCursor[DictRow], ) -> 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( @@ -152,37 +157,103 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio self, list_ops: Sequence[tuple[int, ListNamespacesOp]], results: list[Result], + cur: AsyncCursor[DictRow], ) -> 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 + @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. """ - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - yield cls(conn=conn) + 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. @@ -191,28 +262,33 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio 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 + 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" ) - """ - ) - 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,)) + 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() diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 8bd2b8279..3bb343b0c 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/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 ( diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index 49061b98c..56d199812 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/libs/checkpoint-postgres/tests/conftest.py @@ -24,6 +24,10 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None: await conn.execute("DELETE FROM checkpoint_blobs") await conn.execute("DELETE FROM checkpoint_writes") await conn.execute("DELETE FROM checkpoint_migrations") - await conn.execute("DELETE FROM store_migrations") + except UndefinedTable: + pass + try: + await conn.execute("DELETE FROM store_migrations") + await conn.execute("DELETE FROM store") except UndefinedTable: pass diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index e7a7b31a1..71aaa4e36 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -1,114 +1,76 @@ # type: ignore +import sys import uuid -from datetime import datetime -from typing import Any -from unittest.mock import AsyncMock, MagicMock +from typing import AsyncIterator import pytest from conftest import DEFAULT_URI # type: ignore +from psycopg import AsyncConnection from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp from langgraph.store.postgres import AsyncPostgresStore -class MockAsyncCursor: - def __init__(self, fetch_result: Any) -> None: - self.fetch_result = fetch_result - self.execute = AsyncMock() - self.fetchall = AsyncMock(return_value=self.fetch_result) +@pytest.fixture(scope="function", params=["default", "pipe", "pool"]) +async def store(request) -> AsyncIterator[AsyncPostgresStore]: + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid.uuid4().hex[:16]}" + uri_parts = DEFAULT_URI.split("/") + uri_base = "/".join(uri_parts[:-1]) + query_params = "" + if "?" in uri_parts[-1]: + db_name, query_params = uri_parts[-1].split("?", 1) + query_params = "?" + query_params -class MockAsyncConnection: - def __init__(self) -> None: - self.cursor = MagicMock() - self.pipeline = MagicMock( - return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock()) - ) + conn_string = f"{uri_base}/{database}{query_params}" + admin_conn_string = DEFAULT_URI + async with await AsyncConnection.connect( + admin_conn_string, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string(conn_string) as store: + await store.setup() -@pytest.fixture -def mock_connection() -> MockAsyncConnection: - return MockAsyncConnection() - - -@pytest.fixture -async def store(mock_connection: MockAsyncConnection) -> AsyncPostgresStore: - return AsyncPostgresStore(mock_connection) + if request.param == "pipe": + async with AsyncPostgresStore.from_conn_string( + conn_string, pipeline=True + ) as store: + yield store + elif request.param == "pool": + async with AsyncPostgresStore.from_conn_string( + conn_string, pool_config={"min_size": 1, "max_size": 10} + ) as store: + yield store + else: # default + async with AsyncPostgresStore.from_conn_string(conn_string) as store: + yield store + finally: + async with await AsyncConnection.connect( + admin_conn_string, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") async def test_abatch_order(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_get_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_search_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - ] - ) - mock_list_namespaces_cursor = MockAsyncCursor( - [ - {"truncated_prefix": b"\x01test"}, - ] - ) - - failures = [] - - def cursor_side_effect(binary: bool = False) -> Any: - cursor = MagicMock() - - async def execute_side_effect(query: str, *params: Any) -> None: - # My super sophisticated database. - if "SELECT prefix, key," in query: - cursor.fetchall = mock_search_cursor.fetchall - elif "SELECT DISTINCT ON (truncated_prefix)" in query: - cursor.fetchall = mock_list_namespaces_cursor.fetchall - elif "WHERE prefix = %s AND key" in query: - cursor.fetchall = mock_get_cursor.fetchall - elif "INSERT INTO " in query: - pass - else: - e = ValueError(f"Unmatched query: {query}") - failures.append(e) - raise e - - cursor.execute = AsyncMock(side_effect=execute_side_effect) - return cursor - - mock_connection.cursor.side_effect = cursor_side_effect # type: ignore + # Setup test data + await store.aput(("test", "foo"), "key1", {"data": "value1"}) + await store.aput(("test", "bar"), "key2", {"data": "value2"}) ops = [ - GetOp(namespace=("test",), key="key1"), - PutOp(namespace=("test",), key="key2", value={"data": "value2"}), + GetOp(namespace=("test", "foo"), key="key1"), + PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}), SearchOp( namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 ), ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), GetOp(namespace=("test",), key="key3"), ] + results = await store.abatch(ops) - assert not failures assert len(results) == 5 assert isinstance(results[0], Item) assert isinstance(results[0].value, dict) @@ -118,27 +80,29 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None: assert isinstance(results[2], list) assert len(results[2]) == 1 assert isinstance(results[3], list) - assert results[3] == [("test",)] + assert ("test", "foo") in results[3] and ("test", "bar") in results[3] assert results[4] is None ops_reordered = [ SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), - GetOp(namespace=("test",), key="key2"), + GetOp(namespace=("test", "bar"), key="key2"), ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0), PutOp(namespace=("test",), key="key3", value={"data": "value3"}), - GetOp(namespace=("test",), key="key1"), + GetOp(namespace=("test", "foo"), key="key1"), ] results_reordered = await store.abatch(ops_reordered) - assert not failures assert len(results_reordered) == 5 assert isinstance(results_reordered[0], list) - assert len(results_reordered[0]) == 1 + assert len(results_reordered[0]) == 2 assert isinstance(results_reordered[1], Item) assert results_reordered[1].value == {"data": "value2"} assert results_reordered[1].key == "key2" assert isinstance(results_reordered[2], list) - assert results_reordered[2] == [("test",)] + assert ("test", "foo") in results_reordered[2] and ( + "test", + "bar", + ) in results_reordered[2] assert results_reordered[3] is None assert isinstance(results_reordered[4], Item) assert results_reordered[4].value == {"data": "value1"} @@ -146,26 +110,9 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None: async def test_batch_get_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test",), "key1", {"data": "value1"}) + await store.aput(("test",), "key2", {"data": "value2"}) ops = [ GetOp(namespace=("test",), key="key1"), @@ -184,10 +131,6 @@ async def test_batch_get_ops(store: AsyncPostgresStore) -> None: async def test_batch_put_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor([]) - mock_connection.cursor.return_value = mock_cursor - ops = [ PutOp(namespace=("test",), key="key1", value={"data": "value1"}), PutOp(namespace=("test",), key="key2", value={"data": "value2"}), @@ -198,30 +141,16 @@ async def test_batch_put_ops(store: AsyncPostgresStore) -> None: assert len(results) == 3 assert all(result is None for result in results) - assert mock_cursor.execute.call_count == 2 + + # Verify the puts worked + items = await store.asearch(["test"], limit=10) + assert len(items) == 2 # key3 had None value so wasn't stored async def test_batch_search_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test", "foo"), "key1", {"data": "value1"}) + await store.aput(("test", "bar"), "key2", {"data": "value2"}) ops = [ SearchOp( @@ -233,29 +162,23 @@ async def test_batch_search_ops(store: AsyncPostgresStore) -> None: results = await store.abatch(ops) assert len(results) == 2 - assert len(results[0]) == 2 - assert len(results[1]) == 2 + assert len(results[0]) == 1 # Filtered results + assert len(results[1]) == 2 # All results async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - {"truncated_prefix": b"\x01test.namespace1"}, - {"truncated_prefix": b"\x01test.namespace2"}, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test", "namespace1"), "key1", {"data": "value1"}) + await store.aput(("test", "namespace2"), "key2", {"data": "value2"}) ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)] results = await store.abatch(ops) assert len(results) == 1 - assert results[0] == [("test", "namespace1"), ("test", "namespace2")] - - -# The following use the actual DB connection + assert len(results[0]) == 2 + assert ("test", "namespace1") in results[0] + assert ("test", "namespace2") in results[0] class TestAsyncPostgresStore: diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index add9fb1c5..645c37e9f 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -1,174 +1,118 @@ # type: ignore -import uuid -from datetime import datetime -from typing import Any -from unittest.mock import MagicMock + +from uuid import uuid4 import pytest from conftest import DEFAULT_URI # type: ignore +from psycopg import Connection -from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp +from langgraph.store.base import ( + GetOp, + Item, + ListNamespacesOp, + MatchCondition, + PutOp, + SearchOp, +) from langgraph.store.postgres import PostgresStore -class MockCursor: - def __init__(self, fetch_result: Any) -> None: - self.fetch_result = fetch_result - self.execute = MagicMock() - self.fetchall = MagicMock(return_value=self.fetch_result) +@pytest.fixture(scope="function", params=["default", "pipe", "pool"]) +def store(request) -> PostgresStore: + database = f"test_{uuid4().hex[:16]}" + uri_parts = DEFAULT_URI.split("/") + uri_base = "/".join(uri_parts[:-1]) + query_params = "" + if "?" in uri_parts[-1]: + db_name, query_params = uri_parts[-1].split("?", 1) + query_params = "?" + query_params + conn_string = f"{uri_base}/{database}{query_params}" + admin_conn_string = DEFAULT_URI -class MockConnection: - def __init__(self) -> None: - self.cursor = MagicMock() - self.pipeline = MagicMock() + with Connection.connect(admin_conn_string, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + with PostgresStore.from_conn_string(conn_string) as store: + store.setup() - -@pytest.fixture -def mock_connection() -> MockConnection: - return MockConnection() - - -@pytest.fixture -def store(mock_connection: MockConnection) -> PostgresStore: - return PostgresStore(mock_connection) + if request.param == "pipe": + with PostgresStore.from_conn_string(conn_string, pipeline=True) as store: + yield store + elif request.param == "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: + yield store + finally: + with Connection.connect(admin_conn_string, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") def test_batch_order(store: PostgresStore) -> None: - mock_connection = store.conn - mock_get_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_search_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - ] - ) - mock_list_namespaces_cursor = MockCursor( - [ - {"truncated_prefix": b"\x01test"}, - ] - ) - - failures = [] - - def cursor_side_effect(binary: bool = False) -> Any: - cursor = MagicMock() - - def execute_side_effect(query: str, *params: Any) -> None: - # My super sophisticated database. - if "SELECT prefix, key, value" in query: - cursor.fetchall = mock_search_cursor.fetchall - elif "SELECT DISTINCT ON (truncated_prefix)" in query: - cursor.fetchall = mock_list_namespaces_cursor.fetchall - elif "WHERE prefix = %s AND key" in query: - cursor.fetchall = mock_get_cursor.fetchall - elif "INSERT INTO " in query: - pass - else: - e = ValueError(f"Unmatched query: {query}") - failures.append(e) - raise e - - cursor.execute = MagicMock(side_effect=execute_side_effect) - return cursor - - mock_connection.cursor.side_effect = cursor_side_effect + # Setup test data + store.put(("test", "foo"), "key1", {"data": "value1"}) + store.put(("test", "bar"), "key2", {"data": "value2"}) ops = [ - GetOp(namespace=("test",), key="key1"), - PutOp(namespace=("test",), key="key2", value={"data": "value2"}), + GetOp(namespace=("test", "foo"), key="key1"), + PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}), SearchOp( namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 ), ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), GetOp(namespace=("test",), key="key3"), ] + results = store.batch(ops) - assert not failures assert len(results) == 5 assert isinstance(results[0], Item) assert isinstance(results[0].value, dict) assert results[0].value == {"data": "value1"} assert results[0].key == "key1" - assert results[1] is None + assert results[1] is None # Put operation returns None assert isinstance(results[2], list) assert len(results[2]) == 1 assert isinstance(results[3], list) - assert results[3] == [("test",)] - assert results[4] is None + assert len(results[3]) > 0 # Should contain at least our test namespaces + assert results[4] is None # Non-existent key returns None + # Test reordered operations ops_reordered = [ SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), - GetOp(namespace=("test",), key="key2"), + GetOp(namespace=("test", "bar"), key="key2"), ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0), PutOp(namespace=("test",), key="key3", value={"data": "value3"}), - GetOp(namespace=("test",), key="key1"), + GetOp(namespace=("test", "foo"), key="key1"), ] results_reordered = store.batch(ops_reordered) - assert not failures assert len(results_reordered) == 5 assert isinstance(results_reordered[0], list) - assert len(results_reordered[0]) == 1 + assert len(results_reordered[0]) >= 2 # Should find at least our two test items assert isinstance(results_reordered[1], Item) assert results_reordered[1].value == {"data": "value2"} assert results_reordered[1].key == "key2" assert isinstance(results_reordered[2], list) - assert results_reordered[2] == [("test",)] - assert results_reordered[3] is None + assert len(results_reordered[2]) > 0 + assert results_reordered[3] is None # Put operation returns None assert isinstance(results_reordered[4], Item) assert results_reordered[4].value == {"data": "value1"} assert results_reordered[4].key == "key1" def test_batch_get_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + store.put(("test",), "key1", {"data": "value1"}) + store.put(("test",), "key2", {"data": "value2"}) ops = [ GetOp(namespace=("test",), key="key1"), GetOp(namespace=("test",), key="key2"), - GetOp(namespace=("test",), key="key3"), + GetOp(namespace=("test",), key="key3"), # Non-existent key ] results = store.batch(ops) @@ -182,75 +126,90 @@ def test_batch_get_ops(store: PostgresStore) -> None: def test_batch_put_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor([]) - mock_connection.cursor.return_value = mock_cursor - ops = [ PutOp(namespace=("test",), key="key1", value={"data": "value1"}), PutOp(namespace=("test",), key="key2", value={"data": "value2"}), - PutOp(namespace=("test",), key="key3", value=None), + PutOp(namespace=("test",), key="key3", value=None), # Delete operation ] results = store.batch(ops) - assert len(results) == 3 assert all(result is None for result in results) - assert mock_cursor.execute.call_count == 2 + + # Verify the puts worked + item1 = store.get(("test",), "key1") + item2 = store.get(("test",), "key2") + item3 = store.get(("test",), "key3") + + assert item1 and item1.value == {"data": "value1"} + assert item2 and item2.value == {"data": "value2"} + assert item3 is None def test_batch_search_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + test_data = [ + (("test", "foo"), "key1", {"data": "value1", "tag": "a"}), + (("test", "bar"), "key2", {"data": "value2", "tag": "a"}), + (("test", "baz"), "key3", {"data": "value3", "tag": "b"}), + ] + for namespace, key, value in test_data: + store.put(namespace, key, value) ops = [ - SearchOp( - namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 - ), - SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), + SearchOp(namespace_prefix=("test",), filter={"tag": "a"}, limit=10, offset=0), + SearchOp(namespace_prefix=("test",), filter=None, limit=2, offset=0), + SearchOp(namespace_prefix=("test", "foo"), filter=None, limit=10, offset=0), ] results = store.batch(ops) + assert len(results) == 3 - assert len(results) == 2 + # First search should find items with tag "a" assert len(results[0]) == 2 + assert all(item.value["tag"] == "a" for item in results[0]) + + # Second search should return first 2 items assert len(results[1]) == 2 + # Third search should only find items in test/foo namespace + assert len(results[2]) == 1 + assert results[2][0].namespace == ("test", "foo") + def test_batch_list_namespaces_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - {"truncated_prefix": b"\x01test.namespace1"}, - {"truncated_prefix": b"\x01test.namespace2"}, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data with various namespaces + test_data = [ + (("test", "documents", "public"), "doc1", {"content": "public doc"}), + (("test", "documents", "private"), "doc2", {"content": "private doc"}), + (("test", "images", "public"), "img1", {"content": "public image"}), + (("prod", "documents", "public"), "doc3", {"content": "prod doc"}), + ] + for namespace, key, value in test_data: + store.put(namespace, key, value) - ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)] + ops = [ + ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), + ListNamespacesOp(match_conditions=None, max_depth=2, limit=10, offset=0), + ListNamespacesOp( + match_conditions=[MatchCondition("suffix", "public")], + max_depth=None, + limit=10, + offset=0, + ), + ] results = store.batch(ops) + assert len(results) == 3 - assert len(results) == 1 - assert results[0] == [("test", "namespace1"), ("test", "namespace2")] + # First operation should list all namespaces + assert len(results[0]) == len(test_data) + + # Second operation should only return namespaces up to depth 2 + assert all(len(ns) <= 2 for ns in results[1]) + + # Third operation should only return namespaces ending with "public" + assert all(ns[-1] == "public" for ns in results[2]) class TestPostgresStore: @@ -273,195 +232,111 @@ class TestPostgresStore: assert item.key == item_id assert item.value == item_value - updated_value = { - "title": "Updated Test Document", - "content": "Hello, LangGraph!", - } + # Test update + updated_value = {"title": "Updated Document", "content": "Hello, Updated!"} store.put(namespace, item_id, updated_value) updated_item = store.get(namespace, item_id) assert updated_item.value == updated_value assert updated_item.updated_at > item.updated_at + + # Test get from non-existent namespace different_namespace = ("test", "other_documents") item_in_different_namespace = store.get(different_namespace, item_id) assert item_in_different_namespace is None - new_item_id = "doc2" - new_item_value = {"title": "Another Document", "content": "Greetings!"} - store.put(namespace, new_item_id, new_item_value) - - search_results = store.search(["test"], limit=10) - items = search_results - assert len(items) == 2 - assert any(item.key == item_id for item in items) - assert any(item.key == new_item_id for item in items) - - namespaces = store.list_namespaces(prefix=["test"]) - assert ("test", "documents") in namespaces - + # Test delete store.delete(namespace, item_id) - store.delete(namespace, new_item_id) deleted_item = store.get(namespace, item_id) assert deleted_item is None - deleted_item = store.get(namespace, new_item_id) - assert deleted_item is None - - empty_search_results = store.search(["test"], limit=10) - assert len(empty_search_results) == 0 - def test_list_namespaces(self) -> None: with PostgresStore.from_conn_string(DEFAULT_URI) as store: - test_pref = str(uuid.uuid4()) + # Create test data with various namespaces test_namespaces = [ - (test_pref, "test", "documents", "public", test_pref), - (test_pref, "test", "documents", "private", test_pref), - (test_pref, "test", "images", "public", test_pref), - (test_pref, "test", "images", "private", test_pref), - (test_pref, "prod", "documents", "public", test_pref), - ( - test_pref, - "prod", - "documents", - "some", - "nesting", - "public", - test_pref, - ), - (test_pref, "prod", "documents", "private", test_pref), + ("test", "documents", "public"), + ("test", "documents", "private"), + ("test", "images", "public"), + ("test", "images", "private"), + ("prod", "documents", "public"), + ("prod", "documents", "private"), ] + # Insert test data for namespace in test_namespaces: store.put(namespace, "dummy", {"content": "dummy"}) - prefix_result = store.list_namespaces(prefix=[test_pref, "test"]) - assert len(prefix_result) == 4 - assert all([ns[1] == "test" for ns in prefix_result]) + # Test listing with various filters + all_namespaces = store.list_namespaces() + assert len(all_namespaces) == len(test_namespaces) - specific_prefix_result = store.list_namespaces( - prefix=[test_pref, "test", "documents"] - ) - assert len(specific_prefix_result) == 2 - assert all( - [ns[1:3] == ("test", "documents") for ns in specific_prefix_result] - ) + # Test prefix filtering + test_prefix_namespaces = store.list_namespaces(prefix=["test"]) + assert len(test_prefix_namespaces) == 4 + assert all(ns[0] == "test" for ns in test_prefix_namespaces) - suffix_result = store.list_namespaces(suffix=["public", test_pref]) - assert len(suffix_result) == 4 - assert all(ns[-2] == "public" for ns in suffix_result) + # Test suffix filtering + public_namespaces = store.list_namespaces(suffix=["public"]) + assert len(public_namespaces) == 3 + assert all(ns[-1] == "public" for ns in public_namespaces) - prefix_suffix_result = store.list_namespaces( - prefix=[test_pref, "test"], suffix=["public", test_pref] - ) - assert len(prefix_suffix_result) == 2 - assert all( - ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result - ) + # Test max depth + depth_2_namespaces = store.list_namespaces(max_depth=2) + assert all(len(ns) <= 2 for ns in depth_2_namespaces) - wildcard_prefix_result = store.list_namespaces( - prefix=[test_pref, "*", "documents"] - ) - assert len(wildcard_prefix_result) == 5 - assert all(ns[2] == "documents" for ns in wildcard_prefix_result) - - wildcard_suffix_result = store.list_namespaces( - suffix=["*", "public", test_pref] - ) - assert len(wildcard_suffix_result) == 4 - assert all(ns[-2] == "public" for ns in wildcard_suffix_result) - wildcard_single = store.list_namespaces( - suffix=["some", "*", "public", test_pref] - ) - assert len(wildcard_single) == 1 - assert wildcard_single[0] == ( - test_pref, - "prod", - "documents", - "some", - "nesting", - "public", - test_pref, - ) - - max_depth_result = store.list_namespaces(max_depth=3) - assert all([len(ns) <= 3 for ns in max_depth_result]) - - max_depth_result = store.list_namespaces( - max_depth=4, prefix=[test_pref, "*", "documents"] - ) - assert ( - len(set(tuple(res) for res in max_depth_result)) - == len(max_depth_result) - == 5 - ) - - limit_result = store.list_namespaces(prefix=[test_pref], limit=3) - assert len(limit_result) == 3 - - offset_result = store.list_namespaces(prefix=[test_pref], offset=3) - assert len(offset_result) == len(test_namespaces) - 3 - - empty_prefix_result = store.list_namespaces(prefix=[test_pref]) - assert len(empty_prefix_result) == len(test_namespaces) - assert set(tuple(ns) for ns in empty_prefix_result) == set( - tuple(ns) for ns in test_namespaces - ) + # Test pagination + paginated_namespaces = store.list_namespaces(limit=3) + assert len(paginated_namespaces) == 3 + # Cleanup for namespace in test_namespaces: store.delete(namespace, "dummy") - def test_search(self): + def test_search(self) -> None: with PostgresStore.from_conn_string(DEFAULT_URI) as store: - test_namespaces = [ - ("test_search", "documents", "user1"), - ("test_search", "documents", "user2"), - ("test_search", "reports", "department1"), - ("test_search", "reports", "department2"), - ] - test_items = [ - {"title": "Doc 1", "author": "John Doe", "tags": ["important"]}, - {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]}, - {"title": "Report A", "author": "John Doe", "tags": ["final"]}, - {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]}, + # Create test data + test_data = [ + ( + ("test", "docs"), + "doc1", + {"title": "First Doc", "author": "Alice", "tags": ["important"]}, + ), + ( + ("test", "docs"), + "doc2", + {"title": "Second Doc", "author": "Bob", "tags": ["draft"]}, + ), + ( + ("test", "images"), + "img1", + {"title": "Image 1", "author": "Alice", "tags": ["final"]}, + ), ] - for namespace, item in zip(test_namespaces, test_items): - store.put(namespace, f"item_{namespace[-1]}", item) + for namespace, key, value in test_data: + store.put(namespace, key, value) - docs_result = store.search(["test_search", "documents"]) - assert len(docs_result) == 2 - assert all( - [item.namespace[1] == "documents" for item in docs_result] - ), docs_result + # Test basic search + all_items = store.search(["test"]) + assert len(all_items) == 3 - reports_result = store.search(["test_search", "reports"]) - assert len(reports_result) == 2 - assert all(item.namespace[1] == "reports" for item in reports_result) + # Test namespace filtering + docs_items = store.search(["test", "docs"]) + assert len(docs_items) == 2 + assert all(item.namespace == ("test", "docs") for item in docs_items) - limited_result = store.search(["test_search"], limit=2) - assert len(limited_result) == 2 - offset_result = store.search(["test_search"]) - assert len(offset_result) == 4 + # Test value filtering + alice_items = store.search(["test"], filter={"author": "Alice"}) + assert len(alice_items) == 2 + assert all(item.value["author"] == "Alice" for item in alice_items) - offset_result = store.search(["test_search"], offset=2) - assert len(offset_result) == 2 - assert all(item not in limited_result for item in offset_result) + # Test pagination + paginated_items = store.search(["test"], limit=2) + assert len(paginated_items) == 2 - john_doe_result = store.search( - ["test_search"], filter={"author": "John Doe"} - ) - assert len(john_doe_result) == 2 - assert all(item.value["author"] == "John Doe" for item in john_doe_result) + offset_items = store.search(["test"], offset=2) + assert len(offset_items) == 1 - draft_result = store.search(["test_search"], filter={"tags": ["draft"]}) - assert len(draft_result) == 2 - assert all("draft" in item.value["tags"] for item in draft_result) - - page1 = store.search(["test_search"], limit=2, offset=0) - page2 = store.search(["test_search"], limit=2, offset=2) - all_items = page1 + page2 - assert len(all_items) == 4 - assert len(set(item.key for item in all_items)) == 4 - - for namespace in test_namespaces: - store.delete(namespace, f"item_{namespace[-1]}") + # Cleanup + for namespace, key, _ in test_data: + store.delete(namespace, key) diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 2aacf6db8..43d0c7afe 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -48,8 +48,13 @@ test: make stop-postgres; \ exit $$EXIT_CODE +WORKERS ?= auto +XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,) +MAXFAIL ?= +MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),) + test_watch: - make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index eae7694ff..0381206e3 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -272,6 +272,54 @@ async def _store_postgres_aio(): await conn.execute(f"DROP DATABASE {database}") +@asynccontextmanager +async def _store_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() # Run in its own transaction + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, + pool_config={"max_size": 10}, + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + @asynccontextmanager async def _store_duckdb_aio(): async with AsyncDuckDBStore.from_conn_string(":memory:") as store: @@ -296,6 +344,45 @@ def store_postgres(): conn.execute(f"DROP DATABASE {database}") +@pytest.fixture(scope="function") +def store_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() # Run in its own transaction + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture(scope="function") +def store_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10} + ) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + @pytest.fixture(scope="function") def store_duckdb(): with DuckDBStore.from_conn_string(":memory:") as store: @@ -317,6 +404,12 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: elif store_name == "postgres_aio": async with _store_postgres_aio() as store: yield store + elif store_name == "postgres_aio_pipe": + async with _store_postgres_aio_pipe() as store: + yield store + elif store_name == "postgres_aio_pool": + async with _store_postgres_aio_pool() as store: + yield store elif store_name == "duckdb_aio": async with _store_duckdb_aio() as store: yield store @@ -342,5 +435,17 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ *ALL_CHECKPOINTERS_ASYNC, None, ] -ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"] -ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"] +ALL_STORES_SYNC = [ + "in_memory", + "postgres", + "postgres_pipe", + "postgres_pool", + "duckdb", +] +ALL_STORES_ASYNC = [ + "in_memory", + "postgres_aio", + "postgres_aio_pipe", + "postgres_aio_pool", + "duckdb_aio", +] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fda881b15..fcee1fa20 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,5 +1,6 @@ import enum import json +import logging import operator import re import time @@ -67,16 +68,9 @@ from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInt from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import ( - create_tool_calling_executor, -) +from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import ( - Channel, - GraphRecursionError, - Pregel, - StateSnapshot, -) +from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore @@ -104,6 +98,8 @@ from tests.messages import ( _AnyIdToolMessage, ) +logger = logging.getLogger(__name__) + # define these objects to avoid importing langchain_core.agents # and therefore avoid relying on core Pydantic version @@ -6628,11 +6624,7 @@ def test_message_graph( from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) - from langchain_core.messages import ( - AIMessage, - BaseMessage, - HumanMessage, - ) + from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool @@ -13937,50 +13929,75 @@ def test_store_injected( doc_id = str(uuid.uuid4()) doc = {"some-key": "this-is-a-val"} - - def node(input: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - store.put( - ("foo", "bar"), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": input["count"], - }, - ) - return {"count": 1} - - builder = StateGraph(State) - builder.add_node("node", node) - builder.add_edge("__start__", "node") - graph = builder.compile(store=the_store, checkpointer=checkpointer) - + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") thread_1 = str(uuid.uuid4()) - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) - assert result == {"count": 1} - returned_doc = the_store.get(("foo", "bar"), doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(the_store.search(("foo", "bar"))) == 1 - - # Check update on existing thread - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) - assert result == {"count": 2} - returned_doc = the_store.get(("foo", "bar"), doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1} - assert len(the_store.search(("foo", "bar"))) == 1 - thread_2 = str(uuid.uuid4()) + class Node: + def __init__(self, i: Optional[int] = None): + self.i = i + + def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore): + assert isinstance(store, BaseStore) + store.put( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar"), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} + + builder = StateGraph(State) + builder.add_node("node", Node()) + builder.add_edge("__start__", "node") + N = 500 + M = 1 + if "duckdb" in store_name: + logger.warning( + "DuckDB store implementation has a known issue that does not" + " support concurrent writes, so we're reducing the test scope" + ) + N = M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + + graph = builder.compile(store=the_store, checkpointer=checkpointer) + + results = graph.batch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], + ) + result = results[-1] + assert result == {"count": N + 1} + returned_doc = the_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} + assert len(the_store.search(namespace)) == 1 + # Check results after another turn of the same thread + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) + assert result == {"count": (N + 1) * 2} + returned_doc = the_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len(the_store.search(namespace)) == 1 + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}}) - assert result == {"count": 1} - returned_doc = the_store.get(("foo", "bar"), doc_id).value + assert result == {"count": N + 1} + returned_doc = the_store.get(namespace, doc_id).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc - assert len(the_store.search(("foo", "bar"))) == 1 # still overwriting the same one + assert len(the_store.search(namespace)) == 1 # still overwriting the same one def test_enum_node_names(): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a31e444e1..c5812e896 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1,4 +1,5 @@ import asyncio +import logging import operator import random import re @@ -100,6 +101,8 @@ from tests.messages import ( _AnyIdToolMessage, ) +logger = logging.getLogger(__name__) + pytestmark = pytest.mark.anyio @@ -12272,60 +12275,89 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> doc_id = str(uuid.uuid4()) doc = {"some-key": "this-is-a-val"} + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") + thread_1 = str(uuid.uuid4()) + thread_2 = str(uuid.uuid4()) - async def node(input: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - await store.aput( - ("foo", "bar"), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": input["count"], - }, - ) - return {"count": 1} + class Node: + def __init__(self, i: Optional[int] = None): + self.i = i + + async def __call__( + self, inputs: State, config: RunnableConfig, store: BaseStore + ): + assert isinstance(store, BaseStore) + await store.aput( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar"), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} builder = StateGraph(State) - builder.add_node("node", node) + builder.add_node("node", Node()) builder.add_edge("__start__", "node") + + N = 500 + M = 1 + if "duckdb" in store_name: + logger.warning( + "DuckDB store implementation has a known issue that does not" + " support concurrent writes, so we're reducing the test scope" + ) + N = M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store( store_name ) as the_store: graph = builder.compile(store=the_store, checkpointer=checkpointer) - thread_1 = str(uuid.uuid4()) - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_1}} + # Test batch operations with multiple threads + results = await graph.abatch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], ) - assert result == {"count": 1} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value + result = results[-1] + assert result == {"count": N + 1} + returned_doc = (await the_store.aget(namespace, doc_id)).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len((await the_store.asearch(("foo", "bar")))) == 1 + assert len((await the_store.asearch(namespace))) == 1 - # Check update on existing thread + # Check results after another turn of the same thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_1}} ) - assert result == {"count": 2} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1} - assert len((await the_store.asearch(("foo", "bar")))) == 1 - - thread_2 = str(uuid.uuid4()) + assert result == {"count": (N + 1) * 2} + returned_doc = (await the_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len((await the_store.asearch(namespace))) == 1 + # Test with a different thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_2}} ) - assert result == {"count": 1} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value + assert result == {"count": N + 1} + returned_doc = (await the_store.aget(namespace, doc_id)).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc assert ( - len((await the_store.asearch(("foo", "bar")))) == 1 + len((await the_store.asearch(namespace))) == 1 ) # still overwriting the same one