WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 05:07:51 +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/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 (