WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 17:42:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: Fix race condition in PostgresSaver (#2494)

Browse Source 
Signed-off-by: Tyler Ball <tyleraball@gmail.com>
Co-authored-by: Phoenix Logan <plogan@chanzuckerberg.com>
Co-authored-by: Tyler Ball <2481463+tyler-ball@users.noreply.github.com>
This commit is contained in:
William FH
2024-11-25 20:19:52 +00:00
committed by GitHub
co-authored by Phoenix Logan Tyler Ball
parent 328ef609af
commit 98935e1ffd
13 changed files with 906 additions and 751 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+8 -18
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py
+23
View File
@@ -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)}")
libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py
+21
View File
@@ -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)}")
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+30 -35
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/store/postgres/aio.py
+165 -89
View File
@@ -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()
libs/checkpoint-postgres/langgraph/store/postgres/base.py
+153 -64
View File
@@ -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 (
libs/checkpoint-postgres/tests/conftest.py
+5 -1
View File
@@ -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
libs/checkpoint-postgres/tests/test_async_store.py
+73 -150
View File
@@ -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:
libs/checkpoint-postgres/tests/test_store.py
+188 -313
View File
@@ -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)
libs/langgraph/Makefile
+6 -1
View File
@@ -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
libs/langgraph/tests/conftest.py
+107 -2
View File
@@ -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",
]
libs/langgraph/tests/test_pregel.py
+66 -49
View File
@@ -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():
libs/langgraph/tests/test_pregel_async.py
+61 -29
View File
@@ -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