mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 12:47:53 +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
@@ -2,6 +2,10 @@
|
||||
|
||||
Implementation of LangGraph CheckpointSaver that uses Postgres.
|
||||
|
||||
## Dependencies
|
||||
|
||||
By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
|
||||
|
||||
## Usage
|
||||
|
||||
> [!IMPORTANT]
|
||||
|
||||
@@ -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
|
||||
|
||||
Generated
+2
-2
@@ -266,7 +266,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "1.0.1"
|
||||
version = "1.0.6"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
optional = false
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
@@ -969,4 +969,4 @@ watchmedo = ["PyYAML (>=3.10)"]
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
content-hash = "422b6d716b86db072ea3a612287ad20ff5700c18f22d9e9d59cc4e198514519d"
|
||||
content-hash = "cfa417cbaf126fa847242ef81e9af3e4d13e93f25631bb4c007cdc69926ff3b8"
|
||||
|
||||
@@ -12,7 +12,8 @@ packages = [{ include = "langgraph" }]
|
||||
python = "^3.9.0,<4.0"
|
||||
langgraph-checkpoint = "^1.0.1"
|
||||
orjson = ">=3.10.1"
|
||||
psycopg = {extras = ["binary"], version = ">=3.1.19"}
|
||||
psycopg = "^3.0.0"
|
||||
psycopg-pool = "^3.0.0"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ruff = "^0.1.4"
|
||||
@@ -23,7 +24,7 @@ pytest-asyncio = "^0.21.1"
|
||||
pytest-mock = "^3.11.1"
|
||||
pytest-watch = "^4.2.0"
|
||||
mypy = "^1.10.0"
|
||||
psycopg-pool = "^3.2.2"
|
||||
psycopg = {extras = ["binary"], version = ">=3.0.0"}
|
||||
langgraph-checkpoint = {path = "../checkpoint", develop = true}
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
|
||||
Reference in New Issue
Block a user