mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 20:57:52 +02:00
checkpoint-postgres: allow passing pool (#1452)
* checkpoint-postgres: allow passing pool * make psycopg_pool a non-dev dependency * code review * lockfile * move methods * relax requirements, remove binary * add binary to dev dependencies * update readme
This commit is contained in:
Vadym Barda
@@ -1,12 +1,13 @@
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Iterator, List, Optional
|
||||
from typing import Any, Iterator, List, Optional, Union
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import Connection, Cursor, Pipeline
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
ChannelVersions,
|
||||
@@ -21,16 +22,32 @@ from langgraph.checkpoint.postgres.base import (
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
|
||||
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)}")
|
||||
|
||||
|
||||
class PostgresSaver(BasePostgresSaver):
|
||||
lock: threading.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: Connection,
|
||||
conn: Union[Connection, ConnectionPool],
|
||||
pipe: Optional[Pipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, ConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single Connection, not ConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = threading.Lock()
|
||||
@@ -65,22 +82,21 @@ class PostgresSaver(BasePostgresSaver):
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
with self.lock:
|
||||
with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
try:
|
||||
version = cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
).fetchone()["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
cur.execute(migration)
|
||||
cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
self.pipe.sync()
|
||||
with self._cursor() as cur:
|
||||
try:
|
||||
version = cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
).fetchone()["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
cur.execute(migration)
|
||||
cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
self.pipe.sync()
|
||||
|
||||
def list(
|
||||
self,
|
||||
@@ -333,23 +349,24 @@ class PostgresSaver(BasePostgresSaver):
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
|
||||
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 self.conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
with _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
|
||||
with self.lock, conn.pipeline(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
yield cur
|
||||
else:
|
||||
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
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
|
||||
with self.lock, self.conn.pipeline(), self.conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
yield cur
|
||||
else:
|
||||
with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncIterator, Optional
|
||||
from typing import Any, AsyncIterator, Optional, Union
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
ChannelVersions,
|
||||
@@ -19,16 +20,34 @@ from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _get_connection(
|
||||
conn: Union[AsyncConnection, AsyncConnectionPool],
|
||||
) -> AsyncIterator[AsyncConnection]:
|
||||
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)}")
|
||||
|
||||
|
||||
class AsyncPostgresSaver(BasePostgresSaver):
|
||||
lock: asyncio.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: AsyncConnection,
|
||||
conn: Union[AsyncConnection, AsyncConnectionPool],
|
||||
pipe: Optional[AsyncPipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = asyncio.Lock()
|
||||
@@ -63,25 +82,22 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
async with self.lock:
|
||||
async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
try:
|
||||
results = await cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
version = (await results.fetchone())["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute(
|
||||
f"INSERT INTO checkpoint_migrations (v) VALUES ({v})"
|
||||
)
|
||||
if self.pipe:
|
||||
await self.pipe.sync()
|
||||
async with self._cursor() as cur:
|
||||
try:
|
||||
results = await cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
version = (await results.fetchone())["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
await self.pipe.sync()
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
@@ -290,25 +306,26 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
|
||||
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:
|
||||
async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
async with _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:
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
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
|
||||
async with self.lock, conn.pipeline(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
yield cur
|
||||
else:
|
||||
async with self.lock, conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
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
|
||||
async with self.lock, self.conn.pipeline(), self.conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
yield cur
|
||||
else:
|
||||
async with self.lock, self.conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
yield cur
|
||||
|
||||
Reference in New Issue
Block a user