From c2052d11c25354cfb5effaaba3e015f90eaa831a Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 13 Nov 2024 21:33:42 -0500 Subject: [PATCH] checkpoint-postgres: remove pipeline flag in cursor --- .../langgraph/checkpoint/postgres/__init__.py | 60 ++++++++---------- .../langgraph/checkpoint/postgres/aio.py | 63 ++++++++----------- 2 files changed, 54 insertions(+), 69 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 37e7c2831..7085e107a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -1,5 +1,5 @@ import threading -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from typing import Any, Iterator, Optional, Sequence, Union from langchain_core.runnables import RunnableConfig @@ -308,27 +308,29 @@ class PostgresSaver(BasePostgresSaver): } } - with self._cursor(pipeline=True) as cur: - cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - self._dump_blobs( - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - checkpoint_id, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(metadata), - ), - ) + with self._cursor() as cur: + # Use connection's transaction context manager when not in pipeline mode + with cur.connection.transaction() if self.pipe is None else nullcontext(): + cur.executemany( + self.UPSERT_CHECKPOINT_BLOBS_SQL, + self._dump_blobs( + thread_id, + checkpoint_ns, + copy.pop("channel_values"), # type: ignore[misc] + new_versions, + ), + ) + cur.execute( + self.UPSERT_CHECKPOINTS_SQL, + ( + thread_id, + checkpoint_ns, + checkpoint["id"], + checkpoint_id, + Jsonb(self._dump_checkpoint(copy)), + self._dump_metadata(metadata), + ), + ) return next_config def put_writes( @@ -351,7 +353,7 @@ class PostgresSaver(BasePostgresSaver): if all(w[0] in WRITES_IDX_MAP for w in writes) else self.INSERT_CHECKPOINT_WRITES_SQL ) - with self._cursor(pipeline=True) as cur: + with self._cursor() as cur: cur.executemany( query, self._dump_writes( @@ -364,7 +366,7 @@ class PostgresSaver(BasePostgresSaver): ) @contextmanager - def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: + def _cursor(self) -> Iterator[Cursor[DictRow]]: with _get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently @@ -374,15 +376,7 @@ class PostgresSaver(BasePostgresSaver): 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 + self.pipe.sync() else: with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 59ee7cbf9..eac629605 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,5 +1,5 @@ import asyncio -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, nullcontext from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Union from langchain_core.runnables import RunnableConfig @@ -264,28 +264,29 @@ class AsyncPostgresSaver(BasePostgresSaver): } } - async with self._cursor(pipeline=True) as cur: - await cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - await asyncio.to_thread( - self._dump_blobs, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - await cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - checkpoint_id, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(metadata), - ), - ) + async with self._cursor() as cur: + async with cur.connection.transaction() if self.pipe is None else nullcontext(): + await cur.executemany( + self.UPSERT_CHECKPOINT_BLOBS_SQL, + await asyncio.to_thread( + self._dump_blobs, + thread_id, + checkpoint_ns, + copy.pop("channel_values"), # type: ignore[misc] + new_versions, + ), + ) + await cur.execute( + self.UPSERT_CHECKPOINTS_SQL, + ( + thread_id, + checkpoint_ns, + checkpoint["id"], + checkpoint_id, + Jsonb(self._dump_checkpoint(copy)), + self._dump_metadata(metadata), + ), + ) return next_config async def aput_writes( @@ -316,13 +317,11 @@ class AsyncPostgresSaver(BasePostgresSaver): task_id, writes, ) - async with self._cursor(pipeline=True) as cur: + async with self._cursor() as cur: await cur.executemany(query, params) @asynccontextmanager - async def _cursor( - self, *, pipeline: bool = False - ) -> AsyncIterator[AsyncCursor[DictRow]]: + async def _cursor(self) -> AsyncIterator[AsyncCursor[DictRow]]: async with _get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently @@ -332,15 +331,7 @@ class AsyncPostgresSaver(BasePostgresSaver): 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 + await self.pipe.sync() else: async with self.lock, conn.cursor( binary=True, row_factory=dict_row