Test parallel store writes

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -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,29 +17,17 @@
     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)}")
-
 
 class PostgresSaver(BasePostgresSaver):
     lock: threading.Lock
 
     def __init__(
         self,
-        conn: Conn,
+        conn: _internal.Conn,
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -373,7 +361,7 @@ def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
                 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

libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py

@@ -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

@@ -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

@@ -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,31 +17,17 @@
     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)}")
-
 
 class AsyncPostgresSaver(BasePostgresSaver):
     lock: asyncio.Lock
 
     def __init__(
         self,
-        conn: Conn,
+        conn: _ainternal.Conn,
         pipe: Optional[AsyncPipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -331,7 +317,7 @@ async def _cursor(
                 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